PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.00661v1cs.LGcs.AI

TL;DR

dLLM reasoning with verifiable rewards reuses terminal feedback across denoising decisions whose contexts change, raising questions about accuracy, coverage, and diversity. StepRS-GRPO varies risk-sensitive advantage weighting across denoising states while retaining the underlying trainer. Across mathematical reasoning evaluations, it improves accuracy, sampled-solution coverage, and output diversity over centered GRPO.

  • Problem

    In dLLMs, terminal-reward policy gradients reuse rollout-level advantages even as conditioning contexts change, and accuracy alone does not establish sampled-solution coverage or diversity.

    Source: §1 · pp. 1–2

  • Method

    StepRS-GRPO varies the entropic group-advantage coefficient across denoising states, exactly rescaling centered outcome advantages in the binary-reward setting.

    Source: §1 · p. 2

  • Results

    Across mathematical reasoning evaluations, StepRS-GRPO improves mathematical accuracy, sampled-solution coverage, and output diversity over GRPO.

    Source: §1 · p. 2, §5 · p. 12

  • Takeaways & Limitations

    Ablations support contributions from state allocation and schedule direction, and the gains persist when advantage RMS is matched to centered GRPO.

    Source: §5 · p. 12

    The evaluation focuses on mathematical tasks with binary rewards, and its three inference seeds measure sampling variability at fixed checkpoints rather than variability across training runs.

    Source: §5 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Diffusion language models denoise masked sequences rather than generate strictly left to right, and reinforcement learning with verifiable rewards is a natural next step for improving their reasoning. StepRS-GRPO varies risk-sensitive advantage weighting across denoising states, with the aim of improving accuracy while preserving search coverage and diversity.

1 Introduction

  • dLLMs denoise masked sequences through reverse steps, while RLVR offers a way to improve their reasoning using verifiable outcome rewards.Prior post-training work adapts supervised learning to masked dLLMs or distills autoregressive models into bidirectional masked dLLMs.Source: §1 · p. 1
  • A higher pass@1 alone does not establish better coverage at larger sampling budgets, and neither pass@1 nor coverage establishes diversity of correct reasoning strategies.The paper therefore compares scheduled training with constant-risk training and evaluates accuracy, sampled-solution coverage, and diversity.Source: §1 · p. 2
  • StepRS-GRPO varies the entropic advantage-transformation coefficient across denoising states while retaining terminal rewards and the chosen trainer’s likelihood estimator and loss reduction.For binary rewards, this is exactly a prompt- and state-dependent rescaling of centered outcome advantages.Source: §1 · p. 2
  • The authors use capability-guided coefficient selection and schedule ablations, and report gains in accuracy, solution coverage, and output diversity across dLLM backbones.Their controls examine state allocation and advantage scale, alongside schedule ablations and a masked-dLLM case study.Source: §1 · p. 2

2 Related Work

Related work adapts policy optimization and reward feedback to non-autoregressive denoising, while also studying diversity and risk-sensitive objectives. StepRS-GRPO complements these approaches by changing state-dependent coefficients on terminal rewards while retaining the selected trainer’s loss construction.

2 Related Work

  • Prior dLLM post-training methods adapt policy optimization to denoising, including trace optimization, unconditional clipping, trajectory-estimation changes, and likelihood-estimation approaches.The passage names d1, TraceRL, StableDRL, dTRPO, GDPO, ESPO, AGRPO, SPG, and BGPO, with their cited bibliography markers.Source: §2 · p. 2
  • StepRS complements denoising-stepwise objectives by applying state-dependent coefficients in an entropic group-advantage transformation while retaining the selected trainer’s likelihood estimator and loss construction.The passage contrasts this design with wd1++, which extends weighted log-likelihood to denoising-stepwise optimization, and DMPO, which matches a reward-tilted target through weighted denoising cross-entropy.Source: §2 · p. 2
  • Diversity-related work includes entropy analyses, divergence or explicit diversity rewards, adaptive entropy targets, and direct pass@k optimization; state-aware methods use different feedback procedures.JustGRPO retains parallel decoding at inference; DiSPO scores completions branched from intermediate masked states, whereas IGPO uses inpainting guidance.Source: §2 · p. 3
  • Risk-sensitive methods include entropic objectives and RS-GRPO, while StepRS-GRPO uses exponential tilting to interpret its coefficient and guide schedule endpoints.The paper distinguishes its finite-group analysis of the entropic transformation from the empirical schedule evaluation.Source: §2 · p. 3

