💬 NLP

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

Die Arbeit stellt die Sandwiched Policy Gradient (SPG)-Methode vor, die durch die Nutzung sowohl einer oberen als auch einer unteren Schranke für die Log-Likelihood die Verzerrung bei der Ausrichtung von Diffusions-Sprachmodellen (dLLMs) mittels Reinforcement Learning reduziert und damit in mathematischen Benchmarks signifikant bessere Ergebnisse erzielt als bisherige Ansätze.

Ursprüngliche Autoren: Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

Veröffentlicht 2026-04-16

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

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 Problem: Der „versteckte" Preis

Stell dir vor, du hast einen genialen Koch (das KI-Modell), der Gerichte (Texte) kochen kann. Normalerweise kocht er Schritt für Schritt: Erst die Zwiebeln, dann das Fleisch, dann die Soße. Das ist wie bei herkömmlichen KI-Modellen (autoregressive Modelle).

Aber dieser spezielle Koch, den die Forscher untersucht haben, ist ein Diffusions-Koch. Er fängt mit einem Teller voller roher Zutaten an, die alle durcheinander sind (verdeckt oder „maskiert"), und versucht, das perfekte Gericht zu „herauszaubern", indem er schrittweise die Zutaten sortiert und freilegt. Das Tolle daran: Er kann viele Zutaten gleichzeitig sortieren (Parallelität), was ihn sehr schnell macht.

Das Problem: Wenn man diesen Koch trainieren will, damit er besser wird (z. B. bei Matheaufgaben oder Logikrätseln), braucht man eine Art „Bewertungssystem". Man muss ihm sagen: „Das war ein tolles Gericht!" (Belohnung) oder „Das war ekelhaft!" (Bestrafung).

Um ihn zu trainieren, müsste man normalerweise wissen, wie wahrscheinlich es ist, dass er genau dieses Gericht kocht. Aber bei diesem Diffusions-Koch ist diese Wahrscheinlichkeit wie ein geheimnisvoller Preis in einem verschlossenen Safe. Man kann ihn nicht direkt berechnen.

Bisher haben andere Forscher versucht, den Safe mit einem Schätzwert (einer groben Schätzung, dem „ELBO") zu öffnen.

Wenn das Gericht gut war, hat die Schätzung funktioniert.
Aber: Wenn das Gericht schlecht war und man dem Koch sagen wollte: „Mach das nie wieder!", funktionierte die Schätzung nicht mehr. Es war, als würde man ihm sagen: „Versuche, den Safe zu öffnen", aber der Schlüssel passte nur für die Tür, nicht für den Safe. Das führte dazu, dass der Koch bei Fehlern nicht richtig lernte, sondern verwirrt wurde.

Die Lösung: Das „Sandwich" (SPG)

Die Forscher von Meta, MIT und anderen haben eine clevere Lösung gefunden: Das Sandwiched Policy Gradient (SPG).

Stell dir vor, du willst den genauen Preis eines Gegenstands wissen, aber du hast nur zwei Schätzwerte:

Einen unteren Schätzwert (der garantiert zu niedrig ist).
Einen oberen Schätzwert (der garantiert zu hoch ist).

Der wahre Preis liegt irgendwo dazwischen – wie das Fleisch im Sandwich.

Das SPG-Verfahren nutzt genau diese beiden Schätzwerte, je nachdem, wie gut das Ergebnis war:

Wenn das Ergebnis gut ist (Belohnung):
Der Koch soll das Gericht öfter kochen. Hier nutzen die Forscher den unteren Schätzwert. Da dieser Wert garantiert kleiner als der wahre Preis ist, aber in die richtige Richtung zeigt, können sie den Koch sicher anspornen, ohne Angst zu haben, dass die Mathematik kollabiert. Es ist wie ein sicherer Treppenabsatz nach oben.
Wenn das Ergebnis schlecht ist (Bestrafung):
Der Koch soll das Gericht nicht mehr kochen. Hier nutzen sie den oberen Schätzwert. Da dieser Wert garantiert höher als der wahre Preis ist, können sie den Koch effektiv davon abhalten, diesen Weg zu gehen. Es ist wie ein sicherer Treppenabsatz nach unten.

Das Ergebnis: Indem sie das „Sandwich" nutzen (Untere Grenze für Gutes, Obere Grenze für Schlechtes), umschließen sie den geheimen, unzugänglichen wahren Preis perfekt. Der Koch lernt aus beiden Richtungen – aus Lob und aus Kritik – und wird dadurch viel besser.

Ein weiterer Trick: Das „Block-Weise" Maskieren

Es gab noch ein zweites Problem: Wenn man den Koch trainiert, muss man ihm manchmal Teile des Gerichts verdecken, um zu sehen, ob er sie richtig errät.

Der alte Weg: Man verdeckte zufällige Zutaten (z. B. eine Karotte hier, ein Gewürz dort). Das war chaotisch und passte nicht zu der Art, wie der Koch normalerweise arbeitet.
Der neue Weg (Block-weise): Man verdeckt ganze Blöcke. Stell dir vor, der Koch arbeitet in Reihen. Man verdeckt die letzte Reihe komplett, lässt die erste Reihe sauber und verdeckt die mittlere Reihe nur teilweise.

Das ist wie beim Lesen eines Buches: Man liest den Anfang, überspringt ein paar Seiten (verdeckt), und versucht, den Rest zu erraten. Da der Koch im Training genau so trainiert wird, wie er später arbeitet, funktioniert das viel stabiler und schneller.

Was hat das gebracht?

Die Forscher haben ihr neues „Sandwich-Verfahren" auf harten Tests ausprobiert:

Mathe-Rätsel (GSM8K, MATH500): Der Koch löst die Aufgaben deutlich genauer.
Logik-Rätsel (Countdown, Sudoku): Hier war der Unterschied riesig! Bei Sudoku hat sich die Genauigkeit fast verdreifacht (von ca. 27% auf 94% in einem Szenario).

Zusammenfassend:
Die Forscher haben einen Weg gefunden, wie man eine KI, die Texte nicht Schritt-für-Schritt, sondern parallel „herauszaubert", effektiv trainieren kann. Sie haben das Problem des unzugänglichen „Preises" (der Wahrscheinlichkeit) gelöst, indem sie ein Sandwich aus zwei sicheren Schätzwerten gebaut haben. Das Ergebnis ist eine KI, die bei komplexen Aufgaben wie Mathe und Logik deutlich schlauer und zuverlässiger wird als ihre Vorgänger.

1. Problemstellung

Diffusions-Modelle für große Sprachmodelle (dLLMs) gewinnen als effiziente Alternative zu autoregressiven Modellen an Bedeutung, da sie durch die parallele Dekodierung mehrerer Tokens die Inferenzlatenz reduzieren können. Ein zentrales Hindernis bei der Ausrichtung dieser Modelle auf menschliche Präferenzen oder aufgabenspezifische Belohnungen (z. B. durch Reinforcement Learning, RL) ist jedoch die unberechenbare Log-Likelihood der generierten Sequenzen.

Herausforderung: Standard-Policy-Gradient-Methoden benötigen den Gradienten der Log-Likelihood ( $\log \pi_\theta(x|c)$ ), der bei Diffusionsmodellen nicht direkt berechenbar ist.
Bisherige Ansätze: Vorherige Arbeiten umgingen dies, indem sie die Log-Likelihood durch eine untere Schranke, die Evidence Lower Bound (ELBO), oder durch eine Ein-Schritt-Schätzung approximierte.
Nachteil: Diese einseitigen Approximationen führen zu einem signifikanten Bias im Policy Gradient. Insbesondere führt das Minimieren der ELBO für negative Belohnungen (schlechte Ausgaben) nicht notwendigerweise zur Minimierung der wahren Log-Likelihood. Dies verhindert, dass das Modell effektiv aus negativen Feedback-Signalen lernt, und verzerrt die finale Policy, insbesondere bei Algorithmen, die relative oder negative Belohnungen nutzen.

2. Methodik: Sandwiched Policy Gradient (SPG)

Die Autoren schlagen SPG (Sandwiched Policy Gradient) vor, einen neuen RL-Algorithmus, der eine robustere und weniger verzerrte Schätzung des Policy Gradients ermöglicht.

Kernidee: „Sandwiching" der Log-Likelihood

Anstatt nur eine Schranke zu verwenden, nutzt SPG sowohl eine untere als auch eine obere Schranke der wahren Log-Likelihood, um die Optimierung in Abhängigkeit von der Belohnung zu steuern:

Positive Belohnungen (Advantage $A_j \geq 0$ ): Das Modell maximiert die untere Schranke (ELBO), um die Wahrscheinlichkeit guter Ausgaben zu erhöhen.
Negative Belohnungen (Advantage $A_j < 0$ ): Das Modell minimiert eine obere Schranke (Evidence Upper Bound, EUBO), um die Wahrscheinlichkeit schlechter Ausgaben zu verringern.

Da gilt: $ELBO \leq \log \pi_\theta \leq EUBO$ , stellt die Maximierung dieser „sandwichierten" Zielfunktion einen gültigen Proxy für die Optimierung der wahren Erwartungswert-Belohnung dar.

Theoretische Grundlage: Evidence Upper Bound (EUBO)

Die Autoren leiten eine berechenbare obere Schranke für Masked Diffusion Models her, basierend auf der Rényi-Variationsgrenze. Für einen Parameter $\beta \geq 1$ wird die obere Schranke definiert als:
$\tilde{L}_{EUBO}(x; \theta) = \frac{1}{\beta} \sum_{i=1}^n \log \mathbb{E}_{t, z_t} \left[ w(t) \cdot \mathbb{1}(z_{t,i} = m) \cdot \pi_\theta^\beta(x_i | z_t) \right]$
Im Gegensatz zur ELBO steht der Logarithmus hier außerhalb des Erwartungswerts. Dies führt zwar zu einer leichten Verzerrung bei der Monte-Carlo-Schätzung, bietet aber eine strengere Schranke für die Minimierung negativer Likelihoods.

Praktische Umsetzung: Block-weise Maskierung

Um die Stabilität der Monte-Carlo-Schätzung für diese Schranken zu verbessern, führen die Autoren eine Block-weise Maskierungsstrategie (Block-wise Masking) ein:

Anstatt zufällige Masken über die gesamte Sequenz zu verteilen, wird die Sequenz in Blöcke unterteilt.
Ein zufälliger Block wird ausgewählt, wobei alle vorherigen Blöcke sauber (unmaskiert) und alle folgenden Blöcke vollständig maskiert bleiben. Innerhalb des ausgewählten Blocks werden Tokens zufällig maskiert.
Dies aligniert die Datenverteilung während des Policy-Rollouts besser mit der während der Optimierung, was die Schätzung stabiler macht.

Mischung aus Schranken (Mixture)

Um die Varianz der Gradienten zu reduzieren und die Stabilität zu erhöhen, wird für negative Advantage-Spuren eine Mischung aus ELBO und EUBO verwendet:
$\tilde{L}_{Mix} = \omega \cdot \tilde{L}_{EUBO} + (1-\omega) \cdot L_{ELBO}$
Theoretisch wird gezeigt, dass diese Mischung die Varianz des Gradienten im Vergleich zur Verwendung einer einzelnen Schranke strikt reduziert.

3. Hauptbeiträge

Neuer Policy-Gradient-Algorithmus (SPG): SPG reduziert den Bias in der Optimierung von dLLMs, indem es eine Sandwich-Strategie aus unteren und oberen Schranken basierend auf der Belohnung anwendet.
Block-weise Maskierung: Eine neue Technik zur Monte-Carlo-Schätzung, die die Stabilität des Trainingsziels verbessert und die Effizienz der Policy-Optimierung steigert.
State-of-the-Art Ergebnisse: SPG erzielt auf vier mathematischen und logischen Reasoning-Benchmarks sowie zwei Coding-Benchmarks die besten Ergebnisse unter RL-Algorithmen für Diffusions-Sprachmodelle.

4. Ergebnisse

Die Experimente wurden auf dem Basis-Modell LLaDA-8B-Instruct durchgeführt und verglichen mit SOTA-Baselines wie D1, WD1 und UniGRPO.

Leistung auf Reasoning-Benchmarks (bei Sequenzlänge 256):

GSM8K: SPG verbessert die Genauigkeit um 3,6 % gegenüber dem vorherigen SOTA.
MATH500: Steigerung um 2,6 %.
Countdown: Massive Verbesserung um 18,4 %.
Sudoku: Deutliche Steigerung um 27,0 %.

Coding-Aufgaben (HumanEval & MBPP):

SPG erreichte ebenfalls Verbesserungen von 1,9 % (HumanEval) und 4,7 % (MBPP).

Analyse:

Die Kombination aus ELBO und EUBO (SPG w/ Mixture) erzielte die besten Ergebnisse, was die theoretische Annahme bestätigt, dass beide Schranken für unterschiedliche Optimierungsrichtungen (Maximierung vs. Minimierung) notwendig sind.
Die Block-weise Maskierung führte zu signifikanten Verbesserungen gegenüber zufälliger Maskierung, insbesondere bei komplexen Aufgaben wie Countdown.
SPG zeigte eine schnellere Konvergenz und höhere Stabilität im Reward-Verlauf während des Trainings im Vergleich zu Baselines.

5. Bedeutung und Fazit

Das Paper adressiert ein fundamentales theoretisches und praktisches Problem bei der Anwendung von Reinforcement Learning auf Diffusions-Sprachmodelle: die Intractability der Log-Likelihood.

Theoretischer Durchbruch: Die Einführung einer gültigen oberen Schranke (EUBO) für die Minimierung negativer Belohnungen schließt eine Lücke in der Theorie, die bisherige Methoden (die nur ELBO nutzten) offenließen. Dies ermöglicht es Diffusionsmodellen, effektiv aus negativen Feedback-Signalen zu lernen.
Praktische Relevanz: SPG demonstriert, dass Diffusionsmodelle nicht nur für die Generierung, sondern auch für anspruchsvolle Reasoning-Aufgaben durch RL optimiert werden können, wobei sie autoregressive Modelle in bestimmten Szenarien übertreffen oder zumindest konkurrenzfähig sind.
Zukunftsperspektive: Die Arbeit legt den Grundstein für stabilere und effizientere RL-Verfahren für dLLMs und zeigt, dass die Kombination von variationalen Schranken mit modernen Maskierungsstrategien ein vielversprechender Weg ist, um die Leistungsfähigkeit dieser neuen Modellklasse voll auszuschöpfen.