Technische Samenvatting: Het Verkennen en Benutten van Stabiliteit in Latente Flow Matching

1. Probleemstelling

Het trainen van generatieve modellen, met name Diffusiemodellen en hun Flow Matching (FM) tegenhangers, blijft rekenkundig duur. Deze modellen vereisen doorgaans enorme datasets, uitgebreide rekenmiddelen, lange trainingsduur en, in conditionele opstellingen, uitgebreide annotaties. Hoewel Flow Matching een deterministische formulering biedt voor steekproefgeneratie via Gewone Differentiaalvergelijkingen (ODE's) die vaak het aantal steekproefstappen reduceert in vergelijking met Stochastische Differentiaalvergelijkingen (SDE's), blijft de kostprijs van het trainen van deze modellen op volledige datasets een knelpunt.

De centrale vraag die door dit werk wordt aangepakt, is of de datasetgrootte en modelcapaciteit aanzienlijk kunnen worden verminderd zonder in te leveren op de kwaliteit van de generatie. Specifiek onderzoeken de auteurs de stabiliteit van Latente Flow Matching (LFM)-modellen onder grote verstoringen, waaronder dataset-subsampling (pruning), architecturale wijzigingen en gewijzigde trainingsconfiguraties. Zij stellen dat, indien de transporttrajectoires van LFM inherent stabiel zijn, deze eigenschap kan worden benut om efficiëntere trainings- en inferentie-algoritmen af te leiden.

2. Methodologie

2.1. Theoretische Grondslag: Stabiliteit in Flow Matching

Het werk bouwt voort op de observatie dat Rectified Flow een gesloten-vorm expressie toestaat voor het snelheidsveld. Voor een eindige dataset is het optimale snelheidsveld een gewogen som van trainingssteekproeven, waarbij de gewichten worden bepaald door een softmax-functie over de kwadratische afstanden tussen de huidige staat en geschaalde trainingssteekproeven. De auteurs merken op dat deze softmax vroeg in het transportfase "piekt", wat betekent dat een enkele trainingssteekproef vaak het snelheidsveld domineert voor een gegeven trajectoires.

Deze theoretische inzichten suggereren dat het transporttrajectoires grotendeels wordt bepaald door de aanwezigheid van specifieke "dominerende" steekproeven. Bijgevolg zou het verwijderen van niet-dominante steekproeven (pruning) de meerderheid van de trajectoires onveranderd moeten laten. De auteurs formaliseren dit door een toewijzingsmetriek te definiëren: voor een trajectoires dat begint bij ruis $x_0$, bepalen zij welke trainingssteekproef $x_i$ het eindigt. Zij hypothetiseren dat deze toewijzing invariant blijft, zelfs wanneer een groot deel van de dataset wordt verwijderd, mits de dominerende steekproef voor dat trajectoires behouden blijft.

2.2. Data-Pruning Strategieën

Om deze stabiliteit te benutten, stellen de auteurs drie data-pruning-criteria voor die zijn aangepast van discriminatieve modellen, maar zijn herformuleerd voor de FM-doelstelling (die gedeelde ruispaden en tijdstappen gebruikt):

1. Op gradiënten gebaseerde scoring (G): Steekproeven worden gerangschikt op basis van de grootte van de gradiëntnorm van de loss.
2. Op loss gebaseerde scoring (L): Steekproeven worden gerangschikt op basis van de grootte van de loss die wordt berekend langs gedeelde ruispaden.
3. Op clustering gebaseerde scoring (C): Steekproeven worden geclusterd in een vooraf getrainde embedding-ruimte (bijv. CLIP). De auteurs stellen twee specifieke varianten voor:
   • Proportioneel ($C_p$): Selecteert steekproeven evenredig aan de clustergrootte.
   • Gebalanceerd ($C_b$): Selecteert een gelijke hoeveelheid steekproeven uit elke cluster om distributie-ongelijkheden te mitigeren.
   • Op kernel gebaseerd ($C_b^\kappa$): Selecteert steekproeven gretig om het gemiddelde van de Gaussische kernel van de cluster in de latente ruimte te benaderen.

De auteurs evalueren ook de "inverse" van deze criteria (het selecteren van de laagste scores) om de robuustheid van de stabiliteitsclaim te testen.

2.3. Grof-naar-Fijn (C2F) Inferentie

Door de stabiliteit van trajectoires over modelcapaciteiten heen te benutten, stellen de auteurs een twee-traps inferentiestrategie voor:

1. Groot model: Een lichtgewicht architectuur (bijv. DiT-S/2) wordt getraind op een geprepareerde subset van de data. Het behandelt de eerste fase van het trajectoires ($t \in [0, t_0)$).
2. Fijn model: Een hoge-capaciteit architectuur (bijv. DiT-XL/2), vooraf getraind op de volledige dataset, behandelt de tweede fase ($t \in [t_0, 1]$).

Om een soepele overgang op het "naad"-punt $t_0$ te garanderen, wordt het Groot model fijngeschaald met een continuïteitsloss. Deze loss minimaliseert het verschil tussen de snelheidsvelden van het Groot en Fijn model bij $t_0$. Het trajectoires van het Fijn model wordt gebruikt om de doelstaat $x_{t_0}$ te genereren via ODE-inversie (achterwaartse integratie), zodat het Groot model leert zich aan te passen aan het pad van het Fijn model.

3. Belangrijkste Bijdragen

