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


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-label 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. To address these challenges, we propose a reinforcement learning-based framework that dynamically selects the most suitable retrieval strategy based on query complexity. 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.
Anthology ID:
2025.coling-main.218
Volume:
Proceedings of the 31st International Conference on Computational Linguistics
Month:
January
Year:
2025
Address:
Abu Dhabi, UAE
Editors:
Owen Rambow, Leo Wanner, Marianna Apidianaki, Hend Al-Khalifa, Barbara Di Eugenio, Steven Schockaert
Venue:
COLING
SIG:
Publisher:
Association for Computational Linguistics
Note:
Pages:
3248–3254
Language:
URL:
https://aclanthology.org/2025.coling-main.218/
DOI:
Bibkey:
Cite (ACL):
Xiaqiang Tang, Qiang Gao, Jian Li, Nan Du, Qi Li, and Sihong Xie. 2025. MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity. In Proceedings of the 31st International Conference on Computational Linguistics, pages 3248–3254, Abu Dhabi, UAE. Association for Computational Linguistics.
Cite (Informal):
MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity (Tang et al., COLING 2025)
Copy Citation:
PDF:
https://aclanthology.org/2025.coling-main.218.pdf