Latent Stochastic Interpolants

Probleemstelling

Stochastische Interpolanten (Stochastic Interpolants - SI) zijn een krachtig raamwerk voor generatieve modellering dat in staat is om willekeurige kansverdelingen te transformeren via een continue-tijd stochastisch proces. Echter, de toepassing van SI in gezamenlijk geoptimaliseerde latent variable modellen (waarbij een encoder, decoder en generatieve model samen worden getraind) is tot nu toe onontgonnen terrein.

De huidige beperkingen zijn:

1. Vereiste van directe observatie: SI vereist dat monsters uit zowel de prior-verdeling ($p_0$) als de doel-verdeling ($p_1$) direct beschikbaar zijn. In latent space modellen is de ware posterior $p_1(z)$ echter niet direct waarneembaar; deze moet worden benaderd door een encoder.
2. Vaste latent representaties: Bestaande methoden gebruiken vaak vaste latent representaties die niet optimaal zijn afgestemd op het generatieve proces, wat leidt tot inefficiëntie.
3. Rekenkundige last: Het toepassen van SI direct in hoge-dimensionale observatieruimtes (zoals pixels) is computatief zwaar.
4. Eenvoudige priors: Traditionele diffusion-modellen zijn vaak beperkt tot eenvoudige priors (zoals een standaard Gaussische verdeling), terwijl SI de flexibiliteit biedt om willekeurige priors te gebruiken, maar dit nog niet is geïntegreerd in een latent framework.

Methodologie: Latent Stochastic Interpolants (LSI)

De auteurs introduceren Latent Stochastic Interpolants (LSI), een raamwerk dat SI uitbreidt naar een latent ruimte die gezamenlijk wordt getraind met een encoder en decoder.

Kerncomponenten:

1. Gezamenlijke Optimalisatie: LSI traint drie componenten tegelijkertijd:

   • Een Encoder ($p_\theta(z_1|x_1)$) die hoge-dimensionale observaties afbeeldt op een latent ruimte.
   • Een Decoder ($p_\theta(x_1|z_1)$) die latent representaties reconstrueert naar observaties.
   • Een Latent SI-model ($h_\theta$) dat werkt binnen de geleerde latent ruimte om een prior $p_0(z_0)$ te transformeren naar de geaggregeerde posterior.

2. Continue-tijd ELBO: De belangrijkste theoretische bijdrage is het afleiden van een Evidence Lower Bound (ELBO) direct in continue tijd. Dit doelwit (loss function) combineert:

   • Een reconstructie-term: $\ln p_\theta(x_1|z_1)$.
   • Een dynamische mismatch-term: Een KL-divergentie tussen de variational posterior en het modelproces, uitgedrukt als een integraal over de snelheidsvelden.

