Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models

1. Het Probleem

Hoewel generatieve modellen op basis van diffusie en stroming (flow-based models), zoals Flow Matching en Continuous Normalizing Flows (CNF), uitstekende prestaties leveren in het genereren van beelden, kampen ze met een fundamenteel en langdurig probleem: de berekening van de log-likelihood is uiterst rekenkundig duur.

• Sampling vs. Likelihood: Recent onderzoek heeft succesvol methoden ontwikkeld om het sampling (het genereren van beelden) te versnellen van honderden stappen naar slechts enkele stappen (bijv. 2-8 stappen) via distillatie. Echter, om de log-likelihood (de waarschijnlijkheid van een steekproef onder het model) nauwkeurig te berekenen, is het nog steeds nodig om de divergentie van het snelheidsveld langs de volledige trajectorie te integreren.
• De Bottleneck: Deze integratie vereist doorgaans 100 tot 1000 Neural Function Evaluations (NFEs). Zelfs geavanceerde oplosmethoden kunnen dit niet volledig oplossen zonder de nauwkeurigheid te verliezen.
• Gevolg: Dit maakt likelihood-gebaseerde taken, zoals modelvergelijking, fijnafstelling (fine-tuning) met reinforcement learning (PPO, DPO), en zelfgeleiding (self-guidance), onhaalbaar duur voor moderne snelle generatieve modellen. Bestaande snelle sampling-methoden geven vaak de mogelijkheid tot likelihood-berekening volledig op.

2. Methodologie: Fast Flow Joint Distillation (F2D2)

De auteurs introduceren F2D2, een modulair framework dat zowel het sampling-proces als de likelihood-evaluatie gelijktijdig versnelt door een gezamenlijke distillatie aan te wenden.

Kerninzicht

In CNF's worden de sampling-trajectorie en de evolutie van de log-dichtheid beschreven door een gekoppeld systeem van differentiaalvergelijkingen (ODE's) dat beide afhankelijk is van hetzelfde onderliggende snelheidsveld $v_\theta$:

1. Sampling: $\frac{d}{dt}x_t = v_\theta(x_t, t)$
2. Log-dichtheid: $\frac{d}{dt}\log p_t(x_t) = -\text{div}(v_\theta(x_t, t))$

Omdat de divergentie term ($\text{div}(v_\theta)$) een output is die afgeleid kan worden van hetzelfde snelheidsmodel, stelt F2D2 voor om beide componenten (de trajectorie en de cumulatieve divergentie) gelijktijdig te distilleren binnen één enkel model.

Technische Implementatie

• Gezamenlijke Flow Map: In plaats van het ODE-systeem stap voor stap op te lossen, leert het model een "flow map" $\Phi$ die direct de staat van tijd $t$ naar tijd $s$ voorspelt.
• Parametrisatie: Het model wordt geparametriseerd als een gezamenlijke flow map $\Phi_Y$ voor de toestand $y_t = (x_t, z_t)$, waarbij $z_t$ de log-likelihood is.
  $$\Phi_Y(\hat{y}_t, t, s) = \hat{y}_t + (s - t)f_\theta(\hat{x}_t, t, s)$$
  Hierbij voorspelt $f_\theta$ zowel de gemiddelde snelheid $u_\theta$ als de gemiddelde divergentie $D_\theta$.
• Architectuur: Het model gebruikt een gedeelde backbone (bijv. een U-Net) met twee aparte heads: één voor de snelheid (sampling) en één nieuwe head voor de divergentie (likelihood).
• Trainingsdoelen: Het framework combineert vier verliesfuncties om zowel de sampling als de likelihood te optimaliseren:
  1. Sampling Loss: Zorgt voor nauwkeurige snelheidsvoorspelling (tangent condition).
  2. Flow Map Consistency: Zorgt ervoor dat de flow map consistent is over tijdintervallen (via semigroup eigenschappen of Eulerian/Lagrangian voorwaarden, afhankelijk van de instantiatie zoals Shortcut of MeanFlow).
  3. Divergentie Matching: Zorgt dat de voorspelde divergentie overeenkomt met de ware divergentie van het snelheidsveld.
  4. Divergentie Consistency: Zorgt voor consistentie in de divergentie-berekening over tijdintervallen.

