ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

तकनीकी सारांश: ProbMoE: Mixture-of-Experts के लिए विभेदक संभावabilistic रूटिंग (Differentiable Probabilistic Routing)

समस्या विवरण
Mixture-of-Experts (MoE) मॉडल केवल प्रत्येक टोकन के लिए विशेषज्ञों के एक छोटे उपसमुच्चय (subset) को सक्रिय करके कुशलतापूर्वक स्केल करते हैं। हालाँकि, इन मॉडलों को प्रशिक्षित करना चुनौतीपूर्ण बना हुआ है क्योंकि मानक रूटिंग एक असतत (discrete), गैर-विभेदक (non-differentiable) top-$k$ ऑपरेटर पर निर्भर करती है। पारंपरिक प्रशिक्षण प्रक्रियाएँ इसे दरकिनार करने के लिए चयनित विशेषज्ञों की सॉफ्टमैक्स संभावनाओं (softmax probabilities) के माध्यम से केवल ग्रेडिएंट प्रसारित करती हैं, जिससे राउटर लॉजिट्स (router logits) पर top-$k$ चयन की निर्भरता को अनदेखा कर दिया जाता है। यह सन्निकटन (approximation) वैकल्पिक विशेषज्ञ उपसमुच्चयों के लिए सीखने के संकेत को सीमित करता है, जिससे अक्सर केंद्रित रूटिंग वितरण, खराब विशेषज्ञ उपयोगिता और अस्थिर प्रशिक्षण गतिशीलता उत्पन्न होती है। मौजूदा प्रयास, जैसे कि डेंस बैकप्रोपैगेशन या ह्यूरिस्टिक ग्रेडिएंट अनुमानक, अक्सर विशेषज्ञ उपसमुच्चयों के वितरण को स्पष्ट रूप से मॉडल करने में विफल रहते हैं या प्रशिक्षण के दौरान वैकल्पिक संयोजनों की खोज के लिए सीमित तंत्र प्रदान करते हैं।

कार्यप्रणाली
लेखक ProbMoE का प्रस्ताव करते हैं, जो विशेषज्ञ रूटिंग को कार्डिनैलिटी-प्रतिबंधित (cardinality-constrained) विशेषज्ञ उपसमुच्चयों पर एक संभावabilistic अनुमान (probabilistic inference) समस्या के रूप में पुनर्गठित करता है। विशेषज्ञ चयन को एक नियत (deterministic) ऑपरेटर के रूप में मानने के बजाय, ProbMoE एक वितरण के रूप में विशेषज्ञों के चयन को मॉडल करता है।

1. संभाव्यता सूत्रीकरण (Probabilistic Formulation):

   • Exact-$k$ रूटिंग: ढांचा एक ऐसे उपसमुच्चय $S$ पर एक वितरण को परिभाषित करता है जहाँ $|S|=k$ है। यह वितरण स्वतंत्र बर्नौली चयन संभावनाओं (जो राउटर लॉजिट्स से सिग्मॉइड के माध्यम से प्राप्त होती हैं) को इस शर्त पर आधारित करता है कि ठीक $k$ विशेषज्ञ चुने गए हैं। इस वितरण के लिए नॉर्मलाइजिंग कांस्टेंट $Z_k$ की गणना डायनेमिक प्रोग्रामिंग का उपयोग करके सटीक रूप से की जाती है, जिससे यह सूत्रीकरण सुलभ हो जाता है।
   • Dynamic-$k$ रूटिंग: ढांचा एक रेंज-प्रतिबंधित सेटिंग में विस्तार करता है जहाँ उपसमुच्चय का आकार $|S|$, $[k_{min}, k_{max}]$ के बीच होता है। वितरण इस सीमा पर आधारित है, जो मॉडल को प्रत्येक टोकन के लिए विशेषज्ञ पहचान और रूटिंग कार्डिनैलिटी (routing cardinality) को संयुक्त रूप से अनुमानित करने की अनुमति देता है।

2. फॉरवर्ड पास (Forward Pass):

   • प्रशिक्षण के दौरान, ProbMoE परिभाषित वितरण (exact-$k$ या range-constrained) से एक विशेषज्ञ उपसमुच्चय $S$ का नमूना (sample) लेता है।
   • MoE लेयर केवल नमूना लिए गए विशेषज्ञों का उपयोग करके आउटपुट की गणना करती है, जिससे मानक MoE मॉडलों के स्पार्स निष्पादन पैटर्न (sparse execution pattern) को बनाए रखा जाता है।
   • अनुमान (inference) के समय, मैक्सिमम-ए-पोस्टीरियर (MAP) सेट का चयन किया जाता है, जिससे यह सुनिश्चित होता है कि अनुमान लागत पारंपरिक MoE रूटिंग के समान हो।

