Technische Samenvatting: Stochastische Transitie-Map Distillatie voor Snelle Probabilistische Inferentie

Probleemstelling

Diffusiemodellen hebben state-of-the-art prestaties gevestigd op het gebied van generatiekwaliteit, diversiteit en distributiecoverage in vergelijking met voorgangers zoals GANs en VAEs. Deze prestaties gaan echter vaak gepaard met dure inferentiekosten als gevolg van de noodzaak om continue-tijd gewone differentiaalvergelijkingen (ODE's) of stochastische differentiaalvergelijkingen (SDE's) over vele stappen te integreren.

Hoewel er bestaande versnellingsmethoden zijn, zoals distillatie (bijvoorbeeld Consistency Models) en benaderingen gebaseerd op Optimaal Transport, staan deze vaak voor een afweging tussen generatiekwaliteit/diversiteit en inferentiecomplexiteit. Een kritieke beperking van huidige distillatiemethoden is hun afhankelijkheid van deterministische inferentie (bijvoorbeeld via de probability-flow ODE). Gegeven een vaste ruisstalen is de output uniek. Dit beperkt hun bruikbaarheid in downstream taken die inherent stochastische inferentie vereisen, zoals diffusie-posterior sampling voor inverse problemen, op basis van energie gefundeerde fijne afstelling, of het verbeteren van de afstemming met gebruikersvragen. Bovendien vereisen veel bestaande distillatiekaders een voorgeïnstalleerd teacher-model, bi-niveau optimalisatie, of de simulatie en caching van trajecten, wat de training en schaalbaarheid bemoeilijkt.

Methodologie: Stochastische Transitie-Map Distillatie (STMD)

De auteurs stellen Stochastische Transitie-Map Distillatie (STMD) voor, een teacher-vrij kader dat is ontworpen om de inferentie van diffusiemodellen te versnellen terwijl de stochastische transitiestructuur van het onderliggende proces behouden blijft.

Kernconcept

