Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Die Arbeit stellt einen neuen Ansatz vor, der Diffusionsmodelle durch gezielte Formung der Verteilung auf Zwischen-Rauschstufen (P-GRAFT) und eine inverse Rauschkorrektur effizienter feinabstimmt und dabei sowohl die Leistung bei Text-zu-Bild-Generierung als auch die Bildqualität bei unbedingter Generierung verbessert.

Das Wesentliche

Die große Idee: Wie man einen Künstler perfektioniert, ohne ihn zu überfordern

Stell dir vor, du hast einen genialen, aber etwas ungeduldigen Maler namens Diffusions-Modell. Dieser Maler kann wunderschöne Bilder malen, wenn man ihm eine grobe Skizze (Rauschen) gibt. Er fängt mit einem komplett verschmierten, grauen Bild an und entfernt langsam die Unschärfe, bis das Bild klar ist.

Das Problem ist: Manchmal malt er Dinge, die nicht ganz passen (z. B. ein Bild mit drei Augen bei einer Katze) oder er ignoriert die Anweisungen des Kunden. Wir wollen ihn also „feinjustieren" (fine-tuning), damit er besser wird.

Bisher gab es zwei Hauptmethoden, um ihn zu verbessern, aber beide hatten Haken:

1. Der strenge Lehrer (Policy Gradient): Der Lehrer steht ständig daneben, korrigiert jeden Pinselstrich und sagt: „Nein, das ist falsch!" Das ist sehr anstrengend, instabil und der Maler wird oft verwirrt.
2. Der Auswahler (Rejection Sampling): Der Maler malt 100 Bilder, und der Lehrer behält nur das eine beste Bild und wirft die anderen weg. Dann lernt der Maler nur von diesem einen perfekten Bild. Das funktioniert gut, ist aber sehr ineffizient, weil 99 Bilder weggeworfen werden.

Die Autoren dieses Papers haben nun einen cleveren neuen Weg gefunden, der wie ein Schulbus mit Zwischenstopps funktioniert.

---

1. Der neue Ansatz: P-GRAFT (Der Bus mit Zwischenstopps)

Stell dir den Malprozess als eine Busfahrt vor.

• Start: Der Bus ist voller Nebel (das reine Rauschen).
• Ziel: Der Bus ist am klaren, sonnigen Strand angekommen (das fertige Bild).

Das alte Problem: Wenn der Bus erst ganz am Ziel (dem fertigen Bild) ankommt, um zu entscheiden, ob die Fahrt gut war, ist es oft zu spät. Der Maler hat schon so viele Fehler gemacht, dass es schwer ist, zu verstehen, wo genau er abgekommen ist. Es ist wie wenn man am Ende eines langen Films erst merkt, dass die Handlung unsinnig war – man kann den Regisseur nicht mehr retten.

Die Lösung von P-GRAFT:
Die Autoren sagen: „Warte mal! Wir müssen den Bus nicht erst am Ziel stoppen. Wir halten ihn schon auf halber Strecke an!"

