🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

Dieses Paper führt die Distribution Matching Policy Optimization (DMPO) ein, ein neuartiges Reinforcement-Learning-Framework, das die Reasoning-Fähigkeiten von Diffusion-Large-Language-Models signifikant verbessert, indem es deren Policy-Distribution mit einem optimalen, belohnungs-geteilten Ziel abgleicht und dadurch erhebliche Genauigkeitssteigerungen gegenüber bestehenden Baselines ohne Supervised Fine-Tuning erzielt.

Ursprüngliche Autoren: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

Veröffentlicht 2026-07-14

📖 6 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

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 hätten einen superintelligenten Roboter, der Geschichten schreibt, mathematische Probleme löst und Spiele spielt. Die meisten Roboter von heute arbeiten wie ein Mensch, der ein Buch Wort für Wort von links nach rechts liest. Sie können nicht vorausspringen oder zurückblicken; sie müssen das nächste Wort basierend auf den vorangegangenen Wörtern erraten. Das ist langsam, so als würde man versuchen, ein riesiges Puzzle zu lösen, indem man immer nur das Teil direkt vor sich betrachtet.

Hier kommen Diffusion Large Language Models (dLLMs) ins Spiel. Sie sind die neuen Gesellen auf dem Markt. Anstatt Wort für Wort zu schreiben, beginnen sie mit einem chaotischen, durcheinandergewürfelten Haufen aus „Geheimnis-Token“ (wie ein Puzzle, bei dem jedes Teil von Nebel bedeckt ist) und bereinigen diesen langsam, um die Antwort auf einmal oder in beliebiger Reihenfolge zu enthüllen. Dies macht sie potenziell viel schneller im Denken.

Aber hier ist der Haken: Während diese Diffusions-Roboter schnell sind, sind sie nicht immer die klügsten bei schwierigen Denkaufgaben wie fortgeschrittener Mathematik oder Logikrätseln. Um sie klüger zu machen, nutzen Wissenschaftler normalerweise eine Technik namens Reinforcement Learning (RL). Betrachten Sie dies als ein Videospiel, in dem der Roboter Punkte für eine gute Antwort erhält und Punkte abzieht, wenn die Antwort schlecht ist.

Das Problem: Die „Mode-Seeking“-Falle

Die alte Art, diese Roboter zu trainieren (unter Verwendung von Methoden wie GRPO), war ein wenig wie ein Schüler, der nur die eine Antwort lernt, von der er glaubt, dass sie richtig ist. Wenn der Roboter einen Weg findet, um Punkte zu sammeln, der sich „sicher“ anfühlt, hört er auf zu explorieren. Er bleibt in einer festgefahrenen Routine stecken und ignoriert andere clevere Lösungen, die genauso gut sein könnten. In der Arbeit nennen die Autoren dies „Mode-Seeking“. Es ist wie ein Roboter, der nur lernt, ein mathematisches Problem auf eine ganz bestimmte Weise zu lösen, und wenn dieser Weg blockiert wird, gerät er in Panik. Er neigt auch dazu, die „unordentlichen“, aber korrekten Antworten zu ignorieren, die anders aussehen als die erste gefundene.

Die Autoren argumentieren, dass diese alte Art, nur nach dem höchsten Score zu jagen, fehlerhaft für Diffusionsmodelle ist. Sie legen nahe, dass wir durch die Konzentration auf nur einen einzigen besten Pfad die einzigartige Fähigkeit des Roboters verlieren, viele verschiedene Pfade gleichzeitig zu erkunden.

Die Lösung: DMPO (Der „Karten-Abgleich“-Roboter)

Die Autoren schlagen eine neue Methode namens Distribution Matching Policy Optimization (DMPO) vor.

Anstatt dem Roboter zu sagen: „Finde einfach den höchsten Score“, sagt DMPO: „Hier ist die gesamte Landkarte aller guten Antworten, die du geben könntest. Deine Aufgabe ist es, zu lernen, diese ganze Landkarte abzubilden.“

Stellen Sie sich vor, Sie versuchen, einem Hund das Apportieren beizubringen.

