📊 statistics

One-Step Generative Modeling via Wasserstein Gradient Flows

यह शोध पत्र W-Flow को प्रस्तुत करता है, जो एक वन-स्टेप (one-step) जनरेटिव मॉडलिंग फ्रेमवर्क है जो 1.29 FID के साथ स्टेट-ऑफ-द-आर्ट ImageNet जनरेशन और मल्टी-स्टेप डिफ्यूजन मॉडल्स की तुलना में लगभग 100 $\times$ तेज़ सैंपलिंग प्राप्त करने के लिए वॉसरस्टीन ग्रेडिएंट फ्लो (Wasserstein gradient flows) को एक एकल न्यूरल नेटवर्क इन्फरेंस में संकुचित करता है।

मूल लेखक: Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

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

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

CC BY 4.0

मूल लेखक: Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

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

कल्पना कीजिए कि आप एक रोबोट को बिल्लियों के सटीक चित्र बनाना सिखाना चाहते हैं।

पुराना तरीका (डिफ्यूजन मॉडल्स - Diffusion Models):
अधिकांश आधुनिक AI चित्रकार एक मूर्तिकार की तरह काम करते हैं जो पत्थर के ब्लॉक को तराश कर आकृति बनाता है, या एक फोटोग्राफर की तरह जो एक धुंधली फोटो को धीरे-धीरे साफ करता है। वे रैंडम शोर (noise) से शुरू करते हैं और उस शोर को धीरे-धीरे एक बिल्ली में बदलने के लिए सैकड़ों छोटे कदम उठाते हैं।

समस्या: यह दुकान तक जाने के लिए 100 छोटे, हिचकिचाते हुए कदम उठाने जैसा है, बजाय इसके कि आप सीधे वहां तक चलें। एक सिंगल तस्वीर बनाने में इसमें बहुत समय लगता है और बहुत अधिक ऊर्जा (कंप्यूटिंग पावर) खर्च होती है।

नया तरीका (W-Flow):
यह पेपर W-Flow पेश करता है, एक ऐसी विधि जो रोबोट को एक ही बड़े कदम (single, giant leap) में एक सटीक बिल्ली पेंट करना सिखाती है। 100 कदम उठाने के बजाय, यह "शोर" से "बिल्ली" तक का एक परफेक्ट शॉर्टकट सीख लेता है।

यह कैसे काम करता है? "नदी" का रूपक (Analogy)

W-Flow इस शॉर्टकट को कैसे सीखता है, इसे समझने के लिए, दो समूहों की कल्पना करें:

भीड़ (लक्ष्य/Target): लोगों का एक समूह जो एक आदर्श घेरे में खड़ा है (जो वास्तविक डेटा का प्रतिनिधित्व करता है, जैसे असली बिल्लियों की तस्वीरें)।
तैरने वाले (जनरेटर/Drifters): एक समूह जो एक मैदान में बेतरतीब ढंग से बिखरा हुआ खड़ा है (जो AI के वर्तमान, अस्त-व्यस्त अनुमानों का प्रतिनिधित्व करता है)।

लक्ष्य 'तैरने वालों' को इस तरह से हिलाना है कि वे भीड़ की तरह ही एक आदर्श घेरा बना लें, लेकिन AI को एक ऐसा नियम सीखना होगा जो इसे एक बार में ही कर सके।

1. "ऊर्जा" मानचित्र (ढलान/The Slope)
लेखक 'तैरने वालों' और 'भीड़' के बीच के स्थान की कल्पना एक पहाड़ी परिदृश्य के रूप में करते हैं। "भीड़" घाटी के सबसे निचले बिंदु (सबसे कम ऊर्जा वाले बिंदु) पर है। तैरने वाले किसी पहाड़ी पर कहीं ऊपर हैं।

नियम: यदि आप एक 'तैरने वाले' हैं, तो आप भीड़ तक पहुँचने के लिए जितनी जल्दी हो सके ढलान की ओर नीचे लुढ़कना चाहते हैं।
नवाचार: पिछले तरीकों ने यह पता लगाने के लिए एक "अनुमान" (heuristic) का उपयोग किया कि नीचे का रास्ता किस ओर है। कभी-कभी उन्होंने गलत अंदाजा लगाया, फंस गए, या लोगों को गलत दिशा में बहुत जोर से धकेल दिया।
W-Flow की ट्रिक: वे Sinkhorn Divergence नामक एक गणितीय उपकरण का उपयोग करते हैं। इसे एक सुपर-सटीक GPS के रूप में समझें जो हर एक व्यक्ति के लिए पहाड़ी से नीचे उतरने का सटीक और सबसे तीव्र रास्ता (steepest path) निकालता है, यह ध्यान में रखते हुए कि पूरा समूह एक साथ कैसे चलता है, न कि केवल व्यक्तिगत रूप से।

