Continuous Adversarial Flow Models

Titel: Continuous Adversarial Flow Models (CAFM)

Auteurs: Shanchuan Lin, Ceyuan Yang, Zhijie Lin, Hao Chen, en Haoqi Fan (ByteDance Seed)

---

1. Het Probleem

Flow Matching (FM) heeft de afgelopen jaren aanzienlijke successen geboekt in generatieve modellering (zoals afbeeldingen en video). Echter, er blijft een kritiek probleem bestaan:

• Out-of-Distribution (OOD) samples: Zonder gebruik van "guidance" (sturing) produceren FM-modellen vaak samples die niet goed overeenkomen met de onderliggende data-distributie.
• Beperking van de Loss-functie: FM gebruikt een vastgestelde Euclidische afstand (MSE/L2) als criterium om de snelheidsveld (velocity field) te leren. Hoewel dit theoretisch convergeert naar de ware stroom bij oneindige modelcapaciteit, leidt het bij modellen met beperkte capaciteit tot een incorrecte generalisatie. De Euclidische afstand houdt geen rekening met de "manifold" (de complexe structuur) van de data, wat resulteert in samples die er visueel minder scherp of realistisch uitzien.
• Guidance als workaround: Hoewel guidance (zoals Classifier-Free Guidance) de samplekwaliteit verbetert, verandert het de sampling-distributie en is het niet de oplossing voor het modelleren van de oorspronkelijke data-distributie zelf.
• Bestaande oplossingen: Eerdere pogingen om dit op te lossen (zoals Riemannian Flow Matching of het gebruik van perceptuele loss) hebben nadelen: ze vereisen vaak handmatige definities van de data-manifold of gebruiken vaste netwerken die door de generator kunnen worden "gehackt".

2. Methodologie: Continuous Adversarial Flow Models (CAFM)

De auteurs introduceren CAFM, een uitbreiding van Adversarial Flow Models (AFM) naar continue tijd. Het doel is om de generalisatie van flow-matching-modellen te verbeteren door een geleerde discriminator in te voeren in plaats van een vast Euclidisch criterium.

Kerncomponenten:

• Generator (G): Net als bij standaard Flow Matching leert $G(x_t, t)$ het snelheidsveld $v_t$ om samples van ruis naar data te transporteren via een ODE (Ordinary Differential Equation).
• Discriminator (D): In plaats van een vast criterium, wordt een discriminator $D(x_t, t)$ getraind om het verschil tussen echte en gegenereerde stromen te detecteren.
• JVP (Jacobian-Vector Product): Een cruciale innovatie is dat de discriminatie plaatsvindt in de afgeleide ruimte. Omdat $v_t$ de afgeleide is van de positie $x_t$, wordt de discriminator geëvalueerd via het JVP:
  $$D_{jvp}(x_t, t, \dot{x}_t, \dot{t}) = \frac{\partial D}{\partial x_t}\dot{x}_t + \frac{\partial D}{\partial t}\dot{t}$$
  Hierbij wordt de echte snelheid $\bar{v}_t$ gebruikt als de "tangent" (richting) voor de echte data, en de voorspelde snelheid $G(x_t, t)$ voor de gegenereerde data.
• Doelfunctie: De training volgt een minimax-spel waarbij de generator probeert de discriminator te misleiden door de potentiële energie van de discriminator te maximaliseren in de richting van de gegenereerde stroom.
  • De loss voor de discriminator is gebaseerd op een bounded contrastive function (vergelijkbaar met Least Squares GAN): $f(a, b) = (a-1)^2 + (b+1)^2$.
  • Er wordt een "centering penalty" toegevoegd om te voorkomen dat de absolute waarde van de discriminator drift.
  • Een "optimal transport" regularisatie wordt gebruikt om de generator te dwingen de minimum-norm oplossing te kiezen (vergelijkbaar met het verminderen van de Euclidische norm).

Voor- en nadelen t.o.v. eerdere werk:

• In tegenstelling tot discrete tijd AFM's, werkt CAFM in continue tijd, wat de instabiliteit bij zeer kleine tijdstappen oplost.
• Het elimineert de noodzaak voor zware gradient penalties (zoals R1/R2) die vaak nodig zijn bij GANs, omdat de JVP-structuur het verdwijnen van gradienten (vanishing gradients) inherent vermindert.

3. Belangrijkste Bijdragen

1. Eerste toepassing van adversarial training in continue tijd: CAFM is het eerste model dat adversarial training succesvol integreert in continue-time normalizing flows.
2. Post-training strategie: De methode is primair ontworpen als een efficiënte post-training techniek voor bestaande Flow Matching-modellen. Dit maakt het zeer praktisch, omdat de zware pre-training met FM al is gedaan en CAFM alleen de generalisatie verbetert.
3. Verbeterde generalisatie: Door een geleerde discriminator te gebruiken in plaats van een vaste Euclidische loss, leert het model beter de onderliggende data-manifold, wat leidt tot samples die beter overeenkomen met de echte data-distributie.
4. Efficiënte implementatie: De auteurs tonen aan dat JVP en vectorized maps (vmap) efficiënt kunnen worden berekend met moderne frameworks (PyTorch), zelfs bij grote modellen.

4. Resultaten

De auteurs evalueren CAFM op ImageNet 256px (zowel in latent space als pixel space) en text-to-image taken.

ImageNet 256px (Post-training):

• SiT (Latent Space): De guidance-free FID (Fréchet Inception Distance) daalde drastisch van 8.26 naar 3.63 na slechts 10 epochs van finetuning. De guided FID verbeterde van 2.06 naar 1.53.
• JiT (Pixel Space): De guidance-free FID daalde van 7.17 naar 3.57. De guided FID verbeterde van 1.86 naar 1.80.
• Vergelijking: CAFM presteert beter dan bestaande state-of-the-art methoden (zoals SiD, DDT, REPA) in zowel guided als guidance-free settings, zonder dat de modelarchitectuur zelf is gewijzigd.

Text-to-Image (Z-Image model):

• GenEval: De score steeg van 0.81 naar 0.85.
• DPG-Bench: De score steeg van 83.7 naar 85.2.
• Visuele resultaten tonen duidelijk betere scherpte, textuur en consistentie in de gegenereerde afbeeldingen.

Trainen vanaf scratch:
Hoewel CAFM ook vanaf scratch kan worden getraind, convergeren deze modellen langzamer dan standaard FM-modellen, wat bevestigt dat de post-training strategie de meest efficiënte toepassing is.

5. Betekenis en Conclusie

Dit werk biedt een belangrijk inzicht in de beperkingen van Euclidische loss-functies in generatieve modellen. Door adversarial training te integreren in continue-time flows, slagen de auteurs erin om de "generalization gap" te dichten die vaak optreedt bij Flow Matching.

• Praktische impact: CAFM biedt een krachtige, plug-and-play methode om bestaande Flow Matching-modellen (zoals SiT en JiT) aanzienlijk te verbeteren zonder de noodzaak van volledige hertraining.
• Kwaliteitsverbetering: Het leidt tot samples met hogere fideliteit en een betere statistische overeenkomst met de echte data-distributie, zelfs zonder gebruik van sturing (guidance).
• Toekomstperspectief: Het werk opent de deur voor verdere onderzoek naar het combineren van adversarial objectives met continue dynamische systemen, niet alleen voor generatie, maar ook voor het verbeteren van de distributie-matching in diverse multimodale taken.

Kortom, CAFM bewijst dat het vervangen van een statische, Euclidische metriek door een dynamische, geleerde discriminator in een continue stroommodel leidt tot superieure generatieve prestaties.