Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages

Vishnu Teja Kunde; Fatemeh Doudi; Mahdi Farahbakhsh; Dileep Kalathil; Krishna Narayanan; Jean-Francois Chamberland

cs.LGMar 13, 2026

Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages

Authors: Vishnu Teja Kunde, Fatemeh Doudi, Mahdi Farahbakhsh, Dileep Kalathil, Krishna Narayanan, Jean-Francois Chamberland

Organizations: Department of Electrical & Computer Engineering, Texas A&M University, College Station, TX 77843

Abstract

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.

Panel of Figure 2 — Figure 2: (a) Entropy-guided step selection : $K$ steps with highest entropy $H_{t}$ are selected (solid lines); with $H_{3}>H_{1}>H_{2}>H_{0}$ and $K{=}2$ , steps $3$ and $1$ are chosen. (b) Stepwise advantage estimation : Reward of greedy one-step completion $\hat{\mathbf{x}}_{0\mid t+1}$ from $\mathbf{x}_{t+1}$ approximates $V_{t}$ ; $A_{t}$ measures the additional reward from continuing the full denoising trajectory to $\mathbf{x}_{0}$ .

	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
Model / Seq Len	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.

Panel of Figure 3 — 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.

Panel of Figure 4 — 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.	$\lambda$
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.	$\lambda$
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.

Panel of Figure 5 — 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

\lambda

on Sudoku. Constant

\lambda=1

and the step-varying schedule

\lambda_{t}=1-t/T

converge to similar accuracy; constant

\lambda=0.5

is noticeably slower. The magnitude of

\lambda

matters more than its functional form across denoising steps.

RL for Diffusion Models

Reinforcement learning for diffusion models adapts denoising and sampling to rewards, preferences, and task outcomes.

See all 173 papers

Reinforcement Learning from Denoising Feedback
May 25, 2026Qi He, Huan Chen, Ya Guo +3RL for Language Model Reasoning RL for Diffusion Models
Fudan University · Ant Group · Hong Kong University of Science and Technology
Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models
Jun 7, 2026Yawen Shao, Jie Xiao, Kai Zhu +6RL for Language Model Reasoning Credit Assignment in RL
University of Science and Technology of China · Tongyi Lab · Northeastern University
GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models
May 28, 2026Xiaohang Tang, Keyue Jiang, Che Liu +4Diffusion Model Distillation RL for Diffusion Models
UNIST · University of Basel

May 25, 2026J/K move · Enter open · S save

Reinforcement Learning from Denoising Feedback

Qi He, Huan Chen, Ya Guo, Huijia Zhu+2

Fudan University · Ant Group · Hong Kong University of Science and Technology

RL for Language Model Reasoning RL for Diffusion Models Diffusion Language Models

Open paper AskPDF arXiv

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.

Abstract

RL for Diffusion Models

Explore similar work

Reinforcement Learning from Denoising Feedback

Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models

GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models