Technische Samenvatting: Generatieve Modellering in Één Stap via Wasserstein-Gradiëntstromen

Probleemstelling

Moderne generatieve modellen, met name diffusiemodellen en op stroming gebaseerde methoden, hebben state-of-the-art steekproefkwaliteit bereikt, maar lijden onder hoge rekenkosten en latentie door hun afhankelijkheid van meervoudige iteratieve steekproefnemen. Hoewel methoden voor distillatie in één of enkele stappen bestaan, erven ze vaak trainingsignalen over van paden die oorspronkelijk voor iteratieve generatie zijn ontworpen, of vertrouwen ze op heuristische dynamica (bijv. Drifting Models) die geen principiële convergentiegaranties bieden en vatbaar zijn voor mode-instorting. De centrale uitdaging is het ontwerpen van een trainingsdynamiek die het een generator mogelijk maakt om een directe, één-staps mapping te leren van een referentieverdeling naar een doelgegevensverdeling, terwijl stabiele convergentie en globale mode-coverage worden gewaarborgd.

Methodologie: W-Flow

De auteurs stellen W-Flow voor, een raamwerk dat het ontwerp van trainingsdynamica ontkoppelt van de parametrisering van de generator. De kernidee is om de evolutie van de modelverdeling naar de doelverdeling voor te schrijven met behulp van Wasserstein-Gradiëntstromen (WGF) en vervolgens een statisch neurale netwerk te trainen om deze meervoudige-staps evolutie te comprimeren tot één inferentiestap.

1. Trainingsdynamica via Wasserstein-Gradiëntstroom

In plaats van een reeks denoising-stappen of snelheidsvelden direct te leren, definieert W-Flow de trainingsbaan als een evolutie in continue tijd in de ruimte van waarschijnlijkheidsmaten.

• Energiefunctionaal: De evolutie wordt gedreven door de steilste daling van een energiefunctionaal $F(q) = D(q \| p)$, waarbij $q$ de huidige modelverdeling is en $p$ de doelgegevensverdeling.
• Snelheidsveld: Het snelheidsveld in continue tijd $V_t$ wordt afgeleid uit de eerste variatie van de energiefunctionaal: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discretisatie: De continue stroom wordt gediscretiseerd met behulp van een expliciete Euler-scheme. Bij elke trainingsstap $k$ wordt de modelverdeling $q^{(k)}$ bijgewerkt via een lokaal transportmap $g^{(k)}(x) = x + \eta V^{(k)}(x)$, waarbij $\eta$ de stapgrootte is.

2. Instantiatie met Sinkhorn-divergentie

Om ervoor te zorgen dat het snelheidsveld computationeel hanteerbaar en geometrisch goed gedragen is in hoge dimensies, instantiëren de auteurs de energiefunctionaal met de Sinkhorn-divergentie ($S_\epsilon$), een entropisch geregulariseerde optimale transport (OT) metriek.

• Formulering Snelheidsveld: Het geïnduceerde snelheidsveld wordt gegeven door het verschil tussen twee barycentrische projecties:
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Hier transporteert $T^\epsilon_{q,p}$ massa van de gegenereerde verdeling $q$ naar de doelstelling $p$, terwijl $T^\epsilon_{q,q}$ een zelf-transportterm vertegenwoordigt die consistentie met de huidige verdeling afdwingt.
• Globale Coördinatie: In tegenstelling tot lokale kernel-interacties (bijv. in MMD of Drifting Models) is het Sinkhorn-snelheidsveld afgeleid van een globaal OT-plan, waardoor wordt gegarandeerd dat de beweging van deeltjes over de hele batch gecoördineerd is om mode-instorting te voorkomen.
• Twee-batch Schatter: Om artefacten van zelf-matching te vermijden waarbij deeltjes zichzelf met nul kosten matchen, hanteren de auteurs een twee-batch strategie. Ze berekenen de zelf-transportterm $T^\epsilon_{q,q}$ met behulp van een onafhankelijke tweede batch van steekproeven uit de huidige verdeling, in plaats van de diagonaal van de kostenmatrix van een enkele batch te maskeren.

3. Generator Training

De generator $f_\theta$ wordt getraind om de discrete deeltjesupdates gedefinieerd door de WGF na te bootsen.

• Doelstelling: De verliesfunctie minimaliseert de kwadratische afstand tussen de output van de generator en de "doel"-deeltjespositie afgeleid van de Sinkhorn-update:
  $$\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$$
  waarbij $\text{sg}(\cdot)$ de stop-gradient-operator aanduidt, die de bijgewerkte deeltjespositie behandelt als een vaste regressiedoelstelling.
