Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

यहाँ चarafeddine Mouzouni के शोध पत्र "Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training" का विस्तृत तकनीकी सारांश दिया गया है।

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

Mixture-of-Experts (MoE) आर्किटेक्चर टोकन को विशेषज्ञ नेटवर्क के एक उपसमुच्चय (subset) की ओर रूट करके मॉडल क्षमता को स्केल करता है। एक महत्वपूर्ण इंजीनियरिंग चुनौती लोड बैलेंसिंग (load balancing) है: हस्तक्षेप के बिना, टोकन कुछ उच्च-गुणवत्ता वाले विशेषज्ञों पर केंद्रित होने लगते हैं, जिससे अन्य निष्क्रिय रह जाते हैं (एक्सपर्ट कोलैप्स)।

• वर्तमान स्थिति: मानक उपचार (standard remedies) एक सहायक बैलेंस लॉस (जैसे, Switch Transformers) का उपयोग करते हैं जिसमें लोड एकाग्रता को दंडित करने के लिए एक ट्यून करने योग्य गुणांक $\alpha$ होता है।
• अंतराल (The Gap): हालांकि ये विधियाँ व्यवहार में काम करती हैं, लेकिन इस बात की कोई सैद्धांतिक समझ नहीं है कि प्रशिक्षण प्रक्रिया के दौरान लोड बैलेंस और विशेषज्ञ गुणवत्ता के बीच का ट्रेडऑफ़ कैसे विकसित होता है। मौजूदा विश्लेषण आमतौर पर स्थिर (post-hoc) होते हैं और रूटिंग तंत्र के गतिशील प्रक्षेपवक्र (dynamic trajectory) को पकड़ने में विफल रहते हैं।

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

लेखक MoE टोकन रूटिंग को Mean-Field Games (MFG) के ढांचे के भीतर एक कंजेशन गेम (congestion game) के रूप में मॉडल करते हैं।

• गेम फॉर्मूलेशन:
  • खिलाड़ी (Players): टोकन।
  • संसाधन (Resources): विशेषज्ञ नेटवर्क।
  • लागत फलन (Cost Function): एक टोकन विशेषज्ञ गुणवत्ता ($q_i$) और भीड़भाड़ ($\gamma \mu_i$) के आधार पर लागत वहन करता है, जहाँ $\mu_i$ विशेषज्ञ $i$ पर लोड है और $\gamma$ कंजेशन गुणांक है।
  • संतुलन (Equilibrium): रूटिंग वितरण को एंट्रॉपी रेगुलराइजेशन (softmax) के साथ एक पोटेंशियल गेम के अद्वितीय संतुलन के रूप में मॉडल किया गया है।
• मुख्य पैरामीटर: प्रभावी कंजेशन गुणांक ($\gamma_{eff}$)। यह एकल पैरामीटर किसी भी प्रशिक्षण चेकपॉइंट पर बैलेंस-क्वालिटी ट्रेडऑफ़ की शक्ति को मात्रात्मक रूप से दर्शाता है। यह स्पष्ट सहायक लॉस और ग्रेडिएंट डिसेंट के माध्यम से सीखे गए अंतर्निहित (implicit) बैलेंस दोनों को कैप्चर करता है।
• पहचान (Identification): $\gamma_{eff}$ का अनुमान देखे गए लोड वितरण से सैद्धांतिक MFG संतुलन को फिट करके लगाया जाता है। लेखक देखे गए लोड और सैद्धांतिक बेस्ट-रिस्पॉन्स मैप के बीच $L_1$ दूरी को कम करते हैं।
• विश्लेषित मॉडल:
  • OLMoE-1B-7B: 20 चेकपॉइंट्स (शुरुआती "सर्ज" क्षेत्र में सघन नमूनाकरण)।
  • OpenMoE-8B: 6 चेकपॉइंट्स (अलग आर्किटेक्चर: 32 विशेषज्ञ, 4 MoE लेयर्स)।
• विस्तार (Extensions):
  • Multi-type MFG: विषमता (heterogeneity) को मॉडल करने के लिए टोकन को $K$ प्रकारों में क्लस्टर करता है, जिससे लोड भविष्यवाणी में सुधार होता है।
  • अपघटन (Decomposition): $\gamma_{eff}$ को स्पष्ट ($\gamma_{explicit} = \alpha M$) और अंतर्निहित ($\gamma_{implicit}$) घटकों में विभाजित करता है।

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

