MemSearcher: Training LLMs to Reason, Search and Manage Memory via End-to-End Reinforcement Learning

तकनीकी सारांश: MemSearcher

समस्या विवरण

लार्ज लैंग्वेज मॉडल (LLM)-आधारित खोज एजेंट, विशेष रूप से जो ReAct प्रतिमान (paradigm) का पालन करते हैं, आमतौर पर प्रत्येक टर्न के लिए संपूर्ण इंटरैक्शन इतिहास (विचार, क्रियाएं और अवलोकन) को अपने कॉन्टेक्स्ट विंडो में जोड़ देते हैं। यह डिज़ाइन एक मौलिक बाधा उत्पन्न करता है: जैसे-जैसे इंटरैक्शन टर्न की संख्या बढ़ती है, कॉन्टेक्स्ट की लंबाई रैखिक रूप से बढ़ती जाती है। खोज परिदृश्यों में, इसके परिणामस्वरूप लंबे, शोर वाले इनपुट होते हैं जो रिट्रीव किए गए पैसेजों (passages) के प्रभुत्व में होते हैं, जिनमें से कई उपयोगकर्ता के प्रश्न के लिए अप्रासंगिक होते हैं। टोकनों का यह प्रसार न केवल महत्वपूर्ण जानकारी का उपयोग करने की LLM की क्षमता पर दबाव डालता है, बल्कि पर्याप्त कम्प्यूटेशनल और GPU मेमोरी ओवरहेड भी उत्पन्न करता है, जिससे संसाधन-बाधित सेटिंग्स में खोज एजेंटों की स्केलेबिलिटी और परिनियोजन (deployment) सीमित हो जाता है।

कार्यप्रणाली (Methodology)

लेखक MemSearcher का प्रस्ताव करते हैं, जो एक एजेंट फ्रेमवर्क है जिसे मल्टी-टर्न इंटरैक्शन के दौरान पूर्ण इतिहास को जोड़ने के बजाय, एक संक्षिप्त, पुनरावृत्ति रूप से अपडेट होने वाली मेमोरी बनाए रखने के लिए डिज़ाइन किया गया है।

1. MemSearcher फ्रेमवर्क

ReAct के विपरीत, जहाँ टर्न $i$ पर इनपुट संचयी इतिहास $(q, t_1, a_1, o_1, \dots, t_{i-1}, a_{i-1}, o_{i-1})$ होता है, MemSearcher बैकबोन LLM को केवल दो संक्षिप्त इनपुट प्रदान करता है:

1. उपयोगकर्ता का प्रश्न $q$।
2. एक संक्षिप्त मेमोरी $m_{i-1}$ जिसमें पिछले टर्न की आवश्यक जानकारी शामिल है।

प्रत्येक टर्न पर, LLM इस संक्षिप्त संदर्भ के आधार पर एक रीजनिंग ट्रेस (विचार) और एक क्रिया उत्पन्न करता है। वातावरण (environment) से अवलोकन (जैसे, खोज परिणाम) प्राप्त करने पर, LLM एक मेमोरी मैनेजर के रूप में कार्य करता है। यह वर्तमान अवलोकन से नई प्रासंगिक जानकारी को मौजूदा मेमोरी में एकीकृत करता है और पिछले मेमोरी से आवश्यक विवरणों को सुरक्षित रखते हुए, एक अपडेटेड मेमोरी $m_i$ बनाता है। यह प्रक्रिया सुनिश्चित करती है कि कॉन्टेक्स्ट की लंबाई स्थिर रहे और टर्न की संख्या के बावजूद एक पूर्व-निर्धारित अधिकतम टोकन सीमा (जैसे, 1,024 टोकन) के भीतर सीमित रहे।

2. Multi-Context GRPO

