Technische Samenvatting: Afstemmen van Latente Geometrie voor Sferische Flow Matching bij Beeldgeneratie

Probleemstelling

Latente flow matching (en diffusiemodellen) voor beeldgeneratie werkt doorgaans door Gaussisch ruis naar Variational Autoencoder (VAE) latents te transporteren langs lineaire paden in de Euclidische ruimte. Deze aanpak negeert echter twee kritieke structurele eigenschappen van hoogdimensionale VAE-latents:

1. Concentratie van Maat: Zowel de Gaussische ruisprior als de gecodeerde datalatents concentreren zich in dunne sferische schillen in plaats van het volume van de ruimte te vullen. Een rechte Euclidische lijn (koorde) die twee punten op deze schillen verbindt, loopt door binnenste stralen die door geen enkele verdeling worden bezet, waardoor het model transities moet leren door regio's van de latente ruimte die zelden worden bevolkt.
2. Directionele Dominantie: De output van de decoder is overwegend gevoelig voor de richting (hoekcomponent) van een latent token in plaats van zijn lengte (radiale component). Experimenten waarbij de straal en richting van tokens van buren uit dezelfde klasse worden verwisseld, tonen aan dat het behouden van de richting terwijl de straal wordt gewijzigd, leidt tot minimale perceptuele verandering, terwijl het wijzigen van de richting aanzienlijke semantische verschuivingen veroorzaakt.

Standaard lineaire flow matching negeert deze feiten en wijst een significant deel van het trainingsdoel voor snelheid toe (tot 50-90%, afhankelijk van de tokenizer) aan radiale beweging. Omdat de decoder minder gevoelig is voor de straal, is deze supervisie inefficiënt, wat leidt tot langzamere convergentie en suboptimale beeldkwaliteit (FID) in vergelijking met methoden die de onderliggende geometrie respecteren.

Methodologie

De auteurs stellen Sferische Flow Matching voor, een methode die de latente geometrie afstemt op de sferische structuur van de data zonder de diffusiearchitectuur te wijzigen of hulpencoders te vereisen. De aanpak bestaat uit vier kerncomponenten:

1. Token-gewijze Sferische Projectie:
   De methode projecteert elke latent token $z$ van een voorgeïmplementeerde VAE op een hypersfeer met vaste straal $S^{d-1}(\sqrt{d})$. Dit wordt bereikt via een $L_2$-normalisatiestap die wordt ingevoegd tussen de bevroren encoder en de decoder:
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   Dit zorgt ervoor dat alle tokens op een sfeer liggen met een straal die overeenkomt met de concentratiestraal van de standaard Gaussische prior.

2. Sferische Prior:
   In plaats van ruwe Gaussische ruis te gebruiken, wordt het ruisuiteinde $z_0$ geconstrueerd door isotrope Gaussische ruis $\epsilon \sim \mathcal{N}(0, I)$ te sample en op de sfeer te projecteren:
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   Dit resulteert in een uniforme verdeling op de sfeer, die overeenkomt met de geometrie van de geprojecteerde data.

3. Decoder Finetuning:
   Om de geometrische beperking op te vangen, wordt de decoder enkele epoches (5 epoches in experimenten) gefinetuned terwijl de encoder en projectielaag bevroren blijven. Het reconstructiedoel behoudt standaard pixelverliezen (L1, LPIPS, adversariaal) maar laat de KL-term vallen, aangezien de geprojecteerde latent deterministisch is.

4. Slerp-transport en Tangent Snelheid:
   Het transportpad tussen het ruis- en datauiteinde wordt vervangen door Sferische Lineaire Interpolatie (slerp), dat het geodetische boogpad op de sfeer volgt.

   • Geodetisch Pad: Het pad $z_t$ blijft voor alle $t \in [0, 1]$ op de sfeer.
   • Snelheidsdoel: De conditionele snelheid $u_t$ is de tijdsafgeleide van de geodetische, geprojecteerd op de raakruimte $T_{z_t}S^{d-1}$. Dit zorgt ervoor dat de snelheid puur hoekmatig is, met per constructie een nul radiale component.
   • Training & Sampling: Het model wordt getraind om deze tangent-snelheid te voorspellen. Tijdens inferentie gebruikt de sampler de exponentiële kaart om de snelheid te integreren, wat garandeert dat de trajectorie op de sfeer blijft zonder dat op elke stap radiale projectiecorrecties nodig zijn.