2. "दो-बैच" सुरक्षा जाल (The "Two-Batch" Safety Net)
यह गणना करते समय कि तैरने वालों को कैसे हिलना चाहिए, एक पेचीदा समस्या आती है: यदि आप किसी व्यक्ति को अपने समूह से दूर जाने के लिए कहते हैं, तो वह गलती से स्वयं से दूर जाने की कोशिश कर सकता है, जिसका कोई अर्थ नहीं है।

समाधान: पेपर एक चतुर ट्रिक का उपयोग करता है जिसे "Two-Batch" रणनीति कहा जाता है। कल्पना करें कि तैरने वालों को दो अलग-अलग लाइनों में विभाजित किया गया है। आप गणना करते हैं कि लाइन A को लाइन B के आधार पर कैसे हिलना चाहिए, और इसके विपरीत भी। यह उन्हें खुद के प्रतिबिंब से भ्रमित होने से रोकता है और यह सुनिश्चित करता है कि वे बिना फंसे लक्ष्य की ओर सुचारू रूप से बढ़ें।

3. यात्रा को संकुचित करना (Compressing the Journey)
यहाँ जादू वाला कदम है:

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

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

गति (Speed): पेपर का दावा है कि यह विधि पुराने मल्टी-स्टेप तरीकों की तुलना में लगभग 100 गुना तेज़ है। यदि पुराने तरीके को एक तस्वीर बनाने में 10 सेकंड लगते थे, तो यह एक सेकंड के भी छोटे हिस्से में इसे कर देता है।
गुणवत्ता (Quality): एक स्टेप (one-step) होने के बावजूद, तस्वीरें अविश्वसनीय रूप से उच्च गुणवत्ता वाली हैं। प्रसिद्ध ImageNet टेस्ट पर, इसने एक स्टेप जनरेटर के लिए एक नया रिकॉर्ड स्थापित करते हुए 1.29 का स्कोर (FID) प्राप्त किया। इसका मतलब है कि तस्वीरें लगभग असली तस्वीरों जैसी ही दिखती हैं।
स्थिरता (Stability): क्योंकि इसकी गणितीय नींव 'Wasserstein Gradient Flows' के ठोस सिद्धांत पर आधारित है न कि केवल अनुमान पर, इसलिए AI के "कोलैप्स" (जहाँ वह केवल एक ही प्रकार की बिल्ली बनाना सीख जाता है और बाकी सबको नजरअंदाज कर देता है) होने की संभावना कम है। यह डेटा के सभी अलग-अलग "मोड्स" को बेहतर तरीके से कवर करता है।

सारांश

सोचिए कि W-Flow एक छात्र को गणित की समस्या हल करना सिखाने जैसा है।

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

पेपर यह सिद्ध करता है कि इस विशिष्ट गणितीय "कंपास" (Sinkhorn divergence) का उपयोग करके प्रशिक्षण को निर्देशित करके, आप एक ऐसा जनरेटर बना सकते हैं जो अत्यधिक तेज़ और अत्यधिक सटीक दोनों है।

तकनीकी सारांश: Wasserstein Gradient Flows के माध्यम से वन-स्टेप जेनेरेटिव मॉडलिंग

समस्या विवरण

