💬 NLP

Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts

यह शोध पत्र kNN-MoE को प्रस्तुत करता है, जो एक रिट्रीवल-ऑगमेंटेड रूटिंग फ्रेमवर्क है जो फ्रोजन राउटर की तुलना में वितरण बदलावों (डिस्ट्रीब्यूशन शिफ्ट्स) को अधिक प्रभावी ढंग से संभालने के लिए अनुकूलित पिछले रूटिंग निर्णयों की मेमोरी और कॉन्फिडेंस-ड्रिवन फॉलबैक मैकेनिज्म का लाभ उठाकर मिक्सचर-ऑफ-एक्सपर्ट्स मॉडल्स को गतिशील रूप से बेहतर बनाता है।

मूल लेखक: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

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

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

CC BY 4.0

मूल लेखक: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

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

कल्पना कीजिए कि आपके पास विशेषज्ञों की एक विशाल, अत्यधिक बुद्धिमान लाइब्रेरी है (एक "Mixture-of-Experts" मॉडल)। जब आप इस लाइब्रेरी से कोई प्रश्न पूछते हैं, तो एक Router एक लाइब्रेरियन की तरह काम करता है। इसका काम यह तय करना है कि आपके प्रश्न को देखकर लाइब्रेरी में से कौन सा विशिष्ट विशेषज्ञ उत्तर देने के लिए सबसे उपयुक्त है।

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

यह पेपर kNN-MoE पेश करता है, जो एक चतुर अपग्रेड है जो इस लाइब्रेरियन को पिछले सफल अनुभवों के आधार पर एक "चीट शीट" (cheat sheet) देता है। यह कैसे काम करता है, इसे सरल अवधारणाओं में यहाँ समझाया गया है:

1. "चीट शीट" (मेमोरी निर्माण)

इससे पहले कि सिस्टम वास्तव में कोई वास्तविक प्रश्न हल करे, शोधकर्ता अभ्यास प्रश्नों का एक सेट (reference dataset) लेते हैं और उन्हें लाइब्रेरी के माध्यम से चलाते हैं।

प्रयोग: इन अभ्यास प्रश्नों के प्रत्येक शब्द के लिए, वे पूछते हैं: "यदि हम जादू से लाइब्रेरियन के निर्णय को अभी बदल सकें, तो किस विशेषज्ञ ने अब तक का सबसे अच्छा उत्तर दिया होता?"
परिणाम: वे प्रत्येक विशिष्ट क्षण के लिए "परफेक्ट" विशेषज्ञ को ढूंढते हैं और उसे लिख लेते हैं। वे इन जोड़ों को संग्रहीत करते हैं: "यह विशिष्ट प्रश्न इनपुट" + "परफेक्ट विशेषज्ञ का चुनाव।"
उपमा: कल्पना कीजिए कि लाइब्रेरियन एक अंतिम परीक्षा के लिए पढ़ाई कर रहा है। केवल नियमों को याद करने के बजाय, वे एक विशाल इंडेक्स कार्ड डेक बना रहे हैं। एक तरफ एक कठिन प्रश्न है; दूसरी तरफ वह सटीक विशेषज्ञ है जिसने अतीत में इसे पूरी तरह से हल किया था।

2. "स्मार्ट लुकअप" (इन्फरेंस)

अब, जब एक वास्तविक उपयोगकर्ता कोई प्रश्न पूछता है, तो सिस्टम दो चीजें एक साथ करता है:

स्थिर लाइब्रेरियन (Frozen Librarian): मूल, स्थिर राउटर अपने प्रशिक्षण के आधार पर अपना सर्वश्रेष्ठ अनुमान लगाता है।
चीट शीट: सिस्टम उपयोगकर्ता के प्रश्न को इंडेक्स कार्ड डेक में खोजता है ताकि पिछले सबसे मिलते-जुलते प्रश्नों को पाया जा सके।

