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: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity

基本信息

  • 论文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系统存在以下核心问题：

  1. 检索策略选择不当：大多数RAG框架对所有查询无差别地执行检索，可能引入不必要或离题的段落
  2. 单一方法局限性：对所有查询使用单一检索方法效率低下，简单查询会产生不必要的计算开销，复杂查询可能得不到充分处理
  3. 监督信号不准确：现有自适应方法如AdaptiveRAG使用启发式监督，假设每个查询只有一个最优策略，且倾向于选择检索成本最少的路径

研究动机

本文的核心动机是开发一个能够：

  1. 动态适应查询复杂度：根据问题的复杂程度智能选择检索策略
  2. 平衡准确性与效率：在保证答案质量的同时最小化计算成本
  3. 支持多策略探索：允许多种策略都可能产生正确答案，而不是强制选择单一"最优"路径

核心贡献

  1. 提出MBA-RAG框架：首次将多臂老虎机算法应用于RAG系统的检索策略选择，实现动态自适应检索
  2. 设计动态奖励函数：创新性地将准确性和计算效率结合，通过惩罚高成本方法来优化资源使用
  3. 实现SOTA性能：在6个数据集上取得最佳结果，同时将检索成本降低20%
  4. 提供灵活监督机制：使用部分信息监督替代严格的单标签监督，允许模型探索多种有效策略

方法详解

任务定义

给定查询x，RAG系统需要：

  1. 检索阶段：模块R为查询x检索相关文档D
  2. 生成阶段：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的计算成本(如检索步数)
  • λ：平衡准确性和效率的缩放因子

技术创新点

  1. 多臂老虎机适配：将检索策略选择建模为多臂老虎机问题，每个检索方法对应一个"臂"
  2. 部分信息监督：只对选中的策略提供反馈，不惩罚未选择的策略
  3. 成本感知奖励：动态奖励函数同时考虑准确性和计算效率
  4. 探索-利用平衡：通过ε-贪心策略避免过早收敛到次优解

实验设置

数据集

单跳QA数据集：

  • SQuAD v1.1：阅读理解任务
  • Natural Questions：开放域问答
  • TriviaQA：知识问答

多跳QA数据集：

  • MuSiQue：多步推理问答
  • HotpotQA：多跳推理问答
  • 2WikiMultiHopQA：基于维基百科的多跳问答

评价指标

性能指标：

  • EM (Exact Match)：预测结果与真实答案完全匹配
  • F1：预测答案与真实答案的词汇重叠度
  • Acc (Accuracy)：预测答案是否包含真实答案

效率指标：

  • Step：所选检索策略需要的检索步数

对比方法

  1. No-Retrieval：直接生成答案不进行检索
  2. Adaptive-Retrieval：动态判断是否需要检索
  3. Self-RAG：通过自反思动态决定检索需求
  4. DRAGIN：基于token不确定性激活检索
  5. SEAKR：基于自感知不确定性决定检索
  6. Adaptive-RAG：使用分类器根据查询复杂度选择检索策略

实现细节

  • 查询编码模型：DistilBERT
  • 检索模型：BM25
  • 生成模型：FLAN-T5-XL (3B)
  • 学习率：5e-5
  • 探索策略：ε-贪心算法

实验结果

主要结果

方法EMF1AccStep
No Retrieval14.8721.1215.970.00
Adaptive Retrieval23.8732.2426.730.50
Self-RAG9.9020.7931.570.72
Adaptive-RAG37.1746.9442.102.17
MBA-RAG (Ours)38.8048.6143.571.80

关键发现

  1. 性能提升：MBA-RAG在所有性能指标上均超越基线方法
  2. 效率优化：相比Adaptive-RAG，检索步数减少约17% (从2.17降至1.80)
  3. 单跳数据集表现：在SQuAD和TriviaQA上取得显著改进，检索成本大幅降低
  4. 多跳数据集表现：在2WikiMultiHopQA上取得突出改进，检索成本降低超过20%

分类准确性分析

MBA-RAG的分类准确性达到56.1%，显著高于：

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

消融实验

对比多标签分类器的结果显示，传统多标签方法虽然性能较好，但检索成本过高(Step达到4.514)，而MBA-RAG实现了性能与效率的最佳平衡。

相关工作

RAG系统发展

  1. 传统RAG：Lewis et al. (2020)提出的检索-生成框架
  2. 自适应检索：SEAKR、FLARE等方法实现按需检索
  3. 复杂度感知：AdaptiveRAG根据查询复杂度选择策略

多臂老虎机应用

本文首次将多臂老虎机算法应用于RAG系统，为检索策略选择提供了新的理论框架。

结论与讨论

主要结论

  1. 有效性验证：MBA-RAG在多个数据集上取得SOTA性能
  2. 效率提升：显著降低检索成本，平均减少20%
  3. 适应性强：能够根据查询复杂度动态调整策略

局限性

  1. 算法依赖：框架依赖于特定的多臂老虎机算法结构
  2. 扩展性挑战：在面对新的未见查询类型时可能存在适应性问题
  3. 计算需求：强化学习方法可能引入额外的计算开销

未来方向

  1. 算法优化：探索更高效的算法以减少计算需求
  2. 泛化能力：提升对新查询类型的适应能力
  3. 应用扩展：将方法应用到更广泛的NLP任务

深度评价

优点

  1. 创新性强：首次将多臂老虎机引入RAG系统，理论基础扎实
  2. 实用价值高：同时优化准确性和效率，具有重要应用价值
  3. 实验充分：在6个不同类型数据集上进行全面评估
  4. 方法合理：动态奖励函数设计巧妙，平衡了多个目标

不足

  1. 复杂度增加：相比简单分类方法，引入了额外的算法复杂性
  2. 参数敏感：奖励函数中的平衡参数λ需要针对不同数据集调整
  3. 理论分析不足：缺乏收敛性和最优性的理论保证

影响力

  1. 学术贡献：为RAG系统优化提供了新的研究方向
  2. 实际应用：方法具有较强的实用性，可应用于实际系统
  3. 可复现性：提供了完整的代码实现，便于复现和扩展

适用场景

  1. 知识密集型问答：特别适合需要平衡准确性和效率的场景
  2. 多复杂度查询处理：能够处理从简单到复杂的各种查询
  3. 资源受限环境：在计算资源有限时能够优化检索成本

参考文献

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

总体评价：本文提出了一个创新且实用的RAG优化框架，通过多臂老虎机算法实现了检索策略的动态选择，在保持高准确性的同时显著降低了计算成本。方法理论基础扎实，实验结果令人信服，为RAG系统的进一步发展提供了有价值的思路。