Inference Scaling for Long-Context Retrieval Augmented Generation
Authors: Zhenrui Yue
Honglei Zhuang
Aijun Bai
Kai Hui
Rolf Jagerman
Hansi Zeng
Zhen Qin
Dong Wang
Xuanhui Wang
Michael Bendersky
Venue: ICLR 2025 Oral (2025)
Keywords: inference scaling
long-context LLM
retrieval augmented generation
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: 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.

Published as a conference paper at ICLR 2025
INFERENCE SCALING FOR LONG-CONTER RETRIEVAL AUGMENTED GENERATION
Zhenrui Yue*!, Honglei Zhuang Aijun Bai*, Kai Hui Rolf Jagerman', Hansi Zengf Zhen Qin', Dong Wang', Xuanhui Wang', Michael Bendersky²
'University of Minoiis Urbana- Obbana-Chaamaignn "Google DeepMind, UMass Amherst zhenrui3e 38 11 llincis mrcissedde edu, hizêgoogle. com
ABSTRACT
The scaling of inference computation has unlocked the potential of long-context large language models (LLMs) across diverse settings. For knowledge-intensiin tasks, the increased compute is often allocated to incorporate more external knowl- edge. However, without effectively utilizing such knowledge, solely expanding context does not always enhance performance. In this work, we investigate infer- ence scaling for retrieval augmented generation (RAG), exploring the combination of multiple strategies beyond simply increasing the quantity of knowledge, includ- ing in-context learning and iterative prompting. These strategies provide additional flexibility to scale test-time computation (e-g.. by increasing retrieved documents 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 given budget by modeling the relationship between RAG performance and inference parameters? Our observations reveal that increasing inference computation leads nearly linear gains in RAG performance when optimally allocated, relation- ship 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 benchmark datasets compared to standard RAG.
INTRODUCTION
Long-context large language models (LLMs) are designed to handle extended input sequences, enabling them to process and understand longer context (e.g.. Gemini 1.5 Pro with up to 2M tokens) (Achiam et al. 2023 Team et al. [003 Reid et al. 2024). Combined with increased inference computation, long- context LLMs demonstrate improved performance across various downstream tasks (Agarwal et al. 2024 Snell et al. 2024). For example, many-shot in-context learning (ICL. can match the performance of supervised fine-tuning by providing extensive in-context examples (Bertsch) et al. 2024). Particularly for know ttowledge-intensive tasks that leverage retrieval augmented generation (RAG), increasing the quantity size of retrieved documents up a certain threshold consistently
enhances the performance (Ram et al., 2023; Xu et al. 2024 Jiang et al. 2024
Previous studies on inference scaling for RAG focus on expanding the retrieved knowledge by increasing the number or lengths of retrieved documents (Xu et al. 2024] Jiang et al. (2024) (Shao
al. 2024). However, only emphasizing on the knowledge quantity without providing further guidance presents certain limitations. On one hand, current long-context LLMs still have limited ability to effectively locate relevant information in ultra-long sequences upon challenging tasks (Li
et al. 2024 Kuratov et al. 2024). For instance, the optimal performance of long-context LLMs
"Equal contribution
Work done while Google DeepMind

---

Published as a conference paper at ICLR 2025
RAG
DRAG
Optimal Config
IterDRAG
laod Performncc
Optimal Conf
102
10P
10F 307
Effective Context Langth
Effactive Contect Langth
Figure I: Normalized performance vs. effective context lengths on MuSiQue. Each line represents a fixed configuration, scaled by adjusting the number of documents. Red dots and dash lines represent the optimal configurations and their fitting results. Standard RAG plateaus early at 10 tokens, in contrast, DRAG and IterDRAG show hear-lineal improvement the effective context length grows.
often achieved without fully utilizing the maximum length (Agarwal al. 2024). On the other hand, numerous studies show that retrieving over soft thresholds (e.g., top- 10 documents) leads to performance plateau and may even cause declines (Ram al. 2023 Lee et al. 2024a]) 2024). Such performance drops may be traced back to the increased noise within context, which causes distraction and adversely affects generation (Yoran et al. 2024 Zhang et al. 2024 Leng et al. 2024). As result, inference scaling of long-context RÁG remains challenging for existing methods.
K024a]Kuratovv Kuratov et al.
In this work, we leverage a broader range of strategies to comprehensively explore how RAG benefits from the scaling of inference computation. A straightforward strategy is iitionstraaion-baseed RAG (DRAG), where multiple RAG examples are provided as demonstrations to utilize the long-context capabilities of LLMS (Brown et al. 2020). DRAG allows models to learn (in-context) how to locate relevant information and apply it to response generation"] Nevertheless, the quality of one-step retrieval varies across tasks and often fails provide sufficient information. Inspired by iterative methods (Trivedi et al., 2023 Yoran et al. 2024), we develop iterative demonstration- based RAG (IerDRAG). IRerDRAG learns to decompose input queries into simpler sub-queries and answer them using interleaved retrieval. By iteratively retrieving and generating upon sub sub-queries, LLMs construct reasoning chains that bridge the compositionality gap for multi-hop queries. Together, these strategies provide additional flexibility in scaling inference computation for RAG, allowing long-context LLMs to more effectively address complex mowlledg-intensive queries.
Building on these strategies, we investigate multiple ways to scale up inference computation. Here, we measure computation by considering the total number of input tokens across all iterations, referred to the effective context length. In DRAG, scaling the effective context length can be done by increasing two inference parameters: the number of retrieved documents and in-context examples. In IteD test-time compute can be further extended by introducing additional generation steps. Since different combinations of inference parameters result in varied allocations of computational resources, our goal is to establish the relationship between RAG performance, different scales and allocations of inference computation. Through extensive experiments on benchmark QA datasets, we demonstrate an almost linear relationship between RAG performance and the scale of effective context length by combining both RAG strategies, as shown in Figure [ (right). Moreover, our RAG strategies exhibit improved performance than merely scaling the number of documents, achieving state-of-the--at performance with the compact Gemini 1.5 Flash (See evaluation in Figure 2)
Drawing from our observations, we examine the relationship between RAG performance and in- ference computation, which we quantify as the inference scaling laws for RAG. These observed inference scaling laws reveal that RAG performance consistently improves with the expansion of the effective context length under optimal configurations. Consequently, take a deeper dive into modeling RAG performance with respect to various inference computation allocations. Our goal is to predict the optimal set of inference parameters that maximize the performance across different RAG tasks. To achieve this, we quantitatively model the relationship between RAG performance and varying inference configurations with the computation allocation model for RAG. Using the
Different from in-context RAG that prepends documents/ QA examples (Press et al. 2023 Ram et al.
2023
we leverage multiple examples comprising of documents, questions and answers demonstrate the task.
2

---

Published as a conference paper at ICLR 2025
Zero shot QA
Nany shot
RAG
DR4G
kerDRRS
F0
76.9
68.8
57.6
$6.4
50,952.2
48.4
45.4
34.3
25.220.2
16,8
20.4
23.5
6,6
Baniooogle
HatparQ4
MusiQua
2NTeMMuhHHOOCA
Average
Figure Evaluation accuracy Gemini 1.5 Flash using different methods: zero-shot QA, many-shot QA. RAG (with optimal number of documents), DRAG and IterDRAG on benchmark QA datasets. By scaling up inference compute (up to 5M tokens), DRAG consistently outperforms baselines, while IterDRAG improves upon DRAG through interleaving retrieval and iterative generation.
estimated computation allocation model, the optimal configurations can be empirically determined and generalize well for various scenarios, thereby maximizing the utilization of the computation budget. We summarize our contributions as follows:
We systematically investigate inference scaling for long-context RAG, for which we intro- duce two scaling strategies, DRAG and IterDRAG, to effectively scale inference compute.
We comprehensively evaluate DRAG and IeerDRAG, where they not only achieve state- of-the-art performance, but also exhibit superior scaling properties compared to solely increasing the quantity documents.
Through extensive experiments on benchmark QA datasets, we demonstrate that when test-time compute is optimally allocated, long-context RAG performance can scale almost linearly with the increasing order of magnitude of the computation budget.
We quantitatively model the relationship between RAG performance and different inference parameters, deriving the computation allocation model. This model aligns closely with our experimental results and generalize well across scenarios, providing practical guidance for optimal computation allocation in long-context RAG.
2 RELATED WORK
2.1 LONG-CONTEXT LLMS
Long-context large language models (LLMs) are designed utilize extensive context and thereby improve their generative capabilities. Early works in extending context lengths involve sparse / low- rank kernels to reduce memory requirements (Kitaev et al. 2019 Beltagy et al.) 2020 Zaheer et al. 2020: Choromanski et al. 2020). In addition, recurrent and state space models (SSMs) are proposed as efficient substitutes for transformer- based models (Gu et al.) 2021 Gu & Dao 2023 Peng et al. 2023a Beck et al. 2024). For causal LLMs, extrapolation and interpolation methods have proven effective in expanding context window lengths (Press et al. 2021 Chen et al. 2023 Sun et al. 2023
Peng et al.
2023b). Recent advancements in efficient attention methods (Dao al. 2022 Jacobs
et al.) 2023 Liu et al. 2023) further enable LLMs to train and infer upon input sequences comprising
millions of tokens (Achiam al. 2023 Team al. 2023 Reid al. 2024).
2.2 IN-CONTEXT LEARNING
In-context learning (ICL) offers computationally efficient approach to enhance model performance at inference time by conditioning on a few demonstrations of the task (Brown et al. 2020). To further improve ICL performance, existing works focuses pretraining strategies that optimize the language models to learn in-context (Min et al.) 2022 Wei et al., 2023 Gu et al. (023). In addition, selective usage of few-shot examples shown to be helpful for enhancing downstream task performance (Liu ctt al.) 2022; Rubin et al. 2022 Wang et al. 2024). Notably, reformatting or finding optimal ordering of in-context examples also improves ICI. performance effectiveness (Lu al. 2022 Wu et al. 2023

---

Published as a conference paper at ICLR 2025
Liu et al.
2024a). With the emergence of long-conters LLMs (Achiam al. 2023 Team et al. 2023
Reid et al. 2024), scaling the number of examples becomes possible in ICL (Li et al. 2023 Bertsch
al. 2024: Agarwal et al. 2024). For instance, Agarwal et al. (2024) show that many- shot ICL can
mitigate pretraining biases within LLMs and thus improves ICL performance across various tasks.
2.3 RETRIEVAL AUGMENTED GENERATION
Retrieval augmented generation (RAG) improves language model performance by incorporating
relevant knowledge from external sources (Lewis al. 2020 Guu et al. 2020 Karpukhin al. 2020).
In contrast to naive RAG, optimizing the retrieval stage can effectively enhance context relevance
and improve generation performance (Ma al.) 2023 Trivedi et al. 2023 Jiang al. 2023 Shi
et al.] [2024] Sarthi et al. [024,] Lin et 20244 Mo et al. 2024). An example is REPLUG, in which
Shi et al. (2024) leverage LLM as supervision learn dense retriever model. In addition, encoding
documents can increase knowledge retrieval and improve generation capabilities (Khandelwal et al.
2019] Izacard & Grave 2021; Borgeaud et al. 2022 Izacard et al. 2023). For instance, |zacard
Grave (2021) leverages fusion-in-decot architecture to encode multiple question-passage pairs while maintaining the model efficiency. Alternatively, selectively utilizing knowledge from the documents
improves the robustness of LLMs against irrelevant context ( (Yu et al. 2023 Yoran al. 2024 Yan
itt al. 2024; Yue et al.) 2024) Zhang et al.] 2024). For example, RAFT proposes to train language models with negative documents improve generation quality and relevance (Zhang et al. 2024). Concurrent to our work, long-document retrieval and datastore scaling are proposed to optimize RAG performance (Tiang et al. 2024 Shao et al.) 2024). Despite such progress, inference scaling remains under-explored for long-context RAG methods. As such, we investigate how variations in inference computation impact RAG performance, with the goal of optimizing test-time compute allocation.
3 INFERENCE SCALING STRATEGIES FOR RAG
3.1 PRELIMINARIES
We measure inference computation with effective context length, defined the total number of input tokens across all iterations before the LLM outputs the final answer. For most methods that only call the LLM once, the effective context length equivalent to the number of input tokens in the prompt and is limited by the context window limit of the LLM. For methods that iteratively call the LLM, the effective context length can be extended indefinitely depending on the strategy. We exclude output tokens and retrieval costs from our analysis, as LLMs typically generate significantly fewer tokens (fewer than 10) in nhowledge-intensifne computationally expensive than LLM inference, especially with scalable matching methods (Sun) ct al. 2024). Our objective is to understand how RAG performance changes as we scale up inference computation. In demonstration- RAG (DRAG), we achieve such scaling by incorporating both extensive documents and in-context examples. For further scaling, we increase generation steps through iterative demonstration--aael RAG (ItrDRAG). We introduce both strategies below.
tasks. Additionally, retrieval generally much less
3.2 DAMONSTRATION-BASSIE RAG
Demonstratiinn-asse RAG (DRAG) leverages in-context learning to exploit the capabilities of long- context LLMs by directly generating answers from an extended input context. DRAG builds upon naïve RAG and integrates both documents and in-context examples into the input prompt. This expanded context allows the model to generate answers to the input query within a single inference request (See FigureBleftt, For both in-context examples and the test-time query, we employ a retrieval model to select the top-k retrieved documents from a large corpus (e.g.. Wikipedia). We reverse the order of the retrieved documents, placing higher-ranked documents closer to the query (Liu et al. 2024b). As use instruction-tunet LLMs, we design similar prompt template following Agarwal et al. (2024) and align the formatting with prefixes for retrieved documents, input and output (See Appendix lb. Unlike previous works (Press et al.) 2023 Trivedi et al.) (003), DRAG incorporates extensive retrieved documents within the demonstrations, enabling
LLMs to leam to extract relevant information and answer questions using rich input context.
4

---

Published as a conference paper at ICLR 2025
Sub-Query'1
Sub-Query
Sub-Query
In-Conteat
Documents
Generate
Genenate
Generate
intermediate
Ittrmeclate
meciste
eritt
IteDRAGG
Figure 3: DRAG vs. IteDDRAG. IterDRAG breaks down the input query into sub-queries and answer them to improve the accuracy of the final answer. In test-time, IterDRAG scales the computation through multiple inference steps to decompose complex queries and retrieve documents.
3.3 ITERATIVE DEMONSTRATION- BASED RAG
Despite access to external knowledge, complex multi-hop queries remain challenging due to the compositionality gap. To tackle this issue, we introduce iterative demonstration- based RAG (Iter- DRAG), which handles complex queries by decomposing the query into simpler sub-queries. For each sub-query, retrieval is performed to gather additional contextual information, which then used generate intermediate answers. After all sub-queries resolved, the retrieved context, sub-queries, and their answers are combined to synthesize the final answer (See Figure right).
While multiple existing datasets provide training data with queries and corresponding answers, sub-queries and intermediate answers are often absent. To generate in-context examples with sub- queries and intermediate answers, we prompt LLMs with constrained decoding follow the Self-Ask format (Press et al.) 2023 Koo et al. 2024). In each iteration, LLMs generate either sub-query, an intermediate answer, or the final answer. If sub-query is generated, additional documents are retrieved and interleaved into the prompt before producing the intermediate answer. IterDRAG continues until the final answer is generated or the number of maximum iterations is reached, at which point LLM is forced to generate the final answer. We retain examples with intermediate steps and correct final answers to construct in-context demonstrations. Each example should include the retrieved documents, sub-query and answer pairs, well as the final answer.
During inference, in-context examples are prepended to the initial documents retrieved for the input query. Similarly, each inference request yields sub-query, intermediate answer, or the final answer. Upon sub-queries, additional documents are retrieved and merged with the initial ones to generate intermediate answers. In implementation, decomposition before generating the final answer. This iterative process effectively scales test-time computation, with the input tokens from all iterations summed calculate the effective context length. IteDDAAG facilitates more granular approach by learning to: (1) decompose query into simple and manageable sub-queries; and (2) retrieve and locate relevant information answer quebrees As such, the iterative retrieval and generation strategy helps narrowing the compositionality gap and improves knowledge extraction, thereby enhancing overall RAG performance.
allow up five iterations of query
4 RAG PERFORMANCE AND INFERENCE COMPUTATION SCALE
4.1 FIXED BUDGET OPTIMAL PERFORMANCE
For given budget on inference computation, i.e., a maximum effective context length Lmax, there are multiple ways to optimize the use of computation resources through inference parameters. For example, in DRAG, while the IterDRAG strategy, we additionally introduce the number of iterations for retrieval and generation. Henceforth, use to denote all these inference parameters.
adjust both the number of retrieved documents and in-context examples,
5

