🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

MoE-nD एक प्रति-परत मिश्रण-विशेषज्ञ (mixture-of-experts) रूटिंग फ्रेमवर्क पेश करता है जो एक वैश्विक मेमोरी बजट के तहत प्रत्येक ट्रांसफॉर्मर परत को अद्वितीय टोकन इविक्शन अनुपात और क्वांटाइजेशन बिट-चौड़ाई गतिशील रूप से आवंटित करता है, जिससे लंबी-संदर्भ तर्क कार्यों (long-context reasoning tasks) पर अनकंप्रेस्ड बेसलाइन के तुलनीय सटीकता बनाए रखते हुए 14x तक KV कैश संपीड़न प्राप्त किया जाता है।

मूल लेखक: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

प्रकाशित 2026-04-21

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

CC BY 4.0

मूल लेखक: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

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

यहाँ MoE-nD पेपर का सरल भाषा में अनुवाद दिया गया है:

बड़ी समस्या: "ओवरस्टफ्ड बैकपैक" (ज़रूरत से ज़्यादा भरा हुआ बस्ता)

कल्पना कीजिए कि एक लार्ज लैंग्वेज मॉडल (LLM) एक बहुत ही बुद्धिमान छात्र है जो एक बहुत लंबी कहानी या एक जटिल गणित की समस्या को हल करने की कोशिश कर रहा है। ऐसा करने के लिए, छात्र को अपनी "वर्किंग मेमोरी" (जिसे KV Cache कहा जाता है) में उन सभी चीज़ों की एक चलती हुई सूची रखनी होगी जो उसने अब तक पढ़ी हैं।

छोटी कहानियों के लिए, यह मेमोरी एक बैकपैक में आसानी से समा जाती है। लेकिन लंबे संदर्भों (जैसे कि पूरी किताब या 16,000 शब्दों का दस्तावेज़) के लिए, बैकपैक इतना भारी और बड़ा हो जाता है कि:

यह छात्र की गति को धीमा कर देता है (लेटेंसी)।
यह मानक कंप्यूटरों पर बैकपैक को फाड़ देता है (आउट-ऑफ-मेमोरी एरर)।

पुराना समाधान: "एक ही नियम सबके लिए"

इसे ठीक करने के लिए, शोधकर्ताओं ने बैकपैक को छोटा करने की कोशिश की। उन्होंने दो मुख्य उपकरणों का उपयोग किया:

इविक्शन (Eviction - चीज़ों को बाहर फेंकना): जगह बनाने के लिए पुरानी नोट्स को डिलीट करना।
क्वांटाइजेशन (Quantization - शॉर्टहैंड में लिखना): नोट्स को छोटी, कम सटीक लिखावट में लिखना (जैसे, 16-बिट के बजाय 4-बिट नंबरों का उपयोग करना)।

दोष: पिछले तरीकों ने छात्र के मस्तिष्क के साथ एक "कुकी कटर" (एक ही सांचे) जैसा व्यवहार किया। उन्होंने मस्तिष्क के हर हिस्से पर एक ही नियम लागू किया।

उदाहरण: "पूरी कहानी के नोट्स में से 50% फेंक दो, और बाकी बचे हुए सभी नोट्स को बहुत छोटी शॉर्टहैंड में लिखो।"

यह क्यों विफल होता है: छात्र का मस्तिष्क एक समान नहीं होता।

लेयर 1 (शुरुआत): यह कहानी को समझने के लिए महत्वपूर्ण हो सकती है। यदि आप इन नोट्स का 50% फेंक देते हैं, तो छात्र कहानी पूरी तरह भूल जाएगा।
लेयर 20 (मध्य): यह फालतू की बातों से भरी हो सकती है। यहाँ 50% फेंकने से कोई खास नुकसान नहीं होगा।
लेयर 28 (अंत): इसे अंतिम गणितीय समीकरण को हल करने के लिए उच्च सटीकता की आवश्यकता हो सकती है। यहाँ "छोटी शॉर्टहैंड" में लिखने से गणना में त्रुटियां हो सकती हैं।

