PaperlayerView exact revision on arXiv

Source-linked AI summary

Inference Scaling for Long-Context Retrieval Augmented Generation

Zhenrui Yue, Honglei Zhuang, Aijun Bai, Kai Hui, Rolf Jagerman, Hansi Zeng, Zhen Qin, Dong Wang, Xuanhui Wang, Michael Bendersky

arXiv:2410.04343v2cs.CL

TL;DR

Long-context RAG faces limits when more retrieved knowledge alone does not ensure effective use of context. The paper introduces DRAG and IterDRAG and models compute allocation; under optimal configurations, RAG performance improves almost linearly as test-time compute grows.

  • Problem

    Prior RAG scaling emphasizes more or longer retrieved documents, while inference scaling for long-context RAG remains challenging because relevant information can be hard to locate and excess context can plateau or hurt performance.

    Source: §1 · pp. 1–2

  • Method

    The paper introduces demonstration-based DRAG and iterative IterDRAG, then models performance across inference configurations to estimate optimal compute allocations.

    Source: §1 · pp. 2–3

  • Results

    With optimal inference parameters, long-context RAG performance improves almost linearly with increasing orders of magnitude of test-time compute, and estimated optimal configurations align closely with experiments.

    Source: §1 · p. 3, §7 · p. 14

  • Takeaways & Limitations

    The allocation model can help derive a nearly optimal long-context RAG configuration for a given computation budget.

    Source: §6 · p. 14

    The discussion notes that existing evaluation metrics fail in some cases, such as abbreviations.

    Source: §6 · p. 13

Abstract

from arXiv · show

The scaling of inference computation has unlocked the potential of long-context large language models (LLMs) across diverse settings. For knowledge-intensive tasks, the increased compute is often allocated to incorporate more external knowledge. However, without effectively utilizing such knowledge, solely expanding context does not always enhance performance. In this work, we investigate inference scaling for retrieval augmented generation (RAG), exploring the combination of multiple strategies beyond simply increasing the quantity of knowledge, including in-context learning and iterative prompting. These strategies provide additional flexibility to scale test-time computation (e.g., by increasing retrieved documents or generation steps), thereby enhancing LLMs' ability to effectively acquire and utilize contextual information. We address two key questions: (1) How does RAG performance benefit from the scaling of inference computation when optimally configured? (2) Can we predict the optimal test-time compute allocation for a given budget by modeling the relationship between RAG performance and inference parameters? Our observations reveal that increasing inference computation leads to nearly linear gains in RAG performance when optimally allocated, a relationship we describe as the inference scaling laws for RAG. Building on this, we further develop the computation allocation model to estimate RAG performance across different inference configurations. The model predicts optimal inference parameters under various computation constraints, which align closely with the experimental results. By applying these optimal configurations, we demonstrate that scaling inference compute on long-context LLMs achieves up to 58.9% gains on benchmark datasets compared to standard RAG.

1. Introduction

The introduction frames long-context RAG scaling as a problem of using external knowledge effectively, not simply expanding context. It introduces DRAG and IterDRAG and motivates modeling how inference compute should be allocated.

1. Introduction

  • Prior work reports benefits from longer contexts, many-shot examples, and more retrieved knowledge, but retrieval gains hold only up to a threshold.The introduction positions its inquiry around scaling RAG beyond simply increasing retrieved knowledge.Source: §1 · pp. 1–2
  • DRAG uses RAG examples as demonstrations to teach models to locate relevant information, while IterDRAG decomposes queries and interleaves retrieval with generation.The iterative strategy aims to bridge the compositionality gap in multi-hop queries.Source: §1 · p. 2
  • The paper measures inference computation as effective context length and scales it through retrieved documents, in-context examples, and, for IterDRAG, additional generation steps.Different parameter combinations allocate computation differently.Source: §1 · p. 2
  • The paper distinguishes its demonstration-based RAG examples, which contain documents, questions, and answers, from in-context RAG that prepends documents or QA examples.This distinction is stated in the introduction’s footnote.Source: §1 · p. 2
  • DRAG and IterDRAG outperform simply scaling document quantity, with IterDRAG improving on DRAG through interleaved retrieval and generation on Gemini 1.5 Flash benchmarks.The comparison is shown in Figure 2.Source: Figure 2 · p. 3, §1 · p. 3
  • With optimal inference configurations, long-context RAG performance improves almost linearly as effective context length grows, motivating inference scaling laws and a computation allocation model.The model estimates configurations that maximize performance under different computation constraints.Source: §1 · p. 3

