Research map / Large language models / Retrieval and knowledge
Inference Scaling for Long-Context Retrieval Augmented Generation
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.
Studies inference scaling for long-context RAG with demonstration-based (DRAG) and iterative demonstration-based (IterDRAG) strategies on Gemini 1.5 Flash, finding near-linear gains in performance with effective context length and fitting a computation allocation model.
Problem
Long-context RAG performance plateaus when simply adding more retrieved documents, and it is unclear how to allocate inference compute (documents, in-context examples, iterations) to maximize knowledge-intensive QA performance.
Approach
Two inference strategies, DRAG (many retrieved documents plus many-shot RAG demonstrations in one call) and IterDRAG (query decomposition into Self-Ask sub-queries with interleaved retrieval), are used to scale effective context length; the optimal performance is modeled with a computation allocation model fit by least squares on log-parameters with an inverse sigmoidal output mapping.
Key results
- IterDRAG with Gemini 1.5 Flash reaches 65.6 EM on Bamboogle and 22.2 EM on MuSiQue at 1M effective context, versus 44.0 and 14.5 for DRAG-like best configs at 16k budgets (Table 1 comparisons across budgets).
- Optimal performance of DRAG and IterDRAG rises nearly linearly with the order of magnitude of effective context length, while standard RAG plateaus around 128k tokens.
- The computation allocation model with inverse sigmoid mapping achieves R^2 = 0.903 (MSE 0.085) and reaches 96.6% of oracle performance in domain generalization at 1M.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use IterDRAG instead of DRAG | Gemini 1.5 Flash RAG · Bamboogle, 128k | Acc: 54.4 -> 68.8 | Helped |
| use IterDRAG instead of DRAG | Gemini 1.5 Flash RAG · HotpotQA, 128k | Acc: 52.2 -> 52.8 | Helped |
| use IterDRAG instead of DRAG | Gemini 1.5 Flash RAG · MuSiQue, 128k | Acc: 17.9 -> 24.5 | Helped |
| use IterDRAG instead of DRAG | Gemini 1.5 Flash RAG · 2WikiMultiHopQA, 128k | Acc: 53.1 -> 74.6 | Helped |
| increase effective context budget from 1M to 5M | IterDRAG · HotpotQA | Acc: 55.3 -> 56.4 | Helped |
| increase effective context budget from 1M to 5M | IterDRAG · MuSiQue | Acc: 30.5 -> 30.5 | No clear effect |
| increase effective context budget from 1M to 5M | IterDRAG · 2WikiMultiHopQA | Acc: 76.4 -> 76.9 | Helped |
| add b and task informativeness i to the model (ablation) | Computation allocation model · Computation allocation model | R²: 0.866 -> 0.903 | Helped |
| add b and task informativeness i to the model (ablation) | Computation allocation model · Computation allocation model | MSE: 0.116 -> 0.085 | Helped |
| use linear scaling of P instead of sigmoid scaling (ablation) | Computation allocation model · Computation allocation model | R²: 0.903 -> 0.876 | Hurt |
| use linear scaling of P instead of sigmoid scaling (ablation) | Computation allocation model · Computation allocation model | MSE: 0.085 -> 0.109 | Hurt |
| use a quadratic form of log(θ) instead of the model's standard form (ablation) | Computation allocation model · Computation allocation model | R²: 0.903 -> 0.867 | Hurt |
Surprising findings
- Adding more retrieved documents beyond soft thresholds (around 100-500) can yield diminishing returns or declines, as retrieval NDCG/MRR plateau while recall keeps rising.
- Chain-of-thought with free-form reasoning underperforms IterDRAG substantially, which the authors partly attribute to Gemini 1.5 Flash being relatively small.
Limitations
- Gains become less predictable and smaller beyond 1M effective context; 5M extrapolation error is 5.6%.
- Experiments use a single proprietary model (Gemini 1.5 Flash) and no reranking component; the interaction with refinement remains unknown.
- Evaluation relies on EM/F1/Acc metrics that the authors note can misscore abbreviations and refusals.
Experimental setup
Datasets: Bamboogle HotpotQA MuSiQue 2WikiMultiHopQA TriviaQA Natural Questions StrategyQA Wikipedia (KILT)
Benchmarks: Bamboogle HotpotQA MuSiQue 2WikiMultiHopQA TriviaQA Natural Questions StrategyQA
Metrics: exact match F1 accuracy
Baselines: zero-shot QA many-shot QA standard RAG chain-of-thought
Builds on: Self-Ask many-shot in-context learning Gecko embedding KILT
Closest papers by meaning
- RAPID: Long-Context Inference with Retrieval-Augmented Speculative Decoding (ICML 2025 · similarity 0.94)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.94)
- Accelerating Inference of Retrieval-Augmented Generation via Sparse Context Selection (ICLR 2025 · similarity 0.92)
- Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting (ICLR 2025 · similarity 0.91)
- Influence Guided Context Selection for Effective Retrieval-Augmented Generation (NeurIPS 2025 · similarity 0.91)
- Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.90)
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.