PaperlayerView exact revision on arXiv

Source-linked AI summary

Chain-of-Retrieval Augmented Generation

Liang Wang, Haonan Chen, Nan Yang, Xiaolong Huang, Zhicheng Dou, Furu Wei

arXiv:2501.14342v3cs.IRcs.CL

TL;DR

CoRAG addresses the limitations of single-step RAG retrieval for complex queries by training language models to retrieve and reason iteratively. It uses rejection sampling to create intermediate retrieval chains and configurable decoding to scale test-time compute, achieving strong multi-hop QA and KILT results. The study primarily evaluates short, easy-to-verify tasks and notes that long-form generation remains outside its main scope.

  • Problem

    Single-step RAG retrieval can provide imperfect information for complex queries, while multi-hop QA often requires information that retrievers fail to recall in one step.

  • Method

    CoRAG uses rejection sampling to generate intermediate retrieval chains, fine-tunes an LLM for iterative retrieval and reasoning, and offers decoding strategies that control test-time compute.

  • Results

    CoRAG-8B achieves state-of-the-art performance on multihop QA datasets and KILT, surpassing many baselines built with larger LLMs.

  • Takeaways & Limitations

    CoRAG provides a framework for stepwise retrieval and adjustable test-time computation, with reported gains across multi-hop QA and diverse knowledge-intensive tasks.

  • Takeaways & Limitations

    The study focuses mainly on short, easy-to-verify tasks such as multi-hop QA and entity linking, while long-form generation lacks robust evaluation metrics.

Abstract

from arXiv · show

This paper introduces an approach for training o1-like RAG models that retrieve and reason over relevant information step by step before generating the final answer. Conventional RAG methods usually perform a single retrieval step before the generation process, which limits their effectiveness in addressing complex queries due to imperfect retrieval results. In contrast, our proposed method, CoRAG (Chain-of-Retrieval Augmented Generation), allows the model to dynamically reformulate the query based on the evolving state. To train CoRAG effectively, we utilize rejection sampling to automatically generate intermediate retrieval chains, thereby augmenting existing RAG datasets that only provide the correct final answer. At test time, we propose various decoding strategies to scale the model's test-time compute by controlling the length and number of sampled retrieval chains. Experimental results across multiple benchmarks validate the efficacy of CoRAG, particularly in multi-hop question answering tasks, where we observe more than 10 points improvement in EM score compared to strong baselines. On the KILT benchmark, CoRAG establishes a new state-of-the-art performance across a diverse range of knowledge-intensive tasks. Furthermore, we offer comprehensive analyses to understand the scaling behavior of CoRAG, laying the groundwork for future research aimed at developing factual and grounded foundation models.

1 Introduction

Conventional RAG depends heavily on single-step retrieval quality, which limits performance on complex and multi-hop queries. CoRAG addresses this bottleneck through iterative retrieval, state-based query reformulation, and controllable test-time decoding.

  • Motivation: Conventional RAG systems rely on retrieved information as input to generation, while efficient bi-encoder retrievers can struggle with complex queries.Fixed-size vector representations enable scalable search but constrain retrieval expressiveness.
  • Approach: CoRAG dynamically retrieves relevant information and plans subsequent retrieval steps based on the current state.The framework can adjust retrieval depth and experiment with query rewriting when retrieved information is inadequate.
  • Training: Rejection sampling augments existing RAG datasets with intermediate retrieval chains, after which open-source language models are fine-tuned to retrieve step by step.The method also provides multiple decoding strategies for scaling test-time compute.
  • Results: CoRAG substantially surpasses strong baselines on multi-hop question-answering tasks where single-step retrieval often misses necessary information.Across decoding strategies, the performance-compute Pareto frontier is approximately log-linear, with coefficients varying across datasets.
  • Results: CoRAG achieves new state-of-the-art scores on nearly all tasks in the KILT hidden test set and shows varied scaling behavior across task types.The analysis reports query decomposition, flexible reformulation, and robustness to retrievers of varying quality.

2 Related Work

Related work develops RAG retrieval, iterative retrieval-generation, and test-time reasoning methods. CoRAG builds on these directions while targeting stepwise retrieval and controllable compute scaling.

  • RAG: RAG combines information retrieval with generative models to improve factuality, groundedness, and access to current information.Its effectiveness depends strongly on the relevance and quality of retrieved information.
  • RAG: Fixed-size embedding representations improve retrieval efficiency but limit the handling of complex queries.Recent work therefore investigates stronger general-purpose text embeddings and alternatives to conventional retrieval.
  • Iterative Retrieval: Iterative RAG methods extend the single-retrieval pipeline by interleaving retrieval and generation or letting models decide when and what to retrieve.Examples include FLARE, ITER-RETGEN, and IRCoT.
  • Test-Time Compute: Chain-of-Thought, Tree-of-Thought, and related methods improve reasoning by generating intermediate steps or exploring multiple reasoning paths.These approaches generally increase token consumption and response latency.
  • Test-Time Compute: RAG-specific compute scaling increases retrieved documents or retrieval steps, with prior work studying long-context retrieval, iterative retrieval, and active search.These approaches provide context for CoRAG's controllable retrieval-chain decoding.

