2025-11-13T01:58:10.933950

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

Tang, Gao, Li et al.

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 .

academic

MBA-RAG: 質問複雑度を通じた適応的検索増強生成のためのバンディット手法

基本情報

論文ID: 2412.01572
タイトル: MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity
著者: Xiaqiang Tang, Qiang Gao, Jian Li, Nan Du, Qi Li, Sihong Xie
所属機関: 香港科技大学(広州)、テンセント混元、武漢大学、アイオワ州立大学
分類: cs.AI
発表日: 2025年1月1日 (arXiv v4)
論文リンク: https://arxiv.org/abs/2412.01572
コードリンク: https://github.com/FUTUREEEEEE/MBA

概要

検索増強生成(RAG)は、知識集約的なタスクにおいて言語モデルの生成性能を大幅に向上させている。しかし、既存のRAGフレームワークは、検索を無差別に実行するか、検索方法の選択に硬直した単一分類器に依存しており、異なる複雑度のクエリにおいて効率が低く、性能が最適でない。これらの課題に対処するため、本論文は強化学習に基づくフレームワークを提案し、クエリ複雑度に基づいて最適な検索戦略を動的に選択できる。本手法は多腕バンディットアルゴリズムを活用し、各検索方法を異なる「腕」として扱い、探索と利用のバランスを取ることで選択プロセスに適応する。さらに、精度と効率のバランスを取る動的報酬関数を導入し、正しい結果が得られた場合でも、より多くの検索ステップを必要とする方法にペナルティを与える。本手法は複数の単一ホップおよび多ホップデータセットで新たなSOTA結果を達成し、同時に検索コストを削減した。

研究背景と動機

問題定義

既存のRAGシステムには以下の核心的な問題が存在する：

検索戦略選択の不適切さ：ほとんどのRAGフレームワークはすべてのクエリに対して無差別に検索を実行し、不要または無関係な段落を導入する可能性がある
単一方法の限界：すべてのクエリに対して単一の検索方法を使用することは効率が低く、単純なクエリは不要な計算オーバーヘッドを生じ、複雑なクエリは十分に処理されない可能性がある
監督信号の不正確性：AdaptiveRAGなどの既存の適応的手法はヒューリスティック監督を使用し、各クエリに対して1つの最適戦略のみが存在すると仮定し、検索コストが最も低いパスを選択する傾向がある

研究動機

本論文の核心的な動機は、以下を実現できるシステムを開発することである：

クエリ複雑度への動的適応：問題の複雑さに応じて検索戦略をインテリジェントに選択する
精度と効率のバランス：回答品質を保証しながら計算コストを最小化する
複数戦略の探索をサポート：複数の戦略が正しい答えを生成する可能性を許容し、単一の「最適」パスを強制しない

核心的貢献

MBA-RAGフレームワークの提案：多腕バンディットアルゴリズムをRAGシステムの検索戦略選択に初めて適用し、動的適応的検索を実現
動的報酬関数の設計：精度と計算効率を革新的に組み合わせ、高コスト方法にペナルティを与えることでリソース使用を最適化
SOTA性能の実現：6つのデータセットで最良の結果を達成し、同時に検索コストを20%削減
柔軟な監督メカニズムの提供：厳密な単一ラベル監督の代わりに部分情報監督を使用し、モデルが複数の有効な戦略を探索することを許容

方法の詳細

タスク定義

クエリxが与えられた場合、RAGシステムは以下を必要とする：

検索段階：モジュールRがクエリxの関連文書Dを検索する
生成段階：LLMがxとDを使用して応答ā = LLM(yt|x,D)を生成する

本論文はこれを多腕バンディット問題として再定義し、各検索方法(検索なし、単一検索、複数検索)を「腕」として扱う。

モデルアーキテクチャ

1. クエリエンコーディングと腕選択

エンコーダ：DistilBERTを使用してユーザークエリをエンコードし、アクション分布z = fθ(x)を生成する
選択戦略：ε-貪欲戦略を採用して探索と利用のバランスを取る：
- 確率(1-ε)でa = argmax(z)を選択
- 確率εで生成方法をランダムに選択

2. 学習アルゴリズム

目的関数は実際の報酬raと予測報酬fθ(x)aの間の二乗誤差を最小化することである：

min_θ (ra - fθ(x)a)²

パラメータ更新ルール：

θt+1 = θt - α∇θ((ra - fθ(x)a)²)

3. 動的報酬関数

ra = A(y, ŷa) - λC(a)

ここで：

A(y, ŷa)：生成品質指標(例：完全一致)
C(a)：方法aの計算コスト(例：検索ステップ数)
λ：精度と効率のバランスを取るスケーリング係数

技術的革新点

多腕バンディット適応：検索戦略選択を多腕バンディット問題としてモデル化し、各検索方法を「腕」に対応させる
部分情報監督：選択された戦略に対してのみフィードバックを提供し、選択されなかった戦略にペナルティを与えない
コスト認識報酬：動的報酬関数は精度と計算効率の両方を考慮する
探索-利用バランス：ε-貪欲戦略を通じて準最適解への早期収束を回避

実験設定

データセット

単一ホップQAデータセット：