• Classifier-Free Guidance (CFG): Het raamwerk ondersteunt CFG door het snelheidsveld direct te modificeren. In plaats van verdelingen lineair te mengen, stellen de auteurs snelheidsgeleiding voor, die een term toevoegt die evenredig is met het verschil tussen conditionele en onvoorwaardelijke snelheidsvelden. Deze benadering blijkt de standaard exponentieel getilde doelverdeling te herstellen bij gebruik van KL-divergentie en biedt een geometrisch consistente uitbreiding voor Sinkhorn-divergentie.

Belangrijkste Bijdragen

1. Principiële Één-staps Raamwerk: W-Flow vestigt een nieuw paradigma voor generatieve modellering in één stap door trainingsdynamica te verankeren in Wasserstein-gradiëntstromen, wat een duidelijke variatiele interpretatie biedt die ontbreekt in heuristische benaderingen zoals Drifting Models.
2. Op Sinkhorn gebaseerde Dynamica: De instantiatie van WGF met Sinkhorn-divergentie biedt een efficiënte, op optimale transport gebaseerde update-regel die globale verdelingsverschillen vastlegt en de mode-coverage verbetert ten opzichte van op kernel gebaseerde of lokale interactiemethoden.
3. Theoretische Convergentie: De auteurs bewijzen dat de deeltjesdynamica van W-Flow met eindige steekproeven convergeert naar de dynamica van verdelingen in continue tijd onder geschikte regulariteitsaannames, wat een theoretische rechtvaardiging biedt voor het leerproces.
4. Twee-batch Schatter: Er wordt een nieuwe twee-batch strategie geïntroduceerd voor het schatten van zelf-transporttermen, wat de zelf-matching-bias inherent aan schatters met één batch en heuristieken voor diagonaal masking elimineert.

Experimentele Resultaten

De auteurs evalueren W-Flow op ImageNet 256×256 en FFHQ.

• State-of-the-Art Prestaties: W-Flow bereikt een nieuwe state-of-the-art FID van 1,29 (XL-schaal) en 1,35 (L-schaal) voor generatie in één stap met klassenconditie op ImageNet, wat een aanzienlijke prestatieverbetering is ten opzichte van eerdere één-staps methoden (bijv. Drifting Models, MeanFlow) en veel meervoudige-staps diffusiemodellen (bijv. LightningDiT, SiT) overtreft, ondanks het gebruik van slechts één functiebeoordeling (NFE).
• Steekproefefficiëntie: In vergelijking met meervoudige-staps diffusiemodellen met vergelijkbare FID-scores, levert W-Flow ongeveer 100× snellere steekproefneming op (bijv. 77,88 afbeeldingen/sec versus 0,93 afbeeldingen/sec voor SiT-XL/2 op een enkele H100 GPU).
• Mode Coverage: In experimenten met onevenwichtige verdelingen (bijv. 95% seniorengezichten, 5% kindergezichten) slaagt W-Flow erin minderheidsmodi te vangen, terwijl Drifting Models mode-instorting vertonen en de minderheidsklasse volledig negeren.
• Convergentiesnelheid: W-Flow convergeert aanzienlijk sneller tijdens het trainen en bereikt de prestaties van een Drifting Model dat 1280 epochen is getraind in slechts 384 epochen.
• Domeinoverdracht: De methode demonstreert sterke capaciteit in domeinoverdrachttaken (bijv. leeftijdsvertaling op FFHQ), waarbij kortere en coherenter transporttrajecten worden geproduceerd in vergelijking met heuristische baselines.

Betekenis en Claims

Het artikel beweert dat Wasserstein-gradiëntstromen een principiële en effectieve basis bieden voor snelle en hoogwaardige generatieve modellering. Door het ontwerp van trainingsdynamica te scheiden van de parametrisering van de generator, toont W-Flow aan dat een één-staps generator kan worden getraind om prestaties te bereiken die vergelijkbaar zijn met complexe meervoudige-staps iteratieve processen. De auteurs benadrukken dat de globale coördinatie die wordt geboden door optimale transportdynamica, essentieel is voor het stabiliseren van training en het waarborgen van uitgebreide mode-coverage, waarmee kritieke beperkingen van bestaande één-staps benaderingen worden aangepakt. De resultaten suggereren dat W-Flow kan dienen als een schaalbaar alternatief voor toepassingen die lage latentie en hoogwaardige generatie vereisen, zoals real-time creatieve tools en simulatie.