Inference Scaling for Long-Context Retrieval Augmented Generation

Response to Reviewer FTdU

We thank you for your valuable feedback on our submission and are excited that you find our work "convincing", "well structured" and "easy to understand"! We want to clarify your concerns and questions in the following. *Novelty of the proposed methods is somewhat limited. Both DRAG and IterDRAG are conventional approaches...* - Thank you for your comment. Existing work like ITER-RETGEN, IRCoT and Active-RAG mostly focus on proposing specific inference strategies to improve RAG performance without systematically understanding the dynamic of how RAG performance changes as inference computation increases. - In contrast, our primary goal is not to introduce new retrieval or generation strategies, but to systematically understand and model inference scaling for long-context RAG. To this end, we build on existing paradigms and explore combinations of different strategies (via in-context examples, iterative retrieval with constrained decoding etc.) to more effectively scale inference computation. Using these strategies, we demonstrate that RAG performance can scale almost linearly with increasing magnitude of inference computations when optimally configured, as opposed to the sigmoidal trend as in previous studies [1-2]. Additionally, a key contribution of our work is to quantify the relationship between RAG performance and different combinations of inference parameters using the computation allocation model, deriving near-optimal configurations across diverse scenarios. - Consequently, this distinguishes our work from others by: (1) focusing on the understanding of inference scaling in RAG; and (2) identifying practical solutions to leverage such scaling dynamics in long-context RAG performance. *The experiments were only conducted with Gemini-1.5-Flash and Gecko retrieval...* - We thank you for your suggestion. We provide additional results with GTR XXL, detailed in Appendix E of our rebuttal revision. Overall, we observe similar trends with GTR-retrieved documents, showing that long-context RAG performance consistently improves with effective context lengths. Yet due to resource limitations, we are unable to perform large-scale experiments with further LLMs. As such, we leave further exploration across a broader range of models as future work. *Some experimental phenomena lack more in-depth discussion...* - Thank you for suggesting the need for a more in-depth discussion on the impact of few-shot examples in IterDRAG. As shown by the green line in Figure 5c, IterDRAG can outperform DRAG even with fewer in-context examples (i.e., less demonstrations for in-context knowledge extraction), where we attribute such performance gains to the additional query decomposition process. To provide more concrete results and a comprehensive analysis of IterDRAG configurations, we have updated our observations and discussion in the rebuttal revision. Please refer to the revised Sec 4.4 and the additional ablation studies in Appendix D. These findings further highlight the varying effectiveness of in-context demonstrations and query decomposition (evidenced, for example, through the heatmap analysis). *While increasing the context size can improve RAG performance, it also...* - Thank you for pointing out the trade-off between RAG performance and the associated inference time and token consumption. In our experiments, we quantify inference computation by the number of input tokens across LLM calls (i.e., effective context length), and we have explicitly modeled such trade-off relationships between performance and test-time compute using our computation allocation model (e.g., as shown in Figure 1 and Figure 4). - One of the main discoveries of this paper is that there exists a better trade-off between RAG performance and the inference computation from existing strategies: (1) Existing works that only use one strategy to scale up inference computation (e.g., solely adding more documents or demonstrations) have a limitation: beyond a certain threshold, increasing the inference computation (e.g., by adding 10x more documents) no longer improves RAG performance, which is not an ideal trade-off to consume more tokens. (2) Our work demonstrates that by simply employing a combination of multiple inference strategies, further increasing the inference computation can continue to improve RAG performance almost linearly when the optimal inference parameters are identified, leading to a better trade-off than existing work. [1] Xu, Peng, et al. "Retrieval meets long context large language models." arXiv preprint arXiv:2310.03025 (2023). [2] Leng, Quinn, et al. "Long Context RAG Performance of Large Language Models." arXiv preprint arXiv:2411.03538 (2024).

Response to Reviewer Tea3

