🤖 machine learning

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

Dieser Beitrag stellt TraFL vor, eine Nachtrainierungsmethode auf Basis des Trajektorienausgleichs für Diffusions-Sprachmodelle, die den Fehlermodus des „Trajektorienverriegelns" bei belohnungsmaximierenden Ansätzen überwindet, indem sie die Policy mit einer belohnungsgekippten Zielverteilung in Einklang bringt und dadurch konsistente Leistungssteigerungen bei Benchmarks für mathematisches Schlussfolgern und Codegenerierung erzielt.

Ursprüngliche Autoren: Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

Veröffentlicht 2026-05-15

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

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

Das große Ganze: Eine neue Art, KI beizubringen

Stellen Sie sich einen sehr talentierten Künstler (das KI-Modell) vor, der Bilder malt, indem er mit einer von Rauschen und statischen Störungen erfüllten Leinwand beginnt und diese langsam reinigt, bis ein klares Bild erscheint. So funktionieren Diffusions-Sprachmodelle: Sie schreiben Wörter nicht wie eine Schreibmaschine Wort für Wort (wie ältere KI-Modelle); stattdessen beginnen sie mit einem Durcheinander von Buchstaben und „entrauschen" diese allmählich zu einem kohärenten Satz.

Das Paper argumentiert, dass die aktuelle Methode, diese Künstler darin zu schulen, schwierige Probleme (wie Mathematik oder Programmierung) besser zu lösen, fehlerhaft ist. Die Autoren schlagen eine neue Methode namens TraFL (Trajectory Flow baLancing) vor, die einen spezifischen Mangel namens „Trajectory Locking" (Pfad-Verriegelung) behebt.

Das Problem: Die „Ein-Pfad"-Falle (Trajectory Locking)

Um das Problem zu verstehen, stellen Sie sich ein Labyrinth vor.

Das Ziel: Das Labyrinth hat viele verschiedene korrekte Ausgänge (verschiedene gültige Lösungen für ein mathematisches Problem).
Die alte Methode (Belohnungsmaximierendes RL): Stellen Sie sich vor, Sie trainieren einen Hund, den Ausgang zu finden. Jedes Mal, wenn der Hund einen Ausgang findet, geben Sie ihm einen Leckerbissen.
- Der Fehler: Dem Hund ist egal, welchen Weg er genommen hat, um den Leckerbissen zu erhalten, sondern nur, dass er einen erhalten hat.
- Das Ergebnis: Wenn der Hund zufällig frühzeitig auf einen bestimmten Pfad stößt, der zu einem Leckerbissen führt, erhält er einen. Er merkt sich diesen Pfad. Beim nächsten Mal versucht er diesen Pfad erneut. Er erhält einen weiteren Leckerbissen. Er wird zuversichtlicher in Bezug auf diesen einen Pfad.
- Trajectory Locking: Schließlich hört der Hund auf, das Labyrinth zu erkunden. Er rennt nur noch diesen einen schmalen Pfad entlang, selbst wenn es 50 andere gültige Ausgänge gibt, die er hätte finden können. Er hat sich auf eine einzige Lösung „verriegelt" und vergessen, wie man die anderen findet.

Im Paper bezeichnen die Autoren dies als Trajectory Locking. Da die KI nur für die endgültige Antwort eine „Belohnung" (eine Punktzahl) erhält, ignoriert sie die Tatsache, dass es viele verschiedene Wege (Pfade) gab, um dorthin zu gelangen. Sie kollabiert ihre gesamte Kreativität in eine einzige schmale Route, was sie schlecht darin macht, alternative korrekte Antworten zu finden, wenn Sie sie bitten, es noch einmal zu versuchen.

Die Lösung: Der „Führer"-Ansatz (TraFL)

Die Autoren schlagen TraFL vor, was die Trainingsregeln ändert. Anstatt nur für einen Ausgang einen Leckerbissen zu geben, geben sie der KI einen Führer (ein eingefrorenes Referenzmodell) und eine Karte.

Der Führer (Referenzmodell): Dies ist eine Version der KI, die noch nicht auf die spezifischen Belohnungen trainiert wurde. Sie repräsentiert eine „gesunde", vielfältige Denkweise. Sie weiß, dass es viele Wege gibt, ein Problem zu lösen.
Die Karte (Belohnungs-geneigtes Ziel): Wir sagen der KI: „Sie wollen die Ausgänge finden, die die Belohnung bringen (die korrekten Antworten), ABER Sie müssen Ihre Bewegungsmuster ähnlich wie beim Führer halten."

