Stabilizing Reinforcement Learning for Diffusion Language Models
Abstract
Diffusion Large Language Models (dLLMs) often exhibit severe instability during Group Relative Policy Optimization (GRPO) training, limiting the effectiveness of reinforcement learning for improving reasoning capabilities. In dLLMs, the importance ratios used by GRPO are derived from finite-sample estimates rather than exact likelihoods, making them inherently noisy. In this paper, we show that GRPO is highly sensitive to this noise, which drives training instability. Through theoretical analysis and empirical evidence, we identify a self-reinforcing instability loop in which noisy importance ratios induce gradient spikes and policy drift, further amplifying future importance ratio estimation variance. To address this issue, we propose StableDRL, a novel reinforcement learning framework for dLLMs. StableDRL stabilizes training via (i) unconditional clipping to suppress outlier-induced gradient spikes, and (ii) self-normalization to constrain gradients within the convex hull of per-sample updates. We further extend StableDRL to block-wise diffusion models via a staircase attention mechanism. StableDRL is the first method that enables stable, full-parameter reinforcement learning for dLLMs. It achieves the state-of-the-art performance, outperforming prior best full-attention baselines by 6% on MATH500 and block-diffusion baselines by 25.6% on AIME.
Lay Summary
This paper studies how to train a newer kind of AI language model more reliably. These models can generate text in a more flexible way than standard left-to-right systems, but when we try to improve them through trial-and-error training, the training process can suddenly become unstable and lose progress. We identify one major cause: the training signal is based on noisy estimates, and a few unusually noisy examples can push the model too far in the wrong direction. We propose StableDRL, a training method that limits the effect of these noisy examples and keeps each update better balanced. In experiments on math, logic, and puzzle-solving tasks, StableDRL makes training much more stable and improves performance over prior methods. More reliable training can reduce wasted computation and may help future language models become easier to improve, evaluate, and align with human goals.