CrossFlow: One-Step Generation Across Latent and Pixel Spaces

Technische Samenvatting: CrossFlow: Eén-staps Generatie over Latente en Pixelruimtes

1. Probleemstelling

Huidige state-of-the-art beeldgeneratiesystemen, met name Latent Diffusion Models (LDMs) en flow-matching methoden, werken door de prior, de waarschijnlijkheidspad en het voorspellingsdoel binnen een enkele representatieruimte te definiëren. Hoewel LDMs de computationele efficiëntie verbeteren door de generatieve trajectorie in een gecomprimeerde autoencoder latente ruimte uit te voeren, vertrouwen ze op een tweestaps-pipeline: iteratieve denoising in de latente ruimte gevolgd door een aparte decoderingstap om latents terug naar pixels te mappen.

Deze architectuur introduceert een fundamentele mismatch:

• Optimalisatie Mismatch: De generator wordt geoptimaliseerd om latente verschuivingen te voorspellen, terwijl de uiteindelijke beeldkwaliteit afhangt van een apart getrainde decoder.
• Distributieverschuiving: Tijdens inferentie ontvangt de decoder gegenereerde latents die distributioneel kunnen verschillen van de schone latents geproduceerd door de encoder tijdens de training. Deze verschuiving kan fouten versterken.
• Toezichtbeperking: De generator kan niet direct getraind worden met beeldniveau perceptuele of adversariële verliezen, tenzij de output door de aparte decoder gaat, wat directe supervisie in de pixelruimte verzwakt.

Het artikel stelt de vraag of de ruisruimte en de outputruimte identiek moeten zijn, en stelt een raamwerk voor dat de efficiëntie van latente representaties behoudt terwijl het direct het beeld zelf voorspelt.

2. Methodologie: CrossFlow

CrossFlow is een cross-space flow formulering die ruizige latente inputs direct naar pixel-ruimte beelden mapt in een enkele functie-evaluatie. Het lost de ruimte-mismatch op door doelstellingen af te leiden waarbij de latente-ruimte snelheid wordt geëlimineerd uit het gesuperviseerde doel.

Kerntechnische Formulering:

• Verenigd Raamwerk: De auteurs definiëren een verenigde doelfunctie $Z$ voor single-step generatieve paradigma's. Standaard methoden vertrouwen op latente-ruimte snelheidvelden, wat cross-space generatie verhindert.
• Ontwerpcriteria: Om cross-space supervisie mogelijk te maken, stellen de auteurs twee criteria vast:
  1. Eliminatie van Latente-Snelheidsafhankelijkheden: De doelstelling moet structureel ontkoppeld zijn van latente-ruimte snelheden. De netwerkoutput mag geen latente verschuivingsvector zijn.
  2. Eindpunt Randvoorwaarden: Het raamwerk moet een geldig beeld opleveren bij de bestemming-boundary ($r=0$).
• De CrossFlow Doelstelling: Door aan deze criteria te voldoen, leiden de auteurs een algemene verliesfunctie af (Eq. 10) waarbij het netwerk $F_\theta$ een pixel-ruimte beeld $x_0$ voorspelt in plaats van een latente snelheid.
  • Het latente traject (gedefinieerd door $\alpha(t)$ en $\beta(t)$) bepaat nog steeds hoe schone latents worden gecorrumpeerd om inputs te creëren.
  • Het gesuperviseerde doel is het schone pixelbeeld.
  • De specifieke doelstelling gebruikt in experimenten (Eq. 11) omvat een tijdsafgeleide term berekend via een Jacobian-Vector Product (JVP) langs het latente traject, wat ervoor zorgt dat het model de correcte flow-dynamica leert zonder latente snelheden te voorspellen.

Implementatiedetails:

• Architectuur: Het model gebruikt een vaste, vooraf getrainde beeldencoder (VA-VAE) om afbeeldingen naar latents te mappen en een Vision Transformer (ViT) generator ($F_\theta$) om ruizige latents direct naar pixels te mappen.
• Trainingspipeline: Een tweestaps-aanpak wordt gebruikt om stabiliteit te garanderen:
  1. Train de encoder $E$ met standaard autoencoder doelstellingen (reconstructie + adversarieel).
  2. Bevries $E$ en train de generator $F_\theta$ vanaf nul met de CrossFlow doelstelling.