Die alte Art: Sie werfen einen Ball, und der Hund rennt zu dem einen Ort, an dem der Ball gelandet ist. Wenn der Ball in einem Busch landet, lernt der Hund, nur in Büschen zu suchen.
Die DMPO-Art: Sie zeigen dem Hund eine Karte von jedem möglichen Ort, an dem der Ball landen könnte (Büsche, Gras, Sand, Wasser) und sagen: „Lerne, von jedem dieser Orte zu apportieren, gewichtet danach, wie gut der Ort ist.“ Der Hund lernt, flexibel zu sein und den ganzen Garten zu erkunden, nicht nur eine einzige Ecke.

In technischer Hinsicht nutzt DMPO einen speziellen mathematischen Trick namens Weighted Denoising Cross-Entropy (WDCE). Dies ermöglicht es dem Roboter, aus seinen vergangenen Versuchen zu lernen (selbst aus denen, die nicht perfekt waren), ohne jedes Mal alles von Grund auf neu generieren zu müssen. Es ist wie ein „Replay Buffer“, in dem der Roboter seine alten Spiele immer und immer wieder studieren kann und dabei lernt, aus den guten und schlechten Zügen zu lernen, ohne verwirrt zu werden.

Das Geheimrezept: Die „Gewichts-Baseline“

Die Autoren fanden ein kniffliges Problem beim Training mit kleinen Beispielgruppen (kleinen „Batch Sizes“). Manchmal wurde der Roboter verwirrt und begann, schlechte Antworten zu belohnen, nur weil sie die einzigen waren, die er sah.

Um dies zu beheben, erfanden sie einen cleveren Trick namens Weight Baseline Subtraction.
Stellen Sie sich das wie einen Lehrer vor, der eine Prüfung bewertet. Wenn ein Schüler eine Frage richtig beantwortet, bekommt er einen goldenen Stern. Aber wenn der Lehrer nur diese eine Frage gesehen hat, könnte er denken, dass alles, was der Schüler getan hat, großartig war. Die „Baseline“ ist wie ein „Standard des Durchschnitts“. Der Lehrer zieht den „Durchschnittswert“ vom Ergebnis des Schülers ab.

Wenn der Schüler b besser als der Durchschnitt war, bekommt er einen großen goldenen Stern.
Wenn er schlechter als der Durchschnitt war, erhält er eine „Strafe“ (ein negatives Gewicht), selbst wenn er technisch gesehen einige Punkte erhalten hat.

Dies stellt sicher, dass der Roboter nicht zu begeistert über mittelmäßige Antworten ist und weiterhin nach den wirklich exzellenten strebt.

Die Ergebnisse: Wie viel besser?

Die Autoren testeten diese neue Methode an einigen sehr schwierigen Reasoning-Benchmarks. Sie wandten DMPO auf vortrainierte Modelle (wie LLaDA-Instruct und Dream-Instruct) auf zwei verschiedene Arten an:

Direkte Anwendung: Sie wandten DMPO direkt auf die Basismodelle an, ohne vorherige „Hausaufgaben“ (Supervised Fine-Tuning oder SFT) auf Reasoning-Datensätzen durchgeführt zu haben. Dieser „R1-Zero-ähnliche“ Ansatz wurde verwendet, um das rohe Potenzial von DMPO klar zu demonstrieren.
Erweiterte Anwendung: Sie wandten DMPO auch auf Modelle an, die bereits ein SFT durchlaufen hatten, um zu zeigen, dass es als leistungsstarkes Upgrade für bereits trainierte Modelle funktioniert.

Die Ergebnisse waren über beide Setups hinweg äußerst beeindruckend:

Bei einem mathematischen Rätsel namens GSM8K zeigte das Modell, das direkt auf das vortrainierte Basismodell angewendet wurde (ohate SFT), massive Verbesserungen. Speziell erreichte DMPO eine Steigerung der Genauigkeit um bis zu 39,63 Prozentpunkte gegenüber bisherigen Nicht-DMPO RL-Baselines und eine erstaunliche Steigerung um 67,97 Prozentpunkte gegenüber dem Basismodell selbst.
Bei einem Logikrätsel namens Sudoku sahen sie enorme Gewinne, wobei eine Version ihres Modells von einer Erfolgsquote von 16,41 % auf 80,86 % sprang (ein Unterschied von über 64 Punkten!).
Entscheidend ist, dass DMPO auch dann signifikante Leistungssteigerungen lieferte, wenn es auf Modelle angewendet wurde, die bereits ein SFT durchlaufen hatten, was beweist, dass es eine robuste Methode ist, die funktioniert, egal ob man bei Null anfängt oder auf bestehendem Training aufbaut.

Die Autoren legen nahe, dass DMPO nicht nur eine kleine Anpassung, sondern ein grundlegender Wandel ist. Es ermöglicht dem Roboter, „off-policy“ zu sein, was bedeutet, dass er effizient aus alten Daten lernen kann, und „forward-only“, was bedeutet, dass er keine teuren Rückwärtsberechnungen benötigt, die andere Roboter verlangsamen.

Das Fazit

Die Autoren sind zuversichtlich, dass diese Methode für die spezifischen Modelle, die sie getestet haben (wie LLaDA und Dream), bei diesen spezifischen Reasoning-Aufgaben funktioniert. Sie haben diese Verbesserungen durch rigorose Experimente auf Standard-Datensätzen gemessen und gezeigt, dass DMPO die Leistung steigern kann, egal ob es auf ein Modell angewendet wird, das gerade erst vortrainiert wurde, oder auf eines, das bereits ein Supervised Fine-Tuning durchlaufen hat. Dennoch geben sie zu, dass sie dies noch nicht auf jeden möglichen Modelltyp oder jede Aufgabe getestet haben, sodass das volle Potenzial zwar stark ist, aber noch weiter erforscht wird.

Kurz gesagt: DMPO lehrt Diffusions-Roboter, aufzuhören, sich auf eine einzige „perfekte“ Antwort zu versteifen, und stattdien zu lernen, die gesamte Landschaft guter Lösungen zu schätzen, was sie zu klügeren, schnelleren und kreativeren Problemlösern macht.

Technisches Resümee: Verbesserung der Argumentationsfähigkeit von Diffusions-LLMs mittels Distribution Matching Policy Optimization

Problemstellung

Diffusion Large Language Models (dLLMs) bieten eine vielversprechende Alternative zu autoregressiven LLMs (AR-LLMs), indem sie eine Generierung in beliebiger Reihenfolge und potenziell einen höheren Inferenzdurchsatz ermöglichen. Die Skalierung von dLLMs, um die Argumentationsfähigkeiten (Reasoning) von AR-LLMs zu erreichen, bleibt jedoch eine erhebliche Herausforderung. Bestehende Methoden des Reinforcement Learning (RL) zur Feinabstimmung, wie etwa die Group Relative Policy Optimization (GRPO), sind für dLLMs aus zwei Hauptgründen ungeeignet:

Komplexitäts-Mismatch: dLLMs basieren auf bidirektionaler Verarbeitung, was die Schätzung von Sequenz-Log-Wahrscheinlichkeiten (die für Standard-Policy-Gradient-Methoden erforderlich sind) rechenintensiv und oft unpraktikabel macht im Vergleich zu AR-Modellen.
Algorithmische Einschränkungen: Aktuelle RL-Ansätze für dLLMs verlassen sich oft auf heuristische Approximationen (z. B. Mean-Field-Dekomposition), welche die Korrelationen zwischen Token-Positionen ignorieren. Darüber hinaus konzentrieren sie sich typischerweise auf die Belohnungsmaximierung, was dazu neigt, ein einzelnes Belohnungsmodus zu optimieren. Dieses „Mode-Seeking“-Verhalten kann zu Reward Hacking führen und das Scheitern bei der Exploration der vielfältigen, hochwertigen Argumentationspfade zur Folge haben, zu denen dLLMs theoretisch fähig sind.

Methodik: Distribution Matching Policy Optimization (DMPO)