2. Related Work

Related work covers long-context modeling, in-context learning, and retrieval-augmented generation, then identifies inference-compute allocation for long-context RAG as under-explored.

2.1. Long-Context LLMs

  • Prior long-context work includes sparse or low-rank attention, recurrent and state-space models, context extrapolation or interpolation, and efficient attention for million-token inputs.The section groups methods for reducing memory, replacing transformers, extending context windows, and enabling longer sequences.Source: §2.1 · p. 4

2.2. In-Context Learning

  • In-context learning research studies pretraining, example selection, formatting, ordering, and scaling the number of demonstrations with long-context LLMs.The section notes findings that many-shot ICL can mitigate pretraining biases and improve performance across tasks.Source: §2.2 · p. 4

2.3. Retrieval Augmented Generation

  • RAG research improves retrieval, document encoding, and selective use of knowledge, while inference scaling for long-context RAG remains under-explored.The paper therefore investigates how inference-compute variations affect RAG performance and test-time allocation.Source: §2.3 · p. 4

3. Inference Scaling Strategies for RAG

The section defines effective context length as its measure of inference computation, then presents single-request demonstration-based RAG and iterative query decomposition with interleaved retrieval.

3.1. Preliminaries

  • Effective context length is the total input tokens across iterations before the final answer; it excludes output tokens and retrieval costs.For single-call methods it is the prompt length, while iterative methods can extend it beyond the model’s context-window limit.Source: §3.1 · p. 5

3.2. Demonstration-Based RAG

  • DRAG combines retrieved documents and in-context examples in one expanded prompt, with retrieved documents selected for both demonstrations and the test query.It reverses document order so higher-ranked documents are closer to the query.Source: §3.2 · p. 5

3.3. Iterative Demonstration-Based RAG

  • IterDRAG decomposes complex queries into sub-queries, retrieves context for each, generates intermediate answers, and combines these with the query to synthesize a final answer.Its in-context examples are constructed with sub-queries and intermediate answers, and inference allows up to five decomposition iterations.Source: §3.3 · pp. 5–6
  • Figure 3 contrasts DRAG’s expanded-context answering with IterDRAG’s multiple inference steps for query decomposition and document retrieval.The iterative process is intended to improve final answers by breaking queries into sub-queries.Source: Figure 3 · p. 6

4. RAG Performance and Inference Computation Scale

DRAG and IterDRAG improve RAG performance as inference budgets grow, with the strongest strategy depending on budget length. Optimal performance follows an approximately linear scaling trend, while parameter-specific gains and best configurations vary.

4.2. Overall Performance

  • At 16k and 32k tokens, DRAG typically leads; at 128k and beyond, IterDRAG achieves superior overall results, while both scale better than QA and RAG baselines.Many-shot QA peaks at 16k tokens, and RAG plateaus after 128k.Source: Table 1 · p. 7, §4.2 · p. 8

4.1. Fixed Budget Optimal Performance

  • The study optimizes documents, in-context examples, and generation iterations within token budgets, evaluating Gemini 1.5 Flash on four multi-hop QA datasets.Effective context length counts total input tokens across iterations; tested budgets range from 16k to 5M tokens.Source: §4.1 · p. 7

