PaperlayerView exact revision on arXiv

Source-linked AI summary

Stabilizing Reinforcement Learning for Diffusion Language Models

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

arXiv:2603.06743v1cs.LGcs.AI

TL;DR

GRPO suffers reward collapse on dLLMs because intractable sequence probabilities require noisy importance-ratio estimates, while its conditional clipping and fixed normalization amplify instability. StableDRL reformulates GRPO with unconditional clipping and self-normalization, extends it to block diffusion, and reports stable training and improved reasoning performance. The paper also identifies boundary saturation as a limitation of clipping alone.

  • Problem

    GRPO is effective for autoregressive models but becomes unstable on dLLMs because sequence-probability importance ratios are intractable and must be estimated noisily.

  • Method

    StableDRL combines unconditional clipping and self-normalization, and extends the approach to block diffusion through staircase attention.

  • Results

    StableDRL enables stable full-parameter reinforcement-learning training on full-attention and block dLLMs for over 1,000 steps and achieves state-of-the-art reasoning performance.

  • Takeaways & Limitations

    The proposed reformulation stabilizes dLLM reinforcement-learning training and helps unlock their reasoning potential.

  • Takeaways & Limitations

    Clipping alone can trade unbounded spikes for frequent boundary-saturated updates, while tight bounds may obscure importance-weight signals.

Abstract

from arXiv · show

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.

1 Introduction

GRPO becomes unstable on dLLMs because importance ratios are estimated noisily and standard clipping and normalization do not accommodate that noise. StableDRL breaks the resulting instability loop with unconditional clipping, self-normalization, and a block-diffusion extension.

  • Full-parameter GRPO training on dLLMs exhibits abrupt reward collapse at approximately 300 steps.
  • Estimated importance-ratio noise can bypass conditional clipping, causing gradient spikes, while fixed group-size normalization produces large gradient-magnitude fluctuations.
  • Noisy ratios trigger gradient spikes, policy drift, and higher future ratio variance, forming a self-reinforcing loop that leads to reward collapse.The paper states that exceeding a critical gradient-norm threshold increases the probability of continued divergence.
  • StableDRL uses unconditional clipping to bound importance ratios and prevent outlier-induced gradient spikes.
  • StableDRL uses self-normalization by the sum of clipped importance ratios, constraining updates within the convex hull of per-sample gradients.
  • StableDRL enables stable full-parameter RL training on both full-attention and block dLLMs for over 1,000 steps.Experiments report higher training stability and significant accuracy gains over prior best-in-class methods.

2 Background

dLLMs generate text through discrete diffusion and require likelihood proxies because exact sequence likelihoods are intractable. GRPO uses group-relative advantages and importance sampling, but adapting it depends on noisy ELBO-based estimates.

  • 2.1 Masked Diffusion Language Models: Masked diffusion language models corrupt clean sequences into mask tokens and learn a denoiser to reconstruct the original data.
  • 2.1 Masked Diffusion Language Models: Exact MDLM sequence log-likelihood is intractable because it requires marginalizing over all n! masking trajectories.Training therefore optimizes an Evidence Lower Bound (ELBO).
  • 2.1 Masked Diffusion Language Models: The ELBO expectation is approximated in practice with Monte Carlo sampling, producing a single-sample estimator when one sample is used.
  • 2.2 Reinforcement Learning with dLLMs: GRPO removes the value-function critic by estimating the baseline from the group-average reward.
  • 2.2 Reinforcement Learning with dLLMs: GRPO updates the target policy with standardized group advantages and importance ratios, while clipping ratios to a trust region [1 − ε, 1 + ε].
  • 2.2 Reinforcement Learning with dLLMs: Prior dLLM approaches moved from inaccurate one-step mean-field approximations toward multi-step Monte Carlo ELBO estimates, which remain noisy with limited sampling.

3 Methodology