हर लेयर के साथ एक जैसा व्यवहार करना वैसा ही है जैसे आपने पूरे साल के लिए "औसत" पहनावे का निर्णय लिया हो और गर्मियों में विंटर कोट और सर्दियों में टी-शर्ट पहन रहे हों।

नया समाधान: MoE-nD (एक "स्मार्ट वार्डरोब")

लेखक MoE-nD का प्रस्ताव देते हैं, जिसका अर्थ है Mixture-of-Experts for n-Dimensional routing।

इसे छात्र के बैकपैक के लिए एक स्मार्ट वार्डरोब मैनेजर के रूप में सोचें। हर किसी के लिए एक ही नियम होने के बजाय, मैनेजर मस्तिष्क के हर एक लेयर को व्यक्तिगत रूप से देखता है और पूछता है: "इस विशिष्ट मस्तिष्क भाग को अभी किस चीज़ की आवश्यकता है?"

इसके पास प्रत्येक लेयर के लिए चुनने के लिए तीन "विशेषज्ञ" (विकल्प) हैं:

कचरे का डिब्बा (Eviction): हम इस लेयर की कितनी मेमोरी डिलीट कर सकते हैं?
शॉर्टहैंड पेन (K-Quantization): हम "Key" नोट्स को कितना छोटा लिख सकते हैं?
शॉर्टहैंड पेन (V-Quantization): हम "Value" नोट्स को कितना छोटा लिख सकते हैं?

यह कैसे काम करता है:

मैप (ऑफलाइन कैलिब्रेशन): छात्र की परीक्षा शुरू होने से पहले, सिस्टम एक त्वरित, सस्ता टेस्ट चलाता है ताकि यह देखा जा सके कि प्रत्येक लेयर कितनी संवेदनशील है। यह एक मैप बनाता है: "लेयर 5 को डिलीट होना पसंद नहीं है। लेयर 10 को शॉर्टहैंड में लिखना पसंद नहीं है। लेयर 20 को दोनों में से किसी की भी परवाह नहीं है।"
बजट (ग्लोबल लिमिट): छात्र के पास एक सख्त बैकपैक आकार सीमा है (जैसे, 136 MB)।
ग्रीडी सॉल्वर (Greedy Solver): एक स्मार्ट एल्गोरिदम मैप और बजट को देखता है। यह "बेस्ट वैल्यू" के खेल को खेलता है।
- यह कहता है: "ठीक है, लेयर 5 बहुत संवेदनशील है, इसलिए इसे फुल साइज में रखते हैं। लेकिन ले्यो 10 छोटी शॉर्टहैंड के लिए ठीक है, तो इसे सिकोड़ देते हैं। लेयर 20 अपने नोट्स आधा खो सकता है बिना रोए।"
परिणाम: बैकपैक को उसके मूल आकार के 1/14वें हिस्से तक संकुचित कर दिया जाता है, लेकिन छात्र का प्रदर्शन बिल्कुल वैसा ही होता है जैसा कि उनके पास पूर्ण, बिना संकुचित बैकपैक होता।

मुख्य खोज: "कचरे का डिब्बा ही हीरो है"

पेपर ने यह देखने के लिए एक दिलचस्प प्रयोग किया कि कौन सी रणनीति सबसे अधिक मायने रखती है। उन्होंने तीन संस्करणों का परीक्षण किया:

यूनिफॉर्म (Uniform): सबके लिए एक ही नियम। (भयानक परिणाम)।
स्मार्ट शॉर्टहैंड (Smart Shorthand): अलग-अलग लेयर्स के लिए अलग-अलग शॉर्टहैंड आकार, लेकिन कचरा फेंकने के नियम एक ही। (अभी भी बुरा)।
स्मार्ट ट्रैश (Smart Trash): एक ही शॉर्टहैंड आकार, लेकिन अलग-अलग लेयर्स के लिए अलग-अलग कचरा फेंकने के नियम। (यह विजेता था!)