3. "कॉन्फिडेंस वोट" (एडेप्टिव मिक्सिंग)

यह सबसे महत्वपूर्ण हिस्सा है। सिस्टम चीट शीट पर अंधा विश्वास नहीं करता है। यह जांचता है कि पिछले मामले वर्तमान प्रश्न के कितने समान हैं।

उच्च आत्मविश्वास (High Confidence): यदि पिछले मामले वर्तमान प्रश्न के लगभग समान हैं, तो सिस्टम कहता है, "लाइब्रेरियन इस कठिन प्रश्न के बारे में अनिश्चित है, लेकिन हमारी चीट शीट कहती है कि परफेक्ट विशेषज्ञ X है।" इसके बाद यह लाइब्रेरियन के अनुमान को चीट शीट की सलाह के साथ मिला देता है, और चीट शीट पर अधिक भरोसा करता है।
कम आत्मविश्वास (Low Confidence): यदि वर्तमान प्रश्न पूरी तरह से अद्वितीय है और चीट शीट में कुछ भी मेल नहीं खाता है, तो सिस्टम कहता है, "चीट शीट यहाँ बेकार है; हम शायद थोड़ा शोर (noise) जोड़ सकते हैं।" यह लुकअप को अनदेखा कर देता है और मूल स्थिर लाइब्रेरियन के निर्णय पर भरोसा करता है।

यह क्यों महत्वपूर्ण है

पेपर का दावा है कि यह तरीका दो चरम सीमाओं के बीच का एक "स्वीट स्पॉट" (sweet spot) है:

कुछ न करना (Zero-Shot): केवल स्थिर लाइब्रेरियन का उपयोग करना। यह तेज़ है लेकिन कठिन या नए प्रश्नों पर विफल हो सकता है।
री-ट्रेनिंग (Supervised Fine-Tuning): लाइब्रेरियन को नए नियम शून्य से सिखाना। यह अच्छा काम करता है लेकिन बहुत धीमा, महंगा और हर नए कार्य के लिए पूरी प्रक्रिया को फिर से करने की आवश्यकता होती है।

kNN-MoE बिना भारी लागत के री-ट्रेनिंग वाला प्रदर्शन लाभ प्रदान करता है। यह एक लाइब्रेरियन को स्कूल वापस जाने के लिए मजबूर करने के बजाय, उन्हें उनके पिछले बेहतरीन कदमों की एक गतिशील, खोजने योग्य स्मृति देने जैसा है।

पेपर के मुख्य निष्कर्ष

यह कठिन प्रश्नों पर काम करता है: यह सिस्टम तब सबसे अधिक मदद करता है जब मूल लाइब्रेरियन भ्रमित होता है (उच्च "perplexity")। जब लाइब्रेरियन पहले से ही आश्वस्त होता है, तो सिस्टम गड़बड़ी से बचने के लिए पीछे हट जाता है।
कम ही अधिक है (Less is more): आश्चर्यजनक रूप से, केवल एक पिछले उदाहरण (सबसे करीबी मैच) को देखना कई उदाहरणों का औसत निकालने से बेहतर काम करता है। पेपर सुझाव देता है कि विशेषज्ञ रूटिंग के लिए, अतीत के बहुत अधिक "मत" होने से सिग्नल कमजोर हो जाता है।
गति बनाम सटीकता: इस "चीट शीट" को बनाना (ऑफलाइन) पूरे मॉडल को फिर से प्रशिक्षित करने की तुलना में बहुत तेज़ है। वास्तविक उत्तर देने के चरण के दौरान, यह थोड़ा विलंब (लगभग 3-4% धीमा) जोड़ता है, लेकिन चिकित्सा परीक्षाओं और कोडिंग जैसे कठिन कार्यों पर सटीकता में काफी सुधार करता है।

कमी (सीमाएं)

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

