Escalado de Inferencia para Generación Aumentada con Recuperación de Contexto Largo

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

Resumen

La escalabilidad de la computación de inferencia ha desbloqueado el potencial de los modelos de lenguaje grandes de largo contexto (LLMs) en diversos entornos. Para tareas intensivas en conocimiento, la mayor capacidad de cómputo se destina frecuentemente a incorporar más conocimiento externo. Sin embargo, sin utilizar eficazmente dicho conocimiento, simplemente ampliar el contexto no siempre mejora el rendimiento. En este trabajo, investigamos la escalabilidad de la inferencia para la generación aumentada por recuperación (RAG), explorando estrategias más allá de simplemente aumentar la cantidad de conocimiento. Nos centramos en dos estrategias de escalado de inferencia: aprendizaje en contexto y solicitud iterativa. Estas estrategias brindan flexibilidad adicional para escalar la computación en tiempo de prueba (por ejemplo, aumentando los documentos recuperados o los pasos de generación), mejorando así la capacidad de los LLMs para adquirir y utilizar eficazmente información contextual. Abordamos dos preguntas clave: (1) ¿Cómo se beneficia el rendimiento de RAG del escalado de la computación de inferencia cuando está configurado óptimamente? (2) ¿Podemos predecir la asignación óptima de cómputo en tiempo de prueba para un presupuesto dado modelando la relación entre el rendimiento de RAG y los parámetros de inferencia? Nuestras observaciones revelan que aumentar la computación de inferencia conduce a ganancias casi lineales en el rendimiento de RAG cuando se asigna de manera óptima, una relación que describimos como las leyes de escalado de inferencia para RAG. Basándonos en esto, desarrollamos aún más el modelo de asignación de cómputo para estimar el rendimiento de RAG en diferentes configuraciones de inferencia. El modelo predice los parámetros de inferencia óptimos bajo diversas restricciones de cómputo, que se alinean estrechamente con los resultados experimentales. Al aplicar estas configuraciones óptimas, demostramos que escalar la computación de inferencia en LLMs de largo contexto logra hasta un 58.9% de ganancias en conjuntos de datos de referencia en comparación con RAG estándar.

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.