Belangrijkste Bijdragen

Het artikel identificeert en kwantificeert een "straal-schil-mismatch" in standaard latente flow matching en introduceert een geometrische interventie om deze op te lossen:

1. Identificatie van Mismatch: De auteurs tonen aan dat lineaire paden het model dwingen radiale beweging te voorspellen waarvoor de decoder ongevoelig is, wat capaciteit verspillen en convergentie vertragen.
2. Token-gewijze Projectie: Ze introduceren een methode om voorgeïmplementeerde VAE-latents te beperken tot een sfeer met vaste straal met uitsluitend decoder finetuning, waardoor de noodzaak om de encoder opnieuw te trainen of hulprepresentatie-encoders te gebruiken, wordt vermeden.
3. Slerp-gebaseerde Flow Matching: Ze stellen voor flow matching te trainen langs slerp-geodetische met tangent-snelheidsdoelen, zodat het volledige traject de sferische variëteit respecteert.
4. Empirische Validatie: De methode blijkt prestaties te verbeteren over meerdere tokenizer-families (FLUX.2, VA-VAE, REPA-E FLUX.1) en modelschalen (SiT-B tot SiT-XL) zonder architecturale wijzigingen.

Resultaten

Experimenten op class-conditional ImageNet-256-generatie tonen consistente verbeteringen:

• FID-verbeteringen: De sferische-slerp-methode bereikt lagere FID-scores dan standaard lineaire flow matching over alle geteste configuraties. Bijvoorbeeld, op SiT-B/2 met FLUX.2 daalt de FID van 26,35 (Lineair) naar 20,55 (Sferisch-slerp).
• Trainings-efficiëntie: De methode bereikt een doel-FID van 30 in ongeveer 2,2$\times$ minder trainingsstappen in vergelijking met de standaard lineaire baseline.
• Generalisatie: Winsten worden waargenomen over verschillende leidingschalen (CFG) en modelgroottes (B en XL), wat aangeeft dat het voordeel inherent is aan de latente geometrie in plaats van specifieke hyperparameter-tuning.
• Ablatiestudies:
  • Het verwisselen van decoders tussen standaard en sferische pipelines verslechtert de prestaties in beide richtingen, wat bevestigt dat de geleerde flow nauw gekoppeld is aan de specifieke latente geometrie.
  • Het gebruik van een "schil-gedecomposeerd" pad (interpolatie van straal en richting apart) op standaard latents levert niet dezelfde winsten op, wat bevestigt dat het volledig verwijderen van de radiale vrijheidsgraad cruciaal is.
  • De reconstructiekwaliteit (rFID) van de gefinetuned sferische decoders blijft concurrerend, wat aantoont dat de geometrische beperking de mogelijkheid van de VAE om beelden te reconstrueren niet ernstig ondermijnt.

Betekenis

Het artikel vestigt latente ondersteuningsgeometrie als een kritieke ontwerpas voor generatieve modellering. Het betoogt dat de vorm van de latentverdeling de ruimte definieert die de generator moet doorkruisen, en dat het negeren van deze geometrie (door Euclidische lineariteit aan te nemen) leidt tot inefficiëntie.

De betekenis van dit werk ligt in het vermogen om de kwaliteit van beeldgeneratie en trainsnelheid te verbeteren zonder:

• De diffusiearchitectuur te wijzigen (bijv. de SiT-ruggengraat).
• Hulpencoders toe te voegen (in tegenstelling tot Riemanniaanse flow matching op bevroren DINOv2-features).
• De VAE-encoder vanaf nul opnieuw te trainen.

Door bestaande voorgeïmplementeerde latents simpelweg op een sfeer te projecteren en het flowpad af te stemmen op de resulterende geodetische, bereikt de methode state-of-the-art resultaten met minimale rekenkundige overhead. De auteurs suggereren dat toekomstig werk deze geometrische beperkingen moet verkennen onder sterkere compressieregimes en tokenizers moet ontwerpen die native overeenkomen met deze sferische structuur.