संक्षेप में, kNN-MoE AI विशेषज्ञों को उनके पिछले आदर्श निर्णयों की ओर झांकने की अनुमति देकर उन्हें स्मार्ट बनाने का एक तरीका है, लेकिन केवल तभी जब यह सुरक्षित हो।

तकनीकी सारांश: kNN-MoE

समस्या विवरण

Mixture-of-Experts (MoE) आर्किटेक्चर एक पैरामीट्रिक "रूटर" का उपयोग करके टोकन को विशेषज्ञों के एक विरल (sparse) उपसमूह में भेजकर लार्ज लैंग्वेज मॉडल्स (LLMs) को कुशलतापूर्वक स्केल करते हैं। हालाँकि, मानक रूटर हल्के क्लासिफायर होते हैं जिन्हें एक बार प्रशिक्षित किया जाता है और फिर फ्रीज कर दिया जाता है। इन-डिस्ट्रीबशन डेटा पर प्रभावी होने के बावजूद, वे वितरण बदलावों (जैसे विशेष डोमेन या आउट-ऑफ-डिस्ट्रीब्यूशन कार्यों) का सामना करते समय लचीलेपन की कमी प्रदर्शित करते हैं। ऐसे परिदृश्यों में, रूटर को स्थिर पूर्व धारणाओं (fixed priors) से अनुमान लगाना पड़ता है, जिससे अक्सर उप-इष्टतम विशेषज्ञ चयन और उच्च परप्लेक्सिटी (perplexity) होती है। मौजूदा टेस्ट-टाइम अडैप्टेशन विधियाँ, जैसे कि C3PO, ऑनलाइन रूटिंग लॉजिट्स को फाइन-ट्यून करने का प्रयास करती हैं, लेकिन इसमें बेसलाइन की तुलना में काफी अधिक कम्प्यूटेशनल ओवरहेड (लगभग 5 गुना) होता है। अन्य दृष्टिकोण सुपरवाइज्ड फाइन-ट्यूनिंग (SFT) के माध्यम से स्थायी पैरामीटर अपडेट की मांग करते हैं, जो कम्प्यूटेशनल रूप से गहन और इन्फरेंस के लिए कम कुशल है।

कार्यप्रणाली: kNN-MoE

लेखक kNN-MoE का प्रस्ताव करते हैं, जो एक रिट्रीवल-ऑगमेंटेड रूटिंग फ्रेमवर्क है जो मॉडल के पैरामीटर्स को संशोधित किए बिना विशेषज्ञ असाइनमेंट निर्णयों को परिष्कृत करता है। यह विधि दो अलग-अलग चरणों में कार्य करती है:

1. मेमोरी निर्माण (Offline)

सिस्टम प्रत्येक MoE लेयर के लिए एक नॉन-पैरामीट्रिक मेमोरी स्टोर $\mathcal{M}$ बनाता है, जिसमें की-वैल्यू पेयर्स $(k_t, v_t)$ शामिल होते हैं:

की (Key - $k_t$ ): संदर्भ डेटासेट $D_{ref}$ से रूटर इनपुट (हिडन स्टेट) $x_t$ ।
वैल्यू (Value - $v_t$ ): एक "स्थानीय रूप से इष्टतम" विशेषज्ञ असाइनमेंट, जिसे ग्राउंड-ट्रुथ टोकन $y_t$ $y_{t}$ की संभावना को अधिकतम करने के लिए रूटिंग लॉजिट्स को अनुकूलित करके प्राप्त किया जाता है।
- फ्रोजन रूटर के आउटपुट को स्टोर करने के बजाय, यह विधि प्रत्येक टोकन के लिए एक ऑप्टिमाइज़ेशन समस्या को हल करती है: एक ऐसा लॉजिट वेक्टर $r^*_t$ खोजना जो फ्रोजन नेटवर्क पैरामीटर्स को देखते हुए टारगेट टोकन के नेगेटिव लॉग-लाइक्लीहुड को कम करता हो।
- यह ऑप्टिमाइज़ेशन मूल पैरामीट्रिक लॉजिट्स से शुरू होकर ग्रेडिएंट डिसेंट के कुछ स्टेप्स (आमतौर पर $S=1$ ) के माध्यम से किया जाता है।
- परिणामी अनुकूलित लॉजिट्स $r^{(S)}_t$ मेमोरी में वैल्यू $v_t$ के रूप में कार्य करते हैं।