MemSearcher को प्रशिक्षित करना चुनौतीपूर्ण है क्योंकि एक एकल ट्राजेक्टरी (trajectory) कई टर्न तक फैली होती है, जिनमें से प्रत्येक अलग-अलग संदर्भ (बदलती मेमोरी स्थिति के कारण) के तहत होता है। मानक ग्रुप रिलेटिव पॉलिसी ऑप्टिमाइज़ेशन (GRPO) एक ट्राजेक्टरी को एक एकल इकाई के रूप में मानता है, जो अपर्याप्त है जब प्रत्येक टर्न एक स्वतंत्र ऑप्टिमाइज़ेशन लक्ष्य का प्रतिनिधित्व करता है।

इसे संबोधित करने के लिए, लेखक Multi-Context GRPO पेश करते हैं:

• ट्राजेक्टरी सैंपलिंग (Trajectory Sampling): दिए गए प्रश्न के लिए, ट्राजेक्टरीज का एक समूह सैंपल किया जाता है। प्रत्येक ट्राजेक्टरी में विविध संदर्भों के तहत कई टर्न शामिल होते हैं।
• एडवांटेज प्रोपेगेशन (Advantage Propagation): रिवॉर्ड्स की गणना ट्राजेक्टरी स्तर पर की जाती है। परिणामी एडवांटेज को उस ट्राजेक्टरी के भीतर प्रत्येक टर्न में समान रूप से प्रसारित किया जाता है।
• स्वतंत्र अनुकूलन (Independent Optimization):ism: प्रत्येक टर्न को एक स्वतंत्र ऑप्टिमाइज़ेशन लक्ष्य के रूप में माना जाता है। पॉलिसी को सभी सैंपल की गई ट्राजेक्टरीज के सभी टर्न पर ऑब्जेक्टिव को मैक्सिमाइज़ करके अपडेट किया जाता है, जिससे एंड-टू-एंड सुदृढीकरण शिक्षण (RL) प्रशिक्षण सक्षम होता है।

3. रिवॉर्ड मॉडलिंग (Reward Modeling)

प्रशिक्षण में दो घटकों वाला एक रिवॉर्ड फंक्शन उपयोग किया जाता है:

• फॉर्मेट रिवॉर्ड (Format Reward): यह सुनिश्चित करता है कि मॉडल पूर्व-निर्धारित आउटपुट फॉर्मेट (सही टैग उपयोग और \boxed{} की उपस्थिति) का पालन करता है।
• आंसर रिवॉर्ड (Answer Reward): एक नियम-आधारित मीट्रिक जो \boxed{} के अंदर अंतिम उत्तर और ग्राउंड ट्रुथ के बीच F1 स्कोर की गणना करता है।

मुख्य योगदान

1. MemSearcher फ्रेमवर्क: एक एजेंट आर्किटेक्चर जो एक संक्षिप्त मेमोरी को पुनरावृत्ति रूप से बनाए रखने के लिए बैकबोन LLM का उपयोग एक आंतरिक मेमोरी मैनेजर के रूप में करता है, जिससे पूरे इंटरैक्शन इतिहास को कॉन्टेक्स्ट में जोड़ने की आवश्यकता समाप्त हो जाती है।
2. Multi-Context GRPO: एक नवीन RL प्रशिक्षण एल्गोरिदम जो ट्राजेक्टरी-स्तरीय एडवांटेज को व्यक्तिगत टर्न तक प्रसारित करके MemSearcher ट्राजेक्टरीज के एंड-टू-एंड ऑप्टिमाइज़ेशन को सक्षम बनाता है, जो मल्टी-टर्न इंटरैक्शन के बदलते संदर्भों के अनुकूल होता है।
3. दक्षता और प्रदर्शन: व्यापक प्रयोग जो प्रदर्शित करते हैं कि MemSearcher, ReAct-आधारित बेसलाइन्स की तुलना में बेहतर प्रदर्शन करता है, जबकि मल्टी-टर्न इंटरैक्शन के दौरान लगभग स्थिर कॉन्टेक्स्ट लंबाई (4K टोकन से कम) बनाए रखता है।

प्रयोगात्मक परिणाम