4.3. Inference Scaling Laws for RAG

  • Optimal RAG performance scales nearly linearly with inference compute, but gains diminish beyond 1M tokens and may plateau from 1M to 5M.Figure 4 plots normalized performance against effective context length and shows the optimal configurations with a fitted line.Source: Figure 4 · p. 8, §4.3 · p. 9

4.4. Parameter-Specific Scaling

  • Increasing documents, demonstrations, or generation steps usually improves performance, but the best parameter combinations depend on method, metric, and dataset.More retrieved documents usually yield larger gains than more examples; adding a shot is more helpful for IterDRAG than DRAG, and gains can become marginal or decline beyond soft thresholds.Source: §4.4 · pp. 9–10, Figure 5 · p. 9

5. Inference Computation Allocation for Long-Context RAG

The section formulates a computation-allocation model over documents, demonstrations, and iterations, then tests its predictions across domains and context lengths. Predicted performance tracks observed results closely, including on unseen domains and at target lengths below 1M tokens.

5.1. Formulation and Estimation

  • The model uses document count k, demonstrations m, and maximum iterations n; task informativeness scores cover documents and shots, while iteration informativeness is set to zero.Document and shot informativeness are defined by performance gains from adding one document or example relative to zero-shot QA.Source: §5.1 · p. 10

5.2. Validating the Computation Allocation Model for RAG

  • Figure 6's DRAG predictions track actual metrics as documents and shots increase, with strongest consistency on Bamboogle and greater variability on HotpotQA.Each subplot represents a dataset, each line a document setting, and shots scale context length.Source: Figure 6 · p. 11, §5.2 · p. 11

5.1. Formulation and Estimation

  • Ordinary least squares estimates model-specific coefficients, while task-specific informativeness scores support applying the model to downstream tasks without additional calibration.The model is intended to select inference parameters under a maximum effective context length.Source: §5 · p. 10, §5.1 · p. 11

5.2. Validating the Computation Allocation Model for RAG

  • Adding b with task informativeness i improves relevance and reduces error across tasks, while inverse-sigmoid metric scaling improves estimation over quadratic-input or linear-scaling alternatives.Table 2 reports the ablation results; the complete model uses sigmoid scaling of P.Source: §5.2 · pp. 11–12, Table 2 · p. 11
  • The computation allocation model for RAG achieves 96.6% of optimal performance, while Bamboogle and HotpotQA performance metrics vary by less than 2.5% from oracle results.Source: §5.2 · p. 12
  • From 128k to 1M tokens, predicted results differ from oracle by 2.8% on average; at 5M, the average difference rises to 5.6%.Extrapolation from 32k to 128k is challenging because DRAG performs best around 32k while IterDRAG typically excels at 128k.Source: §5.2 · p. 12, Table 4 · p. 12

6. Discussion

The discussion weighs retrieval quality and long-context limits against the gains from allocating more inference compute. It also identifies persistent error sources and argues that retrieval and reasoning improvements do not resolve every failure mode.

6. Discussion

  • Retrieving more documents improves recall, but discounted ranking metrics show diminishing returns and RAG performance peaks between 100 and 500 documents.How retrieval refinement such as reranking would change the paper's inference-scaling behavior remains unknown.Source: §6 · p. 13
  • IterDRAG substantially improves inaccurate retrieval and reasoning errors, while hallucination, unreliable knowledge, and evaluation failures remain concerns.The discussion also flags outdated retrieval, missing reasoning, unfaithful reasoning, and metric failures such as abbreviation cases.Source: §6 · p. 13
  • Naïvely extending context in each generation step does not always improve RAG: DRAG peaks around 10^5 tokens, while IterDRAG is optimal around 10^6 tokens using multiple generation rounds.The authors suggest improving relevant-information identification in extensive contexts and long-context modeling for lengthy in-context demonstrations.Source: §6 · p. 13
  • With optimal allocation, long-context RAG performance can improve almost linearly as test-time compute increases, so the preferred trade-off depends on the available inference budget.The paper contrasts this with existing methods that often plateau beyond certain compute thresholds.Source: §6 · pp. 13–14

