🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

यह शोध पत्र $f$ -ट्रैजेक्टरी बैलेंस (f-Trajectory Balance) को प्रस्तुत करता है, जो लॉस फंक्शन्स (loss functions) का एक परिवार है जो मीन स्क्वायर एरर (mean square error) दृष्टिकोण को सभी $f$ -डाइवर्जेंस (f-divergences) तक विस्तारित करता है, जिससे जेनरेटिव मॉडल्स और एलएलएम (LLMs) को उनके संबंधित ऑन-पॉलिसी (on-policy) $f$ -डाइवर्जेंस ग्रेडिएंट्स के विशिष्ट अनुकूलन गुणों (जैसे कि मोड कवरिंग) को संरक्षित करते हुए ऑफ-पॉलिसी (off-policy) डेटा के साथ प्रशिक्षित करना सक्षम होता है।

मूल लेखक: Jake Fawkes, Jason Hartford

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

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

CC BY 4.0

मूल लेखक: Jake Fawkes, Jason Hartford

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

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

यह पेपर इन रोबोटों को प्रशिक्षित करने में मदद करने के लिए "शिक्षण उपकरण" (लॉस फंक्शन्स) का एक नया सेट पेश करता है, चाहे वे नई दवाओं के लिए अणु (molecules) बना रहे हों, कला बना रहे हों, या कोड लिख रहे हों।

यहाँ सरल उपमाओं का उपयोग करके पेपर के विचारों का विवरण दिया गया है:

1. समस्या: "गोल्ड रश" बनाम "एक्सप्लोरर" (खोजकर्ता)

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

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

2. पुराना उपकरण: "स्क्वेयर्ड KL" लॉस (Squared KL Loss)

हाल ही में, शोधकर्ताओं ने एक चतुर तरीका खोजा जिससे "स्क्वेयर्ड एरर" विधि (यह मापना कि रोबोट क्या सोचता है और उसे क्या सोचना चाहिए) का उपयोग करके रोबोट को सिखाया जा सकता है।

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

3. नया समाधान: "f-ट्रैजेक्टरी बैलेंस" (f-Trajectory Balance) परिवार

लेखकों ने महसूस किया कि "स्क्वेयर्ड" विधि वास्तव में एक बहुत बड़े शिक्षण उपकरणों के परिवार का केवल एक विशिष्ट स्वाद है। वे इस परिवार को f-ट्रैजेक्टरी बैलेंस कहते हैं।

इस परिवार को एक रेडियो के डायल या वॉल्यूम नॉब के रूप में सोचें।

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

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

पेपर दो मुख्य चीजें सिद्ध करता है:

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

5. वास्तविक दुनिया के परीक्षण (गुफा के मानचित्र)

लेखकों ने तीन अलग-अलग "गुफाओं" में इन उपकरणों का परीक्षण किया:

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

सारांश

यह पेपर एक नए प्रकार का रोबोट नहीं बनाता है। इसके बजाय, यह उन्हें प्रशिक्षित करने के लिए निर्देशों का एक नया सेट बनाता है।

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

यह इंजीनियरों को उनके AI मॉडल की रचनात्मकता या फोकस को नियंत्रित करने के लिए एक सरल नॉब देता है, बिना प्रशिक्षण प्रक्रिया को बाधित किए।

तकनीकी सारांश: f-ट्रैजेक्टरी बैलेंस (f-Trajectory Balance)

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

जेनरेटिव मॉडल्स के अलाइनमेंट और फाइन-ट्यूनिंग, जिसमें लार्ज लैंग्वेज मॉडल्स (LLMs) और जेनेरेटिव फ्लो नेटवर्क्स (GFlowNets) शामिल हैं, आमतौर पर एक मॉडल के डिस्ट्रीब्यूशन और एक टारगेट डिस्ट्रीब्यूशन के बीच विसंगति को कम करने पर निर्भर करते हैं। रिइन्फोर्समेंट लर्निंग (RL) फाइन-ट्यूनिंग में, इसे अक्सर कुलबैक-लीब्लर (KL) डाइवर्जेंस पेनल्टी के साथ अपेक्षित रिवॉर्ड को मैक्सिमाइज करने के रूप में फ्रेम किया जाता है।

