🤖 machine learning

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

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

मूल लेखक: Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

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

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

CC BY 4.0

मूल लेखक: Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

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

यहाँ "Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models" शोध पत्र का सरल भाषा और रचनात्मक उपमाओं के साथ विवरण दिया गया है।

बड़ी तस्वीर: AI को सिखाने का एक नया तरीका

कल्पना कीजिए कि आपके पास एक बहुत ही प्रतिभाशाली कलाकार (AI मॉडल) है जो एक शोर-शराबे वाले, धुंधले कैनवास से शुरू करके धीरे-धीरे उसे साफ करता है जब तक कि एक स्पष्ट चित्र दिखाई न देने लगे। डिफ्यूजन लैंग्वेज मॉडल्स (Diffusion Language Models) इसी तरह काम करते हैं: वे टाइपराइटर की तरह एक-एक शब्द नहीं लिखते (जैसा कि पुराने AI मॉडल करते हैं); इसके बजाय, वे अक्षरों के एक ढेर से शुरू होते हैं और उन्हें एक सुसंगत वाक्य में बदलने के लिए धीरे-धीरे "डिनोइज़" (denoise) करते हैं।

यह शोध पत्र तर्क देता है कि कठिन समस्याओं (जैसे गणित या कोडिंग) को हल करने में बेहतर होने के लिए इन कलाकारों को सिखाने का वर्तमान तरीका त्रुटिपूर्ण है। लेखक एक नई विधि प्रस्तावित करते हैं जिसे TraFL (ट्रैजेक्टरी फ्लो बैलेंसिंग) कहा जाता है, जो एक विशिष्ट दोष को ठीक करती है जिसे "ट्रैजेक्टरी लॉकिंग" (Trajectory Locking) कहा जाता है।

समस्या: "एक-पथ" का जाल (ट्रैजेक्टरी लॉकिंग)

समस्या को समझने के लिए, एक भूलभुलैया (maze) की कल्पना करें।

लक्ष्य: भूलभुलैया में कई अलग-अलग सही निकास (एक गणितीय समस्या के विभिन्न वैध समाधान) हैं।
पुरानी विधि (रिवॉर्ड-मैक्सिमाइजिंग RL): कल्पना कीजिए कि आप एक कुत्ते को निकास खोजने के लिए प्रशिक्षित कर रहे हैं। हर बार जब कुत्ता कोई भी निकास खोज लेता है, तो आप उसे इनाम (treat) देते हैं।
- दोष: कुत्ते को इस बात से फर्क नहीं पड़ता कि वह इनाम पाने के लिए किस रास्ते से आया है, उसे बस इनाम चाहिए।
- परिणाम: यदि कुcha संयोग से शुरुआत में ही इनाम की ओर ले जाने वाले किसी एक विशिष्ट पथ पर stumbled (ठोकर खाकर पहुँच) हो जाता है, तो उसे इनाम मिलता है। उसे वह रास्ता याद हो जाता है। अगली बार, वह फिर से वही रास्ता आजमाता है। उसे फिर से इनाम मिलता है। वह उस एक पथ के प्रति और अधिक आश्वस्त हो जाता है।
- ट्रैजेक्टरी लॉकिंग: अंततः, कुत्ता भूलभुलैया में घूमना पूरी तरह से बंद कर देता है। वह केवल उसी एक संकीले पथ पर दौड़ता है, भले ही वहां 50 अन्य वैध निकास मौजूद हों जिन्हें वह खोज सकता था। उसने एक ही समाधान पर खुद को "लॉक" कर लिया है और अन्य समाधान खोजने की क्षमता खो दी है।

शोध पत्र में, लेखक इसे ट्रैजेक्टरी लॉकिंग कहते हैं। क्योंकि AI को केवल अंतिम उत्तर के लिए "रिवॉर्ड" (स्कोर) मिलता है, इसलिए वह इस तथ्य को अनदेखा कर देता है कि वहां तक पहुँचने के कई अलग-अलग तरीके (पथ) थे। यह उसकी सारी रचनात्मकता को एक संकीले मार्ग में समेट देता है, जिससे वह तब खराब प्रदर्शन करता है जब आपसे दोबारा प्रयास करने के लिए कहा जाता है।