Das Paper schlägt die Distribution Matching Policy Optimization (DMPO) vor, ein prinzipbasiertes RL-Feinabstimmungs-Framework, das auf der stochastischen optimalen Steuerung (Stochastic Optimal Control, SOC) basiert. Im Gegensatz zu traditionellen Methoden, die die erwartete Belohnung maximieren, zielt DMPO darauf ab, die Verteilung der Policy des Modells an eine optimale, belohnungs-getiltte Zielverteilung anzupassen.

Kern-Framework

Das Ziel ist es, eine Policy $\pi_\theta$ zu lernen, sodass die induzierte Sequenzverteilung $p_\theta(o|q)$ der optimalen Verteilung $p^*(o|q)$ entspricht, definiert als:
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
wobei $p_{ref}$ die vortrainierte Policy ist, $r$ die Belohnungsfunktion und $\alpha$ ein Temperaturparameter.

Um dies zu erreichen, formuliert die Autore das Problem als Minimierung der Forward Kullback-Leibler-Divergenz (äquivalent zur Kreuzentropie) zwischen der Zielverteilung und der aktuellen Policy:
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
Dieses „Mass-Covering“-Ziel ermutigt das Modell, diverse hochbelohnte Trajektorien zu explorieren, anstatt zu einem einzelnen Modus zu kollabieren.

Zentrale technische Komponenten

Weighted Denoising Cross-Entropy (WDCE):
Da $p^*$ unpraktikabel ist (aufgrund der unbekannten Partitionsfunktion $Z(q)$ und der Unfähigkeit, direkt daraus zu sampeln), verwendet DMPO Importance Sampling. Es behandelt Samples aus einer Referenz-Policy $p_v$ (eine Kopie des Modells mit gestopptem Gradienten) als gewichtete Samples aus $p^*$ .
Die resultierende Verlustfunktion, WDCE, nutzt den einzigartigen vorwärtsgerichteten (Noising-)Prozess von dLLMs. Sie berechnet den Verlust nur unter Verwendung der finalen, sauberen Samples und der zugehörigen Importance-Gewichte, wobei die intermediären Rollout-Trajektorien verworfen werden. Dies macht WDCE zu einem Off-Policy- und Forward-Only-Ziel.
Weight Baseline Subtraction:
Eine kritische Implementierungsherausforderung, die identifiziert wurde, ist die Instabilität von WDCE bei der Verwendung kleiner Trainings-Batches. In kleinen Batches können positive Gewichte für „schlechte“ Antworten (die dennoch eine Wahrscheinlichkeit $> 0$ unter $p^*$ aufweisen) die Gradienten-Updates dominieren, was zu fehlerhaften Aktualisierungen führt.
Um dies zu adressieren, schlagen die Autoren die Weight Baseline Subtraction vor, bei der ein Basis-Gewicht $w_{base}$ vom Importance-Gewicht subtrahiert wird:
$w_{real} = w - w_{base}$
Es werden drei Strategien für $w_{base}$ vorgeschlagen:

Group Weight Baseline: Ein konstanter Wert (1), um die Konvergenz zu fördern.
Individual Weight Baseline: Basierend auf einem negativen Reward-Tilt.
Model Weight Baseline: Basierend auf den eigenen Importance-Gewichten des aktuellen Modells.
Diese Technik injiziert effektiv negative Gradienten für Antworten mit geringer Belohnung und verhindert so, dass das Modell schlechte Antworten in Szenarien mit geringer Batch-Größe fördert.

Weighted Direct Discriminative Optimization (WDDO):
Als Alternative zur Kreuzentropie untersucht das Paper ein GAN-inspiriertes Ziel, das naturgemäß negative Gradienten für schlechte Samples einbezieht und somit das Distribution Matching weiter verbessert.

Zentrale Beiträge

