RAG-DDR: Optimizing Retrieval-Augmented Generation Using Differentiable Data Rewards

Weaknesses:

Reply to Weaknesses 1: Our proposed DDR method does not increase the computational cost. Both training and inference are performed entirely with white-box models, eliminating any overhead from external API calls. Furthermore, our inference process remains identical to the standard RAG pipeline, without introducing extra computational resource demands. The primary computational cost arises from calculating the reward using the rollout method. However, model inference remains efficient, and the training stage is completed quickly (without training with large amounts of data).

Reply to Weaknesses 2: For these multi-hop QA tasks, the one-time retrieval stage may limit the performance of RAG systems compared to the iterative reasoning methods. The reason may lie in that the one-time retriever may not supply sufficient information to answer the question [3]. Furthermore, RAG-DDR shows much better performance on ASQA, showing its effectiveness in training RAG systems to deal with complex reasoning tasks. The ASQA dataset asks models to explain the source of ambiguity in the question and connect all the valid short answers into a coherent passage, which also needs complex reasoning [2].

References:

[1] Ho X, Nguyen A K D, Sugawara S, et al. Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps. Proceedings of the 28th International Conference on Computational Linguistics.

[2] Stelmakh I, Luan Y, Dhingra B, et al. ASQA: Factoid Questions Meet Long-Form Answers. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. 2022: 8273-8288.

[3] Trivedi H, Balasubramanian N, Khot T, et al. Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions. In Proceedings of ACL. 2023.

Reply to Weaknesses 3: Thank you for your advice. In fact, we have conducted a comprehensive experimental analysis as shown in Figure 2, Table 3, Table 4, and Figure 6 to validate the effectiveness of our method. The experimental results reveal that DDR significantly enhances the utilization of external knowledge and mitigates the issue of knowledge conflict by aligning the data preferences across each module. We will add more theoretical analyses in the next version.

Dear Reviewer LZe5,

Thank you for your time and efforts again in reviewing our paper. We would like to gently remind you that the discussion period is nearing its conclusion. In our submitted rebuttal, we have carefully addressed all your comments and suggestions, and we hope our responses have effectively clarified any concerns or questions you have raised.

If there are any remaining points or if you require further clarification, please feel free to reach out. We sincerely appreciate your attention and valuable feedback.

Best regards,

Authors

Weaknesses

Reply to Weakness 1. The main weakness of the paper is inadequate comparison with related approaches on jointly tuning retriever and generator modules, which make it hard to appreciate the technical contribution of the paper.:

Thanks for your suggestion. We would like to take this opportunity to address and clarify some potential misunderstandings.

a. Our main innovation is to propose an effective method for optimizing multi-agents within the RAG system, rather than merely substituting the SFT loss with the DPO loss. Furthermore, we have conducted experiments in Figure 2 and Figure 7 to demonstrate the effectiveness of our training strategy in helping learn factual knowledge while preventing the loss of previously memorized information.

b. We compare RAG-DDR to the baseline model Filco [1] you mentioned. Please refer to the second point of our General Response. For the other baseline model, Adaptive-RAG [2], we are unable to provide its experimental results due to incomplete open source code.

References:

[1] Wang Z, Araki J, Jiang Z, et al. Learning to filter context for retrieval-augmented generation. 2023.

[2] Jeong S, Baek J, Cho S, et al. Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity. In Proceedings of NAACL. 2024.

Reply to Weakness 2. Other helpful ablations and baselines to experiment with:

Reply to Weakness 2.1: We have already conducted these experiments, as shown in Table 6. The results indicate that the performance of RAG-DDR ($V\text{KR} First$) and RAG-DDR ($V\text{Gen} First$) is indistinguishable across different datasets and $V_\text{Gen}$, demonstrating the robustness of the DDR method to training order variations.

Reply to Weakness 2.2: The challenge of balancing parametric and external knowledge is not unique to smaller LLMs but persists in larger models like ChatGPT. Studies [1] show that as models improve, they develop some awareness of knowledge conflicts and avoid blindly trusting external sources. However, larger LLMs, such as ChatGPT, often prioritize their incorrect internal knowledge over accurate external information, highlighting that this imbalance is a common issue across LLMs.

Due to the high cost of closed-source models, we utilize the larger open-source model Qwen2.5-14B-Instruction as the generation module. As shown in the table below, our proposed DDR method remains effective.

References:

