💬 NLP

Just on Time: Token-Level Early Stopping for Diffusion Language Models

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

मूल लेखक: Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

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

📖 4 मिनट में पढ़ें☕ कॉफ़ी ब्रेक में पढ़ें

CC BY 4.0

मूल लेखक: Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

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

कल्पना कीजिए कि आप एक विशाल जिगसॉ पहेली (jigsaw puzzle) को हल करने की कोशिश कर रहे हैं, लेकिन डिब्बे पर बनी तस्वीर को देखने के बजाय, आपके पास खाली, धूसर (gray) वर्गों से भरा एक डिब्बा है। इसे हल करने के लिए, आपको यह अनुमान लगाना होगा कि प्रत्येक वर्ग का रंग क्या होना चाहिए, फिर पूरी तस्वीर को देखना होगा, यह महसूस करना होगा कि कुछ अनुमान गलत थे, और फिर उन्हें दोबारा पेंट करना होगा। आप इस "अनुमान लगाओ और सुधारो" (guess-and-fix) चक्र को सैकड़ों बार दोहराते हैं जब तक कि तस्वीर अंततः सही न दिखने लगे। एक नया प्रकार का AI, जिसे डिफ्यूजन लैंग्वेज मॉडल (Diffusion Language Model) कहा जाता है, इसी तरह टेक्स्ट लिखता है। पुराने AI के विपरीत जो एक टाइप करने वाले की तरह एक बार में एक शब्द लिखते हैं, ये मॉडल एक खाली पन्ने से शुरुआत करते हैं और पूरे वाक्य को एक साथ परिष्कृत (refine) करते हैं, जैसे एक चित्रकार कैनवास पर धीरे-धीरे विवरण जोड़ता है।

समस्या यह है कि यह पेंटिंग की प्रक्रिया अविश्वसनीय रूप से धीमी है। भले ही AI वाक्य के आसान हिस्सों (जैसे "The" या "cat") को कुछ ही ब्रशस्ट्रोक के बाद समझ लेता है, फिर भी यह सुरक्षित रहने के लिए उन पर बार-बार पेंट करता रहता है। यह एक ऐसे शेफ की तरह है जो सूप चखता है, महसूस करता है कि वह एकदम सही है, लेकिन फिर भी परोसने से पहले एक घंटे तक उसे हिलाता और चखता रहता है। इससे कंप्यूटर की शक्ति और समय की भारी बर्बादी होती है। वैज्ञानिक खुद से सवाल पूछ रहे हैं: हम AI को यह कैसे बता सकते हैं, "हे, तुमने इस हिस्से को सही कर लिया है, इस पर काम करना बंद करो और कठिन चीजों पर ध्यान दो"?

यहीं पर JoT (Just on Time) नामक एक नई विधि आती है। JoT को एक बुद्धिमान सुपरवाइजर के रूप में सोचें जो AI चित्रकार के ऊपर खड़ा है। पूरी टीम को एक ही समय में पेंटिंग रोकने के लिए कहने के बजाय, JoT कैनवास पर प्रत्येक व्यक्तिगत शब्द को देखता है। जैसे ही AI को किसी विशिष्ट शब्द के बारे में पूरा विश्वास हो जाता है—मान लीजिए, वह 99% निश्चित है कि शब्द "cat" है—JoT चिल्लाता है, "इस शब्द को फ्रीज करो! इसे अब और मत छुओ!" और AI का ध्यान अगले धुंधले, अनिश्चित शब्द की ओर मोड़ देता है।

JoT के पीछे के शोधकर्ताओं ने पाया कि यह "टोकन-लेवल अर्ली स्टॉपिंग" (token-level early stopping) चमत्कार की तरह काम करती है। उन्होंने इसे दो शक्तिशाली AI मॉडलों, Dream-7B और LLaDA-8B पर चार अलग-अलग चुनौतियों में परखा: गणित की समस्याओं को हल करना, सामान्य ज्ञान (trivia) के उत्तर देना, वाक्यों को पूरा करना और कोड लिखना। परिणाम चौंकाने वाले थे। गणितीय तर्क की एक परीक्षा, GSM8K पर, JoT ने AI को 5.5 गुना तेज़ बना दिया, जबकि स्कोर में केवल मामूली गिरावट (लगभग 2.3 अंक) आई। कोडिंग चुनौती HumanEval पर, गति में सुधार और भी नाटकीय था, जो लगभग 20 गुना तेज़ (19.6×) तक पहुँच गया, फिर भी AI लगभग सभी उत्तर सही दे रहा था।

