Technische Samenvatting: Generatieve Modellering via Drifting

Probleemstelling

Generatieve modellering is fundamenteel de taak van het leren van een mapping $f$ zodanig dat de pushforward van een prior distributie $p_{prior}$ (bijv. Gaussische ruis) overeenkomt met de data-distributie $p_{data}$, genoteerd als $f\#p_{prior} \approx p_{data}$. Terwijl heersende paradigma's zoals Diffusion Models en Flow Matching dit bereiken door iteratief monsters te verfijnen tijdens de inferentie-tijd (het deconstrueren van de complexe pushforward-map in een keten van haalbare transformaties), vereist deze aanpak meerdere neurale netwerk-evaluaties (NFE) per monster, wat leidt tot hoge computationele kosten.

Bestaande pogingen om one-step generatoren te creëren vertrouwen vaak op het distilleren van vooraf getrainde multi-step modellen of het benaderen van de trajecten van stochastische differentiaalvergelijkingen (SDE's) of gewone differentiaalvergelijkingen (ODE's) tijdens de training. De auteurs stellen een nieuw paradigma voor dat het iteratieve proces fundamenteel verschuift van de inferentie-tijd naar de trainingstijd, wat hoogwaardige generatie mogelijk maakt met een enkele neurale functie-evaluatie (1-NFE) zonder afhankelijk te zijn van SDE/ODE-formuleringen.

Methodologie: Drifting Models

De kernbijdrage is het Drifting Model, dat generatieve modellering formuleert als de evolutie van de pushforward-distributie tijdens het trainingproces zelf.

1. Evolutie tijdens de training

In plaats van training enkel te zien als het bijwerken van gewichten om een statische loss te minimaliseren, beschouwen de auteurs de sequentie van modellen $\{f_i\}$ die voortkomen uit iteratieve optimalisatie (bijv. SGD) als een sequentie van pushforward-distributies $\{q_i\}$. Bij elke iteratie wordt een monster $x_i = f_i(\epsilon)$ impliciet "gedrift" naar $x_{i+1} = f_{i+1}(\epsilon)$ als gevolg van parameterupdates. Het doel is om een mechanisme te definiëren waarbij deze drift de distributie $q$ natuurlijk naar $p_{data}$ leidt.

2. Het Drifting Field

Om deze evolutie te regelen, introduceren de auteurs een drifting field $V_{p,q}(x)$, een functie die een verplaatsingsvector $\Delta x$ berekent voor een monster $x$ op basis van de huidige gegenereerde distributie $q$ en de doel-data-distributie $p$.

• Evenwichtstoestand: Het veld is ontworpen zodanig dat $V_{p,q}(x) = 0$ als en slechts als $p = q$. Wanneer de distributies overeenkomen, stoppen de monsters met driften, wat een evenwicht bereikt.
• Anti-symmetrie: Het veld voldoet aan $V_{p,q}(x) = -V_{q,p}(x)$. Deze eigenschap zorgt ervoor dat als de distributies worden omgewisseld, de richting van de drift omkeert, wat garandeert dat $V=0$ wanneer $p=q$.

3. Instantiatie via Attractie en Repulsie

Het specifieke drifting field dat in het paper wordt gebruikt, is geïnspireerd door het mean-shift algoritme en wordt geconstrueerd met behulp van positieve monsters uit $p$ en negatieve monsters uit $q$:
$$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
Waarbij:

• $V^+_p(x)$ een attractieve kracht is die $x$ naar de data-distributie $p$ trekt.
• $V^-_q(x)$ een repulsieve kracht is die $x$ wegduwt van de gegenereerde distributie $q$.
• Deze krachten worden berekend met behulp van een kernelfunctie $k(x, y)$ (specifiek een exponentiële kernel gebaseerd op $\ell_2$-afstand) en genormaliseerd via softmax-operaties om stabiliteit te garanderen.

4. Trainingsdoel (Training Objective)

Het trainingsdoel is afgeleid van de fixed-point relatie bij evenwicht: $f(\epsilon) = f(\epsilon) + V_{p,q}(f(\epsilon))$. De loss-functie minimaliseert de afstand tussen de huidige voorspelling en een "bevroren doel" (de voorspelling plus de drift):
$$L = \mathbb{E}_{\epsilon} \left[ \| f_\theta(\epsilon) - \text{stopgrad}(f_\theta(\epsilon) + V_{p,q}(f_\theta(\epsilon))) \|^2 \right]$$
Deze formulering stelt de optimizer in staat om de distributie te laten evolueren door monsters te bewegen naar hun gedrifte versies. Cruciaal is dat de gradiënt niet terugpropageert door de distributie $q$ zelf, waardoor de complexiteit van het differentiëren door distributies wordt vermeden.