• Verliesfunctie: De totale doelstelling combineert de analytische CrossFlow loss met pixel-ruimte hulp-verliezen:
  • L1 Reconstructie Verlies: Pixel-wijze uitlijning.
  • Perceptueel Verlies: Structurele gelijkenis.
  • Adversarieel (GAN) Verlies: Hoogfrequente getrouwheid.
  • Opmerking: De analytische weegfactor van de flow-doelstelling wordt weggelaten tijdens de berekening van het perceptuele verlies om trainingsinstabiliteit te voorkomen.

3. Belangrijkste Bijdragen

1. Cross-Space Raamwerk: Introductie van CrossFlow, een one-step generatie raamwerk waarbij de input prior (latent) en de output afbeelding (pixel) in verschillende representatieruimtes verblijven.
2. Velocity-Vrije One-Step Doelstelling: Een nieuwe afleiding die een model in staat stelt zowel een one-step latent-naar-pixel generator als een decoder-vervanger te fungeren, waardoor de noodzaak voor een aparte decoder bij inferentie wordt geëlimineerd.
3. Directe Pixel-Ruimte Supervisie: De mogelijkheid om de generator direct te trainen met perceptuele en adversariële verliezen op het gegenereerde beeld, waarbij de tussenliggende decoder-stap wordt overgeslagen.
4. Empirische Validatie: Demonstratie dat een enkel model concurrerende FID-scores kan bereiken op ImageNet-1k (256×256) met een enkele functie-evaluatie (1-NFE) terwijl het standaard one-step latente methoden overtreft en multi-step latente diffusie baselines evenaart.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op ImageNet-1k bij 256×256 resolutie met klasse-conditionele generatie.

• Hoofdprestaties:

  • CrossFlow-XL (461M parameters) bereikt een FID van 1.62 met 1-NFE.
  • Dit is competitief met multi-step latente diffusie baselines zoals DiT-XL/2 (2.27 FID met 250 stappen) en SiT-XL/2 (2.06 FID met 250 stappen).
  • Het presteert significant beter dan andere one-step latente methoden (bijv. iMF-XL/2 met 1.72 FID) en vermijdt de noodzaak voor een aparte decoder.
  • CrossFlow-B (95M) en CrossFlow-L (315M) bereiken respectievelijk FIDs van 2.54 en 1.69, wat schaalbaarheid aantoont.

• Ablatie-studies:

  • Keuze van Encoder: De VA-VAE encoder (32 kanalen) leverde de beste resultaten op (2.54 FID), wat SD-VAE (4.08 FID) en FLUX.2 (3.81 FID) overtrof, wat suggereert dat de structuur van de latente representatie cruciaal is.
  • Verliescomponenten: Het verwijderen van het perceptuele verlies veroorzaakte een ernstige daling van de FID naar 30.21, en het verwijderen van het GAN verlies liet de FID dalen naar 8.86. De CrossFlow flow-doelstelling levert het structurele signaal, terwijl perceptuele en adversariële verliezen essentieel zijn voor hoogfrequente getrouwheid.
  • Decoder Vervanging: Wanneer gebruikt als decoder voor latents gegenereerd door een LightningDiT-B/1, verbeterde CrossFlow de FID consistent ten opzichte van de standaard VA-VAE decoder gedurende de hele training.

5. Betekenis en Claims

Het artikel claimt dat CrossFlow succesvol de kloof overbrugt tussen de efficiëntie van latente representaties en de getrouwheid van directe pixel-ruimte supervisie. Door de generator en decoder te verenigen in één enkel model getraind op ruizige latente inputs en pixel-ruimte targets, bereikt CrossFlow het volgende:

• Verwijdert de expliciete decoder-stap bij inferentie.
• Blootst het model tijdens de training aan de ruizige/imperfecte latents die optreden tijdens de generatie, in plaats van alleen schone encoder-outputs.
• Demonstreert dat cross-space flow-doelstellingen de computationele voordelen van latente diffusie kunnen combineren met de kwaliteitverbeteringen van directe pixel-niveau adversariële en perceptuele training.

De auteurs hanteren een bescheiden toon, waarbij zij opmerken dat hoewel de latente representatie en hulp-verliezen belangrijk zijn, de kernbijdrage de architecturale en doelstelling-gebaseerde verschuiving is die one-step, cross-space generatie mogelijk maakt zonder dat er aparte decoder-modules nodig zijn.