Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

Die Arbeit stellt SQDF vor, eine neuartige KL-regulierte Reinforcement-Learning-Methode zur Feinabstimmung von Diffusionsmodellen, die durch einen reparametrisierten Policy-Gradienten der Soft-Q-Funktion sowie weitere Innovationen wie einen Diskontfaktor und einen Replay-Puffer das Problem der Reward-Überoptimierung löst und gleichzeitig hohe Zielbelohnungen mit natürlicher Vielfalt in der Bildgenerierung vereint.

Das Wesentliche

Das große Problem: Der "perfekte" aber langweilige Künstler

Stell dir vor, du hast einen genialen KI-Künstler (ein Diffusionsmodell), der wunderschöne Bilder malen kann. Er ist sehr gut darin, Dinge zu erzeugen, die den Menschen gefallen (hohe "Belohnung").

Das Problem ist: Wenn du ihm sagst: "Mache das Bild noch schöner!", wird er oft verrückt. Er beginnt, nur noch grelle Farben zu verwenden, wiederholt sich ständig oder malt Dinge, die zwar technisch perfekt sind, aber keinen Sinn mehr ergeben (wie ein Bild, das nur aus roten Punkten besteht).

In der Fachsprache nennen wir das Reward Over-Optimization (Überoptimierung der Belohnung). Der Künstler hat die Regeln des Spiels so perfekt gelernt, dass er das Spiel kaputt macht. Er verliert seine Vielfalt und wird unnatürlich.

Die Lösung: SQDF – Der kluge Coach

Die Autoren dieses Papiers haben eine neue Methode namens SQDF entwickelt. Stell dir SQDF nicht als strengen Chef vor, der nur auf die Punktzahl schaut, sondern als weisen Coach, der dem Künstler hilft, besser zu werden, ohne ihn zu verrückt zu machen.

Hier sind die drei Geheimwaffen dieses Coaches:

1. Der "Zeit-Verstärker" (Discount Factor)

Stell dir den Malprozess wie das Entfernen von Nebel von einem Bild vor.

• Frühe Schritte: Am Anfang ist das Bild noch voller Nebel. Ein kleiner Strich hier oder da hat kaum Einfluss auf das Endergebnis.
• Späte Schritte: Ganz am Ende, wenn der Nebel fast weg ist, zählt jeder Pinselstrich extrem viel.

Frühere Methoden haben jeden Strich gleich stark bewertet. SQDF sagt aber: "Hey, die frühen Striche sind nicht so wichtig wie die letzten!" Es gewichtet die letzten Schritte höher. Das verhindert, dass der Künstler in den frühen Phasen wild herumrattert und sich in falsche Richtungen verirrt.

2. Der "Kluge Assistent" (Consistency Model)

Um zu wissen, ob ein Strich gut ist, muss der Coach das fertige Bild sehen. Aber das fertige Bild existiert noch nicht!

• Der alte Weg: Der Coach versuchte, das fertige Bild aus dem Nebel zu erraten. Das war oft falsch, besonders wenn noch viel Nebel da war. Das führte zu schlechten Ratschlägen.
• Der neue Weg (SQDF): SQDF nutzt einen speziellen Assistenten (ein Consistency Model). Dieser Assistent ist darauf trainiert, das fertige Bild sofort und sehr genau zu "sehen", selbst wenn noch viel Nebel da ist. Er gibt dem Coach eine klare, präzise Vorhersage, wie das Bild aussehen wird. So bekommt der Künstler sofortiges, korrektes Feedback, ohne dass der Coach raten muss.

3. Der "Erinnerungs-Speicher" (Replay Buffer)

Stell dir vor, der Künstler malt jeden Tag ein Bild.

• Ohne Speicher: Er malt nur das Bild von heute und vergisst alles, was er gestern gelernt hat. Wenn er heute ein tolles Bild malt, aber morgen ein schlechtes, verliert er den guten Stil.
• Mit Speicher (SQDF): SQDF hat einen großen Keller (Replay Buffer), in dem er alle Bilder speichert – die guten, die schlechten und die ganz besonderen. Wenn der Coach trainiert, holt er sich nicht nur das heutige Bild, sondern schaut sich auch alte Meisterwerke an. Das hilft dem Künstler, eine große Vielfalt zu behalten und nicht nur immer das Gleiche zu malen (was man "Mode Collapse" nennt).

Wie funktioniert das Training? (Die "Reparametrisierung")

Normalerweise ist es sehr schwer, einem Künstler zu sagen: "Ändere deinen Pinselstrich hier, damit das Bild später besser wird", weil der Weg vom Anfang bis zum Ende so lang und verworren ist.

SQDF nutzt einen Trick: Es berechnet den Weg so um, dass der Coach den Einfluss eines einzelnen Pinselstrichs direkt auf das Endergebnis sehen kann, ohne den ganzen Weg neu berechnen zu müssen. Das ist wie ein GPS, das dir sofort sagt: "Wenn du jetzt links abbiegst, kommst du schneller ans Ziel", ohne dass du erst die ganze Strecke abfahren musst.

