तकनीकी सारांश: पदानुक्रमित मिश्रण-विशेषज्ञों के साथ दो-चरणीय अनुकूलन (Hi-MoE)

1. समस्या विवरण

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

• रूटिंग कोलैप्स (Routing Collapse): आक्रामक विविधता तंत्र अक्सर "रूटिंग कोलैप्स" की ओर ले जाते हैं, जहाँ कुछ विशेषज्ञ टोकन चयन पर हावी हो जाते हैं जबकि अन्य निष्क्रिय रहते हैं। यह अक्षम क्षमता उपयोग और उपकरणों (devices) के बीच लोड असंतुलन (stragglers) का कारण बनता है।
• ग्रुपिंग डिलेमा (The Grouping Dilemma): हालिया दृष्टिकोण समूहों में विशेषज्ञों को समूहित करके और प्रत्येक समूह से निश्चित चयन लागू करके असंतुलन को कम करते हैं (जैसे, $M$ समूहों में से प्रत्येक से $K/M$ विशेषज्ञों का चयन करना)। जबकि यह डिवाइस-स्तर पर समान गणना सुनिश्चित करता है, यह रूटिंग की स्वतंत्रता को प्रतिबंधित करता है। यह बाधा अक्सर अतिरेक (redundancy) की ओर ले जाती है, जहाँ विभिन्न समूहों के विशेषज्ञ समान कार्यों को सीखते हैं क्योंकि वे समान टोकन मिश्रणों पर अनुकूलित होते हैं, जो प्रभावी रूप से विशेषज्ञ-स्तर के कोलैप्स को समूहों के बीच "विविधता कोलैप्स" (diversity collapse) से बदल देता है।

मौजूदा लॉस-आधारित संतुलन विधियाँ (जैसे GShard) अक्सर कार्य ग्रेडिएंट्स (task gradients) के साथ संघर्ष करती हैं, जिससे अभिसरण (convergence) धीमा हो जाता है, जबकि लॉस-मुक्त विधियाँ रेडंडेंसी को रोकने में विफल हो सकती हैं। एक ऐसे ढांचे की आवश्यकता है जो हार्डवेयर-संरेखित संतुलन और सार्थक, पूरक विशेषज्ञता को एक साथ लागू कर सके।

വMethodology: Hi-MoE

लेखक Hi-MoE का प्रस्ताव करते हैं, जो एक पदानुक्रमित ढांचा है जो रूटिंग नियंत्रण को दो जुड़े हुए स्तरों में विभाजित करता है: इंटर-ग्रुप बैलेंसिंग (inter-group balancing) और इंट्रा-ग्रुप स्पेशलाइजेशन (intra-group specialization)।

आर्किटेक्चर

• समूहीकृत संरचना (Grouped Structure): विशेषज्ञों को $M$ विच्छेदित (disjoint) समूहों में विभाजित किया गया है (अक्सर GPU शेर्ड्स के साथ संरेखित)।
• रूटिंग बाधा (Routing Constraint): प्रत्येक टोकन के लिए, राउटर प्रत्येक समूह $m$ से विशेषज्ञों की एक निश्चित संख्या ($K_m$) चुनता है।
• दो-स्तरीय अनुकूलन (Two-Level Optimization): यह ढांचा पदानुक्रम के प्रत्येक स्तर के लिए विशिष्ट रेगुलराइज़र पेश करता है।

मुख्य घटक

1. इंटर-ग्रुप रेगुलराइजेशन ($R_{inter}$):

   • लक्ष्य: समूहों के बीच निष्पक्ष ट्रैफिक वितरण लागू करना ताकि डिवाइस उपयोग संतुलित रहे।
   • तंत्र: समूहों के बीच पोस्ट-सिलेक्शन रूटिंग वेट्स के $\ell_2$ संकेंद्रण (concentration) को कम करता है। विशिष्ट समूहों पर केंद्रित रूटिंग मास को दंडित करके, यह समूह भार के कोएफिशिएंट ऑफ वेरिएशन (CV) को न्यूनतम करने के प्रॉक्सी के रूप में कार्य करता है।
   • सैद्धांतिक आधार: थ्योरम 4.1 सिद्ध करता है कि इस पद को कम करने से समूह-स्तर के असंतुलन पर एक प्रमाणित ऊपरी सीमा (upper bound) प्राप्त होती है।