Die Analogie:
Stellen Sie sich vor, Sie unterrichten einen Schüler in der Lösung eines mathematischen Problems.

Alte Methode: Sie sagen: „Finden Sie die richtige Antwort, auf welche Weise auch immer!" Der Schüler findet einen Trick, der funktioniert, merkt ihn sich und weigert sich, jede andere Methode zu lernen.
TraFL-Methode: Sie sagen: „Finden Sie die richtige Antwort, aber halten Sie Ihren Denkprozess vielfältig und flexibel, wie ein Spitzenstudent, der viele verschiedene Strategien kennt."

TraFL zwingt die KI, zwei Dinge ins Gleichgewicht zu bringen:

Die Belohnung erhalten: Die richtige Antwort finden.
Vielfältig bleiben: Nicht in einen einzigen, repetitiven Pfad kollabieren. Sie hält die „Wahrscheinlichkeitsmasse" (die Wahrscheinlichkeit, einen Pfad zu wählen) über viele verschiedene gültige Lösungen verteilt, verankert durch den Führer.

Wie sie es funktionierend machten

Das Paper stellt fest, dass Diffusionsmodelle knifflig sind, weil sie ihren „Denkprozess" nicht Schritt für Schritt zeigen wie ältere Modelle. Um dies zu beheben, erfanden die Autoren zwei Werkzeuge:

Ein Surrogat-Score: Da sie die genaue Mathematik jedes Schritts nicht sehen können, erstellten sie einen „Proxy"-Score, der abschätzt, wie gut eine vollständige Antwort ist, sodass sie das Modell trainieren können, ohne perfekte Sichtbarkeit in jeden winzigen Schritt zu benötigen.
Ein gelernter Normalisierer: Sie lehrten der KI einen speziellen „Rechner", der die Schwierigkeit der Aufgabe basierend auf der spezifischen Frage (dem Prompt) anpasst, sodass das Training fair und ausgewogen bleibt.

Die Ergebnisse: Hilft es wirklich?

Die Autoren testeten diese neue Methode auf Mathematik (Lösen von Textaufgaben) und Code (Schreiben von Computerprogrammen).

Die „Ein-Pfad"-Falle war gebrochen: Im Gegensatz zu anderen Methoden, die manchmal schlechter darin wurden, verschiedene Lösungen zu finden, während sie besser darin wurden, eine Lösung zu finden, verbesserte sich TraFL bei jedem einzelnen Test.
Mehr Proben = Bessere Ergebnisse: Als sie die KI baten, 16 verschiedene Antworten zu generieren, anstatt nur eine, wurde TraFL deutlich besser. Dies beweist, dass sie tatsächlich mehr verschiedene korrekte Lösungen fand und nicht nur einen glücklichen Zufallstreffer.
Es funktioniert bei Neuem: Die KI hat nicht nur die Trainingsfragen auswendig gelernt. Als sie an brandneuen mathematischen Problemen (Minerva Math) und neuen Programmierherausforderungen (LiveCodeBench) getestet wurde, die sie noch nie gesehen hatte, war TraFL die stärkste Leistung.
Vielfalts-Check: Sie verwendeten einen „Richter" (eine andere KI), um die Antworten zu betrachten. Der Richter bestätigte, dass TraFL nicht einfach nur die gleiche Antwort in anderen Worten gab; sie verwendete tatsächlich unterschiedliche Denkstrategien und unterschiedliche Algorithmen, um die gleichen Probleme zu lösen.

Zusammenfassung

Das Paper identifiziert einen Fehler, bei dem KI-Modelle auf einer einzigen Art, ein Problem zu lösen, „stecken bleiben" und andere gültige Lösungen ignorieren. Sie haben dies mit TraFL behoben, einer Methode, die der KI beibringt, nach korrekten Antworten zu suchen, während sie eine vielfältige „Erkundung" verschiedener Pfade beibehält, geleitet von einem Referenzmodell. Das Ergebnis ist eine KI, die nicht nur intelligenter ist, sondern auch kreativer und zuverlässiger, wenn sie aufgefordert wird, mehrere Lösungen zu generieren.

Technische Zusammenfassung: Trajectory-Balance-Nachtraining für Diffusions-Sprachmodelle

Problem: Trajektorien-Verriegelung im Diffusions-RL