आधुनिक जेनेरेटिव मॉडल, विशेष रूप से डिफ्यूजन मॉडल और फ्लो-आधारित विधियाँ, अत्याधुनिक सैंपल गुणवत्ता प्राप्त कर चुके हैं, लेकिन वे अपने मल्टी-स्टेप इटरेटिव सैंपलिंग पर निर्भरता के कारण उच्च कम्प्यूटेशनल लागत और लेटेंसी (विलंबता) से ग्रस्त हैं। जबकि वन-स्टेप या फ्यू-स्टेप डिस्टिलेशन विधियाँ मौजूद हैं, वे अक्सर उन पथों से प्राप्त प्रशिक्षण संकेतों (training signals) को विरासत में लेती हैं जो मूल रूप से इटरेटिव जनरेशन के लिए डिज़ाइन किए गए थे, या वे हीयुरिस्टिक डायनेमिक्स (जैसे कि Drifting Models) पर निर्भर करती हैं जिनमें सिद्धांतगत अभिसरण गारंटी (convergence guarantees) का अभाव होता है और जो मोड कोलैप्स (mode collapse) के प्रति संवेदनशील होते हैं। केंद्रीय चुनौती एक ऐसा प्रशिक्षण डायनामिक डिज़ाइन करना है जो एक जनरेटर को संदर्भ वितरण (reference distribution) से लक्षित डेटा वितरण (target data distribution) तक एक सीधा, सिंगल-स्टेप मैपिंग सीखने की अनुमति दे, जबकि स्थिर अभिसरण और वैश्विक मोड कवरेज सुनिश्चित करे।

कार्यप्रणाली: W-Flow

लेखक W-Flow प्रस्तावित करते हैं, जो जनरेटर के पैरामीट्राइजेशन से प्रशिक्षण डायनेमिक्स के डिज़ाइन को अलग करता है। मुख्य विचार यह है कि प्रोबेबिलिटी मेजर्स (probability measures) के स्पेस में लक्षित वितरण की ओर मॉडल वितरण के विकास को Wasserstein Gradient Flows (WGF) का उपयोग करके निर्धारित किया जाए और फिर इस मल्टी-स्टेप विकास को एक सिंगल इन्फरेंस स्टेप में संकुचित करने के लिए एक स्टैटिक न्यूरल नेटवर्क को प्रशिक्षित किया जाए।

1. Wasserstein Gradient Flow के माध्यम से प्रशिक्षण डायनेमिक्स

डिफ्यूजन स्टेप्स या वेलोसिटी फील्ड्स को सीधे सीखने के बजाय, W-Flow प्रशिक्षण प्रक्षेपवक्र (training trajectory) को प्रोबेबिलिटी मेजर्स के स्पेस में एक निरंतर-समय विकास के रूप में परिभाषित करता है।

एनर्जी फंक्शनल (Energy Functional): विकास एक ऊर्जा फंक्शनल $F(q) = D(q \| p)$ के तीव्रतम अवमानन (steepest descent) द्वारा संचालित होता है, जहाँ $q$ वर्तमान मॉडल वितरण है और $p$ लक्षित डेटा वितरण है।
वेलोसिटी फील्ड (Velocity Field): निरंतर-समय वेलोसिटी फील्ड $V_t$ को ऊर्जा फंक्शनल के प्रथम वेरिएशन (first variation) से प्राप्त किया जाता है: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$ ।
विविक्तीकरण (Discretization): निरंतर प्रवाह को एक स्पष्ट यूलर स्कीम (explicit Euler scheme) का उपयोग करके विविक्त (discretize) किया जाता है। प्रत्येक प्रशिक्षण चरण $k$ पर, मॉडल वितरण $q^{(k)}$ को एक स्थानीय ट्रांसपोर्ट मैप $g^{(k)}(x) = x + \eta V^{(k)}(x)$ के माध्यम से अपडेट किया जाता है, जहाँ $\eta$ स्टेप साइज है।

2. Sinkhorn Divergence के साथ इंस्टेंशिएशन

यह सुनिश्चित करने के लिए कि वेलोसिटी फील्ड उच्च आयामों में कम्प्यूटेशनल रूप से सुलभ और ज्यामितीय रूप से व्यवहारजनक हो, लेखक ऊर्जा फंक्शनल को Sinkhorn divergence ( $S_\epsilon$ ) का उपयोग करके इंस्टेंशिएट करते हैं, जो एक एंट्रोपिक रेगुलराइज्ड ऑप्टिमल ट्रांसपोर्ट (OT) मेट्रिक है।

