🤖 AI Summary
This study addresses the context-condition shift problem arising from terminal feedback reuse in reinforcement learning for diffusion language models. To mitigate this issue, we propose StepRS-GRPO, a method that dynamically modulates the risk coefficient of group advantage transformations based on denoising states, which we prove to be equivalent to an adaptive scaling mechanism. By integrating verifiable rewards with the GRPO algorithm, our approach effectively optimizes the training process of diffusion-based large language models. Experimental results demonstrate that StepRS-GRPO significantly improves both accuracy and coverage in mathematical reasoning tasks, while simultaneously enhancing the diversity of generated answers and the richness of reasoning trajectories.
📝 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.