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

arXiv: https://arxiv.org/abs/2505.03075 · HTML: https://mlsanity.com/paper/2505.03075/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Placed in RAG pipeline design and evaluation, one of 15 directions in Retrieval-augmented generation.
- Builds on 55 of its 83 references that are in the map.
- First posted on arXiv 5 May 2025.

## Placed in

- Retrieval-augmented generation / [RAG pipeline design and evaluation](https://mlsanity.com/rag/rag-pipeline-design-and-evaluation/)

## Builds on (most cited first, 30 of 55)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [A Survey on Hallucination in Large Language Models: Principles, Taxonomy, Challenges, and Open Questions](https://mlsanity.com/paper/2311.05232/) (2023)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [Wizard of Wikipedia: Knowledge-Powered Conversational agents](https://mlsanity.com/paper/1811.01241/) (2018)
- [ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction](https://mlsanity.com/paper/2112.01488/) (2021)
- [A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2405.06211/) (2024)
- [AmbigQA: Answering Ambiguous Open-domain Questions](https://mlsanity.com/paper/2004.10645/) (2020)
- [Document Ranking with a Pretrained Sequence-to-Sequence Model](https://mlsanity.com/paper/2003.06713/) (2020)
- [Making Retrieval-Augmented Language Models Robust to Irrelevant Context](https://mlsanity.com/paper/2310.01558/) (2023)
- [Enhancing Retrieval-Augmented Large Language Models with Iterative Retrieval-Generation Synergy](https://mlsanity.com/paper/2305.15294/) (2023)
- [Chain-of-Verification Reduces Hallucination in Large Language Models](https://mlsanity.com/paper/2309.11495/) (2023)
- [RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation](https://mlsanity.com/paper/2310.04408/) (2023)
- [RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs](https://mlsanity.com/paper/2407.02485/) (2024)
- [RA-DIT: Retrieval-Augmented Dual Instruction Tuning](https://mlsanity.com/paper/2310.01352/) (2023)

## Built on by (most cited first, 9 of 9)

- [DeepShop: A Benchmark for Deep Research Shopping Agents](https://mlsanity.com/paper/2506.02839/) (2025)
- [CLaRa: Bridging Retrieval and Generation with Continuous Latent Reasoning](https://mlsanity.com/paper/2511.18659/) (2025)
- [Learning When to Remember: Risk-Sensitive Contextual Bandits for Abstention-Aware Memory Retrieval in LLM-Based Coding Agents](https://mlsanity.com/paper/2604.27283/) (2026)
- [Reflective Personalization Optimization: A Post-hoc Rewriting Framework for Black-Box Large Language Models](https://mlsanity.com/paper/2511.05286/) (2025)
- [Learning What Helps: Task-Aligned Context Selection for Vision Tasks](https://mlsanity.com/paper/2512.00489/) (2025)
- [Tape: A Cellular Automata Benchmark for Evaluating Rule-Shift Generalization in Reinforcement Learning](https://mlsanity.com/paper/2601.04695/) (2026)
- [RASP-Tuner: Retrieval-Augmented Soft Prompts for Context-Aware Black-Box Optimization in Non-Stationary Environments](https://mlsanity.com/paper/2604.18026/) (2026)
- [From Intent to Execution: Composing Agentic Workflows with Agent Recommendation](https://mlsanity.com/paper/2605.03986/) (2026)
- [BRIDGE: Bilevel Retrieval-Credit-Aware Agentic Reinforcement Learning](https://mlsanity.com/paper/2609.36505/) (2026)

## Abstract (arXiv)

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.