Het framework is compatibel met bestaande snelle sampling-methoden zoals Shortcut Models en MeanFlow, en vereist slechts een kleine aanpassing (de divergentie-head) aan bestaande modellen.

3. Belangrijkste Bijdragen

1. Eerste Methode voor Snelle Likelihood: F2D2 is, naar weten van de auteurs, de eerste methode die nauwkeurige log-likelihood-evaluatie mogelijk maakt met slechts een paar NFEs (bijv. 1-8 stappen) in diffusie- en flow-based modellen.
2. Modulair Ontwerp: Het is "plug-and-play" en kan worden toegepast op bestaande few-step modellen zonder de sampling-kwaliteit te verstoren.
3. Maximum Likelihood Self-Guidance: Als toepassing introduceren de auteurs een methode om de initiële ruis $x_0$ te optimaliseren op basis van de voorspelde likelihood. Dit werkt als een zelfgeleidingssysteem zonder externe reward-modellen.
4. Theoretische Onderbouwing: De auteurs bieden een wiskundige karakterisering van gezamenlijke flow maps en bewijzen dat het gezamenlijk distilleren van trajectorie en divergentie leidt tot een geldige flow map voor het gekoppelde ODE-systeem.

4. Resultaten

De methoden zijn getest op CIFAR-10, ImageNet 64x64 en CelebA-64.

• Nauwkeurige Likelihood: Waar bestaande few-step modellen (zoals standaard Shortcut of MeanFlow) onnauwkeurige of ongeldige likelihood-schattingen geven (vaak negatieve BPD-waarden die ver verwijderd zijn van de teacher), levert F2D2 gekalibreerde likelihoods op die dicht bij de teacher-modellen liggen (bijv. ~3.12 BPD op CIFAR-10 in plaats van -124).
• Behoud van Sampling Kwaliteit: De methode behoudt de hoge kwaliteit van de gegenereerde beelden (lage FID-scores). Bijvoorbeeld, MeanFlow-F2D2 met 2 stappen behoudt een FID van 2.59, vergelijkbaar met de originele MeanFlow.
• Superieure Prestaties met Self-Guidance: Een opmerkelijk resultaat is dat een 2-staps MeanFlow met F2D2 en self-guidance een lagere FID bereikt dan een 1024-staps Flow Matching model van dezelfde grootte op CIFAR-10. Dit toont aan dat snelle likelihood-berekening nieuwe optimalisatiemogelijkheden opent.
• 2D Checkerboard: Op synthetische data met een bekende ground truth wordt aangetoond dat F2D2 de dichtheidsverdeling nauwkeurig herstelt met slechts 1 NFE, terwijl traditionele methoden hierin volledig falen.

5. Betekenis en Impact

F2D2 lost een decennia lang bestaand computatiekundig knelpunt op in generatieve modellering. Door de drempel voor likelihood-berekening drastisch te verlagen, opent het de deur voor:

• Efficiëntere Fine-tuning: Likelihood-gebaseerde optimalisatietechnieken (zoals PPO of DPO) worden nu haalbaar voor snelle generatieve modellen.
• Modelanalyse: Het maakt directe vergelijking van modellen mogelijk zonder duizenden ODE-stappen.
• Nieuwe Algoritmen: Het toont aan dat likelihood-gebaseerde zelfgeleiding (self-guidance) de sampling-kwaliteit aanzienlijk kan verbeteren, zelfs met zeer weinig stappen.

Kortom, F2D2 transformeert flow-based modellen van modellen die alleen goed zijn voor generatie naar modellen die ook efficiënt en nauwkeurig kunnen "redeneren" over de waarschijnlijkheid van hun eigen output, wat essentieel is voor de volgende generatie AI-toepassingen.