Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Stabilizing Reinforcement Learning for Diffusion Language Models

About

Group Relative Policy Optimization (GRPO) is highly effective for post-training autoregressive (AR) language models, yet its direct application to diffusion large language models (dLLMs) often triggers reward collapse. We identify two sources of incompatibility. First, GRPO relies on importance ratios defined by sequence probabilities, which are intractable in dLLMs and must be estimated (e.g., via ELBO-based or mean-field likelihood proxies), yielding inherently noisy ratios. Second, standard GRPO's formulation is not designed for estimated ratios: its conditional clipping can be anomalously bypassed by model-agnostic estimation noise, producing gradient spikes, while its fixed group-size normalization amplifies gradient-magnitude fluctuations under high-variance ratio estimates. We show these effects form a self-reinforcing instability loop that drives policy drift and further increases ratio variance. To break this loop, we propose StableDRL, a reformulation of GRPO tailored for dLLMs that uses (i) unconditional clipping to suppress outlier-induced spikes and (ii) self-normalization to constrain updates within the convex hull of per-sample gradients. We further extend StableDRL to block-wise diffusion models via a staircase attention mechanism.

Jianyuan Zhong, Kaibo Wang, Ding Ding, Zijin Feng, Haoli Bai, Yang Xiang, Jiacheng Sun, Qiang Xu• 2026

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningGSM8K--
1398
Mathematical ReasoningGSM8K--
499
Mathematical ReasoningCountdown
Accuracy62.11
252
Mathematical ReasoningMATH 500
pass@179.2
239
Code GenerationHumanEval
Accuracy43.29
217
Mathematical ReasoningGSM8K
GSM8K Accuracy (%)81.35
204
Mathematical ReasoningAIME 24
Pass@1 Accuracy16.7
93
Mathematical ReasoningMATH 500
Pass@1 Rate45.2
82
Code GenerationKodCode
Accuracy28.2
58
ReasoningSudoku
Pass@192.4
26
Showing 10 of 16 rows

Other info

GitHub

Follow for update