Reinforcement Fine-Tuning for History-Aware Dense Retriever in RAG

TL;DR

The paper tackles the misalignment between retriever training objectives and end-to-end RAG performance by introducing HARR, a history-aware reinforcement fine-tuning framework for dense retrievers. It reframes retrieval as an MDP and replaces deterministic top-$k$ selection with stochastic sampling, enabling RL optimization, while adding retrieval history to the state to mitigate multi-hop state aliasing. Using GRPO with sparse terminal rewards and practical approximations, HARR achieves consistent end-to-end QA gains across diverse RAG pipelines, datasets, and retriever scales, often outperforming proxy-objective baselines. The approach offers a lightweight, scalable retriever optimization path that directly tunes retrieval to downstream task rewards, with code available for reproduction. Overall, HARR demonstrates that retriever-centric RL can substantially improve grounding performance in RAG without retraining large LLMs. $r_T = \mathrm{F1}(y, y^*)$ is used as the terminal reward to align retrieval with final answer quality, and history-aware states $s_t = (\mathcal{H}_{t-1}, q_t)$ help distinguish similar queries across retrieval histories.

Abstract

Retrieval-augmented generation (RAG) enables large language models (LLMs) to produce evidence-based responses, and its performance hinges on the matching between the retriever and LLMs. Retriever optimization has emerged as an efficient alternative to fine-tuning LLMs. However, existing solutions suffer from objective mismatch between retriever optimization and the goal of RAG pipeline. Reinforcement learning (RL) provides a promising solution to address this limitation, yet applying RL to retriever optimization introduces two fundamental challenges: 1) the deterministic retrieval is incompatible with RL formulations, and 2) state aliasing arises from query-only retrieval in multi-hop reasoning. To address these challenges, we replace deterministic retrieval with stochastic sampling and formulate RAG as a Markov decision process, making retriever optimizable by RL. Further, we incorporate retrieval history into the state at each retrieval step to mitigate state aliasing. Extensive experiments across diverse RAG pipelines, datasets, and retriever scales demonstrate consistent improvements of our approach in RAG performance.

Figures (3)

• Figure 1: Overview of HARR
• Figure 2: Training reward trajectories on the HotpotQA dataset. Subplot titles denote the specific RAG pipeline and retriever encoder used. All curves are smoothed using Exponential Moving Average (EMA) with a window size of 8.
• Figure 3: Gradient norm trajectories on HotpotQA using the ReAct Agent pipeline. Subplot titles denote the retriever encoder size.