Das Ergebnis: Mehr Belohnung, mehr Spaß

In Tests hat sich gezeigt, dass SQDF:

1. Höhere Punktzahlen erreicht (die Bilder sehen besser aus).
2. Nicht verrückt wird (die Bilder bleiben natürlich und machen Sinn).
3. Vielfältig bleibt (der Künstler malt viele verschiedene Arten von Bildern, nicht nur Kopien).

Zusammenfassend:
SQDF ist wie ein smarter Trainer, der einem KI-Künstler hilft, seine Belohnung zu maximieren, ohne dabei den Verstand zu verlieren. Es nutzt einen klugen Assistenten für bessere Vorhersagen, bewertet die wichtigen Schritte höher und sorgt durch einen Erinnerungsspeicher dafür, dass der Künstler kreativ und vielfältig bleibt.

Technische Zusammenfassung

1. Problemstellung

Diffusionsmodelle haben sich als dominanter Paradigma für generative Aufgaben (z. B. Text-zu-Bild-Synthese) etabliert. Um diese Modelle an spezifische Zielsetzungen (Downstream-Objectives) wie ästhetische Qualität oder menschliche Präferenzen anzupassen, werden Fine-Tuning-Methoden eingesetzt.

Das Hauptproblem besteht jedoch darin, dass bestehende Fine-Tuning-Ansätze (sowohl Reinforcement-Learning-basierte Methoden wie PPO/DDPO als auch direkte Backpropagation-Methoden wie DRaFT/ReFL) stark zu Reward-Over-Optimierung (Belohnungs-Überoptimierung) neigen. Dies führt zu:

• Semantischem Kollaps: Die generierten Bilder verlieren den Bezug zum Eingabe-Prompt und werden zu unerkennbaren Texturen.
• Diversitätskollaps: Die Vielfalt der generierten Samples nimmt drastisch ab.
• Instabilität: Viele Methoden erfordern das Training separater Wertfunktionen (Value Functions), was bei Diffusionsmodellen als bekanntermaßen instabil gilt, oder nutzen Monte-Carlo-Schätzer mit hoher Varianz.

2. Methodik: SQDF (Soft Q-based Diffusion Finetuning)

Die Autoren schlagen SQDF vor, eine neuartige Methode zur Ausrichtung von Diffusionsmodellen, die auf einem KL-regulierten Reinforcement-Learning-Rahmenwerk basiert. Der Kernansatz besteht darin, den Gradienten einer Belohnungsfunktion direkt über eine reparametrisierte Policy-Gradient-Methode zu nutzen, die auf einer trainingsfreien, differentierbaren Schätzung der Soft-Q-Funktion aufbaut.

Kernkomponenten und Innovationen:

1. Trainingsfreie Soft-Q-Schätzung:

   • Anstatt eine separate Q-Netzwerk zu trainieren, approximiert SQDF die Soft-Q-Funktion $Q^*_{soft}$ durch eine Posterior-Mean-Approximation in einem einzigen Schritt.
   • Dies wird durch die Verwendung von Tweedie's Formel oder, wie in SQDF verbessert, durch ein Consistency Model erreicht.
   • Da die Schätzung differentierbar ist, kann der Reward-Gradient direkt zur Aktualisierung der Policy-Parameter verwendet werden, ohne den gesamten Denoising-Pfad rückwärts zu propagieren.

2. Reparametrisierter Policy Gradient:

   • Um die hohe Varianz von Gradienten zu vermeiden, wird die Stochastizität des Denoising-Schritts $x_{t-1} \sim p_\theta(x_{t-1}|x_t)$ durch Reparametrisierung ($x_{t-1} = \mu_\theta(x_t, t) + \sigma_t \epsilon$) handhabbar gemacht.
   • Dies ermöglicht effiziente Updates mit niedrigem Varianzgrad.

3. Drei spezifische Verbesserungen zur Stabilisierung:

   • Diskontfaktor ($\gamma$): Ein Diskontfaktor wird eingeführt, um den Kredit (Credit Assignment) für frühe Denoising-Schritte zu gewichten. Da frühe Schritte (bei hohem Rauschen) weniger Einfluss auf die finale Bildqualität haben, werden ihre Gradienten mit $\gamma^{t-1}$ heruntergewichtet, um das Rauschen in der Approximation zu reduzieren.
   • Consistency Models: Anstatt der oft ungenauen Tweedie'schen Formel (besonders bei hohem Rauschen) wird ein trainiertes Consistency Model verwendet, um $x_0$ aus $x_t$ vorherzusagen. Dies liefert eine präzisere Schätzung der Soft-Q-Funktion über alle Zeitschritte hinweg.
   • Off-Policy Replay Buffer: SQDF nutzt einen Replay-Buffer, um seltene, hochbelohnte und diverse Samples wiederzuverwenden. Dies verbessert die Abdeckung des Modells (Mode Coverage) und hilft, den Trade-off zwischen Belohnung und Diversität zu managen, indem Katastrophales Vergessen verhindert wird.

