Normalizing Flows with Iterative Denoising

1. Problemstellung und Motivation

Normalizing Flows (NFs) sind eine etablierte Klasse von generativen Modellen, die auf der Likelihood-Maximierung basieren. Trotz ihres Potenzials für skalierbares Training und exakte Wahrscheinlichkeitsberechnung haben sie in der Bildgenerierung oft hinter Diffusionsmodellen und diskreten autoregressiven Modellen zurückgeblieben.

Das Papier identifiziert ein zentrales Hindernis, das als „Noise Dilemma" (Rausch-Dilemma) bezeichnet wird:

• Zu wenig Rauschen: Wenn das NF-Modell nur mit geringem Rauschen trainiert wird, neigt es dazu, Bilder mit übermäßig feinen Texturen und hohem Frequenzgehalt zu erzeugen, die jedoch eine schlechte globale Struktur aufweisen. Zudem ist das Training bei sehr geringem Rauschen oft instabil.
• Zu viel Rauschen: Bei hohem Rauschlevel lernt das Modell zwar eine robuste globale Struktur, aber die anschließende „Selbst-Denoising"-Schritt (basierend auf Tweedie's Lemma) führt zu stark verwischten Bildern mit sichtbaren Artefakten und Detailverlust.

Bisherige Ansätze versuchten, dies durch moderate Rauschlevel oder nachgeschaltete VAEs zu umgehen, was jedoch oft die Leistungsfähigkeit einschränkte oder die End-to-End-Natur des Modells unterbrach.

2. Methodik: Iterative TARFlow (iTARFlow)

Die Autoren schlagen iTARFlow vor, eine Erweiterung des TARFlow-Architekturkonzepts, die das Rausch-Dilemma durch eine Kombination aus multi-Noise-Training und iterativem Denoising löst.

Kernkomponenten:

• Multi-Noise Training:
  Im Gegensatz zu herkömmlichen NFs, die oft bei einem festen Rauschlevel trainieren, wird iTARFlow über ein Kontinuum von Rauschleveln $t \in [t_{min}, t_{max}]$ trainiert. Das Modell lernt eine bedingte Wahrscheinlichkeitsverteilung $\log p_\theta(x_t, t)$.

  • Der Trainingsverlust wird über alle Rauschlevel amortisiert, wobei $t$ als zusätzlicher Eingabe-Parameter (Conditioning) dient.
  • Dies ermöglicht dem Modell, sowohl globale Strukturen (bei hohem $t$) als auch feine Details (bei niedrigem $t$) zu lernen.

• Zweistufiger Sampling-Prozess:

  1. Autoregressive Generierung: Das Modell generiert initial ein stark verrauschtes Bild $x_{t_{max}}$ patchweise (autoregressiv), ähnlich wie TARFlow. In dieser Phase wird Classifier-Free Guidance (CFG) angewendet, um die Bildqualität zu steuern.
  2. Iteratives Denoising: Anstatt nur einen einzigen Denoising-Schritt durchzuführen, wird ein iterativer Prozess angewendet. Das Modell nutzt die Score-Funktion $\nabla_x \log p_\theta(x_t, t)$, die durch automatische Differentiation (Auto-Diff) direkt aus dem gelernten Log-Likelihood berechnet wird.
  • Dieser Prozess folgt einer gewöhnlichen Differentialgleichung (ODE), die dem Denoising in Diffusionsmodellen ähnelt, aber auf der Likelihood-Parametrisierung des NF basiert.
  • Der Prozess läuft parallel über das gesamte Bild (nicht mehr autoregressiv), was effizient ist.

• Architektur:
  Das Backbone bleibt ein invertierbarer, kausaler Transformer (TARFlow), der aus gestapelten Transformer-Blöcken besteht. Die Permutation der Token zwischen den Schichten sorgt für bidirektionale Abhängigkeiten.

3. Wichtige Beiträge