समाधान: "गाइडबुक" दृष्टिकोण (TraFL)

लेखक TraFL का प्रस्ताव करते हैं, जो प्रशिक्षण के नियमों को बदल देता है। केवल किसी भी निकास के लिए इनाम देने के बजाय, वे AI को एक गाइडबुक (एक फ्रीज्ड रेफरेंस मॉडल) और एक मानचित्र (Map) देते हैं।

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

उपमा:
कल्प_िए कि आप एक छात्र को गणित का सवाल हल करना सिखा रहे हैं।

पुराना तरीका: आप कहते हैं, "सही उत्तर प्राप्त करो, चाहे जैसे भी हो!" छात्र एक ऐसी ट्रिक ढूंढ लेता है जो काम करती है, उसे रट लेता है, और कोई दूसरा तरीका सीखने से इनकार कर देता है।
TraFL तरीका: आप कहते हैं, "सही उत्तर खोजो, लेकिन अपनी सोचने की प्रक्रिया को विविध और लचीला बनाए रखो, जैसे कि एक शीर्ष छात्र जो कई अलग-अलग रणनीतियों को जानता है।"

TraFL AI को दो चीजों को संतुलित करने के लिए मजबूर करता है:

रिवॉर्ड प्राप्त करना: सही उत्तर खोजना।
विविधता बनाए रखना: एक ही दोहराव वाले पथ में सिमटने से बचना। यह गाइडबुक के सहारे, कई अलग-अलग वैध समाधानों पर अपनी "संभाव्यता द्रव्यमान" (probability mass/चुनने की संभावना) को फैलाए रखता है।

उन्होंने इसे कैसे काम करने के योग्य बनाया

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

एक सरोगेट स्कोर (Surrogate Score): चूंकि वे हर कदम के सटीक गणित को नहीं देख सकते, इसलिए उन्होंने एक "प्रॉक्सी" स्कोर बनाया जो अनुमान लगाता है कि एक पूर्ण उत्तर कितना अच्छा है, जिससे वे हर छोटे कदम की पूर्ण दृश्यता के बिना मॉडल को प्रशिक्षित कर सकें।
एक सीखा हुआ नॉर्मलाइज़र (Learned Normalizer): उन्होंने AI को एक विशेष "कैलकुलेटर" सिखाया जो विशिष्ट प्रश्न (प्रॉम्ट) के आधार पर कार्य की कठिनाई को समायोजित करता है, जिससे प्रशिक्षण निष्पक्ष और संतुलित बना रहता है।

परिणाम: क्या यह वास्तव में मदद करता है?

लेखकों ने इस नई विधि का परीक्षण गणित (शब्द समस्याओं को हल करना) और कोड (कंप्यूटर प्रोग्राम लिखना) पर किया।

"एक-पथ" का जाल टूट गया: अन्य विधियों के विपरीत, जो कभी-कभी एक समाधान खोजने में बेहतर होने के चक्कर में अन्य समाधान खोजने की क्षमता खो देती थीं, TraFL ने हर एक टेस्ट में सुधार किया।
अधिक सैंपल्स = बेहतर परिणाम: जब उन्होंने AI को केवल 1 के बजाय 16 अलग-अलग उत्तर उत्पन्न करने के लिए कहा, तो TraFL काफी बेहतर रहा। यह साबित करता है कि इसने वास्तव में कई अलग-अलग सही समाधान खोजे, न कि केवल एक भाग्यशाली अनुमान।
यह नई चीज़ों पर भी काम करता है: AI ने केवल प्रशिक्षण वाले प्रश्नों को रटा नहीं। जब इसका परीक्षण उन नए गणितीय (Minerva Math) और कोडिंग (LiveCodeBench) चुनौतियों पर किया गया जिन्हें इसने पहले कभी नहीं देखा था, तो TraFL सबसे मजबूत प्रदर्शन करने वाला रहा।
विविधता की जाँच: उन्होंने एक "जज" (एक अन्य AI) का उपयोग किया जो उत्तरों को देखता है। जज ने पुष्टि की कि TraFL केवल एक ही उत्तर को अलग शब्दों में नहीं दे रहा था; बल्कि यह वास्तव में एक ही समस्या को हल करने के लिए विभिन्न तर्क रणनीतियों और विभिन्न एल्गोरिदम का उपयोग कर रहा था।

