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

TL;DR

This paper introduces reinforcement learning-based fine-tuning for history-aware dense retrievers in RAG, addressing objective mismatch issues and improving retrieval performance through stochastic sampling and history incorporation.

Contribution

It proposes a novel RL-based retriever optimization method that incorporates retrieval history and addresses deterministic retrieval limitations in RAG.

Findings

Consistent performance improvements across diverse RAG pipelines.

Effective mitigation of state aliasing through retrieval history.

Enhanced alignment between retriever training and RAG objectives.

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.