🤖 machine learning

N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation

यह शोध पत्र N-vium को प्रस्तुत करता है, जो एक मिश्रण-एग्जिट्स (mixture-of-exits) ट्रांसफार्मर है जो कई गहराइयों से भविष्यवाणियों को मिलाने के लिए टोकन-अनुकूली रूटिंग (token-adaptive routing) का उपयोग करने और ऊपरी परतों की गणना को स्थगित करने के माध्यम से मॉडल की गुणवत्ता से समझौता किए बिना मानक ट्रांसफार्मरों पर 57.9% तक वॉल-क्लॉक स्पीडअप प्राप्त करता है।

मूल लेखक: Aleksander Lorenc, Frédéric Berdoz, Joël Mathys, Roger Wattenhofer

प्रकाशित 2026-05-14

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

CC BY 4.0

मूल लेखक: Aleksander Lorenc, Frédéric Berdoz, Joël Mathys, Roger Wattenhofer

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप एक लंबी, जटिल किताब पढ़ रहे हैं, और आपको एक वाक्य में अगले शब्द का अनुमान लगाना है।

एक मानक AI मॉडल (एक "ट्रांसफॉर्मर") में, प्रक्रिया एक फैक्ट्री असेंबली लाइन की तरह है जिसमें 48 श्रमिकों की एक एकल पंक्ति में खड़ी है। आपके द्वारा टाइप किया गया हर एक शब्द को हर एक श्रमिक के पास से गुजरना होगा, एक-एक करके, प्रत्येक स्टेशन पर थोड़ा सा प्रसंस्करण (processing) प्राप्त करते हुए। केवल तब जब शब्द ने सभी 48 श्रमिकों के पास अपनी यात्रा पूरी कर ली होती है, अंतिम श्रमिक भविष्यवाणी चिल्लाकर बताता है।

समस्या क्या है? कई शब्द सरल होते हैं। "the" या "cat" जैसे शब्द को यह समझने के लिए कि आगे क्या आना चाहिए, 48 श्रमिकों की आवश्यकता नहीं होती। लेकिन मानक मॉडल उन्हें सुरक्षा के लिए फिर भी सभी के पास जाने के लिए मजबूर करता है। यह धीमा है और ऊर्जा बर्बाद करता है।

पुराना समाधान: "अर्ली एग्जिट" (शॉर्टकट)

पिछले तरीकों ने इसे ठीक करने की कोशिश की जिसमें सरल शब्दों को लाइन से जल्दी "बाहर निकलने" दिया जाता था। यदि कोई श्रमिक सोचता है, "मैं यह समझ गया," तो वह शब्द को बाहर जाने देता है।

दोष: यह बिना अंतिम बॉस की जांच किए उत्तर का अनुमान लगाने वाले श्रमिक जैसा था। कभी-कभी वे गलत अनुमान लगाते थे। साथ ही, क्योंकि शब्द जल्दी बाहर निकल गया, फैक्ट्री भविष्य के शब्दों के लिए आवश्यक "स्मृति" (KV कैश) खो देती थी, जिससे सिस्टम को चीजों को नकली बनाने या पुन: गणना करने के लिए मजबूर होना पड़ता था, जिससे उत्तर की गुणवत्ता खराब हो जाती थी।

नया समाधान: N-vium (स्मार्ट मिश्रण)

यह शोध पत्र N-vium पेश करता है, जो खेल के नियमों को पूरी तरह से बदल देता है। एक एकल लाइन के बजाय, कल्पना करें कि फैक्ट्री के पास लाइन के विभिन्न बिंदुओं पर चार अलग-अलग निकास (exits) हैं।

N-vium कैसे काम करता है, इसके कुछ उपमाओं (analogies) का उपयोग यहाँ दिया गया है:

1. "स्मार्ट ट्रैफिक लाइट" (सीखा हुआ रूटिंग)

प्रत्येक चार निकासों पर, एक स्मार्ट ट्रैफिक लाइट (एक राउटर) होती है। यह शब्द को देखती है और निर्णय लेती है:

"यह शब्द सरल है। चलिए इसे निकास 1 पर बाहर जाने देते हैं।"
"यह शब्द कठिन है। इसे आगे बढ़ने दें।"
"यह शब्द बहुत जटिल है। इसे निकास 4 तक भेजें।"

महत्वपूर्ण बात यह है कि AI अपने प्रशिक्षण के दौरान इन लाइटों को चलाना सीखता है। यह केवल अनुमान नहीं लगाता; यह जानता है कि कौन सा निकास उस विशिष्ट शब्द के लिए सबसे अच्छा उत्तर देगा।

2. "परफेक्ट ब्लेंड" (सटीक मिश्रण)

पुराने "अर्ली एग्जिट" तरीकों में, AI को एक निकास चुनना होता था और उम्मीद करनी होती थी कि वह सही है। N-vium में, AI सभी संभावित उत्तरों का एक परफेक्ट स्मूदी बनाता है।

यह निकास 1 से उत्तर का थोड़ा सा हिस्सा लेता है।
यह निकास 3 से थोड़ा सा मिश्रण करता है।
यह निकास 4 का एक चुटकी हिस्सा जोड़ता है।
परिणाम: अंतिम उत्तर गणितीय रूप से सटीक है। यह उतना ही अच्छा है जितना कि शब्द ने सभी 48 श्रमिकों के पास जाने के बाद होता, लेकिन यह तेज़ था क्योंकि "आसान" गणनाओं को शुरुआती निकासों द्वारा संभाल लिया गया था।

3. "पिगीबैक" ट्रिक (याददाश्त का नुकसान नहीं)

यह इस शोध पत्र का सबसे बड़ा जादू है। आमतौर पर, यदि कोई शब्द जल्दी बाहर निकल जाता है, तो फैक्ट्री अगले शब्द के लिए आवश्यक कार्य भूल जाती है।

N-vium का समाधान: जब कोई शब्द जल्दी बाहर निकलता है, तो फैक्ट्री काम को फेंकती नहीं है। इसके बजाय, यह कहती है, "ठीक है, हम इस शब्द के लिए बाकी प्रोसेसिंग बाद में पूरी करेंगे।"
जब अगला शब्द लाइन में आता है, तो फैक्ट्री दो काम एक साथ करती है: यह नए शब्द को प्रोसेस करती है, और पिछले शब्द के अधूरे काम को पिगीबैक (piggyback) करती है।
उपमा: एक बस ड्राइवर की कल्पना करें। आमतौर पर, ड्राइवर यात्रियों को छोड़ने के लिए हर स्टॉप पर रुकता है। N-vium के साथ, ड्राइवर आसान यात्रियों को जल्दी उतार देता है, लेकिन फिर पिछले स्टॉप से "अधूरे" यात्रियों को उठा लेता है और उन्हें अगले स्टॉप पर जाने के रास्ते में छोड़ देता है, और वह भी बिना बस को रोके या धीमा किए।

परिणाम

शोधकर्ताओं ने 1.5 बिलियन "श्रमिकों" (पैरामीटर्स) तक के मॉडल बनाए।

गति: उनका सबसे बड़ा मॉडल उसी आकार के मानक मॉडल की तुलना में 57.9% तेज़ था।
गुणवत्ता: तेज़ होने के बावजूद, उत्तर उतने ही अच्छे थे (गुणवत्ता में कोई कमी नहीं)।
हार्डवेयर: यह विशेष नए हार्डवेयर की आवश्यकता के बिना मानक कंप्यूटर चिप्स (GPUs) पर काम करता है।

सारांश

