Цепочка извлечения с увеличением генерации

Chain-of-Retrieval Augmented Generation

January 24, 2025

Авторы: Liang Wang, Haonan Chen, Nan Yang, Xiaolong Huang, Zhicheng Dou, Furu Wei

cs.AI

Аннотация

Эта статья представляет подход к обучению моделей RAG, подобных o1, которые извлекают и рассуждают над соответствующей информацией пошагово перед генерацией окончательного ответа. Традиционные методы RAG обычно выполняют один шаг извлечения перед процессом генерации, что ограничивает их эффективность в решении сложных запросов из-за неполных результатов извлечения. В отличие от этого, наш предложенный метод, CoRAG (Chain-of-Retrieval Augmented Generation), позволяет модели динамически переформулировать запрос на основе изменяющегося состояния. Для эффективного обучения CoRAG мы используем отборочное сэмплирование для автоматического создания промежуточных цепочек извлечения, тем самым дополняя существующие наборы данных RAG, которые предоставляют только правильный окончательный ответ. На этапе тестирования мы предлагаем различные стратегии декодирования для масштабирования вычислений модели во время тестирования путем контроля длины и количества сэмплированных цепочек извлечения. Экспериментальные результаты на нескольких бенчмарках подтверждают эффективность CoRAG, особенно в задачах вопросно-ответной системы с несколькими шагами, где мы наблюдаем более чем 10-процентное улучшение в оценке EM по сравнению с сильными базовыми моделями. На бенчмарке KILT CoRAG устанавливает новый рекорд производительности в широком спектре задач, требующих знаний. Кроме того, мы предлагаем комплексный анализ для понимания поведения масштабирования CoRAG, заложив основу для будущих исследований, направленных на разработку фактических и обоснованных моделей базовых знаний.

English

This paper introduces an approach for training o1-like RAG models that retrieve and reason over relevant information step by step before generating the final answer. Conventional RAG methods usually perform a single retrieval step before the generation process, which limits their effectiveness in addressing complex queries due to imperfect retrieval results. In contrast, our proposed method, CoRAG (Chain-of-Retrieval Augmented Generation), allows the model to dynamically reformulate the query based on the evolving state. To train CoRAG effectively, we utilize rejection sampling to automatically generate intermediate retrieval chains, thereby augmenting existing RAG datasets that only provide the correct final answer. At test time, we propose various decoding strategies to scale the model's test-time compute by controlling the length and number of sampled retrieval chains. Experimental results across multiple benchmarks validate the efficacy of CoRAG, particularly in multi-hop question answering tasks, where we observe more than 10 points improvement in EM score compared to strong baselines. On the KILT benchmark, CoRAG establishes a new state-of-the-art performance across a diverse range of knowledge-intensive tasks. Furthermore, we offer comprehensive analyses to understand the scaling behavior of CoRAG, laying the groundwork for future research aimed at developing factual and grounded foundation models.