Das Team von Wang Mengdi hat TraceRL vorgestellt: Schritt in Richtung „RL-Einheit“ für Diffusionssprachmodelle

In Anbetracht der Beschränkungen von autoregressiven Large Language Models (LLMs) bei der Rechenleistung und Durchsatz erhalten Diffusions Language Models (DLMs) immer mehr Aufmerksamkeit.

Derzeit fehlt es jedoch in der Branche an einem einheitlichen und effektiven Reinforcement Learning (RL)-Framework, das für verschiedene DLM-Architekturen (wie vollständige Attention-DLMs und Block-Attention-DLMs) geeignet ist. Gleichzeitig wird in den bestehenden Studien die Wichtigkeit der Übereinstimmung von Inferenzpfaden und Trainingszielen vernachlässigt.

Kürzlich hat das Team von Professor Mengdi Wang der Princeton University in einer neuesten Studie das “Trajektorien-sensitives RL”-Framework — TraceRL vorgeschlagen, das auf vollständige Attention- und Block-Attention-Modelle angewendet werden kann und eine schnelle Optimierung ermöglicht.

Link zur Studie: https://arxiv.org/abs/2509.06949

Es ist erwähnenswert, dass das mit TraceRL trainierte 4B-DLM-Modell bei mehreren komplexen mathematischen Inferenztasks besser abgeschnitten hat als das 7B-autoregressive Modell.

Durch Curriculum Learning haben sie auch das erste Long Chain-of-Thought Diffusions Language Model (long-CoT DLM) entwickelt, dessen relative Genauigkeit auf MATH500 um 18,1 % höher ist als die von Qwen2.5-7B-Instruct.

Abbildung | Links: RL-Trainingsdynamik verschiedener Methoden, wobei TraceRL die beste Optimierung erreicht. Rechts: Vergleich der Ergebnisse von komplexen mathematischen Inferenztasks basierend auf KV-Caching und dem LiveCodeBench-V2-Benchmark.

Zusätzlich haben sie ein auf Diffusion basiertes Wertmodell vorgeschlagen, um die Varianz zu verringern und die Trainingsstabilität zu verbessern. Sie haben auch die weiteren potenziellen Anwendungen von TraceRL erforscht, wie die Erhöhung der Blockgröße des Modells und die Beschleunigung der Inferenz.

Darüber hinaus haben sie ein vollständig integriertes Framework — dLLM-RL open source gemacht, um DLMs über verschiedene Architekturen hinweg zu erstellen, zu trainieren und bereitzustellen. Dieses Framework enthält die Implementierungen verschiedener Post-Training-Methoden und Techniken zur Beschleunigung des KV-Cachings und unterstützt sowohl reproduzierbare Forschung als auch praktische Anwendungen.

Adresse: https://github.com/Gen-Verse/dLLM-RL

Die dringende Notwendigkeit, das “Missverhältnis” von DLMs zu lösen

Das Forschungsunternehmen betont, dass es ein erhebliches Missverhältnis zwischen den Zielen, die DLMs im Post-Training-Stadium verfolgen, und den Pfaden, die sie bei der tatsächlichen Inferenz (Textgenerierung) verfolgen, gibt. Standard-Trainingsmethoden, wie die vollständige Zufallsmaskierung, können zwar parallele Decodierung ermöglichen, ignorieren jedoch die inhärente, kontextabhängige logische Reihenfolge der Sprache. Dieser Bruch zwischen Training und Inferenzverhalten führt zu einer ineffizienten Modelloptimierung.

Um diesen Unterschied zu verdeutlichen, haben sie zunächst durch Experimente gezeigt, dass die halb-autoregressive Fine-Tuning-Methode, bei der das Modell trainiert wird, auf der Grundlage des vorherigen Kontexts nachfolgende Inhalte zu generieren, eine deutlich bessere Optimierungsleistung hat als die vollständige Zufallsmaskierungsmethode, selbst wenn die Rechenlast gleich ist. Dies zeigt, dass die Übereinstimmung von Trainingszielen und Inferenzmodus von entscheidender Bedeutung ist.

Um die Wichtigkeit der Übereinstimmung weiter zu validieren, haben sie die “bevorzugten Inferenzpfade” des Modells gesammelt, d. h. die Schrittfolgen, die das Modell tatsächlich bei der Inhaltsgenerierung befolgt. Die Experimentergebnisse zeigen, dass das Fine-Tuning mit diesen echten Inferenzpfaden eine bessere Leistung erzielt als andere Baseline-Methoden und dass die Rechenkosten niedriger oder vergleichbar sind.

