Exploring More, Reasoning Better: Stepwise Risk-Sensitive GRPO for Diffusion Language Models
cs.LG, cs.AI
Submitted: 2026-09-30
Updated: 2026-10-09
Comments: 40 pages, 12 figures, 2 tables. The first two authors contributed equally
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion large language models (dLLMs) generate text by denoising a sequence or successive blocks, allowing several tokens to be revealed in parallel.
Terminology
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.
Sources
- Evaluating Large Language Models Trained on Code
- SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation
- The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models
- Soft Actor-Critic Algorithms and Applications
- Measuring Mathematical Problem Solving With the MATH Dataset
- Risk-Sensitive RL for Alleviating Exploration Dilemmas in Large Language Models
- The Choice of Divergence: A Neglected Key to Mitigating Diversity Collapse in Reinforcement Learning with Verifiable Reward
- Jointly Reinforcing Diversity and Quality in Language Model Generations
- Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion Large Language Models
- The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language Models
- Large Language Diffusion Models
- Diffusion-State Policy Optimization for Masked Diffusion Language Models
- Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective
- RiskPO: Risk-based Policy Optimization via Verifiable Reward for LLM Post-Training
- Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization
- Proximal Policy Optimization Algorithms
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models
- Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems
- SCOPE-RL: Stable and Quantitative Control of Policy Entropy in RL Post-Training
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks