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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Shi, Zhengliang, Yan, Lingyong, Sun, Weiwei, Feng, Yue, Ren, Pengjie, Ma, Xinyu, Wang, Shuaiqiang, Yin, Dawei, de Rijke, Maarten, Ren, Zhaochun
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915275246403584
author Shi, Zhengliang
Yan, Lingyong
Sun, Weiwei
Feng, Yue
Ren, Pengjie
Ma, Xinyu
Wang, Shuaiqiang
Yin, Dawei
de Rijke, Maarten
Ren, Zhaochun
author_facet Shi, Zhengliang
Yan, Lingyong
Sun, Weiwei
Feng, Yue
Ren, Pengjie
Ma, Xinyu
Wang, Shuaiqiang
Yin, Dawei
de Rijke, Maarten
Ren, Zhaochun
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_03075
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models
Shi, Zhengliang
Yan, Lingyong
Sun, Weiwei
Feng, Yue
Ren, Pengjie
Ma, Xinyu
Wang, Shuaiqiang
Yin, Dawei
de Rijke, Maarten
Ren, Zhaochun
Information Retrieval
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.
title Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models
topic Information Retrieval
url https://arxiv.org/abs/2505.03075