उपमा (Analogy):
कल्पthoughtना कीजिए कि आप यात्रा के लिए पैकिंग कर रहे हैं।

यूनिफॉर्म: आप अपने कपड़े 50% फेंक देते हैं और बाकी को सिकोड़ देते हैं। आप ठंड महसूस करते हैं और ठीक से कपड़े नहीं पहन पाते।
स्मार्ट शॉर्टहैंड: आप अपने सभी कपड़े रखते हैं लेकिन कपड़े के फैब्रिक को सिकोड़ देते हैं। आपके पास अभी भी बहुत सारे कपड़े हैं; बैग बहुत भारी है।
MoE-nD (स्मार्ट ट्रैश): आप महसूस करते हैं कि आपके "मोज़े" (कुछ लेयर्स) बेकार हैं और आप उनमें से 90% फेंक देते हैं, लेकिन आप अपनी "जैकेट" (महत्वपूर्ण लेयर्स) को पूरी तरह सुरक्षित रखते हैं। बैग हल्का हो जाता है, लेकिन आप फिर भी गर्म और तैयार रहते हैं।

परिणाम: लंबी कहानियों पर जादू

जब उन्होंने कठिन बेंचमार्क (जैसे 16,000 शब्दों की कहानी पढ़ना या जटिल गणित हल करना) पर इसका परीक्षण किया:

पुराने तरीके: जब उन्होंने मेमोरी को सिकोड़ने की कोशिश की, तो मॉडल भ्रमित हो गया और विफल हो गया (सटीकता लगभग शून्य हो गई)।
MoE-nD: इसने मेमोरी को 14 गुना सिकोड़ दिया (1.9 GB से 136 MB तक) और मॉडल का प्रदर्शन बिल्कुल वैसा ही रहा जैसा कि इसके विशाल, बिना संकुचित संस्करण का था।

यह कब काम नहीं करता?

पेपर अपनी सीमाओं के बारे में ईमानदार है। यह छोटे कार्यों (जैसे 500 शब्दों की साधारण गणित की समस्या) में मदद नहीं करता है।

क्यों? यदि कहानी छोटी है, तो बैकपैक पहले से ही छोटा है। "स्मार्ट मैनेजर" लेयर्स को देखता है और कहता है, "हे, हमारे पास काफी जगह है! चलो सब कुछ रखते हैं।"
इस स्थिति में, सिस्टम बस "सब कुछ रखने" पर डिफ़ॉल्ट हो जाता है, और फैंसी रूटिंग कोई मूल्य नहीं जोड़ती है। यह एक कमरे में नेविगेट करने के लिए GPS का उपयोग करने जैसा है जिसे आप स्पष्ट रूप से देख सकते हैं; आपको जटिल दिशाओं की आवश्यकता नहीं है।

सारांश

MoE-nD एक ऐसी तकनीक है जो AI के मस्तिष्क को मिट्टी के एक एकल ब्लॉक की तरह मानना बंद करती है। इसके बजाय, यह प्रत्येक लेयर को एक अद्वितीय व्यक्ति के रूप में मानती है, कुछ लेयर्स को "डिलीट" बटन देती है, कुछ को "सिकोड़ने" का बटन, और कुछ को "सुरक्षित रखने" का बटन देती है। हर सिंगल लेयर के लिए कंप्रेशन को कस्टमाइज़ करके, यह AI को सुपरकंप्यूटर की मेमोरी की आवश्यकता के बिना लंबी कहानियाँ याद रखने की अनुमति देता है।

यहाँ शोध पत्र "MoE-nD: Multi-Axis KV Cache Compression के लिए Per-Layer Mixture-of-Experts Routing" का विस्तृत तकनीकी सारांश दिया गया है।

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

