Stabilizing Reinforcement Learning for Diffusion Language Models

TL;DR

Standard GRPO collapses when applied to diffusion LLMs because the importance ratio it depends on is intractable in dLLMs and must be estimated (ELBO / mean-field), producing heavy-tailed noisy ratios. The paper diagnoses a self-reinforcing instability loop — noisy ratios → gradient spikes → policy drift → larger ratio variance — and proposes StableDRL: replace GRPO’s conditional clip with an unconditional two-sided clip, and replace fixed-group-size normalization with self-normalization by the sum of clipped ratios (constraining each update to the convex hull of per-sample directions). A staircase attention extension enables the same recipe for block-diffusion LLMs (SDAR). StableDRL is the first method reported to enable stable, full-parameter RL on dLLMs for >1000 steps, and reaches new SOTA on GSM8K / MATH500 / Countdown / Sudoku for LLaDA-8B and on AIME24 for SDAR-8B.

Key claims

GRPO’s instability on dLLMs is decomposable into three stages: (i) ELBO/mean-field importance-ratio estimation has heavy-tailed noise that can drive individual rollout ratios to magnitudes of ~e^7 [§3.1, Fig. 8]; (ii) conditional clipping is bypassed when noise crosses the r > 1+ε threshold under negative advantage, producing unclipped gradient spikes, and group-size normalization amplifies group-level co-fluctuation [§3.1]; (iii) the spikes induce policy drift Δ_t which acts as a multiplicative amplifier of next-step ratio variance [§3.1, Eq. 3].
Lemma 3.1 lower-bounds the per-step spike probability by a non-decreasing function of the drift state; Theorem 3.2 shows that once a spike occurs, the bound on the next-step spike probability cannot decrease — i.e. the loop is provably self-reinforcing [§3.3].
Unconditional clipping alone is insufficient: it merely trades unbounded spikes for high-frequency boundary-hitting which corrupts AdamW momentum and still leads to reward collapse [Lemma 3.3, Theorem 3.4, Fig. 4].
Self-normalization by Σ clip(r) constrains each update to the convex hull of per-sample gradient directions, removing the remaining group-scale randomness [Theorem 3.5, Eq. 10].
For block diffusion (SDAR), naive parallel-attention leakage breaks the ELBO proxy; staircase attention uses a dual-stream input (frozen clean context + corrupted target) with block-lower-triangular masking across blocks and block-diagonal within, giving leakage-free single-pass ELBO estimation [§3.4, Fig. 3].
A numerically stable “clip-then-softmax” log-space implementation is required to avoid under/overflow under mixed-precision training [§3.5].
StableDRL is the first method to enable stable full-parameter RL on both full-attention and block dLLMs for >1000 steps [§1].
Full-attention results on LLaDA-8B-Instruct: average across 128/256/512 lengths — GSM8K 84.2% (vs SPG 83.0, ESPO 82.0), MATH500 41.8% (vs 38.4 / 39.5), Countdown 83.5% (vs 69.9 / 81.0), Sudoku 91.5% (vs 90.0 / 86.0) [Table 1].
Block-diffusion results on SDAR-8B (static sampling): MATH500 79.2 (base 75.4, Trado 78.5), GSM8K 92.4 (base 91.1, Trado 92.3), AIME 2024 16.7 vs base 11.8, Trado 13.3, and AR Qwen3-8B 10.0 — block-diffusion RL beating an AR baseline of the same size [Table 2].
Although RL is run at seq length 256, the trained model preserves quality across 128–512 tokens, suggesting stable full-parameter RL gives better length generalization than LoRA/early-stopping mitigations used by ESPO and SPG [§4.2.1].
A “Relative Gradient Spike Rate” metric (gradient norm exceeds a 32-step moving average by margin m=3) empirically tracks the three regimes: GRPO has rising spike threshold and divergent rewards, unconditional-clipping alone has saturation at the upper bound (oscillatory reward), StableDRL has low/stable spike threshold and monotonic reward improvement [§4.1, Fig. 4].

Method

StableDRL is a reformulation of GRPO targeted at the noisy-importance-ratio regime. Given a group of rollouts with estimated log-ratios ℓ_i = log r_i = Δ_i + ξ_i (where Δ is the true policy drift and ξ is ELBO/mean-field estimation noise), StableDRL (i) applies an unconditional two-sided clip w_i = clip(exp(ℓ_i), 1-ε, 1+ε) regardless of advantage sign, removing the asymmetric “unclipped on negative-advantage” path that GRPO uses to accelerate trust-region return; and (ii) normalizes the policy-gradient surrogate by Σ_i w_i instead of the fixed group size G. The update direction is therefore a convex combination of per-sample directions weighted by clipped importance, bounded in magnitude by the constant per-sample gradient bound. In log-space implementation, w_i / Σ_j w_j is computed via a stable softmax over clipped log-ratios to preserve precision under mixed-precision training.

For block-diffusion models the importance-ratio proxy needs a per-block ELBO that conditions only on clean preceding blocks (not on the current block’s ground truth). Staircase attention encodes this with a dual-stream input: a frozen clean stream and a corrupted target stream sharing positional embeddings. A block-lower-triangular “staircase” mask grants target tokens in block i access to the clean history of blocks <i, while a block-diagonal mask permits parallel independent denoising within block i. This satisfies the strict-causality requirement of the ELBO at training time in a single forward pass (O(n²) vs the O(n_blocks·n²) of the naive iterative implementation).

Results

On LLaDA-8B-Instruct full-attention, StableDRL averages 84.2 / 41.8 / 83.5 / 91.5 on GSM8K / MATH500 / Countdown / Sudoku across decoding lengths 128/256/512, beating the prior best baselines SPG (83.0 / 38.4 / 69.9 / 90.0) and ESPO (82.0 / 39.5 / 81.0 / 86.0). Largest single-cell margin is +13.7 absolute on Countdown@256 over SPG (70.7 → 84.4) [Table 1]. On SDAR-8B block-diffusion, AIME 2024 reaches 16.7 (static) and 13.3 (dynamic), exceeding both Trado (13.3 / 11.0) and the autoregressive Qwen3-8B (10.0); MATH500 reaches 79.2 [Table 2]. Diagnostically, GRPO’s Relative Gradient Spike Rate threshold drifts upward as training proceeds and reward eventually collapses; unconditional-clipping alone saturates at the upper bound and still collapses; StableDRL maintains a low stable spike threshold and shows monotonic reward improvement [Fig. 4].

Why it’s interesting

Two reasons this matters for anyone building reasoning-RL recipes. First, it identifies a specific structural mismatch between GRPO and dLLMs — GRPO’s asymmetric clip was designed to “accelerate return to trust region” under an exact ratio, and that design becomes a footgun the moment the ratio is a noisy proxy. Many filed reasoning-RL papers on this wiki (see Reasoning RL) treat GRPO as a drop-in primitive across model classes; StableDRL is the cleanest argument so far that the primitive does not survive the move from exact to proxy likelihoods. Second, the convex-hull self-normalization construction is portable: any RL setting where importance ratios are estimated (e.g. continuous-action policies with score-function surrogates, or video-diffusion RL where likelihoods are similarly intractable — cf. Human detectors are surprisingly powerful reward models) inherits the same long-tailed-ratio failure mode, and the “clip then convex-combine” fix should generalize directly. The complementarity to WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference is also worth noting: WeDLM makes dLLM inference prefix-cache compatible with AR engines, while StableDRL makes dLLM RL training stable — two pieces of the same “make dLLMs production-grade” puzzle, both shipped in the same quarter.