The methodology diagnoses a self-reinforcing instability loop in dLLM GRPO and proposes StableDRL to interrupt it with unconditional clipping and self-normalization. It also introduces staircase attention for efficient, leakage-free block-diffusion likelihood-proxy estimation.

  • Instability diagnosis: Estimated importance-ratio noise forms a long-tailed distribution, producing high variance and extreme outliers in dLLM training.Exponentiating estimation noise can make a single rollout’s ratio reach magnitudes of 10^5.
  • Instability diagnosis: Conditional GRPO clipping can leave noise-induced, negative-advantage outliers unclipped, generating massive gradients unrelated to true policy alignment.When A < 0 and the ratio exceeds 1 + ϵ, the objective uses the unclipped ratio term.
  • Instability diagnosis: Gradient spikes cause policy drift, which increases the variance and future spike probability of estimated importance ratios, creating a self-reinforcing instability loop.The analysis shows that drift-dependent spike bounds are nondecreasing after spike-induced updates.
  • StableDRL: StableDRL replaces conditional clipping with unconditional bounds on importance ratios, preventing extreme outliers from producing unbounded gradients.The ratio is constrained to [1 − ϵ, 1 + ϵ] regardless of advantage sign.
  • StableDRL: StableDRL replaces fixed group-size normalization with normalization by the sum of clipped ratios, confining updates to the convex hull of per-sample gradients.This removes group-level weight-scale fluctuations, while clipping alone can still produce boundary-saturated updates.
  • Block diffusion extension: For block diffusion, staircase attention uses clean history and corrupted targets to enable single-pass, leakage-free proxy estimation with parallel within-block denoising.The method avoids O(K) iterative evaluation while preventing target tokens from attending to their current block’s ground truth.

4 Experiments

Experiments evaluate StableDRL across full-attention and block diffusion models, diagnose instability mechanisms, and compare its training stability and reasoning performance with prior methods. StableDRL maintains stable learning under noisy ratios and achieves strong accuracy across tasks, decoding lengths, and architectures.

  • Experimental setup: Experiments cover LLaDA-8B-Instruct and SDAR-8B using standard benchmarks, multiple generation lengths, and ablations of clipping and self-normalization.The evaluation includes GSM8K, MATH500, Countdown, and Sudoku, with generation lengths 128, 256, and 512.
  • Instability mechanisms: GRPO develops unbounded importance-ratio drift and rising spike thresholds, while unconditional clipping causes saturation and oscillatory dynamics that destabilize rewards.StableDRL instead maintains a low, stable spike threshold through a structural convex-hull constraint and smooth reward improvement.
  • Full-attention diffusion results: 41.8% average accuracy on MATH500 and +5.2% over SPG at 256 tokens accompany StableDRL’s highest average accuracy across all four LLaDA-8B-Instruct benchmarks.StableDRL also reaches 84.4% on Countdown at 256 tokens, a +13.7% gain, and averages 84.2% on GSM8K and 91.5% on Sudoku.
  • Block diffusion results: 16.7% on AIME 2024 surpasses the base model, Trado, and Qwen3-8B, while StableDRL remains more robust than Trado under dynamic sampling.The block-diffusion evaluation reports 13.3% under dynamic sampling for StableDRL versus 11.0% for Trado.
  • Robustness to proxy noise: StableDRL maintains monotonic improvement under both normal and artificially amplified importance-weight variance, whereas ESPO collapses immediately under exploding weights and SPG degrades in both regimes.The stress test also finds StableDRL’s gradient distribution concentrated at low variance, while ESPO has a heavy right tail of explosive updates.
  • Ablation studies: Removing clipping causes rapid failure, while removing self-normalization produces oscillating update magnitudes and eventual reward collapse; only their combination is stable and monotonic.The threshold study finds stability for ϵ ∈{1, 5}, with ϵ = 5 offering a better exploration-stability trade-off, but catastrophic collapse for ϵ ∈{100, 1000}.

5 Related Work

Related work positions StableDRL amid reinforcement learning for autoregressive and diffusion language models, importance-sampling robustness, and efforts to adapt GRPO to intractable diffusion likelihoods.

  • RL post-training for LMs: Policy-gradient reinforcement learning supports language-model alignment and reasoning, while autoregressive methods use importance-ratio correction to address behavior-target policy staleness.The related work discusses RLHF and reinforcement learning with verifiable rewards before describing autoregressive importance-ratio correction.
  • RL for diffusion LMs: Diffusion-language-model approaches include off-policy RL with smaller performance gains and Markov-process formulations that are computationally expensive and difficult to scale.The passage contrasts these methods with stronger on-policy approaches for diffusion language models.
  • Importance-sampling robustness and off-policy stabilization: Importance-weight variance and heavy tails motivate truncation, clipping, diagnostics, and smoothing in Monte Carlo and off-policy estimation, including methods such as V-trace and Retrace.These techniques provide context for controlling extreme importance weights in reinforcement-learning updates.