3 Methodology

CoRAG generates retrieval chains through rejection sampling, trains an LLM on intermediate and final prediction tasks, and scales inference through configurable decoding strategies.

  • Framework: CoRAG represents the current state as input context and instructions, and the next action as the LLM's response to those instructions.The framework comprises retrieval-chain generation, augmented-data training, and test-time compute scaling.
  • Retrieval Chain Generation: Rejection sampling generates chains of sub-queries and sub-answers from QA datasets that lack intermediate retrieval steps.Generation continues until the maximum chain length is reached or a sub-answer matches the correct answer.
  • Retrieval Chain Generation: The chain with the highest conditional log-likelihood of the correct answer is selected to augment the original QA dataset.This selection evaluates the answer conditioned on the original query and intermediate chain information.
  • Training: Each augmented training instance includes the original query and answer, intermediate queries and answers, and retrieved documents for the original and sub-queries.The LLM is fine-tuned with a standard next-token prediction objective.
  • Training: The model jointly predicts the next sub-query, sub-answer, and final answer.Final-answer prediction additionally receives the top retrieved documents for the original query.
  • Scaling Test-Time Compute: Test-time compute is measured using both prompt and generated tokens, excluding retrieval costs.Prompt and generated tokens are treated equally despite differences in practical computational expense.
  • Decoding Strategies: Greedy decoding generates one retrieval chain sequentially, whereas best-of-N sampling selects among N sampled chains using a penalty score.Tree search expands states with sampled sub-queries and retains states with the lowest average rollout penalty.
  • Scaling Test-Time Compute: The retrieval-chain length L controls compute across decoding methods, while N, rollout count, and expansion size provide additional controls.These parameters allow inference compute to be adjusted without changing model weights.

4 Experiments

CoRAG is evaluated on multi-hop QA and KILT using augmented retrieval-chain datasets, multiple decoding strategies, and test-time compute scaling. It substantially outperforms baselines on multi-hop QA and reaches state-of-the-art performance across nearly all KILT tasks, while scaling gains vary by dataset.

  • Setup: CoRAG is evaluated on multi-hop QA datasets and the broader KILT benchmark using rejection-sampled retrieval chains and E5-large intermediate retrieval.The multi-hop QA collection includes 2WikiMultihopQA, HotpotQA, Bamboogle, and MuSiQue; KILT uses approximately 36 million English Wikipedia passages.
  • Main Results: CoRAG-8B substantially surpasses the compared baselines on multi-hop QA, except on Bamboogle.The comparison includes few-shot and fine-tuned Llama, GPT-4o, Self-RAG, ITER-RETGEN, DRAG, IterDRAG, and Search-o1-32B.
  • Main Results: CoRAG-8B achieves new state-of-the-art performance across KILT tasks except FEVER, where it marginally trails an 11B-parameter model.KILT scores are reported on the hidden test set using the best decoding configuration selected from public validation results.
  • Scaling Test-Time Compute: Increasing retrieval-chain length L improves performance substantially at small values, but gains diminish as L increases.Test-time compute is controlled through retrieval-chain length and the number of sampled chains, with greedy decoding as best-of-N when N = 1 and temperature 0.
  • Scaling Test-Time Compute: The EM-score Pareto frontier approximately follows a log-linear trajectory with token consumption up to 128k tokens, although scaling varies across datasets.The study treats prompt and generated tokens equivalently and ignores retrieval costs.

5 Analysis

The analysis examines training variants, retriever substitutions, task-dependent scaling, learned stopping, and retrieval recall. CoRAG’s benefits are strongest for complex multi-hop tasks and difficult retrieval settings, while early stopping trades token savings for lower performance.

  • Ablations: Iterative training improves 2WikiMultihopQA but slightly lowers performance on other datasets.The mixed results suggest instruction-tuned language models already generate high-quality retrieval chains.
  • Different Retrievers: More test-time compute consistently improves performance with weaker retrievers, although stronger retrievers retain higher absolute performance.The comparison substitutes E5-base and BM25 for E5-large at test time.
  • Ablations: Llama-3B reaches performance close to the 8B model for retrieval-chain generation, whereas Llama-1B causes a noticeable performance drop.Manual inspection attributes the 1B degradation to difficulty following instructions and producing effective retrieval chains.
  • Does Chain-of-Retrieval Always Help?: The advantage of chain-of-retrieval is marginal on mostly single-hop datasets, whereas multi-hop datasets are expected to benefit more.NQ and TriviaQA exemplify tasks where one retrieval step is often sufficient, motivating adaptive decoding based on query complexity.
  • Learning to Stop at Test Time: Learning to stop early can save token quota but causes performance degradation, with the best configuration depending on dataset characteristics and quality requirements.The stopping decision is made after each retrieval step using a Yes/No prediction whose logit bias controls early stopping.
  • Retrieval Recall: CoRAG consistently improves retrieval recall across datasets and recall thresholds, especially on MuSiQue and Bamboogle.The evaluation reports Recall@10, Recall@20, and Recall@100 against standard E5-large retrieval, merging chain results with reciprocal rank fusion.

