Exploring More, Reasoning Better: Stepwise Risk-Sensitive GRPO for Diffusion Language Models

Yue YU; Bowen Zuo; David Crandall; Yinglun Zhu; Dongruo Zhou

cs.LGSep 30, 2026

Exploring More, Reasoning Better: Stepwise Risk-Sensitive GRPO for Diffusion Language Models

Authors: Yue YU, Bowen Zuo, David Crandall, Yinglun Zhu, Dongruo Zhou

Organizations: Indiana University · University of California, Riverside

Abstract

Diffusion large language models (dLLMs) generate text by denoising a sequence or successive blocks, allowing several tokens to be revealed in parallel. Reinforcement learning with verifiable rewards (RLVR) reuses terminal feedback across these decisions, even as their conditioning context changes. We propose stepwise risk-sensitive GRPO (StepRS-GRPO), which varies the risk coefficient of the group-advantage transformation across denoising states while retaining the underlying trainer. For binary rewards, we show that this transformation is exactly a prompt- and state-dependent rescaling of centered outcome advantages. A capability-based calibration suggests a coefficient scale, while endpoint and interpolation ablations guide schedule selection. Across multiple dLLM backbones and mathematical reasoning benchmarks, StepRS-GRPO improves both pass@1 accuracy and pass@k coverage over centered GRPO, while increasing answer diversity. In our ablation studies, mass-matched controls support the contributions of state allocation and schedule direction, and the gains persist after matching the root mean square (RMS) of the advantages to that of centered GRPO. Reasoning-trace diagnostics further show that the diversity gains from StepRS-GRPO extend beyond final-answer strings.

Figures & tables

Figure 1: Pass@

k

on SDAR-1.7B-Chat and SDAR-4B-Chat at the

160

-update checkpoints. Markers show means and capped error bars show one sample standard deviation across inference seeds

42

43

, and

44

, with

64

responses per problem per seed. The selected StepRS schedules are

[1.0\to 0.5]

and

[1.5\to 0.75]

, respectively. Answer diversity is reported in Figure 12 .

Figure 2: SDAR-4B mechanism ablations. Top: StepRS versus prompt-only mass matching (C1) and a mass-normalized reversed schedule (C2). Bottom: GRPO, constant-risk RS-GRPO (

\beta=1.5

), StepRS, and advantage-RMS-matched StepRS (C3). Markers show three-inference-seed means; capped error bars show one sample standard deviation. Controls are defined in the text, with implementation details in Appendix D.1 .

Figure 3: Answer-removed trace diversity on SDAR-4B: (a)

D_{\mathrm{ngram}}

; (b)

D_{\mathrm{emb}}

. Points and capped bars show means and sample standard deviations across three inference seeds, with

64

responses per question per seed. Higher scores indicate greater variation; panel scales differ. Definitions are in Appendix C.3 .

Figure 4: Schedule-shape ablation on SDAR-4B with fixed endpoints

[1.5\to 0.75]

. Linear, cosine, and two-stage schedules are compared at all seven sampling budgets. Markers show means and capped error bars show one sample standard deviation across three inference seeds.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Backbone	Default StepRS schedule	Constant-risk sweep
Qwen3-0.6B-MDLM	$[0.5\to 0.25]$	$\{0.25,0.5,1.0\}$
SDAR-1.7B-Chat	$[1.0\to 0.5]$	$\{0.25,0.5,1.0\}$
SDAR-4B-Chat	$[1.5\to 0.75]$	$\{0.5,1.0,1.5,2.0\}$

Appendix

Table 1: Selected schedules and evaluated constant-risk coefficients.

Figure 5: Qwen3-0.6B-MDLM evaluation curves. The selected

[0.5\to 0.25]

schedule improves pass@

k

over GRPO across GSM8K, MATH500, and AMC23. All curves show means and capped error bars show one sample standard deviation over three seeds.

Hyperparameter	Qwen3-0.6B-MDLM	SDAR-1.7B-Chat	SDAR-4B-Chat
Base model (HuggingFace checkpoint)	dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1	JetLM/SDAR-1.7B-Chat	JetLM/SDAR-4B-Chat
Trainer framework	TraceRL ( Wang et al., 2025c )	StableDRL ( Zhong et al., 2026 )
Training data	GSM8K train	MATH train
Optimizer	AdamW (lr= $5\times 10^{-6}$ )	AdamW (lr= $1\times 10^{-6}$ )
Responses per prompt ( $G$ )	16
Group rollout batch	1024 (64 $\times$ 16)	512 (32 $\times$ 16)

