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

ACM Transactions on Information Systems, vol. 44, pp. 1–30

Abstract

Retrieval-augmented Generation (RAG) integrates Large Language Models (LLMs) 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 LLMs 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 ( DRO ) framework 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 qualitatively analyze the stability, convergence, and variance of DRO. (Code is available on DRO GitHub ).

Authors 10

1. Zhengliang Shi Aachen

   Leiden University

   Affiliation as printed

   Leiden University, Leiden, The Netherlands

2. Lingyong Yan

   Baidu (China)

   Affiliation as printed

   Baidu Inc., Beijing, China

3. Weiwei Sun

   Carnegie Mellon University

   Affiliation as printed

   LTI, Carnegie Mellon University, Pittsburgh, Pennsylvania, USA

4. Yue Feng

   University of Birmingham

   Affiliation as printed

   School of Computer Science, University of Birmingham, Birmingham, UK

5. Pengjie Ren

   Shandong University

   Affiliation as printed

   Shandong University, Qingdao, China

6. Xinyu Ma

   Baidu (China)

   Affiliation as printed

   Baidu Inc., Beijing, China

7. Shuaiqiang Wang

   Baidu (China)

   Affiliation as printed

   Baidu Inc., Beijing, China

8. Dawei Yin

   Baidu (China)

   Affiliation as printed

   Baidu Inc., Beijing, China

9. Maarten de Rijke

   University of Amsterdam

   Affiliation as printed

   University of Amsterdam, Amsterdam, The Netherlands

10. Zhaochun Ren Aachen

    Leiden University

    Affiliation as printed

    Leiden University, Leiden, The Netherlands