2. कॉन्फिडेंस-अवेयर एडेप्टिव मिक्सिंग (Inference)

इन्फरेंस के दौरान, दिए गए रूटर इनपुट $x$ के लिए:

रिट्रीवल (Retrieval): सिस्टम $x$ और संग्रहीत कीज़ के बीच समानता के आधार पर मेमोरी $\mathcal{M}$ से $K$ निकटतम पड़ोसियों $\mathcal{N}(x)$ को प्राप्त करता है।
नॉन-पैरामीट्रिक प्रपोजल (Non-Parametric Proposal): एक मेमोरी-आधारित लॉजिट प्रपोजल $r_{mem}(x)$ की गणना, इनपुट के साथ समानता के आधार पर प्राप्त किए गए अनुकूलित लॉजिट्स को एकत्रित करके की जाती है।
कॉन्फिडेंस-ड्रिवन मिक्सिंग (Confidence-Driven Mixing): अंतिम रूटिंग लॉजिट्स, मूल पैरामीट्रिक रूटर लॉजिट्स $r(x)$ $r (x)$ और मेमोरी-आधारित प्रपोजल $r_{mem}(x)$ $r_{m e m} (x)$ के बीच एक लीनियर इंटरपोलेशन हैं:
$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$
मिक्सिंग कोएफिशिएंट $\lambda(x)$ $λ (x)$ को प्राप्त किए गए पड़ोसियों की औसत समानता द्वारा परिभाषित किया गया है। यह एक कॉन्फिडेंस स्कोर के रूप में कार्य करता है:
- उच्च समानता ( $\lambda \approx 1$ ): मॉडल प्राप्त किए गए, स्थानीय रूप से इष्टतम असाइनमेंट पर भरोसा करता है।
- कम समानता ( $\lambda \approx 0$ ): मॉडल मूल फ्रोजन रूटर पर वापस आ जाता है, जिससे अप्रासंगिक मामलों से होने वाले शोर (noise) से बचाव होता है।
आउटपुट: अंतिम लॉजिट्स को स्पार्स विशेषज्ञ असाइनमेंट उत्पन्न करने के लिए टॉप-K सॉफ्टमैक्स (Top-K Softmax) के माध्यम से पास किया जाता है, जिससे मॉडल की मूल स्पैरसिटी बाधाएं बनी रहती हैं।

मुख्य योगदान

रिट्रीवल-ऑगमेंटेड रूटिंग: एक ऐसा फ्रेमवर्क पेश करता है जो समान पिछले मामलों की मेमोरी से ऑफलाइन-कम्प्यूटेड, स्थानीय रूप से इष्टतम विशेषज्ञ असाइनमेंट का पुन: उपयोग करके रूटिंग निर्णयों में सुधार करता है।
नॉन-पैरामीट्रिक अडैप्टेशन: SFT या ऑनलाइन अडैप्टेशन विधियों के विपरीत, kNN-MoE मॉडल पैरामीटर्स को अपडेट नहीं करता है। यह कम्प्यूटेशनल बोझ को ऑफलाइन मेमोरी निर्माण चरण की ओर स्थानांतरित कर देता है, जिससे कुशल इन्फरेंस संभव होता है।
कॉन्फिडेंस-ड्रिवन फॉलबैक: यह विधि प्राप्त किए गए पड़ोसियों की औसत समानता का उपयोग एक डायनेमिक मिक्सिंग कोएफिशिएंट के रूप में करती है, यह सुनिश्चित करती है कि जब कोई प्रासंगिक ऐतिहासिक मामला न मिले तो मॉडल फ्रोजन रूटर पर वापस लौट जाए।
दक्षता (Efficiency): इस पद्धति के लिए मेमोरी निर्माण के दौरान प्रति टोकन केवल एक ग्रेडिएंट डिसेंट स्टेप की आवश्यकता होती है और यह ऑनलाइन ऑप्टिमाइज़ेशन या पूर्ण मॉडल फाइन-ट्यूनिंग की भारी कम्प्यूटेशनल लागतों से बचती है।

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

