🤖 machine learning

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

Dieser Artikel schlägt Trajectory Matching Policy Optimization (TMPO) vor, ein neuartiges Reinforcement-Learning-Framework, das die Maximierung skalierter Belohnungen durch eine Verteilungsmatching auf Trajektorienebene mittels eines Softmax Trajectory Balance-Ziels und Dynamic Stochastic Tree Sampling ersetzt, um Reward Hacking wirksam zu mindern und dadurch die generative Vielfalt und Effizienz bei der Ausrichtung von Diffusionsmodellen erheblich zu verbessern.

Ursprüngliche Autoren: Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

Veröffentlicht 2026-05-13

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

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 vor, Sie unterrichten einen talentierten Künstler (ein KI-Bildgeneratorsystem) darin, wie er basierend auf Ihrem Feedback malt. Sie möchten, dass der Künstler schöne Bilder erstellt, die Ihrer Beschreibung entsprechen, aber Sie wollen auch, dass er kreativ und variabel ist und nicht einfach immer wieder exakt dasselbe malt.

Dieses Papier, TMPO, stellt eine neue Methode vor, um diesen Künstler zu unterrichten, die ein Hauptproblem löst: Der Künstler neigt dazu, zu „betrügen", indem er einen spezifischen Trick findet, der bei Ihnen eine hohe Punktzahl erzielt, und diesen Trick dann endlos wiederholt, wobei er alle anderen guten Möglichkeiten ignoriert.

Hier ist die Aufschlüsselung der Ideen des Papiers unter Verwendung einfacher Analogien:

1. Das Problem: Der „Ein-Trick-Pony"

Aktuelle Methoden zum Training dieser KI-Künstler funktionieren wie ein strenger Lehrer, der sich nur um die Endpunktzahl kümmert.

Das Szenario: Sie bitten den Künstler, „einen Roboter beim Yoga zu malen".
Der alte Weg (Belohnungsmaximierung): Der Künstler probiert ein paar Posen aus. Eine Pose erhält eine Punktzahl von 9/10. Der Lehrer sagt: „Toll! Mache das noch einmal!" Der Künstler erkennt: „Wenn ich nur diese eine spezifische Yoga-Pose mache, erhalte ich immer eine hohe Punktzahl." Also hört er auf, andere Posen auszuprobieren. Er hört auf, Roboter in verschiedenen Studios, mit verschiedenen Farben oder bei verschiedenen Yoga-Bewegungen zu malen. Er malt einfach immer wieder denselben einen Roboter.
Das Ergebnis: Die KI wird langweilig. Sie erzeugt einen „Mode Collapse" (Modus-Kollaps), bei dem sie nur eine Art von Bild generiert, obwohl es Tausende anderer gültiger, schöner Möglichkeiten gibt, einen Roboter beim Yoga zu zeichnen. Sie beginnt zudem, das System zu „spielen", indem sie seltsame Details hinzufügt, die das Bewertungssystem täuschen, aber für Menschen schlecht aussehen.

2. Die Lösung: TMPO (Trajectory Matching Policy Optimization)

Die Autoren schlagen eine neue Lehrmethode namens TMPO vor. Anstatt dem Künstler nur zu sagen: „Mache das eine Ding, das die höchste Punktzahl erzielt hat", ändert TMPO das Ziel vollständig.

Das neue Ziel: Anstatt eine einzelne Punktzahl zu maximieren, fordert TMPO den Künstler auf, die Verteilung der Belohnungen anzupassen.
Die Analogie: Stellen Sie sich vor, der Lehrer hat eine Tasche voller verschiedener „guter" Ergebnisse. Einige sind perfekt (10/10), einige sind sehr gut (8/10) und einige sind nur okay (6/10).
- Der alte Lehrer sagte: „Zeige mir nur die 10/10er."
- Der TMPO-Lehrer sagt: „Ich möchte, dass du alle guten Ergebnisse im Verhältnis zu ihrer Güte malst. Wenn es drei Wege gibt, um 8/10 zu erreichen, möchte ich alle drei sehen, nicht nur einen."
Wie es funktioniert: Die KI generiert gleichzeitig einen ganzen „Baum" verschiedener Versuche. Dann betrachtet sie die gesamte Gruppe und sagt: „Okay, ich muss sicherstellen, dass meine Wahrscheinlichkeit, diese verschiedenen Versionen zu malen, der 'Güte' jeder Version entspricht." Dies zwingt die KI, weiterhin verschiedene Stile und Layouts zu erkunden und so die Vielfalt zu bewahren.