1. तीन चरणों वाला प्रक्षेपवक्र (The Three-Phase Trajectory): प्रीट्रेनिंग के दौरान $\gamma_{eff}$ के लिए एक नॉन-मोनोटोनिक, इनवर्टेड-U प्रक्षेपवपथ की खोज, जो स्थिर विश्लेषण द्वारा अदृश्य है।
2. सैद्धांतिक ईमानदारी (Theoretical Honesty): यह सिद्ध करना कि अच्छी तरह से संतुलित मॉडलों के लिए, सिंगल-टाइप MFG संतुलन गणितीय रूप से टेम्परेचर-स्केल्ड सॉफ्टमैक्स (temperature-scaled softmax) के समान है। गेम थ्योरी का लेंस स्थिर भविष्यवाणी सटीकता में सुधार नहीं करता है, बल्कि यह संरचनात्मक स्पष्टीकरण प्रदान करता है कि सॉफ्टमैक्स क्यों काम करता है और "टेम्परेचर" क्या दर्शाता है (कंजेशन गुणांक)।
3. प्रभावी कंजेशन अपघटन (Effective Congestion Decomposition): यह प्रदर्शित करना कि ऑप्टिमाइज़र स्पष्ट लॉस सिग्नल से कहीं अधिक अंतर्निहित बैलेंस को आत्मसात करता है। अभिसरण (convergence) पर, $\gamma_{eff}$, $\gamma_{explicit}$ से लगभग $13\times$ बड़ा है।
4. स्कोप डायग्नोस्टिक्स (Scope Diagnostics): ऐसे मेट्रिक्स विकसित करना (जैसे, $K/M$ अनुपात, कंटीन्यूएशन स्प्रेड $\epsilon_l$) जो यह निर्धारित करें कि MFG मॉडल कहाँ लागू होता है और कहाँ विफल हो जाता है (जैसे, उच्च टोकन विभेदीकरण वाले देर के लेयर्स)।

4. मुख्य परिणाम (Key Results)

A. तीन-चरण वाला प्रशिक्षण गतिकी (The Three-Phase Training Dynamics)

OLMoE-1B-7B के 20 चेकपॉइंट्स में $\gamma_{eff}$ को ट्रैक करने से तीन अलग-अलग चरण प्रकट होते हैं:

1. चरण 1: सर्ज (Surge) (स्टेप्स 5K–50K):
   • $\gamma_{eff}$ $\approx 14$ से तेजी से बढ़कर 36–39 के शिखर (स्टेप्स 30K–40K पर) तक पहुँच जाता है।
   • व्याख्या: राउटर लोड को संतुलित करने के लिए आक्रामक रूप से सीखता है। रूटिंग एंट्रॉपी बढ़ती है, और सिस्टम कच्चे विशेषज्ञ गुणवत्ता के बजाय उपयोगिता (utilization) को प्राथमिकता देता है।
2. चरण 2: स्थिरीकरण (Stabilization) (स्टेप्स 100K–400K):
   • $\gamma_{eff}$ $\approx 24–28$ पर स्थिर हो जाता है।
   • व्याख्या: राउटर अपना ऑपरेटिंग पॉइंट पा चुका है। विशेषज्ञ विशिष्टता (specialization) शुरू करते हैं (क्वालिटी स्प्रेड $B_0$ गिरता है), लेकिन बैलेंस बाधा स्थिर रहती है।
3. चरण 3: रिलैक्सेशन (Relaxation) (स्टेप्स 400K–1.2M):
   • $\gamma_{eff}$ $\approx 27$ से काफी घटकर 8.5 हो जाता है।
   • व्याख्या: एक बार जब विशेषज्ञ पर्याप्त रूप से विभेदित हो जाते हैं, तो राउटर गुणवत्ता को प्राथमिकता देने के लिए बैलेंस बाधा को "रिलैक्स" कर देता है। मॉडल एकरूपता के बदले चयनशीलता (selectivity) का व्यापार करता है।

• प्रतिकृति (Replication): इस इनवर्टेड-U पैटर्न को OpenMoE-8B पर भी दोहराया गया, जो पुष्टि करता है कि यह किसी विशिष्ट आर्किटेक्चर का प्रभाव नहीं है।
• फाइन-ट्यूनिंग: तीन-चरण वाला पैटर्न प्रीट्रेनिंग के लिए विशिष्ट है। फाइन-ट्यूनिंग (एनीलिंग) के दौरान, $\gamma_{eff}$ स्थिर रहता है, जो दर्शाता है कि रूटिंग इक्विलिब्रियम पहले ही सेट हो चुका है।

B. स्थिर बनाम गतिशील प्रदर्शन (Static vs. Dynamic Performance)

