SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

Titel: SPG: Sandwiched Policy Gradient voor Masked Diffusion Language Models

Conferentie: ICLR 2026
Auteurs: Chenyu Wang et al. (Meta Superintelligence Labs, MIT, USC, UCLA)

---

1. Het Probleem

Diffusion Large Language Models (dLLMs) zijn een opkomend alternatief voor autoregressieve modellen omdat ze meerdere tokens parallel kunnen decoderen, wat de inferentielatentie aanzienlijk verlaagt. Een cruciale stap om deze modellen te aligneren met menselijke voorkeuren of taakspecifieke beloningen (zoals redeneervermogen) is Reinforcement Learning (RL).

De kernuitdaging bij het toepassen van RL op dLLMs is de onberekenbaarheid (intractability) van de log-likelihood ($\log \pi_\theta(x|c)$).

• Standaard RL-algoritmen (zoals Policy Gradient) vereisen de exacte log-likelihood om de gradiënt te schatten.
• Bestaande werken omzeilen dit probleem door de Evidence Lower Bound (ELBO) te gebruiken als vervanging voor de ware log-likelihood.
• Het tekortkoming: De ELBO is slechts een ondergrens. Het maximaliseren van de ELBO werkt goed voor positieve beloningen, maar het minimaliseren van de ELBO voor negatieve beloningen (straf voor slechte outputs) garandeert niet dat de ware log-likelihood daalt. Dit leidt tot een vooringenomen (biased) policy gradient, waardoor het model niet effectief leert van negatieve feedback en suboptimale prestaties levert.

2. Methodologie: Sandwiched Policy Gradient (SPG)

De auteurs stellen SPG voor, een nieuw RL-algoritme dat de vooringenomenheid oplost door een "sandwich"-benadering te gebruiken die zowel een onder- als een bovengrens van de log-likelihood benut, afhankelijk van de aard van de beloning.

A. Het Sandwich-Principe

In plaats van alleen de ELBO te gebruiken, past SPG een adaptieve strategie toe op basis van de advantage ($A_j$) van een gegenereerde respons:

1. Positieve Advantage ($A_j \geq 0$): Voor goede antwoorden wordt de Evidence Lower Bound (ELBO) gemaximaliseerd. Dit is een geldige benadering om de waarschijnlijkheid van goede sequenties te verhogen.
2. Negatieve Advantage ($A_j < 0$): Voor slechte antwoorden wordt een Evidence Upper Bound (EUBO) geminimaliseerd.
   • De auteurs leiden een nieuwe bovengrens af op basis van de Rényi variational bound.
   • Door de EUBO te minimaliseren voor negatieve voorbeelden, garandeert het algoritme dat de ware log-likelihood daalt, wat een correcte strafmechanisme biedt.

De doelstelling wordt hierdoor een "sandwich" van de ware log-likelihood:
$$J_{SPG}(\theta) = \mathbb{E} \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \geq 0} \cdot A_j \mathcal{L}_{ELBO} + \mathbb{1}_{A_j < 0} \cdot A_j \mathcal{L}_{EUBO} \right) \right]$$

B. Praktische Implementatie en Stabiliteit

Om de EUBO en ELBO stabiel te schatten via Monte Carlo-sampling, introduceert het paper twee belangrijke technieken:

1. Block-wise Masking Strategie:
   • In plaats van willekeurige masking (zoals bij standaard training), wordt de sequentie opgedeeld in blokken.
   • Tijdens de schatting wordt een willekeurig blok geselecteerd: eerdere blokken blijven schoon, latere blokken zijn volledig gemaskeerd, en het geselecteerde blok wordt willekeurig gemaskeerd.
   • Dit sluit beter aan bij de data-distributie die het model tegenkomt tijdens de inferentie (semi-autoregressief), wat de stabiliteit van de gradiënt schatting verbetert.
2. Mixtuur van Grenzen (Mixture):
   • De pure EUBO kan computatievriendelijk zijn maar heeft een hoge variantie in de gradiënt.
   • De auteurs combineren de EUBO en ELBO in een mixtuur-objectief ($\tilde{L}_{Mix}$) voor negatieve advantage traces.
   • Theoretisch bewijs (Propositie 1): Een optimale mixtuur van de boven- en ondergrens resulteert in een strikt lagere variantie van de gradiënt dan het gebruik van alleen de boven- of ondergrens. Dit zorgt voor stabielere training.

3. Belangrijkste Bijdragen

1. SPG Algorithm: Een nieuw RL-algoritme dat de vooringenomenheid van bestaande methoden oplost door een "sandwich" van variational bounds te gebruiken, afhankelijk van de beloningsteken.
2. Evidence Upper Bound (EUBO): Een nieuwe, afgeleide bovengrens voor de log-likelihood van masked diffusion modellen, specifiek ontworpen om negatieve feedback effectief te straffen.
3. Block-wise Masking: Een strategie die de distributie van de input tijdens de policy rollout en de optimalisatie beter aligneert, wat leidt tot stabielere schattingen.
4. Theoretische Analyse: Bewijs dat het mixen van boven- en ondergrenzen de gradiëntvariantie verlaagt en de training stabiliseert.

4. Resultaten

SPG werd getest op vier wiskundige en logische redeneringsbenchmarks en twee coderingsbenchmarks, vergeleken met state-of-the-art (SOTA) baselines zoals D1, WD1, UniGRPO en LLaDA-1.5.

Resultaten op Redeneringsbenchmarks (bij een generatielengte van 256):

• GSM8K: +3.6% verbetering ten opzichte van de vorige SOTA.
• MATH500: +2.6% verbetering.
• Countdown: +18.4% verbetering (een enorme sprong).
• Sudoku: +27.0% verbetering (van ~27% naar ~94% nauwkeurigheid).
• Codering (HumanEval & MBPP): SPG verbeterde de nauwkeurigheid met respectievelijk 1.9% en 4.7%.

Observaties:

• SPG convergeert sneller en bereikt een hoger beloningsniveau dan baselines.
• De methode is robuust over verschillende inferentiestrategieën (bijv. verschillende blokgroottes en decoding-methoden).
• Het gebruik van de mixtuur (ELBO + EUBO) voor negatieve advantage traces bleek superieur aan het gebruik van alleen ELBO of alleen EUBO.

5. Significatie

Dit paper is significant omdat het een fundamenteel theoretisch en praktisch probleem in het veld van Diffusion Language Models oplost.

• Het toont aan dat het simpelweg toepassen van autoregressieve RL-methoden (zoals GRPO/DPO) op diffusion modellen niet werkt vanwege de onberekenbare log-likelihood.
• Het biedt een wiskundig onderbouwde oplossing die het potentieel van dLLMs voor redeneringstaken volledig vrijmaakt.
• De prestaties, vooral op complexe logica- en puzzeltaken zoals Sudoku en Countdown, tonen aan dat dLLMs met SPG concurrerend kunnen zijn met, of zelfs beter kunnen presteren dan, geavanceerde autoregressieve modellen op specifieke taken, terwijl ze profiteren van de parallelle inferentie van diffusion.

Kortom, SPG maakt het mogelijk om diffusion modellen effectief te trainen met Reinforcement Learning, waardoor ze een krachtig alternatief worden voor autoregressieve modellen in scenario's waar snelheid en parallelle verwerking essentieel zijn.