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

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

Most language models today generate text one word at a time. Diffusion language models are a recent alternative that produces multiple words in parallel, enabling faster text generation, and show promise for tasks requiring careful reasoning. In this paper, we introduce d2, a framework for training diffusion language models to reason better using reinforcement learning. Our key insight is that accurately estimating the model's predictions during training is essential for learning, and we develop two algorithms tailored to different types of diffusion language models. Applied to popular open-source models, d2 achieves state-of-the-art performance on math problems (GSM8K, MATH500) and logic puzzles (Sudoku, Countdown), without needing expensive supervised training data.