Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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.

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

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.
Paper Structure (31 sections, 4 theorems, 33 equations, 5 figures, 2 tables, 1 algorithm)

This paper contains 31 sections, 4 theorems, 33 equations, 5 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion Language Models:
  4. RL for LLMs and Reasoning:
  5. RL for Diffusion Image Models:
  6. Preliminaries and Problem Formulation
  7. Masked Diffusion Language Models
  8. Reinforcement Learning for LLMs
  9. Problem Formulation: RL for Diffusion LLMs
  10. Methodology
  11. Diffusion MDP and Policy Gradient Theorem
  12. Entropy-Guided Step Selection
  13. Estimating Stepwise Advantages
  14. From Policy Gradient to GRPO Loss
  15. Practical Implementation:
  16. ...and 16 more sections

Key Result

Theorem 1

The gradient of the objective $J({\boldsymbol{\theta}})$ (c.f. eq:ar_rl_objective) is given by where the step-level advantage is given by

Figures (5)

  • Figure 1: Overview of the performance on coding and reasoning tasks. Our approach outperforms the existing baselines in coding and logical reasoning tasks, while maintaining competitive performance in mathematical reasoning tasks.
  • Figure 2: (a) Illustration of entropy-guided denoising step selection: At each denoising step $t$, the entropy $H_t$ of the unmasking policy distribution is computed and used to identify the $K$ informative steps that have maximum entropy. In the figure, assuming $H_3 > H_1 > H_2 > H_0$ and $K=2$, the two highest-entropy steps $3, 1$ are selected for per-step policy gradient computation (marked by solid lines). (b) Illustration of stepwise advantage estimation: From state $\mathbf{x}_{t+1}$, a greedy one-step completion $\hat{\mathbf{x}}_{0 \mid t+1}$ provides a baseline reward approximating the state value. The stepwise advantage $A_t$ measures the additional reward gained by taking the denoising action at step $t$ and continuing to $\mathbf{x}_0$.
  • Figure 3: Compute efficiency comparison between EGSPO-SA and d1 on Sudoku. (a) FLOPs are accumulated over all forward passes across 8 GPUs. (b) Samples count cumulative prompt--completion pairs seen during training. (c) Gradient steps count optimizer updates (accounting for gradient accumulation). EGSPO-SA dominates d1 under all three compute budgets.
  • Figure 4: Training curves for EGSPO and EGSPO-SA on Sudoku, Countdown, GSM8K, and MATH500 using the training settings described in Section \ref{['sec:hypsetting']}.
  • Figure 5: Ablation studies analyzing (a) entropy-guided step selection (EGSPO) versus uniform step selection (USPO) and (b) the distribution of selected steps.

Theorems & Definitions (8)

  • Theorem 1: Policy Gradient Theorem
  • Remark 2
  • Lemma 3
  • Theorem 4: Policy Gradient Theorem
  • proof
  • Proposition 5: Entropy Bound
  • proof
  • proof : Proof of \ref{['lem: error_bound_using_entropy']}