Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models

TL;DR

This work introduces Direct Retrieval-augmented Optimization (DRO), which jointly trains a list-wise knowledge selector and an LLM generator for end-to-end retrieval-augmented generation (RAG). By treating document permutations as latent variables and using an EM-style alternation with importance sampling, DRO maximizes the evidence lower bound to align knowledge selection with answer generation. The authors provide theoretical links to policy-gradient methods, demonstrate variance reduction with normalized importance weights, and prove convergence of the training process. Empirically, DRO yields 5–15% gains in EM and F1 across five QA datasets and improves document recall, validating stable optimization dynamics and practical impact for knowledge-grounded NLP tasks.

Abstract

Retrieval-augmented generation (RAG) integrates large language models ( LLM s) with retrievers to access external knowledge, improving the factuality of LLM generation in knowledge-grounded tasks. To optimize the RAG performance, most previous work independently fine-tunes the retriever to adapt to frozen LLM s or trains the LLMs to use documents retrieved by off-the-shelf retrievers, lacking end-to-end training supervision. Recent work addresses this limitation by jointly training these two components but relies on overly simplifying assumptions of document independence, which has been criticized for being far from real-world scenarios. Thus, effectively optimizing the overall RAG performance remains a critical challenge. We propose a direct retrieval-augmented optimization framework, named DRO, that enables end-to-end training of two key components: (i) a generative knowledge selection model and (ii) an LLM generator. DRO alternates between two phases: (i) document permutation estimation and (ii) re-weighted maximization, progressively improving RAG components through a variational approach. In the estimation step, we treat document permutation as a latent variable and directly estimate its distribution from the selection model by applying an importance sampling strategy. In the maximization step, we calibrate the optimization expectation using importance weights and jointly train the selection model and LLM generator. Our theoretical analysis reveals that DRO is analogous to policy-gradient methods in reinforcement learning. Extensive experiments conducted on five datasets illustrate that DRO outperforms the best baseline with 5%-15% improvements in EM and F1. We also provide in-depth experiments to qualitatively analyze the stability, convergence, and variance of DRO.

Figures (7)

• Figure 1: Overview of DRO objective. The selection model directly estimate a document permutation for the generator to predict an answer, with both components trained jointly.
• Figure 2: The overall framework for DRO alternates between the (i) E-step: document permutation estimation (Section \ref{['sec:estimation']}); and (ii) M-step: re-weighted maximization (Section \ref{['sec:maximization']}) to progressively optimize the holistic RAG performance.
• Figure 3: Recall@K (k=1, 3, 5) score of the initial retrieval (Colbertv2.0), two re-ranking baselines (i.e., RankVicuna and RankZephyer) and our selection model $\theta_{s}$, respectively.
• Figure 4: Variance during the training process of our method (logarithmic scale).
• Figure 5: Ablation study on five datasets to demonstrate the effectiveness of training the selection model $\theta_s$ and generator $\theta_g$.

Theorems & Definitions (3)

• Lemma 1
• Remark 7.1
• Remark 7.2