Stop-RAG: Value-Based Retrieval Control for Iterative RAG

AI summary

Overview

Research area: Iterative retrieval-augmented generation (RAG) for large language models, agentic tool-use systems, and reinforcement learning for decision control.

Technical level: Intermediate. The paper is readable without deep RL background, but the method section assumes familiarity with MDPs, Q-learning, and temporal-difference targets, and with standard iterative RAG pipelines.

Scope: The paper proposes Stop-RAG, a trained value-based controller that decides at each iteration of an iterative RAG loop whether to stop retrieving or continue, and evaluates it on three multi-hop question-answering benchmarks.

What This Paper Is About

Iterative RAG systems retrieve and generate repeatedly to answer multi-hop questions, but every extra loop adds latency, token cost, and the risk of pulling in distracting documents. Existing systems either run a fixed number of iterations or let an LLM or a confidence signal guess when to stop, and these signals do not reliably reflect whether more retrieval will actually improve the answer. The paper's goal is to replace those heuristics with a learned controller that estimates the future benefit of continuing retrieval before deciding whether to stop.

Key Contributions

1. Frames adaptive stopping as a first-class problem in agentic AI, arguing that deciding when to stop calling tools is an underexplored skill compared with the ability to call tools at all.
2. Formulates iterative RAG as a finite-horizon Markov decision process with two actions, STOP and CONT, where rewards are assigned only at terminal states and reflect answer-generation quality.
3. Trains a Q-network using full-width forward-view Q(λ) targets computed from complete offline trajectories, giving forward-looking estimates of whether continued retrieval will improve answer quality, without requiring online interaction or access to model internals.
4. Demonstrates the method as a plug-and-play stopping module that improves multi-hop QA performance on top of two different iterative RAG pipelines, including one with a fine-tuned generator, while remaining compatible with black-box LLM APIs.

Main Findings

• Stop-RAG beats both fixed-iteration and prompting-based stopping on the authors' own pipeline. On MuSiQue, their pipeline goes from EM 34.5 / F1 44.8 / Acc 41.1 (raw, 10 iterations) and EM 34.2 / F1 44.6 / Acc 41.3 (LLM-Stop) to EM 36.8 / F1 47.0 / Acc 43.9 with Stop-RAG. On 2WikiMultihopQA it moves from EM 64.9 / F1 73.1 / Acc 65.9 (raw) and EM 64.9 / F1 73.3 / Acc 66.1 (LLM-Stop) to EM 68.2 / F1 75.7 / Acc 69.0.

• Gains also appear on HotpotQA for the authors' pipeline, but they are smaller. Raw: EM 51.0 / F1 65.0 / Acc 60.8; LLM-Stop: EM 51.4 / F1 65.2 / Acc 60.7; Stop-RAG: EM 52.4 / F1 66.1 / Acc 62.6.

• Stop-RAG is effective as a plug-in on CoRAG, but with mixed results. On CoRAG (L=10, greedy), Stop-RAG reaches EM 31.5 / F1 43.0 / Acc 36.9 on MuSiQue (vs. 30.9 / 41.9 / 36.3 raw and 31.1 / 42.0 / 36.3 with LLM-Stop) and EM 65.7 / F1 72.5 / Acc 68.2 on 2WikiMultihopQA (vs. 65.3 / 71.7 / 68.0 raw and 65.4 / 72.2 / 67.8 with LLM-Stop). On HotpotQA it slightly underperforms both baselines, at EM 54.7 / F1 67.4 / Acc 64.4 versus 55.0 / 67.5 / 65.4 (raw) and 54.9 / 67.6 / 64.8 (LLM-Stop).

• The authors attribute the HotpotQA exception to the fine-tuned CoRAG generator and to retrieval recall. They argue that a generator trained under a fixed iteration count is more robust to distractors, and that HotpotQA's broader, more challenging corpus makes continuing retrieval for the full iteration count worth the distractor cost.

• Stop-RAG retrieves more than LLM-Stop but less than the full 10 iterations. On MuSiQue in the authors' pipeline, LLM-Stop runs 7.6 steps on average (precision 31.6, recall 64.5), Stop-RAG runs 8.6 steps (precision 22.2, recall 66.5), and the unstopped pipeline runs 10.0 steps (precision 17.4, recall 68.2). On 2WikiMultihopQA, LLM-Stop runs 5.0 steps (precision 55.0, recall 79.1), Stop-RAG runs 5.1 steps (precision 49.5, recall 82.6), and the unstopped pipeline runs 10.0 steps (precision 21.4, recall 88.3).

• Q(λ) is the best training target among the four tested on MuSiQue. Q(λ) achieves EM 36.8 / F1 47.0 / Acc 43.9, versus Monte Carlo returns at EM 36.0 / F1 46.2 / Acc 43.0, one-step TD (Q(0)) at EM 36.0 / F1 46.3 / Acc 42.8, and a binary classification variant at EM 33.3 / F1 43.7 / Acc 40.3. The authors note MC suffers from high variance and Q(0) from being myopic.

• Not reported: absolute latency, token-cost savings, or statistical significance tests are not given. The paper motivates stopping by latency and cost but reports only accuracy, retrieval precision/recall, and average step counts.

Methodology in Plain English

The authors model one round of an iterative RAG pipeline — generating a query, retrieving and reranking documents, and producing an intermediate answer — as a single step in a decision process. The state is the question plus all retrieved documents so far. At each step there are only two choices: stop and answer, or continue and run another RAG round. Rewards are only given at the end, and they measure how good the generated answer is compared with the gold answer.