3. बैकवर्ड पास (Backward Pass):

   • ग्रेडिएंट-आधारित अनुकूलन को सक्षम करने के लिए, ProbMoE SIMPLE ग्रेडिएंट अनुमानक (Ahmed et al., 2023) का उपयोग करता है।
   • डिस्क्रीट सैंपल के माध्यम से विभेदन करने के बजाय, बैकवर्ड पास का उपयोग प्रत्येक विशेषज्ञ के चुने जाने की मार्जिनल संभावनाओं (marginal probabilities) के लिए किया जाता है।
   • रूटिंग वेट्स को स्ट्रैट-थ्रू एस्टीमेटर (straight-through estimator) का उपयोग करके बनाया जाता है: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, जहाँ $z$ नमूना लिया गया मास्क है, $m$ सटीक मार्जिनल संभावनाओं का वेक्टर है, और $\pi$ सॉफ्टमैक्स रूटिंग संभावनाएँ हैं।
   • यह ग्रेडिएंट्स को मार्जिनल संभावनाओं के माध्यम से प्रवाहित होने की अनुमति देता है, जो राउटर को सूचनात्मक अपडेट प्रदान करता है कि चयन संभावनाओं में परिवर्तन अपेक्षित हानि (expected loss) को कैसे प्रभावित करता है, जबकि स्पार्स विशेषज्ञ निष्पादन को बनाए रखता है।

प्रमुख योगदान

1. संभाव्यता अनुमान सूत्रीकरण: शोधपत्र MoE रूटिंग को कार्डिनैलिटी-प्रतिबंधित विशेषज्ञ उपसमुच्चयों पर संभाव्यता अनुमान के रूप में प्रस्तुत करता है, जो ह्यूरिस्टिक ग्रेडिएंट सन्निकटन से आगे बढ़ता है।
2. सुलभ मार्जिनल-आधारित अनुकूलन: SIMPLE अनुमानक का लाभ उठाकर, लेखक स्टोकेस्टिक उपसमुच्चय सैंपलिंग को सटीक मार्जिनल-आधारित राउटर ग्रेडिएंट्स के साथ जोड़ते हैं, जो डेंस विशेषज्ञ-पक्ष गणना के बिना विभेदक प्रशिक्षण को सक्षम बनाता है।
3. एकीकृत Dynamic-$k$ ढांचा: विधि स्वाभाविक रूप से डायनेमिक रूटिंग में विस्तारित होती है, जहाँ मॉडल पूर्व-निर्धारित सीमा के भीतर सक्रिय विशेषज्ञों की संख्या और विशेषज्ञ पहचान दोनों का अनुमान लगाता है, जो किसी भी ह्यूरिस्टिक गेटिंग नियमों पर निर्भर नहीं करता है।
4. अनुभवजन्य सुधार (Empirical Improvements):
   • Exact-$k$: ProbMoE OLMoE और Qwen बैकबोन पर बेंचमार्क (GSM8K, Law, Translation, MBPP, MMLU) में मजबूत प्रदर्शन प्राप्त करता है, जो अक्सर DenseMixer और पारंपरिक Top-$k$ जैसे प्रतिस्पर्धी बेसलाइन के बराबर या उनसे बेहतर होता है।
   • Dynamic-$k$: ProbMoE Dynamic-$k$ औसत रूप से कम विशेषज्ञों को सक्रिय करते हुए Exact-$k$ के समान प्रदर्शन प्राप्त करता है, जो प्रभावी अनुकूलन आवंटन (adaptive allocation) को प्रदर्शित करता है।
5. रूटिंग विविधता: विश्लेषण से पता चलता है कि ProbMoE बेसलाइन की तुलना में व्यापक रूटिंग वितरण और अधिक संतुलित विशेषज्ञ उपयोगिता प्रेरित करता है, जिससे विशेषज्ञ पतन (expert collapse) कम होता है।

परिणाम और महत्व
शोधपत्र प्रदर्शित करता है कि ProbMoE विशेषज्ञ उपयोगिता और रूटिंग विविधता में सुधार करता है। एब्लेशन अध्ययनों में, स्टोकेस्टिक सैंपलिंग और मार्जिनल-आधारित अनुकूलन का संयोजन प्रदर्शन में वृद्धि का प्राथमिक चालक सिद्ध हुआ है, जो उन कॉन्फ़िगरेशन से बेहतर है जो केवल इनमें से एक घटक का उपयोग करते हैं।

लेखक दावा करते हैं कि ProbMoE विशेषज्ञ चयन के लिए एक सिद्धांत-आधारित विधि प्रदान करता है जो प्रशिक्षण स्थिरता और कम्प्यूटेशनल दक्षता में सुधार करता है। रूटिंग को संभाव्यता के आधार पर मॉडल करके, ढांचा प्रशिक्षण के दौरान वैकल्पिक विशेषज्ञ संयोजनों की खोज की अनुमति देता है, जिससे बेहतर-कैलिब्रेटेड रूटिंग वितरण प्राप्त होता है। Dynamic-$k$ संस्करण आगे दिखाता है कि टोकन-स्तरीय जटिलता से मेल खाने के लिए अनुकूलनीय कार्डिनैलिटी (adaptive cardinality) को सीखा जा सकता है, जो कार्य प्रदर्शन से समझौता किए बिना प्रभावी रूटिंग लागत को कम करता है। यह कार्य सुझाव देता है कि भविष्य की दिशाओं में डायनेमिक सेटिंग्स में स्पर्सिटी (sparsity) का पूर्ण लाभ उठाने के लिए प्री-ट्रेनिंग और सिस्टम-लेवल ऑप्टिमाइज़ेशन को शामिल करना शामिल है।