लेखकों ने सात सार्वजनिक ज्ञान-अर्जन डेटासेट्स (Natural Questions, TriviaQA, PopQA, HotpotQA, 2WikiMultiHopQA, Musique, और Bamboogle) पर विभिन्न आकारों (3B, 7B, 14B) के Qwen2.5 मॉडल्स का उपयोग करके MemSearcher का मूल्यांकन किया।

• प्रदर्शन: MemSearcher ने Search-R1, ReSearch, और AutoRefine सहित मजबूत बेसलाइन्स को लगातार पछाड़ दिया। विशेष रूप से, MemSearcher 3B ने 43.8 का औसत स्कोर प्राप्त किया, जो कई 7B पैरामीटर बेसलाइन्स से अधिक था। MemSearcher 7B (48.9 औसत) ने ReSearch 32B को पछाड़ दिया।
• दक्षता: जबकि ReAct-आधारित एजेंटों ने इंटरैक्शन टर्न के साथ टोकन खपत में रैखिक वृद्धि दिखाई, MemSearcher ने काफी कम और स्थिर टोकन संख्या बनाए रखी। यह 4K टोकन से कम के कॉन्टेक्स्ट विंडो के भीतर मल्टी-टर्न इंटरैक्शन की अनुमति देता है।
• प्रशिक्षण की आवश्यकता: एब्लेशन अध्ययनों ने पुष्टि की कि RL प्रशिक्षण महत्वपूर्ण है। वे मॉडल जो RL प्रशिक्षण के बिना MemSearcher फ्रेमवर्क का उपयोग करते हैं, वे Multi-Context GRPO के साथ प्रशिक्षित मॉडलों की तुलना में महत्वपूर्ण प्रदर्शन गिरावट दिखाते हैं। इसके अलावा, RL प्रशिक्षण ने डिस्टिलेशन (distillation) के माध्यम से SFT (Supervised Fine-Tuning) को पछाड़ दिया, क्योंकि SFT इस विशिष्ट फ्रेमवर्क के लिए महंगी मध्यवर्ती एनोटेशन और उप-इष्टतम शिक्षक मॉडल्स पर निर्भर करता है।

महत्व और दावे

पेपर का दावा है कि MemSearcher वर्तमान खोज एजेंटों की स्केलेबिलिटी बाधा को कॉन्टेक्स्ट लंबाई को इंटरैक्शन टर्न की संख्या से अलग करके हल करता है। एक संक्षिप्त मेमोरी बनाए रखकर और प्रशिक्षण के लिए Multi-Context GRPO का उपयोग करके, यह फ्रेमवर्क संसाधन-बाधित वातावरण में प्रभावी खोज एजेंटों के परिनियोजन को सक्षम बनाता है। लेखक इस बात पर जोर देते हैं कि उनका दृष्टिकोण छोटे मॉडल्स (जैसे, 3B या 7B) को पारंपरिक ReAct प्रतिमानों पर निर्भर बड़े मॉडल्स से बेहतर प्रदर्शन करने की अनुमति देता है, जो यह सुझाव देता है कि जटिल रीजनिंग कार्यों के लिए कुशल कॉन्टेक्स्ट प्रबंधन, मॉडल स्केल जितना ही महत्वपूर्ण है।

लेखक सीमाओं को स्वीकार करते हुए कहते हैं कि हालांकि उनका मेमोरी डिज़ाइन प्रभावी है, फिर भी अधिक परिष्कृत तंत्र (जैसे, RAG-शैली या स्ट्रक्चर्ड मेमोरी) में RL-आधारित एजेंटों के लिए आगे की जांच की आवश्यकता है। वे यह भी सुझाव देते हैं कि हालांकि Multi-Context GRPO एक प्रभावी विस्तार है, संभावित लंबाई पूर्वाग्रहों (length biases) को कम करने और RL प्रशिक्षण की क्षमता का पूरी तरह से लाभ उठाने के लिए और अधिक तकनीकों की आवश्यकता है।