🤖 machine learning

Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding

यह शोध पत्र SureLock को पेश करता है, जो Masked Diffusion-LMs के लिए एक ऐसी विधि है जो अनावश्यक गणनाओं को छोड़ने के लिए अभिसरित (converged) टोकन स्थितियों को गतिशील रूप से लॉक करती है और उनके कीज़ (keys) एवं वैल्यूज़ (values) को कैश करती है, जिससे प्रति-इटरेशन जटिलता $O(N^2d)$ से घटकर $O(MNd)$ हो जाती है और जनरेशन की गुणवत्ता से समझौता किए बिना 30–50% FLOP बचत प्राप्त होती है।

मूल लेखक: Daisuke Oba, Danushka Bollegala, Masahiro Kaneko, Naoaki Okazaki

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

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

CC BY 4.0

मूल लेखक: Daisuke Oba, Danushka Bollegala, Masahiro Kaneko, Naoaki Okazaki

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

कल्पना कीजिए कि आप एक विशाल, अराजक नाटक का निर्देशन कर रहे हैं जहाँ पटकथा वास्तविक समय में लिखी जा रही है। आपके पास N अभिनेताओं (वाक्य के टोकन) का एक दल है, और हर एक सेकंड में (जेनरेशन की हर प्रक्रिया में), आप हर एक अभिनेता से कहते हैं:

दूसरों को क्या कह रहे हैं, उसे ध्यान से सुनो।
अपनी अगली लाइन के बारे में सोचो।
अपनी नई लाइन ज़ोर से चिल्लाकर सुनाओ।

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

मिलिए "SURELOCK" से: "बैठ जाओ और शांत रहो" का नियम।

यह पेपर एक चतुर नई विधि पेश करता है जिसे SURELOCK कहा जाता है। इसे एक स्मार्ट निर्देशक के रूप में सोचें जो महसूस करता है कि एक बार जब कोई अभिनेता अपनी लाइन को लेकर आश्वस्त हो जाता है, तो उसे बार-बार रिहर्सल करने की ज़रूरत नहीं है।

यह कैसे काम करता है, यहाँ सरल उपमाओं में दिया गया है:

1. "स्थिर" अभिनेता (लॉकिंग मैकेनिज्म)

नाटक के बीच में, निर्देशक अभिनेताओं की जाँच करता है। यदि कोई अभिनेता कुछ सेकंड से एक ही बात कह रहा है और उसका आत्मविश्वास अधिक है (यानी "पोस्टीरियर स्थिर हो गया है"), तो निर्देशक कहता है:

"ठीक है, तुमने अपनी लाइन चुन ली है। बैठ जाओ। सोचना बंद करो। चिल्लाना बंद करो। बस वहीं बैठे रहो।"

यही है लॉक (Lock)। अभिनेता "लॉक" हो गया है। वह कठिन मानसिक कार्य (भारी गणित जिसे "फीड-फॉरवर्ड" और "क्वेरी प्रोजेक्शन" कहा जाता है) करना बंद कर देता है।

2. "भूतिया" उपस्थिति (Caching K/V)

लेकिन यहाँ एक जादुई ट्रिक है: भले ही अभिनेता बैठा हुआ है और कोई काम नहीं कर रहा है, फिर भी वह दृश्य का हिस्सा है।

निर्देशक उस अभिनेता द्वारा कही गई बातों का एक "भूतिया रिकॉर्डिंग" (उनका Keys और Values) रखता है। जो अन्य अभिनेता अभी भी खड़े हैं और काम कर रहे हैं, वे इस रिकॉर्डिंग को देख सकते हैं और कह सकते हैं, "ओह, मैंने देखा कि उस आदमी ने क्या कहा, इसलिए मैं अपनी लाइन को उसके अनुसार एडजस्ट करूँगा।"

तो, लॉक किए गए अभिनेता कोई काम नहीं करते, लेकिन वे फिर भी कहानी को प्रभावित करते हैं। सक्रिय अभिनेता अभी भी उन पर "अटेंड" (attend) कर सकते हैं।

3. परिणाम: घटता हुआ वर्कफोर्स

जैसे-जैसे नाटक आगे बढ़ता है, अधिक से अधिक अभिनेता लॉक होते जाते हैं।