यह पेपर इस विचार के विरुद्ध तर्क देता है कि "एक ही आकार सबके लिए" (one-size-fits-all) वाला दृष्टिकोण अपनाया जाए, जहाँ AI एक ही समय में सभी शब्दों को रोक देता है। इसके बजाय, JoT एक चतुर तकनीक का उपयोग करता है: यह उन शब्दों के लिए आत्मविश्वास का स्तर कम कर देता है जो उन शब्दों के बगल में हैं जिन्हें AI पहले ही पूरा कर चुका है। यदि AI ने वाक्य की शुरुआत को सफलतापूर्वक सिद्ध कर लिया है, तो वह अगले शब्द पर थोड़ा जल्दी भरोसा कर सकता है। यह AI को अपनी ऊर्जा केवल उन हिस्सों पर केंद्रित करने की अनुमति देता है जो वास्तव में भ्रमित करने वाले हैं, जिससे वह आसान हिस्सों पर किए जाने वाले अनावश्यक काम को छोड़ देता है।

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

तकनीकी सारांश: डिफ्यूजन लैंग्वेज मॉडल्स के लिए जस्ट ऑन टाइम (JoT)

समस्या विवरण

डिफ्यूजन लैंग्वेज मॉडल्स (DLMs) एक पूरी तरह से मास्क किए गए अनुक्रम (sequence) को एक स्वच्छ आउटपुट में पुनरावृत्त रूप से परिष्कृत करके टेक्स्ट उत्पन्न करते हैं। हालांकि यह दृष्टिकोण समानांतर टोकन भविष्यवाणी और द्विदिश संदर्भ एकीकरण (bidirectional context integration) को सक्षम करता है, लेकिन यह डिकोडिंग अक्षमता से ग्रस्त है। मानक DLMs को पूरे अनुक्रम के लिए निश्चित संख्या में रिफाइनमेंट स्टेप्स ( $T$ ) की आवश्यकता होती है, भले ही कई टोकन अंतिम चरण से बहुत पहले ही स्थिर, उच्च-विश्वास भविष्यवाणियों तक पहुँच जाते हैं। इससे अनावश्यक कम्प्यूटेशनल ओवरहेड होता है, क्योंकि मॉडल उन टोकनों को परिष्कृत करना जारी रखता है जो पहले ही अभिसरण (converge) कर चुके हैं। मौजूदा अर्ली-स्टॉपिंग विधियाँ अक्सर वैश्विक मानदंडों को लागू करती हैं या प्रशिक्षण-विशिष्ट अनुकूलन पर निर्भर करती हैं, जो एकल पीढ़ी के भीतर व्यक्तिगत टोकनों की परिवर्तनशील अभिसरण दरों को गतिशील रूप से संबोधित करने में विफल रहती हैं।

कार्यप्रणाली: JoT (जस्ट ऑन टाइम)

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

मुख्य घटक