लेखकों ने तीन MoE मॉडल्स (OLMoE, GPT-OSS, और Qwen3) को छह बेंचमार्क्स (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA, और MBPP) पर kNN-MoE के माध्यम से मूल्यांकित किया।

बेसलाइन्स के विरुद्ध प्रदर्शन: kNN-MoE ने सभी मॉडल्स और डेटासेट्स पर ज़ीरो-शॉट बेसलाइन और "SFT (रूटर ओनली)" बेसलाइन को लगातार पछाड़ दिया।
SFT के साथ प्रतिस्पर्धात्मकता: इस विधि ने फुल सुपरवाइज्ड फाइन-ट्यूनिंग (SFT) के साथ प्रतिस्पर्धी प्रदर्शन हासिल किया और कई डेटा-दुर्लभ सेटिंग्स (जैसे GPQA, USMLE) में SFT से बेहतर प्रदर्शन किया, जो कभी-कभी छोटे रेफरेंस सेट्स पर ओवरफिटिंग का शिकार हो जाता है।
स्थिरता: 5-शॉट प्रॉम्प्टिंग के विपरीत, जो इनपुट लेंथ संबंधी समस्याओं के कारण कुछ मॉडल्स (GPT-OSS, OLMoE) पर प्रदर्शन को कम करती है, kNN-MoE ने स्थिर सुधार प्रदर्शित किया।
एब्लेशन स्टडीज (Ablation Studies):
- पड़ोसी ( $K$ ): विशिष्ट kNN-LM निष्कर्षों के विपरीत, प्रदर्शन $K=1$ पर चरम पर था। अधिक पड़ोसियों ( $K>1$ ) को एकत्रित करने से शोर (noise) बढ़ गया, जो सुझाव देता है कि वास्तव में लाभकारी रूटिंग सिग्नल दुर्लभ और विशिष्ट होते हैं।
- ऑप्टिमाइज़ेशन स्टेप्स ( $S$ ): $S=1$ से अधिक ग्रेडिएंट डिसेंट स्टेप्स बढ़ाने से नगण्य सटीकता लाभ हुआ जबकि निर्माण लागत बढ़ गई।
- रेफरेंस सेट का आकार: प्रदर्शन संदर्भ सेट $|D_{ref}|$ के आकार के साथ निरंतर सुधरा, जिसमें 0.5k–1k उदाहरणों के बाद घटते प्रतिफल (diminishing returns) देखे गए।
दक्षता: 1k उदाहरणों के लिए मेमोरी निर्माण में 0.46 घंटे लगे (बनाम ~1.34h SFT के लिए)। IVFPQ (Inverted File + Product Quantization) वेरिएंट ने बेसलाइन की तुलना में इन्फरेंस लेटेंसी ओवरहेड को ~3.67% और VRAM उपयोग वृद्धि को ~9.43% तक कम कर दिया।

महत्व और दावे

पेपर का दावा है कि kNN-MoE स्टैटिक पैरामीट्रिक रूटिंग और डायनेमिक इन्फरेंस आवश्यकताओं के बीच के अंतर को पाटता है। इसका प्राथमिक महत्व फाइन-ट्यूनिंग के कम्प्यूटेशनल बोझ के बिना MoE अनुकूलन क्षमता को बढ़ाने के लिए एक स्केलेबल मार्ग प्रदान करने में निहित है।

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