नाटक की शुरुआत में: सभी खड़े हैं और काम कर रहे हैं। (उच्च लागत)।
नाटक के बीच में: आधा दल बैठा हुआ है। (मध्यम लागत)।
नाटक के अंत में: केवल कुछ ही अभिनेता जो अंतिम पंचलाइन लिख रहे हैं, खड़े हैं। बाकी चुपचाप बैठे हुए देख रहे हैं। (बहुत कम लागत)।

क्योंकि काम करने वाले लोगों की संख्या वाक्य पूरा होने के साथ कम होती जाती है, कंप्यूटर ऊर्जा की भारी बचत करता है। यह पेपर दिखाता है कि यह बिना कहानी की गुणवत्ता खराब किए, कंप्यूटिंग कार्य को 30% से 50% तक कम कर सकता है।

"सुरक्षा जाल" (जादू के पीछे का गणित)

आप चिंतित हो सकते हैं: "क्या होगा अगर अभिनेता बहुत जल्दी बैठ जाए और फिर उसे एहसास हो कि उसने गलती की है?"

लेखकों ने एक सुरक्षा जाल (Safety Net) बनाया है। वे किसी अभिनेता को केवल इसलिए लॉक नहीं करते क्योंकि वह आत्मविश्वासी दिखता है; वे उन्हें तभी लॉक करते हैं जब उनकी लाइन पूरी तरह से बदलना बंद हो जाती है। वे एक गणितीय परीक्षण का उपयोग करते हैं जिसे KL Divergence (एक "परिवर्तन डिटेक्टर") कहा जाता है।

नियम: यदि किसी अभिनेता की लाइन एक बहुत ही मामूली मात्रा (एक थ्रेशोल्ड जिसे $\epsilon$ कहा जाता है) से कम बदलती है, तो हम जानते हैं कि वे वास्तव में तैयार हैं।
गारंटी: पेपर गणितीय रूप से सिद्ध करता है कि यदि आप तब तक प्रतीक्षा करते हैं जब तक कि परिवर्तन इस स्तर तक छोटा न हो जाए, तो अंतिम कहानी उस कहानी के लगभग समान होगी जहाँ सभी काम कर रहे थे। "त्रुटि" इतनी छोटी है कि वह व्यावहारिक रूप से अदृश्य है।

यह एक बड़ी बात क्यों है?

कल्पना कीजिए कि आप एक केक बना रहे हैं।

पुराना तरीका: आप बैटर को हिलाते रहते हैं, चखते रहते हैं और हर एक मिनट के लिए ओवन का तापमान एडजस्ट करते रहते हैं, भले ही केक बन चुका हो और काउंटर पर रखा हो।
SURELOCK तरीका: एक बार जब केक बन जाता है और टाइमर बज जाता है, तो आप हिलाना बंद कर देते हैं। आप बस उसे वहीं रहने देते हैं जबकि आप केक के बिना पके हिस्सों पर फ्रॉस्टिंग करना पूरा करते हैं।

सारांश

SURELOCK AI के लिए एक ऐसी तकनीक है जो कहती है: "एक बार जब कोई शब्द तय हो जाए, तो उस पर ऊर्जा बर्बाद करना बंद करें।"

यह कंप्यूटर को अनावश्यक गणित करने से बचाता है, जेनरेशन की प्रक्रिया को तेज़ करता है, और लेखन की गुणवत्ता को पहले जैसा ही बनाए रखता है। यह घर के उन कमरों की लाइट बंद करने जैसा है जिनका उपयोग कोई नहीं कर रहा है, जबकि रसोई में जहाँ खाना बन रहा है, वहाँ लाइट चालू रखी जाती है।

यहाँ "Stopping Computation for Converged Tokens in Masked Diffusion-LM Decoding" (SURELOCK) नामक शोध पत्र का विस्तृत तकनीकी सारांश दिया गया है, जिसे ICLR 2026 के लिए एक कॉन्फ्रेंस पेपर के रूप में प्रकाशित किया गया है।

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