Neuartiges RL-Framework: Der Vorschlag von DMPO, einem Distribution-Matching-Framework, das speziell für dLLMs entwickelt wurde und das Paradigma von der Belohnungsmaximierung hin zum Matching der gesamten belohnungs-getilten Policy-Verteilung verschiebt.
Theoretische Fundierung: Herleitung des WDCE-Verlusts mittels Importance Sampling und Stochastic Optimal Control, was eine theoretisch fundierte Methode für das Off-Policy-Training bietet, ohne teure Schätzungen der Sequenz-Log-Wahrscheinlichkeiten zu erfordern.
Praktische Lösungen: Identifizierung des Ausfallmodus bei kleinen Batches in Forward-KL-Objektiven und Einführung von Weight Baseline Subtraction und WDDO, um das Training zu stabilisieren und die Sample-Effizienz zu verbessern.
Effizienz: Demonstration, dass WDCE ein Training allein über den Vorwärtsprozess ermöglicht, was die Wiederverwendung von Rollouts (Off-Policy) erlaubt und die Kompatibilität mit schnellen Inferenztechniken (z. B. KV-Cache) gewährleistet, was das Training im Vergleich zu trajektorienbasierten Methoden signifikant beschleunigt.

Experimentelle Ergebnisse

Die Autoren evaluierten DMPO auf den Modellen LLaDA-8B-Instruct und Dream-7B über vier Reasoning-Benchmarks hinweg: GSM8K, MATH500, Countdown und Sudoku.

Leistungssteigerungen: DMPO übertraf konsistent die Basismodelle und bestehende RL-Baselines (einschließlich diffu-GRPO, cGRPO und d1).
- Bei Sudoku erreichte DMPO-LLaDA eine Verbesserung von 39,63 Prozentpunkten gegenüber der besten Nicht-DMPO RL-Baseline und 67,97 Prozentpunkte gegenüber dem Basismodell.
- Bei Countdown wurden Verbesserungen von bis zu 16,40 Prozentpunkten gegenüber starken Baselines beobachtet.
Kein SFT erforderlich: Die Methode erzielte diese Ergebnisse ohne Supervised Fine-Tuning (SFT) auf kuratierten Datensätzen, was ihre Fähigkeit demonstriert, die Reasoning-Fähigkeiten direkt aus vortrainierten Modellen zu skalieren.
Trainingseffizienz: Aufgrund ihrer Off-Policy-Natur und der Forward-Only-Berechnung erreichte DMPO eine 2- bis 3-mal höhere Sample-Effizienz und eine bis zu 8-mal schnellere Rollout-Sampling-Rate (in Kombination mit Fast-dLLM-Inferenz) im Vergleich zu GRPO-Typen-Methoden. Sie erreichte das Belohnungsniveau von d1 mit nur 31 % des Trainings-Rechenbudgets.
Stabilität: Die Methode blieb selbst bei sehr veralteten Daten stabil (bis zu 24-fache Wiederverwendung desselben Rollout-Batches), eine Eigenschaft, die auf die Stabilität der sequenzbasierten Importance-Gewichte in DMPO zurückzuführen ist.

Bedeutung und Ansprüche

Das Paper behauptet, dass DMPO einen bedeutenden Schritt zur Skalierung von dLLMs für Reasoning-Aufgaben darstellt. Durch die Abkehr vom Paradigma der Belohnungsmaximierung und den computationalen Engpässen trajektorienbasierter RL-Methoden nutzt DMPO die einzigartigen Eigenschaften von Diffusionsmodellen voll aus:

Diversität: Das Forward-KL-Ziel verhindert den Mode Collapse und ermöglicht es dLLMs, diverse Argumentationspfade zu explorieren.
Effizienz: Die Off-Policy- und Forward-Only-Natur von WDCE harmoniert mit den hohen Durchsatzkapazitäten von dLLMs und bietet einen Pfad für skalierbares RL-Feintuning, das für große Modelle rechentechnisch machbar ist.
Generalisierbarkeit: Das Framework wird als ein prinzipierter Ansatz präsentiert, der auf verschiedene dLLM-Backbones angewendet werden kann, wie der Erfolg auf sowohl LLaDA- als auch Dream-Architekturen zeigt.

Die Autoren kommen zu dem Schluss, dass die aktuelle Arbeit sich zwar auf spezifische Modelle und Benchmarks konzentriert, das Distribution-Matching-Framework jedoch neue Wege für die Verbesserung generativer Modelle und das Design effektiverer Weight-Baseline-Techniken für die zukünftige Forschung eröffnet.