Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes d2, a GRPO-style RL framework for masked diffusion LLMs with two trajectory-likelihood estimators: d2-AnyOrder (exact one-pass for any-order decoding models) and d2-StepMerge (approximate, multi-pass for standard MDMs), reporting state-of-the-art reasoning results on LLaDA-8B-Instruct without SFT.

Problem

Policy-gradient RL for masked diffusion language models requires trajectory likelihoods that are intractable to compute exactly, since a naive evaluation needs T forward passes, so prior methods use biased or heuristic likelihood estimates.

Approach

Derive a GRPO-style objective over diffusion trajectories and estimate trajectory likelihood either exactly in one pass via a packed 2L sequence with custom attention for any-order decoders, or by merging diffusion steps into N segments (block composite likelihood) with a KL error bound that shrinks as N grows.

Key results

• d2 (StepMerge, LLaDA-8B-Instruct) reaches 91.9% Sudoku, 56.6% Countdown, 85.0% GSM8K, 41.6% MATH500, versus wd1 at 25.2%, 51.2%, 82.3%, 39.0%.
• On Eso-LM toxicity steering, d2-AnyOrder reaches near-maximum toxicity (about -0.7 to -2.7 vs DDPO around -8.9) at matched FLOPs.
• On Dream 7B 0-shot, d2-StepMerge improves Sudoku from 0.79 (diffu-GRPO) to 0.91 and GSM8K from 0.73 to 0.80.

Surprising findings

• LLaDA-8B-Instruct does not natively support any-order decoding: d2-AnyOrder likelihood per token was -3.051 vs ground truth -0.128.
• diffu-GRPO fails to gain meaningful accuracy on any-order causal LLaDA, while d2-AnyOrder improves steadily.

Limitations

• d2-AnyOrder requires models that support any-order decoding, which most current masked DLMs do not satisfy by default.
• d2-StepMerge is an approximation whose error bound depends on an assumed constant epsilon_block, and the bound is only partially established in the text.
• Main comparisons are at matched FLOPs and 0-shot; several baseline numbers are taken from original papers rather than re-run.

Experimental setup

Datasets: GSM8K MATH500 Countdown Sudoku OpenWebText HumanEval MBPP

Benchmarks: Sudoku Countdown GSM8K MATH500 HumanEval MBPP

Metrics: accuracy toxicity score per-token log-likelihood KL divergence

Baselines: diffu-GRPO DDPO d1 wd1 LLaDA 1.5 LLaDA-8B-Instruct

Builds on: GRPO diffu-GRPO DDPO MDLM d1

Closest papers by meaning

• d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning (NeurIPS 2025 · similarity 0.92)
• The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language Models (ICML 2026 · similarity 0.91)
• Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models (ICLR 2026 · similarity 0.91)
• Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models (ICML 2026 · similarity 0.90)
• Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages (NeurIPS 2026 · similarity 0.90)
• Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.90)