6 Conclusion

The paper attributes dLLM GRPO instability to noisy importance-ratio estimation and proposes StableDRL with unconditional clipping and self-normalization. It also introduces staircase attention for block diffusion, with experiments reporting sustained stability and improved reasoning performance.

  • GRPO instability in dLLMs arises from noisy Monte Carlo importance-ratio estimates that create a self-reinforcing loop of gradient spikes and policy drift.
  • StableDRL uses unconditional clipping and self-normalization to eliminate gradient spikes and stabilize dLLM reinforcement learning.
  • The block-diffusion extension uses staircase attention to support leakage-free, block-wise reinforcement-learning evaluation.

A.2 The Efficiency-Leakage Dilemma

Block-wise diffusion ELBO estimation faces an efficiency–leakage dilemma: sequential masking is expensive, whereas naive single-pass attention invalidates the variational objective. A dual-stream staircase mask resolves this by enabling parallel computation while preserving block-causal conditioning.

  • The Efficiency-Leakage Dilemma: Exact block-wise ELBO estimation conditions each block only on its clean preceding history.
  • The Efficiency-Leakage Dilemma: Naive iterative evaluation requires K forward passes and makes long-sequence reinforcement learning prohibitively expensive.For sequences with 64 blocks, the training cost grows linearly with the number of blocks.
  • The Efficiency-Leakage Dilemma: Naive single-pass bidirectional attention leaks each block’s ground-truth representations, invalidating the variational bound and gradient signal.
  • Dual-Stream Input and Mask Construction: The dual-stream construction concatenates frozen clean context with a corrupted target stream and applies a composite 2n × 2n attention mask.
  • Dual-Stream Input and Mask Construction: The mask combines causal context attention, zero context-to-target attention, intra-block target attention, and block-lower-triangular history access.
  • Dual-Stream Input and Mask Construction: This staircase construction computes gradients for all K blocks simultaneously while preserving the autoregressive factorization required by the objective.

B.1 Formal theorem statements for Sec. 3.3

The formal analysis characterizes standard GRPO’s drift–spike feedback loop, shows that clipping alone can produce boundary saturation, and establishes self-normalization as a structural remedy. The resulting update lies in the convex hull of per-sample gradients and removes the random group-scale factor.

  • The theorem framework proceeds from the GRPO drift–spike feedback loop to clipping saturation and then self-normalization.
  • The formal setup fixes a behavior policy, a rollout group, estimated importance ratios, group-relative advantages, and the implemented GRPO update direction.
  • Theorem B.1 formalizes how standard GRPO’s noisy ratios can generate gradient spikes and continued divergence under stated conditions.
  • Two-sided clipping deterministically bounds the clipping-only update norm by (1 + ϵ)B, preventing unbounded spikes.
  • As drift grows, clipping alone can replace rare extreme spikes with frequent updates saturated at the upper clipping boundary.
  • Self-normalization forms a convex combination of per-sample directions, so bounded per-sample norms imply a deterministically bounded update norm.

B.2 Proof of Theorem B.1

The proof of Theorem B.1 represents GRPO updates through random effective weights and imposes boundedness, tail, independence, gradient, and moment conditions. It then links drift to ratio exceedance and establishes dominance by a drift-maximizing sample.

  • The proof introduces deterministic proxy gradients and an equivalent effective-weight representation of the implemented GRPO direction.
  • For negative advantages, standard GRPO leaves ratios unclipped from above, permitting large effective weights.
  • The theorem assumes bounded per-sample directions, common unbounded right-tail envelopes, conditional sample independence, nontrivial proxy gradients, and residual effective-weight moment control.
  • The drift-maximizer is assumed to retain a nontrivial proxy gradient, while moment control bounds the expected effective-weight mass of the remaining samples.
  • A large drift-maximizer ratio can dominate the group update with constant probability under the residual moment bound.
  • The ratio-exceedance identity shows that the conditional probability of exceeding a fixed threshold is nondecreasing in drift.

Assume Conditions (C1)–(C5). Define

