💬 NLP

Improving Sampling for Masked Diffusion Models via Information Gain

यह शोध पत्र इन्फो-गेन सैंपलर (Info-Gain Sampler) प्रस्तुत करता है, जो मास्क्ड डिफ्यूजन मॉडल्स (Masked Diffusion Models) के लिए एक सिद्धांत-आधारित डिकोडिंग फ्रेमवर्क है जो उनकी नॉन-कॉज़ल प्रकृति का लाभ उठाते हुए तात्कालिक अनिश्चितता और भविष्य के सूचना लाभ के बीच संतुलन बनाने के लिए बनाया गया है, जिससे यह रीजनिंग, कोडिंग और क्रिएटिव राइटिंग जैसे विविध कार्यों में मौजूदा ग्रीडी ह्यूरिस्टिक्स (greedy heuristics) से काफी बेहतर प्रदर्शन करता है।

मूल लेखक: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

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

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

CC BY 4.0

मूल लेखक: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

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

कल्पना कीजिए कि आप एक विशाल, जटिल जिग्सॉ पहेली (jigsaw puzzle) को हल करने की कोशिश कर रहे हैं, लेकिन आपके पास एक जादुई सहायक (AI) है जो केवल सामने रखे टुकड़े को ही नहीं, बल्कि पूरी तस्वीर को एक साथ देख सकता है। मास्क्ड डिफ्यूजन मॉडल्स (MDMs) इसी तरह काम करते हैं। पारंपरिक AI के विपरीत, जो एक कहानी को बाएं से दाएं एक-एक शब्द करके लिखता है (जैसे किताब पढ़ना), MDMs पूरे खाली कैनवास को देख सकते हैं और किसी भी क्रम में टुकड़ों को भर सकते हैं।

हालांकि, एक पेच है: आप यह कैसे तय करेंगे कि पहले कौन सा टुकड़ा भरा जाए?

समस्या: "आसान पहले" का जाल (The "Easy First" Trap)

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

इसे पेपर में ग्रीडी सर्टेन्टी (Greedy Certainty) कहा गया है।

उपमा: कल्पना कीजिए कि आप एक गणित का सवाल हल कर रहे हैं: $2 \times 3 = ?$ $2 \times 3 = ?$ ।
- "आसान पहले" वाला AI देखता है कि उत्तर ($6$) बहुत स्पष्ट है। वह इसे तुरंत भर देता है क्योंकि वह इसे लेकर 100% आश्वस्त है।
- लेकिन फिर वह फंस जाता है। उसने उत्तर भरने से पहले संख्याएँ ($2 $और$ 3$) नहीं समझीं। अब उसे उत्तर के आधार पर कारकों (factors) का अनुमान लगाना पड़ता है, जो अविश्वसनीय रूप से कठिन और त्रुटिपूर्ण है।
- परिणाम: AI आसान हिस्से को तो सही कर लेता है लेकिन पूरी पहेली में विफल हो जाता है क्योंकि उसने आगे की योजना नहीं बनाई। यह घर बनाने जैसा है जहाँ आपने पहले छत लगा दी क्योंकि वह अच्छी दिख रही थी, और बाद में आपको एहसास हुआ कि आपके पास सहारा देने के लिए दीवारें ही नहीं थीं।

समाधान: "इन्फो-गेन" सैंपलर (The "Information Gain" Sampler)

लेखक एक नई रणनीति प्रस्तावित करते हैं जिसे इन्फो-गेन (Info-Gain) सैंपलर कहा जाता है। केवल यह पूछने के बजाय कि "कौन सा टुकड़ा भरना सबसे आसान है?", यह पूछता है, "कौन सा टुकड़ा, यदि मैं इसे अभी भर दूँ, तो मुझे बाकी की पहेली को सबसे तेज़ी से हल करने में मदद करेगा?"