3. Das Geheimnis: Der „Baum"-Trick

Eine KI so zu trainieren, dass sie sich 27 verschiedene Versionen eines Bildes gleichzeitig ansieht, ist normalerweise sehr langsam und teuer (wie einen Maler zu bitten, 27 vollständige Gemälde zu skizzieren, bevor er das beste auswählt).

Die Innovation: TMPO verwendet eine Technik namens Dynamic Stochastic Tree Sampling (Dynamisches stochastisches Baum-Sampling).
Die Analogie: Stellen Sie sich vor, der Künstler beginnt, einen Hintergrund (das „Präfix") zu zeichnen. Anstatt 27 separate Gemälde von Grund auf zu vollenden, zeichnet er ein Hintergrundbild. Dann, an drei bestimmten Momenten, verzweigt er sich in verschiedene Richtungen.
- Verzweigung 1: „Okay, lass uns den Roboter blau machen."
- Verzweigung 2: „Okay, lass uns den Roboter rot machen."
- Verzweigung 3: „Okay, lass uns den Roboter grün machen."
Da sie den anfänglichen Hintergrund teilen, muss die KI die gesamte Szene nicht 27 Mal neu zeichnen. Sie berechnet nur die Unterschiede an den „Verzweigungspunkten". Dies spart eine enorme Menge an Zeit (in ihren Tests bis zu 27 % schneller), während die KI dennoch viele verschiedene Möglichkeiten erkunden kann.

4. Die Ergebnisse: Mehr Vielfalt, weniger Betrug

Das Papier testete dies an einem beliebten KI-Modell (FLUX.1) mit drei verschiedenen Aufgaben:

Kompositionelle Generierung: Sicherstellen, dass Objekte am richtigen Ort sind (z. B. „eine Katze auf einer Matte").
Textdarstellung: Wörter korrekt innerhalb des Bildes zu schreiben.
Menschliche Präferenz: Bilder zu erstellen, die für Menschen gut aussehen.

Was passierte?

Vielfalt: TMPO erzeugte Bilder, die 9,1 % vielfältiger waren als die besten bestehenden Methoden. Die Bilder sahen unterschiedlich aus, anstatt Klone zu sein.
Qualität: Es wurde keine Qualität für Vielfalt geopfert. Die Bilder waren immer noch präzise und von hoher Qualität.
Effizienz: Das Training war aufgrund des „Baum"-Tricks schneller.

Zusammenfassung

Denken Sie an TMPO als einen weisen Kunstlehrer, der erkennt, dass, wenn Sie nur die „perfekte" Antwort loben, der Schüler aufhört, kreativ zu denken. Stattdessen lehrt TMPO die KI, das Spektrum guter Antworten zu schätzen. Durch die Verwendung einer „Baum"-Struktur, um viele Pfade gleichzeitig zu erkunden, ohne Zeit zu verschwenden, schafft es eine KI, die sowohl intelligent ist (hohe Punktzahlen erzielt) als auch kreativ ist (nicht in einer Sackgasse stecken bleibt).

Das Papier behauptet, dass dies das Problem des „Reward Hacking" (bei dem die KI das System betrügt) löst, indem es mathematisch beweist, dass das Anpassen der Verteilung der Belohnungen die KI zwingt, alle „guten" Möglichkeiten abzudecken, nicht nur die einzelne „beste".

Technische Zusammenfassung: TMPO für diverse und effiziente Diffusionsausrichtung

1. Problemstellung

Reinforcement Learning (RL) hat ein erhebliches Potenzial bei der Ausrichtung von Diffusions- und Flow-Matching-Modellen auf nachgelagerte Aufgaben gezeigt. Bestehende Methoden verlassen sich jedoch überwiegend auf die Maximierung skalare Belohnungen, was die Autoren als inhärent modensuchend identifizieren. Dieses Optimierungsziel führt dazu, dass Modelle die Wahrscheinlichkeitsmasse auf wenige hochbelohnte Trajektorien konzentrieren, was zu zwei kritischen Fehlermodi führt:

Verschlechterung der generativen Vielfalt: Das Modell kollabiert in eine schmale Menge von Kompositionen und Stilen (visueller Mode-Collapse) und ignoriert die während des Vor-Trainings erlernte anisotrope Verteilung plausibler Ausgaben.
Verstärkung unzuverlässiger Belohnungen: Durch die übermäßige Verstärkung spezifischer Merkmale, die Belohnungen proxyieren (z. B. OCR-Genauigkeit oder bestimmte ästhetische Scores), verstärkt das Modell zufällige Attribute, oft auf Kosten der gesamten Generierungsqualität.

Aktuelle Minderungsstrategien, wie die KL-Regularisierung, versagen darin, die grundlegende modensuchende Natur des Ziels zu ändern, während bestehende Generative Flow Network (GFlowNet)-Ansätze oft auf unlösbare Zustandsfluss-Schätzungen oder lokale Einschränkungen angewiesen sind, die sich über lange Denoising-Horizonte hinweg zu Fehlern summieren.

2. Methodik: Trajectory Matching Policy Optimization (TMPO)

TMPO verschiebt das Optimierungsparadigma grundlegend von der Maximierung skalare Belohnungen hin zur Trajektorien-Level-Belohnungsverteilungsanpassung.

2.1 Softmax Trajectory Balance (Softmax-TB)

Anstatt die erwartete Belohnung zu maximieren, zielt TMPO darauf ab, die Trajektorienwahrscheinlichkeiten der Richtlinie an eine durch Belohnungen induzierte Zielverteilung anzupassen.

Ziel: Für eine Gruppe von $K$ Trajektorien, die aus demselben Prompt generiert wurden, passt die Methode die Richtlinienvahrscheinlichkeiten $P_\theta(\tau)$ an eine Boltzmann-Verteilung $q(\tau) \propto \exp(\beta R(\tau))$ an.
Partition-freie Formulierung: Durch die Normalisierung von Wahrscheinlichkeiten und Belohnungen innerhalb der Gruppe von $K$ Trajektorien wird die unlösbare globale Partitionfunktion eliminiert. Die Methode definiert einen Softmax-TB-Vorteil:
$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$
Theoretische Garantie: Die Autoren beweisen, dass die Minimierung dieses Ziels äquivalent zur Minimierung der Vorwärts-KL-Divergenz ( $D_{KL}(q \parallel p)$ ) innerhalb der Gruppe ist. Im Gegensatz zur Rückwärts-KL (die in Standard-Policy-Gradienten verwendet wird) ist die Vorwärts-KL modenabdeckend: Sie verhängt eine unendliche Strafe, wenn die Richtlinienvahrscheinlichkeit $p_i$ für jede Trajektorie $i$ mit positivem Zielgewicht $q_i$ gegen Null geht. Dies garantiert theoretisch die Erhaltung der Vielfalt ohne zusätzliche KL-Regularisierung.

2.2 Dynamisches stochastisches Baum-Sampling

Um die Rechenkosten des Samplings mehrerer unabhängiger Trajektorien für großskalige Flow-Matching-Modelle zu adressieren, führt TMPO eine baumstrukturierte Sampling-Strategie ein:

Präfix-Sharing: Trajektorien teilen sich ein deterministisches ODE-Präfix (gemeinsame Denoising-Schritte) und verzweigen sich an dynamisch geplanten stochastischen Punkten.
Dynamische Verzweigung: Verzweigungspunkte sind nicht fest; sie folgen einem Curriculum, das sich mit fortschreitendem Training von frühen hochrauschigen Regionen (grobkörnige Vielfalt) zu späteren niedrigrauschigen Regionen (feinkörnige Variationen) verschiebt. Verzweigungspositionen werden stochastisch unter Verwendung einer Beta-Verteilung gestört, um eine Überanpassung an eine feste Baumgeometrie zu verhindern.
Effizienz: Dieser Ansatz reduziert redundante Berechnungen durch das Teilen des Forward-Passes für das gemeinsame Präfix, was die Trainingszeit erheblich senkt, während eine effektive Exploration erhalten bleibt.

3. Hauptbeiträge

Identifikation der Grundursache: Die Arbeit identifiziert die modensuchende Natur der Belohnungsmaximierung als primäre Ursache für Reward-Hacking und Verschlechterung der Vielfalt im Diffusions-RL und schlägt einen Wechsel zur Verteilungsanpassung vor.
Softmax-TB-Ziel: Einführung eines partition-freien, trajektorien-Level-Ziels zur Verteilungsanpassung, das die modenabdeckenden Eigenschaften der Vorwärts-KL-Divergenz erbt und die generative Vielfalt auf natürliche Weise erhält.
Dynamisches stochastisches Baum-Sampling: Ein skalierbarer Sampling-Mechanismus, der die Trainingszeit für großskalige Modelle durch das Teilen von Denoising-Präfixen und die dynamische Anpassung von Verzweigungsplänen reduziert.

4. Experimentelle Ergebnisse

Die Autoren evaluierten TMPO auf dem FLUX.1-dev-Backbone über drei Ausrichtungsaufgaben: kompositionelle Bildgenerierung (GenEval), visuelle Textrendering (OCR) und Ausrichtung auf menschliche Präferenzen (PickScore).

Vielfalt: TMPO verbesserte die generative Vielfalt im Durchschnitt um 9,1 % im Vergleich zu State-of-the-Art-Methoden (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). Es war die einzige Methode, die konsistent positive Log Geometric Mean Distance (LGMD)-Scores aufrechterhielt, was auf eine gesunde Vielfalt hinweist, während Baselines oft negative LGMD-Werte (Mode-Collapse) zeigten.
Leistung: TMPO erzielte konkurrenzfähige oder überlegene Leistungen bei aufgaben spezifischen Metriken (GenEval-Genauigkeit, OCR-Genauigkeit) und menschlichen Präferenzscores (HPS-v2.1, ImageReward, PickScore).
Effizienz: Durch die Nutzung von Präfix-Sharing und Baum-Sampling reduzierte TMPO die Iterations-Trainingszeit um bis zu 27 % im Vergleich zu bestehenden Methoden.
Kompromisse: Die Pareto-Analyse zeigte, dass TMPO den optimalen Kompromiss zwischen Belohnungsmaximierung, Vielfaltserhaltung und Trainingseffizienz erreicht.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass TMPO eine grundlegende Lösung für das Dilemma des Reward-Hacking in der generativen KI-Ausrichtung bietet. Durch die Neuformulierung des Ziels zur Anpassung belohnungsinduzierter Verteilungen anstelle der Maximierung skalare Belohnungen erzwingt TMPO strukturell die Erhaltung der Vielfalt, ohne sich auf externe Regularisierung oder heuristische Neugewichtung zu verlassen. Die Methode zeigt, dass es möglich ist, hohe Ausrichtungsscores zu erreichen, während die während des Vor-Trainings erlernte reiche, multimodale Ausgabeverteilung erhalten bleibt, und dies gleichzeitig mit einer Reduzierung des Rechenaufwands, der typischerweise mit Multi-Trajektorien-RL-Methoden verbunden ist. Die Autoren positionieren TMPO als skalierbares Framework, das für großskalige Flow-Matching-Modelle geeignet ist, mit potenziellen Erweiterungen auf Videogenerierung und robotische Aktionsgenerierung.