• Die Analogie: Stell dir vor, der Maler malt das Bild. Anstatt zu warten, bis das Bild komplett fertig ist, schauen wir uns das Bild an, wenn es noch zu 75 % oder 50 % unscharf ist.
• Warum das hilft: In diesem „Zwischenzustand" ist das Bild noch nicht so komplex wie am Ende. Es ist einfacher für den Maler zu lernen, wie man von diesem unscharfen Zustand in die richtige Richtung geht.
• Der Trick: Wir nehmen die Belohnung (z. B. „Das Bild ist toll!") vom fertigen Bild und geben sie dem Maler für den unscharfen Zwischenzustand. So lernt er: „Ah, wenn ich in diesem unscharfen Zustand so weitermache, lande ich am Ende bei einem tollen Bild."

Das ist wie beim Lernen eines Musikstücks: Anstatt nur am Ende zu hören, ob du den Song richtig gespielt hast, hörst du dir an, wie du die ersten Takte spielst, und gibst dir Feedback basierend darauf, wie gut der ganze Song am Ende klingen wird. Das macht das Lernen stabiler und schneller.

Das Ergebnis: Der Maler (das Modell) lernt schneller, macht weniger Fehler und produziert bessere Bilder als bei den alten Methoden.

---

2. Der zweite Trick: Inverse Noise Correction (Der Rückwärts-Generator)

Das Paper stellt noch eine zweite Methode vor, die besonders für „Flow-Modelle" (eine spezielle Art von Maler, die Bilder wie einen Fluss fließen lässt) gedacht ist.

Das Problem: Manchmal ist der Maler zwar gut, aber er fängt die Reise immer mit dem falschen Rauschen an. Es ist, als würde er versuchen, ein Haus zu bauen, aber er beginnt mit dem falschen Fundament. Das Haus wird schief, egal wie gut er später baut.

Die Lösung:
Statt den Maler neu zu trainieren, fragen wir uns: „Welches Rauschen müsste er eigentlich haben, damit er am Ende ein perfektes Bild malt?"

• Die Analogie: Stell dir vor, du hast einen perfekten Kuchen, aber du weißt nicht, welche Zutaten genau reinkamen. Du nimmst den fertigen Kuchen und „entschmälst" ihn rückwärts, um herauszufinden, wie die Zutatenmischung aussehen müsste.
• Der Schritt: Wir nehmen die perfekten Bilder, die der Maler schon kann, und laufen den Prozess rückwärts ab. Wir sehen, zu welchem „Rauschen" diese perfekten Bilder führen.
• Der neue Helfer: Wir trainieren einen kleinen, schlauen Assistenten (den „Noise Corrector"). Dieser Assistent lernt, das richtige Rauschen zu erzeugen, das perfekt zu unserem Maler passt.
• Der Gewinn: Wenn wir dann ein neues Bild malen wollen, gibt der Assistent dem Maler das perfekte Start-Rauschen. Der Maler muss dann weniger arbeiten, macht weniger Fehler und braucht weniger Rechenleistung (weniger Strom/Zeit), um ein tolles Bild zu produzieren.

---

Zusammenfassung für den Alltag

Stell dir vor, du willst ein neues Auto lernen zu fahren.

1. Alte Methode: Du fährst los, machst einen riesigen Fehler, und dein Lehrer schreit dich am Ende der Fahrt an. Du weißt nicht, wann genau du abgekommen bist. Oder: Du fährst 100 Runden und dein Lehrer sagt: „Nur die eine Runde war okay, die anderen 99 waren Müll. Versuche es nochmal." (Sehr ineffizient).
2. Die neue Methode (P-GRAFT): Dein Lehrer hält dich schon nach 100 Metern an. Er sagt: „Schau, dein Lenkrad war schon hier etwas schief. Wenn du das korrigierst, kommst du am Ziel perfekt an." Das ist viel einfacher zu lernen.
3. Der Assistent (Inverse Noise): Dein Lehrer gibt dir nicht einfach irgendeinen Schlüssel, sondern einen speziell angepassten Schlüssel, der genau zu deinem Auto passt. Du musst weniger Kraft aufwenden, um zu starten.

Das Fazit:
Die Autoren haben gezeigt, dass man künstliche Intelligenzen nicht nur am Ende des Prozesses bewerten muss, sondern auch in der Mitte. Wenn man die „Zwischenstationen" klug nutzt, werden die KI-Modelle besser, schneller und brauchen weniger Rechenleistung. Das ist ein großer Schritt für die Zukunft der Bildgenerierung, von schönen Kunstwerken bis hin zu chemischen Molekülen für Medikamente.

Technische Zusammenfassung

1. Problemstellung

Diffusionsmodelle sind der aktuelle Standard für generative Aufgaben, erfordern jedoch oft eine Feinabstimmung (Fine-Tuning), um Fehler im Lernprozess zu korrigieren oder sich an spezifische Downstream-Anwendungen (z. B. Text-zu-Bild-Erstellung) anzupassen.

• Herausforderung bei RL-Ansätzen: Herkömmliche Reinforcement-Learning-Methoden wie Proximal Policy Optimization (PPO) mit KL-Regularisierung sind für Diffusionsmodelle schwierig anzuwenden. Im Gegensatz zu autoregressiven Modellen ist die marginale Likelihood für Diffusionsmodelle nicht berechenbar (intractable). Das Ignorieren des KL-Terms führt zu instabilem Training, während Trajektorien-KL-Regularisierung suboptimale Ergebnisse liefert und zu Verzerrungen der Anfangswertfunktion führt.
• Herausforderung bei Rejection Sampling: Bestehende Methoden wie RAFT (Rejection sAmpling based Fine-Tuning) nutzen Ablehnungsstichproben, um die Verteilung zu formen, tun dies jedoch typischerweise nur am Ende des Denoising-Prozesses (am finalen Datenpunkt). Dies ignoriert die strukturellen Vorteile, die durch die Anpassung der Verteilung in intermediären Rauschstufen entstehen könnten.
• Flow-Modelle: Bei Flow-basierten Modellen (die deterministische ODEs nutzen) gibt es oft Fehler in der gelernten Verteilung, die sich nicht durch einfache Reward-Fine-Tuning-Methoden ohne explizite Belohnungsfunktionen korrigieren lassen.

2. Methodik

Das Papier entwickelt einen theoretischen Rahmen, der Ablehnungsstichproben (Rejection Sampling) mit der KL-regulierten Reward-Maximierung verbindet, und leitet daraus zwei Hauptmethoden ab:

A. GRAFT (Generalized Rejection sAmpling Fine-Tuning)

Die Autoren führen ein einheitliches Framework namens GRAFT ein, das verschiedene Ablehnungsstichproben-Strategien (wie Top-K, Best-of-N, RSO) unter einem gemeinsamen Dach vereint.

• Theoretische Grundlage: Sie zeigen mathematisch, dass GRAFT implizit eine KL-regulierte Reward-Maximierung mit umgeformten Belohnungen (reshaped rewards) durchführt, auch wenn die marginale Likelihood nicht berechnet werden kann.
• Mechanismus: Anstatt alle generierten Samples zu nutzen, werden nur die hochbewerteten Samples (basierend auf einer Reward-Funktion) für das Fine-Tuning des Referenzmodells verwendet. Dies entspricht dem Ziehen aus einer „gekippten" (tilted) Verteilung.

B. P-GRAFT (Partial GRAFT) – Die Kerninnovation

Die Hauptinnovation besteht darin, die Verteilung nicht erst am Ende des Denoising-Prozesses ($t=0$), sondern in einem intermediären Rauschzustand ($t > 0$) zu formen.

• Konzept: Das Modell wird nur bis zu einem intermediären Zeitpunkt $t_{NI}$ feinabgestimmt. Die Belohnungen werden basierend auf den vollständig denoisierten Bildern ($X_0$) berechnet, aber das Training erfolgt auf den teilweise denoisierten Latents ($X_{t_{NI}}$).
• Bias-Varianz-Abwägung:
  • Varianz: Die Varianz der Belohnung, bedingt auf den Zustand $X_t$, nimmt mit größerem $t$ (näher am Rauschen) zu. Das macht die Belohnungssignale „verrauschter".
  • Bias (Lernschwierigkeit): Der Score-Funktion (die gelernt werden muss) ist bei späteren Zeitpunkten (näher am Rauschen) einfacher zu lernen, da sie sich der einfachen Gaußschen Score-Funktion annähert.
  • Optimierung: P-GRAFT nutzt diesen Trade-off, indem es einen optimalen Zeitpunkt $t_{NI}$ wählt, an dem das Lernen des Scores einfacher ist, ohne die Varianz der Belohnung zu stark zu erhöhen.

C. Inverse Noise Correction (für Flow-Modelle)

Für Flow-Modelle (die auf ODEs basieren) schlagen die Autoren eine Methode vor, um Fehler in vortrainierten Modellen zu korrigieren, ohne explizite Rewards.

• Idee: Da Flow-Modelle reversibel sind, kann man von einem generierten Bild zurück zum ursprünglichen Rauschen „invertieren".
• Verfahren:
  1. Man nimmt Daten aus der Zielverteilung (z. B. echte Bilder).
  2. Man führt die inverse ODE (Rückwärts-Euler) durch das vortrainierte Modell durch, um eine „inverse Rauschverteilung" ($p^{rev}_1$) zu erhalten.
  3. Ein neues, kleines Adapter-Modell (Noise Corrector) wird trainiert, um von Standard-Gauß-Rauschen ($N(0, I)$) zu dieser inversen Rauschverteilung zu generieren.
  4. Bei der Inferenz wird zuerst das Noise Corrector-Modell genutzt, um korrigiertes Rauschen zu erzeugen, das dann vom ursprünglichen vortrainierten Flow-Modell in ein Bild umgewandelt wird.
• Theoretischer Vorteil: Dies korrigiert systematische Fehler des vortrainierten Modells, indem es die Verteilungsverschiebung am Startpunkt ($t=0$) ausgleicht.

3. Wichtige Beiträge

1. Einheitliches Framework (GRAFT): Beweis, dass Ablehnungsstichproben für Diffusionsmodelle äquivalent zu einer KL-regulierten Reward-Maximierung mit umgeformten Belohnungen sind, was das Problem der nicht berechenbaren marginalen Likelihood umgeht.
2. P-GRAFT: Ein neuartiges Fine-Tuning-Verfahren, das den Denoising-Prozess in zwei Phasen aufteilt (Feinabstimmung bis $t_{NI}$, dann Referenzmodell). Dies führt empirisch zu besseren Ergebnissen als das Fine-Tuning über den gesamten Prozess.
3. Theoretische Begründung: Eine mathematische Herleitung des Bias-Varianz-Trade-offs, die erklärt, warum das Fine-Tuning in intermediären Stufen effektiver ist (einfacherer Score als bei $t=0$).
4. Inverse Noise Correction: Ein parametereffizienter Ansatz zur Verbesserung von Flow-Modellen durch Korrektur der initialen Rauschverteilung, der auch ohne Reward-Funktion funktioniert.

4. Ergebnisse

Die Methoden wurden auf Text-zu-Bild (T2I), Layout-Generierung, Molekülgenerierung und unbedingte Bildgenerierung evaluiert.

• Text-zu-Bild (Stable Diffusion v2):
  • P-GRAFT übertrifft sowohl die Basisversion (SDv2) als auch etablierte Policy-Gradient-Methoden wie DDPO und SDXL-Base signifikant.
  • Auf dem GenAI-Bench und T2ICompBench++ erreichte P-GRAFT (mit $t_{NI} = 0.25N$) die besten VQAScores (z. B. 71.94 vs. 66.87 bei SDv2 auf GenAI-Bench).
  • DDPO zeigte selbst bei höherem Rechenaufwand keine signifikanten Verbesserungen gegenüber der Basis.
• Layout- und Molekülgenerierung:
  • P-GRAFT verbesserte die Ausrichtung (Alignment) und Stabilität von Molekülen.
  • Besonders bemerkenswert: P-GRAFT erreichte mit nur 1x der Stichprobenanzahl (im Vergleich zu 9x bei GRAFT) die besten Ergebnisse bei der Molekülgenerierung, was die Effizienz des Ansatzes unterstreicht.
• Inverse Noise Correction:
  • Bei der unbedingten Bildgenerierung (CelebA-HQ, LSUN-Church) verbesserte die Methode den FID (Fréchet Inception Distance) signifikant.
  • Ein entscheidender Vorteil: Die Methode erreichte bessere FID-Werte bei geringerem Rechenbedarf (FLOPs) pro Bild. Ein Noise Corrector mit 100 Schritten + vortrainiertes Modell mit 100 Schritten war besser als das vortrainierte Modell allein mit 1000 Schritten.

5. Bedeutung und Fazit

Dieses Papier bietet einen Paradigmenwechsel im Fine-Tuning von Diffusions- und Flow-Modellen:

• Effizienz: Es zeigt, dass das Fine-Tuning nicht den gesamten Denoising-Pfad abdecken muss. Durch das Ausnutzen der einfacheren Lernbarkeit in früheren Stufen (höheres Rauschen) lassen sich bessere Ergebnisse mit weniger Rechenaufwand erzielen.
• Stabilität: Die Methode umgeht die Instabilitäten von PPO/DDPO, die durch die Schwierigkeit der KL-Regularisierung bei Diffusionsmodellen entstehen.
• Generalisierung: Der Ansatz funktioniert nicht nur für T2I, sondern auch für diskret-kontinuierliche Modelle (Layouts, Moleküle) und Flow-basierte Modelle.
• Praktische Relevanz: Die „Inverse Noise Correction" bietet einen Weg, vortrainierte Modelle zu verbessern, ohne teure Reward-Modelle oder komplexe RL-Setups zu benötigen, was besonders für ressourcenbeschränkte Szenarien wertvoll ist.

Zusammenfassend etabliert das Paper P-GRAFT als einen überlegenen Ansatz gegenüber bestehenden Policy-Gradient-Methoden und liefert eine theoretisch fundierte Erklärung für die Effektivität des Fine-Tuning in intermediären Zuständen.