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

2026-03-12

AI Summary

Existing reinforcement learning approaches struggle to directly optimize diffusion language models due to the intractability of sequence likelihood, often resorting to biased approximations that disregard the temporal structure of the denoising process. This work formulates diffusion-based generation as a finite-horizon Markov decision process and introduces, for the first time, an unbiased policy gradient method tailored to this framework. By recursively decomposing the advantage function, incorporating an entropy-guided mechanism for selecting denoising steps, and estimating intermediate advantages from single-step denoising rewards—thereby avoiding multi-step rollouts—the method effectively preserves trajectory-level temporal information. Empirically, it achieves state-of-the-art performance on code generation and logical reasoning benchmarks and significantly outperforms existing reinforcement learning post-training approaches for diffusion language models on mathematical reasoning tasks.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Computer Vision: Diffusion Models for Vision

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Reinforcement Learning
• Diffusion Language Models
• Sequence-level Likelihood
• Policy Gradient
• Denoising Trajectory

Innovation

• Diffusion Language Models
• Reinforcement Learning
• Entropy-Guided Step Selection
• Stepwise Advantages
• Policy Gradient