Variational Autoencoding Discrete Diffusion with Enhanced Dimensional Correlations Modeling

1. Het Probleem

Discrete diffusiemodellen, en met name Masked Diffusion Models (MDM), hebben recentelijk veelbelovende resultaten geleverd voor het genereren van discrete data (zoals tekst en pixel-gebaseerde beelden). MDM's werken door data progressief te maskeren en vervolgens meerdere dimensies gelijktijdig te "ontmaskeren" (denoisen) in de omgekeerde richting. Dit biedt een significant voordeel in snelheid ten opzichte van autoregressieve modellen, die dimensies één voor één voorspellen.

Echter, een fundamenteel beperking van bestaande MDM's is hun onafhankelijkheidsaanneming. Bij elke denoising-stap wordt de verdeling over de dimensies gemodelleerd als een product van onafhankelijke categorische verdelingen. Dit betekent dat het model complexe inter-dimensionele afhankelijkheden (correlaties tussen verschillende tokens of pixels) niet effectief kan vangen.

• Gevolg: De kwaliteit van de gegenereerde samples degradeert aanzienlijk wanneer er weinig denoising-stappen worden gebruikt. Omdat veel dimensies dan gelijktijdig moeten worden voorspeld, leiden de onafhankelijkheidsaannames tot cumulatieve benaderingsfouten en een gebrek aan coherentie in het resultaat.

2. Methodologie: VADD

De auteurs stellen Variational Autoencoding Discrete Diffusion (VADD) voor, een nieuw raamwerk dat de beperkingen van MDM's oplost door latent variabele modellering te integreren.

Kernconcepten:

• Latente Structuur: In plaats van de overgangsverdeling $p_\theta(x_s|x_t)$ direct te parametriseren als een product van onafhankelijke verdelingen, definieert VADD deze als een latent variabel model:
  $$p_\theta(x_s|x_t) = \int p_\theta(x_s|x_t, z)p(z)dz$$
  Hierbij is $z$ een latent variabele die hoge-level semantiek of correlaties tussen dimensies controleert. Door $z$ te integreren, kan de marginale verdeling complexe afhankelijkheden modelleren, zelfs als de conditionele verdeling $p_\theta(x_s|x_t, z)$ zelf eenvoudiger is.

• Variational Autoencoding Raamwerk: Omdat de marginale verdeling intractable is (moeilijk te berekenen) door de integratie over $z$, gebruiken de auteurs een Variational Autoencoder (VAE) benadering.

  • Ze introduceren een herkenningsmodel (recognition model) $r_\phi(z|x_0, x_t)$ om de posterior $p_\theta(z|x_0, x_t)$ te benaderen.
  • Het doel is het maximaliseren van een Double Evidence Lower Bound (DELBO). Dit is een ondergrens voor de log-likelihood, die zowel het denoising model ($\theta$) als het herkenningsmodel ($\phi$) optimaliseert.

• KL Annealing: Om het veelvoorkomende probleem van posterior collapse (waarbij het herkenningsmodel de data-afhankelijkheid negeert en de prior aanneemt) te voorkomen, gebruiken ze een KL-annealing strategie. De KL-divergentiestraal in de loss-functie wordt geleidelijk opgevoerd van 0 naar 1 tijdens het trainingproces.

• Architectuur voor Tekst: Voor tekstgeneratie passen de auteurs een aangepaste Transformer-architectuur toe:

  • Denoising Model: Gebruikt AdaLN (Adaptive Layer Normalization) lagen, gestuurd door de latent variabele $z$, om de token-embeddings dynamisch aan te passen.
  • Herkenningsmodel: Gebruikt een Siamese structuur die zowel de originele data $x_0$ als het gemaskerde $x_t$ verwerkt, maar slimme optimalisaties toepast (bijv. het gebruik van een binair masker $M_t$) om de rekentijd te beperken.

3. Belangrijkste Bijdragen

1. Nieuw Raamwerk: VADD is het eerste model dat latent variabele modellering en variational autoencoding succesvol toepast op discrete diffusiemodellen (MDM) om dimensie-afhankelijkheden te vangen.
2. Verbeterde Kwaliteit bij Weinig Stappen: Het model lost het probleem op dat MDM's vaak slechte kwaliteit leveren bij snelle generatie (weinig stappen) door de latent variabele $z$ te gebruiken om de gezamenlijke verdeling van meerdere dimensies te modelleren.
3. Efficiënte Architectuur: De auteurs tonen aan dat de toevoeging van het herkenningsmodel tijdens training de inferentie-efficiëntie niet beïnvloedt, omdat het herkenningsmodel niet nodig is tijdens het generatieproces (sampling). De samplingkosten blijven vergelijkbaar met standaard MDM's.
4. Uitgebreide Validatie: Het paper biedt uitgebreide experimenten op 2D toy-data, pixel-gebaseerde beeldgeneratie (MNIST, CIFAR-10) en tekstgeneratie (LM1B, OpenWebText).

4. Resultaten

De experimenten tonen aan dat VADD consistent beter presteert dan bestaande MDM-baselines (zoals MDLM):

• 2D Toy Data: VADD kan complexe multimodale verdelingen (zoals checkerboard en swissroll) correct modelleren in één stap, terwijl MDM's hierin falen en samenvallen tot een enkel punt.
• Beeldgeneratie (CIFAR-10):
  • VADD bereikt aanzienlijk lagere FID-scores (Fréchet Inception Distance) dan MDM bij een klein aantal stappen (bijv. T=10 of T=20).
  • Bij T=10 stappen haalt VADD een FID van 170.3, terwijl MDM 334.3 haalt.
  • De bits-per-dimension (BPD) verbetert ook, wat aangeeft dat de latent variabele structuur de generatieve capaciteit vergroot.
• Tekstgeneratie:
  • Op de LM1B dataset bereikt VADD een lagere perplexiteit dan zowel autoregressieve modellen als andere diffusiemodellen, zelfs met minder training-iteraties.
  • Op OpenWebText toont VADD een significante verbetering in generative perplexity bij weinig sampling-stappen. Bij 16 stappen is de perplexiteit van VADD (194) veel lager dan die van MDLM (330), wat betekent dat VADD met minder dan 50% van de rekentijd dezelfde kwaliteit bereikt als MDLM met meer stappen.

5. Betekenis en Conclusie

VADD vertegenwoordigt een belangrijke stap voorwaarts in het veld van discrete generatieve modellen. Het bewijst dat het integreren van latent variabele structuren in diffusiemodellen de inherente beperkingen van onafhankelijke dimensie-aannames kan overwinnen.

• Praktische Impact: Het maakt snelle, hoogwaardige generatie mogelijk (cruciaal voor real-time toepassingen) zonder de noodzaak van dure distillatie of iteratieve correcties.
• Toekomstperspectief: Het paper suggereert dat complexere prior-verdelingen (in plaats van een simpele Gaussische prior) en geavanceerde trainingsmethoden verdere verbeteringen kunnen opleveren.

Kortom, VADD combineert de snelheid van parallelle generatie van MDM's met de expressieve kracht van latent variabele modellen, wat resulteert in superieure samplekwaliteit, vooral in scenario's waar inferentiesnelheid prioriteit heeft.