लंबे संदर्भ (long context) को प्रोसेस करने वाले लार्ज लैंग्वेज मॉडल्स (LLMs) एक गंभीर बाधा का सामना करते हैं: KV (Key-Value) कैश प्रत्येक अनुक्रम (sequence) के लिए टेराबाइट्स के बजाय गीगाबाइट्स में मेमोरी का उपभोग करता है, जिससे कमोडिटी हार्डवेयर पर लेटेंसी की समस्या और आउट-ऑफ-मेमोरी (OOM) विफलताएं होती हैं।

मौजूदा संपीड़न रणनीतियाँ (compression strategies) आमतौर पर एक एकल अक्ष (single axis) (जैसे, टोकन इविक्शन, क्वांटाइजेशन, या लो-रैंक प्रोजेक्शन) पर कार्य करती हैं और सभी ट्रांसफार्मर लेयर्स पर एक समान नीति (uniform policy) लागू करती हैं। लेखक तर्क देते हैं कि यह "समरूपता" (homogeneity) उप-इष्टतम (suboptimal) है क्योंकि:

विभिन्न लेयर्स संपीड़न कार्यों (compression operations) के प्रति अलग-अलग प्रतिक्रिया देती हैं (उदाहरण के लिए, कुछ लेयर्स आक्रामक टोकन इविक्शन को सहन कर सकती हैं, जबकि अन्य क्वांटाइजेशन के प्रति संवेदनशील होती हैं)।
इविक्शन (अनुक्रम लंबाई कम करना) और क्वांटाइजेशन (परिशुद्धता कम करना) के बीच इष्टतम संतुलन (trade-off) प्रत्येक लेयर के लिए काफी भिन्न होता है।
वर्तमान विधियाँ वैश्विक मेमोरी बजट के तहत प्रति-लेयर (per-layer) आधार पर इन अक्षों को संयुक्त रूप से अनुकूलित करने में विफल रहती हैं।

2. कार्यप्रणाली: MoE-nD (Methodology: MoE-nD)

लेखक MoE-nD का प्रस्ताव करते हैं, जो एक ऐसा फ्रेमवर्क है जो विभिन्न संपीड़न कॉन्फ़िगरेशन को "विशेषज्ञों" (experts) के रूप में मानता है और प्रत्येक लेयर को उसके इष्टतम कॉन्फ़िगरेशन तक रूट करता है।

मुख्य घटक (Core Components)

Multi-Axis Design Space: प्रत्येक लेयर $\ell$ $ℓ$ के लिए, सिस्टम एक टुपल $(keep\_ratio, k\_bits, v\_bits)$ $(k ee p_r a t i o, k_bi t s, v_bi t s)$ चुनता है:
- Eviction (इविक्शन): एक 'कीप रेश्यो' (रखे गए टोकनों का अंश) जो 10% से 100% तक होता है।
- Quantization (क्वांटाइजेशन): Key (K) और Value (V) टेंसरों के लिए बिट-विड्थ (जैसे, 16, 8, या 4 बिट्स)।
Offline Sensitivity Calibration (ऑफलाइन संवेदनशीलता अंशांकन):
- सिस्टम प्रत्येक लेयर के लिए प्रत्येक संभावित कॉन्फ़िगरेशन के तहत फुल-प्रिसिजन अटेंशन आउटपुट और कंप्रेस्ड आउटपुट के बीच सापेक्ष L2 एरर (relative L2 error) को मापता है।
- यह एक संवेदनशीलता तालिका $S_{\ell, c}$ बनाता है जहाँ $\ell$ लेयर है और $c$ कॉन्फ़िगरेशन है।
- मुख्य निष्कर्ष: आक्रामक इविक्शन (75–90%) लेयर्स के बीच एरर वेरिएंस (error variance) 500–700× पैदा करता है, जबकि V-क्वांटाइजेशन का वेरिएंस न्यूनतम (~1.2×) है। यह पुष्टि करता है कि इविक्शन के लिए प्रति-लेयर रूटिंग आवश्यक है, लेकिन V-क्वांटाइजेशन के लिए यह कम महत्वपूर्ण है।
Greedy Budget Solver (ग्रीडी बजट सॉल्वर):
- वैश्विक इष्टतम कॉन्फ़िगरेशन खोजना एक कठिन असतत अनुकूलन (intractable discrete optimization) समस्या है ( $\approx 10^{42}$ उम्मीदवार एक 28-लेयर मॉडल के लिए)।
- MoE-nD एक ग्रीडी सॉल्वर का उपयोग करता है जो पुनरावृत्ति (iteratively) के माध्यम से उस लेयर को अपग्रेड करता है जो प्रति इकाई मेमोरी बचत पर उच्चतम मार्जिनल रिडक्शन इन सेंसिटिविटी ( $\Delta S / \Delta m$ ) प्रदान करती है, जब तक कि वैश्विक मेमोरी बजट $M$ समाप्त न हो जाए।
- यह प्रक्रिया CPU पर <50ms में पूरी होती है।
Heterogeneous Attention Patch (विषम अटेंशन पैच):
- इन्फरेंस के दौरान अलग-अलग कैश लंबाई और बिट-विड्थ वाली लेयर्स को संभालने के लिए, लेखक एक कस्टम अटेंशन कर्नेल लागू करते हैं।
- यह विशिष्ट लेयर के रिटेन किए गए टोकन के आधार पर टोकन इंडेक्स को ट्रैक करने के लिए प्रति-लेयर पोजीशन एरे का उपयोग करता है और RoPE (Rotary Positional Embedding) इनवर्जन लागू करता है, जिससे विषम कैश अवस्थाओं के बावजूद सही अटेंशन गणना सुनिश्चित होती है।

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