We appreciate your insightful comments on our work and are thrilled that you find our submission "interesting" and "systematic"! We want to clarify your concerns and questions below. *I am concerned that this work is more suitable as a technical report rather than a research-oriented study...* - We appreciate the concern regarding the scope of our work, yet we believe it extends beyond a technical report by providing a structured, empirical framework for understanding inference scaling in long-context RAG [1-2]. Unlike previous works [3-4] that focuses on one specific RAG strategy (e.g., increasing the number of documents or the length of documents), our study provides a comprehensive understanding of inference scaling for RAG by exploring various inference strategies across different compute budgets (measured by effective context lengths). Our experiments reveal a new scaling trend / dynamic, i.e., that RAG performance can scale almost linearly with increased magnitude of inference computation instead of sigmoidal as shown in previous studies [4], as long as one uses the right combination of inference parameters. To identify the "right" inference configuration, we also propose the computation allocation model to quantitatively model the RAG performance across different combinations of inference parameters, offering near-optimal configurations for various scenarios. Considering the above contributions, our work lays a foundation for future research in scaling inference compute in long-context RAG, offering insights that extend beyond a technical report. *Can you provide a straightforward explanation of your findings and how they guide the use of long-context LLMs for RAG?* - In summary, our study investigates inference scaling for long-context RAG, demonstrating that performance can scale nearly linearly with increasing inference compute. Building on these findings, we propose the computation allocation model, which derives near-optimal inference parameters for various knowledge-intensive tasks. - More specifically, we conduct an extensive evaluation of long-context RAG performance across diverse inference configurations and present the following findings: (1) Unlike previous studies showing that RAG performances stop increases when scaling the amount of retrieval information beyond a certain threshold, our findings reveal that RAG performance can continue to increase and scale almost linearly with increased inference computation given optimal inference parameters (e.g., in-context examples and multi-step reasoning). (2) By introducing the computation allocation model, we then provide a systematic framework to predict the right configuration of inference strategies for given budgets, enabling efficient and effective use of long-context LLMs for knowledge-intensive tasks. *If other prompts or methods are used, is your computation allocation model still applicable?* - For improved generalizability, we include supplementary experiment results using the GTR XXL retriever model, as presented in Appendix E in the rebuttal revision. In sum, similar observations are made with GTR-retrieved documents, showing that long-context RAG performance consistently improves with effective context lengths. - While there are numerous ways to scale inference compute for RAG (e.g., more documents), we focus on commonly used methods such as expanding documents, many-shot demonstrations and iterative prompting, refining these approaches for effective long-context scaling. Nonetheless, our main focus is to show that inference scaling and the modeling of such scaling is feasible for long-context RAG, yielding consistent gains with increased computation budgets. Our approach can be easily adapted for further prompt designs / model families, providing a framework to evaluate and understand the inference scaling properties of both existing and emerging RAG strategies. [1] Wu, Yangzhen, et al. "An empirical analysis of compute-optimal inference for problem-solving with language models." (2024). [2] Ruan, Yangjun, Chris J. Maddison, and Tatsunori Hashimoto. "Observational Scaling Laws and the Predictability of Language Model Performance." arXiv preprint arXiv:2405.10938 (2024). [3] Xu, Peng, et al. "Retrieval meets long context large language models." arXiv preprint arXiv:2310.03025 (2023). [4] Jiang, Ziyan, Xueguang Ma, and Wenhu Chen. "Longrag: Enhancing retrieval-augmented generation with long-context llms." arXiv preprint arXiv:2406.15319 (2024).

Response to Reviewer Xv2A

We thank you for your valuable feedback on our work, we are particularly excited that our contributions are regarded as "extremely interesting"! We hope to address your concerns and questions in the following. *It is not clear whether the current analysis may generalize to future SFT/RLHF regimens.* - While our computation allocation model and scaling strategies are designed primarily for long-context retrieval augmented generation during inference, similar settings could be applied to SFT/RLHF settings for improved long-context RAG performance. For example, for a fixed training token budget, one can conduct experiments to figure out the optimal allocation of different tasks in the training data. While this falls outside the scope of this paper which focuses more on inference scaling, we agree that this would be a promising future direction to look into. *More recent models may suffer less from "lost in the middle" issues -- does the current analysis still hold?* - Thank you for pointing out this. Gemini 1.5 (as used in our experiments) is one of the recent models that demonstrates improved performance of "lost in the middle" / long-context modeling [1, 2]. However, we still notice that RAG performance plateaus when merely increasing the number of retrieved documents. Experiments from [3] also show similar observations using a few other recent models with alleviated "lost in the middle" issue. Therefore, solely addressing "lost in the middle" (i.e., enhancing model’s ability to find relevant information in the context) may not result in a linear scaling of RAG performance with respect to the magnitude of inference computation. It is also crucial to enhance the model’s ability to integrate and reason over the "found relevant information". Consequently, we combine multiple additional scaling strategies such as adding demonstrations and iterative querying and show that it is possible to achieve near-linear performance gains in long-context RAG as the magnitude of effective context length increases. [1] Reid, Machel, et al. "Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context." arXiv preprint arXiv:2403.05530 (2024). [2] An, Shengnan, et al. "Make Your LLM Fully Utilize the Context." arXiv preprint arXiv:2404.16811 (2024). [3] Leng, Quinn, et al. "Long Context RAG Performance of Large Language Models." arXiv preprint arXiv:2411.03538 (2024).

Response to Reviewer RvZP

We appreciate your thoughtful feedback and are glad that our contributions are considered "interesting" and "systematic". Regarding your question on our computation allocation model, we agree that grid search is feasible for smaller datasets when context length is limited. However, as context length increases or with larger evaluation sets, grid search becomes extremely expensive. For instance, we estimate that conducting a grid search on MuSiQue for an effective context length budget of 5M tokens using Gemini 1.5 Flash could cost up to $18,879, even with sub-sampling from the evaluation set. While inference generally demands less computation than pretraining, directly searching for the best configuration without guidance can still lead to inefficiencies due to the vast space of inference hyperparameters. Our computation allocation model addresses this challenge by systematically modeling inference compute, delivering near-optimal configurations based on the estimated relationship between performance and inference parameters. This approach not only optimizes long-context RAG solutions within a given compute budget but also generalizes well across tasks and constraints, eliminating the need for costly and exhaustive grid searches during evaluation.