Posterior Augmented Flow Matching

1. Probleemstelling: Sparse Supervisie en Flow Collapse

Het artikel adresseert een fundamentele beperking in Flow Matching (FM), een klasse van continue-tijd generatieve modellen die een tijdsafhankelijk vectorveld leren om samples van een eenvoudige prior (bijvoorbeeld Gaussisch ruis) naar een complexe data-distributie te transporteren.

• De Kernkwestie: In standaard FM berust training op spare, één-op-één supervisie. Voor elke gegeven tussenliggende toestand $z_t$ langs een traject dat een ruissample $\epsilon$ verbindt met een doelpunt $z$, wordt het model gesuperviseerd om de snelheid naar dat enkele specifieke doel $z$ te voorspellen.
• Realiteit in Hoge Dimensies: In ruimten met hoge dimensies (zoals afbeeldingen) kan een enkel tussenliggend punt $z_t$ geldig liggen op vele verschillende trajecten die leiden tot diverse, geldige doelafbeeldingen die aan dezelfde voorwaarde voldoen (bijvoorbeeld "een foto van een hond").
• Gevolgen:
  • Hoge Variantie: De gradiënt op een specifiek $z_t$ wordt geschat vanuit één enkel ruisig signaal.
  • Flow Collapse: Het model neigt om specifieke bron-doelkoppelingen te memoriseren in plaats van de ware continue mapping te leren. Dit leidt tot "flow collapse", waarbij diverse invoer worden gemapt naar overdreven vergelijkbare of verwarden uitvoer, wat faalt in generalisatie.

2. Methodologie: Posterior-Augmented Flow Matching (PAFM)

De auteurs stellen PAFM voor, een theoretisch onderbouwde generalisatie van FM die één-doel supervisie vervangt door een verwachting over de posterior-distributie van geldige doelvolledigingen.

Theoretische Formulering

In plaats van het verlies te minimaliseren tegenover één enkel doel $z$, minimaliseert PAFM de verwachte verliezen over een set plausibele doelen $\{z_j\}$ getrokken uit de posterior-distributie $p(z_j | z_t, y)$, waarbij $y$ de voorwaarde is.

• Doel: De nieuwe verliesfunctie aggregeert gradiënten van meerdere kandidaat-doelen, gewogen naar hun posterior-kans.
• Ongeschatte Schatter: De auteurs bewijzen (Stelling 1) dat PAFM een ongeschatte schatter is van de oorspronkelijke FM-doelstelling. In verwachting geldt: $L_{PAFM} = L_{FM}$.
• Variantiereductie: Cruciaal bewijst Stelling 2 dat door informatie te aggregeren van $K > 1$ plausibele trajecten per tussenliggend punt, PAFM de gradiëntvariantie ondergrenst in vergelijking met standaard FM. De variantiereductiefactor is gerelateerd aan de Effective Sample Size (ESS) van het bemonsteringsschema.

Praktische Implementatie

Aangezien de ware posterior onberekenbaar is, factoriseert PAFM het probleem in twee benaderingen met behulp van Importance Sampling:

1. Kans van de Tussenliggende: $p(z_t | z_j)$: Hoe plausibel is het dat doel $z_j$ de huidige tussenliggende $z_t$ heeft gegenereerd? Dit wordt berekend met behulp van het bekende conditionele waarschijnlijkheidspad (Gaussisch).
2. Kans van de Voorwaarde: $p(y | z_j)$: Hoe goed voldoet doel $z_j$ aan de voorwaarde $y$?
   • Klasse-geconditioneerd: Deterministisch (1 als labels overeenkomen, 0 anders).
   • Tekst-geconditioneerd: Benaderd met visueel-taalmodellen (bijvoorbeeld CLIP/SigLIP) of verondersteld uniform als kandidaten vooraf zijn gefilterd.

Trainingsalgoritme:
Voor elke trainingsstap bemonster het model een tussenliggende $z_t$. In plaats van één doel, haalt het een set van $K$ kandidaat-doelen $\{z_j\}$ op. Het model voorspelt de snelheid naar een gewogen gemiddelde van deze doelen, waarbij de gewichten evenredig zijn met $p(z_t | z_j) \cdot p(y | z_j)$.

