MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity

要約

検索拡張生成 (RAG) は、知識集約型タスクにおける言語モデルの生成パフォーマンスを向上させるのに非常に効果的であることが証明されています。
ただし、既存の RAG フレームワークは、無差別に取得を実行するか、厳格な単一クラス分類子に依存して取得方法を選択するため、さまざまな複雑さのクエリ全体で非効率性や次善のパフォーマンスが発生します。
これらの課題に対処するために、クエリの複雑さに基づいて最適な取得戦略を動的に選択する強化学習ベースのフレームワークを提案します。
% 私たちのソリューション 私たちのアプローチは、各取得方法を個別の「アーム」として扱い、探索と活用のバランスを取ることで選択プロセスを適応させる、マルチアーム バンディット アルゴリズムを活用しています。
さらに、精度と効率のバランスをとる動的な報酬関数を導入し、たとえ正しい結果が得られたとしても、より多くの取得ステップを必要とする方法にペナルティを与えます。
私たちの方法は、取得コストを削減しながら、複数のシングルホップおよびマルチホップのデータセットで新しい最先端の結果を達成します。
私たちのコードは https://github.com/FUTUREEEEEE/MBA で入手できます。

要約(オリジナル)

Retrieval Augmented Generation (RAG) has proven to be highly effective in boosting the generative performance of language model in knowledge-intensive tasks. However, existing RAG framework either indiscriminately perform retrieval or rely on rigid single-class classifiers to select retrieval methods, leading to inefficiencies and suboptimal performance across queries of varying complexity. To address these challenges, we propose a reinforcement learning-based framework that dynamically selects the most suitable retrieval strategy based on query complexity. % our solution Our approach leverages a multi-armed bandit algorithm, which treats each retrieval method as a distinct “arm” and adapts the selection process by balancing exploration and exploitation. Additionally, we introduce a dynamic reward function that balances accuracy and efficiency, penalizing methods that require more retrieval steps, even if they lead to a correct result. Our method achieves new state of the art results on multiple single-hop and multi-hop datasets while reducing retrieval costs. Our code are available at https://github.com/FUTUREEEEEE/MBA .

arxiv情報

提供元, 利用サービス

arxiv.jp, Google