Масштабирование вывода для увеличения поколения с увеличенным извлечением длинного контекста

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 эффективно усваивать и использовать контекстуальную информацию. Мы рассматриваем два ключевых вопроса: (1) Как производительность RAG выигрывает от масштабирования вычислений вывода при оптимальной настройке? (2) Можем ли мы предсказать оптимальное распределение вычислений на тестирование для заданного бюджета, моделируя отношение между производительностью RAG и параметрами вывода? Наши наблюдения показывают, что увеличение вычислений вывода приводит к почти линейному увеличению производительности RAG при оптимальном распределении, отношение, которое мы описываем как законы масштабирования вывода для RAG. На основе этого мы дополнительно разрабатываем модель распределения вычислений для оценки производительности RAG при различных конфигурациях вывода. Модель предсказывает оптимальные параметры вывода при различных ограничениях вычислений, которые тесно соответствуют экспериментальным результатам. Применяя эти оптимальные конфигурации, мы демонстрируем, что масштабирование вычислений вывода для LLM с длинным контекстом дает до 58,9% прироста на эталонных наборах данных по сравнению со стандартным RAG.

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.