मास्क्ड डिफ्यूजन लैंग्वेज मॉडल्स (MDLMs) $T$ चरणों में एक टोकन अनुक्रम को पुनरावर्ती रूप से डिनोइज़िंग (denoising) करके टेक्स्ट जनरेट करते हैं। ऑटोरेग्रेसिव (AR) मॉडल्स के विपरीत, जो स्वाभाविक रूप से की/वैल्यू (K/V) स्टेट्स को कैश करते हैं और केवल अगले टोकन के लिए गणना करते हैं, मानक MDLM सैंपलिंग प्रत्येक चरण में सभी $N$ टोकन स्थितियों के लिए सेल्फ-अटेंशन और फीड-फॉरवर्ड (FFN) सबलेयर्स को पुन: कंप्यूट करती है।

बॉटलनेक (The Bottleneck): भले ही कई टोकन "अनमास्क" (डिकोड) हो चुके हों और उनकी पोस्टीरियर प्रोबेबिलिटी स्थिर हो गई हो, मॉडल अभी भी इन स्थिर टोकनों के लिए पूर्ण $O(N^2d)$ अटेंशन गणना और $O(Nd^2)$ FFN गणना करता रहता है।
अपव्यय (The Waste): इसके परिणामस्वरूप पर्याप्त कम्प्यूटेशनल रेडंडेंसी (redundancy) होती है, क्योंकि मॉडल बार-बार उन टोकनों को प्रोसेस करता है जो प्रभावी रूप से "कन्वर्ज" (converged) हो चुके हैं और अगले चरणों में बदलने की संभावना कम है।
पूर्व कार्य की सीमाएँ (Limitations of Prior Work): मौजूदा त्वरण विधियाँ कुल चरणों ( $T$ ) को कम करने या चरणों के बीच K/V वेक्टर्स के पुन: उपयोग (Reuse) पर ध्यान केंद्रित करती हैं। हालाँकि, ये विधियाँ प्रति-चरण स्थानिक जटिलता (per-step spatial complexity) को कम नहीं करती हैं; वे अंत तक सभी $N$ स्थितियों के लिए अटेंशन कंप्यूट करती रहती हैं, जिससे $O(N^2d)$ लागत बनी रहती है।

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

लेखक SURELOCK प्रस्तावित करते हैं, जो एक विधि है जो टोकन स्थितियों को तब गतिशील रूप से और स्थायी रूप से निष्क्रिय ("लॉक") कर देती है जब उनकी पोस्टीरियर डिस्ट्रीब्यूशन स्थिर हो जाती है।

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

लॉकिंग कंडीशन (Locking Condition): एक टोकन स्थिति $i$ $i$ को चरण $t^*$ $t^{*}$ पर लॉक किया जाता है यदि उसका स्थानीय पोस्टीरियर डिस्ट्रीब्यूशन स्थिर हो गया हो। प्राथमिक मानदंड वर्तमान पोस्टीरियर $p_t^{(i)}$ $p_{t}^{(i)}$ और पिछले पोस्टीरियर $p_{t-1}^{(i)}$ $p_{t - 1}^{(i)}$ के बीच स्टेप-वाइज कुलबैक-लीब्लर (KL) डायवर्जेंस है।
- शर्त: $D_t^{(i)} = \text{KL}(p_t^{(i)} \| p_{t-1}^{(i)}) \leq \epsilon$ ।
- एक माध्यमिक "कॉन्फिडेंस गेट" भी लागू किया जा सकता है ताकि यह सुनिश्चित हो सके कि टोकन अत्यधिक कॉन्फिडेंट (कम एंट्रॉपी वाला) भी है।
स्थायी निष्क्रियता (Permanent Deactivation): एक बार लॉक होने के बाद:
- उस स्थिति को भविष्य के सभी चरणों के लिए सक्रिय टोकनों ( $A_t$ ) के सेट से हटा दिया जाता है।
- मॉडल इस टोकन के लिए क्वेरी प्रोजेक्शन और फीड-फॉरवर्ड नेटवर्क (FFN) गणनाओं को स्किप (skip) कर देता है।
- लॉक किए गए टोकन के लिए की ( $K$ ) और वैल्यू ( $V$ ) वेक्टर्स को कैश (cache) किया जाता है।
निरंतर अटेंशन (Continued Attention): महत्वपूर्ण रूप से, सक्रिय (अनलॉक्ड) टोकन अभी भी लॉक किए गए टोकनों पर अटेंड (attend) कर सकते हैं। अटेंशन मैकेनिज्म लॉक किए गए पदों के लिए कैश किए गए $K/V$ वेक्टर्स का उपयोग करता है, जबकि सक्रिय उपसमूह के लिए नया $Q/K/V$ कंप्यूट करता है।

