तकनीकी सारांश: बेहतर स्पार्स मिक्स्चर ऑफ एक्सपर्ट्स के लिए सिलेक्टिव सिंकहॉर्न रूटिंग (Selective Sinkhorn Routing)

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

स्पार्स मिक्स्चर-ऑफ-एक्सपर्ट्स (SMoE) मॉडल, प्रति टोकन विशेषज्ञों के एक उपसमूह (subset) को सक्रिय करके कम्प्यूटेशनल दक्षता बनाए रखते हुए मॉडल क्षमता बढ़ाने का एक स्केलेबल दृष्टिकोण प्रदान करते हैं। हालांकि, सॉफ्टमैक्स-आधारित रूटिंग पर निर्भर पारंपरिक SMoE मॉडल अक्सर "रूटिंग कोलैप्स" (routing collapse) से ग्रस्त होते हैं, जहाँ विशेषज्ञों का एक छोटा समूह अत्यधिक उपयोग किया जाता है जबकि अन्य कम उपयोग होते हैं।

मौजूदा शमन रणनीतियाँ (mitigation strategies) आमतौर पर इन पर निर्भर करती हैं:

1. सहायक उद्देश्य (Auxiliary Objectives): लोड-बैलेंसिंग लॉस या z-लॉस जो असमान विशेषज्ञ उपयोग को दंडित करते हैं। ये ऑब्जेक्टिव मिसअलाइनमेंट, प्रशिक्षण अस्थिरता और जटिलता को बढ़ा सकते हैं।
2. नॉइजी गेटिंग (Noisy Gating): अन्वेषण (exploration) को प्रोत्साहित करने के लिए शोर (noise) डालना, जो अतिरिक्त ट्रेन करने योग्य पैरामीटर या हाइपरपैरामीटर जोड़ता है।
3. सिंकहॉर्न-आधारित रूटिंग (Sinkhorn-based Routing): ऑप्टिमल ट्रांसपोर्ट (OT) के लिए सिंकहॉर्न एल्गोरिदम का उपयोग करने वाले पूर्ववर्ती तरीके अक्सर ट्रांसपोर्ट मैप का उपयोग केवल विशेषज्ञ चयन (top-k) के लिए करते हैं, न कि वेट असाइनमेंट के लिए। इसके अलावा, वे रूटिंग लचीलेपन को कम कर सकते हैं क्योंकि गेटिंग मैट्रिक्स को सॉफ्टमैक्स विधियों की तरह सीधे ग्रेडिएंट-आधारित लर्निंग के माध्यम से अनुकूलित नहीं किया जाता है, या वे सिंकहॉर्न अपडेट के पुनरावृत्ति (iterative) स्वभाव के कारण पर्याप्त प्रशिक्षण ओवरहेड का सामना करते हैं।

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

यह शोध पत्र सिलेक्टिव सिंकहॉर्न रूटिंग (SSR) प्रस्तावित करता है, जो एक हल्का (lightweight) तंत्र है जो टोकन-टू-एक्सपर्ट असाइनमेंट को एंट्रॉपी-रेगुलराइज्ड ऑप्टिमल ट्रांसपोर्ट समस्या के रूप में पुनर्गठित करता है।

मुख्य सूत्र (Core Formulation)

लेखक $m$ टोकनों को $n$ विशेषज्ञों के असाइनमेंट को एंट्रॉपी-रेगुलराइज्ड ऑप्टिमल ट्रांसपोर्ट समस्या के रूप में मॉडल करते हैं। लक्ष्य एक ट्रांसपोर्ट प्लान $\hat{\Pi}$ खोजना है जो विशिष्ट बाधाओं को पूरा करते हुए अनुकूलता (compatibility) को अधिकतम करता है:

• सकारात्मकता (Positivity): $\Pi > 0$।
• रो सम (Row Sum): प्रत्येक टोकन अपना संपूर्ण द्रव्यमान (mass) विशेषज्ञों को रूट करता है ($\Pi \mathbf{1}_n = \mathbf{1}_m$)।
• कॉलम सम (Column Sum): प्रत्येक विशेषज्ञ समान अपेक्षित कुल भार प्राप्त करता है ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$)।

कॉस्ट मैट्रिक्स $C$ को गेटिंग स्कोर $S$ से प्राप्त किया जाता है। लेखक दो वेरिएंट प्रस्तावित करते हैं:

• लीनियर कॉस्ट (Linear Cost): $C = S$।
• सॉफ्टमैक्स कॉस्ट (Softmax Cost): $C_{i,:} = \text{softmax}(S_{i,:})$, जिसका उपयोग सिंकहॉर्न पुनरावृत्तियों के दौरान संख्यात्मक ओवरफ्लो को रोकने के लिए किया जाता है।

उन पूर्ववर्ती OT-आधारित तरीकों के विपरीत जो ट्रांसपोर्ट प्लान का उपयोग केवल चयन के लिए करते हैं, SSR सीधे गणना करने के लिए $\hat{\Pi}$ के मानों का उपयोग रूटिंग वेट के रूप में करता है। प्रत्येक टोकन के लिए, $\hat{\Pi}$ की संबंधित पंक्ति में उच्चतम प्रविष्टियों के आधार पर शीर्ष-$k$ विशेषज्ञों का चयन किया जाता है, और वेट्स को इन प्रविष्टियों से सामान्यीकृत (normalize) किया जाता है।

चयनात्मक रणनीति और शोर इंजेक्शन (Selective Strategy and Noise Injection)

एक प्रमुख चुनौती जो पहचानी गई है वह यह है कि सिंकहहॉर्न रूटिंग गेटिंग वेट मैट्रिक्स $W_g$ को अपडेट नहीं करता है क्योंकि $W_g$, OT ऑब्जेक्टिव से अलग (decoupled) है। इसे संबोधित करने के लिए, SSR एक हाइब्रिड प्रशिक्षण रणनीति अपनाता है:

• चयनात्मक अनुप्रयोग (Selective Application): प्रशिक्षण के दौरान, प्रत्येक MoE ब्लॉक एक छोटी संभावना $p$ (जैसे, 0.001) के साथ सिंकहॉर्न रूटिंग का उपयोग करता है और अन्यथा मानक सॉफ्टमैक्स गेटिंग का उपयोग करता है। यह सॉफ्टमैक्स गेटिंग को $W_g$ को प्रभावी ढंग से अपडेट करने की अनुमति देता है जबकि सिंकहॉर्न रूटिंग संतुलित विशेषज्ञ उपयोग संकेत प्रदान करता है।
• शोर इंजेक्शन (Noise Injection): विशेषज्ञ पतन (expert collapse) को रोकने और अन्वेषण को प्रोत्साहित करने के लिए, कॉस्ट मैट्रिक्स में गॉसियन शोर जोड़ा जाता है ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$)। सैद्धांतिक विश्लेषण (प्रपोजिशन 4.3) से पता चलता है कि यह सुनिश्चित करता है कि प्रत्येक विशेषज्ञ के पास चयन की गैर-शून्य (non-zero) संभावना हो।

इन्फरेंस व्यवहार (Inference Behavior)

पेपर तर्क देता है कि इन्फरेंस के दौरान संतुलित बाधाओं (जैसे, समान कॉलम सम) को लागू करना प्रतिproductive है, क्योंकि यह इनपुट की विशिष्ट अनुकूलता स्कोर की परवाह किए बिना एक समान रूटिंग को मजबूर करता है। इसलिए, इन्फरेंस के दौरान दोनों सिंकहॉर्न रूटिंग और शोर इंजेक्शन को अक्षम (disable) कर दिया जाता है, जिससे इनपुट-निर्भर भविष्यवाणियों को सुनिश्चित करने के लिए नियत (deterministic) सॉफ्टमैक्स रूटिंग पर वापस आ जाता है।

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

1. नवीन रूटिंग फ्रेमवर्क: एक विधि जो संतुलित विशेषज्ञ उपयोग को बढ़ावा देने के लिए एंट्रॉपी-रेगुलराइज्ड ऑप्टिमल ट्रांसपोर्ट के साथ स्टोकेस्टिक शोर इंजेक्शन को एकीकृत करती है, जो किसी भी सहायक बैलेंसिंग लॉस की आवश्यकता को समाप्त करती है।
2. सैद्धांतिक अंतर्दृष्टि: प्रमाण जो दर्शाते हैं कि सिंकहॉर्न-आधारित रूटिंग और शोर इंजेक्शन प्रशिक्षण में मदद करते हैं क्योंकि वे अन्वेषण और संतुलन को प्रोत्साहित करते हैं, लेकिन विकृत असाइनमेंट से बचने के लिए उन्हें इन्फरेंस के दौरान अक्षम किया जाना चाहिए।
3. दक्षता और प्रदर्शन: व्यापक मूल्यांकन दिखाते हैं कि SSR, स्टेट-ऑफ-द-आर्ट बेसलाइन्स की तुलना में प्रशिक्षण दक्षता, सटीकता और इनपुट भ्रष्टाचार (input corruption) के प्रति मजबूती में सुधार करता है।

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