[1] Jin Z, Cao P, Chen Y, et al. Tug-of-war between knowledge: Exploring and resolving knowledge conflicts in retrieval-augmented language models. 2024.

Reply to Weakness 3. The paper introduces the framework of differentiable data rewards (DDR). However, the framework is only evaluated for simple two agent systems. Consequently, DDR does not come across as a fully developed contribution. Instead, the key contribution, at least to me, is the use of DPO in jointly trying retrieval and generation modules. I would encourage the authors to build on this aspect:

Thank you for your suggestion. We conduct these experiments, please refer to the first point of our General Response.

Questions:

Reply to Question 1: In designing the KR module, we adopt the implementation of kotaemon (https://github.com/Cinnamon/kotaemon/blob/main/libs/kotaemon/kotaemon/indices/rankings/llm.py), and the prompts used for KR are provided in Figure 5.

Specifically, the KR module processes the $n$ retrieved documents in descending order of retrieval scores to assess their usefulness for answering the question. For each document, the module outputs “Yes” to retain it if deemed helpful or “No” to discard it. This continues until $k$ documents are retained, ensuring the top-$k$ most relevant documents are selected.

Reply to Question 2: During the training of the Knowledge Refinement module, we individually feed each document to the Generation Model to obtain the reward used for training the Knowledge Refinement module (line 216). This is because the Knowledge Refinement module determines whether a query is relevant to an individual document and generates Yes/No labels for each query-document pair. Equation 7 shows that the Generation Model is fed $k$ documents for sampling, which serves the same purpose as the function of the generation module—feeding $k$ documents and answering the query.

Reply to Question 3: During training of the Knowledge Refinement module, we randomly select one document with the highest reward score and one with the lowest reward score from the document collection to create preference data pairs for training the entire RAG system.

Reply to Question 4: Thank you for your suggestion and we will explore it in the future. Our main contribution lies in proposing a method to optimize the multi-agent of the RAG system. Therefore, we adhere to standard practices for module selection and design the KR module by following the implementations of kotaemon(https://github.com/Cinnamon/kotaemon/blob/main/libs/kotaemon/kotaemon/indices/rankings/llm.py).

Dear Reviewer 1cKK,

Thank you for your time and efforts again in reviewing our paper. We would like to gently remind you that the discussion period is nearing its conclusion. In our submitted rebuttal, we have carefully addressed all your comments and suggestions, and we hope our responses have effectively clarified any concerns or questions you have raised.

If there are any remaining points or if you require further clarification, please feel free to reach out. We sincerely appreciate your attention and valuable feedback.

Best regards,

Authors

Thank you once again for your insightful review. Your comments provide valuable feedback and offer us an opportunity to improve our work.

We will clarify why the DDR approach outperforms baselines like Filco. Filco optimizes the context selector and generation model independently using SFT. While Falco’s context selector effectively filters out irrelevant information, the retained content may not always benefit the generation model. This is because the context selector in Filco is trained using labeled data rather than feedback from the downstream generation model, causing its outputs to align more with the training data than with the knowledge needs of the generation model. As highlighted by works like AAR [1] and REPLUG [2], incorporating feedback signals from the generation model to train the retrieval model can better align them, ultimately providing higher-quality documents for generation.

References

[1] Yu Z, Xiong C, Yu S, et al. Augmentation-Adapted Retriever Improves Generalization of Language Models as Generic Plug-In. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics. 2023.

[2] Shi W, Min S, Yasunaga M, et al. REPLUG: Retrieval-Augmented Black-Box Language Models. Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers). 2024.

Weaknesses:

Reply to Weakness 1: Thank you for your suggestion. We will revise and update Figure 1 to address these issues in the next version of the paper.

Reply to Weakness 2:

Reply to Weakness 2.1: Thank you for your valuable suggestion. In response, we have conducted experiments with large-scale LLMs. Specifically, we utilized Qwen2.5-14B-Instruct as the generation module and Llama3-8B-Instruct as the refinement module to build a new combination for the RAG system.

Reply to Weakness 2.2: We have implemented the RAG model using the following combinations: (Llama3-8B-Instruct + MiniCPM-2.4B-SFT), (Llama3-8B-Instruct + Llama3-8B-Instruct), and (Llama3-8B-Instruct + Qwen2.5-14B-Instruction). In future work, we aim to explore additional model configurations to further enhance performance.

Questions:

Reply to Question 1: Please refer to the weakness section.

Reply to Question 2: Data preferences refer to the specific input and output data that each module in a RAG system depends on, enabling the entire system to generate more accurate responses. Specifically, each module processes input data and generates output data, which are then passed to subsequent modules. The preferred data for each module are those inputs and outputs that contribute to producing correct answers across the entire RAG system. The DDR training objective seeks to align the input and output data of each module with these preferred data, thereby optimizing the overall performance of the RAG system.

Reply to Question 3:

Reply to Question 3.1: In our experiments, the DDR method consists of two stages: data forward propagation and differentiable data reward. These stages aim to construct the training data and train the RAG module, respectively. For example, during training the generation module (Llama3-8B-Instruct), data forward propagation takes approximately 35 hours, while differentiable data reward requires around 1 hour and 40 minutes on an 8-GPU A100 server. During training the knowledge module (Llama3-8B-Instruct), data forward propagation takes about 12 hours, with differentiable data reward taking approximately 1 hour and 50 minutes on the same setup.

Although the data forward propagation process is time-intensive for generating training data, the actual model training is relatively quick. Moreover, DDR method does not require a large amount of training data and training time. As shown in the table below, using the training data described in this paper, Llama3-8B-Instruct and MiniCPM-2.4B-SFT achieve convergence of dev loss with less than one epoch's worth of training data.

Reply to Question 3.2: Thank you for your suggestion. We conduct these experiments, please refer to the first point of our General Response.

Reply to Question 4: MiniCPM-2.4B-SFT is a smaller parameter model with weaker mathematical reasoning capabilities compared to larger models. We believe its improved performance may be attributed to the inclusion of Math QA and reasoning tasks in the collected training data.

Reply to Question 5: Thank you for your suggestion. We will address this spelling error in the next version of the paper.

Dear Reviewer QfBP,

Thank you for your time and efforts again in reviewing our paper. We would like to gently remind you that the discussion period is nearing its conclusion. In our submitted rebuttal, we have carefully addressed all your comments and suggestions, and we hope our responses have effectively clarified any concerns or questions you have raised.

If there are any remaining points or if you require further clarification, please feel free to reach out. We sincerely appreciate your attention and valuable feedback.

Best regards,

Authors

Thank you for your reply. I have read your reply, I believe most of my concerns have been addressed, and I will raise my score.

Additionally, when comparing Qwen2.5-14B with the other models in Table 1, I noticed that Qwen2.5-14B performs significantly better than LLaMA3-8B in the vanilla RAG setting. However, after training, their performances are comparable. Based on my experience, I personally believe this discrepancy is due to the distinct “tunability” of the base models (e.g., differences in how they learn from fine-tuning data, catastrophic forgetting, etc.). Thus, this is a minor concern on my part.

However, going further and considering your General Response 1, I am curious about the major bottlenecks in the framework that limit its performance. I understand that exploring this experimentally involves significant costs, but I hope the authors could provide some discussion on this topic and potentially explore it further in the future, if feasible.

Thank you once again for your insightful review. Your comments provide valuable feedback and offer us an opportunity to improve our work.

• As you mentioned, Qwen2.5-14B outperforms LLaMA3-8B in the vanilla RAG setting for tasks like NQ but becomes comparable to LLaMA3-8B after training. This discrepancy may stem from inherent differences between the base models. Furthermore, the preference data generated during DPO training can vary significantly across models, potentially affecting generation results. We plan to explore this further by investigating a wider range of model combinations in future work.

• The primary bottleneck of the DDR approach lies in the time cost of generating training data. To maximize the training effect of DDR, we use a rollout method to compute reward scores for each module, ensuring they align with the RAG system’s final output. While this increases the data construction time, the process is conducted offline and does not affect training efficiency.

Weaknesses:

Reply to Weakness 1: Please refer to the Question section.

Reply to Weakness 2: Thank you for your comments and suggestions. We will carefully proofread our paper.

Questions:

Reply to Question 1: Thank you for your suggestion. We have clarified the differences between STEP-DPO [1] and DDR in the final paragraph of the Related Work section and will further highlight this distinction in the next version of our paper.

Specifically, STEP-DPO focuses on optimizing a single agent by splitting the model’s responses into multiple steps and evaluating the correctness of each step individually. In contrast, DDR is tailored for optimizing a multi-agent system, employing a rollout mechanism to gather system-level rewards for each agent and subsequently refining the agents based on these rewards.

References:

[1] Lai X, Tian Z, Chen Y, et al. Step-dpo: Step-wise preference optimization for long-chain reasoning of llms. 2024.

Reply to Question 2: Thank you for your insightful suggestion. To further examine the effectiveness of multi-round optimization, we conducted an analysis of RAG-DDR (2-Round) across the same three scenarios presented in Table 3 of the paper. As illustrated in the table below, RAG-DDR (2-Round) consistently outperforms RAG-DDR (1-Round) across all three scenarios, with particularly significant gains observed in the Internal Knowledge scenario. These findings indicate that 2-round DDR training can further improve the model’s ability to balance internal and external knowledge, especially for larger models.

Reply to Question 3: Thank you for your valuable suggestion. We use positive samples from RAG-DDR to train RA-DIT. Then, we evaluate the downstream performance of RA-DIT and analyze its response length. The corresponding results are presented in the tables below.

a. The model performance is shown in the below table. RAG-DDR consistently outperforms RA-DIT across training setups, whether using ground-truth or DPO-positive samples, highlighting the effectiveness of our DDR training method. Furthermore, RA-DIT trained on DPO-positive samples performs worse than when trained on ground-truth data, suggesting that aligning solely with noisy positive examples is insufficient without the contrast provided by negatives. This emphasizes the advantage of our optimization approach leveraging the DPO loss.

b. The response lengths are presented in the table below. Notably, the response length of RA-DIT trained with DPO-positive samples is longer than that of RA-DIT trained with ground truth. This observation further shows that the SFT training method may lead LLMs to overfit the supervised data, thereby affecting the model's response length.

Dear Reviewer oqk9,

Thank you for your time and efforts again in reviewing our paper. We would like to gently remind you that the discussion period is nearing its conclusion. In our submitted rebuttal, we have carefully addressed all your comments and suggestions, and we hope our responses have effectively clarified any concerns or questions you have raised.

If there are any remaining points or if you require further clarification, please feel free to reach out. We sincerely appreciate your attention and valuable feedback.

Best regards,

Authors

Reply to Weakness 1: Thank you for your comments. We would like to take this opportunity to address and clarify some potential misunderstandings.

a. In our experiments, all models are implemented using the same RAG modules, including the Knowledge Refinement module and Generation Module. RA-DIT also employs the Knowledge Refinement module ($V_\text{KR}$) during evaluation. Thus, we conduct a fair comparison.

b. The notation in Table 2 may cause some misunderstanding; we will revise our paper to ensure all details are presented clearly. In fact, the model performance of the model $w/ V_\text{KR}$ in Table 2 is identical to that of the Vanilla RAG model in Table 1.

Reply to Weakness 2: Thanks for your suggestion. We have included additional baselines for comparison, as detailed in the second point of our General Response.

Dear Reviewer FTLt,

Thank you for your time and efforts again in reviewing our paper. We would like to gently remind you that the discussion period is nearing its conclusion. In our submitted rebuttal, we have carefully addressed all your comments and suggestions, and we hope our responses have effectively clarified any concerns or questions you have raised.

If there are any remaining points or if you require further clarification, please feel free to reach out. We sincerely appreciate your attention and valuable feedback.

Best regards,

Authors

I appreciate the author's efforts in providing detailed clarifications and new experimental results (though Yoran et al. (ICLR 2024) is still missing in the comparison) as additional support, which addressed most of my concerns. To reflect this, I will raise my score.

However, I still believe the positioning of the paper could be improved, since some closely relevant works are not sufficiently discussed in the original manuscript as mentioned earlier. For example, Yoran et al. (ICLR 2024) was published before July 1, 2024, and their code is even publicly available . Unfortunately, there is no comparison or discussion of this work in the paper.

It's good to see the new results and to know the author's plan for revising the manuscript. Please ensure the updated version incorporates these discussions and results to better position and justify the contribution of this paper.

Thank you once again for your insightful review. Your comments provide valuable feedback and offer us an opportunity to improve our work.

For the missing baseline you mentioned, we attempted to reproduce it. Unfortunately, our reproduction was hampered by the lack of some experimental details in the open source code. Other researchers have encountered problems similar to ours [1]. We are trying to solve these problems and reproduce the baseline. For the updated experimental results in the rebuttal, we will add them to subsequent versions of the paper.

References:

[1] https://github.com/oriyor/ret-robust/issues/5