The appendix establishes conditions under which noisy importance-ratio outliers create gradient spikes, increase drift, and make future spikes more likely. It also analyzes how a dominating sample can decrease its loss and amplify cross-sample drift.

  • Spike-probability bound: Under Conditions (C1)–(C5), a sufficiently large estimated importance ratio can force the GRPO gradient norm above a specified spike threshold.The proof uses a negative-advantage sample, a lower-bounded advantage magnitude and gradient contribution, and a tail event for the estimated ratio.
  • Spike-probability bound: The lower bound on spike probability is nondecreasing in the drift state D_i, linking policy drift to persistent instability.The appendix derives this monotonicity from the tail distribution of the estimated importance ratio.
  • Instability loop: Together, the appendix formalizes a feedback mechanism in which gradient spikes worsen drift, while greater drift raises the probability of later spikes.The monotonic drift increment and spike-probability bound provide the two directions of this loop.
  • One-step decrease: When a negative-advantage sample dominates the update and local smoothness conditions hold, one step decreases that sample’s loss by a quantified amount.The update is represented as a dominant gradient step plus a bounded residual, with the decrease guaranteed under a step-size condition.
  • Cross-sample amplification: Anti-aligned gradients from another sample can convert the dominant update into a cross-sample amplification effect that increases the drift measure D_i.This result requires directional anti-alignment and smoothness assumptions for the second sample.

B.3 Proof of Theorem B.2

Theorem B.2 analyzes two-sided unconditional clipping and shows that clipped updates are deterministically bounded, with spike probabilities controlled by the drift state. The section also records experimental configurations and an adversarial protocol for inflating ratio-estimation variance.

  • Theorem B.2: Two-sided clipping defines weights as w_i,j = clip(ˆρ_i,j, 1 − ϵ, 1 + ϵ), imposing a uniform upper bound on every importance weight.Because estimated ratios are positive, the clipped weights remain positive and are bounded above by 1 + ϵ.
  • Theorem B.2: The clipped group update is deterministically bounded and lies within the convex hull of the per-sample gradients.Positive clipped weights make the normalized update a convex combination, preventing unbounded aggregation from ratio outliers.
  • Theorem B.2: The theorem bounds spike probability using the clipping threshold and shows that this bound is nondecreasing in the drift state D_i.This preserves a drift-dependent probabilistic analysis while removing unbounded ratio amplification.
  • Experimental configurations: Experiments configure StableDRL for full-attention and block diffusion models using self-normalized importance sampling and unconditional clipping.The reported implementations use LLaDA-8B-Instruct and SDAR-8B-Chat with architecture-specific sampling and clipping settings.
  • Experimental protocol: The adversarial protocol stresses 70% of each group by using asymmetric easy masking for the numerator and hard masking for the denominator.The resulting ELBO gap drives estimated ratios toward infinity through estimation variance rather than actual policy probabilities.

C.5 Visual Diagnosis of Gradient Instability

The visual diagnosis compares importance-weight magnitudes with gradient norms for ESPO, SPG-IS, and StableDRL. It shows catastrophic outlier-driven behavior for ESPO, drift-linked variance for SPG-IS, and bounded gradients for StableDRL.

  • Figure 8: Figure 8 plots the joint distribution of log10 ρ and log10 ∥ˆg∥ to compare gradient-instability patterns across three methods.The panels correspond to ESPO, SPG-IS, and StableDRL during training.
  • Interpretation: The comparison supports the proposed instability-loop diagnosis by relating ratio noise, gradient spikes, policy drift, and later variance amplification.The adversarial generation procedure constructs exploding weights by separating numerator and denominator masking policies.
  • ESPO: ESPO exhibits a sparse chimney of noise-induced outliers with ρ > 10^6 and gradient norms above 10^4.These outliers bypass clipping for negative advantages and act as unbounded step-size multipliers.
  • SPG-IS: SPG-IS avoids explicit ratios but retains a strong positive correlation between implicit-weight magnitude and update norm.The visualization attributes this sensitivity to rollout-reuse bias that grows as the target policy drifts from the behavior policy.
  • StableDRL: StableDRL keeps gradient norms below 10^1.8 regardless of importance-weight magnitude.The bounded distribution is attributed to unconditional clipping and self-normalization, which confine updates to the samples’ convex hull.
Loading 2603.06743v1…