Joint Multi-Axis Routing: एकल वैश्विक मेमोरी बजट के तहत प्रति-लेयर आधार पर इविक्शन रेश्यो और K/V क्वांटाइजेशन बिट्स को संयुक्त रूप से रूट करने वाला पहला कार्य।
Empirical Discovery of Layer Heterogeneity: यह प्रदर्शित किया कि इविक्शन और क्वांटाइजेशन के बीच का ट्रेड-ऑफ वैश्विक नहीं बल्कि लेयर-विशिष्ट है। उदाहरण के लिए, लेयर 0 इविक्शन (कम एरर लागत) को पसंद कर सकती है, जबकि लेयर 6 समान मेमोरी बचत के लिए क्वांटाइजेशन (कम एरर लागत) को पसंद कर सकती है।
Ablation of Routing Axes: एक एब्लेशन चेन ( $2d_{uniform} \to 2d \to 2d_{hetero}$ ) के माध्यम से, लेखक सिद्ध करते हैं कि प्रति-लेयर इविक्शन रूटिंग प्रदर्शन लाभ का प्राथमिक चालक है, जबकि प्रति-लेयर क्वांटाइजेशन रूटिंग अकेले नगण्य सुधार प्रदान करती है।
Defined Scope of Applicability: पहचाना गया कि यह विधि केवल तभी महत्वपूर्ण लाभ देती है जब इनपुट लंबाई और बजट की कसावट (tightness) गैर-तुच्छ रूटिंग निर्णयों (अर्थात, लॉन्ग-कॉन्टेक्स्ट या टाइट-बजट परिदृश्य) को मजबूर करती है। छोटे-कॉन्टेक्स्ट परिदृश्यों में, सॉल्वर स्वाभाविक रूप से यूनिफॉर्म व्यवहार (keep=1.0) की ओर चला जाता है, जिससे विषमता निरर्थक हो जाती है।

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

विधि का मूल्यांकन DeepSeek-R1-Distill-Qwen-7B पर LongBench-v1 (16k टोकन), AIME (मैथ रीजनिंग), और MATH-500 पर किया गया।

