ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization

यहाँ शोध पत्र "ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization" का विस्तृत तकनीकी सारांश दिया गया है।

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

Mixture-of-Experts (MoE) मॉडल प्रत्येक इनपुट के लिए केवल मापदंडों के एक विरल (sparse) उपसमुच्चय को सक्रिय करके डीप लर्निंग को स्केल करने के लिए एक मानक बन गए हैं। हालाँकि, पारंपरिक MoE आर्किटेक्चर तीन महत्वपूर्ण संरचनात्मक दोषों से ग्रस्त हैं:

• रूटर-एक्सपर्ट मिसअलाइनमेंट (Router-Expert Misalignment): मानक रूटर सीखे गए गेटिंग लॉजिट्स (gating logits) का उपयोग करते हैं जो अक्सर विशेषज्ञों के वास्तविक प्रतिनिधित्व स्थान (representational space) के साथ गलत संरेखित होते हैं। इससे टोकन उप-इष्टतम विशेषज्ञों को रूट किए जाते हैं, जिससे प्रतिनिधित्व पतन (representation collapse - जहाँ विशेषज्ञ रेडंडेंट हो जाते हैं) और अस्थिर रूटिंग (प्रशिक्षण और अनुमान के बीच उतार-चढ़ाव) होता है।
• लोड इम्बैलेंस (Load Imbalance): टोकन-चॉइस रूटिंग एक "रिच-गेट-रिचर" (rich-get-richer) गतिशीलता पैदा कर सकती है जहाँ कुछ विशेषज्ञ ओवरलोडेड (स्ट्रैग्लर्स) हो जाते हैं जबकि अन्य कम उपयोग किए जाते हैं।
• वर्तमान सुधारों की सीमाएँ (Limitations of Current Fixes): मुख्यधारा के समाधान एकरूपता को मजबूर करने के लिए सहायक लोड-बैलेंसिंग लॉस (Auxiliary Load-Balancing Losses - LBL) पर निर्भर करते हैं। हालाँकि, LBL ऐसे इंटरफेरेंस ग्रेडिएंट्स पेश करता है जो विशेषज्ञ विशिष्टता (expert specialization) को कमजोर करते हैं, डाउनस्ट्रीम सटीकता को कम करते हैं, और मिसअलाइनमेंट के मूल कारण को संबोधित करने में विफल रहते हैं।

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

लेखक ERMoE (Eigen-Reparameterized Mixture-of-Experts) का प्रस्ताव करते हैं, जो एक नवीन आर्किटेक्चर है जो यह मौलिक रूप से बदल देता है कि रूटिंग निर्णय कैसे लिए जाते हैं, उन्हें सीधे विशेषज्ञों की आंतरिक ज्यामिति (internal geometry) से जोड़कर।

A. आइगेन-रिपैरामीटरकृत विशेषज्ञ (Eigen-Reparameterized Experts)

मानक डेंस वेट मैट्रिसेस के बजाय, प्रत्येक विशेषज्ञ $e$ को एक सीखे गए ऑर्थोनॉर्मल आइगेनबेसिस (orthonormal eigenbasis) में पैरामीटराइज किया जाता है:
$$\mathbf{W}^{(e)} = \mathbf{U}^{(e)} \,\mathrm{diag}(s^{(e)})\, \mathbf{V}^{(e)\top}$$
जहाँ $\mathbf{U}^{(e)}$ और $\mathbf{V}^{(e)}$ ऑर्थोनॉर्मल बेस हैं, और $s^{(e)}$ ट्रेन करने योग्य गुणांक (coefficients) हैं। यह विशेषज्ञ दिशाओं के बीच लगभग-ऑर्थोगोनैलिटी (near-orthogonality) को लागू करता है, जिससे फीचर रेडंडेंसी कम होती है और विशेषज्ञ सबस्पेस ज्यामितीय रूप से विभाज्य (geometrically separable) हो जाते हैं।

B. आइगेनबेसिस स्कोर रूटिंग (Eigenbasis Score Routing - Content-Aware)

ERMoE मानक सीखे गए गेटिंग लॉजिट्स के स्थान पर एक ज्यामितीय संरेखण स्कोर (geometric alignment score) का उपयोग करता है:

1. प्रोजेक्शन (Projection): एक दिए गए टोकन $x_i$ और उसके अटेंशन-वेटेड कॉन्टेक्स्ट $c_i$ के लिए, मॉडल दोनों वेक्टर्स को प्रत्येक विशेषज्ञ $e$ के आइगेनबेसिस में प्रोजेक्ट करता है।
2. स्कोरिंग (Scoring): रूटिंग स्कोर उस विशिष्ट विशेषज्ञ के बेस के भीतर प्रोजेक्ट किए गए टोकन और प्रोजेक्ट किए गए कॉन्टेक्स्ट के बीच कोसाइन समानता (cosine similarity) है:
   $$\mathrm{Score}_e(i) = \cos\left( \mathbf{U}^{(e)\top} \tilde{x}_i, \, \mathbf{U}^{(e)\top} \tilde{c}_i \right)$$
3. थ्रेशोल्डेड टॉप-k चयन (Thresholded Top-k Selection):
   • एक कॉन्फिडेंस थ्रेशोल्ड $T$ कम-आत्मविश्वास वाले संरेखणों को फ़िल्टर करता है।
   • रूटर $T$ से ऊपर उठने वाले शीर्ष-$k$ विशेषज्ञों का चयन करता है।
   • यदि $k$ से कम विशेषज्ञ थ्रेशोल्ड पार करते हैं, तो एक फॉलबैक मैकेनिज्म वैश्विक रूप से उच्चतम-स्कोरिंग विशेषज्ञों का चयन करता है।
   • मुख्य लाभ: यह सहायक लोड-बैलेंसिंग लॉस की आवश्यकता को समाप्त करता है। संतुलन स्वाभाविक रूप से प्राप्त होता है क्योंकि टोकन केवल उन्हीं विशेषज्ञों को रूट किए जाते हैं जिनके प्रतिनिधित्व सबस्पेस इनपुट कंटेंट के साथ ज्यामितीय रूप से मेल खाते हैं।

C. 3D अनुकूलन (ERMoE-ba)

न्यूरोइमेजिंग के लिए, लेखकों ने 3D ViT बैकबोन का उपयोग करते हुए एक 3D संस्करण ERMoE-ba बनाया है। यह एनाटोमिकली व्याख्या योग्य विशिष्टता (anatomically interpretable specialization) को सीखने की अनुमति देने के लिए विशेष रीजन एक्सपर्ट्स (Region Experts) (जो विशेष रूप से व्हाइट मैटर, ग्रे मैटर और CSF पर प्रशिक्षित हैं) के साथ फ्री एक्सपर्ट्स (Free Experts) को पेश करता है।

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

1. नवीन आर्किटेक्चर: ERMoE का परिचय, जो विशेषज्ञों को ऑर्थोनॉर्मल आइगेनबेसिस में पुन: पैरामीटराइज करता है और फीचर-बेसिस अलाइनमेंट के माध्यम से रूटिंग करता है न कि मुक्त लॉजिट्स के माध्यम से।
2. सहायक लॉस का उन्मूलन: डिज़ाइन उन मुख्य संरचनात्मक दोषों (मिसअलाइनमेंट और कोलैप्स) को हल करता है जो मानक MoE की विशेषता हैं, बिना किसी लोड-बैलेंसिंग लॉस (LBL) के, जिससे इंटरफेरेंस ग्रेडिएंट्स और अत्यधिक एकरूप रूटिंग से बचा जा सके।
3. व्याख्यात्मकता और विशिष्टता (Interpretability & Specialization): रूटिंग को प्रतिनिधित्व सबस्पेस से जोड़कर, मॉडल स्थिर विशिष्टता प्राप्त करता है। 3D डोमेन में, इसके परिणामस्वरूप विशेषज्ञ स्वाभाविक रूप से विशिष्ट मस्तिष्क ऊतकों (जैसे, WM, GM, CSF) में विशेषज्ञता प्राप्त करते हैं।
4. स्टेट-ऑफ-द-आर्ट प्रदर्शन: विधि प्राकृतिक छवि वर्गीकरण, इमेज-टेक्स्ट रिट्रीवल और 3D ब्रेन एज प्रेडिक्शन में SOTA परिणाम प्राप्त करती है।

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

A. प्राकृतिक छवि वर्गीकरण (Natural Image Classification - ImageNet, CIFAR, Tiny-ImageNet)

