SmartRAG: Jointly Learn RAG-Related Tasks From the Environment Feedback

TL;DR

The paper introduces SmartRAG, a joint reinforcement learning framework for retrieval-augmented generation that trains a policy network to decide when to retrieve, how to rewrite queries, and how to generate answers using retrieved observations. It integrates a retriever and uses PPO to optimize end-to-end with an environment-based reward that balances correctness and retrieval cost. Warm-up supervised finetuning provides a reasonable initialization, after which PPO fine-tunes all components, achieving superior performance over modular baselines across multiple open-domain QA datasets. The results highlight the importance of end-to-end inter-module awareness, enabling smarter retrieval decisions, better query formulation, and more effective use of retrieved content, with favorable transfer to unseen datasets.

Abstract

RAG systems consist of multiple modules to work together. However, these modules are usually separately trained. We argue that a system like RAG that incorporates multiple modules should be jointly optimized to achieve optimal performance. To demonstrate this, we design a specific pipeline called \textbf{SmartRAG} that includes a policy network and a retriever. The policy network can serve as 1) a decision maker that decides when to retrieve, 2) a query rewriter to generate a query most suited to the retriever, and 3) an answer generator that produces the final response with/without the observations. We then propose to jointly optimize the whole system using a reinforcement learning algorithm, with the reward designed to encourage the system to achieve the best performance with minimal retrieval cost. When jointly optimized, all the modules can be aware of how other modules are working and thus find the best way to work together as a complete system. Empirical results demonstrate that the jointly optimized SmartRAG can achieve better performance than separately optimized counterparts.

Figures (7)

• Figure 1: Overview of SmartRAG. The policy network takes the original question and optional observations as input and generates an action that can be interpreted as a decision-maker and the corresponding parameters. If the policy decides to retrieve, the retriever will be called to obtain the observation that updates the state. Otherwise, the policy network will output the final answer. SmartRAG is optimized using a reinforcement learning algorithm with a properly designed reward.
• Figure 2: F1 Score of different retrieval percentage across three datasets on Flan-T5-large.
• Figure 3: Influence of different initial policies for SmartRAG.
• Figure 4: Iteration reward, KL reward, policy loss and value loss during training.
• Figure 5: Exact Match, F1 Score, and Retrieval Threshold of different retrieval number K on Flan-T5-large.