एल्गोरिद्मिक जटिलता (Algorithmic Complexity)

बेसलाइन (Baseline): अटेंशन के लिए $O(N^2d)$ और FFN के लिए $O(Nd^2)$ प्रति चरण।
SURELOCK: यदि $t$ $t$ चरण पर अनलॉक्ड (सक्रिय) टोकनों की संख्या $M_t$ $M_{t}$ है (जहाँ $M_t \leq N$ $M_{t} \leq N$ ):
- अटेंशन लागत: $O(M_t N d)$ (सक्रिय क्वेरीज सभी कीज़/वैल्यूज़ पर अटेंड करती हैं)।
- FFN लागत: $O(M_t d^2)$ (केवल सक्रिय टोकन FFN से गुजरते हैं)।
- जैसे-जैसे $t$ बढ़ता है, $M_t$ एकदिष्ट रूप से (monotonically) घटता है, जिससे कम्प्यूटेशनल लोड में महत्वपूर्ण कमी आती है।

सैद्धांतिक औचित्य (Theoretical Justification)

पेपर एक सैद्धांतिक बाउंड (Theorem 1) प्रदान करता है जो यह सिद्ध करता है कि टोकन को जल्दी लॉक करने के कारण अंतिम लॉग-प्रोबेबिलिटी में त्रुटि, लॉकिंग स्टेप पर उपयोग किए गए KL थ्रेशोल्ड के वर्गमूल द्वारा सीमित है:
$\|\log p_T^{(i)} - \log \hat{p}_T^{(i)}\|_\infty \leq C_{tail} \sqrt{\epsilon}$
यह लॉकिंग थ्रेशोल्ड $\epsilon$ और अंतिम आउटपुट डिस्ट्रीब्यूशन के अधिकतम विचलन के बीच एक क्लोज्ड-फॉर्म संबंध स्थापित करता है, जो लॉकिंग मानदंड के रूप में स्थानीय KL के उपयोग को न्यायसंगत बनाता है।

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

ऑर्थोगोनल ऑप्टिमाइजेशन एक्सिस (Orthogonal Optimization Axis): SURELOCK एक नया अनुकूलन आयाम पेश करता है (सक्रिय स्थितियों की एकदिष्ट कमी) जो मौजूदा स्टेप-काउंट रिडक्शन और K/V रियूज विधियों का पूरक है।
सैद्धांतिक गारंटी (Theoretical Guarantee): लॉकिंग स्टेप पर स्थानीय KL डायवर्जेंस और टर्मिनल लॉग-प्रोबेबिलिटी एरर को जोड़ने वाला एक क्लोज्ड-फॉर्म एरर बाउंड का व्युत्पन्न।
कुशल कार्यान्वयन (Efficient Implementation): एक ऐसा तंत्र जो K/V वेक्टर्स को कैश करता है ताकि सक्रिय टोकन लॉक किए गए टोकनों पर अटेंड कर सकें, जिससे यह सुनिश्चित होता है कि मॉडल का ग्लोबल कॉन्टेक्स्ट बना रहे और लॉक किए गए टोकनों की आंतरिक अवस्थाओं को पुन: कंप्यूट न करना पड़े।
अनुभवजन्य सत्यापन (Empirical Validation): LLaDA-8B मॉडल्स पर व्यापक प्रयोग जो जनरेशन क्वालिटी से समझौता किए बिना महत्वपूर्ण FLOP रिडक्शन प्रदर्शित करते हैं।

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

लेखकों ने LLaDA-8B-Base और LLaDA-8B-Instruct पर WikiText-103 (लैंग्वेज मॉडलिंग) और MT-Bench (इंस्ट्रक्शन फॉलोइंग) का उपयोग करके SURELOCK का मूल्यांकन किया।