वेलोसिटी फील्ड फॉर्मूलेशन: प्रेरित वेलोसिटी फील्ड दो बैसेंट्रिक प्रोजेक्शन के अंतर द्वारा दिया जाता है:
$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$
यहाँ, $T^\epsilon_{q,p}$ उत्पन्न वितरण $q$ से लक्षित $p$ तक द्रव्यमान (mass) को ट्रांसपोर्ट करता है, जबकि $T^\epsilon_{q,q}$ एक सेल्फ-ट्रांसपोर्ट टर्म है जो वर्तमान वितरण के साथ निरंतरता लागू करता है।
ग्लोबल कोऑर्डिनेशन (Global Coordination): स्थानीय कर्नेल इंटरैक्शन (जैसे MMD या Drifting Models में) के विपरीत, Sinkhorn वेलोसिटी फील्ड एक ग्लोबल OT प्लान से प्राप्त होता है, जो यह सुनिश्चित करता है कि पूरे बैच में कणों (particles) की गति समन्वित हो ताकि मोड कोलैप्स को रोका जा सके।
टू-बैच एस्टीमेटर (Two-Batch Estimator): सेल्फ-मैचिंग आर्टिफैक्ट्स से बचने के लिए, जहाँ कण स्वयं से मैच होकर शून्य लागत दिखाते हैं, लेखक एक टू-बैच रणनीति का उपयोग करते हैं। वे वर्तमान वितरण के स्वतंत्र दूसरे बैच के नमूनों का उपयोग करके सेल्फ-ट्रांसपोर्ट टर्म $T^\epsilon_{q,q}$ की गणना करते हैं, न कि सिंगल बैच के कॉस्ट मैट्रिक्स के डायगोनल को मास्क करके।

3. जनरेटर ट्रेनिंग

जनरेटर $f_\theta$ को WGF द्वारा परिभाषित डिस्क्रीट पार्टिकल अपडेट्स की नकल करने के लिए प्रशिक्षित किया जाता है।

ऑब्जेक्टिव (Objective): लॉस फंक्शन Sinkhorn अपडेट से प्राप्त "टारगेट" पार्टिकल पोजीशन और जनरेटर के आउटपुट के बीच स्क्वेर्ड डिस्टेंस को कम करता है:
$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$
जहाँ $\text{sg}(\cdot)$ स्टॉप-ग्रेडिएंट ऑपरेटर है, जो अपडेटेड पार्टिकल पोजीशन को एक फिक्स्ड रिग्रेशन टारगेट के रूप में मानता है।
क्लासिफायर-फ्री गाइडेंस (CFG): फ्रेमवर्क वेलोसिटी फील्ड को सीधे संशोधित करके CFG का समर्थन करता है। वितरणों को लीनियरली मिक्स करने के बजाय, लेखक वेलोसिटी गाइडेंस का प्रस्ताव करते हैं, जो एक टर्म जोड़ता है जो कंडीशनल और अनकंडीशनल वेलोसिटी फील्ड के अंतर के समानुपाती होता है। यह दृष्टिकोण दिखाया गया है कि जब KL डाइवर्जेंस का उपयोग किया जाता है तो मानक एक्सपोनेंशियल टिल्टेड टारगेट वितरण को रिकवर करता है और Sinkhorn डाइवर्जेंस के लिए एक ज्यामितीय रूप से सुसंगत विस्तार प्रदान करता है।

मुख्य योगदान

सिद्धांतगत वन-स्टेप फ्रेमवर्क: W-Flow एक नया प्रतिमान स्थापित करता है जो प्रशिक्षण डायनेमिक्स को Wasserstein ग्रेडिएंट फ्लो में ग्राउंड करता है, जो Drifting Models जैसे हीयुरिस्टिक दृष्टिकोणों में अनुपस्थित एक स्पष्ट वेरिएशनल इंटरप्रिटेशन प्रदान करता है।
Sinkhorn-आधारित डायनेमिक्स: WGF को Sinkhorn डाइवर्जेंस के साथ इंस्टेंशिएट करना एक कुशल, ऑप्टिमल-ट्रांसपोर्ट-आधारित अपडेट नियम प्रदान करता है जो ग्लोबल डिस्ट्रिब्यूशनल विसंगतियों को कैप्चर करता है और कर्नेल-आधारित या लोकल इंटरैक्शन विधियों की तुलना में बेहतर मोड कवरेज सुनिश्चित करता है।
सैद्धांतिक अभिसरण (Theoretical Convergence): लेखक सिद्ध करते हैं कि उपयुक्त रेगुलैरिटी धारणाओं के तहत W-Flow के फाइनाइट-सैंपल पार्टिकल डायनेमिक्स निरंतर-समय डिस्ट्रिब्यूशनल डायनेमिक्स की ओर अभिसरित होते हैं, जो सीखने की प्रक्रिया के लिए एक सैद्धांतिक औचित्य प्रदान करते हैं।
टू-बैच एस्टीमेटर: सिंगल-बैच एस्टीमेटर्स और डायगोनल मास्किंग हीयुरिस्टिक्स में निहित सेल्फ-मैचिंग बायस को समाप्त करने के लिए एक नवीन टू-बैच रणनीति पेश की गई है।

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