---

Published as a conference paper at ICLR 2025
Table Optimal performance of different methods with varying maximum effective context lengths Lmux (i.e., the total number of input tokens across all iterations). ZS QA and MS QA refers to zero-shot QA and many-shot QA respectively. Partial results are omitted for methods that do not further scale with increasing Lmax For clarity, we mark the best results for each Lmax in bold.
Bamboogle
HotpotQA
MuSiQue
ZWikiMultiHopQA
Lenxx
Method
EM
F1
Acc
EM
F1
Acc
EM
Fl
Acc
EM
Fl
Acc
ZS QA
16.8
25.9
19.2
22.7
32.0
25.2
5.0
13.2
6.6
28.3
33,5
30.7
MS QA
24.0
30.7
24.8
24.6
34.0
26.2
7.4
16.4
8.5
33.2
37.5
34.3
16k
RAG
44.0
54.5
45.6
44.2
57,9
49.2
12.3
21.5
15.3
42.3
49.3
46.5
DRAG
44,0
55.2
45.6
45.5
58.5
50.2
14.5
24,6
16.9
45.2
53,5
50.5
RerDRAG
46.4
56.2
51.2
36.0
47.4
44.4
8.1
17.5
12.2
33.2
38.8
43.8
RAG
48.8
56.2
49.6
44.2
58.2
49.3
12.3
21.5
15.3
42.9
50.6
48.0
32k
DRAG
48.8
59.2
50.4
46.9
60.3
52.0
15.4
26.0
17.3
45.9
53.7
51.4
IerDRAG
46.4
56.2
52.0
38.3
49.8
44.4
12.5
23.1
19.7
44.3
54,6
56.8
RAG
51.2
60.3
52.8
45.7
59.6
50.9
14.0
23.7
16.8
43,1
50.7
48.4
128k
DRAG
52.8
62.3
54.4
47.4
61.3
52.2
15.4
26.0
17.9
47.5
55.3
53.1
IeeDRAGG
63.2
74.8
68.8
44.8
59.4
52.8
17.3
28.0
24.5
62.3
73.8
74.6
DRAG
56.0
62.9
57.6
47.4
61.3
52.2
15.9
26.0
18.2
48.2
55.7
53.3
IM
ItrDRAGG
65.6
75.6
68.8
48.7
63.3
55.3
22,2
34,3
30.5
65.7
75.2
76.4
5M IteDRAGG
65.6
75.6
68.8
51.7
64.4
56.4
22.5
35.0
30.5
67.0
75.2
76.9
For each input query and its ground-truth answer (I,, y.) E X, we can apply the RAG inference strategy f parameterized by 0. We denote the effective input context length to the LLM as (( 0) and the obtained prediction as O = f (a.; 8). metric P(y. ūs) can then be calculated based on y: and ÿJ- To understand the relationship between RAG performance and inference computation, we sample few different inference computation budgets. For each budget Lmmax we find the optimal average metric P* (mmxx) achievable within this budget by enumerating different 0 E O:
(* = max
{ {
Σ P(u., ((r,:))) Vi, (2,;0) S Lmax {
BEG
Our goal is establish the relationship between the inference computation budget Lmux and the best possible performance within this budget P* ("mmax),, using any possible strategies and parameter configurations allocate the inference computation resources. For simplicity, we also refer to ('mmaa as the optimal performance. We investigate the following factors within the inference parameter set 0: (1) the number of documents k. which are retrieved from a large corpus (e.g.. Wikipedia) based on the input query; (2) the number of in-context examples m, where each of the examples consists of documents, an input query and its label; and (3) the number of generation iterations In DRAG, an answer can be directly generated upon input context, so n = 1. In contrast, IterDRAG involves multiple steps of interleaved retrieval and generation, expanding both the effective context length and inference compute without needing longer context windows.
We evaluate the performance of Gemini 1.5 Flash with context length window up to 1M tokens on knowledge-intensive question answering, including multi-hop datasets Bamboogle. HotpotQA. MuSiQue and 2WikiMultiHopQA (Press et al., 2023] Yang et al., 2018 Trivedi et al. 2022 Ho et al. 2020). Additional results are provided in Appendix C and Appendix To manage the computational costs of extensive experiments, we follow Wu et al. (2024); Gutiérrez et al. (2024) and sample I. 1.2k examples from each dataset for evaluation. The evaluation metrics include exact match (EM), F1 score (F1) and accuracy (Acc), in which the accuracy metric assesses whether the ground truth is located within the prediction. We sample the inference computation budget Lmax as 16k, 32k, 128k, IM and 5M tokens. For the parameter space é of DRAG, we consider the number of documents E (0, 1,2,5 10, 20, 50, 100. 200, 500, 1000), and the number in-context examples m ranging from 0, 20 21_ ..., to 2 For IterDRAG, we further experiment with number of iterations n up to 5. We compare to the following baselines: (1) zero-shot QA (QA), where the model does not leverage any retrieved documents demonstrations: (2) many-shot QA (MS QA), where the model only uses varying number of demonstrations m without any retrieved document; and (3) retrieval augmented generation (RAG), where the model only uses k retrieved documents without demonstrations. We report the optimal performance of each method with different maximum effective context length budgets by examining their performance with different inference parameter configurations.

---

Published as a conference paper at ICLR 2025
Optimal Conti
Narmalzed
107
101
107
Effective Connect Langth
Figure 4: Normalized performance vs. effective context lengths across datasets. Each line represents fixed configuration, scaled by varying the number of documents. Red dots indicate the optimal configurations, with the dashed line showing the fitting results. The observed optimal performance can be approximated by linear relationship with the effective context lengths.
4.2 OVERALL PERFORMANCE
We report the optimal performance (Lmxa) for different inference strategies in Table [ where we identify the optimal inference parameters for each computation budget T Some variants are omitted for certain Tmmax because they do not scale the corresponding context length. For example, the prompt for zero-shot QA cannot be increased, while the number of in-context examples for many-shot QA is capped at 2a, so neither scales to Lmxax = 32k. Similarly, RAG does not scale to Lmux larger than I 128k, and DRAG is limited by the LLM's context window limit of 1M.
Unlike QA and RAG baselines, the performance of DRAG and IterDRAG consistently increase expand the maximum effective context length. More specifically, observe: (1) DRAG and IterDRAG scale better than baselines. Basclines like many-shot QA peak at 16k tokens, while RAG improves until 1 128k, after which performance plateaus. In comparison, DRAG and IterDRAG can find optimal configurations to more effectively utilize test-time compute, exhibiting superior performance and scaling properties. Performance of DRAG consistently improves until IM tokens, while IterDRAG further enhances RAG performance with 5M tokens of computation budget by iteratively calling LLMs. (2) DRAG excels with shorter maximum lengths, while IterDRAG scales more effectively with longer effective context length. At 16k and 32k, DRAG typically delivers the best performance, while 128k and beyond, IterDRAG achieves superior results overall, highlighting the effectiveness of iterative retrieval and generation. These results suggest that increasing 7-max is beneficial for RAG, with DRAG and Ate strategies each excelling at different scales.
4.3 INFERENCE SCALING LAWS FOR RAG
To analyze RAG performance with different effective context lengths, we plot the performance of all configurations across datasets in Figure 4 Similar to Figure [ we visualize DRAG and IterDRAG and highlight the optimal performance pt (Tmmax for different selections T.max- The fitting results are shown grey dashed lines. We provide additional dataset-specific results in Appendix
The optimal performance exhibits consistent gains as the effective context length expands, demon- strating strong linear correlation, which we term the inference scaling laws lowsfor RAG. Combined with dataset-specific results, our key observations are: (1) The optimal performance scales nearly linearly with the order of magnitude of the inference compute. Such linear relationship suggests that RAG performance can be improved by increasing computation, allowing for more accurate predictions of performance given available compute resources. (2) For Timax above 10S, IteDDRAG continues to scale effectively with interleaving retrieval and iterative generation. This aligns with our results in Table where IterDRAG better utilizes computation budgets for effective context lengths exceeding 128k. (3) Gains optimal performance gradually diminish beyond effective context length of IM. Despite dataset variations, the performance follows similar trends up to 1M tokens. Beyond that, improvements from 1M to 5M are less substantial plateau, potentially due to limitations in long-contex: modeling. In summary, while gains are smaller beyond IM tokens, optimal RAG performance scales almost linearly with increasing inference compute through DRAG and IterDAGG.
7

---

Published as a conference paper at ICLR 2025
374
33
433
20000:
400
EM Peformance
F1 Parfomance
Acc Performance
(a) Averaged DRAG performance heatmap for different metrics.
Nerroted Perforano
(b) Performance vs. number of documents.
(c) Performance vs, number of shots.
Figure 5: RAG performance changes with varying number of documents and in-context examples. 5a reports the averaged metric values across datasets, whereas in 5b and [cc each line represents the normalized performance of consistent configuration with progressively increasing documents / shots.
4.4 PARAMETER CPECCIT SCALING
To gain further insights into the dynamics of DRAG and IterDRAG, grid search over different combinations and evaluate the performance. The results are presented in Figure 5 where we visualize DRAG performance using heatmaps (See IterDRAG heatmap in Appendix D). Additionally, we provide further results with varying numbers of documents (k) and shots (m). In summary, scaling retrieval, demonstrations and more generation steps leads to performance gains in most cases, yet such gains vary by effective context length and method. In particular, we note: (1) Documents and in- context examples are not equally helpful. For fixed configuration, increasing the number of retrieved documents k usually leads to more substantial performance gains, as evidenced by the differing slopes Figure 5 (2) Increasing shots m is more helpful for IterDRAG. For example, increase m from 0 to 1 (rather than k) is more helpful for IterDRAG, possibly due to demonstrations that leads to improved in-context query decomposition and knowledge extraction. (3) Scaling saturates differently for DRAG and IterDRAG. An example can be found in the increase of m from 0 to 1, which results in significant improvements for IterDRAG but shows little impact on DRAG. Beyond the soft thresholds, further increases in i m yield marginal gains or even results in performance declines. (4) For given Lmwr the optimal 9 depends on the method, metric and dataset. As illustrated in Figure [a and Figure the optimal combinations are sensitive the metries and located differently, posing challenges for performance modeling w.r.t. 0. conclusion, increasing documents, demonstrations and iterations can enhance RAG performance, but each contributes differently to the overall results. As such, identifying the optimal combination of hyperparameters remains challenging.
5 INFERENCE COMPUTATION ALLOCATION FOR LONG-CONTEXT RAG
After examining the overall performance of different RAG strategies and the varying impacts of different inference parameters, we now quantify the relationship between performance and the hyperparameter set 6. We hypothesize that for long-context RAG, we can model such test-time scaling properties and term it computation allocation model for RAG. This model, in turn, can be used to guide the selection of 0 based on the maximum effective context length T
5.1 FORMULATION AND ESTIMATION
With a slight abuse of notation, we redefine the average performance metric P (c.g., accuracy) on dataset X' function of 0. We consider the number of documents k demonstrations and
8

---

Published as a conference paper at ICLR 2025
Ban boogle
HotzorOA
MuSQqu
Nombbr Shots
Number Shars
Number
Number Shets
Figure 6: The estimated performance using the proposed observational scaling laws vs. actual metric values in DRAG. The subplots represent different datasets, where each line corresponds to a fixed number of documents, we scale the context length by increasing the number of shots.
maximum iterations n within 0, namely 0 == (ki, m, n)=. To account for the variance across methods and tasks, introduce i = (idne; ishot 0)7 1doe and Ssh measure the informativeness of documents and in-context examples respectively. While technically we can also define an liter to measure the informativeness of additional generation steps, applying lieer does not yield improved accuracy, leave it as 0 in our experiments. We formulate the computation allocation model a
o (P(O)) ~ (a b i)T log(O)= +
(2)
where @ refers to glement-wiss product. a, E R and e E I are parameters to estimated, and i can be computed base the specific task. There different ways to define we choose definition to compute i based on the performance difference between selected base configurations. In particular, for each strategy each dataset, dde is defined as the performance gain by only adding one document compared to zero-shot QA. Similarly, i.b is defined as the performance gain by adding only one in-context example compared zero-shot QA. To account for the sub-linearity in extremely long contexts (above IM), we apply an inverse sigmoidal mapping a to scale the values of the metric P. Further implementation details are reported in Appendix
In Equation ( estimations on o, and e are specific to a certain model, reflecting how LLMs improve with varying number of documents and shots (i.e., in-context learning zero-shot capabil- ities). In contrast, i models the performance variations within the selected task (i.c., how external knowledge/ / demonstrations help responding to the query). Therefore, the computation allocation model can be estimated once and applied to various downstream tasks without requiring additional calibration. To estimate the parameters, varying combinations of 0 are evaluated to perform ordinary least squares on b and We report the parameters for Gemini 1.5 Flash in Appendix G
5.2 VALIDATING THE COMPUTATION ALLOCATION MODEL FOR RAG
We evaluate the computation allocation model for RAG by comparing the predicted metrics to the actual values, with normalized results for DRAG visualized in Figure Here, each subplot represents a different dataset, and each line corresponds to document setting (k), we scale the context length by adjusting in-context examples (m). As illustrated, the performance improves with the increase of k and m across datasets, displaying highly consistent trends between the predicted and actual metric values, despite some variations. Notably, each dataset exhibits different levels of consistency: Bamboogle exhibits the highest consistency, while HotpotQA generates more variable results. Our findings demonstrate how external knowledge and in-context learning can effectively enhance RAG performance with long-context capabilities, suggesting the effectiveness of the computation allocation model for RAG and how they may be used predict benchmark results.
Table Ablation study results of the computation allocation model for RAG.
Exclude b
Quadratic
Linear o
Sigmoidal o
R²
MSE
R²
MSE
Rd
MSE
R
MSE
Values 0.866
0.116
0.867
0.117
0.876
0.109
0,903
0.085
"I implementation, we shift the values within 0 by small prevent numerical issues with log(O).

---

Published as a conference paper at ICLR 2025
Ablation Study. To verify the effectiveness of the computation allocation model, we perform
ablation studies and evaluate the fitting performance of different variants. In particular, we assess: (1) estimation without b and i (i.e.. Exclude b); (2) a quadratic form of input log(0) (Quadratie f); (3) linear scaling of P (Linear a); and (4) sigmoid scaling of P (Sigmoidal σ). The R² and MSE values for these variants are reported in Table 2 in which (4) represents the complete design of our computation allocation model. The results indicate that incorporating the additional b with i enhances the relevance and reduces error across all tasks. Moreover, applying inverse sigmoid to P significantly improves the estimation in comparison to quadratic 0 or linear scaling.
Table 3: Domain generalization results of the computation allocation model for RAG.
Barbboogle
HotpotQA
MuSiQue
2WikiMuliHHooOAA
EM
F1
Acc
EM
F1
Acc
EM
F1
Acc
EM
Acc
Baseline
49.6
58.8
51.2
46.3
16.9
46.5
53.7
51.6
Predict
64.0
75.6
68.0
47,8
63,3
55.3
19.3
32.5
29.3
60.8
72.4
74.9
Omclle 65.6
75.6
68.8
48.7
63.3
55.3
22.2
34.3
30.5
65.7
75.2
76.4
Domain Generalization. Wc also examine the generalization of the computation allocation model for RAG for unseen domains. In other words, the parameters of Equation (2) are tested on the target domain but learnt from the remaining domains. For inference, only i is derived from the target domain. We report the results for IM effective context length in Table 3 where we compare to an 8-shot baseline configuration (scaled by increasing retrieved documents) and the optimum results (Oracle). In summary, the results show that computation allocation model significantly outperforms baseline and closely aligns with the oracle results (96.6% of the optimal performance). Notably, Bamboogle and HotpotQA exhibit highly similar target results, with the performance metrics varying by less than 2.5% from the oracle. These results suggest the potential of applying the computation allocation model for RAG to a wider range of knowledge-intensive tasks.
Table 4: Length extrapolation results of the computation allocation model for RAG.
16k -
32k
32k
128k
128k -
1M
1M
5M
EM
F1
Ace
EM
F1
Ace
EM
F1
Acc
EM
F1
Acc
Baseline
37.4
47.6
40.4
39.0
49.5
42.2
39.3
49.3
42.8
44.5
55.4
49.8
Predict
37.4
48.2
41.0
41.2
52.0
45.4
48.0
60.9
56.9
47.9
59,8
55.2
Oracle
39.2
49.8
42.7
46.9
59.0
55.1
50.5
62.1
57.7
51.7
62.6
58.1
Length Extrapolation. In addition to predictability on unseen domains, we explore the extrapola- tion of context length based on the computation allocation model. Here, we estimate the parameters 2 using experiments with shorter context lengths and assess their predictive accuracy
of Equation on longer ones. We assess different extrapolation settings and present the predicted metric values in Table Our observations are: (1) The predictions are accurate and consistently outperform the 8-shot baseline. For instance, the average difference between the predicted and oracle results from 128k to IM tokens is just 2.8%. (2) Extrapolating from 32k to 128k is challenging. This is because DRAG performs best around 32k, while IterDRAG typically excels at a long context of 128k, as evidenced in Figure 4 Consequently, it creates discrepancy between training and predicting performance distribution. (3) 5M context length is less predictable, with the average performance difference between predicted and oracle metrics observed a substantial 5.6%. Overall, length extrapolation with computation allocation model is accurate and more effective for target lengths below 1M.
6 CONCLUSION
In this paper, we explore inference scaling in long-context RAG. By systematically studying the performance with different inference configurations, we demonstrate that RAG performance improves almost linearly with the increasing order of magnitude of the test-time compute under optimal inference parameters. Based on our observations, we derive inference scaling laws for RAG and the corresponding computation allocation model, designed to predict RAG performance on varying hyperparameters. Through extensive experiments, we show that optimal configurations can be accurately estimated and align closely with the experimental results. These insights provide a strong foundation for future research in optimizing inference strategies for long-context RAG.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. GPT-4 technical report. arXiv preprint arXiv:2303.08877. 2023.
Rishabh Agarwal, Avi Singh, Lei Zhang, Bernd Bohnet, Luis Rosias, Stephanie Chan, Biao Zhang, Ankesh Anand, Zaheer Abbas, Azade Nova, John Co-Reyes, Eric Chu, Feryal Behbahani, Aleksan- dra Faust, and Hugo Larochelle. Many-shot in-context learning. Advances in Neural Information Processing Systems, 37:76930-76966. 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Extended long short-term memory. arXiv preprint arXiv:2405.0453 2024.
Iz. Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/555 2020.
Amanda Bertsch, Maor Ivgi, Uri Alon, Jonathan Berant, Matthew R Gormley, and Graham Neu- big. In-context learning with long-context models: An in-depth exploration. arXiv preprint arXiv:2405. 00200, 2024.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptinn Lespiau, Bogdan Damoc, Aidan Clark, al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pp. 2206-2240. PMLR, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020. 2020.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15522 2023.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 2020.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems, 35:16344-163355 2022.
Mor Geva, Daniel Khashabi, Elad Segal, Tushar Khot, Dan Roth, and Jonathan Berant. Did Aristotle use laptop? A question answering benchmark with implicit reasoning strategies. Transactions of the Association for Computational Linguistics, 9:346-361, 2021.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arX1v2232..00752, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 2021.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Pre-training to learn in context. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 4849-4870, 2023.
Bernal Jiménez Gutiérrez, Yiheng Shu, Yu Gu, Michihiro Yasunaga, and Yu Su. HippoRAG: Neurobi- ologically inspired long-ter term memory for large language models, arXiv preprint 1481, 2024.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training In International conference on machine learning, pp. 3929-3938. PMLR, 2020.
11

---

Published as a conference paper at ICLR 2025
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing multi- hop QA dataset for comprchensive evaluation of reasoning steps. In Proceedings of the 28th International Conference Computational Linguistics, pp.
2020.
Gautier Izacard and Édouard Grave. Leveraging passage retrieval with generative models for open domain question answering. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pp. 874-880, 2021.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yo Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research, 24(251):1-44 2023.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhan- dari, and Yuxiong He. DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309. 145009 2023.
Zhengbao Jiang, Frank F Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi- Yu, Yiming Yang, Jamie Callan, and Graham Neubig. Active retrieval augmented generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 7969-7992, 2023.
Ziyan Jiang, Xueguang Ma, and Wenhu Chen. LongRAG: Enhancing retrieval-augmenner generation with long-context LLMs. arXiv preprint arXiv:2406.153: 15319, 2024.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1; Long Pupers), pp. 1601- 1611, 2017.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6769-67868 2020.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference Learning Representations, 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.
Terry Koo, Frederick Liu, and Luheng He. Automata-based constraints for language model decoding. arXiv preprint arXiv:2407.08/000 2024.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. BABILong: Testing the limits of LLMS with long context arXiv preprint 240.::2400000049999 2024.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Séhastien MR Arnold, Vincent Perot, Siddharth Dalmia, et al. Can long-context language models subsume retrieval, RAG, SQL, and more? arXiv preprint arXiv:2406. 13121, 2024a.
Jinhyuk Lee, Zhuyun Dai, Xiaoqi Ren, Blair Chen, Daniel Cer, Jeremy Cole, Kai Hui, Michael Boratko, Rajvi Kapadia, Wen Ding, et al. Gecko: Versatile text embeddings distilled from large language models. arXiv preprint MxXiv:2403.20327 2024b.
Quinn Leng, Jacob Portes, Sam Havens, Matei Zaharia, and Michael Carbin. Long context RAG performance of large language models. arXiv preprint DrXiv:2411.035.2 2024.
12

---

Published as a conference paper at ICLR 2025
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented genera- tion for chowledge-intensive NLP tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
Mukai Li, Shansan Gong, Jiangtao Feng, Yiheng Xu, Jun Zhang, Zhiyong Wu, and Lingpeng Kong. In-context learning with many demonstration examples. arXiv preprint 2302023302.0431,,, 2023.
Tianle Li, Ge Zhang, Quy Due Do, Xiang Yue, and Wenhu Chen. Long-context LLMs struggle with long in-context learning. arXiv preprint wrXiv:2404.02060 2024.
Xi Victoria Lin, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Richard James, Pedro
Rodrigue7, Jacob Kahn, Gergely Szilvasy, Mike Lewis, al. RA-DIT: Rerrieval-augmented dual instruction tuning. In The Twelfth International Conference on Learning Representations, 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context. arXiv preprini arXiv:2310.01889 2023.
Jiachang Liu, Dinghan Shen, Yizhe Zhang. William B Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for GPT-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Inlegration for Deep Learning Architectures, pp. 100-1 14, 2022.
Sheng Liu, Haotian Ye, Lei Xing, and James Y Zou. In-context vectors: Making in context learning more effective and controllable through latent space steering, In Forty-first International Conference on Machine Learning, 2024a.
Zihan Liu, Wei Ping, Rajarshi Roy, Peng Xu, Chankyu Lee, Mohammad Shoeybi, and Bryan Catanzaro. ChatQA: Surpassing GPT-4 on conversational QA and RAG. arXiv preprint arXiv:2401.10225, 2024b.
Yao Lu, Max Bartolo. Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1086-8098 2022.
Xinbei Ma, Yeyun Gong, Pengcheng He, Hai Zhao, and Nan Duan. Query rewriting in retrieval- augmented large language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5303-5315, 2023.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2791-2809, 2022.
Fengran Mo, Kelong Mao, Ziliang Zhao, Hongjin Qian, Haonan Chen, Yiruo Cheng, Xiaoxi Li, Yutao Zhu, Zhicheng Dou, and Jian- Yun Nie. A survey of conversational search. arXiv preprint 1rXiv:2410.15576 2024.
Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Ábrego, Ma, Vincent Y Zhao, Yi Luan, Keith B Hall, Ming-Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, 2022.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak. Samuel Arcadinho, Stella Biderman, Huanqi Cao. Xin Cheng, Michael Chung, Leon Derezynski, et al. RWKV; Reinventing mns for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 14048-14077, 2023a.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yam: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023b.
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thome Yacine Jernite, Vladimir Karpukhin, Jean Maillard, et al. KILT: a benchmark for knowledge intensive language tasks. arXiv preprint arXiv:2009.0225. 2020.
13

---

Published as a conference paper at ICLR 2025
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2 12409, 2021.
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 5687-5711, 2023.
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Browe and Yoav Shoham. In-context retrieval-auggeette language models. Transactions of the Association for Computational Linguistics, 11:1316-1331, 2023.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.
Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in in-context learning. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2655-2671, 2022.
Parth Sarthi, Salman Abdullah, Aditi Tuli, Shubh Khanna, Anna Goldie, and Christopher D Man- ning. RAPTOR: Recursive abstractive processing for tree-organized retrieval. In The Twelfth International Conference on Learning Representations, 2024.
Rulin Shao, Jacqueline He, Akari Asai, Weijia Shi, Tim Dettmers, Sewon Min, Luke Zettlemoyer, and Pang Wei Koh. Scaling retrieval-basee language models with a trillion-token datastore. arXiv preprint arXiv:2407. 222.:24407..2844, 2024.
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Richard James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. REPLUG: Retrieval-augmmnted black-box language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pp. 8364 8377, 2024.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling LLM test-time compute optimally can more effective than scaling model parameters. arXiv preprint 0rXiv:2408.03314, 2024.
Philip Sun, David Simcha, Dave Dopson, Ruiqi Guo, and Sanjiv Kumar. SOAR: improved indexing for approximate nearest neighbor search. Advances in Neural Information Processing Systems, 36, 2024.
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A lengthhextrappolatble transformer, In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Lonng Papers), pp. 14590- 14604, 2023.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. MuSiQue: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539-554, 2022.
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Interleaving retrieval with chain-of-thought reasoning for knowledge-intense inteesivve multi-step questions. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10014-10037, 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thoughe prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-2433 2022.
Jerry Wei, Le Hou, Andrew Lampinen, Xiangning Chen, Da Huang, Yi Tay, Xinyun Chen, Yifeng Lu, Denny Zhou, Tengyu Ma, et al. Symbol tuning improves in-context learning in language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. pp. 968-979, 2023.
Kevin Wu, Eric Wu, and James Zou. How faithful are RAG models? quantifying the tug-of-war between RAG and LLMs' internal prior. arXiv preprint arXiv:2404. 10198, 2024.
Zhiyong Wu, Yaoxiang Wang, Jiacheng Ye, and Lingpeng Kong. Self-adaptive in in-context learning: An information compression perspective for in-context example selection and ordering. In Proceedings of the sl Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1423-1436, 2023.
Peng Xu, Wei Ping, Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Subramanian, Evelina Bakhturina, Mohammad Shoeybi, and Bryan Catanzaro. Retrieval meets long context large language models. The Twelfth International Conference on Learning Representations, 2024.
Shi-Qi Yan, Jia-Chen Gu, Yun Zhu, and Zhen-Hua Ling. Corrective retricval augmented generation. arXiv preprint 1584 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D Manning. HotpotQA: A dataset for diverse, explainable multi-hop question answering. Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2380, 2018.
Ori Yoran, Tomer Wolfson, Ori Ram, and Jonathan Berant. Making retrieval-augmented language models robust to irrelevant context. In The Twelfth International Conference on Learning Repre- sentations, 2024.
Wenhao Yu, Hongming Zhang, Xiaoman Pan, Kaixin Ma, Hongwei Wang, and Dong Yu. Chain-of-note: Enhancing robustness in retrieval-augmented language models. arXiv preprint arXiv:2311.09210, 2023.
Zhenrui Yue, Huimin Zeng, Yimeng Lu, Lanyu Shang, Yang Zhang, and Dong Wang. Evidence- driven retrieval augmented response generation for online misinformation. arXiv preprint arXiv:2403. 14952, 2024.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 12833-3777 2020.
Tianjun Zhang, Shishir G Patil, Naman Jain, Sheng Shen, Matei Zaharia, Ion Stoica, and Joseph Gonzalez. RAFT: Adapting language model domain specific rag. arXiv preprint arXiv:2403100331, 2024.
15

---

Published as a conference paper at ICLR 2025
A DISCUSSION
Retrieval. One critical factor in improving performance of RAG lies in the quality of the retrieved documents. To study how retrieval impacts final accuracy, we analyze retrieval performance and report the results across different document sizes in Appendix B In all datasets, recall scores demonstrate improvements as the number of documents increases, approaching near-perfect scores with large document sets (e.g.. ~Ik). Despite consistent gains in recall, the results show diminishing returns on discounted ranking metrics like NDCG, indicating increasing distraction within the context. This trend is also evident in Figure 5b where RAG performance peaks between 100 and 500 documents. Our observations suggest the necessity of refining retrieval (e.g., through re-ranking) to further optimize the document relevance, particularly in cases of complex, multi-hop queries. However, how the inference scaling behavior discovered in this paper would change in the presence of such a refining component remains unknown. Alternatively, iterative retrieval, as seen in IterDRAG, improves recall performance by using simpler, straightforward sub-queries to collect additional context for each intermediate answer. In summary, retrieving more documents improves recall but does not necessarily lead to better generation quality if the documents are not effectively ranked or filtered. This highlights the need for retrieval methods that dynamically adjust to minimize irrelevant content.
Error Analysis. Despite overall improvements, our error analysis in Appendix H reveals that certain errors persist, particularly in cases of compositional reasoning tasks where multiple hops of reasoning required. The common errors fall into four categories: (I) inaccurate or outdated retrieval; (2) incorrect or lack of reasoning: (3) hallucination or unfaithful reasoning; and (4) evaluation issues refusal to answer. The first category highlights the need for enhancing retrieval methods and maintaining a reliable & up-to-date knowledge base, specially for complex questions that rely on multiple supporting facts. In addition, incorrect or missing reasoning steps often result in errors or partially correct answers. In our experiments, we observe that both (1) and (2) are substantially improved with IterDRAG, suggesting the importance of interleaving retrieval and iterative generation for multi-hop queries. Moreover, developing faithful LLMs and strategies to mitigate hallucination could further enhance RAG performance. Finally, we note that existing metrics fail in certain cases (e.g., abbreviations), underscoring the need for more robust and reliable evaluation methods.
Long-Context Modeling. We also discuss the impact of long-context modeling w.r.t. RAG per- formance. In summary, we find that retrieving more documents is generally beneficial for RAG performance, demonstrated in Section 4 Nevertheless, naively extending the context length in each generation step does not always lead to better results. Specifically, DRAG performance peaks at around 10a tokens, while IterDRAG achieves optimal performance around 10 tokens by leveraging multiple rounds of generation. For instance, as seen in the performance plateau in Figure [nd Figure [1] LLMs struggle to effectively utilize very long contexts (2 10y tokens) in each iteration, potentially due to inherent limitations of long-context modeling. Our observations suggest that: (1) the model's ability to identify relevant information from extensive context remains to be improved, especially when presented with large quantity of "similar" documents; (2) the long-context modeling should be further refined to enhance in-context learning capabilities, where multiple lengthy demonstrations are provided.
Trade-Off Between Inference Compute and RAG Performance. In our experiments, we observe consistent benefits of inference scaling using DRAG and IterDRAG, potentially changing the optimal trade-off between inference compute and RAG performance. Existing methods often exhibit dimin- ishing returns when scaling inference compute beyond certain thresholds where RAG performance plateaus. As a result, the optimal trade-off between inference compute and RAG performance is unlikely to be found beyond these thresholds, as further investment in scaling inference compute becomes inefficient. In contrast, our findings demonstrate that long-context RAG performance can improve almost linearly with increased test-time compute when optimally allocated. Therefore, the optimal trade-off in our setting largely depends on the inference budget, with higher budgets consistently yielding steady gains. Combined with the computation allocation model for RAG, this approach enables the derivation of a (nearly) optimal solution for long- contect RAG given computation constraints.
16

---

Published as a conference paper at ICLR 2025
B RETRIEVAL QUALITY
We assess the retrieval quality of DRAG and IterDRAG using the Gecko- model (Lee et al. 2024b) and evaluate their impact on final RAG performance. Specifically, we retrieve varying numbers of documents per input query and measure the retrieval quality using three metries: Recall, NDCG, and MRR, with document counts ranging from I to 2k. The retrieval results of DRAG are shown in Figure 7 In addition, we evaluate the quality of iterative retrieval, where maximum of five interleaving retrieval steps are performed. Here, we retrieve 50 documents at each step and use 2-shot setting, with the results in comparison to DRAG in Table 5
MRR
Bamboogic
HotpooQA
MuSiQuc
ZWialMulUHippAA
Scoree
04-
08-
0.4
10'
12
10"
30
30
30"
Number of Decumants
Number Decuments
Number of Decumonts
Number of Documents
Figure 7: Retrieval performance of DRAG on different datasets.
In Figure 7 recall demonstrates consistent improvements as the number of documents increases, approaching near-perfect scores when large document sets (e.g.. 1k) are retrieved. However, both NDCG and MRR metrics plateau early at around 100 documents, with diminishing gains as the document count further rises. This divergence suggests that while more documents lead to better recall, the relevance and ranking quality (captured by NDCG and MRR) do not improve proportionally, and even introduce extensive noise. Therefore, higher recall doesn' 't necessarily translate into better final answer quality when the retrieved documents aren't 't effectively ranked or filtered.
Table 5: Retrieval performance of DRAG and IteDDRAG (k = 50 documents, m 2 shots).
Bamboogle
HotpotQA NDCG
MuSiQue
2WikiMultiHoOQA
Recall
NDCG
MRR
Rocall
MRR
Recall
NDCG
MRR
Recall
NDCG
MRR
DRAG
0.632
0.321
0.239
0.783
0.535
0.465
0.509
0.255
0.188
0.722
0.421
0.336
IteDRAGG
0.736
0.420
0.346
0.855
0.549
0.478
0.670
0.365
0.291
0.935
0.605
0.528
Unlike the one-step retrieval in DRAG, iterative retrieval based on query decomposition often yields simpler sub-queries, facilitating more effective retrieval. In addition, merging the retrieved documents from different steps typically results in higher overall retrieval performance, evidenced in Table 5 With AG, the performance gains are consistent and reach the average of 30.5%. Specifically, we observe higher gains for complex multi-hop queries (e.g., 2WikiMultHHopAA) where metric improvements can be as high as 57.1%. Moreover, the gains ranking dissouedd metrics (30.7% in NDCG and 39.9% MRR) show greater improvements compared to recall (21.7%). In summary, these findings highlight the superiority of iterative retrieval with query decomposition over one-step methods, which effectively contribute to the overall performance of IterDRAG.
C CHAIN-OF-THOUGHHT vS. ITERDRAG.
Table 6: Chain-of-thought (CoT) vs. IterDRAG results (k documents, m = shots).
HotpotQA
MuSiQue
2WikiMulttHopQA
EM
F1
Acc
EM
F1
Acc
EM
F1
Acc
CoT
40.2
51.3
45.6
8.9
16.1
10.8
33.0
37.9
36.7
TeeDRAG
44.8
59.4
52.8
17.9
30.1
25.9
57.5
69.9
72.3
17

---

Published as a conference paper at ICLR 2025
To evaluate different iterative strategies, we compare the commonly used chain-of-thoughe (CoT) with ItA (Wei et al., (022). In particular, we generate the CoT examples following Trivedi et al. 2023) and adopt the 4-shot setting with 5 documents. The results on three larger datasets (HotporQA, MuSiQue and 2wiiiMuliHopQQA) as reported in Table highlight the performance differences between these strategies, in which IterRAGG consistently outperforms CoT with signifi-
cant improvements. Such difference can be traced back to three key factors: (1) the retrieval quality of CoT is limited without interleaving retrieval as IterDRAG; (2) Gemini 1.5 Flash is relatively small and may not perform well in free-form reasoning in comparison to larger LLMs; and (3) the generated CoT examples are less informative than handcrafted ones and underperform compared to constrained decoding with Self-Ask (Press et al.) 2023; Koo et al. 2024). Consequently, IterDRAG
demonstrates its effectiveness as a scalable method for knowledge-intenee tasks.
D ADDITIONAL RAG RESULTS
11
00.0
7 0M
18:
02
200
2-0
2-0
515
40.00
7 Me
13
42.1 1eaaara
27 ma
46.0
05.000
2
20.0
48.0
54.4
2
41
00
2-
07
8:
50.0
37.0000
200
40.
0033
EM Performance
F1 Performance
Acc Performance
Figure 8: IterDRAG performance heatmap for different metrics averaged across datasets.
We report the IterDRAG results averaged across datasets in Figure shown as heatmaps where the x- axis represents the number of documents and the y-axis represents the number of shots. Performance is color-coded, with blue indicating lower values and red indicating higher values. The best-performing combinations are located toward the bottom right of each heatmap, which corresponds to longer context lengths. In comparison to DRAG, as reported in Figure the optimal number of in-context examples higher 32, which highlights the importance in-context demonstrations in enabling better query decomposition and interleaved retrieval. Combined with multiple generation steps, IterDRAG further improves RAG performance over DRAG.
2
2-06
514
6:5
56.9
62.7
60.1
64/
65.7
152
18.5
52.0
53.2
51.4
41.5
14.0
533
66.1
57.3
67.5
66.2
Stol
60.0
49:4
522
52.7
52.0
635
17.3
67.6
18.4
51.
TriviaQA Arc
NaturalQ,
Figure 9: Evaluation accuracy of DRAG on TriviaQA and Natural Questions (NaturalQ.).
In addition to multi-hop question answering datasets, we also report results on one-hop datasets, specifically TriviaQA and Natural Questions (Joshi et al. 2017, Kwiatkowski et al.] (019). The evaluations for one-hop datasets are performed with DRAG and presented in Figure 9 similar to Figure 8 For TriviaQA, increasing the number of documents generally leads to improved accuracy, where the highest accuracy of 69.0% is achieved with 50 documents. Natural Questions, perfor- mance increases with the number of documents up about 10 20 documents, but further increases in the document count lead to diminishing returns or even slight declines in accuracy. The highest accuracy of 54.6% is achieved with 20 documents in 1-shot, and performance drops slightly when
18

---

Published as a conference paper at ICLR 2025
more documents are included. In summary, the optimal number of shots falls between and 4. While increasing the number of shots and documents leads to initial performance gains, these improvements platcau beyond certain thresholds. This trend, in contrast multi-hop datasets, may be partially attributed to the nature of the one-hop questions and retrieval relevance.
Table 7: Strategy QA aceuracy results.
StrategyQA
Zero-shot QA
shny QA
RAG
DRAG
IteDDR
Acc 61.1
74.7
74.7
83.4
We also include the multi-hop and binary StrategyQA dataset in our experiments, see Table 7 Geva et al. 202T). Despite being binary questions, we observe similar trends to our main experiments. For example, DRAG consistently outperforms the baseline QA and RAG methods, with 29.3% accuracy improvement to for the baseline QA model. Furthermore, the performance is boosted with 83.4 accuracy using the iterative Iee These results demonstrate that even for binary, multi-hop tasks, iterative approaches provide substantial gains, confirming the effectiveness of both long-context and iterative strategies for inference scaling in RAG.
E ADDITIONAL RESULTS ON INFERENCE SCALING LAWS WITH GTR RETRIEVER
Optimal Confi
IterERAG
Optimal Carf
Normalled
10a
30
10f
307
107
10at
Effective Contest Length
Effective Context Length
Figure 10: Normalized performance with increasing effective context lengths MuSiQue, the results are obtained with GTR XXL. retriever and Gemini 1.5 Flash.
To enhance the generalizability of our scaling observations and validate the findings with an alternative retriever model, we conduct additional experiments using the open-source GTR XXL retriever and Gemini 1.5 Flash (Ni et al.] 2022). Figure 10 shows the results on MuSiQue, evaluated using 100 sampled examples from the dataset for computational efficiency. Different from the performance plateau in standard RAG, DRAG and IterDR yields consistent performance gains with increasing context length, especially with IterDRAG at longer context lengths. Overall, the results demonstrate consistent patterns in inference scaling even with a different retriever model, highlighting the potential of expanding text-time compute in long-context RAG for more generalized scenarios.
F ADDITIONAL RESULTS ON INFERENCE SCALING LAWS FOR RAG
We present data-specific results on the relationship between the performance and the effective context length. Figure [] presents the results on the other three datasets other than MuSiQue (See Figure [ for visualized results on MuSiQue). We observe different behavior depending the datasets. For instance, the gains are more linear and consistent on Bamboogle and MuSiQue, and almost linear on 2WikiMultHHpQAA until 1M tokens. However, HotpotQA and 2WikiMultiHopQA with effective context length longer than 100k tokens exhibit more sigmoidal patterns, likely due to the difficulty of the datasets and the quality of the retrieved documents.
19

---

Published as a conference paper at ICLR 2025
RAG
DRAG IterDRAG
Optimal Config
Optimal Conf
102
105
10F
307
Effective Context Langth
Effactive Contect Langth
(a) Normalized performance effective context lengths Bamboogle.
Contg
Normallend Performance
Cpttimal g
Length
Length
(b) Normalized performance vs. effective context lengths on HotpotQA.
Optimal Config
IterDRAG
Optimal Conf
Normalized
107
20
107
Effective Contest Length
Effective Context Length
(c) Normalized performance vs. effective context lengths on 2WikiMultHHoPQA.
Figure 11: Normalized performance with increasing effective context lengths on different datasets.
G ADDITIONAL RESULTS ON COMPUTATION ALLOCATION MODEL FOR RAG
Table 8: Computation allocation mode of Gemini 1.5 Flash with p-value, R² and MSE statistics.
b
R²
MSE
Value 0.325
0.101
0.177
-0.067
-0.008
0
-0.730
0.903
0.085
y-value
0.000
0.000
0.000
0.000
0.092
N/A
0.000
N/A
N/A
We further explore the findings on the computation allocation model. In particular, we report the estimated parameters along with p-values, R², and MSE statistics in Table In our implementation, we constrain the last element b, leaving six learnable parameters in total. Our analysis shows that all parameters are statistically significant, except for b,. which has p-value slightly above 0.05. Nonetheless, our experiments suggest that retaining b, improves generalization in many cases, such IterDRAG multi-hop datasets. For sigmoid scaling, we fit custom function between the predicted P and ground truth P values, defined a (xx) I4c 10.40 2.18.
3.30
We also visualize the predictions for IterDRAG across different datasets in Figure 12 where each subplot represents a dataset and each line corresponds to a document setting (k). The inference compute is scaled by increasing the number of in-context examples (m) and generation iterations (n)..
20

---

Published as a conference paper at ICLR 2025
100c
220-Bo6
Bemboogge
HotporGA
MuSiQue
Normalted Parforanic anmmaancc
Figure 12:
The estimated performance using the proposed computation allocation model vs. actual metric values in IterDRAG, The subplots represent different datasets, where each line corresponds to a fixed number of documents, we scale the context length by increasing the number of shots.
(a) Performance vs. predicted surface for DRAG.
(b) Performance vs. predicted surface for IterDRAG.
Figure 13: Normalized performance vs. predicted surface for DRAG and IterDRAG.
Here, we find similar trends to those in Figure although IterDRAG shows larger variations compared to DRAG. HotpotQA and 2WikiMultiHopQAA show more consistent trends with the predictions, likely due to the predominance of multi-hop queries. In summary, our findings are consistent for both DRAG and IterDRAG, demonstrating that RAG performance can be accurately modeled by computation allocation model for RAG. For Bamboogle, HotpotQA and 2WikiMuliHHppA. provide the normalized performance with increasing effective context lengths in Figure IT which observe similar trends to the results MuSiQue (See Figure I). We also illustrate the prediction surface for both DRAG and IterDRAG in Figure
H ERROR ANALYSIS
Despite the performance gains from scaling effective context length, RAG performance challenging datasets like MuSiQue remain moderate, even IterDRAG. address this, we analyze the mistakes in both DRAG and IterDRAG to examine the limitations and errors inherent in these approaches. In the following, we explore common failure cases (See Figure 15) to understand where each method falls short and how they could be further improved.
We provide selected example mistakes from Figure 15a to Figure 15d with retrieved documents omitted for brevity. The reasons for common errors be grouped into four categories: (1) inaccurate or outdated retrieval; (2) incorrect or lack of reasoning; (3) hallucination or unfaithful reasoning; and (4) evaluation issues refusal to answer. We elaborate these categories below:
21

---

Published as a conference paper at ICLR 2025
In-context examples (m)
Documents (k)
Documents (k)
Documents (k)
Documents
Example Query
Exemple Ansver
Documents
Example Query
Example Anower
Test Documents
Query
Generate with RAG Iterative RAG
Findll Answer
Exampla
Exampee
Prompt
Figure 14: Input prompt that comprises of m in-context examples, the test documents and query, in which each document chunk consists of k retrieved documents. For IterDRAG, the example answers additionally provide sub-queries and intermediate answers as demonstrations.
Inaccurate or outdated retrieval: A major source of RAG errors stems from the retrieval process, where relevant knowledge is not correctly retrieved. For example, in the first question of Figure 15c the top-50 retrieved documents do not contain the correct answer. similar issue occurs in the second QA pair, where outdated retrieval results fail to provide useful information. In the third case, although both battles are retrieved, the initial documents overly focus the Battle of Manila, leading to incorrect response.
Incorrect lack of reasoning: Beyond retrieval issues, incorrect reasoning chains are another common source of errors. For example, in the first case in Figure 15b although the correct documents are retrieved, the reasoning process is incomplete (i.e., no explicit comparison of the mountain heights), leading to an incorrect answer in DRAG. Similarly, second and third cases, the reasoning is either absent (as DRAG) or flawed. result, reasoning-related errors tend to occur more frequently in difficult questions and in the one-step DRAG approach.
Hallucination or unfaithful reasoning: Other than retrieval and reasoning, hallucination and unfaithful reasoning also contribute errors in knowledge-intensive tasks. In the first case, the prediction is incorrect and cannot be found in the retrieved documents, As for the cases, while the answers related, certain steps in the reasoning chain are flawed and cause errors in the final answers. These highlight the persistent challenge of hallucination in LLMs, particularly long-context generation tasks.
Evaluation issues or refusal to answer: Finally, we observed several evaluation issues that may lead to inaccurate evaluation. For instance, the of abbreviations variations in date format can result in incorrect scoring across all metrics. Moreover, our experiments do not account for abstaining from answering, which could cause unfair scores.
I IMPLEMENTATION
In our experiments, we utilize the Gecko-1B (en) embedding model to index both the documents and input queries (Lee et al. 2024b). using Wikipedia passages from the KILT benchmark the document source (Petroni et al. 2020). In test-time, the input query is compared against all embeddings in the corpus, and the top-k neighbors are selected for inference. Each document then truncated on the right side to maximum of 1024 tokens using whitespace tokenization. For each example, we arrange the elements in the following order: documents, query, and label, with the retrieved documents listed in reverse order, placing the higher-ranked documents closer to the query (Liu et al. 2024b). Consequently, the prompt comprises of multiple in-context examples, followed by the test documents and test query, as illustrated in Figure 16
For generation, we utilize Gemini 1.5 Flash for more efficient experiments. In DRAG, inference scaling achieved by increasing the context length through the combination of documents (k) and
22

---

Published as a conference paper at ICLR 2025
in-context examples (m). Then, the prompt (See Figure 16a is provided to the model for one-time generation using the default generation parameters. For IteDDAGG the input prompt is constructed in similar fashion, with the example answers consisting of assembled sub-queries, intermediate answers, and the final answer (See Figure 16b). Here, we scale test-time compute by incorporating iterative retrieval and generation, along with the increase of documents and demonstrations, In each iteration, we restrict the generation to adhere to the Self-Ask format, in which the response should start with "Follow up: "Intermediate answer: or "So the final answer is: (Koo al.) 2024). Each iteration begins with the generation of a sub-query and concludes with the production of an intermediate answer. If sub-query is generated, additional documents are retrieved and appended to the initial set (i.e., Test Documents in Figure 16). after which the model generates an intermediate answer. We allow up to five iterations, after which the model forced to produce the final answer.
To evaluate the estimated parameters within computation allocation model for RAG, we normalized the performance metries by subtracting the mean and dividing by the standard deviation for each dataset and metric. For DRAG, the effective context length is calculated by counting the tokens in the prompt, while for IterDRAG, it determined by summing the context tokens across all inference requests. We constrain the last parameter in b and perform ordinary least squares estimate rest six parameters in Equation ( To prevent numerical instability, we shift the values by a small constant of 0.01. When computing At2 and MSE, manage noisy data by excluding peak and valley outliers in our experiments. However, for domain generalization and length extrapolation, all data points are included in the evaluation. To predict downstream task performance. i should be computed for each task. Specifically, in each strategy and task: àoc P(k = 1, m 0, P 0, m O, 1). Shx P(k 0, m I, = 1) P(k m 0, 1). For the predicted optimal hyperparameters, we present the actual metric values to validate the efficacy of computation allocation model for RAG.
23

---

Published as a conference paper at ICLR 2025
Inaccurate or outdated retrieval
Question: What the lowest clevation of the longest railway tunnel? irodiction: 500 meters
Annotation: 312
Question: According to QS World University Rankings, where does the college that Ibrahim Shihata attended rank?
Prediction: 3rd
Annotation: 551-600
Question: Which battle occurred first, Battle Manila Battle Guam?
Prediction: Battle of Manila
Annotation: Battle of Guam
(a) Example mistakes due to inaccurate outdated retrieval.
Incorrect lack of reasoning
Queation:
Which mountain, Masherbrum or Khunyang Chhish, is taller moun-
tain?
Prediction: Masherbrum
Annotation:
Khunyang Chhish
Question: What date death the director film The Organization (Film)?
Prediction: April 15, 2018
Annotation: December 12. 2012
Question: Who introduced system of musical notation the 14th century that is used in the area where most of the invasion of the eastern Roman Empire took place? Predi in: Philippe de Vitry
Annotation: John Kukuzelis
(b) Example mistakes due to incorrect or lack of reasoning.
Hallucination or unfaithful reasoning
Question: Who was the last emperor of the dynasty that succeeded the Song dynasty?
Predi ction: Emperor Yuanzhen
Annotation: Toghon Temür
Question: What is another notable work by the illustrator of Sylvester and the Magic Pebble?
Prediction Shrek!
Annotation Doctor De Sato
Question; In what movie did a Kenyan- MMxxicn actress, who graduated from Hampshire College, star in in 2015?
Predi ct inno Queen of Katwe
Annotation: Star Wars: The Force Awakens
(c) Example mistakes due to hallucination or unfaithful reasoning.
Evaluation issues or refusal to answer
Question: The most populous city in Punjab is how large (area wise)? Prediction 310 sq. km
Annotation: 310 square kilometers
Question ion: Renáta Tomanová and Larisa Neiland are former professional athletes for what sport?
Predi in: Tennis
Annotation: Professional tennis
(d) Example mistakes due to evaluation issues or refusal answer.
Figure 15: Example mistakes of DRAG and IterDRAG across datasets.
24

---

Published as a conference paper at ICLR 2025
Prompt for DRAG
You are an expert in question answering. am going to give you one or more example triples of context, question and answer, in which the context may or may not be relevant to the question. The examples will be written.
Context (which may may not be relevant): <Retriever documents:
Question: What is the place of birth the director of film Servant's Entrance?
Answer: Helsingfors
<Further remonstrat ons> inns>
After the examples, I am going provide another pair of context and question, in which context may may not relevant the question. want you answer the question. Give only the answer, and no extra commentary, formatting, or chattiness. Answer the question.
Context (which may may not be relevant):
<Retrieved documents>
Question: Who was born out ('Tomaa Henry Holland and Jean-Mandé Sigogne? Answer:
(a) Example prompt for DRAG. The prompt comprises of instructions and varying number of demonstrations, followed by test example.
Prompt for IterDRAG
You are an expert in question answering. am going to give you more example sets context, question, potential follow up questions and their respective answers, which the context may or may not be relevant to the questions. The examples will be written.
Context:
<Retrieved documents>>
Question: What nationality is director film Boggy Creek And The Legend Continues?
Follow up: Who is the director of the film Boggy Creek II: And The Legend Continues? Intermediate
The director of the film Boggy Creek II: And The Legend Continues Charles B. Pierce.
Follow up: What the nationality of Charles B. Pierce? Intermediate answer: The nationality of Charles B. Pierce is American. So the final answer is: American
<Further demonstrats
After the examples, am going provide another pair of context and question, which the context may or may not be relevant to the question. want you to answer the question. When needed, generate follow up question(s) using the format FFollw up: X', where X is the follow up question. Then, answer each follow up question using 'Intermediate answer: X' with being the answer. Finally, answer question with the format the final answer is: X', where X the final answer.
main
Context:
<Retrieved documents (with interl eereingeviig retrieval)>
Question: Where the director of film Death Of A Friend born? Follow up: Intermediate answer: So the final answer
Example prompt for IterRAGG. The prompt comprises of instructions and varying number demonstrations. followed by a test example. each iteration, we control the generation follow the Self-Ask format with constrained decoding.
Figure 16: Example prompts for DRAG and IterDRAG.
25


Reviews:
[{"review_id": "7Q6KZbhhal", "reviewer": "Reviewer_RvZP", "summary": "This paper studies the inference scaling behaviors of two retrieval augmented generation (RAG) methods, demonstration-based RAG (DRAG) and iterative demonstration-based RAG (IterDRAG). The inference computation can be scaled in multiple ways, including increasing the number of retrieved documents, in-context examples, or introducing additional generation steps in IterDRAG. Experimental results show DRAG and IterDRAG achieve scaling properties with the proposed configurations, and demonstrate the performance of DRAG and IterDRAG can scale almost linearly with an increasing computation budget. Besides, the paper also learns a computational allocation model that could provides configuration guidance for DRAG and IterDRAG.", "questions": "Please see the question above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper studies two interesting research questions including the scaling behavior and the prediction of test-time computation allocation long-context RAG methods. The paper conducts systematical experiments on inference scaling of long-context RAG models, and reveals the scaling properties of DRAG and IterDRAG, i.e., the performance improves almost linearly with optimal configuration. Besides, the computational allocation model generalizes well across domains and context lengths, which potentially helps the community to better configure RAGs.", "weaknesses": "I have a question on the application of the computational allocation model. When pretraining LLMs, computational allocation models are crucial since pretraining is extremely resource-intensive. However, inference is typically much less costly by comparison. So, why not determine the best configuration by simply searching it?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G7Tsl7FiAK", "reviewer": "Reviewer_Xv2A", "summary": "This paper explores inference scaling in long-context RAG; it analyses downstream QA results in different configurations, showing that RAG performance improves ~linearly with increasing test-time compute under optimal inference parameters. Furthermore, based on these observations, authors derive a set of inference scaling laws for RAG and a computation allocation model.\n\nMy only concern is to what extent the current model generalises -- from Fig. 4b, it seems like the considered models might suffer from \"lost in the middle\" (https://arxiv.org/abs/2307.03172) problems, while more recent long-context models seem to suffer from this significantly less (e.g., https://arxiv.org/abs/2404.16811)", "questions": "More recent models may suffer less from \"lost in the middle\" issues -- does the current analysis still hold?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Inference-time scaling laws for RAG systems -- extremely interesting, and the community really needs an analysis like this one.", "weaknesses": "It is not clear whether the current analysis may generalise to future SFT/RLHF regimens.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5NMfXdyPXV", "reviewer": "Reviewer_Tea3", "summary": "The paper explores inference scaling for retrieval augmented generation (RAG) in long-context large language models (LLMs), focusing on strategies beyond simply expanding the knowledge base. The authors investigate how in-context learning and iterative prompting can optimize the use of additional compute resources during inference. They address two primary questions: the benefit of inference computation scaling for RAG and the prediction of optimal compute allocation within a given budget. Their findings reveal that optimal allocation of inference computation results in nearly linear performance gains for RAG, a phenomenon described as inference scaling laws. The authors also develop a computation allocation model that accurately predicts the optimal inference parameters, with experimental results showing up to 58.9% performance improvements on benchmark datasets compared to standard RAG configurations.", "questions": "1. Can you provide a straightforward explanation of your findings and how they guide the use of long-context LLMs for RAG?\n2. If other prompts or methods are used, is your computation allocation model still applicable?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The research question is quite interesting, as there is not much work on inference time scaling for RAG; this study systematically explores this area and may draw some attention.", "weaknesses": "I am concerned that this work is more suitable as a technical report rather than a research-oriented study. There is considerable related work combining long-context LLMs and RAG, and the main contribution of this work is mainly the proposed RAG inference scaling law. However, this conclusion is method-specific and may not apply to other methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "edwmDvaKSF", "reviewer": "Reviewer_FTdU", "summary": "This paper systematically investigates the performance of RAG systems as inference computation resources scale up, demonstrating an almost linear improvement in RAG performance with optimal test-time compute allocation. Furthermore, the authors derive an empirical scaling law that can predict the optimal inference parameter configuration for RAG systems under various computational budgets.", "questions": "See Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. This paper studies a significant issue in the LLM community.\n2. The authors conducts extensive experiments across various datasets and evaluation metrics, yielding convincing results.\n3. This paper is well structured, clearly written, and easy to understand.", "weaknesses": "1. Novelty of the proposed methods is somewhat limited. Both DRAG and IterDRAG are conventional approaches in academia and industry, similar to some classic methods like ITER-RETGEN, IRCoT, and Active-RAG.\n\n2. The experiments were only conducted with Gemini-1.5-Flash and Gecko retrieval. Gemini-1.5-Flash is a relatively small LLM, and different conclusions might be drawn on larger, more powerful LLMs. Moreover, if the scaling laws derived could be generalized to other LLMs and retrieval methods, it would add greater value to this work.\n\n3. Some experimental phenomena lack more in-depth discussion. For example, on line 371, the impact of few-shot examples on IterDRAG should be assessable, at least through ablation studies to determine whether it is the in-context query decomposition or the knowledge extraction capability that is more influential.\n\n4. While increasing the context size can improve RAG performance, it also leads to greater inference time and token consumption, especially when using iterative retrieval. The authors did not discuss this trade-off.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)