SQuAD v1.1：読解理解タスク
Natural Questions：オープンドメイン質問応答
TriviaQA：知識質問応答

多ホップQAデータセット：

MuSiQue：複数ステップ推論質問応答
HotpotQA：多ホップ推論質問応答
2WikiMultiHopQA：ウィキペディアベースの多ホップ質問応答

評価指標

性能指標：

EM (完全一致)：予測結果と真の答えが完全に一致
F1：予測答えと真の答えの語彙重複度
Acc (精度)：予測答えが真の答えを含むかどうか

効率指標：

Step：選択された検索戦略に必要な検索ステップ数

比較手法

No-Retrieval：検索を行わずに直接答えを生成
Adaptive-Retrieval：検索が必要かどうかを動的に判定
Self-RAG：自己反省を通じて検索ニーズを動的に決定
DRAGIN：トークン不確実性に基づいて検索を活性化
SEAKR：自己認識不確実性に基づいて検索を決定
Adaptive-RAG：分類器を使用してクエリ複雑度に基づいて検索戦略を選択

実装詳細

クエリエンコーディングモデル：DistilBERT
検索モデル：BM25
生成モデル：FLAN-T5-XL (3B)
学習率：5e-5
探索戦略：ε-貪欲アルゴリズム

実験結果

主要結果

手法	EM	F1	Acc	Step
No Retrieval	14.87	21.12	15.97	0.00
Adaptive Retrieval	23.87	32.24	26.73	0.50
Self-RAG	9.90	20.79	31.57	0.72
Adaptive-RAG	37.17	46.94	42.10	2.17
MBA-RAG (提案手法)	38.80	48.61	43.57	1.80

主要な発見

性能向上：MBA-RAGはすべての性能指標においてベースライン手法を上回る
効率最適化：Adaptive-RAGと比較して、検索ステップ数が約17%削減される(2.17から1.80へ)
単一ホップデータセットの性能：SQuADおよびTriviaQAで顕著な改善を達成し、検索コストが大幅に削減される
多ホップデータセットの性能：2WikiMultiHopQAで優れた改善を達成し、検索コストが20%以上削減される

分類精度分析

MBA-RAGの分類精度は56.1%に達し、以下を大幅に上回る：

Adaptive Retrieval：42.0%
Self-RAG：41.5%
Adaptive-RAG：54.0%

アブレーション実験

マルチラベル分類器の結果との比較により、従来のマルチラベル手法は良好な性能を示すものの、検索コストが過度に高い(Stepが4.514に達する)ことが示され、MBA-RAGは性能と効率の最適なバランスを実現している。

結論と考察

主要な結論

有効性の検証：MBA-RAGは複数のデータセットでSOTA性能を達成
効率向上：検索コストを大幅に削減し、平均20%削減
強い適応性：クエリ複雑度に基づいて戦略を動的に調整できる

限界

アルゴリズム依存性：フレームワークは特定の多腕バンディットアルゴリズム構造に依存
スケーラビリティの課題：未知のクエリタイプに直面する際に適応性の問題が生じる可能性がある
計算要件：強化学習手法は追加の計算オーバーヘッドを導入する可能性がある

今後の方向性

アルゴリズム最適化：計算要件を削減するためのより効率的なアルゴリズムを探索
汎化能力：新しいクエリタイプへの適応能力を向上させる
応用拡張：より広範なNLPタスクへの手法の応用

深層的評価

利点

革新性が高い：多腕バンディットをRAGシステムに初めて導入し、理論的基礎が堅実
実用価値が高い：精度と効率を同時に最適化し、重要な応用価値を持つ
実験が充分：6つの異なるタイプのデータセットで包括的な評価を実施
手法が合理的：動的報酬関数の設計が巧妙で、複数の目標のバランスを取る

不足点

複雑性の増加：単純な分類手法と比較して、追加のアルゴリズム複雑性を導入
パラメータ感度：報酬関数のバランスパラメータλは異なるデータセットに対して調整が必要
理論分析の不足：収束性と最適性の理論的保証が欠けている

影響力

学術的貢献：RAGシステム最適化に新しい研究方向を提供
実用的応用：手法は実用性が高く、実際のシステムに応用可能
再現性：完全なコード実装を提供し、再現と拡張が容易

適用シーン

知識集約的な質問応答：特に精度と効率のバランスが必要なシーンに適している
複数複雑度クエリの処理：単純から複雑までの様々なクエリを処理できる
リソース制約環境：計算リソースが限定される場合に検索コストを最適化できる

参考文献

Lewis, P., et al. (2020). Retrieval-augmented generation for knowledge-intensive nlp tasks. NeurIPS.
Jeong, S., et al. (2024). Adaptive-rag: Learning to adapt retrieval-augmented large language models through question complexity. arXiv preprint.
Katehakis, M. N., & Veinott Jr, A. F. (1987). The multi-armed bandit problem: decomposition and computation. Mathematics of Operations Research.

総合評価：本論文は革新的で実用的なRAG最適化フレームワークを提案し、多腕バンディットアルゴリズムを通じて検索戦略の動的選択を実現している。高い精度を維持しながら計算コストを大幅に削減している。手法は理論的基礎が堅実で、実験結果は説得力があり、RAGシステムのさらなる発展に価値のある知見を提供している。