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 কাঠামো হয় নির্বিচারে পুনরুদ্ধার সম্পাদন করে অথবা পুনরুদ্ধার পদ্ধতি নির্বাচনের জন্য কঠোর একক-শ্রেণী শ্রেণীবিভাজকের উপর নির্ভর করে, যার ফলে বিভিন্ন জটিলতার প্রশ্নগুলিতে অদক্ষতা এবং সর্বোত্তম নয় এমন কর্মক্ষমতা দেখা যায়। এই চ্যালেঞ্জগুলি সমাধান করার জন্য, এই পত্রটি একটি শক্তিশালী শিক্ষার ভিত্তিতে একটি কাঠামো প্রস্তাব করে যা প্রশ্নের জটিলতার উপর ভিত্তি করে সবচেয়ে উপযুক্ত পুনরুদ্ধার কৌশল গতিশীলভাবে নির্বাচন করতে পারে। এই পদ্ধতিটি মাল্টি-আর্মড ব্যান্ডিট অ্যালগরিদম ব্যবহার করে, প্রতিটি পুনরুদ্ধার পদ্ধতিকে বিভিন্ন "বাহু" হিসাবে বিবেচনা করে, অন্বেষণ এবং শোষণের ভারসাম্য রেখে নির্বাচন প্রক্রিয়ার সাথে খাপ খাইয়ে নেয়। অতিরিক্তভাবে, নির্ভুলতা এবং দক্ষতার ভারসাম্য রাখে এমন একটি গতিশীল পুরস্কার ফাংশন প্রবর্তন করা হয়েছে, যা সঠিক ফলাফল পাওয়ার ক্ষেত্রেও আরও বেশি পুনরুদ্ধার পদক্ষেপের প্রয়োজন এমন পদ্ধতিগুলিকে শাস্তি দেয়। এই পদ্ধতিটি একাধিক একক-হপ এবং মাল্টি-হপ ডেটাসেটে নতুন SOTA ফলাফল অর্জন করেছে, একই সাথে পুনরুদ্ধার খরচ হ্রাস করেছে।
বিদ্যমান RAG সিস্টেমগুলি নিম্নলিখিত মূল সমস্যাগুলির সম্মুখীন হয়:
পুনরুদ্ধার কৌশল নির্বাচনে ত্রুটি: বেশিরভাগ RAG কাঠামো সমস্ত প্রশ্নের জন্য নির্বিচারে পুনরুদ্ধার সম্পাদন করে, যা অপ্রয়োজনীয় বা প্রাসঙ্গিক নয় এমন অনুচ্ছেদ প্রবর্তন করতে পারে
একক পদ্ধতির সীমাবদ্ধতা: সমস্ত প্রশ্নের জন্য একটি একক পুনরুদ্ধার পদ্ধতি ব্যবহার করা অদক্ষ, সহজ প্রশ্নগুলি অপ্রয়োজনীয় গণনামূলক ওভারহেড তৈরি করে, এবং জটিল প্রশ্নগুলি পর্যাপ্ত প্রক্রিয়াকরণ পায় না
তত্ত্বাবধান সংকেত অনির্ভুলতা: AdaptiveRAG এর মতো বিদ্যমান অভিযোজনযোগ্য পদ্ধতিগুলি অনুমানমূলক তত্ত্বাবধান ব্যবহার করে, প্রতিটি প্রশ্নের জন্য শুধুমাত্র একটি সর্বোত্তম কৌশল আছে বলে অনুমান করে, এবং সর্বনিম্ন পুনরুদ্ধার খরচের পথ নির্বাচনের দিকে ঝুঁকে থাকে
MBA-RAG কাঠামো প্রস্তাব: প্রথমবারের মতো মাল্টি-আর্মড ব্যান্ডিট অ্যালগরিদম RAG সিস্টেমের পুনরুদ্ধার কৌশল নির্বাচনে প্রয়োগ করে, গতিশীল অভিযোজনযোগ্য পুনরুদ্ধার বাস্তবায়ন করে
গতিশীল পুরস্কার ফাংশন ডিজাইন: সৃজনশীলভাবে নির্ভুলতা এবং গণনামূলক দক্ষতা একত্রিত করে, উচ্চ খরচ পদ্ধতিগুলিকে শাস্তি দিয়ে সম্পদ ব্যবহার অপ্টিমাইজ করে
SOTA কর্মক্ষমতা অর্জন: ৬টি ডেটাসেটে সর্বোত্তম ফলাফল অর্জন করে, একই সাথে পুনরুদ্ধার খরচ ২০% হ্রাস করে
নমনীয় তত্ত্বাবধান প্রক্রিয়া প্রদান করে: কঠোর একক-লেবেল তত্ত্বাবধানের পরিবর্তে আংশিক তথ্য তত্ত্বাবধান ব্যবহার করে, মডেলকে একাধিক কার্যকর কৌশল অন্বেষণ করতে অনুমতি দেয়
পুনরুদ্ধার পর্যায়: মডিউল R প্রশ্ন x এর জন্য প্রাসঙ্গিক নথি D পুনরুদ্ধার করে
প্রজন্ম পর্যায়: LLM x এবং D ব্যবহার করে প্রতিক্রিয়া ā = LLM(yt|x,D) তৈরি করে
এই পত্রটি এটিকে একটি মাল্টি-আর্মড ব্যান্ডিট সমস্যা হিসাবে পুনর্সংজ্ঞায়িত করে, যেখানে প্রতিটি পুনরুদ্ধার পদ্ধতি (কোন পুনরুদ্ধার নেই, একক পুনরুদ্ধার, একাধিক পুনরুদ্ধার) একটি "বাহু" হিসাবে কাজ করে।
মাল্টি-আর্মড ব্যান্ডিট অভিযোজন: পুনরুদ্ধার কৌশল নির্বাচনকে মাল্টি-আর্মড ব্যান্ডিট সমস্যা হিসাবে মডেল করে, প্রতিটি পুনরুদ্ধার পদ্ধতি একটি "বাহু" এর সাথে সামঞ্জস্যপূর্ণ
আংশিক তথ্য তত্ত্বাবধান: শুধুমাত্র নির্বাচিত কৌশলের জন্য প্রতিক্রিয়া প্রদান করে, অনির্বাচিত কৌশলগুলিকে শাস্তি দেয় না
খরচ-সচেতন পুরস্কার: গতিশীল পুরস্কার ফাংশন নির্ভুলতা এবং গণনামূলক দক্ষতা উভয়ই বিবেচনা করে
অন্বেষণ-শোষণ ভারসাম্য: ε-লোভী কৌশলের মাধ্যমে সর্বোত্তম নয় এমন সমাধানে প্রাথমিক সংমিশ্রণ এড়ায়
মাল্টি-লেবেল শ্রেণীবিভাজকের ফলাফলের সাথে তুলনা দেখায় যে, ঐতিহ্যবাহী মাল্টি-লেবেল পদ্ধতি যদিও ভাল কর্মক্ষমতা প্রদান করে, তবে পুনরুদ্ধার খরচ অত্যধিক (Step ৪.৫১৪ এ পৌঁছায়), যখন MBA-RAG কর্মক্ষমতা এবং দক্ষতার সর্বোত্তম ভারসাম্য অর্জন করে।
এই পত্রটি প্রথমবারের মতো মাল্টি-আর্মড ব্যান্ডিট অ্যালগরিদম 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 সিস্টেমের আরও উন্নয়নের জন্য মূল্যবান অন্তর্দৃষ্টি প্রদান করে।