Inference Scaling for Long-Context Retrieval Augmented Generation

TL;DR

This paper explores how scaling inference computation in retrieval augmented generation (RAG) with long-context LLMs improves performance, introduces inference scaling laws, and predicts optimal compute allocation for enhanced knowledge utilization.

Contribution

It introduces inference scaling laws for RAG, models optimal test-time compute allocation, and demonstrates significant performance gains through scaled inference configurations.

Findings

Scaling inference compute yields nearly linear performance gains.

The proposed model accurately predicts optimal inference parameters.

Scaling inference achieves up to 58.9% improvements on benchmarks.

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.