# MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity

arXiv: https://arxiv.org/abs/2412.01572 · HTML: https://mlsanity.com/paper/2412.01572/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Placed in Agentic RAG, agents and memory, one of 15 directions in Retrieval-augmented generation.
- Builds on 13 of its 31 references that are in the map.
- First posted on arXiv 2 Dec 2024.

## Placed in

- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)

## Builds on (most cited first, 13 of 13)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions](https://mlsanity.com/paper/2212.10509/) (2022)
- [Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity](https://mlsanity.com/paper/2403.14403/) (2024)
- [CO2: Efficient Distributed Training with Full Communication-Computation Overlap](https://mlsanity.com/paper/2401.16265/) (2024)
- [Spatial-Temporal Interplay in Human Mobility: A Hierarchical Reinforcement Learning Approach with Hypergraph Representation](https://mlsanity.com/paper/2312.15717/) (2023)

## Built on by (most cited first, 29 of 29)

- [A Survey of Graph Retrieval-Augmented Generation for Customized Large Language Models](https://mlsanity.com/paper/2501.13958/) (2025)
- [When to use Graphs in RAG: A Comprehensive Analysis for Graph Retrieval-Augmented Generation](https://mlsanity.com/paper/2506.05690/) (2025)
- [KABB: Knowledge-Aware Bayesian Bandits for Dynamic Expert Coordination in Multi-Agent Systems](https://mlsanity.com/paper/2502.07350/) (2025)
- [UniversalRAG: Retrieval-Augmented Generation over Corpora of Diverse Modalities and Granularities](https://mlsanity.com/paper/2504.20734/) (2025)
- [METIS: Fast Quality-Aware RAG Systems with Configuration Adaptation](https://mlsanity.com/paper/2412.10543/) (2024)
- [LAG: Logic-Augmented Generation from a Cartesian Perspective](https://mlsanity.com/paper/2508.05509/) (2025)
- [Fast or Better? Balancing Accuracy and Cost in Retrieval-Augmented Generation with Flexible User Control](https://mlsanity.com/paper/2502.12145/) (2025)
- [RAGRouter-Bench: A Dataset and Benchmark for Adaptive RAG Routing](https://mlsanity.com/paper/2602.00296/) (2026)
- [Bayesian Active Learning with Gaussian Processes Guided by LLM Relevance Scoring for Dense Passage Retrieval](https://mlsanity.com/paper/2604.17906/) (2026)
- [BalanceRAG: Joint Risk Calibration for Cascaded Retrieval-Augmented Generation](https://mlsanity.com/paper/2605.20084/) (2026)
- [OkraLong: A Flexible Retrieval-Augmented Framework for Long-Text Query Processing](https://mlsanity.com/paper/2503.02603/) (2025)
- [KAG-Thinker: Interactive Thinking and Deep Reasoning in LLMs via Knowledge-Augmented Generation](https://mlsanity.com/paper/2506.17728/) (2025)
- [Confident RAG: Enhancing the Performance of LLMs for Mathematics Question Answering through Multi-Embedding and Confidence Scoring](https://mlsanity.com/paper/2507.17442/) (2025)
- [Online-Optimized RAG for Tool Use and Function Calling](https://mlsanity.com/paper/2509.20415/) (2025)
- [RAGTrack: Language-aware RGBT Tracking with Retrieval-Augmented Generation](https://mlsanity.com/paper/2603.03617/) (2026)
- [Budget-Constrained Online Retrieval-Augmented Generation: The Chunk-as-a-Service Model](https://mlsanity.com/paper/2604.26981/) (2026)
- [Maximizing Rollout Informativeness under a Fixed Budget: A Submodular View of Tree Search for Tool-Use Agentic Reinforcement Learning](https://mlsanity.com/paper/2605.05262/) (2026)
- [POLYRAG: Integrating Polyviews into Retrieval-Augmented Generation for Medical Applications](https://mlsanity.com/paper/2504.14917/) (2025)
- [Feedback-Induced Performance Decline in LLM-Based Decision-Making](https://mlsanity.com/paper/2507.14906/) (2025)
- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)
- [Diagnosing and Repairing Factual Errors in RAG under Budget Constraints](https://mlsanity.com/paper/2606.29377/) (2026)
- [A Dynamic Retrieval-Augmented Generation System with Selective Memory and Remembrance](https://mlsanity.com/paper/2601.02428/) (2026)
- [Deep GraphRAG: A Balanced Approach to Hierarchical Retrieval and Adaptive Integration](https://mlsanity.com/paper/2601.11144/) (2026)
- [Is a Picture Worth a Thousand Words? Adaptive Multimodal Fact-Checking with Visual Evidence Necessity](https://mlsanity.com/paper/2604.04692/) (2026)
- [MAB-DQA: Addressing Query Aspect Importance in Document Question Answering with Multi-Armed Bandits](https://mlsanity.com/paper/2604.08952/) (2026)
- [ADEPT: An Entropy-Driven Dual-Strategy Agent for Interactive Video Retrieval](https://mlsanity.com/paper/2606.28326/) (2026)
- [LODESTAR: Robust Entropy-Based Answer Selection in Retrieval-Augmented Generation for Question Answering -- Directing Frozen-LLM Entropy with a Reinforcement-Learned Prompt Polarizer under Misleading Passages](https://mlsanity.com/paper/2608.11922/) (2026)
- [Beyond the Query: Do Retrieval Signals Improve Adaptive Multimodal RAG Routing?](https://mlsanity.com/paper/2609.12437/) (2026)
- [Fusion is the New Mutation: Bandit-Guided Evolution on Workflow Graphs](https://mlsanity.com/paper/2610.05284/) (2026)

## Abstract (arXiv)

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 .