Schließlich erfordert die Sammlung dieser Pfade zwar viel zusätzliche Arbeit, obwohl das Fine-Tuning mit bevorzugten Pfaden gute Ergebnisse erzielt. Im Gegensatz dazu werden diese Inferenzpfade bei RL natürlich während des “Rollouts” (d. h. der Generierung von Stichproben durch das Modell) erzeugt. Daher ist RL eine praktischere und effektivere Post-Training-Strategie, die diese Pfade natürlich nutzen kann, um das Modell zu optimieren.

TraceRL: “Kleines Diffusionsmodell” > “Großes autoregressives Modell”

In dieser Arbeit konzentriert sich TraceRL auf die von DLMs generierten Zwischenpfade und kann über verschiedene Architekturen hinweg angewendet werden.

Abbildung | Überblick über TraceRL. In diesem Beispiel werden die Parameter s = 2, L = 6 und B = 3 verwendet. Durch die Aggregation von jeweils s benachbarten Schritten wird Trajektorien-sensitives RL realisiert. Die Zahlen in den Kästchen entsprechen der Ausführungsreihenfolge des Strategie-Inferenzprozesses.

Bei den Daten hat das Forschungsunternehmen verschiedene Datenquellen verwendet:

(1) Es wurde der Math-Trainingsdatensatz ausgewählt, um 8000 anspruchsvolle Aufgaben zu generieren;

(2) GSM8K, MATH500 und AIME2024 wurden als Testbenchmarks ausgewählt, um Inferenztasks im Bereich Mathematik und Programmierung zu bewerten;

(3) Im Szenario des Coding-Reinforcement Learning wurden 6000 verifizierte Aufgaben von der PrimeIntellect-Plattform verwendet, um die Validierung durchzuführen;

(4) Für die Programmiertests wurden LiveCodeBench-V2 und LiveBench als Testbenchmarks ausgewählt.

Tabelle | Hauptbenchmark-Ergebnisse für verschiedene mathematische und programmatische Aufgaben. “Statisch” bedeutet statische Stichprobenahme, “Dynamisch” bedeutet dynamische Stichprobenahme. Hier wird die dynamische Stichprobenahme mit einem Schwellenwert von 0,9 verwendet, um das Long-CoT-Modell TraDo-8B-Instruct zu bewerten.

Beim Modelltraining wurden sowohl vollständige Attention-Modelle als auch Block-Attention-Modelle verwendet, und bei der Bewertung wurden sowohl die Ergebnisse der statischen Stichprobenahme als auch der dynamischen Stichprobenahme berichtet. Der Prozess ist wie folgt:

Erster Schritt: Das Modell wird separat mit TraceRL trainiert;

Zweiter Schritt: Das Long-CoT-Modell wird gemeinsam trainiert;

Drritter Schritt: TraceRL wird mit anderen RL-Methoden in Vergleichsexperimenten getestet;

Vierter Schritt: TraceRL wird in vollständigen Attention-Modellen und Coding-Tasks validiert;

Fünfter Schritt: Es werden Experimente zur Erweiterung der Blockgröße durchgeführt.

Auf der Grundlage der Experimentergebnisse haben sie die Effektivität und starke Leistung von TraceRL bewiesen. Die vollständigen Ergebnisse sind wie folgt:

Zunächst haben sie auf der Grundlage des SDAR-Basismodells mit TraceRL zwei Modelle, TraDo-4B-Instruct und TraDo-8B-Instruct, entwickelt. Bei der Bewertung in mathematischen, programmatischen und 5 Inferenzdatensätzen haben diese Modelle nicht nur effektiv mit starken Diffusionssprachenmodellen und autoregressiven Modellen verglichen, sondern auch eine bemerkenswerte Überlegenheit in der Generierungsfähigkeit gezeigt.

Abbildung | Trainingskurven von 4B- und 8B-Modellen in mathematischen Aufgaben mit TraceRL. Die rote Kurve zeigt die Genauigkeit der dynamischen Stichprobenahme, die schneller ist; die blaue Kurve zeigt die Genauigkeit der statischen Stichprobenahme, die eine höhere Genauigkeit liefert. Das 4B-Modell wird mit einem Wertmodell trainiert, während das 8B-Modell direkt mit Jpolicy trainiert wird.