2. इंट्रा-ग्रुप रेगुलराइजेशन ($R_{intra}$):

   • लक्ष्य: समूहों के भीतर पूरक विशेषज्ञता को बढ़ावा देना और विशेषज्ञों को रेडंडेंट फंक्शन सीखने से रोकना।
   • तंत्र: प्री-सिलेक्शन रूटिंग वितरण के $\ell_2$ संकेंद्रण को अधिकतम करता है (प्रभावी रूप से एक एंटी-ओवरलैप टर्म)। यह राउटर को एक समूह के भीतर कई विशेषज्ञों के साथ कम आत्मविश्वास के साथ मिश्रण करने के बजाय, विशिष्ट विशेषज्ञों के प्रति निर्णायक, स्पार्स प्रतिबद्धता बनाने के लिए प्रोत्साहित करता है।
   • सैद्धांतिक आधार: थ्योरम 4.2 दिखाता है कि रूटिंग ओवरलैप को कम करने से विशेषज्ञों के बीच ग्रेडिएंट कपलिंग कम होती है, जिससे वे अलग-अलग ऑप्टिमाइजेशन ट्रेजेक्टरी का पालन कर पाते हैं। थ्योरम 4.3 प्रदर्शित करता है कि संतुलित मार्जिनल्स के तहत, रूटिंग संकेंद्रण को बढ़ाने से इनपुट और चयनित विशेषज्ञ के बीच म्यूचुअल इंफॉर्मेशन अधिकतम होती है, जो विशेषज्ञता का एक प्रॉक्सी है।

3. स्थिरीकरण तंत्र (Stabilization Mechanism):

   • इंट्रा-ग्रुप विविधता सिग्नल को डिजेनरेट रिजीम में जाने से रोकने के लिए, लेखक राउटर लॉजिट्स के एक्सपोनेंशियल मूविंग एवरेज (EMA) का उपयोग करके एक बायस-करेक्टेड रूटिंग गणना का उपयोग करते हैं। यह बिना किसी अलग बैलेंसिंग लॉस के, ऐतिहासिक रूप से पसंदीदा विशेषज्ञों के निरंतर ओवर-सिलेक्शन को रोकता है।

ऑब्जेक्टिव फंक्शन

कुल प्रशिक्षण उद्देश्य कार्य हानि (task loss), मानक लोड-बैलेंसिंग लॉस और प्रस्तावित पदानुक्रमित पदों को जोड़ता है:
$$\mathcal{L} = \mathcal{L}_{task} + \mathcal{L}_{load} + R_{intra} + R_{inter}$$
लेखक इसे एक बाधा-युक्त अनुकूलन समस्या के रूप में फ्रेम करते हैं जहाँ गुणांक $\lambda_{intra}$ और $\lambda_{inter}$ क्रमशः विविधता और संतुलन बाधाओं के लिए लैग्रेंज मल्टीप्लायर के रूप में कार्य करते हैं।

3. मुख्य योगदान

• तनाव की पहचान: पेपर यह पहचानता है कि सख्त समूह बाधाओं के माध्यम से डिवाइस-स्तर के संतुलन को लागू करने से उपयोगिता तो स्थिर हो सकती है लेकिन समूहों के बीच रेडंडेंसी बढ़ सकती है।
• पदानुक्रमित ढांचा: Hi-MoE एक मॉड्यूलर, ड्रॉप-इन ऑब्जेक्टिव पेश करता है जो इन चिंताओं को अलग करता है:
  • हार्डवेयर-संरेखित संतुलन के लिए एक इंटर-ग्रउंट ऑब्जेक्टिव।
  • कार्यात्मक विविधता और विशेषज्ञता के लिए एक इंट्रा-ग्रुप ऑब्जेक्टिव।
• सैद्धांतिक गारंटी: पेपर ग्रुप असंतुलन बाउंड्स, ग्रेडिएंट डिकपलिंग और बढ़ी हुई रूटिंग इंफॉर्मेशन (विशेषज्ञता) से जुड़े औपचारिक प्रमाण प्रदान करता है।
• अनुभवजन्य सत्यापन (Empirical Validation): विजन और लैंग्वेज डोमेन में सुसंगत सुधार, जो दर्शाता है कि Hi-MoE, पर्सप्लेक्सिटी/एक्यूरेसी और लोड बैलेंस (CV) के बीच के पैरेटो फ्रंटियर (Pareto frontier) को बढ़ाता है।

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

लेखकों ने तीन अलग-अलग प्री-ट्रेनिंग सेटिंग्स पर Hi-MoE का मूल्यांकन किया:

A. कंप्यूटर विजन (Tiny ImageNet पर Swin-MoE)