सारांश

यह शोध पत्र एक दोष की पहचान करता है जहाँ AI मॉडल समस्या को हल करने के एक तरीके पर "अटक" जाते हैं, और अन्य वैध समाधानों को अनदेखा कर देते हैं। उन्होंने इसे TraFL के साथ ठीक किया, जो एक ऐसी विधि है जो AI को एक रेफरेंस मॉडल द्वारा निर्देशित करते हुए, विभिन्न पथों के विविध "एक्सप्लोरेशन" को बनाए रखते हुए सही उत्तर खोजना सिखाती है। परिणाम एक ऐसा AI है जो न केवल स्मार्ट है, बल्कि जब कई समाधान उत्पन्न करने के लिए कहा जाता है, तो अधिक रचनात्मक और विश्वसनीय भी है।

तकनीकी सारांश: डिफ्यूजन लैंग्वेज मॉडल्स के लिए ट्राजेक्टरी-बैलेंस पोस्ट-ट्रेनिंग

समस्या: डिफ्यूजन RL में ट्राजेक्टरी लॉकिंग (Trajectory Locking)

डिफ्यूजन लैंग्वेज मॉडल्स (dLLMs) रीजनिंग और कोड जनरेशन के लिए ऑटोरिग्रेसिव मॉडल्स के एक आशाजनक विकल्प के रूप में उभर रहे हैं। हालाँकि, dLLMs के लिए मौजूदा पोस्ट-ट्रेनिंग विधियाँ मुख्य रूप से ऑटोरिग्रेसिव मॉडल्स के लिए डिज़ाइन किए गए रिवॉर्ड-मैक्सिमाइजिंग रिइन्फोर्समेंट लर्निंग (RL) उद्देश्यों (जैसे PPO, GRPO) को अनुकूलित करती हैं। लेखक इस अनुकूलन में एक केंद्रीय विफलता मोड की पहचान करते हैं जिसे ट्राजेक्टरी लॉकिंग (trajectory locking) कहा जाता है।

डिफ्यूजन मॉडल्स में, रिवॉर्ड केवल अंतिम डिनोइज़्ड कंप्लीशन ( $y$ ) को दिया जाता है, जबकि मॉडल डिस्ट्रीब्यूशन पूरे डिनोइज़िंग ट्राजेक्टरी ( $\tau$ ) पर प्रेरित होता है। क्योंकि टर्मिनल रिवॉर्ड्स "पाथ-इंडिफरेंट" (path-indifferent) होते हैं—अर्थात, कोई भी ट्राजेक्टरी जो एक ही सही $y$ तक पहुँचती है, उसे समान रिवॉर्ड मिलता है—सैंपलड पॉलिसी-ग्रेडिएंट अपडेट्स उन विशिष्ट ट्राजेक्टरीज को सुदृढ़ करने की प्रवृत्ति रखते हैं जो प्रशिक्षण के शुरुआती चरणों में सैंपल हुई होती हैं। यह एक स्व-सुदृढ़ीकरण फीडबैक लूप बनाता है जहाँ प्रोबेबिलिटी मास (probability mass) डिनोइज़िंग पथों के एक संकीर्ण सेट पर केंद्रित हो जाता है। फलस्वरूप, भले ही मॉडल एक सही उत्तर उत्पन्न करना सीख जाए, वह बार-बार सैंपलिंग के तहत वैध समाधानों का एक विविध सेट (modes) उत्पन्न करने की क्षमता खो देता है। लेखक इसे थ्योरम 1 (Theorem 1) के माध्यम से औपचारिक रूप देते हैं, जो सिद्ध करता है कि टर्मिनल मोड कवरेज, ट्राजेक्टरी विविधता द्वारा ऊपरी रूप से सीमित (upper-bounded) है; अतः, ट्राजेक्टरी लॉकिंग मॉडल के सॉल्यूशन कवरेज को सीधे सीमित करती है।