7. Conclusion

The conclusion reports that optimally configured long-context RAG improves almost linearly as test-time compute increases. It also presents a computation allocation model that predicts performance across configurations and closely matches experiments.

7. Conclusion

  • Long-context RAG performance improves almost linearly with increasing orders of magnitude of test-time compute under optimal inference parameters.The paper describes this relationship as inference scaling laws for RAG.Source: §7 · p. 14
  • The paper derives inference scaling laws and a computation allocation model to predict RAG performance across varying hyperparameters.The model is based on the observed relationship between performance and inference computation.Source: §7 · p. 14
  • Experiments show that optimal configurations can be estimated accurately and align closely with experimental results.Source: §7 · p. 14

A. Retrieval Quality

The appendix examines retrieval quality for DRAG and IterDRAG and compares IterDRAG with chain-of-thought prompting. It finds that iterative retrieval improves ranking-sensitive retrieval metrics, while IterDRAG consistently outperforms the evaluated CoT strategy.

A. Retrieval Quality

  • The retrieval evaluation varies documents per query from 1 to 2k and measures Recall, NDCG, and MRR; iterative retrieval uses up to five steps with 50 documents per step and 2 shots.Source: A. Retrieval Quality · p. 20
  • Recall rises toward near-perfect levels as document counts approach 1k, but NDCG and MRR plateau around 100 documents, so higher recall need not improve answer quality.Figure 7 presents DRAG retrieval performance across datasets.Source: A. Retrieval Quality · p. 20, Figure 7 · p. 20
  • IterDRAG retrieval gains average 30.5%, reach 57.1% on some complex multi-hop queries, and are larger on NDCG and MRR than on recall.Reported average gains are 30.7% in NDCG and 39.9% in MRR, compared with 21.7% in recall.Source: A. Retrieval Quality · p. 20, Table 5 · p. 20

B. Chain-of-Thought vs. IterDRAG.

  • IterDRAG consistently outperforms the evaluated CoT strategy, with the authors linking the difference to retrieval quality, model size, and CoT-example informativeness.The comparison uses three larger datasets and a 4-shot setting with 5 documents.Source: B. Chain-of-Thought vs. IterDRAG · p. 21, Table 6 · p. 21

C. Additional RAG Results

Additional results show IterDRAG improving on DRAG for multi-hop QA, while one-hop gains plateau at different document counts across datasets. Scaling patterns also persist with another retriever, though their shape varies by dataset.

C. Additional RAG Results

  • IterDRAG further improves RAG performance over DRAG, with best-performing document-and-shot combinations toward longer contexts and an optimal 32 shots.Figure 8 averages results across datasets; its axes are document count and shot count, and color encodes performance.Source: C. Additional RAG Results · p. 21, Figure 8 · p. 21
  • 69.0% accuracy is TriviaQA’s highest with 50 documents, while Natural Questions peaks at 54.6% with 20 documents in 1-shot and then slightly declines.Across these one-hop datasets, the optimal number of shots falls between 1 and 4; gains plateau beyond dataset-specific thresholds.Source: C. Additional RAG Results · p. 22, Figure 9 · p. 22
  • On binary multi-hop StrategyQA, DRAG improves baseline QA accuracy by 29.3%, and IterDRAG reaches 83.4 accuracy.The reported results are in Table 7.Source: C. Additional RAG Results · p. 22, Table 7 · p. 22

D. Additional Results on Inference Scaling Laws with GTR Retriever

  • With GTR XXL on 100 sampled MuSiQue examples, DRAG and IterDRAG gain consistently as context grows, especially IterDRAG at longer lengths.The authors report this as a consistent scaling pattern with an alternative retriever, unlike the plateau in standard RAG.Source: D. Additional Results on Inference… · p. 22, Figure 10 · p. 23

