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
2024-12-02
Abstract: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 <a class="link-external link-https" href="https://github.com/FUTUREEEEEE/MBA" rel="external noopener nofollow">this https URL</a> .
Artificial Intelligence
What problem does this paper attempt to address?