Normalizing Flows with Iterative Denoising

Probleemstelling

De auteurs identificeren twee dominante paradigma's in beeldgeneratie en hun respectievelijke beperkingen:

1. Diffusiemodellen: Bereiken state-of-the-art visuele kwaliteit, maar vereisen tientallen tot honderden iteratieve stappen tijdens het genereren, wat leidt tot hoge rekenkosten. Bovendien is hun schaalbaarheid minder voorspelbaar dan die van autoregressieve modellen.
2. Autoregressieve modellen (discreet): Bieden snelle steekproeven en goede schaalbaarheid, maar vereisen kwantisatie van pixels naar tokens. Dit introduceert een informatiebottleneck die de fideliteit kan verminderen.
3. Normalizing Flows (NF): Zijn een klassieke likelihood-gebaseerde methode die exacte log-likelihood-training toestaat. Echter, eerdere NF-varianten (zoals TARFlow) worstelden met een "ruis-dilemma":
   • Bij lage ruisniveaus tijdens training: Het model genereert beelden met rijke lokale texturen, maar vaak met een gebrekkige globale structuur.
   • Bij hoge ruisniveaus: Het model leert betere globale structuren, maar de resulterende beelden zijn onscherp en vervormd door het "Tweedie's lemma" effect (posterieur gemiddelde wordt gladgestreken).

Het centrale probleem is dus het vinden van een generatief model dat de schaalbaarheid van autoregressieve architecturen combineert met de expressiviteit van continue pixelrepresentaties, zonder vast te lopen in het ruis-dilemma.

Methodologie: iTARFlow

De auteurs introduceren iTARFlow (iterative TARFlow), een uitbreiding van de bestaande TARFlow-architectuur. De kernideeën zijn:

1. Training over een continuüm van ruisniveaus:
   In plaats van te trainen op één vast ruisniveau, wordt het model getraind om de log-likelihood $\log p_\theta(x_t, t)$ te maximaliseren voor een continuüm van ruisniveaus $t \in [t_{min}, t_{max}]$. Hierbij wordt $x_t = x + t\epsilon$ gebruikt, waarbij $\epsilon$ Gaussische ruis is.

   • Dit stelt het model in staat om zowel stabiele optimalisatie (bij hogere ruis) als hoge fideliteit (bij lagere ruis) te leren.
   • De loss-functie omvat een gewichtsterm $\gamma_t$ (in dit werk gekozen als $t$) voor stabiliteit.

2. Iteratieve Denoising tijdens Sampling:
   Het generatieproces verloopt in twee fasen:

   • Fase 1 (Autoregressief): Het model genereert een ruw, hoog-ruisig beeld $x_{t_{max}}$ patch-voor-patch, gebruikmakend van de causale Transformer-architectuur. Op dit punt bevat het beeld een goede globale structuur maar is het onscherp.
   • Fase 2 (Parallelle Denoising): Het ruwe beeld wordt verfijnd via een iteratief denoising-proces. In tegenstelling tot diffusiemodellen die een apart score-netwerk nodig hebben, gebruikt iTARFlow de automatische differentiatie van de gelijktijdig getrainde likelihood-functie om de score $\nabla_x \log p_\theta(x_t, t)$ direct te berekenen.
   • Dit proces volgt een probabilistische ODE (vergelijkbaar met DDIM/Flow Matching) en wordt volledig parallel uitgevoerd, wat het computatie-efficiënt maakt.

3. Architectuur:
   Het model gebruikt een stack van causale Transformer-blokken (TARFlow). Tijdens inferentie wordt Classifier-Free Guidance (CFG) toegepast op de eerste Transformer-laag voor de autoregressieve stap, maar niet tijdens de parallelle denoising-stap.

Belangrijkste Bijdragen

• Identificatie van het Ruis-Dilemma: De auteurs karakteriseren systematisch het compromis tussen lokale textuur en globale structuur in NF-training en tonen aan dat hogere ruisniveaus paradoxalerwijs helpen bij het leren van globale coherentie.
• Iteratieve TARFlow (iTARFlow): Een nieuw trainings- en inferentieprotocol dat NF's in staat stelt om iteratief te denoisen, vergelijkbaar met diffusiemodellen, maar binnen een puur likelihood-gebaseerd raamwerk.
• Unificatie van Architecturen: iTARFlow toont aan dat NF's naadloos kunnen worden gecombineerd met bestaande DiT (Diffusion Transformer) denoisers, wat een weg vrijmaakt voor hybride autoregressieve-diffusie frameworks.

Resultaten

De prestaties zijn geëvalueerd op ImageNet bij resoluties van 64, 128 en 256 pixels (gemeten via FID - Fréchet Inception Distance):

• ImageNet-64: iTARFlow-B (770M parameters) bereikt een FID van 1.68, wat aanzienlijk beter is dan eerdere TARFlow-varianten (2.66) en vergelijkbaar met de beste diffusiemodellen (bijv. EDM-SDE: 1.55), maar met minder parameters.
• ImageNet-128: iTARFlow-L bereikt een FID van 2.44, wat een sterke verbetering is ten opzichte van TARFlow (5.03) en concurrerend is met diffusiemodellen.
• ImageNet-256: iTARFlow-XL bereikt een FID van 3.32. Hoewel dit nog iets achterblijft bij de absolute state-of-the-art diffusiemodellen (bijv. SiD2: 1.38), is het een enorme stap voorwaarts voor NF's en sluit het de kloof met andere generatieve modellen aanzienlijk.
• Efficiëntie: iTARFlow convergeert sneller (600 epochs) dan vergelijkbare DiT-modellen (1400 epochs) en vereist minder iteratieve stappen tijdens sampling (5-10 stappen) voor hoge kwaliteit.

Significantie en Beperkingen

• Significantie: iTARFlow bewijst dat Normalizing Flows een krachtig alternatief kunnen zijn voor diffusiemodellen, met name vanwege hun snelle sampling en goed begrepen schaalwetten. Het paper toont aan dat het combineren van autoregressieve structuur met iteratieve denoising de beperkingen van beide benaderingen kan overwinnen.
• Beperkingen:
  • Artefacten: Er worden twee specifieke fouten geobserveerd: "collapsed backgrounds" (zwarte achtergronden) en "low-frequency blur" (onscherpte), vooral bij hoge resoluties. De auteurs vermoeden dat dit samenhangt met de eerste autoregressieve token die zonder context wordt gegenereerd en de gebruikte CFG-schaal.
  • Rekengeheugen: Het gebruik van automatische differentiatie tijdens de iteratieve denoising-stap vereist veel GPU-geheugen, omdat de volledige backpropagation door het netwerk moet worden uitgevoerd.
  • Hyperparameters: Er is nog geen theoretische basis voor het kiezen van de optimale maximale ruis ($t_{max}$), wat empirisch tuning vereist.

Concluderend biedt iTARFlow een veelbelovende richting voor de toekomst van generatieve modellering, waarbij de voordelen van likelihood-gebaseerde training en autoregressieve schaalbaarheid worden gecombineerd met de visuele kwaliteit van diffusie-achtige denoising.