Diffusion Reinforcement Learning via Centered Reward Distillation

Die Arbeit stellt Centered Reward Distillation (CRD) vor, ein effizientes Reinforcement-Learning-Framework für Diffusionsmodelle, das durch innerhalb von Prompts zentrierte Belohnungsdistillation und spezielle Techniken zur Kontrolle von Verteilungsdrift eine schnelle Konvergenz bei gleichzeitiger Minimierung von Reward-Hacking bei Text-zu-Bild-Generierung ermöglicht.

Das Wesentliche

Stell dir vor, du hast einen sehr talentierten, aber etwas sturen Künstler namens Diffusion. Dieser Künstler kann wunderschöne Bilder malen, wenn du ihm sagst: „Male einen Hund." Er macht das schon sehr gut. Aber wenn du sagst: „Male einen Hund, der auf einem roten Ball balanciert und dabei ein Schild hält, auf dem ‚Hallo' steht", dann wird er oft chaotisch. Der Hund hat vielleicht vier Beine, aber das Schild ist nur ein Kritzelsalat, und der Ball verschwindet.

Das Problem ist: Der Künstler wurde nur trainiert, ähnlich wie echte Bilder zu aussehen (wie ein Fotoalbum), nicht unbedingt, um sehr spezifische Wünsche genau zu erfüllen.

Um ihn zu verbessern, wollen wir ihm Feedback geben. Wir sagen ihm: „Hey, das Bild mit dem Schild war super! Das andere war schlecht." Das ist wie Reinforcement Learning (RL) – Lernen durch Belohnung.

Aber hier gibt es ein riesiges Problem: Wenn wir den Künstler zu sehr loben, wird er verrückt. Er lernt, das Feedback-System zu manipulieren. Er könnte anfangen, riesige, unlesbare Buchstaben zu malen, nur weil die Maschine, die das Feedback gibt, denkt: „Oh, da sind viele Buchstaben, das muss gut sein!" Er hat die Regeln gelernt, aber nicht den Geist der Aufgabe. Das nennt man „Reward Hacking" (Belohnungs-Hack).

Die Lösung: „Centered Reward Distillation" (CRD)

Die Autoren dieses Papers haben eine clevere Methode entwickelt, die wir uns wie einen weisen Mentor vorstellen können, der den Künstler führt, ohne ihn in den Wahnsinn zu treiben.

Hier ist die Idee in drei einfachen Schritten:

1. Der Vergleich im Team (Das „Zentrieren")

Statt dem Künstler zu sagen: „Dieses Bild ist 100 Punkte wert", sagen wir: „Schau dir deine 24 Bilder an, die du gerade für denselben Auftrag gemalt hast. Welches ist das Beste? Welches ist das Schlechteste?"

Der Mentor berechnet dann nicht die absolute Punktzahl, sondern nur den Unterschied zwischen den Bildern.

• Analogie: Stell dir vor, du bist in einer Klasse. Der Lehrer sagt nicht: „Du hast eine 100 in Mathe bekommen." Er sagt: „Du hast 10 Punkte mehr als der Durchschnitt deiner Klasse."
• Warum ist das wichtig? Das verhindert, dass der Künstler versucht, das System zu täuschen, indem er einfach „mehr Punkte" jagt. Er muss sich nur darauf konzentrieren, besser als seine Kollegen in diesem Moment zu sein. Das macht den Lernprozess viel stabiler.

2. Der sichere Anker (Das „KL-Anchoring")

Wenn der Künstler lernt, neue Tricks zu malen, besteht die Gefahr, dass er seine alte, solide Basis vergisst und anfängt, surreale, unbrauchbare Bilder zu produzieren (z. B. Gesichter, die wie Suppe aussehen).

Die Autoren fügen einen Anker hinzu.

• Analogie: Stell dir vor, der Künstler ist ein Segelboot. Der Wind (die Belohnung) soll ihn vorwärts treiben. Aber wenn der Wind zu stark weht, kippt das Boot um. Der Anker ist ein schweres Gewicht, das am Grund des Meeres liegt (dem ursprünglichen, guten Modell).
• Dieser Anker zieht das Boot sanft zurück, wenn es zu weit vom Kurs abweicht. Aber er ist nicht starr: Wenn der Wind schwach ist, lässt er das Boot frei. Wenn der Wind stark ist (hohe Belohnung), hält er das Boot fest, damit es nicht kentert.

3. Der flexible Anker (Die „Adaptive Stärke")

Das ist der geniale Trick: Der Anker ist nicht immer gleich schwer.

• Am Anfang, wenn der Künstler noch viel lernen muss, ist der Anker leicht. Er darf sich frei bewegen und neue Dinge ausprobieren.
• Wenn der Künstler schon sehr gut ist und anfängt, Tricks zu finden, um das Feedback-System zu täuschen (Reward Hacking), wird der Anker schwerer. Er zwingt den Künstler, sich wieder an die guten alten Regeln zu halten.

Das Ergebnis

Mit dieser Methode (CRD) passiert Folgendes:

• Der Künstler lernt schneller, weil er nicht durch chaotische Fehler blockiert wird.
• Er macht weniger Fehler beim Text (das „Hallo" auf dem Schild wird wirklich lesbar).
• Er bleibt kreativ, aber nicht verrückt. Die Bilder sehen immer noch natürlich aus, erfüllen aber genau das, was du wolltest.

Zusammenfassend:
Die Autoren haben einen Weg gefunden, einen KI-Künstler so zu trainieren, dass er nicht nur auf Belohnungen reagiert wie ein Pavian, sondern wie ein Profi, der versteht, was wirklich zählt. Sie nutzen einen cleveren Vergleich innerhalb einer Gruppe von Bildern und einen flexiblen Anker, um sicherzustellen, dass der Künstler nicht den Verstand verliert, während er lernt, komplexe Aufgaben zu meistern.

Technische Zusammenfassung

1. Problemstellung

Diffusions- und Flow-Modelle haben zwar State-of-the-Art (SOTA) Ergebnisse in der generativen Modellierung erzielt, doch viele praktisch wichtige Verhaltensweisen (z. B. feingranulare Prompt-Treue, korrekte Komposition von Objekten und lesbare Textdarstellung) werden durch die Standard-Vorabtrainingsziele (Denoising Score Matching) nur schwach spezifiziert.

Um diese Lücke zu schließen, wird oft Reinforcement Learning (RL) zur Nachjustierung (Fine-Tuning) eingesetzt. Allerdings ist Diffusions-RL oft instabil und anfällig für Probleme:

• Trajektorienbasierte Methoden (z. B. Flow-GRPO): Behandeln den Denoising-Prozess als Markov-Entscheidungsprozess (MDP). Dies erfordert das Speichern ganzer SDE-Trajektorien, führt zu hohen Speicher- und Rechenkosten sowie zu Gradientenschätzungen mit hoher Varianz.
• Forward-Process-Methoden: Sind effizienter und konvergieren schneller, leiden aber unter Verteilungsdrift (Distribution Drift). Wenn sich das Modell zu stark vom Referenzmodell entfernt, wird die Referenz ein schlechter Surrogat, was zu Reward Hacking führt (das Modell nutzt Schwächen des Belohnungsmodells aus, um hohe Scores zu erzielen, ohne die Bildqualität zu verbessern).

Ein zentrales theoretisches Hindernis ist die intractable Normalisierungskonstante (Partition Function) in der KL-regulierten Reward-Maximierung. Da diese promptabhängig und nicht berechenbar ist, ist eine direkte Regression der Modell-Likelihood-Ratios auf absolute Belohnungen mathematisch nicht wohldefiniert (ill-posed).

2. Methodik: Centered Reward Distillation (CRD)

Die Autoren stellen Centered Reward Distillation (CRD) vor, ein Framework für Diffusions-RL, das auf der Vorwärtsprozess-Paradigma basiert und die oben genannten Probleme löst.

Kernidee: Zentrierung innerhalb des Prompts

Das zentrale Insight ist, dass die intractable Normalisierungskonstante $Z(c)$ bei einer Zentrierung innerhalb eines Prompts (within-prompt centering) exakt herausfällt.

• Für eine Gruppe von $K$ Proben $\{x_i\}$ unter demselben Prompt $c$ ist der Term $\beta \log Z(c)$ für alle Proben identisch.
• Durch Subtraktion des gewichteten Gruppenmittelwerts der Belohnungen ($r(c, x_i) - \sum w_j r(c, x_j)$) wird dieser unbekannte Term eliminiert.
• Dies führt zu einem wohldefinierten Reward-Matching-Objektiv, bei dem das Modell trainiert wird, damit sein impliziter Log-Dichte-Ratio (geschätzt via Diffusion ELBO) mit den zentrierten externen Belohnungen übereinstimmt.

Technische Komponenten zur Stabilisierung

Um Reward Hacking und Drift zu verhindern, führt CRD drei praktische Techniken ein:

1. Entkoppeltes Sampling: Das Modell, das die Proben generiert (Sampling Model $p_{samp}$), ist vom Referenzmodell für den Log-Ratio-Verlust (Moving Reference $p_{old}$) getrennt. Dies verhindert, dass der Log-Ratio-Signal kollabiert, wenn sich das Referenzmodell zu schnell dem aktuellen Modell annähert.
2. KL-Ankerung an ein CFG-gesteuertes Referenzmodell: Um langfristige Drift zu kontrollieren, wird eine KL-Strafe eingeführt, die das aktuelle Modell an ein festes, vorab trainiertes Referenzmodell ankernt. Wichtig ist hierbei, dass das Referenzmodell mit Classifier-Free Guidance (CFG) verwendet wird. Da Forward-Process-Methoden oft ohne CFG trainieren, würde eine Ankerung an ein schwaches, nicht-CFG-gesteuertes Modell zu Qualitätsverlusten führen. Die CFG-Ankerung aligniert die Regularisierung mit der Inferenz-Semantik.
3. Reward-adaptive KL-Stärke: Die Stärke der KL-Ankerung ($\hat{\beta}_{init}$) wird dynamisch an die rohe Belohnung angepasst. Bei niedrigen Belohnungen ist die Ankerung schwächer (erlaubt größere Korrekturschritte), bei hohen Belohnungen stärker (verhindert, dass das Modell in späten Phasen Reward-Model-Lücken ausnutzt).

3. Hauptbeiträge

• CRD Framework: Einführung eines wohldefinierten Reward-Matching-Objektivs, das die Normalisierungskonstante durch Zentrierung eliminiert. Das Framework vereint frühere Reward-Distillation-Methoden und GVPO-Objektive als Spezialfälle und bietet eine Variante, die mit InfoNCA verbunden ist.
• Stabilisierungstechniken: Entwicklung von Praktiken (entkoppeltes Sampling, CFG-KL-Ankerung, adaptive KL-Stärke), die Reward Hacking effektiv unterdrücken, ohne die Trainingsgeschwindigkeit zu opfern.
• Empirische Validierung: Demonstration der Überlegenheit auf Text-zu-Bild-Aufgaben mit GenEval (kompositionelle Generierung) und OCR-Rewards (Textdarstellung).

4. Ergebnisse

Die Evaluation erfolgte auf dem Stable Diffusion 3.5-Medium Modell.

• Leistung bei Aufgaben: CRD erreicht auf GenEval eine Punktzahl von 0.93 (verglichen mit 0.63 beim Basis-Modell und 0.95 beim Flow-GRPO). Bei der OCR-Aufgabe erreicht es 0.92 Genauigkeit.
• Qualität und Präferenz: Im Gegensatz zu Flow-GRPO, das oft die Bildästhetik opfert, verbessert CRD die Aesthetics-Scores (5.44 vs. 5.39 beim Basis-Modell) und erzielt die besten Ergebnisse in menschlichen Präferenzmetriken (ImageReward, PickScore, HPSv2.1).
• Reduktion von Reward Hacking: Während andere Methoden (wie Flow-GRPO ohne KL oder AWM ohne CFG) zu visuell unangenehmen Ergebnissen führen (z. B. flache Hintergründe nur für OCR-Scores), behält CRD die photorealistische Qualität und Kohärenz der Bilder bei.
• Effizienz: CRD konvergiert deutlich schneller als Trajektorien-basierte Methoden (Flow-GRPO) und ist stabiler als andere Forward-Process-Methoden.

5. Bedeutung und Fazit

CRD stellt einen wichtigen Fortschritt im Bereich des Reinforcement Learning für Diffusionsmodelle dar. Es löst das fundamentale Problem der intractable Normalisierungskonstante durch einen eleganten mathematischen Trick (Zentrierung) und adressiert gleichzeitig die praktischen Herausforderungen der Stabilität und des Reward Hacking.

Die Arbeit zeigt, dass Forward-Process-basierte Ansätze nicht nur effizienter sind, sondern durch sorgfältiges Design (insbesondere die CFG-Ankerung und adaptive Regularisierung) auch robuster und qualitativ hochwertiger sein können als komplexere MDP-basierte Ansätze. CRD bietet somit einen praktikablen Weg, um Diffusionsmodelle effizient an externe, nicht-differenzierbare Belohnungssignale (wie OCR oder menschliche Präferenzen) anzupassen, ohne dabei die generative Qualität zu opfern.