In tegenstelling tot score-gebaseerde diffusiemodellen die typisch het gemiddelde van de posterior-distributie modelleren, distilleert STMD de volledige transitie-map die geassocieerd is met de bemonsterings-SDE. De methode leert een conditionele generator die de transitiekans $p(x_{t'} | x_t)$ voorspelt voor $t' < t$ in één of enkele stappen, waardoor de kostbare numerieke integratie van de omgekeerde SDE wordt vermeden.

Technische Formulering

1. Conditionele Stroom (MF): Het kader parameteriseert de SDE-transities met behulp van een conditioneel Mean Flow-model. In plaats van een deterministische stroom-map te leren, leert het het gemiddelde snelheid langs het traject van het diffusieproces.
2. x0-Parametrisatie: Om de transitie $p(x_{t'} | x_t)$ te modelleren, gebruiken de auteurs een $x_0$-parametrisatie. Ze drukken de transitie uit als een integraal over de posterior $p(x_0 | x_t)$ en de bekende brugmaat $p(x_{t'} | x_0, x_t)$. Aangezien de brugmaat voor variantiebehoudende (VP) SDE's in gesloten vorm beschikbaar is, reduceert het probleem zich tot het leren van een conditioneel model $p_\theta(x_0 | x_t)$.
3. Trainingsdoel: Het model $u_\theta$ wordt getraind om de conditionele gemiddelde snelheid te benaderen. De auteurs definiëren een Conditional Mean Flow (CMF)-verliesfunctie:
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   waarbij $z_s$ een rechte lijn-interpolant is tussen de data $x_0$ en een prior-stalen $z_1$, geconditioneerd op de diffusiestaat $x_t$ en tijd $t$. De $\text{sg}$-operator duidt de stop-gradient aan.
4. Teacher-Vrije Training: Het trainingsproces vereist geen voorgeïnstalleerd teacher-model. Het vertrouwt op stalen uit de gezamenlijke distributie $p(x_0, x_t)$, die kunnen worden gegenereerd door $x_0$ te bemonsteren uit de data-distributie en de forward SDE-kernel toe te passen (die in gesloten vorm bekend is).

Inferentie

Inferentie wordt uitgevoerd door de geleerde conditionele Mean Flow terug in de tijd te integreren. Het proces omvat:

1. Het bemonsteren van een initiële ruis $x_1$ (of $z_1$).
2. Het iteratief toepassen van het geleerde snelheidsveld $u_\theta$ om te stappen van tijd $s$ naar $s-\Delta s$.
3. Op elke diffusietijdstap $t$, het bemonsteren uit de brugdistributie $p(x_{t-\Delta t} | x_0, x_t)$ om de staat bij te werken, waardoor de stochastische transitie effectief in grote stappen wordt gesimuleerd.

Belangrijkste Bijdragen

1. Nieuw Kader: Het voorstellen van STMD, een distillatiekader gebaseerd op conditionele Mean Flows dat het leren van transitie-maps van hoog-dimensionale diffusie-SDE's mogelijk maakt met behulp van een generator in één of enkele stappen via een regressiedoel, zonder enig voorgeïnstalleerd teacher-model.
2. Theoretische Convergentie: De afleiding van convergentiegrenzen voor het Mean Flow-algoritme in de 2-Wasserstein-afstand. Naar kennis van de auteurs is dit de eerste convergentieanalyse van het Mean Flow-algoritme in deze metriek. Ze breiden dit uit om te bewijzen dat hun geleerde conditionele transitie-map convergeert naar de transitiekern van de onderliggende omgekeerde SDE.
3. Empirische Validatie: Demonstratie van de methode op taken voor beeldgeneratie en inpainting over de MNIST-, CIFAR-10- en CelebA-datasets, waarbij concurrerende prestaties worden getoond ten opzichte van baselines terwijl stochastiek behouden blijft.

Experimentele Resultaten

De auteurs evalueerden STMD op drie datasets:

• MNIST: Met behulp van een U-Net-ruggengraat werd het model vergeleken met het originele Mean Flow-algoritme en een Denoising Diffusion Probabilistic Model (DDPM). De auteurs maakten gebruik van de Fréchet-afstand (FD) in de latente ruimte van een classifier (aangezien FID onbetrouwbaar is voor MNIST). Resultaten toonden aan dat STMD vergelijkbaar presteert met de baselines terwijl het stochastische bemonstering ondersteunt.
• CIFAR-10: Vergelijkbare experimenten werden uitgevoerd met een U-Net. Het model behaalde concurrerende Fréchet Inception Distance (FID)-scores in vergelijking met DDPM- en Mean Flow-baselines voor een vast aantal trainingiteraties, terwijl de volledige stochastiek van de diffusie-sampler behouden bleef.
• CelebA: Experimenten maakten gebruik van een latente DiT-architectuur. Het model genereerde onvoorwaardelijke stalen met een FID van 8,28 (met 4 inferentiestappen en 2 Mean Flow-stappen) en voerde met succes taken voor beeldinpainting uit.

In alle gevallen werden de modellen vanaf nul getraind met dezelfde netwerkruggengraten en iteratieaantallen als de baselines om een eerlijke vergelijking te waarborgen.

Betekenis en Claims

Het artikel stelt dat STMD de kloof overbrugt tussen snelle deterministische distillatie en de behoefte aan stochastische inferentie in downstream taken. Door de transitie-map te distilleren in plaats van alleen de flow-map, behoudt STMD de stochastische structuur van het diffusieproces.

De auteurs benadrukken dat hun aanpak teacher-vrij is, waardoor de behoefte aan bi-niveau optimalisatie of traject-caching wordt geëlimineerd, wat het trainingsproces efficiënter en schaalbaarder maakt. De theoretische bijdrage van het vaststellen van 2-Wasserstein-convergentiegrenzen biedt een rigoureuze basis voor de Mean Flow-aanpak, waardoor een brug wordt geslagen tussen theoretische garanties en praktische toepassing.

Het werk positioneert zichzelf als een oplossing voor taken die probabilistische inferentie vereisen, zoals inverse problemen en posterior sampling, waar deterministische distillatiemethoden tekortschieten. De auteurs blijven bescheiden wat betreft de huidige prestaties, met de opmerking dat het beperkte aantal trainingiteraties state-of-the-art resultaten verhindert, maar de vergelijkende studie valideert de effectiviteit van de methode ten opzichte van bestaande baselines.