[Paper Review] Zhang et al.(2025)_REARANK: Reasoning Re-ranking Agent via Reinforcement Learning

Dobby·2026년 4월 4일

RAG

목록 보기
3/3
post-thumbnail

Introduction

  • 최근 연구들은 리랭킹 작업에 LLM (reranking agents)으로 점수가 아니라 직접적으로 아웃풋을 출력하는 방식을 활용함.
  • 그러나 다음과 같은 도전 과제가 남아있음.
    1. LLM 자체는 랭킹이라는 목적에 최적화되어 있지 않음.
      또한, zero-shot 방식은 순위를 매기는 과정에서 출력한 결과가 좋은지 나쁜지에 대한 피드백을 받아 모델의 가중치를 업데이트하지 못함.
    2. 좋은 성능을 달성하기 위해서는 종종 SFT가 필요하며, 이때 라벨링된 랭킹 데이터를 얻기가 어려움.
    3. LLM이 최종 결정을 내리는 과정이 불투명함.
    4. SOTA reranking agent들은 거대 모델이나 비공개 모델(e.g., GPT-4)에 의존함. 이러한 의존은 고비용과 지연을 초래함.
  • 이 연구에서는 강화학습을 활용한 reasoning listwise reranking agent인 REARANK를 제안한다.
  • 또한 197개의 질문만으로 listwise ranking 데이터셋 증강을 하였다.
  • 추론 과정에서 명시적이고 해석 가능한 추론을 보여준다.
  • 작은 사이즈의 모델에 강력한 추론 능력을 부여하여 저비용으로도 운영할 수 있다.

main contritubutions

  1. 강화학습을 활용하여 추론 기반 listwise 리랭킹 에이전트를 학습시켰다.
  2. GPT-4, Qwen3 등을 능가하는 성능을 달성하였다.
  3. 추론의 전이, 추론 길이와 랭킹 성능 간의 관계를 조사하였다.

Related Work

Large Reasoning Models

  • 강화학습을 통해 긴 CoT 능력을 향상한 deepseek-R1, OpenAI o1, o3
  • 이러한 발전은 LLM이 수학이나 플래닝과 같은 복잡한 도메인에 LLM을 활용하는 것을 가능케 하였다.
  • 이 연구에서는 이러한 진보된 추론 능력을 리랭킹에 활용하고자 한다.

LLMs for Re-ranking

  • pointwise, pairwise, setwise, listwise 방법
  • few-shot/zero-shot 프롬프팅, SFT
  • 이 연구에서는 listwise LLM reranker를 강화학습을 활용하여, out-of-domain, reasoning-intensive 시나리오에서 강건한 성능과 추론 능력을 갖도록 한다.
  • Rank-R1에서도 강화학습으로 setwise reranker를 학습하였다. Rank-R1은 가장 관련 있는 문단의 인덱스 1개를 찾는 것으로 task를 단순화하였고, binary reward signal에 의존하였다. 그러나 이러한 신호는 listwise ranking에서 나타나는, 순서 기반의 풍부한 정보는 결여되어 있다. 이것은 결국 상당한 양의 학습 데이터를 필요로 하게끔 만든다. 게다가, setwise 추론은 한 번에 하나의 문서만 순위매길 수 있어서 매우 비효율적이다. (Fig. 2)
  • 따라서 우리는 listwise reranking 전략을 활용한다.

3. Method

3.1. Listwise Re-ranking Agent

Objective

  • qq: query
  • nn: 검색된 문서 집합 PP의 개수
  • KnK_n: 문서 집합 PP에서 가능한 모든 순열(랭킹)
  • σ\sigma: 특정한 순열(랭킹)
  • pσ(i)p_{\sigma(i)}: σ\sigma에 의해 정의된 순열에, i번째 순위에 위치한 문서
  • rr: 순열의 품질을 측정하는 점수 계산 함수