3 Method: From GRPO to StepRS-GRPO

StepRS-GRPO changes how terminal binary-reward feedback is weighted across denoising decisions, without changing the selected trainer’s underlying update construction. Its finite-group analysis identifies the transformation as a rescaling of centered advantages, and a capability-based scale guides empirical schedule selection.

3 Method: From GRPO to StepRS-GRPO

  • The method specifies terminal-reward updates, introduces risk-sensitive weighting, then varies that weighting across denoising states while leaving the rollout procedure and trainer components unchanged.The rollout procedure, verifier, likelihood estimator, loss reduction, optimizer, and reference regularization remain as selected.Source: §3 · p. 3

3.1 Diffusion language models

  • A dLLM trajectory comprises reverse denoising states containing partially revealed text and decoder information; maskedness and available information change as tokens are revealed.Masked dLLMs use bidirectional context, while block dLLMs generate blocks left to right and denoise within each block.Source: §3.1 · p. 4

3.2 From centered advantages to state-dependent risk

  • Centered GRPO reuses each response’s terminal binary reward across active decisions, while the native loss and its reduction and reference-regularization terms are retained.The trainer’s native surrogate can depend on the full sampled advantage array through its normalization; its construction is not replaced by the PPO example.Source: §3.2 · pp. 4–5
  • StepRS-GRPO applies a fixed schedule to each decision’s state coordinate, so the same terminal feedback can receive different weights at different denoising states.Setting the schedule to zero recovers centered GRPO; a positive constant recovers capped constant-risk RS-GRPO.Source: §3.2 · p. 5

3.3 Binary-reward interpretation and schedule selection

  • For binary rewards, the entropic transformation factors into a centered outcome advantage times a positive multiplier that decreases with the group’s solve fraction.The multiplier preserves advantage signs; at a fixed coefficient, groups with fewer successes receive larger multipliers than groups with more successes.Source: §3.3 · p. 6
  • Exponential tilting interprets β as adding β to success log-odds, and matching the fractional failure-odds reduction to the initial policy’s average solve rate provides a candidate coefficient scale.For SDAR-4B, the passage reports an average solve rate of approximately 0.66 and a candidate β⋆ of approximately 1.08; this calibration does not establish an optimal coefficient.Source: §3.3 · p. 6
  • The linear schedule assigns βmax at the more-masked endpoint and βmin at the less-masked endpoint; half-decay is used as a practical default and schedule shape is tested separately.The endpoint ratio γ controls decay, with γ = 0 giving full decay and γ = 1 a constant coefficient.Source: §3.3 · p. 7

4 Experiments

Experiments evaluate StepRS-GRPO on SDAR diffusion models and mathematical reasoning benchmarks, then test its state allocation, risk schedule, and advantage scaling. StepRS-GRPO improves accuracy and sampled-solution coverage over GRPO, while ablations and trace diagnostics examine the source and breadth of those gains.

4.1 Setup

  • The experiments train 1.7B- and 4B-parameter SDAR-Chat models on 7,500 MATH problems and evaluate four mathematical benchmarks with repeated inference seeds.Training uses 16 rollouts per prompt, binary correctness rewards, and 160 optimizer steps; the selected schedules are [1.0 →0.5] and [1.5 →0.75].Source: §4.1 · pp. 7–8, Figure 1 · p. 8

4.2 Main Results

  • StepRS-GRPO improves mean pass@k over GRPO at every evaluated budget across all eight SDAR model–benchmark pairs.At 4B, AIME2024 and AIME2025 gains grow from 1.17 and 1.87 points at k = 1 to 8.89 and 10.01 points at k = 64.Source: Figure 1 · p. 8, §4.2 · pp. 8–9

