d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

2025-09-25 · Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, Volodymyr Kuleshov arxiv

While diffusion language models (DLMs) have achieved competitive performance in text generation, improving their reasoning ability with reinforcement learning remains an active research area. Here, we introduce d2, a reasoning framework tailored for masked DLMs. Central to our framework is a new policy gradient algorithm that relies on accurate estimates of the sampling trajectory likelihoods. Because computing these likelihoods naively is computationally expensive for masked DLMs, we develop a family of estimators tailored to distinct model classes. For DLMs that support a sampling algorithm called any-order decoding, we propose d2-AnyOrder, which achieves exact trajectory likelihood with a single model pass. Through an empirical study of widely used DLMs, we show that any-order decoding is not universally supported in practice. For standard masked diffusion models, we propose d2-StepMerge, which approximates the trajectory likelihood, trading off compute for approximation accuracy in an analytically tractable manner. Empirically, d2 significantly outperforms widely-used RL baselines when applied to popular DLMs, and sets a new state-of-the-art performance for DLMs on logical reasoning tasks (Countdown and Sudoku) and math reasoning benchmarks (GSM8K and MATH500). We provide the code along with a blog post on the project page: https://guanghanwang.com/d2

📄 PDF Abstract BibTeX arXiv:2509.21474

Code (0)

등록된 구현이 없습니다.

Tasks

Reinforcement Learning Logical Reasoning Text Generation

Improved Techniques for Maximum Likelihood Estimation for Diffusion ODEs

2023-05-06 · Kaiwen Zheng, Cheng Lu, Jianfei Chen, Jun Zhu

Diffusion models have exhibited excellent performance in various domains. The probability flow ordinary differential equation (ODE) of diffusion models (i.e., diffusion ODEs) is a particular case of continuous normalizin…

Density Estimation Image Generation

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

2026-03-13 · Vishnu Teja Kunde, Fatemeh Doudi, Mahdi Farahbakhsh, Dileep Kalathil 외 arxiv

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 likel…

Reinforcement Learning Mathematical Reasoning Logical Reasoning

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

2026-05-08 · Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade, Lokesh Boominathan 외 arxiv

Diffusion large language models are a compelling alternative to autoregressive models, yet existing RL methods for diffusion treat all denoising steps as equally important and rely on biased, high-variance likelihood est…

Reinforcement Learning Mathematical Reasoning Code Generation

Mask-Aware Policy Gradients for Diffusion Language Models

2026-07-16 · Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl arxiv

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likeliho…

Mathematical Reasoning Reinforcement Learning

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

2026-01-28 · Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang 외 arxiv

Diffusion Language Models (DLMs) offer order-agnostic generation that can explore many possible decoding trajectories. However, current decoding methods commit to a single trajectory, limiting exploration in trajectory s…

Mathematical Reasoning