🤖 machine learning

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

Dieses Papier stellt Reinforce Adjoint Matching (RAM) vor, eine skalierbare RL-Nachtrainierungsmethode für Diffusions- und Flow-Matching-Modelle, die eine überlegene Ausrichtung mit Belohnungen erreicht, indem sie einen einfachen Konsistenzverlust ableitet, der die Vortrainierungsziele korrigiert, ohne kostspielige SDE-Rollouts, rückwärtige adjungierte Durchläufe oder Belohnungsgradienten zu erfordern.

Ursprüngliche Autoren: Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

Veröffentlicht 2026-05-12

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich einen Meisterkünstler vor, der Jahre damit verbracht hat, das Malen zu erlernen, indem er Tausende von Fotos kopierte. Dieser Künstler ist hervorragend darin, das Gesehene nachzubilden, weiß aber nicht unbedingt, wie man spezifische, knifflige Anweisungen befolgt, wie etwa „malt ein rotes Zebra neben einen blauen Lastwagen" oder „schreiben Sie das Wort 'HALLO' deutlich auf einen Schild".

Um dem Künstler diese neuen Fähigkeiten beizubringen, verwenden wir üblicherweise eine Methode namens Reinforcement Learning (RL). Denken Sie daran wie an einen strengen Kunstkritiker, der das Gemälde betrachtet, ihm eine Note gibt und dem Künstler sagt: „Machen Sie es beim nächsten Mal besser."

Der alte Weg: Der teure, langsame Kritiker

Früher war das Trainieren dieser KI-Künstler mit einem Kritiker so, als würde man jemandem das Schwimmen beibringen, indem man ihn ins Meer wirft und zusehen muss, wie er um sich schlägt.