E. Additional Results on Inference Scaling Laws for RAG

  • Scaling curves vary by dataset: gains are more linear and consistent on Bamboogle and MuSiQue, nearly linear on 2WikiMultiHopQA through 1M tokens, and more sigmoidal on HotpotQA and longer 2WikiMultiHopQA contexts.The authors suggest dataset difficulty and retrieved-document quality as possible explanations for the sigmoidal patterns.Source: E. Additional Results on Inference… · p. 23

F. Additional Results on Computation Allocation Model for RAG

Additional analyses report statistically significant model parameters apart from b1 and compare predicted with actual IterDRAG performance across datasets. The model’s predictions are reported as consistent with observed results for both DRAG and IterDRAG.

F. Additional Results on Computation Allocation Model for RAG

  • The computation allocation model’s parameters are statistically significant except b1, whose p-value is slightly above 0.05; retaining b1 improved generalization in many cases.The report includes parameter estimates, p-values, R2, and MSE statistics in Table 8.Source: F. Additional Results on Computation… · p. 23, Table 8 · p. 23
  • The model’s predictions align with actual IterDRAG performance across datasets, with larger prediction variations than DRAG and more consistent trends on HotpotQA and 2WikiMultiHopQA.Figure 12 plots datasets by subplot and document settings by line while scaling compute through shots and generation iterations.Source: F. Additional Results on Computation… · p. 23, Figure 12 · p. 25
  • The authors report computation-allocation modeling as accurate for both DRAG and IterDRAG, and visualize prediction surfaces for both methods.Figure 13 shows normalized performance against the predicted surface.Source: F. Additional Results on Computation… · p. 24

G. Error Analysis

The error analysis groups mistakes into retrieval, reasoning, hallucination or unfaithful reasoning, and evaluation or refusal issues. It also notes scoring and abstention limitations in the experiments.

G. Error Analysis

  • On challenging datasets such as MuSiQue, performance remains moderate even with IterDRAG, motivating analysis of common failure cases.Examples include missing answer-bearing retrievals, incomplete reasoning despite correct documents, and answers unsupported by retrieved documents.Source: G. Error Analysis · p. 25
  • Errors fall into four categories: inaccurate or outdated retrieval, incorrect or missing reasoning, hallucination or unfaithful reasoning, and evaluation issues or refusal to answer.Reasoning errors tend to occur more frequently on difficult questions and with one-step DRAG; retrieval may fail even among the top 50 documents.Source: G. Error Analysis · p. 25
  • Scoring can be inaccurate for abbreviations or date-format variations, and the experiments do not account for abstaining from answers.The authors note that omitted abstention could lead to unfair scores.Source: G. Error Analysis · p. 26

H. Implementation

The implementation uses retrieved Wikipedia passages and Gemini 1.5 Flash, with DRAG scaling through documents and demonstrations and IterDRAG adding iterative retrieval and generation. The allocation model is fit and evaluated using normalized metrics and task-specific informativeness estimates.

H. Implementation

  • Documents and queries are indexed with Gecko-1B (en) over KILT Wikipedia passages; top-k documents are selected and truncated to 1024 tokens.Prompts place retrieved documents, query, and label in that order, with higher-ranked documents closer to the query.Source: H. Implementation · p. 26, Figure 15 · p. 29
  • DRAG scales Gemini 1.5 Flash prompts through document count and demonstrations, while IterDRAG adds Self-Ask iterative retrieval and generation for up to five iterations.Each iteration generates a sub-query and intermediate answer; retrieved documents are appended before the next answer step.Source: H. Implementation · p. 26, Figure 16 · p. 29, Figure 17 · p. 30
  • The model fitting normalizes performance per dataset and metric, estimates six parameters by ordinary least squares, and uses task-specific document and shot informativeness.For R2 and MSE, peak and valley outliers are excluded; domain generalization and length extrapolation use all data points.Source: H. Implementation · p. 27
Loading 2410.04343v2…