3. Simulatie-vrije Training:

   • Om de hoge rekenkosten van het simuleren van Stochastische Differentiaalvergelijkingen (SDE's) tijdens elke training te omzeilen, construeren de auteurs een specifieke variational posterior.
   • Door aan te nemen dat de drift lineair is en het ruisadditief, kan de conditional density $p(z_t | z_1, z_0)$ analytisch worden afgeleid als een Gaussische verdeling.
   • Dit maakt het mogelijk om $z_t$ direct te reparameteriseren zonder SDE-simulatie:
     $$z_t = \eta_t \epsilon + \kappa_t z_1 + \nu_t z_0$$
     Waarbij $\epsilon \sim \mathcal{N}(0, I)$, en $\eta, \kappa, \nu$ functies zijn van de tijd $t$.

4. InterpFlow Parameterisatie:

   • De auteurs introduceren een specifieke parameterisatie genaamd InterpFlow om de stabiliteit van de training te verbeteren en de variatie in de gradiënten te verminderen. Deze methode vermijdt de numerieke instabiliteit die vaak optreedt bij termen met $\sqrt{1-t}$ in de noemer.

5. Flexibiliteit:

   • LSI ondersteunt willekeurige prior-verdelingen ($p_0$), niet alleen Gaussische.
   • Het biedt Classifier-Free Guidance (CFG) voor betere steekproefkwaliteit.
   • Het ondersteunt inversie (van beeld naar latent code) en flexibele steekproefstrategieën (deterministisch vs. stochastisch).

Belangrijkste Bijdragen

1. LSI Framework: Een nieuw, flexibel raamwerk voor het schaalbaar trainen van latent variable generatieve modellen met continue-tijd dynamiek, waarbij encoder, decoder en generatief model gezamenlijk worden geoptimaliseerd.
2. Unificerend Perspectief: Een nieuwe kijk op het integreren van de continue-tijd formulering van SI binnen latent variable modellen, gebruikmakend van inzichten uit continue-tijd stochastische processen.
3. Principiële ELBO-doelstelling: Een nieuwe ELBO die de sterke punten van SI behoudt (simulatie-vrije training, flexibele prior-keuze) maar deze combineert met de voordelen van gezamenlijke training in een latent ruimte.

Resultaten

De auteurs evalueren LSI op het standaard ImageNet (2012) dataset voor beeldgeneratie op verschillende resoluties (64x64, 128x128, 256x256).

• Kwaliteit (FID): LSI-modellen bereiken Frechet Inception Distance (FID) scores die vergelijkbaar zijn met observation-space SI-modellen van vergelijkbare grootte. Bijvoorbeeld, voor 128x128 bereikt LSI een FID van 3.12, wat vergelijkbaar is met observation-space baselines.
• Efficiëntie: Hoewel het totale aantal parameters vergelijkbaar is, is LSI aanzienlijk efficiënter tijdens het sampling (generatie). Omdat de encoder niet wordt gebruikt tijdens sampling en de decoder slechts één keer wordt aangeroepen (terwijl het latent model herhaaldelijk wordt uitgevoerd), leidt dit tot aanzienlijke besparingen in FLOPs (rekenkracht). Bij 256x256 resolutie wordt tot 64.9% minder FLOPs verbruikt tijdens sampling vergeleken met observation-space modellen.
• Gezamenlijke Training: Experimenten tonen aan dat gezamenlijke training (met een gewicht $\beta > 0$) cruciaal is. Het dwingt de encoder om latent representaties aan te passen die beter passen bij het generatieve proces, wat leidt tot betere FID-scores vergeleken met modellen waarbij encoder en decoder los van het generatieve model worden getraind.
• Prior Flexibiliteit: LSI werkt effectief met diverse priors (Uniform, Laplace, Gaussian Mixture), hoewel de Gaussische prior de beste prestaties levert. Dit bevestigt de flexibiliteit van het raamwerk.
• Parameterisatie: De InterpFlow parameterisatie presteert consistent beter dan alternatieven zoals OrigFlow, NoisePred en Denoising, met name in termen van trainingstabiliteit en FID-verbetering.

Betekenis en Impact

Dit paper is significant omdat het een brug slaat tussen twee krachtige maar tot nu toe gescheiden gebieden: Stochastic Interpolants en Latent Variable Modellen.

• Overcoming Limitations: Het lost het probleem op dat SI niet direct toepasbaar was op latent spaces door een principieel ELBO te bieden dat simulatie-vrije training mogelijk maakt.
• Schaalbaarheid: Het biedt een oplossing voor de rekenkundige last van generatieve modellen in hoge dimensies door het generatieve proces te verplaatsen naar een compactere latent ruimte, terwijl de flexibiliteit van SI behouden blijft.
• Toekomstperspectief: LSI opent de deur voor het gebruik van complexe, niet-Gaussische priors in latent generatieve modellen en biedt een unificerend theoretisch fundament voor continue-tijd generatieve processen. Het bewijst dat het gezamenlijk leren van representatie en generatie leidt tot efficiëntere en kwalitatief betere modellen dan de traditionele "pre-train encoder, then train generator" aanpak.

Kortom, LSI combineert de wiskundige elegantie en flexibiliteit van Stochastic Interpolants met de efficiëntie van latent space modellering, wat resulteert in een schaalbaar en krachtig raamwerk voor next-generation generatieve AI.