लेखकों ने भाषा मॉडलिंग (WikiText-103, Enwik-8) और इमेज क्लासिफिकेशन (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R) पर SSR का मूल्यांकन किया।

• भाषा मॉडलिंग (Language Modeling):
  • WikiText-103 पर, SSR वेरिएंट्स ने वैनिला SMoE और लोड-बैलेंसिंग लॉस, z-loss, या नॉइज़ का उपयोग करने वाले बेसलाइन्स को लगातार पछाड़ दिया।
  • SSR-L w/ noise ने 34.367 का टेस्ट परप्लेक्सिटी (PPL) प्राप्त किया, जो वैनिला SMoE (35.550) की तुलना में 1.183 की कमी है, जो अगले सबसे अच्छे बेसलाइन से काफी आगे है।
  • महत्वपूर्ण रूप से, SSR में न्यूनतम प्रशिक्षण ओवरहेड (0.33% – 0.65%) हुआ, जबकि पूर्ण सिंकहॉर्न-आधारित SMoE में 72.47% ओवरहेड हुआ।
  • मोमेंटम सेटिंग्स के तहत, SSR ने वैनिला SMoE के विशिष्ट नॉर्म इवोल्यूशन को संरक्षित करते हुए प्रदर्शन में सुधार किया, जबकि अन्य बैलेंसिंग विधियों ने प्रदर्शन को कम कर दिया।
• बाइट-लेवल मॉडलिंग (Enwik-8):
  • SSR-S w/ noise ने 1.128 का सर्वश्रेष्ठ टेस्ट बिट्स-पर-कैरेक्टर (BPC) प्राप्त किया, जो वैनिला SMोE से 0.010 BPC बेहतर है।
  • यह पूर्ण सिंकहॉर्न-आधारित SMoE की तुलना में 2.28x तेज़ था।
• विज़न टास्क (Vision Tasks):
  • ImageNet-1K पर, SSR-L w/ noise ने 77.420 की टॉप-1 सटीकता प्राप्त की, जो वैनिला SMoE पर +2.368 की वृद्धि है।
  • इसने ImageNet-O और ImageNet-R पर भी मजबूत मजबूती (robustness) दिखाई।
• एब्लेशन स्टडीज (Ablation Studies):
  • $p$ की संभावना (सिंकहॉर्न रूटिंग लागू करने की) महत्वपूर्ण है; बहुत छोटे मान (जैसे, $10^{-3}$ या $10^{-4}$) राउटर को अत्यधिक बाधित किए बिना संतुलित सिग्नल प्रदान करने के लिए पर्याप्त थे।
  • शोर इंजेक्शन ने विभिन्न सेटिंग्स में लगातार प्रदर्शन में सुधार किया।
  • इन्फरेंस में बैलेंसिंग मैकेनिज्म को अक्षम करने से सर्वोत्तम परिणाम मिले, जो सैद्धांतिक दावों की पुष्टि करते हैं।

महत्व और दावे (Significance and Claims)

पेपर का दावा है कि सिलेक्टिव सिंकहॉर्न रूटिंग SMoE डिज़ाइन के लिए एक व्यावहारिक और प्रभावी समाधान प्रदान करता है। जटिल सहायक लॉस के स्थान पर कुशल, रुक-रुक कर होने वाले सिंकहॉर्न-आधारित अनुकूलन का उपयोग करके, SSR प्राप्त करता है:

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

लेखक इस बात पर जोर देते हैं कि यह विधि लाइटवेट है, इसमें कोई अतिरिक्त ट्रेन करने योग्य पैरामीटर नहीं है (शोर वाले संस्करण में), और यह संतुलित उपयोग सुनिश्चित करते हुए विशेषज्ञ चयन के लचीलेपन को बनाए रखती है।