MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity
Tang, Gao, Li et al.
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 .
academic
MBA-RAG: प्रश्न जटिलता के माध्यम से अनुकूली पुनर्प्राप्ति-संवर्धित पीढ़ी के लिए एक बैंडिट दृष्टिकोण
पुनर्प्राप्ति-संवर्धित पीढ़ी (RAG) ज्ञान-गहन कार्यों में भाषा मॉडल के जनरेशन प्रदर्शन को महत्वपूर्ण रूप से बेहतर बनाता है। हालांकि, मौजूदा RAG ढांचे या तो भेदभाव के बिना पुनर्प्राप्ति करते हैं या पुनर्प्राप्ति विधि चुनने के लिए कठोर एकल-वर्ग वर्गीकारक पर निर्भर करते हैं, जिससे विभिन्न जटिलता वाली क्वेरीज़ में अक्षमता और उप-इष्टतम प्रदर्शन होता है। इन चुनौतियों को संबोधित करने के लिए, यह पेपर एक सुदृढ़ीकरण सीखने-आधारित ढांचा प्रस्तावित करता है जो क्वेरी जटिलता के आधार पर सबसे उपयुक्त पुनर्प्राप्ति रणनीति को गतिशील रूप से चुन सकता है। यह विधि बहु-सशस्त्र दस्यु (Multi-Armed Bandit) एल्गोरिदम का उपयोग करती है, प्रत्येक पुनर्प्राप्ति विधि को विभिन्न "भुजाओं" के रूप में मानती है, और अन्वेषण और शोषण को संतुलित करके चयन प्रक्रिया को अनुकूलित करती है। इसके अतिरिक्त, एक गतिशील पुरस्कार फ़ंक्शन पेश किया जाता है जो सटीकता और दक्षता को संतुलित करता है, यहां तक कि सही परिणाम प्राप्त करने की स्थिति में भी अधिक पुनर्प्राप्ति चरणों की आवश्यकता वाली विधियों को दंडित करता है। यह विधि कई एकल-हॉप और बहु-हॉप डेटासेट पर नए SOTA परिणाम प्राप्त करती है, साथ ही पुनर्प्राप्ति लागत को कम करती है।
मौजूदा RAG सिस्टम निम्नलिखित मुख्य समस्याओं का सामना करते हैं:
पुनर्प्राप्ति रणनीति का अनुचित चयन: अधिकांश RAG ढांचे सभी क्वेरीज़ के लिए भेदभाव के बिना पुनर्प्राप्ति करते हैं, जो अनावश्यक या विषय से हटे हुए पैराग्राफ़ को शामिल कर सकते हैं
एकल विधि की सीमा: सभी क्वेरीज़ के लिए एकल पुनर्प्राप्ति विधि का उपयोग अक्षम है, सरल क्वेरीज़ अनावश्यक कम्प्यूटेशनल ओवरहेड उत्पन्न करती हैं, जटिल क्वेरीज़ को पर्याप्त रूप से संभाला नहीं जा सकता है
अनुचित पर्यवेक्षण संकेत: AdaptiveRAG जैसी मौजूदा अनुकूली विधियां अनुमानी पर्यवेक्षण का उपयोग करती हैं, यह मानते हुए कि प्रत्येक क्वेरी के लिए केवल एक इष्टतम रणनीति है, और सबसे कम पुनर्प्राप्ति लागत वाले पथ को चुनने की प्रवृत्ति रखती हैं
MBA-RAG ढांचा प्रस्तावित करता है: RAG सिस्टम की पुनर्प्राप्ति रणनीति चयन के लिए बहु-सशस्त्र दस्यु एल्गोरिदम को पहली बार लागू करता है, गतिशील अनुकूली पुनर्प्राप्ति को लागू करता है
गतिशील पुरस्कार फ़ंक्शन डिज़ाइन करता है: सटीकता और कम्प्यूटेशनल दक्षता को नवीन तरीके से संयोजित करता है, उच्च लागत वाली विधियों को दंडित करके संसाधन उपयोग को अनुकूलित करता है
SOTA प्रदर्शन प्राप्त करता है: 6 डेटासेट पर सर्वश्रेष्ठ परिणाम प्राप्त करता है, साथ ही पुनर्प्राप्ति लागत को 20% कम करता है
लचीली पर्यवेक्षण तंत्र प्रदान करता है: कठोर एकल-लेबल पर्यवेक्षण के बजाय आंशिक जानकारी पर्यवेक्षण का उपयोग करता है, मॉडल को कई प्रभावी रणनीतियों का अन्वेषण करने की अनुमति देता है
क्वेरी x दिए गए, RAG सिस्टम को निम्नलिखित की आवश्यकता है:
पुनर्प्राप्ति चरण: मॉड्यूल R क्वेरी x के लिए संबंधित दस्तावेज़ D को पुनः प्राप्त करता है
जनरेशन चरण: LLM x और D का उपयोग करके प्रतिक्रिया ā = LLM(yt|x,D) उत्पन्न करता है
यह पेपर इसे बहु-सशस्त्र दस्यु समस्या के रूप में पुनः परिभाषित करता है, जहां प्रत्येक पुनर्प्राप्ति विधि (कोई पुनर्प्राप्ति नहीं, एकल पुनर्प्राप्ति, बहु-पुनर्प्राप्ति) एक "भुजा" के रूप में कार्य करती है।
बहु-सशस्त्र दस्यु अनुकूलन: पुनर्प्राप्ति रणनीति चयन को बहु-सशस्त्र दस्यु समस्या के रूप में मॉडल करता है, प्रत्येक पुनर्प्राप्ति विधि एक "भुजा" के अनुरूप है
आंशिक जानकारी पर्यवेक्षण: केवल चयनित रणनीति के लिए प्रतिक्रिया प्रदान करता है, अचयनित रणनीतियों को दंडित नहीं करता है
लागत-जागरूक पुरस्कार: गतिशील पुरस्कार फ़ंक्शन सटीकता और कम्प्यूटेशनल दक्षता दोनों पर विचार करता है
अन्वेषण-शोषण संतुलन: ε-लालची रणनीति के माध्यम से उप-इष्टतम समाधान में जल्दी अभिसरण से बचता है
बहु-लेबल वर्गीकारक के परिणामों की तुलना से पता चलता है कि पारंपरिक बहु-लेबल विधि अच्छा प्रदर्शन करती है, लेकिन पुनर्प्राप्ति लागत बहुत अधिक है (Step 4.514 तक पहुंचता है), जबकि MBA-RAG प्रदर्शन और दक्षता का सर्वोत्तम संतुलन प्राप्त करता है।
Lewis, P., et al. (2020). Retrieval-augmented generation for knowledge-intensive nlp tasks. NeurIPS.
Jeong, S., et al. (2024). Adaptive-rag: Learning to adapt retrieval-augmented large language models through question complexity. arXiv preprint.
Katehakis, M. N., & Veinott Jr, A. F. (1987). The multi-armed bandit problem: decomposition and computation. Mathematics of Operations Research.
समग्र मूल्यांकन: यह पेपर एक नवीन और व्यावहारिक RAG अनुकूलन ढांचा प्रस्तावित करता है, जो बहु-सशस्त्र दस्यु एल्गोरिदम के माध्यम से पुनर्प्राप्ति रणनीति के गतिशील चयन को लागू करता है, उच्च सटीकता बनाए रखते हुए कम्प्यूटेशनल लागत को महत्वपूर्ण रूप से कम करता है। विधि में ठोस सैद्धांतिक आधार है, प्रयोगात्मक परिणाम प्रेरक हैं, और RAG सिस्टम के आगे विकास के लिए मूल्यवान अंतर्दृष्टि प्रदान करता है।