• स्थिर समानता (Static Equivalence): अभिसरित मॉडल पर, सिंगल-टाइप MFG और टेम्परेचर-स्केल्ड सॉफ्टमैक्स सांख्यिकीय रूप से अविभेद्य हैं (Held-out $L_1$: 0.199 बनाम 0.200)। गेम थ्योरी मॉडल एकल स्नैपशॉट के लिए कोई अतिरिक्त भविष्य कहने वाली शक्ति नहीं जोड़ता है।
• गतिशील मूल्य (Dynamic Value): इसका मूल्य पूरी तरह से डायनामिक्स में है। नॉन-मोनोटोनिक प्रक्षेपवक्र यह प्रकट करता है कि शुरुआती प्रशिक्षण संतुलन को प्राथमिकता देता है, जबकि देर का प्रशिक्षण गुणवत्ता को प्राथमिकता देता है।

C. अंतर्निहित बैलेंस का प्रभुत्व (Implicit Balance Dominance)

• अभिसरण पर, स्पष्ट सहायक लॉस $\gamma_{explicit} = 0.64$ में योगदान देता है।
• फिट किया गया प्रभावी कंजेशन $\gamma_{eff} = 8.5$ है।
• परिणाम: ऑप्टिमाइज़र स्पष्ट लॉस द्वारा प्रदान किए गए $\gamma_{explicit}$ की तुलना में $13\times$ अधिक प्रभावी कंजेशन को आत्मसात करता है। सहायक लॉस एक "बीज" (seed) के रूप में कार्य करता है, लेकिन ग्रेडिएंट डायनेमिक्स बैलेंस को आंतरिक रूप से विकसित करते हैं।

D. मल्टी-टाइप विस्तार (Multi-Type Extension)

• टोकन को प्रकारों ($K=4$) में क्लस्टर करके, मल्टी-टाइप MFG सभी 16 लेयर्स में सिंगल-टाइप मॉडल की तुलना में औसत रूप से लोड भविष्यवाणी सटीकता में 30% सुधार करता है।
• यह पुष्टि करता है कि टोकन विषमता (heterogeneity) महत्वपूर्ण है, हालांकि कपलिंग प्रभाव अच्छी तरह से संतुलित मॉडलों में कम स्पष्ट है।

5. महत्व और निहितार्थ (Significance and Implications)

• नया डायग्नोस्टिक टूल: $\gamma_{eff}$ MoE प्रशिक्षण के लिए एक वास्तविक समय के मॉनिटर के रूप में कार्य करता है। $\gamma_{eff}$ में समय से पहले गिरावट या "सर्ज" शिखर तक पहुँचने में विफलता विशेषज्ञ पतन (expert collapse) या प्रशिक्षण अस्थिरता का संकेत दे सकती है।
• सहायक लॉस का पुनर्गठन (Reframing the Auxiliary Loss): यह पेपर इस विचार को चुनौती देता है कि सहायक लॉस बैलेंस का प्राथमिक चालक है। इसके बजाय, यह सुझाव देता है कि लॉस एक उत्प्रेरक (catalyst) है जो एक बहुत बड़े आंतरिक बैलेंसिंग तंत्र को जन्म देता है जिसे ऑप्टिमाइज़र द्वारा संचालित किया जाता है।
• अनुकूलन रणनीति (Optimization Strategy): ये निष्कर्ष एक संभावित एडेप्टिव बैलेंस शेड्यूल (adaptive balance schedules) का सुझाव देते हैं। अभ्यासकर्ता सैद्धांतिक रूप से "रिलैक्सेशन" चरण (चरण 3) के दौरान स्पष्ट बैलेंस गुणांक ($\alpha$) को कम कर सकते हैं ताकि राउटर को सिस्टम को अस्थिर किए बिना गुणवत्ता पर ध्यान केंद्रित करने की अनुमति मिल सके।
• सैद्धांतिक अंतर्दृष्टि: यह MoE राउटर में सॉफ्टमैक्स की व्यापकता के लिए एक कठोर गेम-थ्योतिक औचित्य प्रदान करता है, जो "टेम्परेचर" पैरामीटर को सीधे पोटेंशियल गेम के कंजेशन गुणांक से जोड़ता है।

निष्कर्ष (Conclusion)

यह पेपर स्थापित करता है कि MoE प्रशिक्षण एक निश्चित उद्देश्य का स्थिर अनुकूलन नहीं है, बल्कि तीन विशिष्ट शासन (regimes) के माध्यम से एक गतिशील विकास है: सर्ज, स्थिरीकरण और रिलैक्सेशन। एक कंजेशन गेम के रूप में रूटिंग को मॉडल करके, लेखक प्रकट करते हैं कि राउटर का "टेम्परेचर" एक गतिशील पैरामीटर है जो लोड बैलेंसिंग और विशेषज्ञ विशिष्टता के बीच तनाव को दर्शाता है, एक ऐसी घटना जो मानक पोस्ट-ट्रेनिंग विश्लेषण में अदृश्य है।