MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity
This addresses the problem of suboptimal performance and inefficiency in retrieval-augmented generation for knowledge-intensive tasks, representing a novel method for a known bottleneck.
The paper tackles the inefficiency of existing RAG frameworks by proposing a reinforcement learning-based approach that dynamically selects retrieval strategies based on query complexity, achieving new state-of-the-art results on multiple datasets while reducing retrieval costs.
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 .