4.3 Ablation Studies

  • StepRS has higher mean pass@k than prompt-only mass-matched C1 in every comparison, supporting state-dependent allocation beyond prompt-level weighting.C1 replaces state-dependent multipliers with each prompt’s loss-weighted mean, preserving total multiplier mass before the common cap.Source: Figure 2 · p. 9, §4.3 · p. 10
  • StepRS is higher in all 28 comparisons than the mass-matched reversed schedule, with pass@64 gains of 5.58 and 5.59 points on AIME2024 and AIME2025.The positive rescaling factor preserves reference multiplier mass before the common cap; reversal alone need not do so because the transformation is nonlinear.Source: Figure 2 · p. 9, §4.3 · p. 10
  • RMS-matched C3 stays within 0.24 points of StepRS on pass@1, and both variants exceed GRPO in all comparisons.C3 matches overall advantage RMS to centered GRPO without recentering the StepRS advantages.Source: Figure 2 · p. 9, §4.3 · pp. 10–11
  • At SDAR-4B, the calibrated [1.08 →0.54] schedule beats [1.0 →0.5] at every benchmark–budget combination and is within 0.47 percentage points of [1.5 →0.75] on pass@1.The calibration uses a solve-rate proxy of approximately 0.66 to suggest β⋆≈1.08 as a starting scale.Source: §4.3 · p. 11
  • Linear and cosine schedules differ by at most 0.47 percentage points in mean pass@k; both exceed the tested two-stage schedule at every benchmark–budget combination.With endpoints fixed at [1.5 →0.75], the authors retain linear half-decay as a simple, effective default.Source: §4.3 · p. 11, Figure 4 · p. 11
  • Across four benchmarks, StepRS exceeds GRPO by 2.3–6.6 Dngram points and 1.3–3.5 Demb points on answer-removed reasoning traces.These diagnostics show that increased variation extends beyond final-answer strings; the two measures use different panel scales.Source: Figure 3 · p. 11, §4.3 · p. 12

5 Conclusion and Limitations

The paper concludes that StepRS-GRPO improves mathematical accuracy, sampled-solution coverage, and output diversity, with gains retained under RMS matching. Its appendices describe masked and block diffusion conventions and note that the evaluation does not establish an optimal schedule or training-run consistency.

5 Conclusion and Limitations

  • StepRS-GRPO improves mathematical accuracy, sampled-solution coverage, and output diversity, with gains over GRPO retained by the advantage-RMS-matched variant.Ablations support state allocation and schedule direction beyond prompt-level weighting; calibration and schedule-shape experiments guide coefficient selection.Source: §5 · p. 12
  • Evaluation is limited to mathematical tasks with binary rewards, and three inference seeds measure sampling variability at fixed checkpoints rather than variability across training runs.The capability rule guides endpoint selection but does not establish an optimal schedule; broader reward settings and additional training replications are identified as useful extensions.Source: §5 · p. 12

A Diffusion Language Model Families

  • Appendix A defines the clean-token vocabulary, mask token, padded output length, and point-mass indicator used in its diffusion notation.The mask token is outside the clean-token vocabulary, and the indicator equals one when its condition holds.Source: A Diffusion Language Model Families · p. 18

A.1 Masked diffusion

  • Masked diffusion independently masks positions in an absorbing forward process, whose expected masked fraction is α(ζ); a confidence-based decoder may use a different reverse transition.The denoising objective assumes absolutely continuous α with α(ζ) > 0 for 0 < ζ ≤ 1, and uses α′/α because α denotes masking probability.Source: §A.1 · p. 18

A.2 Block diffusion

  • Block diffusion partitions the sequence into blocks, uses a prefix before each block, and applies causal attention across blocks with bidirectional attention within the active block.The risk coordinate is not automatically continuous diffusion time or expected masking probability; the SDAR experiments use block length 4.Source: §A.2 · pp. 18–19

B Additional Theoretical Analysis and Calibration

The analysis distinguishes three probability quantities by their roles in the method: forming advantages, describing a fixed rollout distribution, and suggesting a coefficient scale.

B Additional Theoretical Analysis and Calibration

B.1 Proof of Lemma 3.1

The proof gives an exact finite-group factorization for binary rewards and interprets its multiplier, then describes capability-guided coefficient calibration and the scope of state-wise scheduling.

B.1 Proof of Lemma 3.1

  • For each realized binary-reward group, the finite-group factorization holds without an independence assumption or large-group limit, and reduces continuously to centered advantages at β = 0.When all rewards in a group are identical, all advantages are zero; the finite-group identity does not by itself establish an unbiased gradient estimator.Source: §B.1 · pp. 19–20
  • At fixed β, groups with lower observed solve fractions receive larger multipliers, but increasing β need not increase advantage magnitude uniformly across solve rates.With a common coefficient, group advantages sum to zero; action-dependent coefficients and token-reduction weights need not preserve a zero action-weighted mean.Source: §B.1 · p. 20

