長い文脈検索増強生成のための推論スケーリング

Inference Scaling for Long-Context Retrieval Augmented Generation

October 6, 2024

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

cs.AI

要旨

推論計算のスケーリングは、さまざまな設定で長いコンテキストを持つ大規模言語モデル（LLM）の潜在能力を開放しました。知識集約的なタスクでは、増加した計算リソースはしばしばより多くの外部知識を組み込むために割り当てられます。しかし、そのような知識を効果的に活用しない限り、単にコンテキストを拡張するだけではパフォーマンスが常に向上するわけではありません。本研究では、リトリーバを拡張した生成（RAG）のための推論スケーリングを調査し、知識の量を単純に増やすことを超えた戦略を探ります。具体的には、インコンテキスト学習と反復プロンプティングに焦点を当てます。これらの戦略は、テスト時の計算をスケーリングするための追加の柔軟性を提供し（たとえば、取得される文書や生成ステップを増やすことにより）、LLMがコンテキスト情報を効果的に取得および利用する能力を向上させます。我々は、次の2つの主要な問いに取り組みます：（1）最適に構成された場合、RAGのパフォーマンスは推論計算のスケーリングからどのように利益を得るか？（2）RAGのパフォーマンスと推論パラメータの関係をモデル化することで、所与の予算に対する最適なテスト時計算の割り当てを予測できるか？我々の観察結果は、推論計算を増やすことが最適に割り当てられた場合、RAGのパフォーマンスがほぼ線形的に向上することを示し、これをRAGの推論スケーリング則として説明します。さらに、異なる推論構成にわたるRAGのパフォーマンスを推定するための計算割り当てモデルを開発します。このモデルは、さまざまな計算制約下で最適な推論パラメータを予測し、実験結果と密接に一致します。これらの最適な構成を適用することで、長いコンテキストを持つLLMにおける推論計算のスケーリングは、標準のRAGに比べてベンチマークデータセットで最大58.9％の利益をもたらすことを示します。

English

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 strategies beyond simply increasing the quantity of knowledge. We focus on two inference scaling strategies: 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.