1. Empirische Demonstratie van LFM-Stabiliteit: De auteurs tonen aan dat onafhankelijk getrainde LFM-modellen convergeren naar zeer consistente transporttrajectoires onder dataset-subsampling, architecturale wijzigingen en verschillende willekeurige initialisaties. Deze stabiliteit is gekoppeld aan de theoretische eigenschappen van de FM-doelstelling.
2. Praktische Pruning-algoritmen: Zij introduceren drie pruning-criteria (Gradiënt, Loss, Clustering) die geschikt zijn voor FM. Zij tonen aan dat agressieve pruning (bijv. het verwijderen van 50-95% van de data) de prestaties niet verslechtert en zelfs de convergentiesnelheid en generatiekwaliteit kan verbeteren onder beperkte rekenbudgetten.
3. Versnelling van Grof-naar-Fijn Inferentie: Zij stellen een generatiestrategie met twee modellen voor die het trajectoires splitst. Deze aanpak bereikt een aanzienlijke versnelling van de inferentie (ongeveer 2,15$\times$) terwijl de generatiekwaliteit behouden blijft, mits de stabiliteitsaanname geldt voor de modelarchitecturen.
4. Bias-mitigatie via Gebalanceerde Pruning: De auteurs tonen aan dat gebalanceerde clustering-pruning ($C_b$) biases in gegenereerde data (bijv. geslacht, huidskleur) kan mitigeren zonder expliciete label-bewuste balancing te vereisen, door de representatie van semantische clusters in de trainingsset te egaliseren.

4. Experimentele Resultaten

De methoden werden geëvalueerd op CelebA-HQ, FFHQ en ImageNet-1K.

• Stabiliteit onder Pruning: Op CelebA-HQ produceerden modellen die waren getraind op disjuncte subsets (gepruned op 50%) visueel vergelijkbare afbeeldingen wanneer ze startten vanuit dezelfde ruisseed $x_0$. De ArcFace-cosine-ähnlichkeit tussen geprepareerde en ongepruneerde modellen varieerde van 0,79 tot 0,83, significant hoger dan willekeurige paren (0,37).
• Prestatiemetrics:
  • CelebA-HQ: Het verwijderen van 50% van de data met behulp van gebalanceerde clustering ($C_b$) verbeterde de Fréchet Inception Distance (FID) licht van 24,24 (ongepruned) naar 22,80.
  • ImageNet: Het verwijderen van 75% van de data ($pr=0.75$) resulteerde in snellere convergentie en behield of verbeterde licht de FID in vergelijking met de volledige dataset. Zelfs bij 95% pruning convergeerde het model aanvankelijk sneller, hoewel de prestaties na een bepaald aantal iteraties verslechterden.
• Robuustheid tegen Verstoringen:
  • Architectuur: Het vergroten van de modelcapaciteit (DiT-S/2 naar DiT-XL/2) resulteerde in minimale trajectoiresveranderingen (ähnelijkheid $s=0,81$). Overschakelen naar een U-Net-architectuur verlaagde de ähnelijkheid ($s=0,55$) maar behield grove attributen.
  • Latente Ruimte: De stabiliteit is gevoelig voor veranderingen in de parametrisatie van de latente ruimte (bijv. het wijzigen van de VAE-seed of het toepassen van invertibele transformaties), wat de stabiliteit verbrak.
  • Doelstelling: Het overschakelen van Flow Matching naar score-gebaseerde diffusie veranderde de trajectoires volledig, wat aangeeft dat stabiliteit specifiek is voor de FM-doelstelling.
• Efficiëntie van Grof-naar-Fijn: Op CelebA-HQ reduceerde de C2F-aanpak met $t_0=0,7$ de inferentietijd van 93,95 ms/afbeelding (alleen Fijn) naar 43,53 ms/afbeelding (ongeveer 2,15$\times$ versnelling) terwijl de FID behouden bleef.
• Fairness: Gebalanceerde pruning ($C_b$) verminderde de KL-divergentie naar een uniforme verdeling over attributen (geslacht, leeftijd, huidskleur) effectiever dan ongepruneerde modellen, ondanks dat het label-agnostisch was.

5. Betekenis en Claims

De auteurs claimen dat de waargenomen stabiliteit een karakteristieke eigenschap is van het Rectified Flow Matching-paradigma onder vaste latente representaties, en niet een artefact van specifieke trainingsalgoritmen.

• Efficiëntie: De primaire betekenis is het vermogen om trainingsduur en data-eisen te verminderen zonder in te leveren op kwaliteit. De auteurs tonen aan dat voor grote datasets zoals ImageNet pruning de convergentie kan versnellen, en voor kleinere datasets zoals CelebA-HQ het de generatiekwaliteit kan verbeteren.
• Versnelling van Inferentie: De C2F-aanpak biedt een principiële manier om inferentiekosten te verlagen door de vroege, grove fasen van generatie uit te besteden aan lichtgewicht modellen, waarbij wordt vertrouwd op de stabiliteit van het trajectoires om ervoor te zorgen dat het model met hoge capaciteit alleen de laatste details hoeft te verfijnen.
• Bias-Controle: Het werk suggereert dat dataset-pruning, specifiek via gebalanceerde clustering, kan dienen als een mechanisme om distributie-ongelijkheden in gegenereerde data te controleren en te mitigeren, en zo een weg biedt naar eerlijkere generatieve modellen zonder complexe label-bewuste herweging.

Het artikel concludeert dat stabiliteit, hoewel robuust tegen veranderingen in data en modelcapaciteit, gevoelig is voor veranderingen in de structuur van de latente ruimte en de generatieve doelstelling zelf. Deze inzichten stellen praktici in staat om meer hulpbronnen-efficiënte pijplijnen te ontwerpen terwijl hoge-fideliteit generatie behouden blijft.