R3-RAG: Learning Step-by-Step Reasoning and Retrieval for LLMs via Reinforcement Learning

2025-05-26

AI Summary

In RAG systems, dense retrievers—limited by small parameter counts and weak reasoning capabilities—constitute a key bottleneck for factual accuracy; meanwhile, existing prompt-driven iterative methods rely on hand-crafted pipelines and suffer from poor generalization. This paper proposes the first end-to-end RAG framework trained via Proximal Policy Optimization (PPO), enabling LLMs to autonomously learn synergistic retrieval-and-reasoning strategies through a two-stage RL procedure. We introduce a novel joint reward mechanism: correctness of the final answer serves as the outcome reward, while relevance verification at intermediate steps acts as the process reward—yielding an interpretable, generalizable iterative closed loop. Our method supports interleaved retrieval-reasoning action modeling and incorporates cold-start pretraining. It achieves significant improvements over SOTA across multiple benchmarks and demonstrates cross-retriever transferability. The code is publicly available.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Retrieval-Augmented Generation (RAG) integrates external knowledge with Large Language Models (LLMs) to enhance factual correctness and mitigate hallucination. However, dense retrievers often become the bottleneck of RAG systems due to their limited parameters compared to LLMs and their inability to perform step-by-step reasoning. While prompt-based iterative RAG attempts to address these limitations, it is constrained by human-designed workflows. To address these limitations, we propose $ extbf{R3-RAG}$, which uses $ extbf{R}$einforcement learning to make the LLM learn how to $ extbf{R}$eason and $ extbf{R}$etrieve step by step, thus retrieving comprehensive external knowledge and leading to correct answers. R3-RAG is divided into two stages. We first use cold start to make the model learn the manner of iteratively interleaving reasoning and retrieval. Then we use reinforcement learning to further harness its ability to better explore the external retrieval environment. Specifically, we propose two rewards for R3-RAG: 1) answer correctness for outcome reward, which judges whether the trajectory leads to a correct answer; 2) relevance-based document verification for process reward, encouraging the model to retrieve documents that are relevant to the user question, through which we can let the model learn how to iteratively reason and retrieve relevant documents to get the correct answer. Experimental results show that R3-RAG significantly outperforms baselines and can transfer well to different retrievers. We release R3-RAG at https://github.com/Yuan-Li-FNLP/R3-RAG.

Problem

• Improving RAG systems by enabling step-by-step reasoning and retrieval
• Overcoming dense retriever limitations with reinforcement learning
• Enhancing answer correctness via iterative reasoning and document relevance

Innovation

• Reinforcement learning for reasoning and retrieval
• Cold start for iterative reasoning and retrieval
• Dual rewards for answer correctness and relevance