लेखक ImageNet 256×256 और FFHQ पर W-Flow का मूल्यांकन करते हैं।

स्टेट-ऑफ-द-आर्ट प्रदर्शन: W-Flow ImageNet पर वन-स्टेप क्लास-कंडीशनल जनरेशन के लिए एक नया SOTA FID 1.29 (XL स्केल) और 1.35 (L स्केल) प्राप्त करता है, जो पिछले वन-स्टेप तरीकों (जैसे Drifting Models, MeanFlow) को काफी पीछे छोड़ देता है और केवल एक सिंगल फंक्शन इवैल्यूएशन (NFE) के बावजूद कई मल्टी-स्टेप डिफ्यूजन मॉडल्स (जैसे LightningDiT, SiT) से बेहतर प्रदर्शन करता है।
सैंपलिंग दक्षता: समान FID स्कोर वाले मल्टी-स्टेप डिफ्यूजन मॉडल्स की तुलना में, W-Flow लगभग 100× तेज़ सैंपलिंग प्रदान करता है (उदाहरण के लिए, एक सिंगल H100 GPU पर 77.88 इमेज/सेकंड बनाम SiT-XL/2 के लिए 0.93 इमेज/सेकंड)।
मोड कवरेज: असंतुलित वितरणों (जैसे 95% सीनियर फेस, 5% चाइल्ड फेस) के प्रयोगों में, W-Flow सफलतापूर्वक अल्पसंख्यक मोड (minority modes) को कैप्चर करता है, जबकि Drifting Models मोड कोलैप्स प्रदर्शित करते हैं और अल्पसंख्यक क्लास को पूरी तरह से अनदेखा कर देते हैं।
अभिसरण गति: W-Flow प्रशिक्षण के दौरान काफी तेज़ी से अभिसरित होता है, जो 1280 एपोच के लिए प्रशिक्षित Drifting Model के प्रदर्शन को केवल 384 एपोच में प्राप्त कर लेता है।
डोमेन ट्रांसफर: विधि डोमेन ट्रांसफर कार्यों (जैसे FFHQ पर एज ट्रांसलेशन) में मजबूत क्षमता प्रदर्शित करती है, जो हीयुरिस्टिक बेसलाइन्स की तुलना में छोटे और अधिक सुसंगत ट्रांसपोर्ट ट्रेजेक्टरीज बनाती है।

महत्व और दावे

पेपर का दावा है कि Wasserstein ग्रेडिएंट फ्लो तेज़ और उच्च-निष्ठा (high-fidelity) जेनेरेटिव मॉडलिंग के लिए एक सिद्धांतगत और प्रभावी आधार प्रदान करते हैं। प्रशिक्षण डायनेमिक्स के डिज़ाइन को जनरेटर पैरामीट्राइजेशन से अलग करके, W-Flow यह प्रदर्शित करता है कि एक सिंगल-स्टेप जनरेटर जटिल मल्टी-स्टेप इटरेटिव प्रक्रियाओं के बराबर प्रदर्शन प्राप्त करने के लिए प्रशिक्षित किया जा सकता है। लेखक इस बात पर जोर देते हैं कि ऑप्टिमल ट्रांसपोर्ट डायनेमिक्स द्वारा प्रदान किया गया ग्लोबल कोऑर्डिनेशन प्रशिक्षण को स्थिर करने और व्यापक मोड कवरेज सुनिश्चित करने के लिए महत्वपूर्ण है, जो मौजूदा वन-स्टेप दृष्टिकोणों की महत्वपूर्ण सीमाओं को संबोधित करता है। परिणाम बताते हैं कि W-Flow वास्तविक समय के क्रिएटिव टूल्स और सिमुलेशन जैसे कम-लेटेंसी, उच्च-गुणवत्ता वाले जनरेशन के लिए आवश्यक अनुप्रयोगों के लिए एक स्केलेबल विकल्प के रूप में कार्य कर सकता है।