• ImageNet-1K: ERMoE ने 88.03% टॉप-1 सटीकता प्राप्त की, जो मजबूत V-MoE बेसलाइन (87.41%) और अन्य MoE वेरिएंट्स (DeepMoE, Single-gated MoE) को काफी पीछे छोड़ते हुए उत्कृष्ट प्रदर्शन किया।
• फ्यू-शॉट लर्निंग (Few-Shot Learning): CIFAR-100 (5-शॉट) पर, ERMoE ने 93.47% प्राप्त किया, जो V-MoE (89.49%) से लगभग 4% बेहतर है। यह दर्शाता है कि ERMoE अधिक रैखिक रूप से विभाज्य (linearly separable) और कम रेडंडेंट फीचर्स सीखता है।
• लोड बैलेंस (Load Balance): V-MoE के विपरीत, जो LBL के बावजूद विशेषज्ञों के उपयोग में भारी-पूंछ वितरण (heavy-tailed distribution) दिखाता है, ERMoE बिना किसी बैलेंसिंग लॉस के एक सपाट, लगभग-समान उपयोग वक्र (flat, near-uniform utilization curve) प्रदर्शित करता है।

B. इमेज-टेक्स्ट रिट्रीवल (Image-Text Retrieval - COCO, Flickr30K)

• CLIP फ्रेमवर्क में एकीकृत होकर (CLIP-ERMoE), मॉडल ने COCO (65.4% R@1) और Flickr30K (63.4% R@1) पर SOTA इमेज-टू-टेक्स्ट (I2T) रिट्रीवल हासिल किया, जो समर्पित मल्टीमॉडल MoE मॉडल जैसे CLIP-MoE से बेहतर है।
• यह प्रदर्शित करता है कि विविध विशेषज्ञ प्रतिनिधित्व मल्टीमॉडल अलाइनमेंट कार्यों में अच्छी तरह से सामान्यीकृत (generalize) होते हैं।

C. 3D ब्रेन MRI (ADNI Dataset)

• ब्रेन एज प्रेडिक्शन (Brain Age Prediction): ERMoE-ba ने ADNI टेस्ट सेट पर 2.31 वर्ष का मीन एब्सोल्यूट एरर (MAE) प्राप्त किया।
• तुलना: यह 3D CNNs (3.13), 3D ViTs (3.52) और 3D Swin Transformers (2.83) की तुलना में एक महत्वपूर्ण सुधार है।
• व्याख्यात्मकता (Interpretability): विश्लेषण से पता चला कि जैसे-जैसे प्रशिक्षण आगे बढ़ा, रूटर ने विशिष्ट इनपुट्स को एनाटॉमिकली प्रासंगिक विशेषज्ञों को सौंपना सीखा (जैसे, CSF इनपुट्स को 0.91 के स्कोर के साथ CSF विशेषज्ञ को रूट किया गया), जो एनाटॉमिकली सार्थक विशिष्टता (anatomically meaningful specialization) की पुष्टि करता है।

5. महत्व और प्रभाव (Significance and Impact)

• सैद्धांतिक बदलाव (Theoretical Shift): ERMoE, MoE रूटिंग के प्रतिमान को "रूट करने के लिए सीखना" (लॉजिट्स के माध्यम से) से "ज्यामिति द्वारा रूटिंग" (सबस्पेस अलाइनमेंट के माध्यम से) में बदल देता है। यह उस स्थिरता-विशिष्टता ट्रेड-ऑफ को हल करता है जिसने MoE अनुसंधान को परेशान किया है।
• दक्षता (Efficiency): यह विधि V-MoE की तुलना में कम्प्यूटेशनल ओवरहेड (FLOPs में 10.3% की गिरावट) और सक्रिय मापदंडों को कम करते हुए सटीकता में सुधार करती है।
• नैदानिक अनुप्रयोग (Clinical Applicability): ब्रेन एज प्रेडिक्शन में ERMoE-ba की सफलता यह दर्शाती है कि स्पार्स, व्याख्यात्मक MoE मॉडल उच्च-आयामी 3D मेडिकल डेटा को संभाल सकते हैं, जो भारी कम्प्यूटेशनल संसाधनों या जटिल बैलेंसिंग ह्यूरिस्टिक्स की आवश्यकता के बिना कुशल, नैदानिक रूप से प्रासंगिक AI टूल बनाने का मार्ग प्रदान करते हैं।
• मजबूती (Robustness): सहायक लॉस को हटाना पाइपलाइन को सरल बनाता है, जिससे MoE मॉडल को तैनात करना आसान हो जाता है और यह बैलेंसिंग कोएफिशिएंट्स के संबंध में हाइपरपैरामीटर ट्यूनिंग के प्रति कम संवेदनशील हो जाता है।