Inference Scaling for Long-Context Retrieval Augmented Generation

2024-10-06

arXiv.org

AI Summary

This work addresses the inefficient allocation of computational resources in long-context retrieval-augmented generation (RAG). Unlike conventional approaches that merely scale up retrieved document volume, we propose a dynamic computation scaling paradigm for the inference phase. Our method introduces multi-step iterative prompting and in-context learning, enabling— for the first time—the formulation of an inference scaling law for RAG. We further develop a computation-budget-aware model that predicts optimal inference configurations, automatically allocating computational capacity across inference steps. This yields near-linear improvements in both performance and total compute cost. Experiments across multiple benchmark datasets demonstrate up to 58.9% improvement over standard RAG, empirically validating the critical role of inference-level computation optimization. Our approach establishes a principled, predictive framework for efficient and controllable retrieval-augmented generation in the era of large language models.

Technology Category

• Search and Optimization: Learning to Search
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• Optimizing inference computation for long-context retrieval augmented generation (RAG).
• Exploring strategies beyond knowledge quantity to enhance RAG performance.
• Predicting optimal test-time compute allocation for RAG under budget constraints.

Innovation

• Combines in-context learning and iterative prompting
• Develops computation allocation model for RAG
• Achieves 58.9% gains with optimal inference scaling