Diffusions-Sprachmodelle (dLLMs) bieten eine vielversprechende Alternative zu autoregressiven Modellen für das logische Schlussfolgern und die Codegenerierung. Allerdings passen bestehende Nachtrainingsmethoden für dLLMs überwiegend belohnungsmaximierende Reinforcement-Learning-Ziele (RL, z. B. PPO, GRPO) an, die für autoregressive Modelle entwickelt wurden. Die Autoren identifizieren einen zentralen Fehlermodus bei dieser Anpassung, der als Trajektorien-Verriegelung (trajectory locking) bezeichnet wird.

Bei Diffusionsmodellen wird die Belohnung nur der finalen entrauschten Vervollständigung ( $y$ ) zugewiesen, während die Modellverteilung über die gesamte Entrauschungs-Trajektorie ( $\tau$ ) induziert wird. Da Endbelohnungen „pfadindifferent" sind – was bedeutet, dass jede Trajektorie, die dieselbe korrekte $y$ erreicht, dieselbe Belohnung erhält – neigen geschätzte Policy-Gradient-Updates dazu, die spezifischen Trajektorien zu verstärken, die zufällig früh im Training gesampelt werden. Dies erzeugt einen sich selbst verstärkenden Feedback-Loop, bei dem sich die Wahrscheinlichkeitsmasse auf eine schmale Menge von Entrauschungspfaden kollabiert. Folglich verliert ein Modell, selbst wenn es lernt, eine korrekte Antwort zu produzieren, die Fähigkeit, unter wiederholtem Sampling eine vielfältige Menge gültiger Lösungen (Modi) zu generieren. Die Autoren formalisieren dies durch Theorem 1 und beweisen, dass die Abdeckung terminaler Modi durch die Trajektorienvielfalt nach oben beschränkt ist; somit begrenzt die Trajektorien-Verriegelung direkt die Lösungsabdeckung des Modells.

Methodik: TraFL (Trajectory Flow Balancing)

Um die Trajektorien-Verriegelung zu adressieren, schlagen die Autoren TraFL vor, ein Nachtrainingsziel, das auf dem Prinzip des Trajektorien-Gleichgewichts aus Generative Flow Networks (GFlowNets) basiert. Anstatt die Belohnung direkt zu maximieren, trainiert TraFL die Policy, um eine belohnungs-geneigte Zielverteilung anzunähern, die an ein eingefrorenes Referenzmodell angelehnt ist.

Das Ziel

Die Zielverteilung $p^*(\tau | x)$ ist definiert als:
$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$
wobei:

$p_{\text{ref}}$ ein eingefrorenes Referenzmodell ist.
$r(x, y)$ die Endbelohnung ist.
$\beta$ die Schärfe der Belohnung steuert.
$Z(x)$ ein prompt-abhängiger Normalisierungsterm (Partitionfunktion) ist.

Das Trainingsziel minimiert das quadratische Residuum zwischen dem Log-Wahrscheinlichkeitswert des aktuellen Modells und dem Ziel:
$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$
$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$

Praktische Implementierung für Diffusion

Da dLLMs keine exakten Trajektorien-Log-Wahrscheinlichkeiten offenlegen, führen die Autoren zwei Schlüsselkomponenten ein, um TraFL praktikabel zu machen:

Diffusions-kompatibler Surrogat auf Sequenzebene: Sie leiten einen handhabbaren Log-Wahrscheinlichkeits-Surrogat ( $\log \hat{p}_\theta(y|x)$ ) basierend auf einer unteren Schranke für maskierte Rekonstruktion ab. Dies beinhaltet das Sampling eines Maskierungsgrades $l$ , das Korruptieren von $l$ Tokens und das Mitteln der Log-Wahrscheinlichkeiten der rekonstruierten Tokens. Ein Faktor $1/l$ wird verwendet, um den Score unabhängig von der Anzahl der maskierten Tokens zu normalisieren.
Gelernte prompt-abhängige Normalisierung: Im Gegensatz zu Methoden, die $Z(x)$ über eine Softmax-Schätzung über einen Batch schätzen (was zu Konzentration führen kann), lernt TraFL eine Partitionfunktion $Z_\phi(x)$ über ein kleines MLP. Dieser Term erhält Gradientensignale in jedem Schritt, unabhängig davon, welche Vervollständigungen gesampelt werden, und verhindert so das bei anderen Verteilungsanpassungsansätzen beobachtete Versagen der „batch-lokalen" Normalisierung.

Hauptbeiträge

