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

TL;DR

This work introduces Differentiable Data Rewards (DDR) to end-to-end optimize Retrieval-Augmented Generation (RAG) by aligning data preferences across RAG modules using rollout-based rewards and Direct Preference Optimization. DDR jointly tunes a knowledge refinement module and a generation module, using perturbations and system-level rewards to overcome knowledge conflicts between external knowledge and parametric memory. Empirical results show DDR outperforms supervised fine-tuning baselines, especially for smaller LLMs, and exhibits robustness to noisy documents, with ablations confirming the central importance of the generation module. The approach generalizes across tasks and model scales, offering a practical path to more reliable RAG systems with better information extraction and calibrated output length.

Abstract

Retrieval-Augmented Generation (RAG) has proven its effectiveness in mitigating hallucinations in Large Language Models (LLMs) by retrieving knowledge from external resources. To adapt LLMs for the RAG systems, current approaches use instruction tuning to optimize LLMs, improving their ability to utilize retrieved knowledge. This supervised fine-tuning (SFT) approach focuses on equipping LLMs to handle diverse RAG tasks using different instructions. However, it trains RAG modules to overfit training signals and overlooks the varying data preferences among agents within the RAG system. In this paper, we propose a Differentiable Data Rewards (DDR) method, which end-to-end trains RAG systems by aligning data preferences between different RAG modules. DDR works by collecting the rewards to optimize each agent in the RAG system with the rollout method, which prompts agents to sample some potential responses as perturbations, evaluates the impact of these perturbations on the whole RAG system, and subsequently optimizes the agent to produce outputs that improve the performance of the RAG system. Our experiments on various knowledge-intensive tasks demonstrate that DDR significantly outperforms the SFT method, particularly for LLMs with smaller-scale parameters that depend more on the retrieved knowledge. Additionally, DDR exhibits a stronger capability to align the data preference between RAG modules. The DDR method makes the generation module more effective in extracting key information from documents and mitigating conflicts between parametric memory and external knowledge. All codes are available at https://github.com/OpenMatch/RAG-DDR.

Figures (7)

• Figure 1: The Illustration of End-to-End Retrieval-Augmented Generation (RAG) Training with Our Differentiable Data Reward (DDR) Method. During training, we iteratively optimize the Generation module ($V_\text{Gen}$) and Knowledge Refinement module ($V_\text{KR}$).
• Figure 2: Characteristics of the Generation Module in RAG Optimized with Different Training Strategies. We use MiniCPM-2.4B to build the generation module ($V_\text{Gen}$) and then train it using different strategies. The performance of $V_\text{Gen}$ is shown in a zero-shot setting, along with the generation module optimized using the RA-DIT and DDR methods.
• Figure 3: Effectiveness of Different RAG Models in Defending Noisy Information. We use MiniCPM-2.4B to build the generation module ($V_\text{Gen}$). Then we retain one informative document and randomly replace $n$ top-retrieved documents with noisy ones.
• Figure 4: Prompt Templates of Different Training and Evaluating Tasks.
• Figure 5: Prompt Templates Used for Knowledge Refinement.