Der Prozess: Die KI würde ein vollständiges Bild generieren (das „Schwimmen"), der Kritiker würde es bewerten, und dann würde das System versuchen herauszufinden, genau welcher Pinselstrich für die gute oder schlechte Bewertung verantwortlich war.
Das Problem: Um dies zu tun, musste die KI den gesamten Malprozess von Anfang bis Ende immer wieder simulieren, nur um ein einziges Feedback zu erhalten. Es war langsam, rechenintensiv und oft instabil (wie der Versuch, die Windgeschwindigkeit zu berechnen, während das Boot sinkt).

Der neue Weg: RAM (Reinforce Adjoint Matching)

Die Autoren dieses Papers, Andreas Bergmeister und sein Team, erkannten, dass es einen viel schlaueren Weg gibt, dies zu tun. Sie fanden einen „Abkürzungsweg", der das Training schnell und einfach hält, ähnlich wie die ursprüngliche Vorab-Trainingsphase.

So funktioniert RAM, anhand einer einfachen Analogie:

1. Der Trick mit dem „sauberen Endpunkt"
Stellen Sie sich den Malprozess wie einen Film vor, der rückwärts läuft. Der Film beginnt mit einer leeren Leinwand (Rauschen) und endet mit einem fertigen Gemälde.

Alte Methode: Die KI musste den ganzen Film ansehen, das Ende bewerten und dann den Film Bild für Bild zurückspulen, um zu sehen, wo es schiefging.
RAM-Methode: Die Autoren erkannten, dass man, wenn man das endgültige Gemälde (den „sauberen Endpunkt") kennt, nicht den ganzen Film ansehen muss, um den Prozess zu verstehen. Man kann einfach dieses fertige Bild nehmen und mathematisch sofort „Rauschen hinzufügen", um eine unordentliche Version zu erstellen, genau wie beim ursprünglichen Training.

2. Das „Einmal-und-fertig"-Feedback
Anstatt den gesamten Malprozess zu simulieren, um Feedback zu erhalten, macht RAM Folgendes:

Generieren: Die KI malt ein fertiges Bild (den Endpunkt).
Bewerten: Der Kritiker gibt ihm eine Note (z. B. „Großartige Arbeit beim Text!").
Sofortiges Zurückspulen: Anstatt das Zurückspulen zu simulieren, erstellt die Mathematik sofort eine „verrauschte" Version dieses Bildes.
Lernen: Die KI lernt, diese spezifische verrauschte Version zurück in das hochbewertete Bild zu verwandeln.

Die magische Erkenntnis:
Das Paper beweist, dass die Regeln dafür, wie man einem Bild „Rauschen hinzufügt", sich nicht ändern, selbst wenn man versucht, das Bild zu verbessern. Das Einzige, was sich ändert, ist, welche Bilder die KI überhaupt zuerst zu malen beschließt. Da die „Rauschregeln" gleich bleiben, kann die KI dieselbe einfache Mathematik verwenden, die sie während ihres ursprünglichen Trainings genutzt hat, nur mit einem neuen Ziel.

Warum das eine große Sache ist

Geschwindigkeit: Das Paper behauptet, RAM sei 34 bis 50 Mal schneller als frühere Methoden. Es erreicht das gleiche Fähigkeitsniveau in einem Bruchteil der Zeit.
Einfachheit: Es erfordert keine komplexen, instabilen Berechnungen (wie „SDE-Rollouts" oder „backward adjoint sweeps"). Es ist einfach eine Regressionsaufgabe, ähnlich dem ursprünglichen Training.
Qualität: Als sie es an Stable Diffusion 3.5M testeten, wurde die KI deutlich besser in:
- Komponierbarkeit: Objekte an die richtigen Stellen zu setzen (z. B. ein Lastwagen links von einem Kühlschrank).
- Textdarstellung: Wörter klar in Bildern zu schreiben.
- Menschliche Präferenz: Bilder zu erstellen, die Menschen tatsächlich gerne ansehen.

Das Fazit

Denken Sie an RAM als daran, dem Künstler einen „magischen Radiergummi" und eine „magische Wertungskarte" zu geben. Anstatt den Künstler zu zwingen, die gesamte Leinwand neu zu malen, um zu sehen, was schiefgelaufen ist, zeigt der magische Radiergummi ihnen sofort eine unordentliche Version ihrer besten Arbeit, und die Wertungskarte sagt ihnen genau, wie sie diesen spezifischen Unfug beheben sollen. Dies ermöglicht es der KI, komplexe neue Fähigkeiten zu erlernen, ohne die hohen Rechenkosten früherer Methoden.

Technische Zusammenfassung: Reinforce Adjoint Matching (RAM)

Problemstellung

Diffusions- und Flow-Matching-Modelle erzielen eine hochfidele Generierung, da ihr Pretraining auf überwachter Regression basiert: Ein sauberes Sample wird analytisch verrauscht, und das Modell regrediert gegen ein geschlossenes Ziel. Das Abstimmen dieser Modelle auf Belohnungsfunktionen (z. B. für kompositionelle Korrektheit, Textrendering oder menschliche Präferenzen) mittels Reinforcement Learning (RL) im Nachtraining stört jedoch diese effiziente Struktur.

Bestehende RL-Nachtrainingsmethoden stehen vor erheblichen rechnerischen und strukturellen Hürden:

Policy-Gradient-Methoden: Sie behandeln das Denoising als Markov-Entscheidungsprozess, erfordern stochastische SDE-Rollouts und optimieren variationsuntere Schranken. Dies führt zu hoher Varianz und erfordert viele Sampling-Schritte.
Stochastische Optimalsteuerung (Adjoint)-Methoden: Sie zielen exakt auf das KL-regularisierte Belohnungsmaximierungsziel ab, erfordern jedoch rückwärtsgerichtete Adjoint-Scans entlang gespeicherter SDE-Trajektorien und durch den Sampling-Prozess zurückpropagierte Belohnungsgradienten. Dies ist numerisch instabil und rechnerisch teuer, insbesondere für Flow-Matching-Modelle, bei denen die Rauschpläne nahe dem verrauschten Ende divergieren.
Surrogat-Verluste: Neuere Ansätze ersetzen nicht handhabbare Wahrscheinlichkeiten durch Surrogate (z. B. ELBOs oder kontrastive Endpunkte) und opfern dabei die saubere Regressionsstruktur, die das Skalieren des Pretrainings ermöglicht.

Die Kernherausforderung besteht darin, ein RL-Nachtrainingsziel abzuleiten, das die Einfachheit und Skalierbarkeit des Pretraining-Regressionsrahmens beibehält, während es gleichzeitig effektiv Belohnungen optimiert.

Methodik: Reinforce Adjoint Matching (RAM)

Die Autoren schlagen Reinforce Adjoint Matching (RAM) vor, eine Methode, die die Pretraining-Regressionsstruktur auf das RL-Nachtraining erweitert, ohne SDE-Rollouts, rückwärtsgerichtete Adjoint-Scans oder Belohnungsgradienten zu erfordern.

Theoretische Grundlage

Die Methode basiert auf KL-regularisierter Belohnungsmaximierung, formuliert als stochastisches Optimalsteuerungsproblem.

Struktur des optimalen Prozesses: Unter KL-Regularisierung kippt der optimale generative Prozess die Verteilung sauberer Endpunkte ( $X_0$ ) hin zu hochbelohnten Samples ( $p^*_0 \propto p_{ref}^0 \exp(r)$ ), lässt jedoch das bedingte Gesetz verrauschter Zustände gegeben einen sauberen Endpunkt unverändert. Insbesondere bleibt die Trajektorie von einem festen $X_0$ zu einem verrauschten Zustand $X_t$ derselbe analytische Gauß-Kernel, der im Pretraining verwendet wurde.
Adjoint-Matching-Bedingung: Die optimale Steuerung $u_t$ erfüllt eine Fixpunktbedingung, bei der die Steuerung gleich dem negativ skalierten Gradienten der Wertfunktion (dem Adjunkt) ist: $u_t = -\sigma_t \nabla_x V_t(x)$ .
REINFORCE-Zerlegung: Um die Berechnung von $\nabla r$ (die nicht differenzierbar sein kann) zu vermeiden, wenden die Autoren die Log-Derivative- (REINFORCE-) Identität auf die Wertfunktion an. Dies zerlegt den Adjunkt in einen Belohnungsterm, gewichtet mit dem Score des rückwärtsgerichteten Brückens, und einen Pfadkosten-Korrekturterm.

Das RAM-Ziel

RAM approximiert den Adjunkt, indem er die Pfadkosten-Korrektur weglässt (gerechtfertigt durch geringe Varianz und Stabilität im großen Maßstab) und die strukturelle Erkenntnis nutzt, dass der Rauschkernel analytisch bleibt.

Endpunktabtastung: Anstatt eine gesteuerte SDE zu simulieren, zieht RAM einen sauberen Endpunkt $X_0$ aus der Verteilung des aktuellen Modells unter Verwendung eines beliebigen handelsüblichen ODE-Samplers.
Analytisches Rauschen: Der abgetastete $X_0$ wird analytisch verrauscht, um Trainingszustände $X_t = (1-t)X_0 + t\epsilon$ zu generieren, genau wie im Pretraining.
Geschlossenes Ziel: Das Regressionsziel für das Geschwindigkeitsfeld $v_\theta$ wird abgeleitet, indem der Belohnungsterm und der Bayes-Brücken-Score (der für lineares Rauschen einen geschlossenen Ausdruck hat) in die Adjunkt-Bedingung substituiert werden.

Die resultierende Verlustfunktion lautet:
$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$
wobei $X_0 \sim p_\theta^0$ , $\epsilon \sim \mathcal{N}(0, I)$ und $\text{sg}(\cdot)$ den Stop-Gradient-Operator bezeichnet.

Schlüsselmechanismen für Effizienz

Keine SDE-Rollouts: Trainingszustände werden generiert, indem ein einzelner Endpunkt $K$ -mal mit unabhängigen Rauschvektoren verrauscht wird. Dies liefert $K$ bedingt unabhängige Trainingsproben pro Vorwärtsdurchlauf des Modells und bietet mehr Gradientensignal pro Schritt als Methoden, die auf einzelnen SDE-Trajektorien basieren.
Keine Belohnungsgradienten: Die Methode funktioniert mit Black-Box-Belohnungen.
Keine Rückwärts-Scans: Die analytische Natur des Rauschkernels eliminiert die Notwendigkeit einer rückwärtsgerichteten Adjoint-Integration.

Hauptbeiträge

Strukturelle Erweiterung: Die Arbeit zeigt, dass sich die Regressionsstruktur des Pretrainings unter KL-Regularisierung auf das RL-Nachtraining erstreckt. Der optimale Prozess ändert nur die Endpunktverteilung, nicht jedoch das bedingte Rauschgesetz.
RAM-Algorithmus: Ein neuartiger Konsistenzverlust, der das Pretraining-Ziel mit Belohnungsinformationen unter Verwendung eines geschlossenen Bayes-Brücken-Scores korrigiert und so die rechnerischen Engpässe bestehender Methoden umgeht.
Theoretische Begründung: Eine Herleitung, die zeigt, dass der Fixpunkt von RAM dem KL-regularisierten Optimum entspricht, mit Exaktheit garantiert für Gaußsche Referenzen unter linearen Belohnungen und allgemein erster Ordnung Übereinstimmung.
Varianzanalyse: Eine Analyse von Pfadkosten-Schätzern, die zeigt, dass das Weglassen der Pfadkosten-Korrektur (wie RAM es tut) die Varianz im Vergleich zu Schätzern mit vollem Horizont oder lokalen Adjunkten signifikant reduziert, was ihn zum einzigen für das Nachtraining im Bildmaßstab viablen Schätzer macht.

Experimentelle Ergebnisse

Die Autoren evaluierten RAM auf Stable Diffusion 3.5 Medium (SD3.5M) über drei Aufgaben hinweg:

Kompositionelle Bildgenerierung (GenEval): Messung der Korrektheit von Objekten, Attributen und räumlichen Beziehungen.
Visuelles Textrendering (OCR): Messung der Lesbarkeit von Text in generierten Bildern.
Ausrichtung auf menschliche Präferenzen (PickScore): Ausrichtung an menschlichen Präferenzen.

Leistung:

Belohnungsmaximierung: RAM erzielte bei allen drei Aufgaben die höchsten Trainingsbelohnungen und übertraf Baselines einschließlich Flow-GRPO, DiffusionNFT und AWM.
Trainingseffizienz: RAM erreichte die Spitzenbelohnung von Flow-GRPO in 50× weniger Schritten für GenEval, 48× weniger für OCR und 34× weniger für PickScore.
Bildqualität: Im Gegensatz zu einigen Baselines, die unter „Reward Hacking" litten (Verschlechterung der ästhetischen Qualität zugunsten von Aufgabenmetriken), behielt RAM die Bildqualitätsmetriken (Aesthetic, DeQA, ImageReward, HPSv2) bei oder verbesserte sie, während er die Aufgabenbelohnungen maximierte.
CFG-Kompatibilität: RAM bleibt nach dem Nachtraining mit Classifier-Free Guidance (CFG) kompatibel, wohingegen Baselines wie DiffusionNFT und AWM einen signifikanten Rückgang ihrer Aufgabenbelohnungen verzeichneten, wenn CFG angewendet wurde.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass RAM einen Paradigmenwechsel im RL-Nachtraining für Diffusions- und Flow-Matching-Modelle darstellt. Durch die Nutzung der spezifischen strukturellen Eigenschaften des Optimalsteuerungsproblems stellt RAM das „einfache Regression"-Paradigma wieder her, das das Pretraining skalierbar machte.

Die Autoren behaupten, dass:

RAM die Notwendigkeit kostspieliger SDE-Rollouts und rückwärtsgerichteter Adjoint-Scans eliminiert, die die Hauptengpässe für die Skalierung von RL auf die Generierung hochauflösender Bilder waren.
Die Methode „einfach und skalierbar" ist und State-of-the-Art-Ergebnisse mit deutlich weniger Trainingsschritten und vergleichbaren Berechnungskosten pro Schritt im Vergleich zu bestehenden Methoden erzielt.
Der Ansatz Modelle erfolgreich mit komplexen Belohnungen (Komponierbarkeit, Text, Präferenzen) abstimmt, ohne sichtbare Qualitätsverschlechterung oder Reward Hacking, und zeigt, dass die Pretraining-Struktur robust genug ist, um Nachtrainingsausrichtung zu handhaben, wenn die korrekten theoretischen Bedingungen erfüllt sind.