Sliding window

  • context limit을 피하기 위해 sliding window 활용

  • τ\tau: PP의 초기 랭킹

  • hh: LLM 기반 순열 함수 (Fig. 2)

  • ww: window size

  • 뒤에서부터 각 window마다 hh를 반복 적용함으로써, 최종 top-k 리스트를 얻는다.

  • 윈도우는 w/2w/2 스텝만큼 움직여서 overlap을 만든다. (LLM call은 총 O(2n/w)O(2n/w)

3.2. RL for Listwise Re-ranking

  • Markov Decision Process (MDP) 프레임워크의 관점에서 보면,
    • agent: LLM 정책 πθ\pi_\theta
    • environment: query qq, 초기 랭킹 τ\tau
    • state space: 현재의 query에 대한 랭킹 순열 S=(pτ(1),,pτ(n),q)S = ({p_{\tau(1)}, …, p_{\tau(n)}}, q)
    • action space: agent가 취할 수 있는 행동. LLM이 현재 상태에 적용할 수 있는 permutation function hh의 가능한 집합
    • transition function TT: 에이전트가 현재 상태에서 특정 행동을 취했을 때, 다음 상태로 넘어갈 확률(0에서 1 사이의 값)을 나타내는 함수
    • rr: relevance metric에 기반하여 리랭킹 액션의 품질을 계량화하여 에이전트에게 피드백을 제공하는 보상 함수

GRPO

  • 입력 xx에 대해 출력 시퀀스 그룹 G=o1,o2oGG = {o_1, o_2… o_G}를 뽑고, 각 출력에서 보상 rir_i을 계산하고 그룹 GG 안에서 정규화하여 advantage A^i\hat A_i 를 얻는다.

  • G|G|: 입력 xx와 현재 정책 πθ\pi_\theta로부터 추출된 출력 시퀀스의 개수
  • oi|o_i|: 시퀀스 길이
  • li,tl_{i,t}: 토큰별 loss
  • A^i,t=rimean(r)std(r)\hat A_{i, t} = \frac{r_i - mean(r)}{std(r)}
  • rr: batch 안의 보상 점수들

3.3. Reward Design

Ranking reward

  • 출력 시퀀스 GiG_i는 추론(…<|think|>), 랭킹(…<|answer|>)의 요소들로 구조화되어 있음.
  • 총 보상 rr은 랭킹 성능과 동시에 출력 형식을 지키도록 유도함
  • 보상 신호는 NDCG@10 (실제로 관련성 있는 문서를 선순위에 위치시킬수록 더 높은 점수를 준다.)
    • DCG@10IDCG@10\frac{DCG@10}{IDCG@10}
    • DCG@10=i=110relilog2(i+1)DCG@10 = \sum^{10}_{i=1} \frac{rel_i}{log_2 (i+1)}
    • IDCG@10=i=110reliideallog2(i+1)IDCG@10 = \sum^{10}_{i=1} \frac {rel_{i} ^{ideal}}{log_2 (i+1)}
  • i: LLM이 예측한, 문서의 순위
  • 각 query는 사람 전문가가 관련성을 판단하여 매긴 점수로 라벨링되어 있음.
  • relirel_i는 생성된 순열 내의 요소 pσ(i)p_{\sigma(i)} 에 대한 라벨링된 점수. 즉, 모델이 생성한 순위에서 i번째에 위치한 문서가 인간 전문가로부터 부여 받은 관련성 점수
  • reliidealrel_i ^{ideal}은 해당 query에 대해 가장 관련성이 높은 문서부터 내림차순으로 완벽하게 정렬했을 때 i번째 자리에 와야 할 정답 문서의 관련성 점수
  • 무작위로 추출된 문서 집합 PP가 관련된 문서를 얼마나 포함하는가에 따라 NDCG@10의 가능한 최댓값이 달라질 수 있기 때문에, 우리는 랭킹 보상 rrankr_{rank}를 relative improvement score이라고 정의한다.

- $r^*$: 특정 문서 집합 $P$에서 얻을 수 있는 최고 NDCG@10 점수
- $r_{init}$: 문서 집합 $P$에 대한 초기 랭킹 $\tau$(초기 검색 결과)에 대한 NDCG@10
- $r_{rerank}$: 모델이 생성한 순위의 NDCG@10 점수

Format reward

  • rforamt1r_{foramt1}: , 태그가 모두 출력되면 1 부여
  • rforamt2r_{foramt2}: 태그 내의 랭킹 리스트 형식이 [3] > [1] > [2] 처럼 되어있을 때 1 부여

Final reward

3.4. Initial State Expansion

  • query와, 관련성 점수가 라벨링된 학습 데이터가 부족하다는 문제를 해결하기 위해 multi sampling data augmentation 기법을 도입
  • MSMARCO-V2에서 관련성 점수(0-3)가 라벨링된 179개의 query를 활용한다.
  • 각각의 q에 대해 BM25 검색으로 top 100 문서 집합을 수집한다.
  • 위에서 수집한 집합에서 문서를 랜덤하게 추출하여 다수의 PP를 만든다.
  • 동일한 질문에 대해 다양한 초기 랭킹 목록을 만듦으로써, 다양한 상황에 대해 모델을 강건하게 학습시킬 수 있음.

4. Experiments

4.1. Experimental Setup

Training Details

  • 데이터
    • 179개의 쿼리에 대해 BM25로 뽑은 top 100 pool에서, 각 쿼리마다 20개의 후보 문서들을 무작위순서로 추출함. 이 과정을 50회 반복
    • 이때 nDCG@10 < 0.1 인 샘플들은 제외함. (GRPO의 zero reward problem 회피)
    • 총 12k 개의 학습 데이터
  • 베이스라인 모델
    • Qwen2.5-7B-Instruct
  • 학습
    • VeRL
    • batch size 128, 1스텝 당 32개의 roolout
    • SFT 없이 RL으로만 학습
    • 8개의 H100 으로 160 스텝을 학습

Sliding window reranking

  • BM25로 100개의 문서를 검색함. (n=100)
  • window size k = 20
  • query 당 10번의 LLM 호출 (2 x 100 / 20 = 10) 을 통해 top 10 문서 추출

Baselines

  • Qwen2.5-7B, GPT-4,
  • reasoning: Qwen3-32B, Qwen3-235B-A22B
  • LLM baseline 모델들에 대해서도 동일한 슬라이딩 윈도우, 동일한 프롬프트 사용함. (RankQwen, RankGPT)
  • SFT: RankZephyr
  • RL: Rank-R1 (setwise)

Benchmarks

  • TREC-DL19, DL-20, MS-MARCO-V1, BEIR
  • BRIGHT

4.2. In-domain & OOD Retrieval Results

  • GPT-4가 모든 벤치마크에서 가장 좋은 성능을 달성함 → 강력한 언어 이해 성능 때문

  • REARANK → 7B 모델에 기반한 강화학습을 통해 GPT-4와 견줄 만한 성능 달성

  • REARANK는 RANKQwen2.5-7B 베이스라인에 비해, in-domain에서 6.5%, OOD에서 4.5% 성능 향상 달성.

  • SFT 모델인 RankZephyr-7B와 비교하였을 때는 TREC-DL에서 견줄 만 했고, BEIR에서는 능가함. → SFT는 in-domain에는 강하지만, RL을 하면 OOD task에서도 강건함

  • 179개 쿼리(Rank-R1의 0.2%)만으로 Rank-R1을 능가함.

  • Qwen3-32B, Qwen3-235B도 능가함

4.3. Reasoning-intensive Retrieval Results


BRIGHT (https://arxiv.org/pdf/2407.12883)

  • BRIGHT 벤치마크 실험 결과, GPT-4를 능가하는 성능 → RL을 통해 추론 능력이 향상됨에 따른 것으로 보임
  • SFT 모델은 BM25보다 성능 하락
  • setwise인 RankR1보다 REARANK-7B가 더 좋은 성능을 달성 → 동일한 RL 기반 모델 대비 우리의 방법이 더 효율적
  • Qwen3과 같이 일반적인 추론 학습(math, coding, agent)을 거치지 않았음에도 견줄 만한 성능 달성
  • listwise 방식이 setwise에 비해 더 풍부한 신호를 주기에 더 효과적이고, 더 적은 LLM 호출을 하기 때문에 더 효율적

4.4. Ablation Studies

  • REARANK의 추론 프롬프트를 Qwen2.5-7B에 넣었을 때의 효과를 보자.
    • 조금의 성능 향상은 있지만 REARANK 만큼은 못함. → 프롬프팅만으로는 불충분하다.
  • 데이터 필터링의 효과를 보자.
    • 필터링 하지 않으면 full model에 비해 성능 하락 → 저품질의 후보셋은 zero reward를 생산하여, 제대로 된 보상이 주어지지 않아 성능 하락
  • 보상 함수 설계의 효과를 보자.
    • 바닐라 NDCG@10 (w/rrank=Srerankw/r_{rank} = S_{rerank}): 분산이 너무 커서 성능 하락
    • 바닐라 NDCG@10 - 초기 점수 (w/rrank=SrerankSinitw/r_{rank} = S_{rerank} - S_{init}): 안정성이 향상되지만 성능은 여전히 하락 → 보상 점수의 크기 자체가 작아져서 제대로 된 학습이 안 되었을 것.
    • full model의 정규화 보상 함수가 더 효과적인 가이드를 제공하여 좋은 성능이 나왔을 것임.
  • RL의 효과를 보자.
    • 고품질 데이터에 대해 SFT만 수행함 (w/direct SFT) → in-domain, OOD에서 미미한 성능 향상. 그러나 BRIGHT에서는 성능 감소 → 소량의 고품질 데이터로 리랭킹 에이전트를 학습시키기 위해서는 RL 필요

5. Analysis

Analysis of Reward Functions

  • nDCG@10의 다양한 변이형을 비교
    • Normalized nDCG (이 연구의 방법)
    • Absolute nDCG (SrerankS_{rerank})
    • Difference nDCG (SrerankSinitS_{rerank} - S_{init})

Mean Reward Score Curve

  • Normalized nDCG는 reward가 꾸준히 우상향 (100스텝에는 0.8 근접). → 이상적인 랭킹 품질에 근접
  • Absolute nDCG → 50 스텝에서 포화됨
  • Difference nDCT → 0.15 부근에서 포화됨

→ 정규화를 적용한 nDCG가 더 효과적인 신호를 주어 성능 향상에 기여했다.

Mean Response Length Curve

  • Normalized nDCG가 일반적으로 가장 긴 길이(~850 토큰)를 출력

→ Normalized nDCG가 더 상세한 추론 수행

Is reasoning helpful?

  • Qwen3-32B는 thinking 모드 on/off를 지원하는데, thinking을 켜도 성능 향상은 제한적이다. → Qwen3-32B의 강력한 성능은 reasoning에서 비롯된 게 아니라 pre-training에서 비롯된다.
  • 반면, REARANK에서는 추론 능력이 성능 향상에 핵심적으로 기여한다.
    • reasoning 프롬프트를 적용하지 않더라도 베이스라인을 이미 뛰어넘기는 한다. → 이미 내재된 리랭킹 능력
    • 그러나 추론 능력을 활용해야만 최적의 성능을 달성할 수 있다. 그만큼 RL로 추론 능력을 학습시키는 것이 중요하다.

Reasoning Pattern

  • 우리의 모델은 문서와 쿼리 간의 관련성에 대해 추론하고, 핵심 키워드를 추출한다.
  • 이전 문서에 대한 비교와 동일한 비교이면 “same”과 같은 단어를 활용하는 모습을 보이는데, 이를 통해 verbosity를 줄이면서도 핵심 키워드는 보존할 수 있다.

Is improved reasoning transferable?

  • 우리의 추론 기반 리랭킹 강화학습이 다른 능력으로 전이될 수 있는가를 확인하기 위해 수학 추론 데이터셋에 대해 평가함. (temperature=0.7로 16개의 샘플에 대해 pass@1 계산) → rearanking task에 대해 추론 능력을 학습하는 것이 또 다른 reasoning task의 성능에도 전이될 수 있다.

Impact of Reasoning Length on Performance

  • BRIGHT의 biology에 대해 한 쿼리 당 50회씩 반복적으로 리랭킹 성능을 측정하였다. (temp=0.7) 0점인 경우들은 제외. → 기존 연구(Deepseek-r1, Marjanovic et al., 2025)와는 다르게 reasoning length와 reranking 성능 간의 유의미한 상관관계는 보이지 않았다. Deepseek-r1: reasoning token이 늘수록 성능이 올라가다가 특정 지점을 넘으면 다시 떨어진다.


6. Conclusion

  • 강화학습 기반의 listwise reranking agent REARANK를 학습했다.
  • 소량의 고품질 데이터셋으로 학습하여, 리랭킹 task에서 GPT-4, Qwen3과 같은 모델과 견줄 만한 성능을 달성하거나, 뛰어넘기도 했다.

Limitations

  • 최종 랭킹을 결정하기 위해 출력하는 추론에 할루시네이션이 있을 수 있다. 이에 대한 평가는 진행하지 못했다.
  • BM25가 제공하는 초기 후보 문서 집합의 품질에 따라 전체 성능이 영향을 받는다. 즉, 리랭킹 전의 초기 검색의 품질에 의존적이다.

Reference

REARANK: Reasoning Re-ranking Agent via Reinforcement Learning (https://arxiv.org/abs/2505.20046)

profile
언어를 좋아하는 AI 엔지니어

0개의 댓글