B.2 Capability-guided calibration

  • The capability-guided rule provides a search scale, not an optimizer-derived optimum or upper bound; for SDAR-4B, p̄ ≈ 0.66 gives β⋆ ≈ 1.08.The rounded candidate is compared with neighboring endpoints, which are then held fixed during each run.Source: §B.2 · p. 20

B.3 Scope of denoising-state scheduling

  • StepRS applies one scalar group transformation per action coefficient across that group’s normalizer; coefficients are not mixed within a denominator.The schedule supplies advantages to the selected native surrogate, preserving the chosen trainer’s loss construction rather than deriving a fixed-coefficient trajectory objective.Source: §B.3 · pp. 20–21

C.1 Additional results: 0.6B masked diffusion LM

On Qwen3-0.6B-MDLM, the selected StepRS schedule improves pass@k over GRPO across three benchmarks and also increases unique-answer ratios; this case study is complementary portability evidence.

C.2 Training configuration

  • The model is fine-tuned on GSM8K with 16 responses per prompt and binary correctness rewards, and evaluated at the 60-step checkpoint using three inference seeds.The appendix points to Table 2 for detailed training recipes across the reported model scales.Source: §C.1 · p. 21, §C.2 · p. 21

C.1 Additional results: 0.6B masked diffusion LM

  • Across all 21 benchmark–budget combinations, [0.5 →0.25] has higher mean pass@k than GRPO; at pass@1, GSM8K rises from 40.70 to 44.57, MATH500 from 5.20 to 8.13, and AMC23 from 0.20 to 1.20.At k = 64, the corresponding values are 84.25 →89.20, 26.67 →40.20, and 10.00 →28.33.Source: §C.1 · p. 21
  • Unique-answer ratios increase from 18.03 to 28.00 on GSM8K, 11.40 to 33.47 on MATH500, and 17.30 to 24.83 on AMC23.Because this case study uses a different backbone family and training corpus, it is complementary portability evidence rather than a controlled model-size comparison.Source: §C.1 · p. 21

C.3 Evaluation protocol

The protocol evaluates fixed checkpoints across inference seeds and sampling budgets, using verifier-based pass@k alongside separate answer-string and reasoning-trace variation measures. It also specifies prompt and control procedures.

C.3 Evaluation protocol

  • Each question is sampled 64 times per inference seed at temperature 1.0; pass@k is computed per seed and averaged over questions, without pooling seeds into an N = 192 estimate.Capped error bars show the sample standard deviation across inference seeds, not a standard error or confidence interval.Source: §C.3 · pp. 21–22
  • Unique-answer ratio counts normalized extracted final-answer strings, while answer-removed trace metrics retain incorrect responses and repeated texts and do not measure distinct correct solution strategies.Failed extractions form one shared empty-answer class; traces are assessed within each question.Source: §C.3 · p. 22
  • Lexical variation uses Self-BLEU: lower average BLEU-4 overlap with the other sampled traces indicates greater lexical variation, and it is not distinct-4.Embedding variation averages pairwise cosine distances between normalized sentence embeddings; the two measures have different scales and are not numerically comparable.Source: §C.3 · pp. 23–24

C.4 Prompt templates

  • The ChatML prompt requests a boxed final answer; the listed evaluations use no system message, few-shot examples, or additional reasoning scaffold.The verifier extracts the final answer from the boxed span in the assistant continuation.Source: §C.4 · p. 24

D.1 Mechanism-control definitions and matching properties

  • The C1 control replaces action-level multipliers with a prompt-level weighted average, preserving multiplier mass while removing state variation.The defined mass is a sum of nonnegative multiplier weights, not probability mass, signed advantage, or gradient magnitude.Source: §D.1 · p. 25
  • C3 scales the active StepRS advantage array by one positive scalar to match centered-GRPO RMS, preserving signs and relative weighting rather than matching the centered-GRPO gradient.The scale does not match the KL gradient, optimizer state, or final parameter update; applying another cap or non-uniform normalization afterward changes the matching identity.Source: §D.1 · p. 26

D.2 Calibration and endpoint choice

