Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

Titel: Post-Hoc Guidance für Konsistenzmodelle durch gemeinsames Lernen von Flussverteilungen (Joint Flow Distribution Learning)

Autoren: Chia-Hong Hsu und Randall Balestriero (Brown University)

1. Problemstellung

Konsistenzmodelle (Consistency Models, CMs) sind eine neue Klasse generativer Modelle, die Bilder in nur einem oder wenigen Schritten erzeugen können, was sie deutlich effizienter als herkömmliche Diffusionsmodelle (Diffusion Models, DMs) macht. Ein zentrales Problem bei der Anwendung von CMs ist jedoch die Steuerung der Generierung (Guidance).

• Herausforderung bei Diffusionsmodellen: Bei DMs ist die Classifier-Free Guidance (CFG) ein Standardverfahren, das es ermöglicht, einen Kompromiss zwischen der Treue (Fidelity) zu einem Prompt und der Vielfalt (Diversity) der generierten Bilder einzustellen. Dies geschieht durch eine lineare Interpolation zwischen bedingten und unbedingten Vorhersagen während der Inferenz, ohne das Modell neu trainieren zu müssen.
• Herausforderung bei Konsistenzmodellen: Bisherige Methoden zur Guidance bei CMs erfordern eine Wissensdistillation (Knowledge Distillation) von einem separaten, trainierten DM als „Lehrer". Dies bindet die Guidance-Fähigkeit an den Ansatz der Consistency Distillation (CD).
• Die Lücke: Es fehlte bisher eine Methode, um bereits trainierte CMs, die durch Consistency Training (CT) entstanden sind (also rein datengetrieben ohne DM-Lehrer), nachträglich (post-hoc) mit Guidance-Fähigkeiten auszustatten. Ohne diese Fähigkeit können CT-Modelle keine Guidance-Skalen nutzen, um die Bildqualität zu steuern.

2. Methodik: Joint Flow Distribution Learning (JFDL)

Die Autoren schlagen JFDL vor, eine leichte Nachtrainierungsmethode, die es ermöglicht, Guidance in einem vortrainierten CM zu aktivieren, ohne auf einen DM-Lehrer zurückzugreifen.

Kernkonzept:
Die Methode betrachtet den vortrainierten CM als Löser einer gewöhnlichen Differentialgleichung (ODE) für den Diffusionspfad. Sie nutzt die Verbindung zwischen Diffusionsmodellen und Flow-based Models (FMs).

Der Algorithmus (Naive JFDL):

1. Annahme: Ein vortrainierter CM kann sowohl bedingte (class-conditioned) als auch unbedingte (unconditional, $\emptyset$) ODE-Pfade lösen.
2. Hybrider Fluss: Für ein gegebenes Bild $x_0$ einer Klasse $c$ wird ein „hybrider" Lösungsweg konstruiert:
   • Zuerst wird das Bild vom Datenraum zum Rauschraum transportiert (bedingter Pfad).
   • Anschließend wird es vom Rauschraum zurück in den Datenraum transportiert, aber unter Verwendung der unbedingten ODE ($\emptyset$).
   • Das Ergebnis ist ein Punkt $y_{0}^{\emptyset, t}$, der als „Anker" für die unbedingte Verteilung dient.
3. Richtungsvektoren: Es werden zwei Denoising-Richtungen berechnet:
   • $u_{cls}$: Die Richtung zum ursprünglichen bedingten Bild.
   • $u_{\emptyset}$: Die Richtung zum hybriden unbedingten Anker.
4. Extrapolation: Diese Richtungen werden mit einem Guidance-Faktor $\omega$ interpoliert: $u_{\omega} = \omega \cdot u_{cls} + (1-\omega) \cdot u_{\emptyset}$.
5. Verlustfunktion: Das Modell wird so nachtrainiert, dass es konsistent mit dieser extrapolierten Richtung ist. Dies geschieht durch eine Kombination aus dem ursprünglichen CT-Verlust ($L_{ECT}$) und dem neuen JFDL-Verlust ($L_{JFDL}$), gewichtet durch GradNorm.

Theoretische Fundierung & Pseudo-Rauschen:
Ein kritischer theoretischer Schritt ist die Analyse des „Pseudo-Rauschens". Die Autoren beweisen (Proposition 1 & 2), dass die Verteilung des hybriden Ankers $y_{0}^{\emptyset, t}$ über alle Klassen hinweg die marginale Datenverteilung $p_0(x)$ erhält. Zudem zeigen sie durch Normalitätstests (Shapiro-Wilk, Anderson-Darling), dass das verbleibende Signal im Vergleich zum dominanten Rauschterm ($\sigma_{max} z$) vernachlässigbar ist. Das bedeutet: Der unbedingte Anker verhält sich effektiv wie reines Gaußsches Rauschen, was die Gültigkeit der Guidance-Interpolation mathematisch untermauert.