4. Ziel-Funktion:
   Die Verlustfunktion kombiniert die negative belohnte Schätzung mit einer KL-Divergenz-Strafe, um das feinabgestimmte Modell nahe am vor-trainierten Modell zu halten:
   $$\mathcal{L}_{SQDF}(\theta) = \mathbb{E}_{x_t \sim \mathcal{D}, x_{t-1} \sim p_\theta} [-\gamma^{t-1} r(f_\psi(x_{t-1})) + \alpha D_{KL}(p_\theta(x_{t-1}|x_t) || p'(x_{t-1}|x_t))]$$

3. Wichtige Beiträge

• Vermeidung von Instabilität: SQDF eliminiert die Notwendigkeit, separate Wertfunktionen (Value Networks) zu trainieren, was eine der Hauptquellen für Instabilität in Diffusions-RL ist.
• Effizienz: Durch die Nutzung von Gradienten aus differenzierbaren Belohnungen (z. B. Aesthetic Score, HPS) in Kombination mit Reparametrisierung wird eine hohe Sample-Effizienz erreicht.
• Robustheit gegen Over-Optimierung: Die Kombination aus KL-Regularisierung, dem Diskontfaktor und dem Replay-Buffer verhindert effektiv den semantischen und Diversitätskollaps, der bei anderen Methoden häufig auftritt.
• Black-Box-Optimierung: Die Methode wurde erfolgreich auf Szenarien mit begrenztem Query-Budget (Online Black-Box Optimization) angewendet, wo sie sowohl hohe Belohnungen als auch natürliche Ergebnisse liefert.

4. Ergebnisse

Die Autoren evaluieren SQDF in zwei Hauptszenarien:

1. Text-zu-Bild Fine-Tuning (Stable Diffusion 1.5 & XL):

   • Metriken: Aesthetic Score, HPSv2 (Human Preference Score), ImageReward, sowie Diversitätsmetriken (LPIPS, DreamSim).
   • Vergleich: SQDF wurde gegen DDPO, DRaFT, ReFL und deren KL-regularisierte Varianten verglichen.
   • Ergebnis: SQDF erreicht bei gleicher Belohnung deutlich bessere Alignment-Scores (Prompt-Treue) und Diversität als die Baselines. Während Methoden wie DRaFT bei hohen Belohnungen in semantischen Kollaps verfallen, behält SQDF die Bildqualität und Vielfalt bei. Die Ergebnisse liegen auf der Pareto-Grenze (bessere Belohnung bei gleicher Diversität).

2. Online Black-Box Optimierung:

   • In einem Setting mit begrenzten Abfragen an eine "Oracle"-Belohnungsfunktion (simuliert durch einen Proxy) übertrifft SQDF den State-of-the-Art SEIKO.
   • SQDF erzielt höhere Zielbelohnungen bei gleichzeitig besserer Erhaltung der natürlichen Eigenschaften und Diversität des Modells.

3. Ablationsstudie:

   • Das Entfernen des Diskontfaktors führt zu langsamerer Konvergenz und schlechteren Diversitätswerten.
   • Das Entfernen des Consistency Models verschlechtert die Trainingsstabilität und die Zielerreichung.
   • Das Entfernen des Replay Buffers reduziert die Diversität der generierten Samples.

5. Bedeutung und Fazit

SQDF stellt einen signifikanten Fortschritt im Bereich des Fine-Tunings von Diffusionsmodellen dar. Es löst das kritische Problem der Reward-Over-Optimierung, ohne dabei auf die Effizienz von Gradienten-basierten Methoden verzichten zu müssen.

• Praktische Relevanz: Die Methode ermöglicht es, Diffusionsmodelle sicher an komplexe menschliche Präferenzen anzupassen, ohne dass die generierten Inhalte "kaputtgehen" (kollabieren).
• Theoretischer Beitrag: Die Arbeit zeigt, wie Soft-Q-Funktionen durch Posterior-Mean-Approximationen trainingsfrei und stabil geschätzt werden können, was den Bedarf an komplexen RL-Infrastrukturen (wie Value-Function-Training) reduziert.
• Zukunftsausblick: Die Autoren sehen die Verwendung fortschrittlicherer One-Step-Distillationsmodelle und ausgefeilterer Buffer-Management-Techniken als vielversprechende Richtungen für zukünftige Arbeiten.

Zusammenfassend bietet SQDF einen robusten, effizienten und stabilen Rahmen, um die Pareto-Grenze zwischen Belohnungsoptimierung und der Erhaltung der generativen Qualität (Diversität/Natürlichkeit) bei Diffusionsmodellen zu erweitern.