Advancing Reasoning in Diffusion Language Models with Denoising Process Rewards

1 Introduction

• Mask-based diffusion language models refine masked sequences iteratively, but complex multi-step reasoning remains challenging.Diffusion language models offer potential inference acceleration through parallel generation.Source: §1 · p. 1
• Outcome-only rewards provide no direct denoising guidance, allowing repetitive or vacuous intermediate steps and unstable, poorly interpretable reasoning.GRPO and diffu-GRPO typically use sparse rewards based on final-answer correctness, leaving the final denoising steps to determine the output.Source: §1 · p. 1, Figure 1 · p. 2
• Autoregressive process reward models use human-annotated reasoning steps, but applying them to diffusion models faces partial-state mismatch, reward hacking, and computational overhead.Scoring intermediate states with a separate reward model increases GPU memory use and slows training.Source: §1 · pp. 1–2
• DPR estimates denoising-interval contributions by comparing final task outcomes across timesteps, without external annotations or reward models.The approach provides process-level supervision compatible with the denoising dynamics of mask-based diffusion models.Source: §1 · p. 2
• DPR consistently improves reasoning stability, interpretability, and overall task performance on challenging reasoning benchmarks.The method reuses standard rollout outcomes to estimate interval contributions, enabling integration into existing reinforcement learning pipelines with minimal additional cost.Source: §1 · p. 2

Mask-based diffusion-based large language models

• Mask-based diffusion work includes LLaDA, which gradually removes mask tokens, and diffu-GRPO, which assumes generated-token independence for likelihood estimation in reinforcement learning.Other approaches address sampling, remasking, weighted likelihood estimation, and efficiency or long-context generation.Source: §2 · p. 2
• Related approaches include partial-trace guidance for exploration, preferred-trajectory optimization with a diffusion value model, and methods for improving inference or estimating process rewards.The paper contrasts its approach with trajectory supervision and auxiliary value models, deriving rewards directly from denoising outcomes.Source: §2 · p. 3

Process reward model

• Process reward models localize errors within reasoning, but step-wise feedback can be costly to collect; automatic methods estimate process correctness from downstream answer outcomes.Examples include Monte Carlo completion or search methods, multimodal process rewards, tree-search policies, and implicit token-level rewards.Source: §2 · p. 3

3.1 Preliminary: Outcome-based Group Relative Policy Optimization (GRPO)

• GRPO samples responses, scores final-answer correctness, and optimizes clipped, KL-regularized advantages; in MdLLMs, diffu-GRPO approximates sequence likelihoods with independent per-token probabilities.The mean-normalized advantage is distributed uniformly across response tokens.Source: §3.1 · p. 3

3.2 Policy Optimization with Denoising Process Rewards

• DPR addresses outcome-only training’s lack of signal about intermediate denoising, which can leave reasoning trajectories uninformative or unstable despite correct answers.Directly applying process reward models to masked diffusion language models is also unreliable because intermediate states are partially specified and unordered.Source: §3.2 · p. 3
• DPR segments the denoising trajectory and scores each segment by comparing final-answer accuracy after completing from its two endpoint states.Contributions are aggregated across sub-intervals to reward trajectories that consistently guide generation toward correct solutions.Source: §3.2 · p. 4
• The stochastic estimator samples an intermediate timestep, compares its completion outcomes with standard rollouts from the fully masked state, and requires one additional intermediate-state rollout.Sampling different timesteps across updates estimates supervision over the full trajectory in expectation.Source: §3.2 · pp. 4–5, Figure 2 · p. 4
• Up-weighted advantages condition process supervision on correctness, increasing the gradient contribution of correct rollouts with better intermediate reasoning while leaving other advantages unchanged.This avoids over-penalizing correct responses or rewarding incorrect ones.Source: §3.2 · p. 5

4.1 Setup

• The evaluation uses GSM8K, MATH, COUNTDOWN, and SUDOKU, comparing against LLaDA, Diffu-GRPO, TSE, and WINO.The setup follows diffu-GRPO unless otherwise specified.Source: §4.1 · p. 5

4.2 Results