इस डोमेन में एक महत्वपूर्ण चुनौती मोड-सीकिंग (mode-seeking) और मोड-कवरिंग (mode-covering) व्यवहारों के बीच का ट्रेड-ऑफ है। मानक ऑब्जेक्टिव्स, जैसे कि रिवर्स KL डाइवर्जेंस (जिसे अक्सर ट्रैजेक्टरी बैलेंस या स्क्वेयर्ड KL लॉस के माध्यम से लागू किया जाता है), मोड-सीकिंग होते हैं, जिससे मॉडल्स कुछ उच्च-संभाव्यता वाले मोड्स पर कोलैप्स हो जाते हैं। यह उन कार्यों के लिए उप-इष्टतम (suboptimal) है जिनमें अन्वेषण (exploration) की आवश्यकता होती है, जैसे कि ड्रग डिस्कवरी या विविध कैंडिडेट जनरेशन, जहाँ "मोड-कवरिंग" (टारगेट डिस्ट्रीब्यूशन की पूरी विविधता को कैप्चर करना) अधिक वांछनीय है।

इसके अलावा, मौजूदा तरीके अक्सर ऑफ-पॉलिसी वैधता (off-policy validity) के साथ संघर्ष करते हैं। जबकि REINFORCE या PPO जैसे एस्टिमेटर्स सामान्य हैं, वे उच्च वेरिएंस (high variance) से ग्रस्त होते हैं। हाल के दृष्टिकोण जो स्क्वेयर्ड KL लॉस ( $KL_{sq}$ ) का उपयोग करते हैं, ऑन-पॉलिसी पर लो-वेरिएंस ग्रेडिएंट्स और ऑफ-पॉलिसी वैलिडिटी प्रदान करते हैं, लेकिन वे स्वाभाविक रूप से रिवर्स KL डाइवर्जेंस के गुणों से बंधे होते हैं, जो एक्सप्लोरेशन-एक्सप्लोइटेशन ट्रेड-ऑफ को नियंत्रित करने की क्षमता को सीमित करते हैं।

कार्यप्रणाली (Methodology)

लेखक f-ट्रैजेक्टरी बैलेंस प्रस्तावित करते हैं, जो सरोगेट लॉस फंक्शन्स का एक परिवार है जो स्क्वेयर्ड KL लॉस को संपूर्ण f-डाइवर्जेंस परिवार तक सामान्यीकृत (generalize) करता है।

सैद्धांतिक आधार (Theoretical Foundation)

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

डाइवर्जेंस से लॉस तक: किसी भी फलन $f$ के लिए जो एक f-डाइवर्जेंस को परिभाषित करता है $D_f(p_\theta \| p^*)$ , लेखक एक लॉस फंक्शन $L_f$ व्युत्पन्न करते हैं जो लॉग-प्रोबेबिलिटीज के अंतर $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ पर आधारित है:
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
इस लॉस के दो महत्वपूर्ण गुण हैं:
- ऑन-पॉलिसी ग्रेडिएंट इक्विलेंस: जब सैंपल मॉडल पॉलिसी ( $\mu = p_\theta$ ) से लिए जाते हैं, तो $L_f$ के अपेक्षित ऑटो-डिफरेंशिएटेड ग्रेडिएंट्स, संबंधित f-डाइवर्जेंस के ग्रेडिएंट्स से मेल खाते हैं।
- ऑफ-पॉलिसी वैलिडिटी: जब इसे ऑफ-पॉलिसी डेटा (व्यवहार वितरण $\mu$ से सैंपल किया गया) पर मूल्यांकित किया जाता है, तो लॉस एक वैध ऑब्जेक्टिव बना रहता है जिसका समान ग्लोबल मिनिमाइज़र ( $p_\theta = p^*$ ) होता है, बशर्ते कि टारगेट का सपोर्ट, व्यवहार पॉलिसी के सपोर्ट के भीतर हो।
लॉस से डाइवर्जेंस तक: इसके विपरीत, लॉग-प्रोबेबिलिटीज पर कोई भी स्ट्रिक्टली कॉन्वेक्स, ट्रांसलेशन-इनवेरिएंट लॉस $\ell$ , एक विशिष्ट f-डाइवर्जेंस को मिनिमाइज करने के अनुरूप होता है। पेपर सिद्ध करता है कि ये मैपिंगिंग्स एक दूसरे के इनवर्स हैं।

