[論文レビュー] Chain-of-Retrieval Augmented Generation

Liang Wang, Haonan Chen+4 more|ArXiv.org|Jan 24, 2025

Speech and dialogue systems被引用数 3

ひとこと要約

CoRAG は情報を連鎖的に取得・推論する retrieval-augmented 生成モデルを訓練し、情報を多段階で取得・推論することで多段階の QA や知識集約タスクの性能を向上させ、推論時の計算量を制御する戦略を提供します。

ABSTRACT

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.

研究の動機と目的

• 反復的な取得と推論を可能にして RAG の改善を目指し、単一の取得ステップに依存しない動機づけ
• 中間の取得チェーンを拒否サンプリングを用いて QA データセットに付与
• 取得と生成のチェーンの次の行動を予測するように LLM を訓練
• チェーンベースの取得をスケールさせるためのテスト時デコード戦略を検討
• CoRAG をマルチホップ QA データセットと KILT ベンチマークで評価し、一般化とスケーラビリティを評価

提案手法

• 拒否サンプリングを用いて sub-query および sub-answer の系列を持つ QA データセットを拡張する取得チェーンを生成
• 拡張データでマルチタスク目的を用いて open-source LLM をファインチューニングし、sub-query・sub-answer・最終回答の予測を含む
• 各サブクエリに対して top-k ドキュメントを取得するリトリーバを用い、チェーンの対数尤度でチェーン品質を評価
• トークン消費を制御するためのテスト時デコード戦略（グリーディ、Best-of-N サンプリング、木探索）を提供
• データセット・リトリーバ・一般化シナリオ間でのスケーリング挙動とロバスト性を分析
• 任意で、現在の情報が十分かどうかを予測してチェーンをテスト時に停止させることを学習することも可能

(a) Test-time scaling behavior of CoRAG.

実験結果

リサーチクエスチョン

• RQ1 反復的な取得と推論は、単一ステップの RAG と比較して複雑な多段階 QA タスクの性能を向上させるか？
• RQ2 チェーン長とチェーン数といったテスト時計算量が性能と効率に与える影響は？
• RQ3 チェーンオブリトリーバルは、QA を超える多様な知識集約タスクへ一般化できるか？
• RQ4 より弱いリトリーバとモデルを使う場合と強いものを使う場合で、CoRAG の効果はどう変わるか？
• RQ5 推論時に取得チェーンを早期停止する機構を学習することは有益か？

主な発見

• CoRAG は複数のマルチホップ QA データセットで強力なベースラインを大きく上回り、いくつかのデコード戦略下で EM/F1 の顕著な改善を示した
• KILT ベンチマークでは、最も難しい FEVER タスクを除く多様なタスクで最先端の性能を達成した
• 短いチェーンから開始した場合は長いチェーンを用いると性能が向上するが、チェーン長が長くなると利得は次第に低下する
• テスト時のスケーリング挙動は、トークン消費と性能の間に対数線形の関係を示すデータセットが複数ある
• ロバストネス実験では、強力なリトリーバの使用による利益があり、弱いものでもいくらかの利益を示しつつ、タスクタイプ全体で一般化は好ましい傾向にある
• アブレーション要因としての反復的トレーニングは混在した結果を示し、インストラクション適合型の LLM はしばしば高品質な取得チェーンを生成することを示唆した

(b) An example of CoRAG in action.