The endpoint and constant-risk comparisons support the selected schedules across backbones, while showing that their advantage varies by benchmark and sampling budget.

D.2 Calibration and endpoint choice

  • The calibrated [1.08 →0.54] schedule improves mean pass@1 over [1.0 →0.5] by 0.40, 0.83, 0.33, and 1.27 points across four benchmarks, and improves mean pass@k at every evaluated budget.Its pass@1 remains close to the selected [1.5 →0.75] schedule, supporting calibration as a useful starting point for endpoint selection.Source: §D.2 · pp. 26–27, Figure 6 · p. 27
  • At a fixed upper endpoint of 2.0, half-decay [2.0 →1.0] improves pass@1 over full decay [2.0 →0] by 0.53, 2.07, 0.57, and 0.77 points.Its mean pass@k is higher in 27 of 28 comparisons and equal in the remaining one, supporting a nonzero late-stage coefficient in this setting.Source: Figure 6 · p. 27, §D.2 · p. 27

D.3 Full constant-risk sweep

  • Against constant β = 1.5, the selected SDAR-4B schedule has higher mean pass@k in 26 of 28 comparisons and ties in the other two.It improves all seven budgets on AMC23, AIME2024, and AIME2025; the best constant coefficient varies by benchmark and budget.Source: §D.3 · p. 27, Figure 7 · p. 28

D.4 Risk and endpoint sensitivity on smaller backbones

  • On SDAR-1.7B, [1.0 →0.5] matches or improves on constant β = 1.0 at all 28 benchmark–budget combinations; on Qwen3-0.6B-MDLM, [0.5 →0.25] leads at every GSM8K and MATH500 budget.On AMC23, constant β = 0.5 leads by 0.10 points at pass@1, but the scheduled variant leads from k = 2 through 64, reaching 28.33 versus 27.50 at pass@64.Source: §D.4 · p. 28

D.5 Answer and trace diversity

StepRS-GRPO increases mean answer variation over GRPO on every model–benchmark pair, complementing accuracy and coverage results. The section also reports answer-removed trace diagnostics.

D.5 Answer and trace diversity

  • StepRS-GRPO increases mean answer variation over GRPO on every model–benchmark pair, complementing accuracy and coverage results.Figure 12 reports unique-answer ratios across all three backbones, with 64 responses per problem per inference seed.Source: §D.5 · p. 28, Figure 12 · p. 30
  • Schedule comparisons include constant β values of 0.25, 0.5, and 1.0, plus linearly interpolated endpoint schedules for Qwen3-0.6B-MDLM and SDAR-1.7B.Source: Figure 10 · p. 29, Figure 9 · p. 29, Figure 11 · p. 30
  • Answer-removed trace diagnostics are presented in the main text as Figure 3.Source: §D.5 · p. 29
  • Appendix C.3 specifies preprocessing, Self-BLEU-4, embedding distance, and aggregation across questions and inference seeds under one evaluation protocol.Source: §D.5 · p. 29

E Solution Examples

This section presents solution examples comparing base, GRPO, and StepRS-GRPO across three benchmarks, including supplied GSM8K and MATH-500 cases and a worked check of candidate values against the original equations.

GSM8K on Qwen3-0.6B-MDLM

  • For Lee’s farm, one solution concludes 15 sheep; the base-model and GRPO solution examples are labeled incorrect, and cited solutions state 10 sheep and 10 geese.Source: E Solution Examples · pp. 31–33

MATH-500 on SDAR-1.7B-Chat

  • The StepRS-GRPO finetuned model’s solution example is marked correct, whereas the base-model and GRPO-finetuned examples are marked incorrect; the median QM length is 30.Source: E Solution Examples · pp. 33–35

AMC’23 on SDAR-4B-Chat

  • The excerpts label the Base Model Solution Example and GRPO Finetuned Model Solution Example incorrect, while labeling the StepRS-GRPO Finetuned Model Solution Example correct.Source: E Solution Examples · pp. 36, 37, 39

E Solution Examples

  • Checking both candidate sums, 36 and 2, against the original equations selects x = 27, y = 9, and rejects x = 1, y = 1.For (27, 9), both conditions equal 324; for (1, 1), the second condition fails because 4 ≠ 0.Source: E Solution Examples · p. 40

Related papers

Loading 2610.00661v1…