1. Identifikation und Charakterisierung des Noise Dilemmas: Das Papier liefert eine systematische Analyse, die zeigt, wie die Wahl des Rauschlevels den Kompromiss zwischen globaler Struktur und lokaler Detailtreue bestimmt.
2. Entwicklung von iTARFlow: Einführung eines Trainings- und Inferenzpipelines, das NFs die Fähigkeit verleiht, hochverrauschte, aber strukturell kohärente Bilder zu generieren und diese schrittweise zu verfeinern.
3. Überwindung der Leistungslücke: iTARFlow schließt die Lücke zwischen Normalizing Flows und State-of-the-Art-Modellen (Diffusion und autoregressive Modelle) signifikant.
4. Analyse von Artefakten: Die Autoren identifizieren spezifische Fehlermodi (z. B. kollabierte Hintergründe und Unschärfen im ersten generierten Token) und geben Hinweise für zukünftige Verbesserungen.

4. Ergebnisse

Die Evaluation erfolgte auf dem ImageNet-Datensatz mit Auflösungen von 64, 128 und 256 Pixeln.

• Quantitative Leistung (FID):
  • ImageNet-64: iTARFlow (770M Parameter) erreicht einen FID von 1.68, was nahe an die besten Diffusionsmodelle heranreicht und deutlich besser ist als der Vorgänger TARFlow (2.66).
  • ImageNet-128: Mit einem FID von 2.44 (1.2B Parameter) übertrifft iTARFlow TARFlow (5.03) und konkurriert mit Diffusionsmodellen.
  • ImageNet-256: iTARFlow-XL erreicht einen FID von 3.32 im Pixelraum, was eine deutliche Verbesserung gegenüber TARFlow (5.56) und STARFlow (4.69) darstellt.
• Effizienz:
  • Das Modell benötigt weniger Parameter als vergleichbare Diffusionsmodelle für ähnliche Ergebnisse.
  • Das Training ist effizienter als bei Diffusionstransformern (DiT), da iTARFlow in weniger Epochen konvergiert (600 vs. 1400).
• Kompatibilität:
  • Ein überraschendes Ergebnis ist, dass die von iTARFlow geschätzten Score-Funktionen mit denen von DiT-Modellen vergleichbar sind. iTARFlow kann nahtlos mit DiT-Denoisern kombiniert werden, was eine neue Richtung für hybride AR-Diffusions-Frameworks eröffnet.

5. Bedeutung und Limitationen

Bedeutung:
iTARFlow demonstriert, dass Normalizing Flows durch die Integration von iterativem Denoising und Multi-Noise-Training zu einer starken Alternative zu Diffusionsmodellen werden können. Es bewahrt die Vorteile von NFs (exakte Likelihood, skalierbares Training, autoregressive Struktur) und kombiniert sie mit der visuellen Qualität von Diffusionsmodellen. Die Fähigkeit, Score-Funktionen direkt aus der Likelihood abzuleiten, ohne separate Score-Netzwerke zu trainieren, ist ein wesentlicher methodischer Fortschritt.

Limitationen:

• Speicherbedarf: Die iterative Denoising-Phase erfordert automatische Differentiation über das gesamte Netzwerk, was den GPU-Speicherverbrauch erheblich erhöht.
• Artefakte: Es treten spezifische Fehler auf, wie kollabierte schwarze Hintergründe (vermutlich durch CFG-Skalierung bei großen Patches verursacht) und Unschärfen im ersten autoregressiven Token (da dieser ohne Kontext generiert wird).
• Hyperparameter-Tuning: Es gibt derzeit keine theoretische Grundlage für die optimale Wahl der Rauschgrenzen ($t_{min}, t_{max}$), was eine empirische Suche erfordert.

Fazit:
Das Paper stellt einen bedeutenden Schritt voran für die Generativmodellierung mit Normalizing Flows und zeigt, dass diese Klasse von Modellen durch iterative Verfeinerung wettbewerbsfähige Bildqualität erreichen kann, ohne die Vorteile der Likelihood-basierten Modellierung zu opfern.