• Our method achieves the best results on most benchmarks and sequence lengths, with substantially higher reasoning–answer alignment on COUNTDOWN and GSM8K.Table 1 covers COUNTDOWN, GSM8K, SUDOKU, and MATH; Figure 5 reports alignment ratios.Source: §4.2 · p. 5, Table 1 · p. 6, Figure 5 · p. 8
• On GSM8K, diffu-GRPO+PRM scores 71.7, 80.9, and 81.5 versus diffu-GRPO’s 72.6, 79.8, and 81.9; ours reaches 72.9, 82.2, and 82.4 at lengths 128, 256, and 512.The comparison reports that our method has better downstream performance than the pretrained-PRM variant.Source: Table 1 · p. 6, §4.2 · pp. 6–7
• Pretrained PRM training is slower than our method, can produce NaNs on irregular responses, and shows rising reward while task performance fails to improve.Training takes 7.62 seconds per iteration with PRM versus 3.42 seconds for our method.Source: §4.2 · p. 6, Figure 4 · p. 7
• Ours-NoUpweight substantially outperforms diffu-GRPO; Ours-Cover has highest accuracy but higher cost, while Ours matches Cover and Random at diffu-GRPO-like speed.Ours-Random also has comparable accuracy but is slower because it needs additional forward passes and cannot reuse final completions.Source: §4.2 · p. 7
• Our method consistently achieves higher total rewards, combining accuracy and format rewards, during training.The authors connect these curves with the performance gains on the evaluated datasets.Source: §4.2 · p. 7, Figure 7 · p. 8

Our model demonstrates strong generalization

• On unseen SVAMP and ARC datasets, our method improves performance and achieves higher intermediate accuracy across both datasets.The authors suggest this intermediate accuracy may offer advantages when combined with MdLLM acceleration techniques.Source: Our model demonstrates strong… · p. 8
• With N ≥ 3 samples, the reward improves over LLaDA and diffu-GRPO across sequence lengths 128, 256, and 512; N = 1 is noisy and suboptimal.The ablation varies the number of samples across N ∈ {1, 3, 6, 9}.Source: Our model demonstrates strong… · p. 8

• The approach targets complex-reasoning training where denoising steps may be used poorly without process-level supervision.Source: §5 · p. 8
• Denoising-trajectory supervision is derived without external reward models or human annotation.Source: §5 · p. 8
• Experiments report more coherent reasoning, more reliable generation, and consistent improvements on challenging reasoning benchmarks.Source: §5 · p. 8

• The method adopts a mean-field assumption for response likelihood, treating token predictions as conditionally independent and neglecting token-level dependencies.Source: Limitations · p. 9
• More expressive likelihood estimators that capture token interactions would significantly increase inference cost and memory, undermining reinforcement-learning scalability for large diffusion models.Source: Limitations · p. 9
• Future work could relax the mean-field assumption while preserving efficient and stable training, including through structured variational approximations or alternative parameterizations.Source: Limitations · p. 9

• Denoising-derived intrinsic rewards may reinforce spurious correlations or reasoning patterns that perform well on benchmarks but are misaligned with human values or real-world norms.Source: Ethical considerations · p. 9
• DPR does not inherently encode ethical constraints or safety objectives; future work could integrate explicit safety-aware or value-aligned objectives.Source: Ethical considerations · p. 9

• The training code is provided in the supplementary material.Source: A Implementation · p. 12
• Training uses diffu-GRPO with LLaDA-8B-Instruct, six rollouts per problem, and LoRA rank 128 in 4-bit precision; each dataset takes around 72 hours on eight modern GPUs.Source: A Implementation · p. 12
• For reward estimation, one intermediate timestep is sampled per rollout, and 64 mask tokens are appended before denoising continues.Source: A Implementation · p. 12
• Denoising process reward is the difference between accuracy from intermediate-state completions and accuracy from original fully masked rollouts.Source: A Implementation · p. 12

B Additional Examples

• On COUNTDOWN and SUDOKU, Ours produces more coherent, higher-quality reasoning than diffu-GRPO at mask lengths 128 and 512, respectively.Diffu-GRPO often produces irrelevant reasoning before abruptly giving a final answer.Source: B Additional Examples · p. 12
• In a COUNTDOWN example, diffu-GRPO’s unrelated reasoning leads to the incorrect expression 92 −17 −17 rather than 92 −17 −77.The passage presents this as an example of a final answer that fails despite being numerically close to the correct expression.Source: B Additional Examples · p. 12
• At GSM8K step 32, Ours produces higher prediction accuracy than diffu-GRPO under early decoding.At that stage, Ours already produces substantial correct, coherent reasoning, while diffu-GRPO often has incomplete or flawed reasoning.Source: B Additional Examples · p. 12

C License of Used Code and Datasets

• GSM8K and MATH are released under the MIT License, while COUNTDOWN and the diffu-GRPO codebase are licensed under Apache License 2.0.The passage states that these licenses permit use for academic research purposes.Source: C License of Used Code and Datasets · p. 12
• The authors state that their use of the datasets and diffu-GRPO codebase is consistent with intended research, benchmarking, and method-development use.Source: C License of Used Code and Datasets · p. 12
• The examples include COUNTDOWN with 128 mask tokens, SUDOKU with 512, and GSM8K intermediate decoding at step 32.The captions describe comparisons between Ours and diffu-GRPO on these examples.Source: Table 5 · p. 13, Table 6 · p. 14, Table 7 · p. 15