कार्यप्रणाली: TraFL (ट्राजेक्टरी फ्लो बैलेंसिंग)

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

ऑब्जेक्टिव (The Objective)

टारगेट डिस्ट्रीब्यूशन $p^*(\tau | x)$ को इस प्रकार परिभाषित किया गया है:
$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$
जहाँ:

$p_{\text{ref}}$ एक फ्रीज़्ड रेफरेंस मॉडल है।
$r(x, y)$ टर्मिनल रिवॉर्ड है।
$\beta$ रिवॉर्ड शार्पनेस (sharpness) को नियंत्रित करता है।
$Z(x)$ एक प्रॉम्प्ट-डिपेंडेंट नॉर्मलाइजेशन टर्म (partition function) है।

ट्रेनिंग ऑब्जेक्टिव वर्तमान मॉडल के लॉग-प्रोबेबिलिटी और टारगेट के बीच स्क्वेयर्ड रेसिड्यूल (squared residual) को कम करता है:
$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$
$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$

डिफ्यूजन के लिए व्यावहारिक कार्यान्वयन

चूँकि dLLMs सटीक ट्राजेक्टरी लॉग-प्रोबेबिलिटी प्रकट नहीं करते हैं, इसलिए लेखक TraFL को व्यावहारिक बनाने के लिए दो प्रमुख घटक पेश करते हैं:

डिफ्यूजन-कम्पैटिबल सीक्वेंस-लेवल सरोगेट: वे एक मास्क-रीकंस्ट्रक्शन लोअर बाउंड पर आधारित एक सुलभ लॉग-प्रोबेबिलिटी सरोगेट ( $\log \hat{p}_\theta(y|x)$ ) व्युत्पन्न करते हैं। इसमें एक मास्क लेवल $l$ को सैंपल करना, $l$ टोकन्स को करप्ट करना, और पुनर्गठित (reconstructed) टोकन्स के लॉग-प्रोबेबिलिटी का औसत लेना शामिल है। स्कोर को मास्क किए गए टोकन्स की संख्या से स्वतंत्र रूप से नॉर्मलाइज़ करने के लिए $1/l$ फैक्टर का उपयोग किया जाता है।
लर्नड प्रॉम्प्ट-डिपेंडेंट नॉर्मलाइजेशन: बैच पर सॉफ्टमैक्स के माध्यम से $Z(x)$ का अनुमान लगाने वाली विधियों के विपरीत (जो कंसंट्रेशन पैदा कर सकती हैं), TraFL एक छोटे MLP के माध्यम से एक पार्टीशन फंक्शन $Z_\phi(x)$ को सीखता है। यह टर्म हर स्टेप पर ग्रेडिएंट सिग्नल प्राप्त करता है, चाहे जो भी कंप्लीशन सैंपल की गई हो, जो अन्य डिस्ट्रीब्यूशन-मैचिंग दृष्टिकोणों में देखी जाने वाली "बैच-लोकल" नॉर्मलाइजेशन विफलता को रोकता है।

मुख्य योगदान