उपमा: इसे एक जासूस की तरह सोचें जो रहस्य सुलझा रहा है।
- पुराना तरीका (Greedy): जासूस एक संदिग्ध को देखता है जो दोषी लग रहा है और उसे तुरंत गिरफ्तार कर लेता है। लेकिन बाद में, उन्हें एहसास होता है कि वह संदिग्ध निर्दोष था, और असली सुराग छूट गया था।
- नया तरीका (Info-Gain): जासूस पूरे अपराध स्थल को देखता है। उसे एहसास होता है कि गायब हथियार (एक कठिन सुराग) को ढूंढ लेना ही तुरंत अपराधी का पता लगा देगा। भले ही हथियार ढूंढना संदिग्ध को देखने से अधिक कठिन हो, लेकिन ऐसा करने से उसे सबसे अधिक इन्फॉर्मेशन गेन (सूचना प्राप्ति) मिलती है। यह पूरे मामले की उलझन को दूर कर देता है।

यह कैसे काम करता है (जादुई ट्रिक)

पेपर MDMs की एक विशेष शक्ति को उजागर करता है: वे भविष्य देख सकते हैं।

क्योंकि MDMs पूरे वाक्य या छवि को एक साथ देखते हैं (bidirectional), वे सिमुलेशन कर सकते हैं: "यदि मैं इस विशिष्ट शब्द को अभी भर दूँ, तो बाकी का वाक्य कितना स्पष्ट हो जाएगा?"

पुराना तरीका: "मुझे 90% यकीन है कि यह शब्द 'cat' है। चलिए 'cat' लिख देते हैं।" (यह अनदेखा करता है कि 'cat' लिखने से अगला वाक्य अर्थहीन हो सकता है)।
नया तरीका: "मुझे केवल 60% यकीन है कि यह शब्द 'cat' है, लेकिन यदि मैं इसे लिखता हूँ, तो पूरे वाक्य की अनिश्चितता 50% कम हो जाएगी। चलिए 'cat' लिखते हैं।"

नया सैंपलर तत्काल आत्मविश्वास (हम अभी कितने आश्वस्त हैं) और भविष्य की स्पष्टता (यह हमें बाद में कितनी मदद करेगा) के बीच संतुलन बनाता है।

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

शोधकर्ताओं ने कई कठिन कार्यों पर इसका परीक्षण किया:

गणित और तर्क (Math & Logic): इसने जटिल समीकरणों को हल किया क्योंकि इसने आसान उत्तरों के पीछे भागने के बजाय कठिन हिस्सों को पहले समझा।
कोडिंग (Coding): इसने विवरण भरने से पहले संरचना की योजना बनाकर बेहतर कोड लिखा।
रचनात्मक लेखन (Creative Writing): इसने अधिक सुसंगत कहानियाँ लिखीं जो भटकती नहीं थीं या विरोधाभासी नहीं थीं।
छवियाँ (Images): इसने ऐसी तस्वीरें बनाईं जहाँ वस्तुएं सही स्थानों पर थीं और उनके रंग भी सही थे।

निचोड़ (The Bottom Line)

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

यह उस छात्र के बीच का अंतर है जो परीक्षा पास करने के लिए उत्तर रटता है (Greedy) और उस छात्र के बीच का अंतर है जो किसी भी समस्या को हल करने के लिए अंतर्निहित अवधारणाओं को समझता है (Info-Gain)। परिणाम? अधिक स्मार्ट, अधिक सटीक और अधिक रचनात्मक AI।

यहाँ शोध पत्र "Improving Sampling for Masked Diffusion Models via Information Gain" का विस्तृत तकनीकी सारांश दिया गया है।

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

मास्क्ड डिफ्यूजन मॉडल्स (MDMs) डिस्क्रीट सीक्वेंस जनरेशन के लिए ऑटोरेग्रेसिव मॉडल्स (ARMs) के एक शक्तिशाली विकल्प के रूप में उभरे हैं, जो द्विदिश ध्यान (bidirectional attention) और गैर-कारण डिकोडिंग पथों (non-causal decoding paths) के माध्यम से लचीलापन प्रदान करते हैं। हालाँकि, एक महत्वपूर्ण प्रशिक्षण-अनुमान विसंगति (training-inference mismatch) मौजूद है:

प्रशिक्षण (Training): MDMs को रैंडम मास्किंग पैटर्न के तहत प्रशिक्षित किया जाता है, जिससे वे किसी भी संदर्भ के आधार पर किसी भी टोकन की भविष्यवाणी करना सीखते हैं।
अनुमान (Inference): जनरेशन के लिए एक बहु-चरणीय, क्रम-संवेदनशील डिकोडिंग प्रक्रिया की आवश्यकता होती है जहाँ सैंपलर को यह निर्णय लेना होता है कि कौन से मास्क किए गए स्थानों को अगला प्रकट करना है और किन टोकन को असाइन करना है।

मुख्य सीमा: मौजूदा सैंपलर लालची, स्थानीय निश्चितता अनुमानी (greedy, local certainty heuristics) पर निर्भर करते हैं (जैसे, उच्चतम आत्मविश्वास या निम्नतम एंट्रॉपी वाले स्थानों का चयन करना)।

अल्पदर्शिता (Myopia): ये विधियाँ भविष्य के चरणों पर वर्तमान निर्णयों के प्रभाव पर विचार किए बिना तत्काल स्थानीय निश्चितता को प्राथमिकता देती हैं।
विफलता के मामले: जैसा कि पेपर के "वन-वे मल्टीप्लिकेशन" और "बाइनरी जजमेंट" केस स्टडीज में प्रदर्शित किया गया है, लालची सैंपलर अक्सर ऐसे पथ चुनते हैं जो स्थानीय रूप से निश्चित दिखते हैं (जैसे, तर्क चरणों से पहले बाइनरी उत्तर टोकन को डिकोड करना) लेकिन बाद में उच्च संचयी अनिश्चितता (cumulative uncertainty) और तार्किक त्रुटियों का कारण बनते हैं।
अनछुआ потенциаल (Untapped Potential): ARMs के विपरीत, जहाँ भविष्य के प्रभावों का मूल्यांकन करना "नेक्स्ट-टोकन बॉटलनेक" के कारण कम्प्यूटेशनल रूप से निषेध है, MDMs की द्विदिश प्रकृति सैद्धांतिक रूप से कुशलता से यह मूल्यांकन करने की अनुमति देती है कि एक डिकोडिंग निर्णय पूरे शेष अनुक्रम में अनिश्चितता को कैसे नया आकार देता है।

2. कार्यप्रणाली: इन्फो-गेन सैंपलर (Methodology: The Info-Gain Sampler)

लेखक इन्फो-गेन सैंपलर (Info-Gain Sampler) प्रस्तावित करते हैं, जो एक सिद्धांत-आधारित डिकोडिंग ढांचा है जो तत्काल डिकोडिंग लागत और शेष अनुक्रम के लिए अपेक्षित सूचना लाभ (information gain) के बीच संतुलन बनाता है।

मुख्य अवधारणाएँ

अवस्था अनिश्चितता ( $H(z_t)$ ): वर्तमान में मास्क किए गए सभी स्थानों पर औसत मार्जिनल एंट्रॉपी के रूप में परिभाषित है।
**सूचना लाभ ($IG $):** एक विशिष्ट डिकोडिंग क्रिया$ a_t$ द्वारा प्रेरित अवस्था अनिश्चितता में कमी।
$IG(a_t; z_t) = H(z_t) - H(z_{t-1})$
उद्देश्य फलन (Objective Function): सैंपलर एक स्कोर $J_{IG}$ को अनुकूलित करता है जो तत्काल लागत (डिकोड किए जा रहे टोकन की अनिश्चितता) और सूचना लाभ (शेष टोकनों के लिए अनिश्चितता में कमी) के बीच संतुलन बनाता है:
$J_{IG}(a_t | z_t) = \underbrace{IG(a_t; z_t)}_{\text{Information Gain}} - \underbrace{C(a_t | z_t)}_{\text{Immediate Cost}}$
जहाँ $C(a_t | z_t)$ चयनित टोकनों की मार्जिनल एंट्रॉपी का योग है।

