Trust Region Masking for Long-Horizon LLM Reinforcement Learning

यहाँ शोध पत्र "Trust Region Masking for Long-Horizon LLM Reinforcement Learning" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

यह शोध पत्र लार्ज लैंग्वेज मॉडल्स (LLMs) के लिए सुदृढीकरण शिक्षण (Reinforcement Learning - RL) में एक महत्वपूर्ण सैद्धांतिक और व्यावहारिक अंतराल को संबोधित करता है, विशेष रूप से लॉन्ग-होराइजन कार्यों (जैसे जटिल तर्क, बहु-चरणीय समस्या समाधान) के लिए जहाँ अनुक्रम की लंबाई ($T$) हजारों टोकन तक पहुँच सकती है।

• मुख्य मुद्दा: आधुनिक LLM-RL पाइपलाइन अनिवार्य ऑफ-पॉलिसी मिसमैच ($\pi_{roll} \neq \pi_\theta$) से ग्रस्त हैं। यह विचलन तीन प्राथमिक स्रोतों से उत्पन्न होता है:
  1. बैकएंड विसंगतियाँ (Backend Discrepancies): अटेंशन कर्नेल, परिशुद्धता (FP8 बनाम BF16), और ऑपरेटर फ्यूजन के संबंध में इन्फरेंस इंजन (जैसे vLLM, SGLang) और ट्रेनिंग फ्रेमवर्क (जैसे Megatron-LM, PyTorch FSDP) के बीच अंतर।
  2. MoE रूटिंग विच्छिन्नता (MoE Routing Discontinuities): Mixture-of-Experts मॉडल्स में, मामूली संख्यात्मक उतार-चढ़ाव (numerical jitter) विशेषज्ञ चयन (expert selection) को बदल सकता है, जिससे टोकन संभावनाओं में भारी उछाल आता है।
  3. डिस्ट्रीब्यूटेड स्टेलेनेस (Distributed Staleness): एसिंक्रोनस आर्किटेक्चर में डेटा जनरेशन (actor) और ग्रेडिएंट अपडेट (learner) के बीच विलंबता (latency)।
• सैद्धांतिक विफलता: शास्त्रीय ट्रस्ट रीजन बाउंड्स (जैसे Kakade & Langford, 2002) अनुक्रम लंबाई के साथ $O(T^2)$ के रूप में स्केल करते हैं। लंबे अनुक्रमों (जैसे $T=4096$) के लिए, ये बाउंड्स रिक्त (vacuous) हो जाते हैं (उदाहरण के लिए, त्रुटि सीमा > 1000 जब अधिकतम रिवॉर्ड 1 हो), जिसका अर्थ है कि वे कोई गारंटी नहीं देते कि सरोगेट ऑब्जेक्टिव को अनुकूलित करने से वास्तव में वास्तविक ऑब्जेक्टिव में सुधार होगा।
• वर्तमान विधियों की सीमा: मानक PPO टोकन-स्तरीय क्लिपिंग (clipping) का उपयोग करके अपडेट को नियंत्रित करता है। हालाँकि, क्योंकि ऑटोरेग्रेसिव जनरेशन क्रमिक (sequential) होता है, एक अनुक्रम के प्रारंभ में होने वाला छोटा सा संभाव्यता परिवर्तन आगे चलकर संचयी प्रभाव डालता है। टोकन-स्तरीय क्लिपिंग अधिकतम टोकन-स्तरीय विचलन ($D_{tok,max}^{KL}$ या $D_{tok,max}^{TV}$) को नियंत्रित नहीं कर सकती, जो कि एरर बाउंड निर्धारित करने वाला महत्वपूर्ण कारक है।

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

लेखक दो-आयामी दृष्टिकोण प्रस्तावित करते हैं: कड़े सैद्धांतिक बाउंड्स प्राप्त करना और एक अनुक्रम-स्तरीय मास्किंग तंत्र लागू करना।

A. सैद्धांतिक विश्लेषण: कड़े बाउंड्स (Theoretical Analysis: Tighter Bounds)

लेखक अनुमान त्रुटि $|J(\pi_\theta) - J(\pi_{roll}) - L(\pi_\theta)|$ के लिए नए बाउंड्स का एक परिवार व्युत्पन्न करते हैं जो शास्त्रीय $O(T^2)$ बाउंड की तुलना में काफी कड़े हैं। वे KL-डाइवर्जेंस और टोटल वेरिएशन (TV) दोनों मार्गों का उपयोग करते हैं।

