Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models

Probleemstelling

Latente Diffusiemodellen (LDMs) zijn de huidige standaard voor hoogwaardige en efficiënte beeldgeneratie. Ze werken door diffusieprocessen te trainen in een compacte, geleerde latente ruimte in plaats van direct in de pixelruimte. Hoewel eerdere onderzoek zich vooral richtte op reconstructie-accuraatheid en semantische uitlijning van deze latente ruimte, identificeerden de auteurs een kritiek, maar vaak over het hoofd gezien probleem: de robuustheid tegenover steekproefperturbaties (sampling perturbations).

De auteurs observeren een tegenintuïtief fenomeen: modellen met een tokenizer die bijna perfecte reconstructie biedt en een lage diffusieverlies (loss), genereren vaak visueel inferieure beelden vergeleken met modellen met iets hogere reconstructiefouten.

• De oorzaak: De veelgebruikte $\beta$-VAE-tokenizers (met een zeer kleine KL-divergentie-term) leiden tot een overmatig compacte latente manifold.
• Het gevolg: Omdat de latenten extreem dicht bij elkaar liggen (variatie $\sigma^2$ is verwaarloosbaar klein), zijn ze uiterst gevoelig voor de stochastische ruis die inherent is aan het diffusie-samplingproces. Kleine perturbaties duwen de latenten al snel buiten de geldige data-manifold, wat leidt tot decodeerfouten en visuele degradatie.

Methodologie: Variance Expansion Loss (VE Loss)

Om dit probleem op te lossen, stellen de auteurs een nieuwe trainingsstrategie voor die een latente ruimte creëert die robuust is tegen perturbaties, zonder de reconstructiekwaliteit te offeren.

1. Theoretische Analyse:

   • De auteurs tonen aan dat de reconstructieloss ($L_{rec}$) een sterke gradiënt creëert die de latent-variatie ($\sigma^2$) naar nul duwt (variatie-inzakking).
   • De standaard KL-regularisatie in VAE's is in de context van diffusiemodellen niet noodzakelijk en kan zelfs schadelijk zijn omdat hij de expressiviteit van de latente ruimte onderdrukt en de reconstructie verstoort.

2. De Oplossing: Variance Expansion (VE) Loss:
   In plaats van de KL-term te gebruiken, introduceren ze een nieuwe loss-functie die de variatie actief uitbreidt:

   • $L_{var}$: Een term die de variatie $\sigma^2$ expliciet aanmoedigt om niet te instorten. De formule is $L_{var}(\sigma) = \frac{1}{\sigma^2 + \delta}$. Dit creëert een sterke tegengestelde gradiënt wanneer $\sigma$ klein is, waardoor instorting wordt voorkomen.
   • $L_{reg}$: Een empirische regularisatieterm ($e^{|z|-\tau}$) om te voorkomen dat de totale magnitude van de latenten onbeperkt groeit.
   • Totaal Doel: $L = L_{rec} + \lambda_1 L_{var} + \lambda_2 L_{reg}$.

3. Adaptief Evenwicht:
   Door het "adversariale" spel tussen de reconstructieloss (die kleine variatie wil) en de VE-loss (die voldoende variatie wil), leert het model adaptieve varianties. De variatie is groot genoeg om stochastische perturbaties tijdens het samplingproces op te vangen, maar klein genoeg om de reconstructie-accuraatheid te behouden. Dit is adaptief afhankelijk van de lokale gevoeligheid van de decoder.

Kernbijdragen

1. Identificatie van een fundamenteel probleem: De auteurs tonen aan dat bestaande LDMs vaak falen in de generatieve kwaliteit vanwege een gebrek aan robuustheid in de latente ruimte tegenover sampling-perturbaties, veroorzaakt door een te compacte manifold.
2. Nieuwe Loss-functie: Ze introduceren de Variance Expansion (VE) loss, een simpele maar effectieve methode om een robuuste latente ruimte te construeren zonder de noodzaak van zware KL-regularisatie.
3. Empirisch bewijs: Uitgebreide experimenten tonen aan dat deze methode de generatiekwaliteit consistent verbetert over verschillende architecturen (zoals DiT, SiT, LightningDiT) en tokenizers, terwijl de reconstructiekwaliteit competitief blijft.

Resultaten

De auteurs hebben hun methode getest op de ImageNet-dataset (256x256) met verschillende state-of-the-art baselines:

• Generatieve Kwaliteit: Het gebruik van VE-loss leidt tot een aanzienlijke verbetering in de Fréchet Inception Distance (FID). In vergelijking met de beste bestaande methoden (zoals VA-VAE en RAE), bereikt hun model een FID van 1.18 met slechts 530 trainings-epochs, wat beter is dan veel concurrenten die 800+ epochs nodig hebben.
• Reconstructie: Hoewel er een zeer lichte daling in reconstructiekwaliteit kan optreden bij strikt identieke setups, blijft de prestatie (gemeten in rFID, PSNR, SSIM) competitief en vaak zelfs beter dan de baselines in de praktijk.
• Robuustheid: De visualisaties tonen aan dat de latente ruimte met VE-loss "verspreider" is en minder snel uit de manifold raakt tijdens sampling, wat leidt tot stabielere en faithfulere generaties.

Significantie

Dit paper biedt een nieuw perspectief op het ontwerp van latente ruimten voor generatieve modellen. Het benadrukt dat reconstructie-accuraatheid niet het enige doel mag zijn; de geometrie van de latente ruimte moet ook robuust zijn tegen de stochastische aard van het diffusieproces.

• Het paper suggereert dat de traditionele KL-regularisatie in LDMs vaak contraproductief is.
• De voorgestelde VE-loss biedt een systematische, theoretisch onderbouwde oplossing voor het probleem van "over-compactheid" in tokenizers.
• De bevindingen zijn breed toepasbaar en kunnen leiden tot stabielere en hogere kwaliteit generatieve modellen in de toekomst, zowel voor afbeeldingen als video.