एल्गोरिदम वर्कफ़्लो

एक्शन सैंपलिंग: शीर्ष-k निश्चित टोकन को नियत रूप से चुनने के बजाय, सैंपलर $N$ $N$ क्रियाओं का एक विविध उम्मीदवार सेट $\mathcal{C}$ $C$ उत्पन्न करता है। इसमें शामिल है:
- टोकन सैंपलिंग: टेम्परेचर स्केलिंग के साथ उम्मीदवार टोकन खींचना।
- पोजीशन सैंपलिंग: निश्चितता स्कोर के आधार पर स्थानों का चयन करना।
समानांतर मूल्यांकन (Parallel Evaluation): प्रत्येक उम्मीदवार क्रिया $\mathcal{C}$ के लिए, मॉडल सभी उम्मीदवारों के लिए परिणामी अवस्था अनिश्चितता ( $H(z_{t-1})$ ) की गणना करने के लिए एक सिंगल बैच्ड फॉरवर्ड पास करता है। यह महंगे ट्री सर्च (ARMs के विपरीत) के बिना कुशलता से सूचना लाभ का अनुमान लगाने के लिए द्विदिश आर्किटेक्चर का लाभ उठाता है।
संक्रमण (Transition): $J_{IG}$ को अधिकतम करने वाली क्रिया का चयन किया जाता है, अवस्था को अपडेट किया जाता है, और प्रक्रिया तब तक दोहराई जाती है जब तक कि अनुक्रम पूर्ण न हो जाए।

दक्षता अनुकूलन

ब्लॉक-वाइज कंप्यूटेशन: KV-कैशिंग को सक्षम करने के लिए सूचना लाभ गणना को वर्तमान सक्रिय ब्लॉक तक सीमित करता है।
हाई-कॉन्फिडेंस बाईपास: यदि किसी टोकन की संभावना एक थ्रेशोल्ड $\gamma$ से अधिक हो जाती है, तो उसे सीधे एक्शन सेट में फिक्स कर दिया जाता है, जिससे हाई-कॉन्फिडेंस पोजीशन के लिए मूल्यांकन चरण को छोड़कर लेटेंसी कम हो जाती है।

3. मुख्य योगदान

सीमाओं की पहचान: अनुभवजन्य रूप से प्रदर्शित किया कि लालची निश्चितता-आधारित सैंपलर संचयी अनिश्चितता को कम करने में विफल रहते हैं, जिससे अक्सर जटिल तर्क कार्यों में त्रुटि प्रसार होता है।
नया डिकोडिंग ढांचा: इन्फो-गेन सैंपलर पेश किया, जो तत्काल लागतों के साथ भविष्य के सूचना लाभ को संतुलित करके कुशल, वैश्विक योजना (global planning) करने के लिए MDMs की अद्वितीय द्विदिश प्रकृति का उपयोग करता है।
व्यापक मूल्यांकन: विभिन्न आर्किटेक्चर (फुल-अटेंशन और सेमी-ऑटोरेग्रेसिव) और कार्यों (रीजनिंग, कोडिंग, प्लानिंग, क्रिएटिव राइटिंग और इमेज जनरेशन) में विधि को मान्य किया, जो मौजूदा बेसलाइन्स पर निरंतर श्रेष्ठता दिखाता है।

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

इन्फो-गेन सैंपलर का मूल्यांकन यूनिफॉर्म, एंट्रॉपी, कॉन्फिडेंस, मार्जिन, KLASS, PC-सैंपलर और एक समवर्ती लुकअहेड विधि (LookUM) जैसे बेसलाइन्स के विरुद्ध किया गया।

