Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

यहाँ शोध पत्र "Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing" का विस्तृत तकनीकी सारांश दिया गया है।

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

Mixture-of-Experts (MoE) आर्किटेक्चर पैरामीटर गणना और कम्प्यूटेशनल लागत को अलग करके भाषा मॉडलों को कुशलतापूर्वक स्केल करने के लिए महत्वपूर्ण हैं। हालाँकि, मौजूदा रूटिंग तंत्र लोड बैलेंसिंग (load balancing) और कौज़ैलिटी (causality) (यह आवश्यकता कि अनुमान संबंधी निर्णय केवल पिछले टोकन पर निर्भर होने चाहिए) के बीच एक मौलिक समझौते (trade-off) का सामना करते हैं:

• टोकन चॉइस (TC) रूटिंग: प्रत्येक टोकन को एक निश्चित संख्या में टॉप-$G$ विशेषज्ञों (experts) तक रूट करता है। यह कौज़ल (causal) तो है, लेकिन लोड असंतुलन के साथ संघर्ष करता है, जिसे अक्सर "रूटिंग कोलैप्स" (जहाँ कुछ विशेषज्ञों का अत्यधिक उपयोग होता है) को रोकने के लिए सहायक लॉस (auxiliary losses) या PID कंट्रोलर्स की आवश्यकता होती है। इसमें डायनेमिक कम्प्यूटेशन एलोकेशन (प्रत्येक टोकन के लिए निश्चित कंप्यूट) का भी अभाव है।
• एक्सपर्ट चॉइस (EC) रूटिंग: प्रत्येक विशेषज्ञ एक बैच से अपने शीर्ष-$k$ टोकनों का चयन करता है। यह सटीक लोड बैलेंसिंग प्राप्त करता है और डायनेमिक कम्प्यूटेशन की अनुमति देता है (टोकन को 0, 1, या कई विशेषज्ञों को रूट किया जा सकता है)। हालाँकि, यह ऑटोरेग्रेसिव जनरेशन के लिए नॉन-कौज़ल (non-causal) है क्योंकि बैच के भीतर शीर्ष-$k$ टोकनों का चयन करने के लिए पूरे बैच तक पहुँच की आवश्यकता होती है, जिसमें भविष्य के टोकन शामिल होते हैं जो इन्फरेंस (inference) के दौरान मौजूद नहीं होते। यह एक गंभीर 'ट्रेन-इन्फरेंस मिसमैच' (train-inference mismatch) पैदा करता है।

यह शोध पत्र इस तनाव को हल करने के लिए एक ऐसा रूटिंग तंत्र बनाने का लक्ष्य रखता है जो डायनेमिक कम्प्यूटेशन एलोकेशन और लोड बैलेंसिंग का समर्थन करता है, जबकि ऑटोरेग्रेसिव लैंग्वेज मॉडलिंग के लिए पूर्णतः कौज़ल (fully causal) बना रहता है।

2. कार्यप्रणाली: एक्सपर्ट थ्रेशोल्ड (ET) रूटिंग

लेखक एक्सपर्ट थ्रेशोल्ड (ET) रूटिंग का प्रस्ताव करते हैं, जो वैश्विक टोकन वितरण (global token distribution) पर एक स्टोकेस्टिक एक्सपेक्टेशन (stochastic expectation) के पक्ष में EC की सख्त प्रति-बैच बाधाओं को शिथिल करता है।

मुख्य तंत्र (Core Mechanism)

एक विशिष्ट बैच के भीतर शीर्ष-$k$ टोकनों को चुनने के बजाय, ET प्रत्येक विशेषज्ञ $i$ के लिए एक एक्सपोनेंशियल मूविंग एवरेज (EMA) थ्रेशोल्ड ($c_i$) बनाए रखता है। यह थ्रेशोल्ड सक्रियण के लिए आवश्यक राउटर स्कोर के वैश्विक क्वांटाइल (global quantile) का प्रतिनिधित्व करता है।

• रूटिंग नियम: किसी भी टोकन $t$ और विशेषज्ञ $i$ के लिए, टोकन को तब रूट किया जाता है यदि उसका राउटर स्कोर $r_{t,i}$ विशेषज्ञ के थ्रेशोल्ड से अधिक हो:
  $$z_{t,i} = \mathbb{1}\{r_{t,i} > c_i\}$$
