Mise à l'échelle de l'inférence pour la génération augmentée de récupération à long contexte

Inference Scaling for Long-Context Retrieval Augmented Generation

October 6, 2024

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

cs.AI

Résumé

L'échelle de calcul d'inférence a débloqué le potentiel des grands modèles de langage à long contexte (LLMs) dans divers contextes. Pour les tâches intensives en connaissances, la puissance de calcul accrue est souvent allouée pour incorporer davantage de connaissances externes. Cependant, sans utiliser efficacement de telles connaissances, l'expansion du contexte seule n'améliore pas toujours les performances. Dans ce travail, nous étudions l'échelle de calcul d'inférence pour la génération augmentée par récupération (RAG), explorant des stratégies allant au-delà de simplement augmenter la quantité de connaissances. Nous nous concentrons sur deux stratégies d'échelle d'inférence : l'apprentissage en contexte et les invitations itératives. Ces stratégies offrent une flexibilité supplémentaire pour augmenter le calcul au moment du test (par exemple, en augmentant les documents récupérés ou les étapes de génération), améliorant ainsi la capacité des LLMs à acquérir et utiliser efficacement des informations contextuelles. Nous abordons deux questions clés : (1) Comment les performances de RAG bénéficient-elles de l'échelle de calcul d'inférence lorsqu'elles sont configurées de manière optimale ? (2) Pouvons-nous prédire l'allocation de calcul optimale au moment du test pour un budget donné en modélisant la relation entre les performances de RAG et les paramètres d'inférence ? Nos observations révèlent qu'augmenter le calcul d'inférence entraîne des gains presque linéaires en performances de RAG lorsqu'il est alloué de manière optimale, une relation que nous décrivons comme les lois d'échelle d'inférence pour RAG. Sur cette base, nous développons davantage le modèle d'allocation de calcul pour estimer les performances de RAG à travers différentes configurations d'inférence. Le modèle prédit les paramètres d'inférence optimaux sous diverses contraintes de calcul, qui correspondent étroitement aux résultats expérimentaux. En appliquant ces configurations optimales, nous démontrons que l'augmentation du calcul d'inférence sur les LLMs à long contexte permet d'obtenir jusqu'à 58,9 % de gains sur des ensembles de données de référence par rapport à RAG standard.

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.