Research map / Reinforcement learning / RL from feedback
Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives a policy gradient for masked diffusion LLMs as a finite-horizon MDP with stepwise advantages, and proposes entropy-guided selection of denoising steps (EGSPO) plus one-step greedy completion advantages (EGSPO-SA) for RL post-training of LLaDA and Dream.
Problem
RL post-training of diffusion language models is hard because sequence-level likelihoods are intractable, and existing surrogate-likelihood methods introduce bias and ignore per-step structure of denoising.
Approach
Formulate masked diffusion generation as a finite-horizon MDP over denoising steps, derive a policy gradient decomposing into stepwise advantages, select the highest-entropy denoising steps for gradient computation (entropy-based bound on gradient error), and estimate stepwise advantages with a one-step greedy completion available from the rollout forward pass.
Key results
- On LLaDA-8B-Instruct, EGSPO-SA reaches 94.3% on Sudoku and 85.2% on Countdown (0-shot, best across lengths), beating ESPO (72.9%, 82.0%) under the uniform protocol.
- On LLaDA-8B-Instruct coding, EGSPO-SA achieves 51.1% on MBPP (vs. 44.7% for d1 and 49.0% for ESPO) and 44.5% on HumanEval (vs. 37.8% for d1).
- On Dream-7B-Instruct, EGSPO-SA reaches 90.6% on Sudoku and 82.4% on Countdown, surpassing ESPO by 18 and 14 points.
Surprising findings
- Entropy-based step selection concentrates on intermediate timesteps rather than early heavily masked steps, and late steps are rarely selected.
- On several GSM8K configurations EGSPO without stepwise advantages slightly outperforms EGSPO-SA, indicating the biased one-step value approximation can mildly hurt.
Limitations
- Stepwise advantage uses a one-step greedy completion as value proxy, which is biased when many tokens remain masked.
- Evaluated only on two DLM backbones and with 2-token/step decoding; ESPO baseline was re-evaluated by the authors under a unified protocol, not its original one.
- Gradient treats top-k unmasking positions as fixed, an empirical design choice without full theoretical justification.
Experimental setup
Datasets: GSM8K MATH 4x4 Sudoku dataset Countdown KodCode-Light-RL-10k
Benchmarks: GSM8K MATH500 Sudoku Countdown HumanEval MBPP
Metrics: accuracy
Baselines: d1 d2 wd1 SPG ESPO TraceRL DiFFPO
Builds on: GRPO d1 ESPO DDPO
Closest papers by meaning
- DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models (NeurIPS 2026 · similarity 0.91)
- d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning (NeurIPS 2025 · similarity 0.91)
- d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation (ICML 2026 · similarity 0.90)
- MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization (NeurIPS 2025 · similarity 0.90)
- Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models (ICLR 2025 · similarity 0.90)
- Breaking AR’s Sampling Bottleneck: Provable Acceleration via Diffusion Language Models (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