रीजनिंग कार्य (मैथ और कोड):
- Dream-7B मॉडल पर, इन्फो-गेन ने सर्वश्रेष्ठ लालची बेसलाइन पर औसत सटीकता में 3.6% सुधार प्राप्त किया।
- इसने संचयी एंट्रॉपी (Cumulative Entropy) (जो डिकोडिंग गुणवत्ता का एक प्रॉक्सी है) को काफी कम कर दिया, जो विशिष्ट रीजनिंग सेटिंग्स में 78.4 (सर्वश्रेष्ठ लालची बेसलाइन) से गिरकर 48.6 हो गई।
- इसने कोड और प्लानिंग कार्यों पर समवर्ती LookUM विधि को पछाड़ा, जिनमें उच्च टोकन-स्तरीय परिशुद्धता की आवश्यकता होती है।
क्रिएटिव राइटिंग:
- लंबाई-नियंत्रित मूल्यांकनों में बेसलाइन्स के विरुद्ध 63.1% जीत-दर (win-rate) प्राप्त की।
- टेम्परेचर स्केलिंग के प्रति बेहतर मजबूती प्रदर्शित की; उच्च तापमान ( $\tau=1.5$ ) पर, इसने एंट्रॉपी बेसलाइन के मुकाबले 80.3% जीत-दर बनाए रखी, जिससे रचनात्मकता और सुसंगतता के बीच संतुलन बना रहा।
इमेज जनरेशन (Text-to-Image):
- GenEval पर, उच्चतम औसत स्कोर (56.3 के अगले सर्वश्रेष्ठ के मुकाबले 58.2) प्राप्त किया।
- ImageNet-512 पर, FID (43.3 के मुकाबले 38.1) और इनसेप्शन स्कोर (53.3 के मुकाबले 63.0) में महत्वपूर्ण सुधार किया।
दक्षता: समानांतर मूल्यांकन और कैशिंग अनुकूलन के कारण, यह विधि लालची बेसलाइन्स की तुलना में केवल मामूली जनरेशन समय वृद्धि (+24%) और GPU मेमोरी (+20%) का भार डालती है।

5. महत्व और प्रभाव

स्थानीय और वैश्विक योजना के बीच सेतु: यह पेपर वर्तमान MDM सैंपलर की "अल्पदर्शिता" के लिए एक सिद्धांत-आधारित समाधान प्रदान करता है, यह सिद्ध करता है कि गैर-ऑटोरेग्रेसिव मॉडल ARMs से जुड़ी कम्प्यूटेशनल लागत के बिना प्रभावी वैश्विक योजना कर सकते हैं।
डेटा मैनिफोल्ड संरक्षण: लेखक तर्क देते हैं कि इन्फो-गेन सैंपलर डिकोडिंग प्रक्षेपवक्र (trajectory) को प्रशिक्षण डेटा मैनिफोल्ड के करीब रखने में प्रभावी है। भविष्य की अनिश्चितता (तार्किक सुसंगतता से विचलन) पैदा करने वाली क्रियाओं को दंडित करके, यह मॉडल को "भ्रम" (hallucination) की स्थिति में जाने से रोकता है।
सामान्यीकरण (Generalizability): यह ढांचा ट्रेनिंग-फ्री है और विभिन्न MDM आर्किटेक्चर (फुल-अटेंशन, सेमी-एआर, ब्लॉक डिफ्यूजन) के साथ संगत है, जो इसे वर्तमान पीढ़ी के डिस्क्रीट डिफ्यूजन मॉडल्स के लिए एक बहुमुखी अपग्रेड बनाता है।
नया मीट्रिक: पेपर संचयी एंट्रॉपी (Cumulative Entropy) को जनरेशन सटीकता के एक विश्वसनीय, नकारात्मक रूप से सह-संबंधित प्रॉक्सी के रूप में रेखांकित करता है, जो डिकोडिंग रणनीतियों के मूल्यांकन के लिए एक नया दृष्टिकोण प्रदान करता है।

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