# Stochastic RAG: End-to-End Retrieval-Augmented Generation through Expected Utility Maximization

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

## Key facts

- Cited by 29 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 37 of its 51 references that are in the map.
- First posted on arXiv 5 May 2024.

## 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 37)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [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)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [FEVER: a large-scale dataset for Fact Extraction and VERification](https://mlsanity.com/paper/1803.05355/) (2018)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (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)
- [LaMP: When Large Language Models Meet Personalization](https://mlsanity.com/paper/2304.11406/) (2023)
- [Zero-Shot Relation Extraction via Reading Comprehension](https://mlsanity.com/paper/1706.04115/) (2017)
- [FreshLLMs: Refreshing Large Language Models with Search Engine Augmentation](https://mlsanity.com/paper/2310.03214/) (2023)
- [On the Planning Abilities of Large Language Models : A Critical Investigation](https://mlsanity.com/paper/2305.15771/) (2023)
- [Distilling Knowledge from Reader to Retriever for Question Answering](https://mlsanity.com/paper/2012.04584/) (2020)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Retrieving and Reading: A Comprehensive Survey on Open-domain Question Answering](https://mlsanity.com/paper/2101.00774/) (2021)
- [A Survey on Retrieval-Augmented Text Generation](https://mlsanity.com/paper/2202.01110/) (2022)
- [Autoregressive Search Engines: Generating Substrings as Document Identifiers](https://mlsanity.com/paper/2204.10628/) (2022)
- [Optimization Methods for Personalizing Large Language Models through Retrieval Augmentation](https://mlsanity.com/paper/2404.05970/) (2024)
- [Evaluating Retrieval Quality in Retrieval-Augmented Generation](https://mlsanity.com/paper/2404.13781/) (2024)
- [Retrieve and Refine: Improved Sequence Generation Models For Dialogue](https://mlsanity.com/paper/1808.04776/) (2018)
- [End-to-End Training of Neural Retrievers for Open-Domain Question Answering](https://mlsanity.com/paper/2101.00408/) (2021)
- [FiD-Light: Efficient and Effective Retrieval-Augmented Text Generation](https://mlsanity.com/paper/2209.14290/) (2022)
- [Retrieval-Enhanced Machine Learning](https://mlsanity.com/paper/2205.01230/) (2022)

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

- [Searching for Best Practices in Retrieval-Augmented Generation](https://mlsanity.com/paper/2407.01219/) (2024)
- [Improving Retrieval-Augmented Generation through Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2501.15228/) (2025)
- [Retrieval-Augmented Generation: A Comprehensive Survey of Architectures, Enhancements, and Robustness Frontiers](https://mlsanity.com/paper/2506.00054/) (2025)
- [Claim Verification in the Age of Large Language Models: A Survey](https://mlsanity.com/paper/2408.14317/) (2024)
- [Towards Fair RAG: On the Impact of Fair Ranking in Retrieval-Augmented Generation](https://mlsanity.com/paper/2409.11598/) (2024)
- [Learning to Rank for Multiple Retrieval-Augmented Models through Iterative Utility Maximization](https://mlsanity.com/paper/2410.09942/) (2024)
- [Accelerating Retrieval-Augmented Generation](https://mlsanity.com/paper/2412.15246/) (2024)
- [Training a Utility-based Retriever Through Shared Context Attribution for Retrieval-Augmented Language Models](https://mlsanity.com/paper/2504.00573/) (2025)
- [Automatic Database Configuration Debugging using Retrieval-Augmented Language Models](https://mlsanity.com/paper/2412.07548/) (2024)
- [Retrieval-Enhanced Machine Learning: Synthesis and Opportunities](https://mlsanity.com/paper/2407.12982/) (2024)
- [Iterative Self-Incentivization Empowers Large Language Models as Agentic Searchers](https://mlsanity.com/paper/2505.20128/) (2025)
- [Agentic-R: Learning to Retrieve for Agentic Search](https://mlsanity.com/paper/2601.11888/) (2026)
- [Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models](https://mlsanity.com/paper/2505.03075/) (2025)
- [Mitigating Bias in RAG: Controlling the Embedder](https://mlsanity.com/paper/2502.17390/) (2025)
- [The Next Phase of Scientific Fact-Checking: Advanced Evidence Retrieval from Complex Structured Academic Papers](https://mlsanity.com/paper/2506.20844/) (2025)
- [Utility-Focused LLM Annotation for Retrieval and Retrieval-Augmented Generation](https://mlsanity.com/paper/2504.05220/) (2025)
- [LLM-Specific Utility for Retrieval-Augmented Generation](https://mlsanity.com/paper/2510.11358/) (2025)
- [Learning to Retrieve from Agent Trajectories](https://mlsanity.com/paper/2604.04949/) (2026)
- [Reinforcement Fine-Tuning for History-Aware Dense Retriever in RAG](https://mlsanity.com/paper/2602.03645/) (2026)
- [How Retrieved Context Shapes Internal Representations in RAG](https://mlsanity.com/paper/2602.20091/) (2026)
- [Binary Hypothesis Testing for Softmax Models and Leverage Score Models](https://mlsanity.com/paper/2405.06003/) (2024)
- [Oreo: A Plug-in Context Reconstructor to Enhance Retrieval-Augmented Generation](https://mlsanity.com/paper/2502.13019/) (2025)
- [RAGtifier: Evaluating RAG Generation Approaches of State-of-the-Art RAG Systems for the SIGIR LiveRAG Competition](https://mlsanity.com/paper/2506.14412/) (2025)
- [Improving End-to-End Training of Retrieval-Augmented Generation Models via Joint Stochastic Approximation](https://mlsanity.com/paper/2508.18168/) (2025)
- [FLOWREADER: Min-Cost Flow Optimization for Multi-Modal Long Document Q&A](https://mlsanity.com/paper/2606.07235/) (2026)
- [Leveraging Retrieval-Augmented Generation for Persian University Knowledge Retrieval](https://mlsanity.com/paper/2411.06237/) (2024)
- [Towards Efficient Educational Chatbots: Benchmarking RAG Frameworks](https://mlsanity.com/paper/2503.00781/) (2025)
- [Plan-and-Refine: Diverse and Comprehensive Retrieval-Augmented Generation](https://mlsanity.com/paper/2504.07794/) (2025)
- [Critic-R: Improving Agentic Search using Instruction-tuned Retrievers with Natural Language Introspective Feedback](https://mlsanity.com/paper/2606.00590/) (2026)

## Abstract (arXiv)

This paper introduces Stochastic RAG--a novel approach for end-to-end optimization of retrieval-augmented generation (RAG) models that relaxes the simplifying assumptions of marginalization and document independence, made in most prior work. Stochastic RAG casts the retrieval process in RAG as a stochastic sampling without replacement process. Through this formulation, we employ straight-through Gumbel-top-k that provides a differentiable approximation for sampling without replacement and enables effective end-to-end optimization for RAG. We conduct extensive experiments on seven diverse datasets on a wide range of tasks, from open-domain question answering to fact verification to slot-filling for relation extraction and to dialogue systems. By applying this optimization method to a recent and effective RAG model, we advance state-of-the-art results on six out of seven datasets.