cs.LGMar 6, 2026
SaveStabilizing Reinforcement Learning for Diffusion Language Models
Organizations: Huawei Foundation Model Department · The Chinese University of Hong Kong · The Hong Kong University of Science and Technology
Abstract
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.
Figures & tables
Figure 1: StableDRL is the first method to enable stable full-parameter RL training on both full-attention and block dLLMs, better unlocking reasoning capability for dLLMs. The left panel reports performance on full-attention dLLMs (LLaDA-8B Nie et al. (2025a) ). Based on Table 1, Best Prior corresponds to WD1 Tang et al. (2025) , and Best SOTA corresponds to the best performance between ESPO and SPG Ou et al. (2025a) ; Wang and others (2025) for each task. The right panel demonstrates results for block diffusion models (SDAR-8B Cheng et al. (2025) ).
Figure 2: (a) Training instability. Naive integration of noisy importance ratios into GRPO leads to severe instability under full-parameter RL training with dLLMs. Notably, reward collapse occurs even with Policy Gradient, where the importance ratio is fixed at 1. (b) Instability loop. Estimation noise triggers gradient spikes and policy drift, creating a self-reinforcing cycle that amplifies the variance of future importance ratios. (c) StableDRL. To address this, we propose a reformulated GRPO for noisy importance ratios. By employing unconditional clipping and self-normalization, StableDRL effectively breaks the instability loop.
Figure 3: Staircase Attention for Efficient Proxy Estimation. To evaluate the ELBO for block diffusion in a single pass ( ), we use a dual-stream construction. The Clean Context (top rows) provides immutable history. The Corrupted Target stream (bottom rows) uses a “staircase” mask ( , bottom-left) to attend to valid history without peeking at the ground truth of the current block. The target self-attention ( , bottom-right) is block-diagonal, ensuring independent parallel denoising.
Figure 4: Verification of Instability Mechanisms across Methods. Left Column (GRPO): Unbounded drift (bottom) fuels an accelerating spike rate (middle), causing reward collapse (top). Middle Column (Unconditional Clipping): Clipping saturates the drift (bottom) but induces a high-frequency, stochastic spike regime (middle) that destabilizes learning (top). Right Column (StableDRL): Our method maintains a low, stable spike rate (middle) decoupled from drift (bottom), resulting in monotonic reward improvement (top).
| GSM8K | MATH500 | Countdown | Sudoku | |||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Model / Seq Len | 128 | 256 | 512 | Avg | 128 | 256 | 512 | Avg | 128 | 256 | 512 | Avg | 128 | 256 | 512 | Avg |
| LLaDA-8B-Inst. | 69.5 | 77.2 | 79.8 | 75.5 | 28.2 | 32.4 | 34.6 | 31.7 | 18.8 | 16.8 | 16.8 | 17.5 | 5.7 | 27.7 | 26.2 | 19.9 |
| LLaDA-1.5 | 70.4 | 80.5 | 81.9 | 77.6 | 26.8 | 32.2 | 35.8 | 31.6 | 31.9 | 21.1 | 21.5 | 24.8 | 7.4 | 26.9 | 29.0 | 21.1 |
| D1 | 72.2 | 80.6 | 81.3 | 78.0 | 31.4 | 36.0 | 39.4 | 35.6 | 30.9 | 30.9 | 34.4 | 32.1 | 7.2 | 32.5 | 29.3 | 23.0 |
| WD1 | 74.6 | 81.5 | 83.0 | 79.7 | 31.0 | 37.4 | 39.0 | 35.8 | 48.8 | 52.3 | 50.8 | 50.6 | 33.1 | 32.1 | 22.5 | 29.2 |
| UniGRPO | 74.9 | 82.5 | 82.7 | 80.0 | 32.4 | 37.4 | 39.4 | 36.4 | 44.5 | 43.0 | 57.0 | 48.2 | 59.0 | 67.0 | 62.9 | 63.0 |
Table 1: State-of-the-Art Reasoning Performance on LLaDA-8B-Instruct. We report pass@1 accuracy under three decoding budgets ( ) and the mean performance ( Avg ) for each dataset. Bold denote the best result and underline the second best. StableDRL achieves the highest average accuracy on all four benchmarks, demonstrating superior consistency across generation lengths.
| Method | MATH500 | GSM8K | AIME 24 |
|---|---|---|---|
| Autoregressive (AR) Baselines | |||
| Qwen3-4B | 74.1 | 90.7 | 12.9 |
| Qwen3-8B | 78.4 | 92.8 | 10.0 |
| Block Diffusion (Dynamic Sampling) | |||
| SDAR-8B (Base) | 70.6 | 90.4 | 0 8.3 |
| Trado | 75.0 | 91.2 | 11.0 |
Table 2: Block Diffusion Reasoning Performance. Pass@1 accuracy on MATH500, GSM8K, and AIME, comparing StableDRL (SDAR-8B backbone) against AR baselines (Qwen3) and prior Block Diffusion methods. StableDRL notably outperforms the strong Qwen3-8B AR model on the rigorous AIME benchmark.
Figure 5: Robustness to Proxy Noise: The “Exploding Weight” Stress Test (GSM8K). We compare training stability under standard conditions (“Normal”, solid lines) versus an adversarial regime where importance weight variance is artificially amplified (“Exploding”, dashed lines; see App. C.2 ). (Left) Reward Trajectories: StableDRL ( Green ) demonstrates invariant stability , maintaining monotonic improvement in both regimes. In contrast, ESPO ( Orange ) suffers immediate, noise-accelerated collapse , confirming its sensitivity to ratio outliers. SPG ( Blue ) degrades in both settings, indicating that avoiding ratios (to reduce variance) fatally exposes the model to off-policy bias. (Right) Gradient Norm Density: Visualizing the failure mechanism. StableDRL maintains a condensed, low-variance gradient distribution. Conversely, ESPO exhibits a heavy right tail of explosive updates (log-norm ), confirming that the “Asymmetric Clipping Failure” allows noise spikes to propagate unchecked.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
| Category | Value |
|---|---|
| Model & Initialization | |
| Base Model | LLaDA-8B-Instruct |
| Precision | bfloat16 |
| Activation Checkpointing | Whole Layer |
| Generation (Rollout) | |
| Decoding Strategy | Iterative (128 steps) |
Table 3: Hyperparameter Configuration for Full-Attention Diffusion (LLaDA-8B-Instruct)
| Category | Value |
| Model & Initialization | |
| Base Model | JetLM/SDAR-8B-Chat |
| Architecture | Block Diffusion ( ) |
| Precision | bf16 (TF32 enabled) |
| Generation (Rollout) | |
| Sampling Strategy | Dynamic ( ) |
Table 4: Hyperparameter Configuration for Block Diffusion (SDAR-8B-Chat)
Figure 8: Diagnosing Gradient Instability in dLLM Training. We visualize the joint distribution of importance weights ( ) and gradient norms ( ) during training. (Left) ESPO: Exhibits a characteristic “chimney” failure where rare, noise-induced outliers bypass clipping on negative advantages, acting as unbounded step-size multipliers that drive gradients to explosion ( ). (Center) SPG-IS: Despite avoiding explicit ratios, the update variance is strongly correlated with policy drift, confirming that rollout-reuse bias accumulates to destabilize training. (Right) StableDRL (Ours): By enforcing strict clipping and self-normalization, our method decouples update magnitude from proxy noise, confining gradients to the convex hull of the samples (Theorem 3.1) and maintaining deterministic stability.