Knowing You Don't Know: Learning When to Continue Search in Multi-round RAG through Self-Practicing

TL;DR

The paper tackles the challenge of knowing when to stop retrieval in multi-round RAG. It introduces SIM-RAG, a modular framework with a Reasoner, a Retriever, and a lightweight Critic that provides information-sufficiency judgments, guided by a Self-Practicing data-generation stage and reasoning-enhanced in-context reinforcement learning. Empirical results across TriviaQA, HotPotQA, and 2WikiMultiHopQA show SIM-RAG achieves state-of-the-art performance, with the GPT-4 Reasoner delivering the strongest results and even lightweight Critics delivering substantial gains. The approach is data-efficient and plug-and-play, not requiring changes to the LLMs or search engines, and opens avenues for using the Critic as a reward model for further RAG optimization and broader AI reasoning tasks.

Abstract

Retrieval Augmented Generation (RAG) has shown strong capability in enhancing language models' knowledge and reducing AI generative hallucinations, driving its widespread use. However, complex tasks requiring multi-round retrieval remain challenging, and early attempts tend to be overly optimistic without a good sense of self-skepticism. Current multi-round RAG systems may continue searching even when enough information has already been retrieved, or they may provide incorrect answers without having sufficient information or knowledge. Existing solutions either require large amounts of expensive human-labeled process supervision data or lead to subpar performance. This paper aims to address these limitations by introducing a new framework, SIM-RAG, to explicitly enhance RAG systems' self-awareness and multi-round retrieval capabilities. To train SIM-RAG, we first let a RAG system self-practice multi-round retrieval, augmenting existing question-answer pairs with intermediate inner monologue reasoning steps to generate synthetic training data. For each pair, the system may explore multiple retrieval paths, which are labeled as successful if they reach the correct answer and unsuccessful otherwise. Using this data, we train a lightweight information sufficiency Critic. At inference time, the Critic evaluates whether the RAG system has retrieved sufficient information at each round, guiding retrieval decisions and improving system-level self-awareness through in-context reinforcement learning. Experiments across multiple prominent RAG benchmarks show that SIM-RAG is an effective multi-round RAG solution. Furthermore, this framework is system-efficient, adding a lightweight component to RAG without requiring modifications to existing LLMs or search engines, and data-efficient, eliminating the need for costly human-annotated mid-step retrieval process supervision data.

Figures (5)

• Figure 1: A key challenge in multi-round RAG systems is determining the optimal stopping point for retrieval and then generating the answer. This figure illustrates two typical patterns that hurt performance: Over-Confidence (e.g., stopping too early, as seen in 0 and 1 Round RAG, where the system provides an incorrect answer based on limited context) and Over-Retrieval (e.g., retrieving unnecessary information in 3 Rounds RAG which introduces a long and overly complex context that confuses the LLM and leads to incorrect answers.).
• Figure 2: SIM-RAG overview at inference time, consisting of three main components: (1) Answer Generation, where the Reasoner generates an answer and rationale; (2) Sufficiency Inspection, where the Critic decides whether the generated answer to accept the answer or trigger refinement; and (3) Information Retrieval, where, if deemed insufficient, a query is generated to retrieve additional documents from the corpus to update the context. This iterative process continues until the Critic deems the answer sufficient or the maximum number of retrieval rounds is reached. The orange line represents the information collection path, while the blue line represents the answer generation path, visualizing the information flow between components.
• Figure 3: Confusion matrices of Critic predictions in SIM-RAG with GPT-4.
• Figure 4: Ablation study on different Critic model choices. Both follow the SIM-RAG setting with GPT-4 as the Reasoner, but the blue bar uses learned Flan-T5-783M as a Critic, while the red bar uses GPT-4 as the Critic (a.k.a., self-critique).
• Figure 5: Ablation study on different Inner Monologue turns on HotPotQA. The green line represents the Naive Generation baseline, while the blue line shows the performance of SIM-RAGlite. The red line indicates the EM score using the ground-truth label as an Oracle Critic.