{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/mba-rag-a-bandit-approach-for-adaptive","title":"MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity","arxiv_id":"2412.01572","date":"2024-12-02","proceeding":null,"authors":["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-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 .","url_abs":"https://arxiv.org/abs/2412.01572v4","url_pdf":"https://arxiv.org/pdf/2412.01572v4.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"mba-rag-a-bandit-approach-for-adaptive","repo_url":"https://github.com/futureeeeee/mba","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"pytorch","reach":null}],"tasks":[{"task_slug":"language-modeling","task_name":"Language Modeling"},{"task_slug":"language-modelling","task_name":"Language Modelling"},{"task_slug":"rag","task_name":"RAG"},{"task_slug":"retrieval","task_name":"Retrieval"},{"task_slug":"retrieval-augmented-generation","task_name":"Retrieval-augmented Generation"}],"methods":[{"method_slug":"adam","method_name":"Adam"},{"method_slug":"attention","method_name":"Attention"},{"method_slug":"attention-dropout","method_name":"Attention Dropout"},{"method_slug":"bart","method_name":"BART"},{"method_slug":"bert","method_name":"BERT"},{"method_slug":"bpe","method_name":"BPE"},{"method_slug":"dense-connections","method_name":"Dense Connections"},{"method_slug":"dropout","method_name":"Dropout"},{"method_slug":"layer-normalization","method_name":"Layer Normalization"},{"method_slug":"linear-layer","method_name":"Linear Layer"},{"method_slug":"linear-warmup-with-linear-decay","method_name":"Linear Warmup With Linear Decay"},{"method_slug":"multi-head-attention","method_name":"Multi-Head Attention"},{"method_slug":"rag","method_name":"RAG"},{"method_slug":"residual-connection","method_name":"Residual Connection"},{"method_slug":"softmax","method_name":"Softmax"},{"method_slug":"weight-decay","method_name":"Weight Decay"},{"method_slug":"wordpiece","method_name":"WordPiece"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/2412.01572","atlas_url":"https://app.syntology.ai/?focus=2412.01572","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2412.01572"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-25T09:33:49+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/futureeeeee/mba","reach":null}],"summary":{"ran_draft_wrong":2,"ran_violates":1},"by_repo_kind":{"official":{"samples":3,"ran":3,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"19ee4683eda506c6","entry":"list_of_strings","repo":"futureeeeee/mba","repo_kind":"official","path":"MAB/train_mab_mo_multiple.py","file_url":"https://github.com/futureeeeee/mba/blob/HEAD/MAB/train_mab_mo_multiple.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"19ee4683eda506c6"}},{"code_sha256_prefix":"4c94fc4279f6ba53","entry":"running_mean","repo":"futureeeeee/mba","repo_kind":"official","path":"MAB/train_mab_mo_multiple.py","file_url":"https://github.com/futureeeeee/mba/blob/HEAD/MAB/train_mab_mo_multiple.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4c94fc4279f6ba53"}},{"code_sha256_prefix":"69e582d84850bb60","entry":"softmax","repo":"futureeeeee/mba","repo_kind":"official","path":"MAB/train_mab_mo_multiple.py","file_url":"https://github.com/futureeeeee/mba/blob/HEAD/MAB/train_mab_mo_multiple.py","link_basis":"first_harvest_node","language":"python","status":"ran_violates","verification_level":1,"contract_check":"VIOLATES","metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"69e582d84850bb60"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}