• थ्रेशोल्ड अपडेट: प्रशिक्षण के दौरान, थ्रेशोल्ड $c_i$ को बैच में देखे गए $k$-वें सबसे बड़े राउटर लॉजिट (logit) के EMA का उपयोग करके अपडेट किया जाता है (जहाँ $k = N/E$ लक्षित क्षमता है)।
  $$c_i \leftarrow \beta c_i + (1 - \beta) \cdot \text{k-th-largest}(\{r_{t,i}\}_{t=1}^N)$$
• कौज़ैलिटी (Causality): चूंकि टोकन $t$ के लिए निर्णय केवल उसके वर्तमान स्कोर $r_{t,i}$ और ऐतिहासिक वैश्विक थ्रेशोल्ड $c_i$ पर निर्भर करता है, इसलिए यह तंत्र पूर्णतः कौज़ल है। इसके लिए भविष्य के टोकन या वर्तमान बैच की संरचना के ज्ञान की आवश्यकता नहीं है।

वॉर्मअप रणनीति (Warmup Strategy)

प्रशिक्षण की शुरुआत में, राउटर लॉजिट अस्थिर होते हैं, और EMA थ्रेशोल्ड अभिसरित (converge) नहीं हुआ होता है। इससे "एक्सपर्ट स्टार्वेशन" (जहाँ अधिकांश टोकन थ्रेशोल्ड को पार करने में विफल रहते हैं) हो सकता है। इसे संबोधित करने के लिए, लेखक एक वॉर्मअप चरण का उपयोग करते हैं:

1. पहले 4,000 स्टेप्स के लिए मानक EC रूटिंग (top-$k$ चयन) का उपयोग करें।
2. नियंत्रित लोड बैलेंसिंग के तहत EMA थ्रेशोल्ड को स्थिर होने दें।
3. एक बार जब थ्रेशोल्ड सार्थक हो जाए, तो ET रूटिंग पर स्विच करें।

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

1. कौज़ल डायनेमिक रूटिंग: ET पहला ऐसा तंत्र है जो एक्सपर्ट चॉइस (डायनेमिक कम्प्यूटेशन और अपेक्षा में पूर्ण लोड बैलेंसिंग) के लाभों को प्राप्त करता है, जबकि पूर्ण कौज़ैलिटी बनाए रखता है, जिससे यह ऑटोरेग्रेसिव LLMs के लिए उपयुक्त बनता है।
2. सहायक लॉस का उन्मूलन: TC रूटिंग के विपरीत, ET स्वाभाविक रूप से थ्रेशोल्ड तंत्र के माध्यम से लोड बैलेंसिंग प्राप्त करता है, जिससे उन सहायक लोड-बैलेंसिंग लॉस की आवश्यकता समाप्त हो जाती है जो राउटर लॉजिट्स को बायस (bias) कर सकते हैं।
3. ट्रेन-इन्फरेंस निरंतरता: क्योंकि प्रशिक्षण और इन्फरेंस दोनों के लिए समान थ्रेशोल्ड-आधारित तर्क का उपयोग किया जाता है, ET उन अंतरालों को समाप्त करता है जो EC मॉडल्स में निहित होते हैं (जिन्हें अक्सर जटिल प्रेडिक्टर्स या रिट्रेनिंग की आवश्यकता होती है)।
4. पॉपुलेशन-लेवल एप्रोक्सिमेशन: सैद्धांतिक रूप से, ET एक्सपर्ट चॉइस रूटिंग के रूप में उस सीमा (limit) का अनुमान लगाता है जहाँ बैच का आकार अनंत हो जाता है, जहाँ चयन थ्रेशोल्ड विशिष्ट बैच संरचना से स्वतंत्र हो जाता है।

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

लेखकों ने 2.4B पैरामीटर्स (0.56B एक्टिव) तक स्केल किए गए मॉडलों का उपयोग करके FineWeb-Edu पर ET का मूल्यांकन किया।

