Research map / Label-efficient and robust learning / Recognition and long-tail
A Statistical Framework for Data-dependent Retrieval-Augmented Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a statistical framework for retrieval-augmented models with a learned retriever and predictor, trained end-to-end with a log-loss marginalized over the retriever distribution, and proves excess risk bounds that separate retriever and predictor contributions.
Problem
Retrieval-augmented models are hard to train end-to-end because retrieval is discrete, and their statistical properties (generalization, expressivity, role of data-store size) are poorly understood.
Approach
Write the natural objective as the expected prediction loss under the retriever-induced distribution over the data-store, show it upper-bounds the EMDR2 marginal likelihood via Jensen, and decompose excess risk into generalization, retriever approximation, and predictor approximation errors, yielding log dependence on data-store size.
Key results
- On NQ, jointly training predictor and retriever with Reverse Cross-Entropy + TopK reaches 46.4 exact match with large predictor and large retriever, versus 44.8 for the no-retriever-training baseline variants at base size.
- PDist achieves the highest retrieval recall, yet other objectives such as Reverse Cross-Entropy + TopK give better end-to-end exact match than PDist.
- Equivalent NQ accuracy (>=38.8) is reachable with large predictor plus small retriever (135 QPS), base-base (333 QPS), or small predictor plus large retriever (800 QPS).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add a retrieval component | predictor trained without retrieval · NQOpen | exact match accuracy: 29.1 -> 37.8 | Helped |
| add a retrieval component | predictor trained without retrieval · TriviaQA | exact match accuracy: 28.0 -> 49.7 | Helped |
| train the retriever while keeping the predictor fixed | RAM with a fixed retriever and trained predictor · NQOpen | exact match accuracy: 37.8 -> 45.0 | Helped |
| train the retriever while keeping the predictor fixed | RAM with a fixed retriever and trained predictor · TriviaQA | exact match accuracy: 49.7 -> 62.3 | Helped |
| jointly train the predictor and retriever using Reverse Cross-Entropy + TopK | RAM · NQOpen | exact match accuracy: 45.0 -> 46.4 | Helped |
| jointly train the predictor and retriever using Reverse Cross-Entropy + TopK | RAM · TriviaQA | exact match accuracy: 60.3 -> 62.4 | Mixed |
| use PDist instead of Reverse Cross-Entropy + TopK for joint training | RAM · TriviaQA | exact match accuracy: 62.4 -> 63.3 | Helped |
| use PDist to train the retriever | RAM · NQOpen and TriviaQA | recall | Helped |
Surprising findings
- Higher retrieval recall (PDist) does not translate into better end-to-end exact match than other objectives.
- Different predictor/retriever size combinations achieve the same accuracy at very different query throughput.
Limitations
- Analysis restricted to k=1 retrieved evidence per instance in the theory.
- Efficient gradient computation over the full data-store requires approximations (top-K or REINFORCE).
- Excess risk bounds rely on smoothness (Sobolev) assumptions on the loss gap and retrieval quality.
Experimental setup
Datasets: NaturalQuestions (NQOpen) TriviaQA chunked Wikipedia 2018
Benchmarks: NaturalQuestions (NQOpen) TriviaQA
Metrics: exact match recall of answer string in retrieved documents queries per second
Baselines: EMDR2 PDist Reverse Cross-Entropy + PG Reverse Cross-Entropy + TopK Fixed retriever training
Builds on: EMDR2 PDist REINFORCE RAG
Closest papers by meaning
- Making Retrieval-Augmented Language Models Robust to Irrelevant Context (ICLR 2024 · similarity 0.88)
- RECOMP: Improving Retrieval-Augmented LMs with Context Compression and Selective Augmentation (ICLR 2024 · similarity 0.86)
- USTAD: Unified Single-model Training Achieving Diverse Scores for Information Retrieval (ICML 2024 · similarity 0.86)
- Benchmarking Large Language Models in Retrieval-Augmented Generation (AAAI 2024 · similarity 0.85)
- Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models (ICLR 2025 · similarity 0.85)
- Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.