Reinforcement learning (RL) has been effective for post-training autoregressive (AR) language models, but extending these methods to diffusion language models (DLMs) is challenging due to intractable sequence-level likelihoods. Existing approaches therefore rely on surrogate likelihoods or heuristic approximations, which can introduce bias and obscure the sequential structure of denoising. We formulate diffusion-based sequence generation as a finite-horizon Markov decision process over the denoising trajectory and derive an exact, unbiased policy gradient that decomposes over denoising steps and is expressed in terms of intermediate advantages, without requiring explicit evaluation of the sequence likelihood. To obtain a practical and compute-efficient estimator, we (i) select denoising steps for policy updates via an entropy-guided approximation bound, and (ii) estimate intermediate advantages using a one-step denoising reward naturally provided by the diffusion model, avoiding costly multi-step rollouts. Experiments on coding and logical reasoning benchmarks demonstrate state-of-the-art results, with strong competitive performance on mathematical reasoning, outperforming existing RL post-training approaches for DLMs. Code is available at https://github.com/vishnutez/egspo-dllm-rl.
Figures & tables
Figure 1: Our methods achieve state-of-the-art results on nearly all coding, logical, and mathematical reasoning benchmarks. The red bar reports the better of EGSPO and EGSPO-SA per task (LLaDA-8B-Instruct, 0-shot, 2-token/step); see Tables 1 , 3 for the full breakdown.
Figure 2: (a) Entropy-guided step selection : K steps with highest entropy Ht are selected (solid lines); with H3>H1>H2>H0 and K=2 , steps 3 and 1 are chosen. (b) Stepwise advantage estimation : Reward of greedy one-step completion x^0∣t+1 from xt+1 approximates Vt ; At measures the additional reward from continuing the full denoising trajectory to x0 .
Sudoku
Countdown
GSM8K
MATH500
Model / Seq. Len.
128
256
512
Best
128
256
512
Best
128
256
512
Best
128
256
512
Best
LLaDA-8B-Instruct
11.7
6.7
5.5
11.7
20.7
19.5
16.0
20.7
68.7
76.7
78.2
78.2
26.0
32.4
36.2
36.2
d1
22.1
16.7
9.5
22.1
34.8
32.0
42.2
42.2
73.2
81.1
82.1
82.1
33.8
38.6
40.2
40.2
wd1
-
76.4
62.8
76.4
-
51.2
46.1
51.2
-
80.8
82.3
82.3
-
34.4
39.0
39.0
SPG
82.9 †
94.0 †
93.1 †
94.0 †
68.8
71.5
70.3
71.5
78.5
86.1
84.5
86.1
33.4
40.0
41.8
41.8
ESPO
72.9
58.6
54.5
72.9
81.6
82.0
79.3
82.0
76.7
79.9
80.8
80.8
32.0
34.6
39.4
39.4
Table 1: Results on reasoning benchmarks for LLaDA-8B-Instruct. We report accuracy at different generation lengths and the best performance across them. d2 reports only the best performance; wd1 does not report evaluation at length 128. For Sudoku, † denotes 3 -shot evaluation, whereas all other entries are 0 -shot. ESPO results are re-evaluated from the publicly released ESPO checkpoints under the uniform protocol shared by all other entries in the table, since ESPO’s original protocol uses 1-shot evaluation on Sudoku and decodes 1 token/step on GSM8K and MATH500; their originally published numbers under that protocol are shown in gray for completeness.
Model / Seq Len
Sudoku
Countdown
GSM8K
MATH500
128
256
512
Best
128
256
512
Best
128
256
512
Best
128
256
512
Best
Dream
9.3
2.1
14.0
14.0
8.5
7.8
17.4
17.4
75.8
81.3
80.7
81.3
38.2
45.7
48.0
48.0
ESPO
71.7
72.3
71.3
72.3
68.8
66.8
64.8
68.8
79.6
82.3
82.0
82.3
40.3
47.4
50.3
50.3
EGSPO (ours)
86.4
90.3
87.0
90.3
76.6
73.4
72.3
76.6
77.4
82.1
82.3
82.3
42.4
49.2
50.4
50.4
EGSPO-SA (ours)
90.6
90.3
89.2
90.6
79.7
82.4
78.5
82.4
77.6
84.9
86.1
86.1
41.2
47.2
47.8
47.8
Table 2: Results on reasoning benchmarks for Dream-7B-Instruct . We report accuracy at different generation lengths and the best performance across them. For GSM8K and MATH500, all models are evaluated using 1 token/step decoding to match ESPO’s protocol; ESPO is also trained with 1 token/step decoding, whereas our method is trained with 2 token/step decoding. For Sudoku and Countdown, both training and evaluation use 2 token/step decoding for all methods.
HumanEval
MBPP
Model
128
256
512
Best
128
256
512
Best
LLaDA
27.4
35.3
37.8
37.8
36.2
41.2
40.4
41.2
d1
31.1
32.9
37.8
37.8
40.5
44.7
42.8
44.7
ESPO
25.0
31.7
36.0
36.0
49.0
48.6
49.0
49.0
ESPO
28.1
42.1
50.0
50.0
47.4
44.6
44.2
47.4
EGSPO
32.3
40.2
39.6
40.2
49.9
50.6
50.3
50.6
Table 3: Results on coding benchmarks for LLaDA-8B-Instruct. We report accuracy at different generation lengths and the best performance across them. ESPO uses full fine-tuning, whereas d1, EGSPO, and EGSPO-SA use LoRA-based fine-tuning. ESPO results are re-evaluated from the publicly released ESPO checkpoints under the uniform protocol shared by all other entries in the table, since ESPO’s original protocol decodes 1 token/step on coding tasks; their originally published numbers under that protocol are shown in gray for completeness.
Figure 3: Compute efficiency comparison between EGSPO-SA and d1 on Sudoku. (a) FLOPs accumulated over all forward passes across 8 GPUs. (b) Cumulative prompt–completion pairs seen during training. (c) Optimizer updates (accounting for gradient accumulation). EGSPO-SA dominates d1 under all three compute budgets.
Figure 4: Ablation studies analyzing (a) entropy-guided step selection (EGSPO) versus uniform and random step selection, (b) the distribution of selected denoising steps, and (c) performance as a function of the step budget K .
K
Approx. Grad. Steps
128
256
512
4
4400
76.5
83.4
83.6
8
3600
80.7
86.0
86.1
16
2400
78.2
83.4
83.6
Table 4: GSM8K accuracy (EGSPO, LLaDA) for different step budgets K and generation length
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Task
B
G
K
Grad. Accum.
λ
Sudoku
48
8
8
2
1 (constant)
Countdown
96
8
8
2
1 (halved every 500 steps)
GSM8K
48
8
8
8
0.2 (constant)
MATH500
96
8
8
8
1 (halved every 500 steps)
Coding
96
8
16
4
1 (halved every 500 steps)
Appendix
Table 5: Task-specific hyperparameters for LLaDA-8B-Instruct . B : batch size (prompt-completion pairs per update), G : completions per prompt, K : step budget.
Task
B
G
K
Grad. Accum.
λ
Sudoku
192
16
8
2
1 (constant)
Countdown
192
16
8
2
0.2 (halved every 500 steps)
GSM8K
192
16
8
2
0.5 (halved every 500 steps)
MATH500
192
16
8
2
0.5 (halved every 500 steps)
Appendix
Table 6: Task-specific hyperparameters for Dream-7B-Instruct . B : batch size (prompt-completion pairs per update), G : completions per prompt, K : step budget.
Figure 5: Training curves for EGSPO and EGSPO-SA on Sudoku, Countdown, GSM8K, and MATH500 using the training settings described in Section D .
Figure 6: Ablation on the advantage weight λ on Sudoku. Constant λ=1 and the step-varying schedule λt=1−t/T converge to similar accuracy; constant λ=0.5 is noticeably slower. The magnitude of λ matters more than its functional form across denoising steps.
Policy loss estimation remains a fundamental and long-standing challenge in reinforcement learning (RL) for diffusion language models (DLMs). We introduce Reinforcement Learning from Denoising Feedback (RLDF), a novel training paradigm that leverages feedback obtained from rollout and training processes to facilitate accurate and efficient policy loss estimation. To balance the trade-off between computational efficiency and estimation effectiveness, RLDF optimizes the model toward the clipped clean state from intermediate noisy states, combined with weighted timestep sampling over denoising timesteps. Extensive experiments demonstrate that RLDF achieves consistent and substantial improvements in both performance and generalizability across two representative DLM architectures, LLaDA and Dream, on multiple reasoning benchmarks. Our work lays a principled foundation for scalable reinforcement learning in diffusion language models. We build Drift, a training framework for DLMs, available at https://github.com/ant-research/Drift.
Qi He, Huan Chen, Ya Guo +3
Fudan University · Ant Group · Hong Kong University of Science and Technology
Reinforcement learning (RL) holds immense promise for enhancing the reasoning capabilities of diffusion large language models (dLLMs). However, progress is fundamentally constrained by a dual misalignment between authentic generation trajectory and the gradient update process: (i) Process-reward misalignment. Sparse, terminal rewards are indiscriminately assigned to all intermediate steps of the generation process, failing to provide discriminative credit assignment. (ii) State-trajectory misalignment. Policy updates are often diverted toward artificial, out-of-trajectory states, squandering gradients on less informative samples. To address these limitations, we introduce Process Aligned Policy Optimization (PAPO), a novel framework that holistically aligns the RL update with the dLLM's generative trajectory via Step-Aware Process Rewards (SPR) that transform sparse terminal rewards into dense, step-wise credit, and Entropy-Guided Historical Re-enactment (EHR) that replays authentic trajectories at high-uncertainty steps. Extensive experiments on four benchmarks demonstrate that PAPO significantly outperforms baselines, achieving gains of up to 4.5% on GSM8K, 4.8% on MATH500, 42.2% on Countdown and 16.1% on Sudoku.
Yawen Shao, Jie Xiao, Kai Zhu +6
University of Science and Technology of China · Tongyi Lab · Northeastern University
Reinforcement learning (RL) can be used to improve the policy (denoiser) of diffusion large language models (dLLMs), while being hindered by the intractability of the policy likelihood. A dominant and efficient family of methods replaces the likelihood in standard RL with its evidence lower bound (ELBO), estimated from randomly masked sequences. Despite being well aligned with pre-training, these approaches introduce bias through training--inference mismatch by using the ELBO as a likelihood surrogate, which can degrade performance. In this work, we propose Guided Denoiser Self-Distillation (GDSD) to directly distill the denoiser of dLLMs from an advantage-guided self-teacher, derived from the closed-form optimum of reverse-KL regularized RL. GDSD matches the dLLM's denoiser logits to the teacher's via a normalization-free objective, which reduces RL to likelihood-free self-distillation and thus bypasses the TIM biases. Recent ELBO-based methods emerge as instances of applying different distillation divergences, but with diagnosable pathologies that GDSD avoids. On planning, math, and coding benchmarks with LLaDA-8B and Dream-7B, GDSD consistently outperforms prior state-of-the-art ELBO-based methods with a more stable training reward dynamics, achieving test-accuracy improvements of up to +19.6%. These results suggest that direct denoiser self-distillation, without relying on an ELBO likelihood surrogate, can provide a more stable and effective RL procedure for dLLMs. Code is available at https://github.com/GaryBall/GDSD.