Identifikation der Trajektorien-Verriegelung: Die Arbeit identifiziert und beweist formal, dass belohnungsmaximierende Ziele in dLLMs zu einem Kollaps der Trajektorienvielfalt führen, was wiederum die Abdeckung unterschiedlicher korrekter Lösungsmodi einschränkt.
TraFL-Ziel: Die Einführung eines Trajektorien-Gleichgewichts-Ziels, das die Policy an ein Referenzmodell anlehnt und die Rolle der Belohnung (Geneigung der terminalen Vervollständigungen) von der des Referenzmodells (Zuweisung der Wahrscheinlichkeitsmasse über plausible Trajektorien) trennt.
Praktischer Surrogat und Normalisierung: Die Entwicklung eines diffusions-kompatiblen Surrogats auf Sequenzebene und eines gelernten Normalisierungsterms, die ein stabiles Training ohne exakte Trajektorienwahrscheinlichkeiten ermöglichen.
Empirische Überlegenheit: TraFL ist die einzige evaluierte Nachtrainingsmethode, die über alle Benchmark-Längeneinstellungen (GSM8K, MATH-500, HumanEval, MBPP) und Sampling-Budgets hinweg konsistent eine Verbesserung gegenüber dem Basismodell erzielt.

Ergebnisse

Die Autoren evaluieren TraFL an Benchmarks für mathematisches Schlussfolgern und Codegenerierung gegenüber dem Basismodell (LLaDA-8B-Instruct) und starken Baselines (ESPO, JUSTGRPO).

In-Distribution-Leistung: TraFL verbessert Pass@k über alle Sampling-Budgets ( $k \in \{1, 2, 4, 8, 16\}$ ) und Temperaturen hinweg. Bemerkenswert ist, dass es Pass@1 (Genauigkeit bei Einzelsampling) verbessert und gleichzeitig bei erhöhten Sampling-Budgets besser skaliert, im Gegensatz zu Baselines, die oft die Einzelsampling-Genauigkeit für Gewinne bei Mehrfach-Sampling opfern.
Generalisierung auf Held-Out-Daten: Die Gewinne übertragen sich auf Out-of-Distribution-Benchmarks. Bei Minerva Math bleibt TraFL über dem Basismodell. Bei LiveCodeBench ist TraFL die stärkste Methode über alle Schwierigkeits-Splits hinweg (Einfach, Mittel, Schwer) und übertrifft die Baselines bei mittleren und schwierigen Problemen signifikant.
Vielfaltsanalyse: Unter Verwendung eines LLM-as-a-judge-Protokolls zeigen die Autoren, dass die Gewinne von TraFL mit einer breiteren Abdeckung unterschiedlicher Lösungsstrategien (z. B. verschiedene Beweisansätze oder algorithmische Zerlegungen) verbunden sind und nicht nur mit einer Schärfung des Scores eines einzelnen Modus. Baselines wie ESPO zeigen, obwohl sie die Genauigkeit verbessern, keine signifikanten Vielfaltsgewinne gegenüber dem Basismodell, was mit der Hypothese der Trajektorien-Verriegelung übereinstimmt.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass TraFL eine prinzipielle Lösung für das „Modus-suchende" Versagen des standardmäßigen RL-Nachtrainings für Diffusionsmodelle bietet. Durch den Wechsel von der Belohnungsmaximierung zum Trajektorien-Gleichgewicht erhält die Methode die Vielfalt der Entrauschungspfade, die notwendig ist, um mehrere gültige Lösungen abzudecken.

Die Autoren stellen fest, dass TraFL die einzige evaluierte Methode ist, die in jedem getesteten Benchmark-Längensetting eine Verbesserung gegenüber dem Basismodell erzielt. Sie betonen, dass die Verbesserungen robust gegenüber Erhöhungen des Sampling-Budgets sind und sich auf Held-Out-Datensätze verallgemeinern, was darauf hindeutet, dass das Trajektorien-Gleichgewichts-Ziel das Belohnungssignal effektiv vom Pfad-Allokationsmechanismus entkoppelt. Die Arbeit plädiert für Verteilungsanpassungs-Ziele gegenüber der reinen Belohnungsmaximierung für dLLM-Nachtraining und bietet einen praktischen Rahmen, der die Konzentration-Versagensmodi vermeidet, die in standardmäßigen Policy-Gradient-Ansätzen inhärent sind.

Die Arbeit bleibt bezüglich der Einschränkungen bescheiden und stellt fest, dass die Methode auf einem approximativen Surrogat auf Sequenzebene statt auf exakten Trajektorienwahrscheinlichkeiten beruht und grobe binäre Belohnungen (exakte Übereinstimmung oder Testbestehen) verwendet, was die Lerneffizienz in Szenarien mit seltenen Erfolgen einschränken kann.