Appendix

Table 2: Training hyperparameters used for the three model scales.

Figure 6: Calibration and endpoint ablations on SDAR-4B. Top: linear half-decay schedules with upper endpoints

1.0

1.08

1.5

, and

2.0

. Bottom: full decay versus half-decay at a fixed upper endpoint of

2.0

. Columns correspond to the four benchmarks. Curves show means and one-standard-deviation bars across three inference seeds.

Figure 7: The selected

[1.5\to 0.75]

StepRS schedule versus constant risk coefficients on SDAR-4B. Each panel shows the complete pass@

1

–pass@

64

profile for one benchmark, with mean and one-standard-deviation bars over three inference seeds.

Figure 8: SDAR-1.7B constant-risk comparison:

\beta\in\{0.25,0.5,1.0\}

versus StepRS

[1.0\to 0.5]

. Curves show means with one-standard-deviation bars across three inference seeds.

Figure 9: SDAR-1.7B endpoint comparison:

[0.25\to 0]

[0.5\to 0.25]

[1.0\to 0]

, and the default

[1.0\to 0.5]

. All schedules are linear; curves show three-seed means and one-standard-deviation bars.

Figure 10: Qwen3-0.6B-MDLM constant-risk comparison. The default StepRS schedule is

[0.5\to 0.25]

. Curves show three-seed means with capped one-standard-deviation error bars. All three constant coefficients

\beta\in\{0.25,0.5,1.0\}

are shown.

Figure 11: Qwen3-0.6B-MDLM endpoint comparison:

[0.5\to 0]

[1.0\to 0]

, and the default

[0.5\to 0.25]

, all with linear interpolation. Curves show means with one-standard-deviation error bars across three inference seeds.

Figure 12: Unique-answer ratios for Base, GRPO, and StepRS across all three backbones. Points show means and capped error bars show sample standard deviations over inference seeds

42

43

, and

44

, with

64

responses per problem per seed.

Explore similar work

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models
May 8, 2026Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade +4Diffusion-Based Reinforcement Learning Methods Diffusion Language Models
The Ohio State University · Apple
Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models
Jun 7, 2026Yawen Shao, Jie Xiao, Kai Zhu +6Large Language Model Reinforcement Learning Diffusion Language Models
University of Science and Technology of China · Tongyi Lab · Northeastern University
Relative Score Policy Optimization for Diffusion Language Models
May 11, 2026Zichao Yu, Shengze Xu, Bingqing Jiang +2Large Language Model Reinforcement Learning Diffusion Language Models
University of Science and Technology of China · The Chinese University of Hong Kong · The University of Hong Kong

May 8, 2026, cs.LGJ/K move · Enter open · S save

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade, Lokesh Boominathan+3

The Ohio State University · Apple

Diffusion-Based Reinforcement Learning Methods Diffusion Language Models Large Language Model Reinforcement Learning Reward Gradients Credit Assignment Denoising Trajectory

Open paper AskPDF arXiv

Diffusion large language models are a compelling alternative to autoregressive models, yet existing RL methods for diffusion treat all denoising steps as equally important and rely on biased, high-variance likelihood estimates. We identify two fundamental weaknesses: the absence of temporal credit assignment across the denoising trajectory, and the systematic bias of mean-field likelihood estimates used for policy optimization. To address these, we propose Denoising-Aware Credit Assignment for GRPO (DACA-GRPO), a lightweight, plug-and-play enhancement for any GRPO-style trainer. DACA-GRPO introduces two complementary mechanisms: Denoising Progress Scores, which extract per-token importance weights from intermediate predictions at no additional forward cost, and Stratified Masking Likelihood, which partitions token positions into strata so that each token is predicted with most of the sequence as context, reducing the mean-field bias. Applied on top of three GRPO base methods, DACA-GRPO achieves consistent improvements across seven benchmarks spanning mathematical reasoning, code generation, constraint satisfaction, and constrained generation, with gains of up to 5.6pp on math reasoning, 7.4pp on code generation, 36.3pp on constraint satisfaction, and 5.9pp on JSON schema adherence.

Abstract

Explore similar work

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models

Relative Score Policy Optimization for Diffusion Language Models