• सेटअप: 64 विशेषज्ञों (4 समूहों में 16) के साथ 1B पैरामीटर मॉडल।
• परिणाम: Vanilla GShard, Loss-free balancing, और MoGE की तुलना में Hi-MoE ने उच्चतम Top-1 सटीकता (59.13%) और सर्वश्रेष्ठ लोड बैलेंस (CV 0.29) प्राप्त किया।
• अवलोकन: अटेंशन मैप्स के विज़ुअलाइज़ेशन ने खुलासा किया कि Hi-MoE समूहों ने विशिष्ट विजुअल विशेषज्ञता विकसित की (जैसे, एक समूह शरीर के अंगों पर ध्यान केंद्रित करता है, दूसरा बैकग्राउंड सतहों पर), जबकि बेसलाइन अधिक रेडंडेंसी दिखाते हैं।

B. लैंग्वेज मॉडलिंग (OpenWebText पर nanoGPT)

• सेटअप: 8 विशेषज्ञों (4 समूहों में 2) के साथ 1.06B पैरामीटर मॉडल।
• परिणाम: Hi-MoE ने सबसे कम पर्सप्लेक्सिटी (2.947) और सर्वश्रेष्ठ बैलेंस (CV 0.23) प्राप्त किया।
• एब्लेशन (Ablation): $R_{inter}$ को हटाने से बैलेंस कम हुआ; $R_{intra}$ को हटाने से क्वालिटी कम हुई। बायस करेक्शन ने एक स्टेबलाइज़र के रूप में कार्य किया।
• दक्षता: पदानुक्रमित राउटर ने प्रति लेयर न्यूनतम ओवरहेड (1.3ms से 1.6ms) जोड़ा, जबकि फ्लैट रूटिंग की तुलना में सिंक्रोनाइज़ेशन ओवरहेड को ~3x कम कर दिया।

C. बड़े पैमाने पर प्री-ट्रेनिंग (OLMoE-7B बनाम Hi-MoE-7B)

• सेटअप: 58.54B टोकन (Dolma मिश्रण) पर प्रशिक्षित 7B पैरामीटर मॉडल।
• पर्सप्लेक्सिटी: Hi-Moel-7B ने मानक OLMoE-7B बेसलाइन की तुलना में 11 वैलिडेशन डोमेन में औसतन 5.6% पर्सप्लेक्सिटी कम की।
• डाउनस्ट्रीम टास्क: Hi-MoE ने MMLU स्लाइस (जैसे, +22.2% STEM, +11.9% Humanities) और कॉमन सेंस रीजनिंग कार्यों पर महत्वपूर्ण लाभ दिखाया।
• लोड बैलेंस: Hi-MoE ने विशेषज्ञ संतुलन (CV) में 40% का सुधार किया (0.153 बनाम 0.257), जो सीधे तौर पर बड़े पैमाने पर प्रशिक्षण में संसाधन दक्षता की बाधा को संबोधित करता है।

5. महत्व और दावे

पेपर का दावा है कि Hi-MoE स्पार्स MoE ट्रेनिंग में निहित "बैलेंस-डाइवर्सिटी" तनाव को हल करता है। रूटिंग के पदानुक्रम (ग्रुप बनाम एक्सपर्ट) को स्पष्ट रूप से मॉडल करके, यह विधि मॉडलों को अनुमति देती है:

1. कुशलता से स्केल करना: मॉडल की कंडीशनल कैपेसिटी से समझौता किए बिना समान GPU यूटिलाइजेशन बनाए रखना।
2. विशेषीकृत प्रतिनिधित्व सीखना: उस "डाइवर्सिटी कोलैप्स" को रोकना जहाँ ग्रुप किए गए विशेषज्ञ कार्यात्मक रूप से रेडंडेंट हो जाते हैं।
3. गुणवत्ता में सुधार: विजन और लैंग्वेज टास्क में पर्सप्लेक्सिटी और डाउनस्ट्रीम टास्क प्रदर्शन में निरंतर लाभ प्राप्त करना।

लेखक Hi-MoE को मौजूदा स्टेबिलाइज़ेशन मैकेनिज्म के विकल्प के रूप में नहीं, बल्कि एक पूरक, मॉड्यूलर एन्हांसमेंट के रूप में देखते हैं जो विशेषज्ञ ग्रुपिंग को केवल एक सिस्टम कंस्ट्रेंट के बजाय एक सिद्धांत आधारित ऑप्टिमाइजेशन मैकेनिज्म बनाता है। यह कार्य सुझाव देता है कि पदानुक्रमित समन्वय (hierarchical coordination) ट्रिलियन-पैरामीटर मॉडलों में विश्वसनीय गुणवत्ता लाभ को अनलॉक करने का एक व्यवहार्य मार्ग है।