3. Doelselectiestrategieën

Het artikel onderzoekt verschillende praktische strategieën om de kandidaatset $\{z_j\}$ te construeren:

1. Dichtstbijzijnde Buren in Latente Ruimte (Primair): Het gebruik van FAISS om $K$ dichtstbijzijnde buren van de trainingsafbeelding in de latente ruimte op te halen. Dit waarborgt geometrische nabijheid en semantische gelijkenis.
2. Willekeurige Augmentaties: Het creëren van doelen door willekeurige crops, herschalingen en rotaties toe te passen op de bronafbeelding. Dit verhoogt de ruimtelijke diversiteit.
3. VAE Momentbemonstering: Het $K$ keer bemonsteren uit de Gaussische distributie (moment) van de VAE-encoding van de bronafbeelding. Dit biedt kleine verstoringen rond de gemiddelde encoding.

4. Belangrijkste Bijdragen

• Theoretische Generalisatie: PAFM is bewezen een ongeschatte generalisatie van FM te zijn die gradiëntvariantie reduceert door supervisie te aggregeren over de volledige posterior van compatibele doelen.
• Mitigatie van Flow Collapse: Door een dichter supervisiessignaal te bieden, voorkomt PAFM dat het model specifieke koppelingen memoriseert, wat leidt tot robuustere vectorvelden.
• Flexibiliteit: Het raamwerk is agnostisch over hoe kandidaten worden geselecteerd, wat integratie toelaat met diverse datamodi en architecturen.
• Efficiëntie: De methode introduceert verwaarloosbare rekenkundige overhead omdat kandidaatsets vooraf kunnen worden berekend (bijvoorbeeld via FAISS-indexen) en tijdens training kunnen worden geladen.

5. Experimentele Resultaten

De auteurs hebben PAFM geëvalueerd over verschillende schalen, architecturen en datasets:

• Datasets:
  • ImageNet-1K (Klasse-geconditioneerd): Getest op SiT-B/2 en SiT-XL/2 modellen.
  • CC12M (Tekst-naar-Afbeelding): Getest op MMDiT-architectuur.
• Metrieken: FID50K (Fréchet Inception Distance), sFID, Precisie, Recall.
• Prestatiewinsten:
  • ImageNet: PAFM verbeterde FID met maximaal 3,4 punten (bijvoorbeeld SiT-XL/2 verbeterde van 11,14 naar 9,85 FID).
  • CC12M: PAFM verlaagde FID van 10,37 (FM) naar 9,45.
  • Consistentie: Verbeteringen werden waargenomen over alle modelgroottes (SiT-B/2, SiT-XL/2) en architecturen (SiT, MMDiT).
• Gradiëntvariantie: Empirische analyse toonde aan dat PAFM de mini-batch gradiëntvariantie met ongeveer 4x reduceerde in vergelijking met FM, wat leidt tot stabielere trainingsdynamiek.
• Rekenkosten:
  • Doorvoer nam slechts af met 6,6%.
  • Piekmemorygebruik nam toe met 0,4%.
  • FLOPs per stap bleven onveranderd.

6. Betekenis

PAFM vertegenwoordigt een significante stap voorwaarts in generatieve modellering door de "spare supervisie"-bottleneck inherent aan flow-gebaseerde en diffusiemodellen aan te pakken.

• Robuustheid: Het verbetert fundamenteel de stabiliteit van het trainen van generatieve modellen met hoge dimensies, wat het risico op mode collapse en artefacten vermindert.
• Efficiëntie: Het bereikt state-of-the-art prestatieverbeteringen met minimale extra rekenkracht, waardoor het een zeer praktische drop-in vervanging is voor standaard Flow Matching-doelstellingen.
• Generaliseerbaarheid: De aanpak is toepasbaar op zowel klasse-geconditioneerde als tekst-naar-afbeelding generatie, wat suggereert dat het een breed scala aan continue generatieve modellen ten goede kan komen, verder dan de specifieke geteste architecturen.