🤖 machine learning

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

The paper introduces d2, a novel reinforcement learning framework for masked diffusion language models that employs specialized trajectory likelihood estimators (d2-AnyOrder and d2-StepMerge) to significantly enhance reasoning capabilities on logical and mathematical benchmarks, achieving new state-of-the-art performance.

Original authors: Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, Volodymyr Kuleshov

Published 2026-06-02✓ Author reviewed ⓘ

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, Volodymyr Kuleshov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented artist who can paint a masterpiece by starting with a blank canvas and gradually adding details until the image is complete. This is how Diffusion Language Models (DLMs) work. Unlike traditional AI writers that build sentences one word at a time (like stacking bricks), these models start with a jumbled mess of "noise" (like static on an old TV) and slowly "denoise" it, revealing a coherent story or solution step-by-step.

The paper introduces a new training method called d2 to teach these artists how to solve hard puzzles (like math problems or logic games) much better. Here is the breakdown of how it works, using simple analogies.

The Problem: The "Blind" Critic

To teach an AI to reason, researchers use Reinforcement Learning (RL). Think of this as a game where the AI tries to solve a puzzle, and a "Critic" (a reward system) gives it a score. If the AI gets a good score, it learns to do that again.

However, there's a catch with Diffusion models. To learn effectively, the AI needs to know exactly how likely it was to generate that specific sequence of words.

For traditional AI (Autoregressive): It's like writing a letter. You know exactly what you wrote first, second, and third. Calculating the "likelihood" is easy and fast.
For Diffusion AI: It's like watching a sculpture emerge from a block of stone. The AI makes thousands of tiny adjustments. To know the exact probability of the final statue, you'd theoretically have to rewind the video and replay every single frame of the sculpting process. Doing this for every training step is computationally impossible (too slow and expensive).

Because they couldn't calculate this "score" accurately, previous methods were guessing, leading to poor reasoning skills.

The Solution: d2 (The Smart Scorekeeper)

The authors created d2, a framework that acts like a super-efficient scorekeeper. It figures out the "likelihood" of the AI's journey without having to replay the whole movie every time. They built two different tools for two different types of artists:

Tool 1: d2-AnyOrder (The "Magic Mirror")

Some diffusion models are special; they can reveal the final image in any order. You could paint the eyes first, then the nose, then the background, or vice versa, and it still works.

The Analogy: Imagine you have a magic mirror. Instead of watching the painting happen step-by-step, you can look at the finished painting and instantly see exactly what the artist would have said at every single step, all at once.
The Result: This tool, called d2-AnyOrder, calculates the perfect score in a single glance (one computer pass). It's incredibly accurate and makes the AI learn very fast.

Tool 2: d2-StepMerge (The "Chunking Strategy")

Most popular diffusion models (like the famous LLaDA) can't paint in any order; they have to follow a specific sequence. They don't have the "magic mirror."

The Analogy: Since we can't see the whole movie at once, we break the movie into chunks. Instead of watching every single frame, we watch the movie in 10-second segments. We estimate the score for each segment and add them up.
The Result: This tool, called d2-StepMerge, is an approximation. It's not perfect, but the authors proved mathematically that the error is small and controllable. It trades a tiny bit of accuracy for a massive gain in speed, making it practical for standard models.

The Results: From Novice to Grandmaster

The researchers tested these tools on some of the hardest puzzles for AI:

Logic Games: Sudoku and Countdown (a math game).
Math Problems: GSM8K and MATH500 (standard benchmarks for math reasoning).

The Outcome:

Without any extra "cheat sheets" (Supervised Fine-Tuning): The d2 framework taught the models to reason better than any previous method.
Beating the Competition: On the Sudoku and Countdown tests, the d2-trained models achieved scores near 92% and 56% respectively, crushing the previous best methods which were stuck around 22% and 42%.
Math Mastery: On complex math benchmarks, d2 set a new "State-of-the-Art" record, proving that diffusion models can be just as good at reasoning as traditional text-generating models, but with the added benefits of being faster and more controllable.

Summary

The paper says: "We found a way to teach Diffusion AI models how to think clearly."

They realized the old way of grading the AI's work was too slow and inaccurate.
They invented d2, a new grading system that is either perfectly accurate (for special models) or smartly efficient (for standard models).
Using this new system, the AI models learned to solve logic and math puzzles significantly better than before, without needing to be pre-taught the answers.

In short, d2 gives diffusion models the ability to "look back" at their own thinking process accurately, allowing them to learn from their mistakes and become true reasoning experts.

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

Problem Statement

Diffusion Language Models (DLMs), particularly masked DLMs, have emerged as competitive alternatives to autoregressive (AR) models due to their controllability and parallel generation capabilities. However, enhancing their reasoning abilities via Reinforcement Learning (RL) remains challenging. While RL has become the standard for inducing reasoning in AR models, applying it to DLMs is non-trivial. The core difficulty lies in the policy gradient formulation: unlike AR models where likelihoods factorize cleanly across token positions, the exact likelihood of a diffusion trajectory is computationally intractable. Naively computing these likelihoods requires $T$ forward passes (where $T$ is the number of diffusion steps), rendering standard policy gradient methods like Group Relative Policy Optimization (GRPO) computationally prohibitive. Existing approximations often introduce significant bias, leading to suboptimal policy updates.