कॉन्फिडेंस मेट्रिक (Confidence Metric):
JoT शीर्ष दो अनुमानित संभावनाओं ( $p_1$ और $p_2$ ) के अनुपात का उपयोग करके प्रत्येक मास्क किए गए स्थान $i$ पर भविष्यवाणी की निश्चितता को मापता है:
$r_i = \frac{p_1}{p_2 + \epsilon}$
यह मेट्रिक अनस्केल्ड लॉजिट्स (unscaled logits) से प्राप्त होता है, जो इसे सैंपलिंग टेम्परेचर के प्रति अपरिवर्तनीय बनाता है। एक उच्च अनुपात एक निर्णायक प्राथमिकता को दर्शाता है, जबकि 1 के करीब का अनुपात अनिश्चितता का संकेत देता है।
स्थानिक मॉड्यूलेशन (Spatial Modulation):
यह पहचानते हुए कि पहले से हल किए गए (unmasked) संदर्भ के निकट स्थित टोकन समृद्ध स्थानीय जानकारी और कम प्रभावी शोर (noise) से लाभान्वित होते हैं, JoT एक स्थानिक रूप से अनुकूलित थ्रेशोल्ड पेश करता है।
- एक ज्योमेट्रिक कर्नेल एक विंडो $D$ के भीतर अनमास्क किए गए टोकनों से स्थिति $i$ की निकटता के आधार पर एक भार $w_i$ की गणना करता है।
- इस भार को एक सॉफ्टनिंग फैक्टर $\phi_i \in [0, 1]$ में सामान्यीकृत किया जाता है।
- स्थिति $i$ के लिए स्वीकृति थ्रेशोल्ड $\tau_i$ को एक अधिकतम थ्रेशोल्ड $\tau_{max}$ और एक न्यूनतम $\tau_{min}$ के बीच इंटरपोलेट किया जाता है:
  $\tau_i = \tau_{max} - (\tau_{max} - \tau_{min}) \cdot \phi_i$
- हल किए गए संदर्भ के पास के टोकन कम थ्रेशोल्ड प्राप्त करते हैं (जल्दी फाइनल होने में आसान), जबकि मास्क किए गए क्षेत्रों के केंद्र में स्थित टोकन सख्त थ्रेशोल्ड बनाए रखते हैं।
अर्ली-एग्जिट निर्णय (Early-Exit Decision):
प्रत्येक चरण में, यदि किसी मास्क किए गए टोकन का कॉन्फिडेंस रेशियो $r_i$ उसके अनुकूलित थ्रेशोल्ड $\tau_i$ से अधिक हो जाता है, तो टोकन को फाइनल (argmax या सैंपलिंग के माध्यम से चुना जाता है) कर दिया जाता है और मास्क किए गए स्थानों के सेट से हटा दिया जाता है। यह प्रक्रिया तब तक चलती है जब तक कि सभी टोकन फाइनल न हो जाएं या अधिकतम स्टेप काउंट तक न पहुँच जाएँ।

सैद्धांतिक विश्लेषण

पेपर "लीडिंग-पेयर रिवर्सल" (leading-pair reversals)—उन मामलों के संबंध में (जहाँ अधिक संदर्भ प्रकट होने पर एक प्रतिबद्ध टोकन को बाद में रनर-अप द्वारा पीछे छोड़ दिया जाता है) एक औपचारिक विश्लेषण (प्रपोजिशन 1) प्रदान करता है। यह मानते हुए कि मॉडल का अनुमान वास्तविक कंडीशनल डिस्ट्रीब्यूशन (कैलिब्रेशन) से मेल खाता है, ऐसे रिवर्सल की संभावना $2/(1+\tau)$ द्वारा सीमित है। लेखक नोट करते हैं कि उच्च थ्रेशोल्ड पर अवशिष्ट त्रुटियां (residual errors) कैलिब्रेशन की विफलता से उत्पन्न हो सकती हैं, न कि स्वयं स्टॉपिंग तंत्र से।

प्रमुख योगदान

JoT फ्रेमवर्क: एक ट्रेनिंग-फ्री, प्रति-टोकन अर्ली-स्टॉपिंग विधि जो हल किए गए संदर्भ की स्थानिक निकटता के आधार पर स्वीकृति थ्रेशोल्ड को अनुकूलित करती है।
व्यापक मूल्यांकन: दो अत्याधुनिक मॉडल्स (Dream-7B और LLaDA-8B) पर चार विविध बेंचमार्क: GSM8K (गणित), MMLU (ज्ञान), HellaSwag (सामान्य ज्ञान), और HumanEval (कोड) पर प्रयोग किए गए।
एब्लेशन अध्ययन (Ablation Studies): थ्रेशोल्ड चयन और स्थानिक मॉड्यूलेशन के प्रभावों को अलग करने वाला विस्तृत विश्लेषण, जो हाइपरपैरामीटर कॉन्फ़िगरेशन के लिए व्यावहारिक मार्गदर्शन प्रदान करता है।

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