6 Conclusion

CoRAG teaches language models to perform iterative retrieval and reasoning for complex queries, using automatically generated retrieval chains and test-time decoding strategies to balance performance and compute.

  • CoRAG teaches LLMs to conduct iterative retrieval and reasoning for complex queries.
  • Rejection sampling automatically generates intermediate retrieval chains, avoiding manual annotation.
  • Multiple decoding strategies manage the trade-off between performance and test-time compute.
  • CoRAG-8B achieves state-of-the-art performance on multihop QA datasets and the KILT benchmark, surpassing many baselines built with larger LLMs.

7 Limitations and Broader Impacts

The study focuses on short, easy-to-verify RAG tasks, while long-form generation remains difficult to evaluate because robust metrics are lacking. The framework aims to improve factuality and groundedness, but hallucination risk persists in deployment.

  • The study primarily investigates short, easy-to-verify tasks such as multi-hop QA and entity linking.
  • Long-form generation remains challenging because robust evaluation metrics are absent.
  • The framework aims to improve the factuality and groundedness of language-model outputs, while hallucination risk persists in practical deployments.

A Implementation Details

The implementation generates retrieval chains through rejection sampling, trains on multi-hop QA and KILT data, and evaluates with task-specific metrics and official benchmark procedures.

  • Rejection Sampling: Rejection sampling draws up to 16 retrieval chains with maximum lengths randomly selected from 1 to 5.
  • Rejection Sampling: Each sub-query retrieves the top-5 documents using E5-large from a 36-million-passage KILT Wikipedia corpus.
  • Training Data: The multi-hop QA training set combines 2WikiMultihopQA, HotpotQA, and MuSiQue for 125k samples, while Bamboogle is evaluation-only.
  • Training Data: KILT training omits ELI5 and WoW, caps large datasets at 100k samples, and adds 100k BLINK samples for entity linking.
  • Evaluation: The KILT evaluation uses official scripts, with EM for open QA and accuracy for other task types.

B Additional Results

Additional experiments examine benchmark decoding, training-data and test-time compute, sampling temperature, model-family transfer, qualitative behavior, evaluation uncertainty, and prompt designs.

  • Decoding Strategies: KILT decoding strategies show less performance disparity because most tasks are easier for strong dense retrievers than multi-hop QA.
  • Training-Data Compute: Increasing sampled chains during rejection sampling shows no definitive trend toward better performance.
  • Test-Time Compute: Targeted fine-tuning is essential for improving the scaling upper bound: untuned Llama-3.1-8B-Instruct has similar curves but significantly lower absolute performance.
  • Sampling Temperature: Sampling temperature has no consistent performance effect; lower temperatures generally work better on MuSiQue and HotpotQA.
  • Model Families: CoRAG consistently outperforms fine-tuned baselines across Qwen3-4B and Qwen3-8B datasets, with improvements over 10 EM points on average.
  • Case Analysis: CoRAG decomposes complex multi-hop queries into simpler sub-queries and can self-correct through additional retrieval steps.
  • Statistical Analysis: CoRAG-8B is significantly better than baselines on all datasets except Bamboogle under the reported bootstrap confidence analysis.

2. Limitations

The paper states that its limitations are discussed in Section 7, while the supplied material provides no specific limitation beyond general review guidance.

  • The paper reports that its limitations are discussed in Section 7.
  • The review guidance asks authors to state strong assumptions, robustness to violations, and the scope of empirical claims.
  • The supplied theoretical-results assessment marks the question as not applicable because the paper provides no theoretical results.

4. Experimental result reproducibility

The paper reports that implementation and training details are provided in Sections 4.1 and A, with code and data planned for release after publication.

  • The paper states that implementation details are provided in Section A.
  • The reproducibility guidance recognizes that replication may use detailed instructions, hosted models, checkpoints, or other appropriate means.
  • The authors state that they will release the code and data after publication.
  • The paper directs readers to Sections 4.1 and A for training and test details, including data splits and hyperparameters.

7. Experiment statistical significance

The paper reports confidence intervals for its experiments and describes statistical-reporting expectations, while also stating that its work conforms to the NeurIPS Code of Ethics and addresses societal-impact review questions.

  • Statistical reporting: The paper reports confidence intervals in Section C.
  • Statistical reporting: The statistical-reporting guidance requires authors to explain what variability error bars capture and how they are calculated.
  • Ethics and impacts: The paper states that its research conforms to the NeurIPS Code of Ethics.
  • Ethics and impacts: The paper directs readers to Section 7 for its discussion of positive and negative societal impacts.
  • Responsible release: The authors state that no safeguards are needed because the work does not involve high-risk data or models.
Loading 2501.14342v3…