5. Feature Space en Guidance

• Feature Space: Om hoogdimensionale data te verwerken, wordt de drifting loss berekend in een feature space met behulp van een vooraf getrainde zelfgesuperviseerde encoder (bijv. ResNet-stijl MAE). Dit zorgt ervoor dat semantisch vergelijkbare monsters dicht bij elkaar blijven, wat rijkere gradiëntsignalen biedt.
• Classifier-Free Guidance (CFG): De methode ondersteunt van nature CFG door ongeconditioneerde data-monsters te mengen in de negatieve monster-distributie tijdens de training, wat controleerbare generatie bij inferentie mogelijk maakt zonder hertraining.

Belangrijkste Bijdragen

1. Nieuw Paradigma: Introduceert "Drifting Models", een framework waarbij de pushforward-distributie evolueert tijdens de trainingstijd, waardoor de noodzaak voor iteratieve inferentie wordt geëlimineerd.
2. Drifting Field: Stelt een nieuw loss-mechanisme voor gebaseerd op een drifting field dat evenwicht bereikt wanneer de gegenereerde en data-distributies overeenkomen, gestuurd door attractie naar data en repulsie van de huidige generatie.
3. State-of-the-Art One-Step Generation: Demonstreert dat een single-pass netwerk prestaties kan leveren die competitief zijn met, of superieur aan, multi-step diffusion en flow-gebaseerde modellen.
4. Robuustheid tegen Mode Collapse: Toy-experimenten suggereren dat de methode robuust is tegen mode collapse; als de gegenereerde distributie naar één modus convergeert, zal de attractie van andere modi in de data-distributie de drift blijven aansturen, waardoor de distributie wordt gedreven om te evolueren.

Experimentele Resultaten

De auteurs evalueren Drifting Models op ImageNet 256×256 in zowel latent als pixel-ruimte.

• Latent Space: Gebruikmakend van een standaard SD-VAE tokenizer, bereikt het model een 1-NFE FID van 1.54 (met een grotere L/2 model). Dit presteert beter dan alle voorgaande one-step methoden (bijv. iMeanFlow, AdvFlow) en is competitief met multi-step diffusion modellen (bijv. SiT-XL/2 met 250 stappen bereikt 2.06 FID).
• Pixel Space: Zonder latent tokenizer bereikt het model een 1-NFE FID van 1.61, wat aanzienlijk beter is dan voorgaande pixel-ruimte one-step methoden (bijv. StyleGAN-XL met 2.30 FID) en concurreert met multi-step pixel-ruimte modellen.
• Robotics Control: In de Diffusion Policy setting komt de one-step Drifting Model de prestaties van multi-step diffusion policies (100 NFE) op diverse robotica-manipulatietaken evenaren of overtreffen.

Betekenis en Claims

Het paper claimt nieuwe mogelijkheden te openen voor hoogwaardige, efficiënte generatieve modellering. Door iteratieve neurale netwerk-training te herformuleren als een mechanisme voor distributie-evolutie, daagt dit werk de noodzaak van SDE/ODE-formuleringen voor generatieve taken uit.

De auteurs benadrukken dat hun methode:

• Natuurlijk one-step inferentie (1-NFE) toelaat, wat de computationele kosten bij implementatie aanzienlijk vermindert.
• Niet afhankelijk is van distillatie van vooraf getrainde multi-step modellen, maar vanaf nul traint.
• Een conceptueel verschillend alternatief biedt voor diffusion en flow matching, waarbij de focus ligt op de update-regel ($x_{i+1} = x_i + \Delta x_i$) tijdens de training in plaats van tijdens de inferentie.

De auteurs blijven bescheiden over theoretische garanties, waarbij zij opmerken dat hoewel $p=q \implies V=0$ bewezen is, het omgekeerde ($V \to 0 \implies p \to q$) empirisch standhoudt en onder specifieke niet-degeneratieve aannames geldt, maar niet gegarandeerd is voor willekeurige vectorenvelden. Zij beschouwen hun werk als een startpunt voor het verkennen van andere realisaties van mechanismen voor de evolutie van distributies.