Chain-of-Retrieval Augmented Generation

Thank you for your positive feedback and constructive suggestions!

"I am a little confused about the terms "Current State" and "Next Action" in Figure 1. How many times will the LLM be trained? Does the "State" refer to the token?"

In Figure 1, "Current State" denotes the input context and instructions provided to the LLM. The "Next Action" refers to the LLM output responding to the given instruction. To illustrate, consider sub-query generation: the "Current State" encompasses the original query along with all historical sub-queries and their corresponding answers, while the "Next Action" is the newly generated sub-query. This conceptualization similarly applies to both sub-answer generation and final answer generation. For detailed examples of our prompt templates, please refer to Page 19 of the manuscript.

To clarify the training process, the LLM is trained only once. We mix the training data for all three tasks, enabling simultaneous multi-task learning.

"Some recent papers have proposed R1-Zero-like RAG methods. It would be interesting to compare them. I am aware that they should be considered contemporary work, and this will not influence my evaluation of this paper."

We appreciate your suggestion! We have incorporated a brief discussion of Search-R1 in the related work section. RL holds great promises for agentic systems; however, it is also well-known for high computational demands and training instability. Given the emergence of numerous concurrent works applying RL to RAG in recent months, we intend to substantially expand upon this discussion in the final version of our paper.

Thank you for your positive feedback and constructive suggestions!

"the results suggest that rejection sampling had limited impact, which weakens this contribution. It would have been valuable to explore this direction further—e.g., through preference learning or alternative training approaches."

Our findings indicate that the choice of LLM for rejection sampling significantly impacts performance (Table 3), whereas merely increasing the number of sampled trajectories yielded limited gains (Figure 6 in Appendix B). We hypothesize that the diversity and quality of sampled trajectories are currently constrained by the inherent capabilities of the LLM. We agree that exploring approaches like online rejection sampling or reinforcement learning, as you suggested, holds substantial promise for addressing this limitation, and we plan to investigate these directions in our future work.

"Did the authors observe any failure cases where longer chains hurt performance, or where irrelevant subquestions derailed the retrieval process?"

Yes, we indeed observed common failure patterns. A prominent one is the generation of sub-queries that are mere paraphrases of previous ones. The model sometimes persists in asking similar questions even when the retriever fails to find relevant information, rather than exploring alternative avenues. This issue is particularly pronounced with smaller models such as Llama-3.2-1B. We will provide a more detailed analysis of these failure cases in the final version of the paper.

"Some relevant (albeit slightly older work was missed) eg [1] [1] Buck, Christian, et al. "Ask the right questions: Active question reformulation with reinforcement learning." arXiv preprint arXiv:1705.07830 (2017)."

Thanks for the pointer! We agree that reinforcement learning is a crucial area for developing more effective RAG agents. We will discuss this work in the related work section.

Thanks for your constructive feedback!

"CoRAG achieves an F1 of 69.8 on HotpotQA, whereas leaderboard models exceed 80.5 [1]. Similarly, on MuSiQue, CoRAG’s best F1 is 42.4, while top models score above 78.4 [2]"

The leaderboard scores you cited for HotpotQA and MuSiQue are based on distinct and easier settings, making them not directly comparable to our results. Specifically:

• The MuSiQue-Ans leaderboard and HotpotQA Distractor setting provide gold paragraphs within the context, simplifying the task to reading comprehension. In contrast, our work focuses on an open-domain QA setting, where no gold paragraphs are assumed.
• The HotpotQA Full-Wiki setting, by design, only requires retrieval from the abstracts of Wikipedia articles. Our approach does not incorporate this specific prior knowledge, as it does not generalize.

We will clarify these differences in the paper to provide a more accurate context for our evaluations.

"it would be helpful to include standard retrieval metrics (e.g., recall@k or supporting fact accuracy) to better assess the effectiveness of the proposed method in retrieving supporting paragraphs"

We have incorporated retrieval metrics for 9 KILT tasks in Table 7 of Appendix B. For multihop QA tasks, we calculated answer recall@{10, 20, 100} for your reference:

As demonstrated, CoRAG consistently and significantly improves recall across all datasets, underscoring its effectiveness in retrieving relevant information.

Notes: We followed DPR [1] in calculating recall based on answer matches, as not all datasets provide gold supporting paragraphs. For CoRAG, we utilized reciprocal rank fusion to merge multiple retrieval results into a single ranked list, from which recall was then calculated. For baselines, we use $E5_{\text{large}}$ to retrieve the top 100 passages.

[1] Karpukhin, V., Oğuz, B., Min, S., Lewis, P., Wu, L.Y., Edunov, S., Chen, D., & Yih, W. (2020). Dense Passage Retrieval for Open-Domain Question Answering. ArXiv, abs/2004.04906.

"Does the proposed framework improve performance when applied to other LLMs beyond LLaMA-3?"

Our proposed method is designed to be model-agnostic and can be seamlessly applied to various LLMs beyond LLaMA-3. For instance, we have successfully applied CoRAG to the Qwen3 family of models, observing similar performance improvements:

"The idea of multi-step, dynamic retrieval has been extensively studied in prior work, such as IRCoT, Self-RAG, ITER-RETGEN, and others."

We acknowledge that multi-step retrieval has been explored in prior works as you mentioned. However, to our knowledge, none of the existing literature has investigated the test-time scaling of RAG within a fine-tuning setting. We believe this is a crucial aspect for emerging applications like Deep Research.

An apt analogy would be the distinction between zero-shot CoT prompting and o1-like training for long CoT: while LLMs can generate CoT in a zero-shot manner, explicitly training LLMs to produce longer and higher-quality CoT traces still yields significant benefits.

"Is there any human evaluation to verify the validity of generated sub-questions and sub-answers?"

Our primary evaluation relies on automatic task metrics. However, for deeper analysis, we also conduct manual inspections of generated sub-questions and sub-answers. Representative examples are provided in Table 9 of the Appendix and in the supplementary material.

"In the ablation study on retrievers (Line 234), performance gains are attributed to stronger retrievers. Can the authors isolate the contribution of the CoRAG mechanism from that of the retrieval encoder?"

The ablation results concerning retrievers illustrate that stronger retrievers contribute to improved performance as an orthogonal factor, rather than being the primary source of CoRAG's gains. In our main results (Table 1), we maintained a consistent retriever (E5$_{\text{large}}$) across various decoding strategies and baselines, including SelfRAG and fine-tuned LLaMA-8B, to isolate and demonstrate the specific contributions of the CoRAG mechanism itself.

Dear Reviewer gW6e, thank you for your initial review.

We have submitted a rebuttal addressing your concerns, including the clarification of the leaderboard numbers you mentioned and additional results on recall metrics and generalization to other model families. We would greatly appreciate it if you could take a moment to review our response.

Thank you again for your time and consideration.

Thank you for your encouraging feedback!

"what model is used during the rejection sampling stage?"

We use the off-the-shelf instruct model (e.g., Llama-3.1-8B-Instruct) for rejection sampling. For iterative training, we use the model trained in the previous iteration. This approach allows the model to generate its own training data in a bootstrapping manner.

"It would bring more insights to the audience to conduct ablations over the training objective."

Thanks for the suggestion! One factor we studied is the data sampling ratio for $L_{\text{subquery}}$ and $L_{\text{subanswer}}$. We found that a sampling ratio of 0.2 works well and accelerates training. Increasing the sampling ratio to 1.0 resulted in only a negligible performance drop, as shown in the table below. Both models were evaluated with greedy decoding and a maximum of 6 retrieval steps. We will include this ablation in the revised paper.

"I would also like to see more experiments regarding the rejection sampling stage, e.g., how different models choices affect the final performance? How about other design choices compared to rejection sampling such as distilling from larger models or using larger models as a judge?"

In Table 3, we present an ablation study comparing the use of smaller LLMs (Llama-3.2-1B and Llama-3.2-3B) with GPT-4o for rejection sampling. The results align with expectations: larger models consistently lead to better performance. Using larger models as a judge is an interesting idea that we will explore in future work.