TraDo-4B-Instruct hat auf Inferenztasks ein SOTA-Niveau erreicht, was die Effektivität von TraceRL beweist. Sowohl die dynamische Stichprobenahme (schneller) als auch die statische Stichprobenahme (genauer) haben die Leistung des Modells erheblich verbessert. Es ist bemerkenswert, dass TraDo-4B-Instruct in allen mathematischen Aufgaben sogar bessere Ergebnisse erzielt hat als starke autoregressive Baselinemodelle wie Qwen2.5-7B Instruct.

Obwohl sie bei der RL-Trainingsphase eine Strategie der dynamischen Stichprobenahme angewandt haben, haben sich sowohl die dynamische als auch die statische Genauigkeit stetig verbessert, und es deutet sich an, dass das Modell noch Potential zur Erweiterung hat. Dieses RL-Training hat die mathematische Inferenzfähigkeit des Modells erheblich verbessert:

Im MATH500-Test hat das TraDo-4B Instruct-Modell eine Verbesserung der statischen Genauigkeit um 5,4 % und der dynamischen Genauigkeit um 4,2 % erzielt und nach der Optimierung Qwen2.5-7B-Instruct übertroffen; das TraDo-8B-Instruct-Modell hat dagegen eine Verbesserung der statischen Genauigkeit um 4,2 % und der dynamischen Genauigkeit um 4,8 % erreicht.

Abbildung | Ablationsstudie von RL-Methoden auf der Grundlage von Block-Diffusionsmodellen und mathematischen RL-Tasks. Die rote und die gelbe Kurve entsprechen den Trainingsergebnissen von TraceRL mit und ohne Wertmodell, die blaue Kurve verwendet ein Zufallsmaskierungsziel ähnlich der halb-autoregressiven Trainingsmethode für das Training innerhalb des Blocks, und die grüne Kurve erzielt zusätzliche Trainingsergebnisse durch die Hinzufügung einer komplementären Maskierung innerhalb des Blocks.

Darüber hinaus haben sie TraceRL mit bestehenden RL-Methoden verglichen, wobei sie sich auf Block-Diffusionsmodelle konzentriert haben. Obwohl die aktuellen RL-Methoden hauptsächlich für vollständige Attention-Modelle entwickelt wurden, haben sie sie direkt auf die Blockstruktur angepasst. Bei der Zufallsmaskierungsmethode haben sie die Stichprobenahme innerhalb jedes Blocks eingeschränkt, so dass sie der halb-autoregressiven Methode ähnelt. Bei der gekoppelten RL haben sie in jedem Trainingsmodul eine komplementäre Zielfunktion eingeführt, um ein stabileres und effizienteres Training zu erzielen. Die Experimentergebnisse auf der Grundlage von mathematischen Aufgaben zeigen, dass TraceRL unabhängig von der Optimierungsstrategie die beste Leistung zeigt.

Abbildung | Ablationsstudie des RL-Trainings des vollständigen Attention-Modells Dream-7B-Coder-Instruct, mit Fokus auf Coding-Tasks; der Vergleich zwischen der Verwendung und Nicht-Verwendung eines Wertmodells zeigt, dass die Einführung eines Wertmodells die Schwankungen während des Trainings effektiv reduzieren kann. Alle Experimente wurden in mathematischen Aufgaben mit einem 4B-Modell durchgeführt.

Darüber hinaus haben sie um die breite Anwendbarkeit von TraceRL zu validieren, auch Experimente mit vollständigen Attention-Modellen in Coding-RL-Tasks durchgeführt. Auf der Grundlage von Dream-7B-Coder-Instruct wurde nach dem Fine-Tuning mit Destillationsdaten in der Cold-Start-Phase das RL-Training durchgeführt. Um den Trainingsvorgang zu beschleunigen, wurde der Kontraktionsparameter auf s = 8 festgelegt. Die Experimente zeigen, dass TraceRL schneller konvergiert und eine bessere Leistung hat.

Zukunftstraum

Auf der Grundlage der obigen Experimentergebnisse hat das Forschungsunternehmen die Effektivität von TraceRL in verschiedenen RL-Tasks validiert. Gleichzeitig haben sie die Vorteile von TraceRL bei der Beschleunigung der Inferenz und der Erweiterung der Blockgröße gezeigt, was vielversprechende Richtungen für zukünftige Forschungen bietet.

Insbesondere die Kombination der Beschleunigungsfähigkeit der Inferenz von Diffusionsmodellen mit ihrer potenziell starken Inferenzfähigkeit repräsentiert eine vielversprechende Forschungsrichtung. Obwohl die aktuellen Long-CoT-