कम्प्यूटेशनल दक्षता (Computational Efficiency):
- FLOP रिडक्शन: SURELOCK ने बेसलाइन सैंपलर की तुलना में एल्गोरिद्मिक FLOPs को 30% से 50% तक कम कर दिया।
- टेम्पोरल डायनेमिक्स (Temporal Dynamics): सैंपलिंग के साथ प्रगति करने पर सक्रिय टोकनों का अनुपात ( $M_t/N$ ) एकदिष्ट रूप से घटता है, जिससे बाद के चरणों में तेजी से बचत होती है (चित्र 2)।
- थ्रूपुट (Throughput): कंप्यूट-बाउंड परिदृश्यों (बड़े बैच साइज और लंबी सीक्वेंस) में एंड-टू-एंड टोकन-पर-सेकंड (TPS) में 1.39x तक सुधार हुआ, हालांकि कार्यान्वयन ओवरहेड्स (अनियमित मेमोरी एक्सेस) के कारण मेमोरी-बाउंड या छोटे-बैच वाले परिदृश्यों में लाभ मामूली थे।
जनरेशन क्वालिटी (Generation Quality):
- इंस्ट्रक्शन फॉलोइंग: MT-Bench पर, LLM-as-a-judge स्कोर विभिन्न सेटिंग्स में लगभग अपरिवर्तित रहा (अंतर $\leq 0.1$ अंक)।
- लैंग्वेज मॉडलिंग: WikiText-103 पर, जनरेटेड परप्लेक्सिटी (Gen.-PPL) में न्यूनतम गिरावट देखी गई। छोटे जनरेशन ( $N_{gen}=64$ ) में PPL में थोड़ी अधिक वृद्धि देखी गई, जो बताता है कि यह विधि लंबे कॉन्टेक्स्ट के लिए सबसे प्रभावी है।
- कोड जनरेशन: HumanEval बेंचमार्क पर, SURELOCK ने कम्प्यूटेशनल लागत को आधा करते हुए भी अपना Pass@1 स्कोर (18.6% बनाम 18.6-20.3%) बनाए रखा, जो सिद्ध करता है कि लॉकिंग सिंिटैक्टिक (syntactic) या लॉजिकल संरचनाओं को तोड़ता नहीं है।

5. महत्व (Significance)

लॉन्ग कॉन्टेक्स्ट के लिए स्केलेबिलिटी (Scalability for Long Contexts): चूंकि डिफ्यूजन LMs उच्च प्रति-चरण लागत के कारण लंबे कॉन्टेक्स्ट जनरेशन में संघर्ष करते हैं, SURELOCK प्रति-चरण लागत को एडेप्टिव और घटते हुए बनाकर इन मॉडल्स को लंबी सीक्वेंस तक स्केल करने का मार्ग प्रशस्त करता है।
हार्डवेयर एग्नोस्टिक (Hardware Agnosticism): यह विधि सैद्धांतिक एल्गोरिद्मिक जटिलता (FLOPs) को कम करती है, जो एक हार्डवेयर-एग्नोस्टिक मीट्रिक है। हालांकि वर्तमान PyTorch कार्यान्वयन मेमोरी एक्सेस पैटर्न के कारण रूढ़िवादी वॉल-क्लॉक स्पीडअप दिखाते हैं, पेपर स्पष्ट रूप से हार्डवेयर-विशिष्ट अनुकूलन (जैसे फ्यूज्ड कर्नेल, कंटिगस कैश लेआउट) के रास्ते बताता है ताकि गति के लाभों को पूरी तरह से प्राप्त किया जा सके।
पूरक प्रकृति (Complementary Nature): SURELOCK को मल्टीप्लिकेटिव एफिशिएंसी गेन प्राप्त करने के लिए अन्य त्वरण तकनीकों (जैसे समानांतर अनमास्किंग, K/V रियूज, या सिलेक्शन-आधारित स्पारसिफिकेशन) के साथ जोड़ा जा सकता है।

संक्षेप में, SURELOCK बुद्धिमानी से कन्वर्ज्ड टोकन की पहचान करके और उनके ग्लोबल अटेंशन मैकेनिज्म में योगदान को सुरक्षित रखते हुए उनके कंप्यूटेशन को फ्रीज करके, मास्क डिफ्यूजन डिकोडिंग में "वेस्टेड कंप्यूट" (wasted compute) की समस्या को प्रभावी ढंग से संबोधित करता है, जिससे आउटपुट क्वालिटी पर नगण्य प्रभाव के साथ पर्याप्त दक्षता प्राप्त होती है।