PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding
यह शोध पत्र PRESTO को प्रस्तुत करता है, जो एक सिद्धांत-आधारित ढांचा है जो डिफ्यूजन मार्जिनल्स और ऑटोरेग्रेसिव वेरिफिकेशन के बीच के बेमेल को हल करने के लिए प्रीफिक्स-अलाइन्ड स्कोरिंग और प्रायोरिटी-बेस्ड ट्री सर्च को लागू करके डिफ्यूजन-आधारित स्पेक्युलेटिव डिकोडिंग को बढ़ाता है, जिससे एंड-टू-एंड थ्रूपुट में महत्वपूर्ण सुधार होता है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक कहानी में अगले शब्द की भविष्यवाणी करने की कोशिश कर रहे हैं। लंबे समय तक, सबसे स्मार्ट कंप्यूटर (जिन्हें लार्ज लैंग्वेज मॉडल्स कहा जाता है) एक बार में एक शब्द करके ये कहानियाँ लिखते थे, जैसे कोई व्यक्ति किताब को ज़ोर से पढ़ रहा हो, और हर एक शब्द के बाद रुककर यह सोचता हो कि आगे क्या आना चाहिए। यह सटीक तो है, लेकिन धीमा है। हाल ही में, वैज्ञानिकों ने "डिफ्यूजन" (diffusion) मॉडल्स का उपयोग करके इन कहानियों को लिखने का एक नया तरीका खोजा है। इसे ऐसे समझें जैसे एक मूर्तिकार जो एक पत्थर के ब्लॉक से शुरू करता है और एक ही बार में पूरी मूर्ति को तराश देता है, बजाय इसके कि वह एक-एक छोटा टुकड़ा काटता रहे। यह कंप्यूटर को कई शब्दों का अनुमान एक साथ लगाने की अनुमति देता है, जो अविश्वसनीय रूप से तेज़ है।
हालाँकि, इसमें एक पेंच है। जब आप एक साथ कई शब्दों का अनुमान लगाते हैं, तो आप कुछ गलतियाँ कर सकते हैं। इसे ठीक करने के लिए, एक चतुर तकनीक है जिसे "स्पेक्टिवल डिकोडिंग" (speculative decoding) कहा जाता है। यह एक तेज़, जूनियर असिस्टेंट की तरह है जो अगले कुछ शब्दों का अनुमान लगाता है, और फिर एक धीमा, सुपर-स्मार्ट बॉस यह जाँचता है कि उसके अनुमान सही हैं या नहीं। यदि बॉस सहमत होता है, तो वह शब्दों के पूरे समूह को तुरंत स्वीकार कर लेता है, जिससे बहुत समय बचता है। समस्या यह है कि जूनियर असिस्टेंट (डिफ्यूजन मॉडल) व्यक्तिगत शब्दों का अनुमान लगाने में तो माहिर है, लेकिन उसे हमेशा यह नहीं पता होता कि वे शब्द एक विशिष्ट क्रम में कैसे फिट बैठते हैं। यह वैसा ही है जैसे असिस्टेंट केक के लिए अलग-अलग सामग्री चुनने में तो बहुत अच्छा है, लेकिन उसे यह नहीं पता कि सामग्रियों का कौन सा संयोजन वास्तव में स्वादिष्ट होगा।
यहीं पर एक नया शोध पत्र आता है। शोधकर्ताओं ने, जो ज़ेंग वांग और उनके सहयोगियों के नेतृत्व में थे, महसूस किया कि इन तेज़ सहायकों का उपयोग करने का वर्तमान तरीका बहुत सारी गति को बर्बाद कर रहा है। उन्होंने पाया कि जबकि असिस्टेंट शब्दों के संयोजन की एक विशाल विविधता उत्पन्न कर सकता है, वर्तमान विधि केवल एक एकल पथ (path) की जाँच करती है, जैसे एक ही गलियारे में चलना और यह उम्मीद करना कि अंत में दरवाज़ा खुला होगा। लेखक एक नया सिस्टम प्रस्तावित करते हैं जिसे PRESTO (प्रिफिक्स-अलाइन्ड ट्री ड्राफ्टिंग) कहा जाता है। एक ही गलियारे में चलने के बजाय, PRESTO संभावनाओं का एक पेड़ (tree) बनाता है, और एक साथ कई अलग-अलग रास्तों की खोज करता है। लेकिन यहाँ जादू है: यह इस बात के तरीके को ठीक करता है कि असिस्टेंट के आत्मविश्वास (confidence) को कैसे मापा जाता है। असिस्टेंट का मूल आत्मविश्वास "प्रिफिक्स-ब्लाइंड" (prefix-blind) है, जिसका अर्थ है कि उसे इससे कोई फर्क नहीं पड़ता कि पहले कौन से शब्द आए थे। PRESTO एक "प्रिफिक्स-अलाइन्ड" (prefix-aligned) स्कोर जोड़ता है, जो एक दिशा-सूचक यंत्र (compass) की तरह काम करता है, यह सुनिश्चित करता है कि खोज के लिए चुने गए पथ वे हैं जिन्हें बॉस द्वारा स्वीकार किए जाने की सबसे अधिक संभावना है।
परिणामस्वरूप, यह प्रणाली काफी तेज़ है। अपने परीक्षणों में, PRESTO ने कंप्यूटर को प्रत्येक अनुमान के दौर में अधिक शब्दों को स्वीकार करने में मदद की। मौजूदा बेहतरीन सेटअपों में से कुछ पर, इसने पूरी प्रक्रिया को 1.5 गुना तेज़ बना दिया। अन्य सेटअपों पर, इसने 1.12 गुना की गति वृद्धि प्रदान की। शोध पत्र सुझाव देता है कि अनुमान लगाने की प्रक्रिया को एक सीधी रेखा के बजाय एक पेड़ जैसी शाखाओं वाली साहसिक यात्रा की तरह मानकर, हम सर्वश्रेष्ठ दोनों चीजें पा सकते हैं: डिफ्यूजन की गति और सावधानीपूर्वक जाँच की सटीकता।
समस्या: "एक-पथ" का जाल (The "One-Path" Trap)
यह समझने के लिए कि PRESTO की आवश्यकता क्यों है, कल्पना कीजिए कि आप अपने एक दोस्त के साथ "मैड लिब्स" (Mad Libs) खेल रहे हैं जो गायब शब्दों का अनुमान लगाने की कोशिश कर रहा है। आपका दोस्त एक डिफ्यूजन मॉडल है। वे खाली जगह को देखकर यह कहने में माहिर हैं, "मुझे लगता है कि यहाँ शब्द 'बिल्ली' होगा!" या "शायद 'कुत्ता' होगा?" या "या शायद 'रॉकेट'?" वे इन सभी विकल्पों को एक ही समय में चिल्लाकर बता सकते हैं।
हालाँकि, इस दोस्त का उपयोग करने का वर्तमान तरीका बहुत कठोर है। यह दोस्त के शीर्ष अनुमान को लेता है, उसे लिखता है, और फिर "बॉस" (टारगेट मॉडल) से पूछता है कि क्या उसका अनुमान सही है। यदि बॉस कहता है "नहीं," तो पूरी चीज़ को फेंक दिया जाता है, और आपको फिर से शुरू करना पड़ता है। यदि बॉस कहता है "हाँ," तो आप अगले शब्द पर जाते हैं और दोहराते हैं। इसे लीनियर ड्राफ्टिंग (linear drafting) कहा जाता है। यह एक जंगल में चलने जैसा है जहाँ आप केवल अपने ठीक सामने वाले रास्ते को देखते हैं।
लेखकों ने देखा कि यह दृष्टिकोण अक्षम है। क्योंकि डिफ्यूजन मॉडल एक साथ कई विकल्प उत्पन्न करता है, इसलिए संभावनाओं का एक विशाल "कॉम्बिनेटोरियल स्पेस" (combinatorial space) मौजूद होता है। यह एक मानचित्र की तरह है जिसमें हज़ार अलग-अलग रास्ते हैं, लेकिन आपको केवल एक ही रास्ते पर चलने की अनुमति है। शोध पत्र दिखाता है कि केवल एक पथ पर टिके रहकर, सिस्टम कई वैध रास्तों को खो देता है। वास्तव में, GSM8K जैसे गणित के सवालों पर, वर्तमान विधि औसतन लगभग 6.5 शब्द स्वीकार करती थी, लेकिन शोधकर्ताओं ने गणना की कि यदि वे सभी सर्वोत्तम पथों की जाँच कर पाते, तो वे लगभग 10 शब्द स्वीकार कर सकते थे। यह एक बहुत बड़ा अंतर है!
बेमेल: "अंधा" दिशा-सूचक यंत्र (The "Blind" Compass)
शोधकर्ताओं ने गहराई से जांच की और पाया कि केवल अधिक पथों की जाँच करना (एक पेड़ बनाना) पुराने तरीकों के साथ पूरी तरह से काम क्यों नहीं कर रहा था। उन्होंने एक "मौलिक बेमेल" (fundamental mismatch) की पहचान की।
मानक AI (ऑटोरेग्रेसिव मॉडल्स) की दुनिया में, एक शब्द के लिए आत्मविश्वास स्कोर इस बात पर बहुत अधिक निर्भर करता है कि उससे पहले कौन से शब्द आए थे। यदि वाक्य है "बिल्ली ... पर बैठी है", तो मॉडल जानता है कि "मैट" एक बहुत ही संभावित अगला शब्द है, लेकिन "पिज्जा" नहीं है। यह प्रिफिक्स-अलाइन्ड (prefix-aligned) है।
लेकिन डिफ्यूजन मॉडल अलग तरह से काम करते हैं। वे प्रत्येक स्थिति के लिए स्वतंत्र रूप से एक "मार्जिनल" (marginal) प्रायिकता उत्पन्न करते हैं। यह ऐसा है जैसे मॉडल कहता है, "स्थिति 5 पर, 'बिल्ली' की संभावना 80% है," बिना इस पर ध्यान दिए कि स्थिति 4 पर "द" (The) था या "द क्विक ब्राउन" (The quick brown)। यह प्रिफिक्स-ब्लाइंड (prefix-blind) है।
जब आप अनुमानों का एक पेड़ बनाने के लिए इन अंधे स्कोर का उपयोग करते हैं, तो आपको रैंकिंग की समस्या होती है। आप एक ऐसा पथ चुन सकते हैं जो पहले शब्द के लिए बहुत अच्छा दिखता है लेकिन दूसरे शब्द के लिए बहुत बुरा है क्योंकि मॉडल को यह एहसास नहीं हुआ कि पहले शब्द ने संदर्भ बदल दिया है। यह एक GPS की तरह है जो केवल वर्तमान सड़क के आधार पर आपको निर्देश देता है, इस तथ्य को अनदेखा करते हुए कि आपने अभी-पात मोड़ लिया है और अब आप एक वन-वे सड़क पर हैं। शोध पत्र का तर्क है कि इन अंधे स्कोरों का उपयोग करके एक पेड़ बनाना "अविश्वसनीय पथ रैंकिंग" (unreliable path ranking) की ओर ले जाता है, जिसका अर्थ है कि सिस्टम गलत शाखाओं की खोज करता है और समय बर्बाद करता है।
समाधान: PRESTO
PRESTO (प्रिफिक्स-अलाइन्ड स्कोरिंग एंड प्रायोरिटी-बेस्ड ट्री सर्च फॉर डिफ्यूजन स्पेकुलेटिव डिकोडिंग) डिफ्यूजन मॉडल के स्कोर में एक "सुधार" जोड़कर इसे हल करता है।
- प्रिफिक्स-अलाइन्ड स्कोरिंग: लेखकों ने महसूस किया कि उन्हें डिफ्यूजन मॉडल के मजबूत "मार्जिनल" सिग्नल (कैसे एक शब्द अपने आप में संभावित है) को "प्रिफिक्स-कंडीशन्ड" सिग्नल (पिछले शब्दों के आधार पर वह कितना संभावित है) के साथ जोड़ने की आवश्यकता है। उन्होंने एक नया स्कोरिंग फॉर्मूला बनाया जो डिफ्यूजन प्रायिकता को एक सरल n-ग्राम मॉडल (एक हल्का उपकरण जो शब्दों के संयोजन को देखता है) से प्राप्त सुधार कारक के साथ गुणा करता है। यह एक ऐसा स्कोर बनाता है जो कहानी के प्रवाह का सम्मान करता है।
- प्रायोरिटी-बेस्ड ट्री सर्च: केवल शीर्ष पथ चुनने के बजाय, PRESTO एक पेड़ बनाता है। यह नए, सुधारे गए स्कोर का उपयोग यह तय करने के लिए करता है कि कौन सी शाखाएँ बढ़ानी हैं। यह उन पथों को प्राथमिकता देता है जिनमें बॉस द्वारा स्वीकार किए जाने की उच्चतम क्षमता होती है। यह एक हाइकर (पर्वतारोही) की तरह है जो, केवल सीधे चलने के बजाय, मानचित्र देखता है और उस रास्ते को चुनता है जिसके शिखर तक पहुँचने की सबसे अधिक संभावना है, भले ही वह रास्ता शुरुआत में सबसे स्पष्ट न हो।
शोध पत्र ने इस पेड़ को विकसित करने के दो तरीके होते देखे: बीम सर्च (प्रत्येक चरण में निश्चित संख्या में शीर्ष पथ रखना) और बेस्ट-फर्स्ट सर्च (हमेशा अब तक मिले सबसे अच्छे एकल पथ का विस्तार करना)। उन्होंने पाया कि उनके विशिष्ट सेटअप के लिए, बीम सर्च उतना ही अच्छा काम करता है जितना कि अधिक जटिल बेस्ट-फर्स्ट सर्च, इसलिए वे सरल और अधिक कुशल विकल्प पर टिके रहे।
परिणाम: तेज़ और स्मार्ट
लेखकों ने गणितीय समस्याओं (GSM8K, Math500), कोडिंग चुनौतियों (HumanEval, LiveCodeBench), और चैट बातचीत सहित विभिन्न कार्यों पर PRESTO का परीक्षण किया। उन्होंने दो प्रकार के सिस्टम का उपयोग किया:
- डेडीकेटेड डिफ्यूजन ड्राफ्टर्स: एक बड़ा ऑटोरेग्रेसिव मॉडल के लिए अनुमान लगाने वाला एक छोटा, तेज़ डिफ्यूजन मॉडल (जैसे dFlash)।
- सेल्फ-स्पेक्टिवल डिफ्यूजन LLMs: एक ही डिफ्यूजन मॉडल जो खुद ही अनुमान लगाता है और खुद ही जाँच करता है (जैसे Nemotron-Labs-Diffusion)।
परिणाम सभी क्षेत्रों में सुसंगत थे। PRESTO ने लगातार औसत स्वीकृति लंबाई (Average Acceptance Length) को बढ़ाया, जो शब्दों की वह संख्या है जिसे बॉस एक बार में स्वीकार करता है।
- dFlash सिस्टम (Qwen3-8B का उपयोग करते हुए) पर, औसत स्वीकृति लंबाई लगभग 6.6 शब्दों से बढ़कर 9.6 शब्द हो गई। इसने एंड-टू-एंड गति में 1.5 गुना की वृद्धि की।
- Nemotron-Labs-Diffusion सिस्टम पर, स्वीकृति लंबाई 8.8 से बढ़कर 9.9 शब्द हो गई, जिसके परिणामस्वरूप 1.12 गुना की गति वृद्धि हुई।
सबसे प्रभावशाली बात यह है कि शोध पत्र ने दिखाया कि PRESTO तब भी काम करता है जब सिस्टम "स्टोकेस्टिक" (यानी रैंडम/अनिश्चित) होता है, जो आमतौर पर अनुमान लगाना कठिन होता है। इन मामलों में, गति वृद्धि और भी अधिक स्पष्ट थी, जिसमें कुछ बेंचमार्क में थ्रूपुट लगभग दोगुना देखा गया।
लेखकों ने इस नए तरीके की "लागत" की भी जाँच की। उन्होंने पाया कि पेड़ बनाने और नए स्कोर की गणना करने के लिए आवश्यक अतिरिक्त कार्य बहुत कम था—कुल समय का 4% से भी कम। अधिकांश समय (90% से अधिक) अभी भी बॉस द्वारा सत्यापन (verification) पर खर्च होता है। इसका मतलब है कि PRESTO एक अत्यधिक कुशल अपग्रेड है जो सिस्टम को ओवरहेड के साथ धीमा नहीं करता है।
PRESTO क्या नहीं है
यह ध्यान रखना महत्वपूर्ण है कि शोध पत्र क्या दावा नहीं करता है। लेखक स्पष्ट रूप से कहते हैं कि केवल एक साधारण (naive) ट्री संरचना लागू करना (उनके प्रिफिक्स-अलाइन्ड स्कोरिंग के बिना) उप-इष्टतम (suboptimal) है। यदि आप केवल डिफ्यूजन मॉडल के कच्चे स्कोर का उपयोग करके एक पेड़ बनाते हैं, तो आपको पूर्ण लाभ नहीं मिलता क्योंकि स्कोर "ब्लाइंड" प्रकृति के होते हैं। PRESTO विशेष रूप से उस स्कोरिंग बेमेल को ठीक करने के बारे में है।
इसके अलावा, शोध पत्र यह दावा नहीं करता है कि उसने डिफ्यूजन मॉडल्स की समस्या को पूरी तरह से हल कर दिया है। वे स्वीकार करते हैं कि उनका तरीका सुधार करने के लिए एक "ट्रैक्टेबल प्रिफिक्स-अलाइन्ड सिग्नल" (जैसे कि उनके द्वारा उपयोग किया गया n-gram मॉडल) पर निर्भर करता है। वे सुझाव देते हैं कि भविष्य के कार्य और भी समृद्ध संकेतों की खोज कर सकते हैं, लेकिन फिलहाल, उनका सरल सुधार बड़े लाभ देखने के लिए पर्याप्त है।
यह क्यों महत्वपूर्ण है
AI को तेज़ और अधिक कुशल बनाने की दौड़ में, गति का हर अंश मायने रखता है। स्पेक्टिवल डिकोडिंग एक चर्चित विषय रहा है क्योंकि यह हमें सर्वोत्तम, सबसे सटीक मॉडल का उपयोग करने की अनुमति देता है बिना उनके पूर्ण समय दंड (time penalty) का भुगतान किए। हालाँकि, वर्तमान विधियाँ इस तथ्य से सीमित थीं कि वे डिफ्यूजन मॉडल्स को रैखिक मशीनों की तरह मानते थे, उनके एक साथ कई विकल्प उत्पन्न करने की अनूठी क्षमता को अनदेखा करते थे।
PRESTO गेम बदल देता है क्योंकि यह डिफ्यूजन मॉडल्स को उनके वास्तविक स्वरूप—एक मल्टी-पाथ एक्सप्लोरर—के रूप में देखता है। स्कोरिंग को बॉस के काम की जाँच करने के तरीके के साथ संरेखित करके, यह डिफ्यूजन मॉडल के समानांतर उत्पादन (parallel generation) की पूरी क्षमता को अनलॉक करता है। परिणाम एक ऐसी प्रणाली है जो न केवल थोड़ी तेज़ है, बल्कि काफी अधिक कुशल भी है, जो टेक्स्ट, गणित की समस्याओं और कोड को ऐसी गति से उत्पन्न करने की अनुमति देती है जो पहले इन प्रकार के मॉडल्स के लिए असंभव मानी जाती थी। जैसा कि लेखकों ने कहा, उन्होंने एक "सिंगल-पाथ" वॉक को एक "ट्री-बेस्ड" अभियान में बदल दिया है, यह सुनिश्चित करते हुए कि उठाया गया हर कदम सही उत्तर की ओर एक कदम है।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।