To train the controller, they first run the pipeline to its maximum of 10 iterations with no stopping, producing full trajectories. Every prefix of a trajectory becomes a training state. For each state they run eight independent answer-generation trials, score each output against the gold answer with F1, and average — that average is the reward for stopping there. They discard terminal states and states that give zero signal for both actions.

The controller itself is a DeBERTa-v3-large encoder with two feed-forward prediction heads, one for STOP and one for CONT. It sees only the main question plus the retrieved documents, separated by [SEP]; intermediate queries and answers are deliberately excluded, because the authors found the retrieved evidence alone sufficient for learning a stopping policy. Training uses full-width forward-view Q(λ) targets: because there are only two actions and STOP always ends the episode, the paper can compute targets over all actions at a state rather than sampling them. Training starts with λ = 1 (long-horizon, low-bias targets) and anneals λ down to 0.1, trading bias for variance reduction as learning proceeds. This avoids any online interaction — everything is derived from the pre-collected offline dataset.

At inference the system does not simply take the argmax action. Instead it uses a margin rule: stop only if the STOP value exceeds the CONT value by a threshold tuned on the validation set.

Experimental setup: evaluation is on MuSiQue, HotpotQA, and 2WikiMultihopQA. Training uses the official train partitions; 1000 questions are sampled from dev for validation and another 1000 for testing. The pipeline uses Llama-3.1-8B-Instruct as query and answer generator, Contriever-MSMARCO as retriever, and bge-reranker-v2-m3 as reranker, retrieving 10 documents per iteration and selecting the top-ranked one, with a maximum of 10 iterations. Metrics are exact match, F1, and answer accuracy. Q-network hyperparameters include prediction-head hidden size 4096, AdamW, a cosine-decayed learning rate from 5e-5 to 0.0, warmup ratio 0.1, weight decay 0.01, batch size 128, bfloat16 mixed precision, and 3 epochs for MuSiQue versus 1 epoch for HotpotQA and 2WikiMultihopQA.

Why This Matters

Impact on research. The paper argues that adaptive stopping is a missing component in agentic systems, where most work focuses on enabling more tool calls rather than judging when enough information has been gathered. It contributes a concrete RL formulation and a training recipe that operates on offline trajectories and black-box model APIs, which lowers the barrier to studying stopping as a standalone learned skill. The ablation also provides evidence about which value-based target works best, with Q(λ) ahead of Monte Carlo, one-step TD, and binary classification on F1 and accuracy.

Real-world applications:

• Enterprise search and knowledge assistants that iterate over internal document stores, where each extra retrieval round costs money and time and can pull in stale or misleading documents.
• Customer support and technical troubleshooting bots that must decide when accumulated documentation is enough to answer, rather than looping indefinitely.
• Clinical or legal research assistants over large document corpora, where excessive retrieval can surface irrelevant evidence that degrades a final answer.
• Multi-hop fact-checking and investigative tools that need to balance retrieval depth against precision when tracing claims across sources.

Industry relevance. Because the controller is trained offline and only reads retrieved documents, it can be bolted onto existing RAG deployments without exposing model internals or retraining the generator. The paper shows this works on CoRAG, a pipeline with a fine-tuned generator, as well as on the authors' own training-free pipeline, suggesting the module transfers across architectures — though the HotpotQA result indicates the benefit depends on how well the underlying generator handles distractors.

Future Directions

• Extending the controller to other iterative RAG frameworks. The authors explicitly flag applying Stop-RAG to other high-performing iterative RAG systems as future work, since only CoRAG and their own pipeline were tested.
• Understanding when adaptive stopping does not help. The HotpotQA underperformance on CoRAG suggests that fine-tuned generators may already absorb the distractor cost that stopping is meant to avoid; characterizing this interaction is an open question.
• Moving beyond a state representation built only from retrieved documents. The current design omits intermediate queries and answers; whether including them helps, and under what conditions, is unexplored.
• Quantifying the efficiency payoff. The paper frames latency and token cost as the motivation for stopping but reports no wall-clock or token measurements, leaving the practical savings unmeasured.

Target Audience

Researchers and graduate students in natural language processing and reinforcement learning who work on retrieval-augmented generation, multi-hop question answering, or agentic LLM systems. It is also useful for machine learning engineers and applied scientists building production RAG pipelines who need a principled, modular stopping mechanism and can train a small encoder-based controller offline. Readers without any exposure to MDPs or value-based RL will find the methodology section demanding, though the experimental results and motivation are accessible on their own.

Authors’ abstract

Iterative retrieval-augmented generation (RAG) enables large language models to answer complex multi-hop questions, but each additional loop increases latency, costs, and the risk of introducing distracting evidence, motivating the need for an efficient stopping strategy. Existing methods either use a predetermined number of iterations or rely on confidence proxies that poorly reflect whether more retrieval will actually help. We cast iterative RAG as a finite-horizon Markov decision process and introduce Stop-RAG, a value-based controller that adaptively decides when to stop retrieving. Trained with full-width forward-view Q($λ$) targets from complete trajectories, Stop-RAG learns effective stopping policies while remaining compatible with black-box APIs and existing pipelines. On multi-hop question-answering benchmarks, Stop-RAG consistently outperforms both fixed-iteration baselines and prompting-based stopping with LLMs. These results highlight adaptive stopping as a key missing component in current agentic systems, and demonstrate that value-based control can improve the accuracy of RAG systems.