Random JFDL (Praktische Erweiterung):
Da viele vortrainierte CMs nicht explizit für das Lösen der unbedingten ODE ($\emptyset$) trainiert wurden, schlagen die Autoren eine Variante vor, bei der der $\emptyset$-Label durch eine zufällige Klasse $c'$ aus dem Datensatz ersetzt wird. Dies funktioniert überraschend gut und erzeugt Bilder mit höherem Kontrast, die noch näher an der klassischen CFG liegen.

3. Wichtige Beiträge

1. Erste post-hoc Guidance für CMs ohne DM-Lehrer: JFDL ist die erste Methode, die es ermöglicht, einem rein durch CT trainierten CM nachträglich Guidance-Fähigkeiten zu verleihen, ohne ein separates Diffusionsmodell zu benötigen.
2. Theoretische Verbindung zu Flow-Modellen: Die Arbeit liefert theoretische Beweise und experimentelle Verifikationen (Normalitätstests), warum die Interpolation von bedingten und unbedingten Flussrichtungen in CMs funktioniert und Gaußsches Rauschen impliziert.
3. Flexibilität: Die Methode funktioniert sowohl mit Modellen, die explizit unbedingte Pfade gelernt haben (Naive JFDL), als auch mit solchen, die dies nicht getan haben (Random JFDL).
4. Effizienz: Das Nachtrainieren ist extrem ressourcenschonend (ca. 7,5 % der Trainingsdaten und 30 % der GPU-Zeit des ursprünglichen Trainings).

4. Ergebnisse

Die Methode wurde auf den Datensätzen CIFAR-10 und ImageNet 64x64 evaluiert.

• Quantitative Ergebnisse (FID):
  • Auf CIFAR-10 verbesserte sich der 1-Schritt-FID von 3,40 (Baseline ECT) auf 3,24 (Random JFDL) bzw. 3,29 (Naive JFDL).
  • Auf ImageNet 64x64 sank der 1-Schritt-FID von 5,84 auf 4,68 (Random) bzw. 4,38 (Naive).
  • Die Kurvenverläufe der FID-Werte in Abhängigkeit von der Guidance-Stärke $\omega$ ähneln stark dem Verhalten von CFG in Diffusionsmodellen (Verbesserung bei kleinem $\omega$, Verschlechterung bei zu großem $\omega$).
• Qualitative Ergebnisse:
  • Mit steigendem $\omega$ verbessert sich die visuelle Schärfe und der Kontrast der Bilder (z. B. klarere Details bei „Hotdog" oder „Vogel" auf ImageNet).
  • Wie bei CFG nimmt die Vielfalt (Diversity) der generierten Bilder bei hohen $\omega$-Werten ab.
• Vergleich: Die Ergebnisse zeigen, dass JFDL die Lücke schließt, die bisher bestand, da CT-Modelle nun Guidance nutzen können, ohne dass ein DM-Lehrer vorhanden sein muss.

5. Bedeutung und Ausblick

• Brückenschlag: JFDL verbindet die Welt der schnellen Konsistenzmodelle mit der etablierten Technik der Classifier-Free Guidance. Dies macht CMs für Anwendungen attraktiver, bei denen eine feine Kontrolle der Generierung (z. B. Text-zu-Bild) erforderlich ist.
• Ressourceneffizienz: Da keine teure Wissensdistillation von einem DM-Lehrer mehr nötig ist, um Guidance zu ermöglichen, wird der Zugang zu hochwertigen, steuerbaren Generationsmodellen demokratisiert.
• Zukunft: Die Autoren sehen Potenzial in der Entwicklung adaptiver Guidance-Schemata für Mehrschritt-Verfahren (Multi-step Sampling), da die aktuelle Theorie für diesen Fall noch nicht vollständig geklärt ist.

Fazit: Das Paper stellt einen signifikanten Fortschritt dar, indem es die Kontrolle über Konsistenzmodelle revolutioniert und diese von der Abhängigkeit gegenüber Diffusionsmodellen befreit, während es gleichzeitig die Effizienz und die Bildqualität (FID) verbessert.