Methodology

The authors introduce d2, a reasoning framework tailored for masked DLMs. The framework centers on a new policy gradient algorithm derived from the GRPO objective, which explicitly requires accurate estimation of sampling trajectory likelihoods. To address the computational bottleneck, the paper proposes a family of estimators tailored to distinct model classes:

1. Theoretical Foundation

The authors derive a GRPO-style objective for masked DLMs by marginalizing the likelihood over time latents and employing importance sampling. The resulting gradient decomposition (Theorem 3.1) depends on the ratio of trajectory likelihoods between the current policy $\pi_\theta$ and a stale policy $\pi_{old}$ . The key challenge is evaluating $\pi(x_{0:T})$ efficiently.

2. The Estimators

The d2 framework offers two specific estimators based on the model's architecture:

d2-AnyOrder (Exact Estimator):
- Target: DLMs that support any-order decoding (AO-dLLMs). These models allow tokens to be decoded in any order, provided the attention mask respects specific causality constraints (independent masks and order causality).
- Mechanism: It constructs a sequence of length $2L$ (concatenating the clean tokens and mask tokens) with specific positional encodings and attention masks. This allows the model to compute the exact conditional likelihoods of all tokens in a single forward pass.
- Limitation: Standard masked DLMs (like LLaDA) do not natively support any-order decoding; their training objectives do not guarantee the necessary attention properties for this estimator to be unbiased.
d2-StepMerge (Approximate Estimator):
- Target: Standard masked diffusion models (MDMs) that do not support any-order decoding.
- Mechanism: Inspired by block composite likelihood, this estimator approximates the trajectory likelihood by dividing the $T$ -step trajectory into $N$ contiguous time segments. Instead of evaluating every step, it evaluates the likelihood of tokens unmasked within each segment using a single forward pass per segment.
- Trade-off: It reduces the number of forward passes from $T$ to $N$ . The paper provides a theoretical upper bound (Theorem 4.1) on the approximation error (KL divergence), showing that the error decreases monotonically as $N$ increases, quantifying a compute-bias trade-off.

Key Contributions

d2 Framework: A principled RL post-training framework for masked DLMs that derives a GRPO objective dependent on accurate trajectory likelihood estimation.
d2-AnyOrder: An exact, single-pass likelihood estimator for AO-dLLMs, demonstrating that unbiased likelihood estimation significantly outperforms existing baselines in RL settings.
d2-StepMerge: A deterministic, multi-pass estimator for standard MDMs with analytically bounded approximation error, offering a superior efficiency-accuracy trade-off compared to previous methods like diffu-GRPO.
Empirical Validation: Extensive experiments showing that d2 achieves state-of-the-art performance on reasoning benchmarks without relying on supervised chain-of-thought (CoT) fine-tuning.

Experimental Results

The authors evaluated d2 on logical reasoning (Countdown, Sudoku) and mathematical reasoning (GSM8K, MATH500) benchmarks, as well as coding tasks and toxicity steering.

Performance on AO-dLLMs: When applied to models supporting any-order decoding (e.g., Eso-LM, Set Diffusion, and a custom Any-Order Causal LLaDA), d2-AnyOrder significantly outperformed baselines like DDPO and diffu-GRPO. For instance, on GSM8K with a Qwen3-1.7B base, d2-AnyOrder achieved 67% accuracy compared to 63% for diffu-GRPO.
Performance on Standard MDMs: Applied to LLaDA-8B-Instruct (a standard masked DLM), d2-StepMerge (with $N=16$ $N = 16$ ) set new state-of-the-art results:
- Sudoku: 91.9% (vs. 22.1% for d1 with SFT).
- Countdown: 56.6% (vs. 42.2% for d1).
- GSM8K: 85.0% (vs. 82.1% for d1).
- MATH500: 41.6% (vs. 39.0% for wd1).
- Notably, these results were achieved without supervised fine-tuning on extra CoT data, whereas some baselines relied on such data.
Generalization: The method generalized to other architectures (Dream 7B) and tasks (HumanEval, MBPP), showing improved length generalization compared to diffu-GRPO.
Ablation: Experiments confirmed that exact likelihoods (d2-AnyOrder) provide higher-fidelity signals than approximations, and that d2-StepMerge with an optimal $N$ (e.g., 16) balances compute cost and accuracy better than extreme values.

Significance and Claims

The paper claims that accurate trajectory likelihood estimation is central to the success of RL for reasoning in DLMs. The authors argue that previous methods failed to achieve optimal performance because they relied on biased or computationally expensive likelihood estimates.

Theoretical Contribution: The work provides the first rigorous derivation of a GRPO objective for masked DLMs and establishes theoretical bounds on the approximation error of step-merging strategies.
Practical Impact: d2 demonstrates that diffusion models can achieve reasoning capabilities competitive with, or superior to, autoregressive models when properly trained with RL, even without extensive supervised CoT data.
Future Outlook: The authors note that while d2-AnyOrder is superior, its applicability is currently limited to models specifically trained for any-order decoding. They view the development of stronger general-purpose any-order dLLMs as a natural next step to fully leverage the exact estimator.

The paper concludes that d2 sets a new state-of-the-art for DLMs on logical and mathematical reasoning tasks, validating the efficacy of their likelihood estimation strategies in unlocking the reasoning potential of diffusion-based language models.