N-vium को एक ऐसे शॉर्टकट के रूप में न सोचें जो कोनों को काटता है, बल्कि एक स्मार्ट पुनर्गठन के रूप में देखें। यह महसूस करता है कि हर शब्द को पूरे फैक्ट्री टूर की आवश्यकता नहीं होती है। यह सरल शब्दों को जल्दी बाहर जाने देता है, उनके उत्तरों को जटिलों के साथ पूरी तरह से मिलाता है, और यह सुनिश्चित करने के लिए "पिगीबैकिंग" का उपयोग करता है कि कोई भी काम बर्बाद न हो। परिणाम एक ऐसा AI है जो कम सोचे बिना तेज़ी से सोचता है।

तकनीकी सारांश: N-vium – त्वरित सटीक पीढ़ी (Exact Generation) के लिए Mixture-of-Exits ट्रांसफॉर्मर

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

मानक ऑटो-रिग्रेसिव लार्ज लैंग्वेज मॉडल्स (LLMs) टोकन को डिकोडर परतों के एक स्टैक के माध्यम से क्रमिक रूप से प्रोसेस करते हैं। यह फॉरवर्ड पास दो जुड़े हुए कार्यों को पूरा करता है: (1) भविष्य के अटेंशन के लिए प्रत्येक परत पर की-वैल्यू (KV) रिप्रजेंटेशन बनाना, और (2) अगले टोकन की भविष्यवाणी करने के लिए अंतिम हिडन स्टेट को प्रोजेक्ट करना। इसके परिणामस्वरूप, प्रत्येक टोकन को सभी परतों के माध्यम से एक पूर्ण फॉरवर्ड पास की आवश्यकता होती है, भले ही प्रेडिक्शन डिस्ट्रीब्यूशन पहले ही अभिसरित (converge) क्यों न हो गया हो।

इन्फरेंस दक्षता में सुधार के लिए मौजूदा दृष्टिकोण, जैसे कि अर्ली एग्जिट (Early Exit) विधियाँ, इन कार्यों को मध्यवर्ती परतों पर प्रेडिक्शन हेड्स लगाकर अलग करने का प्रयास करती हैं। हालाँकि, ये विधियाँ आमतौर पर मध्यवर्ती भविष्यवाणियों को अंतिम परत के वितरण के एक सन्निकटन (approximation) के रूप में मानती हैं। इससे कई कमियाँ उत्पन्न होती हैं:

सन्निकटन त्रुटियाँ (Approximation Errors): मध्यवर्ती भविष्यवाणियाँ अक्सर पूर्ण-डेप्थ वाली भविष्यवाणियों की तुलना में कमजोर होती हैं, जिससे ट्रेन-टेस्ट मिसमैच होता है।
ह्यूरिस्टिक रूटिंग (Heuristic Routing): कब बाहर निकलना (exit) है, इसके निर्णय अक्सर हाथ से बनाए गए कॉन्फिडेंस थ्रेशोल्ड पर निर्भर करते हैं जिन्हें कैलिब्रेशन की आवश्यकता होती है।
अपूर्ण KV कैश (Incomplete KV Caches): कंप्यूट बचाने के लिए, अर्ली-एग्जिट विधियाँ अक्सर बायपास की गई परतों के लिए KV कैश की गणना को छोड़ देती हैं, जिसके लिए अनुमानित कॉपी करने या छोड़ने की आवश्यकता होती है, जो गुणवत्ता को कम कर देता है।
गुणवत्ता में गिरावट (Quality Degradation): गति अक्सर परप्लेक्सिटी (perplexity) और जनरेशन क्वालिटी की कीमत पर प्राप्त की जाती है।

पेपर का तर्क है कि जबकि अनुभवजन्य अवलोकन (जैसे, "लॉगिट लेंस" के माध्यम से) दिखाते हैं कि प्रेडिक्शन डिस्ट्रीब्यूशन अक्सर अंतिम परत से पहले स्थिर हो जाता है, वर्तमान विधियाँ गुणवत्ता हानि के बिना सटीक (exact) जनरेशन के लिए इसका लाभ उठाने में विफल रहती हैं।

2. कार्यप्रणाली: N-vium

लेखक N-vium पेश करते हैं, जो एक 'मिक्सचर-ऑफ-एग्जिट्स' फ्रेमवर्क है जो मानक ट्रांसफॉर्मर आर्किटेक्चर का सख्ती से सामान्यीकरण (generalization) करता है। अर्ली एग्जिट्स को केवल सन्निकटन मानने के बजाय, N-vium मॉडल के आउटपुट को कई एग्जिट्स के ऊपर एक सीखे हुए मिश्रण वितरण (learned mixture distribution) के रूप में पुनर्परिभाषित करता है।

आर्किटेक्चर

मिक्सचर ऑफ एग्जिट्स (Mixture of Exits): मॉडल बैकबोन की $L$ परतों को $N_V$ ब्लॉक्स में विभाजित करता है। प्रत्येक ब्लॉक के अंत में (जंक्शन $k$ ), एक राउटर उस परत के वितरण $\pi_k$ से अगले टोकन को सैंपल करने की संभावना निर्धारित करता है।
टोकन-एडेप्टिव रूटिंग (Token-Adaptive Routing): एक राउटर MLP जंक्शन $k$ पर हिडन स्टेट $h^{(\ell_k)}_t$ को प्रोसेस करता है ताकि एक संभाव्यता $w_k$ आउटपुट की जा सके। जंक्शन $k$ पर एग्जिट करने की संभावना $p_k$ इस प्रकार गणना की जाती है:
$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$
जहाँ कन्वेंशन के अनुसार $w_{N_V} = 1$ है।
मिक्सचर डिस्ट्रीब्यूशन (Mixture Distribution): अंतिम नेक्स्ट-टोकन डिस्ट्रीब्यूशन सभी एग्जिट डिस्ट्रीब्यूशन का एक भारित योग (weighted sum) है:
$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$
शेयर्ड हेड (Shared Head): सभी एग्जिट्स एक ही फाइनल लैंग्वेज मॉडलिंग (LM) प्रोजेक्शन हेड $W_{lm}$ को साझा करते हैं, जिसमें हल्के वेट एडैप्टर MLPs मध्यवर्ती हिडन स्टेट्स को हेड के अपेक्षित ज्योमेट्री से मेल खाने के लिए ट्रांसफॉर्म करते हैं।

ट्रेनिंग

ऑब्जेक्टिव (Objective): मॉडल को व्यक्तिगत हेड्स को अलग से ऑप्टिमाइज़ करने के बजाय, संयुक्त मिश्रण वितरण $\pi_{mix}$ के क्रॉस-एन्ट्रॉपी लॉस को कम करने के लिए प्रशिक्षित किया जाता है।
कंप्यूट पेनल्टी (Compute Penalty): एक सहायक लॉस टर्म $\mathcal{L}_{compute}$ जोड़ा जाता है जो राउटर को शुरुआती एग्जिट्स को संभाव्यता द्रव्यमान (probability mass) असाइन करने के लिए प्रोत्साहित करता है, जिससे मॉडल दक्षता की ओर मुड़ता है:
$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$
जहाँ $\mathcal{L}_{compute}$ अपेक्षित नॉर्मलाइज्ड लेयर डेप्थ है।
वॉर्म-अप (Warm-up): राउटर को "राउटर कोलैप्स" (जहाँ सभी टोकन अंतिम परत की ओर रूट होते हैं) से बचाने के लिए एक राउटर वॉर्म-अप चरण का उपयोग किया जाता है, जो शुरुआत में यूनिफॉर्म एग्जिट डिस्ट्रीब्यूशन से विचलन को दंडित करता है।

इन्फरेंस: पिगीबैकिंग (Piggybacking)

एक महत्वपूर्ण नवाचार पिगीबैकिंग तकनीक है जो यह सुनिश्चित करती है कि बिना कंप्यूट स्किप किए सटीक KV कैश रिकवरी हो सके:

डिफर्ड कंप्यूटेशन (Deferred Computation): जब एक टोकन जंक्शन $k$ पर एग्जिट करता है, तो परतों $k+1$ से $L$ के लिए गणना को स्थगित (defer) कर दिया जाता है।
पैरेलल बैचिंग (Parallel Batching): जैसे-जैसे भविष्य के टोकन प्रोसेस किए जाते हैं, उनके फॉरवर्ड पासेज स्थगित परतों तक पहुँचते हैं। इन डेप्थ्स पर, स्थगित टोकन्स के हिडन स्टेट्स को वर्तमान टोकन के स्टेट के साथ कॉन्कैटिनेट किया जाता है और समानांतर में प्रोसेस किया जाता है।
परिणाम: यह एक सिंगल पैरेलल पास में सभी स्थगित टोकन्स के लिए पूर्ण और सटीक KV स्टेट्स को भर देता है, जिससे यह सुनिश्चित होता है कि KV कैश पूर्ण, सटीक और सैंपल किए गए टोकन के अनुरूप है। कोई भी कंप्यूट स्किप नहीं किया जाता है; इसे केवल क्रमिक (sequential) से पैरेलल बैच प्रोसेसिंग में पुनर्गठित किया जाता है।

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

एक्सैक्ट सैंपलिंग मिक्सचर (Exact Sampling Mixture): N-vium कई एग्जिट्स के ऊपर सीखे हुए मिश्रण वितरण के साथ सटीक सैंपलिंग सक्षम करने वाला पहला तरीका है, जो मानक अर्ली-एग्जिट विधियों में निहित सन्निकटन त्रुटियों से बचता है।
फुल KV रिकवरी (Full KV Recovery): स्थगित और बैचिंग के माध्यम से ऊपरी-परत की गणनाओं को प्रोसेस करके, N-vium पूर्ण और सटीक KV कैश प्राप्त करता है, जो अन्य एडेप्टिव कंप्यूटेशन विधियों में अक्सर खो जाता है।
ट्रेनिंग फ्रेमवर्क: लेखक एक व्यावहारिक ट्रेनिंग पाइपलाइन प्रदान करते हैं जो लर्नबल कंप्यूट पेनल्टी के माध्यम से दक्षता की ओर मॉडल को मोड़ने के लिए मिक्सचर ऑब्जेक्टिव को सीधे ऑप्टिमाइज़ करती है।
स्केलेबिलिटी (Scalability): विधि को 1.5 बिलियन पैरामीटर्स तक के स्केल पर मान्य किया गया है, जो यह दर्शाता है कि यह दृष्टिकोण छोटे मॉडल्स तक सीमित नहीं है।

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

लेखकों ने C4 डेटासेट पर LLaMA 2 आर्किटेक्चर का उपयोग करके N-vium मॉडल्स (विशेष रूप से "Quadrivium" जिसमें 4 एग्जिट्स हैं) को प्री-ट्रेन किया, और उनकी तुलना पैरामीटर- और डेटा-मैच्ड डेंस बेसलाइन्स से की।

स्पीडअप बनाम परप्लेक्सिटी (Speedup vs. Perplexity):
- सबसे बड़े मॉडल (1.5B पैरामीटर्स, 48 लेयर्स, 1536 हिडन विड्थ) ने कोई परप्लेक्सिटी गिरावट नहीं (0% अंतर) के साथ डेंस बेसलाइन पर 57.9% वॉल-क्लॉक स्पीडअप हासिल किया।
- छोटे मॉडल्स ने 23% से 41% तक निरंतर स्पीडअप दिखाया, जिसमें थोड़ी परप्लेक्सिटी सुधार (जैसे, 36-लेयर मॉडल ने 36.4% स्पीडअप के साथ 3.7% परप्लेक्सिटी सुधार दिखाया) भी देखा गया।
पारेटो फ्रंटियर (Pareto Frontier): क्वालिटी-स्पीड प्लॉट्स में, N-vium "ऊपरी-दाएं" (upper-right) चतुर्थांश में स्थित है (तेज़ और समान या बेहतर क्वालिटी), जबकि CALM (CALM* के रूप में पुन: कार्यान्वित) जैसी विधियाँ "निचले-दाएं" (lower-right) में हैं (तेज़ लेकिन खराब क्वालिटी) और LayerSkip बेसलाइन के नीचे एक स्थिर परप्लेक्सिटी लाइन पर है।
डाउनस्ट्रीम परफॉरमेंस (Downstream Performance): Tulu 3 मिक्सचर पर सुपरवाइज्ड फाइन-ट्यूनिंग (SFT) के बाद, N-vium मॉडल्स ने डाउनस्ट्रीम बेंचमार्क्स (MMLU, ARC, HellaSwag, आदि) पर डेंस बेसलाइन्स के बराबर या उनसे बेहतर प्रदर्शन किया, जो पुष्टि करता है कि मॉडल स्टीयरेबल (steerable) बने रहते हैं और टास्क परफॉरमेंस की कीमत पर क्रॉस-एन्ट्रॉपी ऑब्जेक्टिव पर ओवरफिट नहीं होते हैं।
हार्डवेयर रोबस्टनेस (Hardware Robustness): A100, RTX A6000, और RTX 3090 GPUs पर स्पीडअप मापेशन ने सुसंगत परिणाम (~57%) दिए, जबकि पुराने ट्यूरिंग आर्किटेक्चर वाले GPUs ने नेटिव bfloat16 सपोर्ट की कमी के कारण कम लाभ दिखाया।

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

पेपर का दावा है कि N-vium इन्फरेंस दक्षता में एक प्रतिमान परिवर्तन (paradigm shift) का प्रतिनिधित्व करता है। अधिकांश विधियों के विपरीत जो प्रति टोकन FLOPs को कम करके (अक्सर सन्निकटन के माध्यम से) टोकन प्रति सेकंड में सुधार करती हैं, N-vium प्रति सेकंड FLOPs को बढ़ाकर दक्षता में सुधार करता है। यह क्रमिक गणना को पैरेलल बैचिंग पास में पुनर्गठित करके ऐसा करता है, जिससे हार्डवेयर पैरेललिज्म का बेहतर उपयोग होता है, बिना किसी गणना को हटाए।

लेखक जोर देते हैं कि N-vium एक स्ट्रिक्ट जनरलाइजेशन है: जब मध्यवर्ती एग्जिट्स के लिए रूटिंग वेट शून्य होते हैं, तो मॉडल बिल्कुल एक मानक ट्रांसफॉर्मर की तरह व्यवहार करता है। यह विधि चार महत्वपूर्ण मानदंडों को एक साथ पूरा करती है, जिसे लेखक दावा करते हैं कि कोई अन्य मौजूदा विधि नहीं करती है:

सटीक इन्फरेंस (Exact Inference): सन्निकटन के बिना प्रशिक्षित वितरण से सैंपलिंग।
फुल KV कैश रिकवरी (Full KV Cache Recovery): कोई भी लेयर्स या अनुमानित कैश फिलिंग नहीं।
लर्न्ड रूटिंग (Learned Routing): प्रति-टोकन गणना एक प्रशिक्षित राउटर द्वारा आवंटित की जाती है, न कि ह्यूरिस्टिक्स द्वारा।
डिकोड-टाइम स्पीडअप (Decode-time Speedup): ऑटो-रिग्रेसिव जनरेशन के दौरान वॉल-क्लॉक लेटेंसी को कम करता है।

पेपर निष्कर्ष निकालता है कि हालांकि N-vium को ट्रेन करने में प्रति टोकन उच्च FLOP लागत आती है (कई जंक्शन्स पर हेड्स को इवैल्यूएट करने के कारण), यह ओवरहेड प्रोडक्शन डिप्लॉयमेंट के लिए उचित है जहाँ इन्फरेंस लागत हावी होती है, और सापेक्ष ओवरहेड मॉडल साइज बढ़ने के साथ कम होता जाता है।