ट्राजेक्टरी लॉकिंग की पहचान: यह पेपर औपचारिक रूप से पहचानता है और सिद्ध करता है कि dLLMs में रिवॉर्ड-मैक्सिमाइजिंग ऑब्जेक्टिव्स ट्राजेक्टरी विविधता के पतन (collapse) की ओर ले जाते हैं, जो बदले में अलग-अलग सही समाधान मोड के कवरेज को सीमित करता है।
TraFL ऑब्जेक्टिव: एक ट्राजेक्टरी-बैलेंस ऑब्जेक्टिव का प्रस्ताव जो पॉलिसी को एक रेफरेंस मॉडल से जोड़ता है, जिससे रिवॉर्ड (टर्मिनल कंप्लीशंस को टिल्ट करना) और रेफरेंस मॉडल (प्लॉसिबल ट्राजेक्टरीज पर प्रोबेबिलिटी मास आवंटित करना) की भूमिका अलग हो जाती है।
व्यावहारिक सरोगेट और नॉर्मलाइजेशन: एक डिफ्यूजन-कम्पैटिबल सीक्वेंस-लेवल सरोगेट और एक लर्नड नॉर्मलाइजेशन टर्म का विकास जो सटीक ट्राजेक्टरी प्रोबेबिलिटी के बिना स्थिर प्रशिक्षण सक्षम करता है।
अनुभवजन्य श्रेष्ठता (Empirical Superiority): TraFL एकमात्र मूल्यांकित पोस्ट-ट्रेनिंग विधि है जो सभी बेंचमार्क-लेंथ सेटिंग्स (GSM8K, MATH-500, HumanEval, MBPP) और सैंपलिंग बजट के तहत बेस मॉडल से लगातार सुधार करती है।

परिणाम

लेखक गणितीय तर्क और कोड जनरेशन बेंचमार्क पर बेस मॉडल (LLaDA-8B-Instruct) और स्ट्रॉन्ग बेसलाइन्स (ESPO, JUSTGRPO) के विरुद्ध TraFL का मूल्यांकन करते हैं।

इन-डिस्ट्रीब्यूशन परफॉरमेंस: TraFL सभी सैंपलिंग बजट ( $k \in \{1, 2, 4, 8, 16\}$ ) और टेम्परेचर के लिए Pass@k में सुधार करता है। विशेष रूप से, यह Pass@1 (सिंगल-सैंपल एक्यूरेसी) में सुधार करता है और साथ ही सैंपलिंग बजट बढ़ने के साथ बेहतर स्केल करता है, जबकि बेसलाइन्स अक्सर मल्टी-सैंपल गेन्स के लिए सिंगल-सैंपल एक्यूरेसी का त्याग कर देती हैं।
हेल्ड-आउट जनरलाइजेशन (Held-Out Generalization): लाभ आउट-ऑफ-डिस्ट्रीब्यूशन बेंचमार्क में भी स्थानांतरित होते हैं। Minerva Math पर, TraFL बेस मॉडल से ऊपर रहता है। LiveCodeBench पर, TraFL सभी कठिनाई स्प्लिट्स (Easy, Medium, Hard) में सबसे मजबूत विधि है, जो मीडियम और हार्ड समस्याओं पर बेसलाइन्स को महत्वपूर्ण रूप से पीछे छोड़ देता है।
विविधता विश्लेषण (Diversity Analysis): LLM-as-a-judge प्रोटोकॉल का उपयोग करते हुए, लेखक दिखाते हैं कि TraFL के लाभ केवल एक सिंगल मोड के स्कोर को तेज करने के बजाय, विशिष्ट समाधान रणनीतियों (जैसे, विभिन्न प्रमाण दृष्टिकोण या एल्गोरिथमिक डिकम्पोजिशन) के व्यापक कवरेज से जुड़े हैं। ESPO जैसे बेसलाइन्स, हालांकि सटीकता में सुधार करते हैं, बेस मॉडल पर महत्वपूर्ण विविधता लाभ नहीं दिखाते हैं, जो ट्राजेक्टरी-लॉकिंग हाइपोथेसिस के अनुरूप है।

महत्व और दावे

यह पेपर दावा करता है कि TraFL डिफ्यूजन मॉडल्स के लिए मानक RL पोस्ट-ट्रेनिंग की "मोड-सीकिंग" विफलता का एक सिद्धांतगत समाधान प्रदान करता है। रिवॉर्ड मैक्सिमाइजेशन से ट्राजेक्टरी-बैलेंस की ओर शिफ्ट होकर, यह कई वैध समाधानों को कवर करने के लिए आवश्यक डिनोइज़िंग पथों की विविधता को सुरक्षित रखता है।

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

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