• टोकन चॉइस (TC) बनाम प्रदर्शन:
  • ET ने TC रूटिंग की तुलना में 0.067 कम क्रॉस-एन्ट्रॉपी (CE) लॉस प्राप्त किया।
  • यह सुधार 1.6$\times$ कम टोकन के साथ समान प्रदर्शन स्तर तक पहुँचने के बराबर है।
  • ET ने CORE बेंचमार्क (इन-कॉन्टेक्स्ट लर्निंग कार्यों का एक समूह) पर भी TC से काफी बेहतर प्रदर्शन किया।
• एक्सपर्ट चॉइस (EC) बनाम प्रदर्शन:
  • ET ने बहुत बड़े बैच साइज (512k टोकन) के साथ प्रशिक्षित EC के प्रदर्शन की बराबरी की।
  • EC का प्रदर्शन 64k–512k के बैच साइज के आसपास संतृप्त (saturate) हो जाता है; ET बिना बड़े-बैच समन्वय के इस "लार्ज-बैच" प्रदर्शन को प्राप्त करता है।
  • ट्रेन-इवैल्यूएशन गैप: छोटे बैच (जैसे 2k) के साथ EC ने नॉन-कौज़ल टॉप-$k$ चयन के कारण प्रशिक्षण और मूल्यांकन लॉस के बीच एक महत्वपूर्ण अंतर दिखाया। ET ने न्यूनतम अंतर दिखाया, जो इसकी स्थिरता की पुष्टि करता है।
• विशेषज्ञ विशेषज्ञता (Expert Specialization):
  • विश्लेषण से पता चला कि ET ने बड़े-बैच EC के समान विशेषज्ञ विशेषज्ञता पैटर्न विकसित किए (जैसे, कोड बनाम गणित को संभालने वाले विशिष्ट विशेषज्ञ), जो दर्शाता है कि पॉपुलेशन-लेवल थ्रेशोल्ड डोमेन-विशिष्ट रूटिंग संरचनाओं को सफलतापूर्वक कैप्चर करता है।
• डायनेमिक कम्प्यूटेशन:
  • ET ने कठिन टोकन (जैसे गणित की समस्याओं में नंबर, कोड में कीवर्ड) को अधिक और आसान टोकन को कम कंप्यूट आवंटित करने में सफलता प्राप्त की, जो प्रभावी डायनेमिक रिसोर्स एलोकेशन को प्रदर्शित करता है।

5. महत्व (Significance)

यह शोध पत्र MoE अनुसंधान में एक महत्वपूर्ण अंतर को पाटता है। वर्षों से, एक्सपर्ट चॉइस रूटिंग को लोड बैलेंसिंग और डायनेमिक कम्प्यूटेशन के लिए सैद्धांतिक रूप से श्रेष्ठ माना जाता था, लेकिन कौज़ैलिटी उल्लंघन के कारण ऑटोरेग्रेसिव जनरेशन के लिए व्यावहारिक रूप से अनुपयोगी माना जाता था। टोकन चॉइस ऑटोरेग्रेसिव मॉडल्स के लिए मानक था लेकिन लोड असंतुलन और स्टेटिक कंप्यूट एलोकेशन से ग्रस्त था।

एक्सपर्ट थ्रेशोल्ड (ET) इसे बैच-लेवल ऑप्टिमाइज़ेशन से पॉपुलेशन-लेवल एस्टीमेशन की ओर शिफ्ट करके हल करता है। EMA-ट्रैक किए गए थ्रेशोल्ड का उपयोग करके, ET मॉडलों को अनुमति देता है:

• सहायक लॉस के बिना कुशलतापूर्वक स्केल करना।
• टोकन की कठिनाई के आधार पर कंप्यूट को डायनेमिक रूप से आवंटित करना।
• अपेक्षा में पूर्ण लोड बैलेंसिंग बनाए रखना।
• पूर्णतः कौज़ली संचालित होना, जिससे बिना किसी आर्किटेक्चरल बदलाव या रिट्रेनिंग के निर्बाध प्रशिक्षण और इन्फरेंस संभव होता है।

यह तंत्र स्केलेबल, कुशल और उच्च-प्रदर्शन वाले ऑटोरेग्रेसिव लैंग्वेज मॉडल्स के निर्माण के लिए नए मार्ग खोलता है।