व्यावहारिक कार्यान्वयन (Practical Implementations)

प्रशिक्षण की व्यावहारिक चुनौतियों को संबोधित करने के लिए, लेखक कई विस्तार पेश करते हैं:

DevGrad लॉस (बैच-वाइज नॉर्मलाइजेशन): इंट्राक्टेबल पार्टीशन फंक्शन्स (जो GFlowNets और अननॉर्मलाइज्ड टारगेट्स में सामान्य हैं) को संभालने के लिए, लेखक VarGrad एस्टिमेटर का सामान्यीकरण करते हैं। बैच वेरिएंस के बजाय, वे लॉग-प्रोबेबिलिटी अंतर के बैच-वाइज जनरलाइज्ड डेविएशन का उपयोग करते हैं। इसमें स्कोर फंक्शन कोएफिशिएंट्स को सेंटर करने के लिए लॉस को मिनिमाइज करके $\log Z$ का अनुमान लगाना शामिल है, जिससे वेरिएंस कम होता है।
टेम्पर्ड लॉस (Tempered Loss): छोटे टेम्परेचर पैरामीटर $\beta$ (जहाँ टारगेट डिस्ट्रीब्यूशन $p^* \propto \pi_{ref} \exp(r/\beta)$ है) के साथ काम करते समय न्यूमेरिकल ओवरफ्लो को रोकने के लिए, लेखक एक टेम्पर्ड लॉस प्रस्तावित करते हैं। यह लॉस को $1/\beta$ से स्केल करता है और इसे टेम्पर्ड डिस्ट्रीब्यूशन पर लागू करता है, जिससे ग्रेडिएंट्स बाउंडेड और स्टेबल रहते हैं।
GFlowNets के लिए f-ट्रैजेक्टरी बैलेंस: इस फ्रेमवर्क को GFlowNets पर लागू किया जाता है, जहाँ $L_f$ लॉस का उपयोग करके ट्रैजेक्टरी बैलेंस कंस्ट्रेंट्स को लागू किया जाता है। यह मानक ट्रैजेक्टरी बैलेंस (जो रिवर्स KL के अनुरूप है) को मनमाने f-डाइवर्जेंसेस तक सामान्यीकृत करता है।

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

सामान्यीकृत लॉस फैमिली: ट्रांसलेशन-इनवेरिएंट सरोगेट लॉस ( $L_f$ ) के एक सामान्य परिवार का व्युत्पत्ति, जिनके ऑन-पॉलिसी ग्रेडिएंट्स संबंधित f-डाइवर्जेंस से मेल खाते हैं, जबकि ऑफ-पॉलिसी वैलिडिटी बनी रहती है।
सैद्धांतिक समानता: ट्रांसलेशन-इनवेरिएंट लॉग-प्रोबेबिलिटी लॉस और f-डाइवर्जेंस के बीच व्युत्क्रम संबंध का प्रमाण, जो वांछित गुणों (जैसे, मोड-कवरिंग) के साथ नए लॉस फंक्शन्स के व्यवस्थित डिजाइन की अनुमति देता है।
अननॉर्मलाइज्ड टारगेट्स तक विस्तार: DevGrad लॉस के माध्यम से VarGrad तकनीक का सामान्यीकरण, जो RL फाइन-ट्यूनिंग में इंट्राक्टेबल पार्टीशन फंक्शन की समस्या को संबोधित करता है।
एक्सप्लोरेशन का नियंत्रण: यह प्रदर्शन कि f-डाइवर्जेंस का चयन (जो $\alpha$ -डाइवर्जेंसेस में $\alpha$ द्वारा पैरामीटराइज्ड है) मोड-सीकिंग बनाम मोड-कवरिंग व्यवहार को सीधे नियंत्रित करता है। कम $\alpha$ मान (जैसे, फॉरवर्ड KL, हेलिंगर) मोड-कवरिंग को बढ़ावा देते हैं, जबकि उच्च $\alpha$ मान (जैसे, पियर्सन $\chi^2$ ) मोड-सीकिंग को बढ़ावा देते हैं।

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

लेखक चार अलग-अलग कार्यों में लॉस फैमिली का मूल्यांकन करते हैं:

सिंथेटिक ग्रिड (GFlowNets): मोड कवरेज का परीक्षण करने के लिए डिज़ाइन किए गए 2D ग्रिड टास्क पर, कम $\alpha$ वाले लॉस (फॉरवर्ड KL, हेलिंगर) ने मानक ट्रैजेक्टरी बैलेंस (रिवर्स KL) की तुलना में बहुत तेज़ी से और अधिक विश्वसनीयता के साथ टारगेट डिस्ट्रीब्यूशन के चारों मोड्स को सफलतापूर्वक खोज लिया। इसके विपरीत, पियर्सन $\chi^2$ ( $\alpha=2$ ) ने मजबूत मोड-सीकिंग व्यवहार दिखाया, जो पहले मिले मोड में ही फंस गया।
मॉलिक्यूल डिस्कवरी (SynFlowNet): सिंथेटिक रूप से सुलभ अणुओं को उत्पन्न करने के लिए SynFlowNet फ्रेमवर्क पर लागू किया गया। लेखकों ने पाया कि $\alpha$ को ट्यून करने से उत्पन्न अणुओं की विविधता को नियंत्रित किया जा सकता है। एक एनीलिंग शेड्यूल (एक्सप्लोरेशन के लिए कम $\alpha$ से शुरू होकर एक्सप्लोइटेशन के लिए $\alpha=1$ तक बढ़ना) ने सबसे अच्छा ट्रेड-ऑफ दिया, जिससे मानक ट्रैजेक्टरी बैलेंस की तुलना में अधिक विविधता के साथ उच्च-रिवॉर्ड वाले अणु उत्पन्न हुए।
डिफ्यूजन मॉडल्स: MNIST डिजिट्स (विषम बनाम सम) के कंडीशनल सैंपलिंग में, मानक ट्रैजेक्टरी बैलेंस सभी टारगेट मोड्स में समान रूप से सैंपल करने में विफल रहा (विशिष्ट डिजिट्स जैसे 0 और 6 का ओवरसैंपलिंग)। f-ट्रैजेक्टरी बैलेंस लॉस, विशेष रूप से कम $\alpha$ वाले, टारगेट मोड्स के बीच अधिक समान सैंपलिंग प्राप्त करने में सफल रहे।
एसिंक्रोनस LLM ट्यूनिंग: लॉस को गणितीय समस्याओं (GSM8k, MATH) पर 50-स्टेप एसिंक्रोनस डिले के साथ LLMs (Qwen और OLMo फैमिली) को फाइन-ट्यून करने के लिए लागू किया गया। f-डाइवर्जेंस लॉस ने स्थिर प्रशिक्षण और सुसंगत एंट्रॉपी-रिवॉर्ड ट्रेड-ऑफ प्रदर्शित किया, जबकि मानक PPO कार्यान्वयन ने समान एसिंक्रोनस स्थितियों के तहत अस्थिरता दिखाई।

महत्व और दावे (Significance and Claims)

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

प्राथमिक महत्व इस क्षमता में निहित है कि केवल एक अलग सरोगेट लॉस फंक्शन (या एक $\alpha$ पैरामीटर को ट्यून करके) को चुनकर एक्सप्लोरेशन-एक्सप्लोइटेशन ट्रेड-ऑफ को स्पष्ट रूप से नियंत्रित किया जा सकता है, बिना अंतर्निहित प्रशिक्षण एल्गोरिदम को बदले। यह अभ्यासकर्ताओं को मोड-सीकिंग (रिवॉर्ड मैक्सिमाइजेशन के लिए) और मोड-कवरिंग (विविधता और अन्वेषण के लिए) व्यवहारों के बीच चयन करने की अनुमति देता है, जबकि आधुनिक, बड़े पैमाने पर एसिंक्रोनस प्रशिक्षण (LLMs और GFlowNets के लिए) के लिए आवश्यक स्थिरता और ऑफ-पॉलिसी कम्पैटिबिलिटी बनाए रखता है। यह कार्य सुझाव देता है कि यह फ्रेमवर्क विशेष रूप से ड्रग डिस्कवरी और एक्सप्लोरेटरी एजेंट्स जैसे अनुप्रयोगों के लिए मूल्यवान है जहाँ उच्च-रिवॉर्ड कैंडिडेट के स्पेस को कवर करना महत्वपूर्ण है।