Escalonamento de Inferência para Geração Aprimorada de Recuperação de Contexto Longo

Inference Scaling for Long-Context Retrieval Augmented Generation

October 6, 2024

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

cs.AI

Resumo

A escalabilidade da computação de inferência desbloqueou o potencial de modelos de linguagem grandes de longo contexto (LLMs) em diversos cenários. Para tarefas intensivas em conhecimento, a computação aumentada é frequentemente alocada para incorporar mais conhecimento externo. No entanto, sem utilizar efetivamente esse conhecimento, simplesmente expandir o contexto nem sempre melhora o desempenho. Neste trabalho, investigamos a escalabilidade da inferência para geração aumentada por recuperação (RAG), explorando estratégias além de simplesmente aumentar a quantidade de conhecimento. Concentramo-nos em duas estratégias de escalabilidade da inferência: aprendizado em contexto e solicitações iterativas. Essas estratégias fornecem flexibilidade adicional para escalar a computação no momento do teste (por exemplo, aumentando os documentos recuperados ou os passos de geração), melhorando assim a capacidade dos LLMs de adquirir e utilizar efetivamente informações contextuais. Abordamos duas questões-chave: (1) Como o desempenho do RAG se beneficia da escalabilidade da computação de inferência quando configurado de forma otimizada? (2) Podemos prever a alocação ótima de computação no momento do teste para um determinado orçamento modelando a relação entre o desempenho do RAG e os parâmetros de inferência? Nossas observações revelam que aumentar a computação de inferência leva a ganhos quase lineares no desempenho do RAG quando alocados de forma otimizada, uma relação que descrevemos como as leis de escalabilidade da inferência para o RAG. Com base nisso, desenvolvemos ainda mais o modelo de alocação de computação para estimar o desempenho do RAG em diferentes configurações de inferência. O modelo prevê parâmetros de inferência ótimos sob várias restrições de computação, que estão alinhados de perto com os resultados experimentais. Ao aplicar essas configurações ótimas, demonstramos que a escalabilidade da computação de inferência em LLMs de longo contexto alcança ganhos de até 58,9% em conjuntos de dados de referência em comparação com o RAG padrão.

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.