LongBench-v1 (Long Context):
- MoE-nD ( $2d_{hetero}$ ) ने 14× कंप्रेशन (1.9 GB फुल कैश के मुकाबले 136 MB) प्राप्त किया और इसमें कोई पता लगाने योग्य सटीकता हानि नहीं (12.0 बनाम 11.5 बेसलाइन) हुई।
- समान मेमोरी बजट पर प्रतिस्पर्धी विधियाँ (सिंगल-अक्स इविक्शन या यूनिफॉर्म क्वांटाइजेशन) भारी सटीकता गिरावट का शिकार हुईं (8/100 से नीचे रहीं)।
- मैच किए गए मेमोरी (136 MB) पर, MoE-nD ने सर्वश्रेष्ठ नॉन-हेटरो बेसलाइन ( $2d$ ) को 6.1 अंकों से पछाड़ा।
AIME Reasoning Benchmarks:
- MoE-nD ने सभी बजटों में सभी बेसलाइन्स को लगातार पछाड़ दिया।
- टाइट बजट (जैसे, $b=64$ ) पर, MoE-nD ने 16.7–30% सटीकता प्राप्त की, जबकि अन्य विधियाँ लगभग 0% पर गिर गईं।
- विशेष रूप से, AIME-25 पर $b=256$ पर, MoE-nD ने फुल-कैश बेसलाइन को भी पीछे छोड़ दिया (36.7 बनाम 30.0), जो संभवतः आक्रामक इविक्शन के माध्यम से सेल्फ-डिस्ट्रैक्टिंग इंटरमीडिएट टोकनों को हटाने के कारण हुआ।
Null Results (स्कोप वैलिडेशन):
- MATH-500 और LongBench TREC (छोटे इनपुट) पर, MoE-nD ने यूनिफॉर्म बेसलाइन्स की तुलना में कोई लाभ नहीं दिखाया।
- कारण: ग्रीडी सॉल्वर ने निर्धारित किया कि ढीले मेमोरी बजट को देखते हुए अधिकांश लेयर्स के लिए 100% टोकन रखना (keep=1.0) इष्टतम था, जिससे हेटेरोजेनियस विधि यूनिफॉर्म व्यवहार में बदल गई। यह विधि की सैद्धांतिक सीमाओं को मान्य करता है।

5. महत्व और निहितार्थ (Significance and Implications)

दक्षता (Efficiency): बिना गुणवत्ता से समझौता किए उपभोक्ता हार्डवेयर (जैसे, 1.9 GB फुल कैश के मुकाबले 136 MB KV कैश) पर लंबे संदर्भ वाले LLMs को चलाने में सक्षम बनाता है।
पैराडाइम शिफ्ट (Paradigm Shift): KV कैश कंप्रेशन को "वन-साइज़-फिट्स-ऑल" से अनुकूली, लेयर-विशिष्ट अनुकूलन (adaptive, layer-specific optimization) की ओर ले जाता है।
व्यावहारिक अंतर्दृष्टि (Practical Insight): यह पेपर स्पष्ट करता है कि प्रति-लेयर क्वांटाइजेशन रूटिंग, प्रति-लेयर इविक्शन रूटिंग की तुलना में कम महत्वपूर्ण है, जो भविष्य के शोध को टोकन चयन रणनीतियों को अनुकूलित करने की ओर निर्देशित करता है।
सीमाएँ (Limitations): वर्तमान में इसके लिए एक कस्टम अटेंशन कर्नेल की आवश्यकता है (अभी तक FlashAttention संगत नहीं है) और विभिन्न मॉडल परिवारों के लिए पुन: अंशांकन (re-calibration) की आवश्यकता है (हालांकि लेखकों को उम्मीद है कि विषमता का पैटर्न आर्किटेक्चर केAcross बना रहेगा)।

निष्कर्षतः, MoE-nD यह प्रदर्शित करता है कि मल्टी-अक्स (multi-axis) के तहत इंटेलिजेंट, प्रति-लेयर मेमोरी संसाधन आवंटन ही संसाधन-बाधित उपकरणों पर उच्च-प्रदर्शन, लंबे-संदर्भ वाले इन्फरेंस को अनलॉक करने की कुंजी है।