लेखकों ने फुल डिकोडिंग और अन्य अर्ली-स्टॉपिंग बेसलाइन्स (Prophet, KLASS, Fast-dLLM) के विरुद्ध JoT का मूल्यांकन किया।

स्पीड-क्वालिटी ट्रेड-ऑफ:
- Dream-7B: JoT ने GSM8K पर 5.54× स्पीडअप और HumanEval पर 19.60× स्पीडअप प्राप्त किया। सटीकता में गिरावट मामूली थी (जैसे, GSM8K पर -2.3 अंक, HumanEval पर -0.6), प्रदर्शन अधिकांश बेंचमार्क पर फुल डिकोडिंग के ~3 प्रतिशत अंकों के भीतर बना रहा।
- LLaDA-8B: JoT ने सभी बेंचमार्क में फुल डिकोडिंग के 3.1 अंकों के भीतर सटीकता बनाए रखी, जबकि प्रतिस्पर्धी स्पीडअप (जैसे, GSM8K पर 3.75×) प्रदान किया।
तुलना: JoT आम तौर पर Prophet (जो बड़ी सटीकता गिरावट से ग्रस्त था) से बेहतर प्रदर्शन करता है और कई कार्यों पर KLASS और Fast-dLLM की तुलना में बेहतर स्पीड-एक्यूरेसी पारेटो फ्रंटियर्स (Pareto frontiers) प्रदान करता है।
कंपोजेबिलिटी (Composability): JoT को सफलतापूर्वक Fast-dLLM के DualCache (जो प्रति-चरण विलंबता को कम करता है) के साथ एकीकृत किया गया, जो यह दर्शाता है कि स्टेप रिडक्शन और प्रति-चरण अनुकूलन ऑर्थोगोनल (orthogonal) हैं और उन्हें जोड़ा जा सकता है।
विफलता मोड (Failure Modes): 500 GSM8K नमूनों के विश्लेषण से पता चला कि त्रुटियां मुख्य रूप से तब हुईं जब मॉडल ने एक महत्वपूर्ण रीजनिंग ब्रांच पर समय से पहले एक टोकन को चुन लिया (उदाहरण के लिए, "/ 2" के बजाय "× 2" चुनना)। ये त्रुटियां कोई नई विफलता नहीं थीं बल्कि फुल डिकोडिंग की तुलना में अलग दर पर हुईं।

महत्व और दावे

पेपर का दावा है कि JoT बिना री-ट्रेनिंग या टास्क-विशिष्ट फाइन-ट्यूनिंग के डिफ्यूजन लैंग्वेज मॉडल्स की डिकोडिंग अक्षमता का एक सम्मोहक समाधान प्रदान करता है। जैसे ही टोकन पर्याप्त आत्मविश्वास प्रदर्शित करते हैं, उन्हें गतिशील रूप से फाइनल करके, JoT आवश्यक कुल डिफ्यूजन स्टेप्स की संख्या को काफी कम कर देता है।

लेखक JoT को एक ऐसी विधि के रूप में देखते हैं जो Dream-7B पर परीक्षण किए गए अर्ली-स्टॉपिंग तरीकों के बीच सर्वश्रेष्ठ स्पीड-क्वालिटी ट्रेड-ऑफ प्राप्त करती है और LLaDA-8B पर प्रतिस्पर्धी बनी रहती है। वे इस बात पर जोर देते हैं कि यह विधि विशेष रूप से प्रभावी है क्योंकि यह मॉडल को "स्पष्ट" संदर्भ टोकनों के अनावश्यक रिफाइनमेंट को छोड़ने की अनुमति देती है जबकि जटिल रीजनिंग चरणों के लिए आवश्यक कम्प्यूटेशनल प्रयास को सुरक्षित रखती है। पेपर निष्कर्ष निकालता है कि हालांकि इष्टतम थ्रेशोल्ड कार्य के अनुसार भिन्न हो सकते हैं, हाइपरपैरामीटर के प्रति विधि की मजबूती इसे विविध अनुप्रयोगों के लिए व्यावहारिक बनाती है।