根据问题复杂度动态选择检索策略，提升生成效率与准确率
MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity
- 用多臂赌博机算法动态选择最优检索方法
- 在多跳和单跳数据集上达到新最好性能，降低检索开销
- 适合需要高效知识增强生成的场景，如问答系统
检索增强生成（RAG）在知识密集型任务中显著提升了语言模型的生成能力。然而，现有框架或对所有查询统一执行检索，或依赖固定的单分类器选择检索方式，导致在不同复杂度查询下效率低下、性能不佳。为此，我们提出一种基于强化学习的自适应框架，根据查询复杂度动态选择最合适的检索策略。该方法将每种检索方式视为一个“臂”，利用多臂赌博机算法，在探索与利用之间进行平衡。同时引入动态奖励函数，兼顾准确性与效率，对需更多检索步骤的方法施加惩罚，即使结果正确也予以扣分。实验表明，该方法在多个单跳与多跳数据集上达到新最优表现，同时显著降低检索成本。代码已开源：https://github.com/FUTUREEEEEE/MBA。
原文摘要 · Abstract (English)
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 .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。