1. पिन्सकर-मार्जिनल बाउंड्स ($O(T^{3/2})$):
   • मार्जिनल KL डाइवर्जेंस पर पिन्सकर असमानता (Pinsker's inequality) लागू करता है।
   • यह कॉन्टेक्स्ट शिफ्ट के लिए सबलीनियर स्केलिंग ($\sqrt{t}$) प्राप्त करता है लेकिन यदि विचलन कुछ ही टोकन पर केंद्रित है तो यह ढीला (loose) हो सकता है।
2. मिक्स्ड बाउंड्स ($O(T)$):
   • अनुक्रम-स्तरीय डाइवर्जेंस ($D_{seq}^{KL}$ या $D_{seq}^{TV}$) का उपयोग करता है जो $t$ के साथ नहीं बढ़ता है।
   • यह तब अधिक सटीक होता है जब विचलन विरल (sparse) हो (अर्थात, $D_{seq} \ll T \cdot \delta$)।
3. कपलिंग बाउंड्स ($O(T)$ विद कैप्स):
   • शुद्ध TV-आधारित कपलिंग। यह रैखिक रूप से स्केल करता है लेकिन 1 पर सीमित (capped) है, जो उच्च विचलन होने पर $O(T^2)$ विस्फोट को रोकता है।
4. अनुकूली बाउंड्स (Adaptive Bounds - डेटा-निर्भर):
   • सबसे सटीक ज्ञात गारंटी। यह प्रत्येक स्थिति (position) के लिए अपेक्षित प्रति-स्थिति TV ($\bar{D}_t$) का उपयोग करके त्रुटि को विभाजित करता है न कि वर्स्ट-केस डाइवर्जेंस का।
   • यह शेष क्षितिज (remaining horizon) के आधार पर प्रत्येक स्थिति पर कड़े मार्ग (Pinsker बनाम Coupling) का गतिशील रूप से चयन करता है।
   • एकीकृत बाउंड (Unified Bound $B^*$): उपरोक्त सभी बाउंड्स में से न्यूनतम। यह सभी डाइवर्जेंस व्यवस्थाओं में सबसे सटीक ज्ञात गारंटी प्रदान करता है।

मुख्य अंतर्दृष्टि: सभी बाउंड्स अधिकतम टोकन-स्तरीय विचलन ($D_{tok,max}$) पर निर्भर करते हैं। केवल औसत विचलन को नियंत्रित करना प्रमाणित रूप से अपर्याप्त है।

B. ट्रस्ट रीजन मास्किंग (TRM)

चूंकि टोकन-स्तरीय विधियाँ (जैसे PPO क्लिपिंग) $D_{tok,max}$ को नियंत्रित नहीं कर सकतीं, इसलिए लेखक ट्रस्ट रीजन मास्किंग (TRM) का प्रस्ताव करते हैं।

• तंत्र (Mechanism): व्यक्तिगत टोकन के लिए ग्रेडिएंट को क्लिप करने के बजाय, TRM पूरे अनुक्रम का मूल्यांकन करता है।
• मानदंड (Criterion): एक अनुक्रम $(x, y)$ तभी स्वीकार किया जाता है यदि उस अनुक्रम के सभी टोकन में अधिकतम KL डाइवर्जेंस एक थ्रेशोल्ड $\delta$ से नीचे हो:
  $$M(x, y) = \mathbb{I}\left[ \max_{t} D_{KL}(c_t) \leq \delta \right]$$
• कार्यान्वयन (Implementation):
  • फॉरवर्ड पास के दौरान, मॉडल $\pi_\theta$ के लिए लॉजिट्स (logits) की गणना करता है।
  • चूंकि $\pi_{roll}$ लॉजिट्स रोलआउट चरण से संग्रहीत होते हैं, इसलिए पूर्ण वोकैबुलरी पर सटीक $D_{KL}$ बिना किसी अतिरिक्त इन्फरेंस लागत के निकाला जा सकता है।
  • यदि शर्त का उल्लंघन होता है, तो पूरे अनुक्रम को मास्क (रिजेक्शन सैंपलिंग) कर दिया जाता है, जिससे अपडेट में शून्य ग्रेडिएंट योगदान मिलता है।
• सैद्धांतिक गारंटी: यदि वैश्विक शर्त $D_{tok,max}^{KL} \leq \delta$ बनी रहती है, तो एकीकृत बाउंड $B^*(\delta)$ लागू होता है, जो गैर-रिक्त निरंतर सुधार ($J(\pi_\theta) > J(\pi_{roll})$) सुनिश्चित करता है।

3. मुख्य योगदान (Key Contributions)

1. कड़े सैद्धांतिक बाउंड्स: बाउंड्स के एक परिवार का व्युत्पत्ति (Pinsker-Marginal, Mixed, Adaptive) जो $O(T^{3/2})$ या $O(T)$ के रूप में स्केल करते हैं, जो लंबे अनुक्रमों के लिए शास्त्रीय रिक्त $O(T^2)$ बाउंड्स में सुधार करते हैं।
2. नियंत्रण चर की पहचान: यह सिद्ध करना कि एरर बाउंड अधिकतम टोकन-स्तरीय विचलन पर निर्भर करता, न कि औसत पर, जो मानक टोकन-स्तरीय क्लिपिंग को अप्रभावी बनाता है।
3. ट्रस्ट रीजन मास्किंग (TRM): एक नवीन एल्गोरिदम जो अनुक्रम-स्तर पर ट्रस्ट रीजन को लागू करता है। यह लॉन्ग-होरइजन LLM-RL के लिए गैर-रिक्त निरंतर सुधार गारंटी प्रदान करने वाली पहली विधि है।
4. प्रायोगिक सत्यापन: गणितीय तर्क बेंचमार्क पर TRM की सफलता का प्रदर्शन, जहाँ मानक PPO विफल हो जाता है।

4. प्रयोगात्मक परिणाम (Experimental Results)

लेखकों ने गणितीय तर्क कार्यों (AIME25) पर Zero-RL सेटअप के साथ Qwen3-8B-Base का उपयोग करके TRM का मूल्यांकन किया।

• सेटअप: वास्तविक मिसमैच को सिम्युलेट करने के लिए इन्फरेंस के लिए vLLM और ट्रेनिंग के लिए PyTorch FSDP का उपयोग किया गया।
• मेट्रिक्स: लॉग एब्सोल्यूट परप्लेक्सिटी (PPL) गैप ($\pi_{roll}$ बनाम $\pi_\theta$ विचलन का माप) और AIME25 स्कोर।
• निष्कर्ष:
  • PPO क्लिपिंग की विफलता: मानक PPO क्लिपिंग अस्थिरता को बढ़ाता है, जिससे PPL गैप बढ़ता है और प्रदर्शन स्कोर गिर जाता है।
  • TRM की सफलता: TRM-Max (कठोर मैक्स थ्रेशोल्ड) और TRM-Avg (औसत थ्रेशोल्ड) दोनों वेरिएंट्स ने प्रशिक्षण के दौरान एक सीमित PPL गैप बनाए रखा।
  • प्रदर्शन: TRM वेरिएंट्स ने AIME25 स्कोर पर निरंतर सुधार दिखाया, जबकि PPO बेसलाइन्स या तो ध्वस्त हो गए या स्थिर हो गए।
  • संयुक्त मानदंड: मैक्स और औसत बाधाओं के संयोजन का उपयोग करने से सबसे मजबूत परिणाम मिले, जिसने आउटलेयर्स को पकड़ा और साथ ही संचयी ड्रिफ्ट को भी सीमित किया।

5. महत्व (Significance)

• सैद्धांतिक सफलता: यह शोध पत्र ऑटोरेग्रेसिव जनरेशन पर ट्रस्ट रीजन विधियों को लागू करने की एक लंबे समय से चली आ रही सैद्धांतिक सीमा को हल करता है। यह सिद्ध करता है कि बिना अधिकतम प्रति-टोकन विचलन को नियंत्रित किए, लॉन्ग-होरइजन RL सैद्धांतिक रूप से अनगारंटीड है।
• व्यावहारिक स्थिरता: यह उस "ट्रेनिंग-इन्फरेंस मिसमैच" समस्या के लिए एक ठोस, लागू करने योग्य समाधान (TRM) प्रदान करता है जो आधुनिक डिस्ट्रिब्यूटेड LLM ट्रेनिंग को प्रभावित करती है।
• स्केलेबिलिटी: गैर-रिक्त बाउंड्स प्रदान करके, TRM शॉर्ट टास्क से जटिल, लॉन्ग-चेन रीजनिंग टास्क तक RLHF/RL को सुरक्षित रूप से स्केल करने में सक्षम बनाता है, जो एजेंटिक AI के लिए आवश्यक हैं।
• सामान्यीकरण (Generalizability): यह दृष्टिकोण मौजूदा फ्रेमवर्क्स (जैसे GRPO/PPO) के साथ संगत है और इसे लंबे अनुक्रमों के प्रति पूर्वाग्रह को रोकने के लिए लेंथ-न्यूट्रल वेरिएंट्स (LN-TRM) में विस्तारित किया जा सकता है।

संक्षेप में, यह कार्य "व्यक्तिगत टोकन अपडेट को क्लिप करने" से "संपूर्ण ट्राजेक्टरीज को मान्य करने" के प्रति प्रतिमान (paradigm) को बदल देता है, जो स्थिर, लॉन्ग-होरइजन LLM सुदृढीकरण शिक्षण के लिए पहला कठोर सैद्धांतिक आधार प्रदान करता है।