Elamur Research⌕
☰
home / papers / 2025 / d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

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

Paper · 2025 · Guanghan Wang, Gilad Turok, Yair Schiff et al.

Summary

Reinforcement learning is the standard way to add reasoning skills to autoregressive LLMs. It is harder to apply to masked diffusion language models (DLMs) because policy-gradient updates need the likelihood of the whole sampling trajectory. Computing that likelihood naively takes one forward pass per denoising step. The authors propose d2, a GRPO-style RL framework for masked DLMs, and derive its objective from a policy-gradient decomposition over diffusion steps. They pair it with two trajectory-likelihood estimators:

  • d2-AnyOrder gives the exact likelihood in a single pass. It only works for models that truly support any-order decoding (AO-dLLMs), and the authors show empirically that LLaDA does not support it by default.
  • d2-StepMerge works with standard MDMs. It scores only selected steps, and a theoretical bound links the number of steps to the approximation error, giving an explicit compute–bias tradeoff.

Applied to LLaDA-8B-Instruct without supervised chain-of-thought finetuning, d2 beats earlier diffusion RL methods on Sudoku, Countdown, GSM8K and MATH500 under matched FLOPs. The work shows that accurate likelihood estimation is the main bottleneck for RL on DLMs.

Key points

  • The paper is by Cornell Tech authors and appears at ICML 2026 (PMLR 306). Code and a blog post are available at guanghanwang.com/d2.
  • On AO-dLLMs, d2-AnyOrder clearly outperforms the RL baselines diffu-GRPO and DDPO.
  • With LLaDA-8B-Instruct, d2-StepMerge outperforms d1 (which uses SFT) and wd1 on all four reasoning benchmarks, setting a new DLM state of the art.
  • On coding with output length L=512, d2-StepMerge scores 40.9 on HumanEval and 45.5 on MBPP. diffu-GRPO scores 34.8 and 41.6, falling below the base LLaDA on HumanEval (37.8). This suggests better generalization to outputs longer than those used in training (L=256).
  • The authors attribute bias in prior methods to their likelihood proxies. diffu-GRPO and wd1 use all-mask logits, and coupled-GRPO uses the MDLM ELBO.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 4,145 worksReasoning · 1,948 works

Cite

Cite as Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola and Volodymyr Kuleshov. 2025. d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation. arXiv:2509.21474.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2509.21474 · arXiv 2509.21474

@misc{wang2025d2,
  title = {{d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation}},
  author = {Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola and Volodymyr Kuleshov},
  year = {2025},
  doi = {10.48550/arxiv.2509.21474},
  eprint = {2509.21474},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2509.21474}
}

References

53 references, 19 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ahmadian, A., Cremer, C., Gallé, M., Fadaee, M., Kreutzer, J., Pietquin, O., Üstün, A., and Hooker, S. Back to basics: Revisiting REINFORCE-style optimization for learning from human feedback in LLMs. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volu… find
  2. Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open-WebText Corpus. http://Skylion007.github.io/OpenWebTextCorpus, 2019. ↗
  3. Gong, S., Zhang, R., Zheng, H., Gu, J., Jaitly, N., Kong, L., and Zhang, Y. DiffuCoder: Understanding and improving masked diffusion models for code generation. In International Conference on Learning Representations (ICLR), 2026. find
  4. Arriola, M. and Kuleshov, V. Set Diffusion: Interpolating token orderings between autoregression and diffusion for fast and flexible decoding. In International Conference on Machine Learning (ICML), 2026. find
  5. Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981–17993, 2021a.
  6. Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Rı́o, J. F., Wiebe, M., Peterson, P., Gérard-Marchant, P., Sheppard, K., Re… — Array Programming with NumPy
  7. Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021b.
  8. Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., Berg, R. v. d., and Salimans, T. Autoregressive diffusion models. In International Conference on Learning Representations (ICLR), 2022. find
  9. Ashz, N. star-sft-intellect-instruct-3. https://huggingface.co/neginashz/ star-sft-intellect-instruct-3, 2024. ↗
  10. Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W., et al. LoRA: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
  11. Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. Large language diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), 2025.
  12. Huang, Z., Chen, Z., Wang, Z., Li, T., and Qi, G.-J. Reinforcing the diffusion chain of lateral thought with diffusion language models. In Advances in Neural Information Processing Systems (NeurIPS), 2025. find
  13. Nisonoff, H., Xiong, J., Allenspach, S., and Listgarten, J. Unlocking guidance for discrete state-space diffusion and flow models. In International Conference on Learning Representations (ICLR), 2025. find
  14. Hunter, J. D. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 9(3):90–95, 2007. doi: 10.1109/MCSE.2007.55. — [doi:10.1109/mcse.2007.55](https://doi.org/10.1109/mcse.2007.55) find
  15. Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., and Li, C. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In International Conference on Learning Representations (ICLR), 2025. find
  16. Kakade, S. and Langford, J. Approximately optimal approximate reinforcement learning. In Proceedings of the nineteenth international conference on machine learning, pp. 267–274, 2002. find
  17. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022.
  18. Khanna, S., Kharbanda, S., Li, S., Varma, H., Wang, E., Birnbaum, S., Luo, Z., Miraoui, Y., Palrecha, A., Ermon, S., et al. Mercury: Ultra-fast language models based on diffusion. arXiv preprint arXiv:2506.17298, 2025. pandas development team, T. pandas-dev/pandas: Pandas, February 2020. URL https:/… find
  19. Li, Z., Xu, T., Zhang, Y., Lin, Z., Yu, Y., Sun, R., and Luo, Z.-Q. ReMax: A simple, effective, and efficient reinforcement learning method for aligning large language models. In International Conference on Machine Learning (ICML), 2024. find
  20. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An Imperative Style, High-Perfor… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
  21. Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let’s verify step by step. In International Conference on Learning Representations (ICLR), 2024. — Let's Verify Step by Step
  22. Lindsay, B. G. Composite likelihood methods. In Contemporary Mathematics, volume 80, pp. 221–239. American Mathematical Society, 1988. find
  23. Logacheva, V., Dementieva, D., Ustyantsev, S., Moskovskiy, D., Dale, D., Krotova, I., Semenov, N., and Panchenko, A. ParaDetox: Detoxification with parallel data. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 6804–6818, Dublin…
  24. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36: 53728–53741, 2023.
  25. Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J. T., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
  26. Ma, Y., Du, L., Wei, L., Chen, K., Xu, Q., Wang, K., Feng, G., Lu, G., Liu, L., Qi, X., et al. dInfer: An efficient inference framework for diffusion language models. arXiv preprint arXiv:2510.08666, 2025. find
  27. Sahoo, S. S., Deschenaux, J., Gokaslan, A., Wang, G., Chiu, J., and Kuleshov, V. The diffusion duality. In International Conference on Machine Learning (ICML), 2025. find
  28. Muennighoff, N., Yang, Z., Shi, W., Li, X. L., Fei-Fei, L., Hajishirzi, H., Zettlemoyer, L., Liang, P., Candès, E., and Hashimoto, T. s1: Simple test-time scaling. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 20275– 20321, 2025.
  29. Sahoo, S. S., Yang, Z., Akhauri, Y., Liu, J., Singh, D., Cheng, Z., Liu, Z., Xing, E., Thickstun, J., and Vahdat, A. Esoteric language models. In International Conference on Machine Learning (ICML), 2026. find
  30. Schiff, Y., Sahoo, S. S., Phung, H., Wang, G., Boshar, S., Dalla-torre, H., de Almeida, B. P., Rush, A., Pierrot, T., and Kuleshov, V. Simple guidance mechanisms for discrete diffusion models. In International Conference on Learning Representations (ICLR), 2025. find
  31. Gallouédec, Q. TRL: Transformer reinforcement learning. https://github.com/huggingface/trl, 2020. ↗
  32. Wang, G., Schiff, Y., Sahoo, S. S., and Kuleshov, V. Remasking discrete diffusion models with inference-time scaling. In Advances in Neural Information Processing Systems (NeurIPS), 2025. find
  33. Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. Trust region policy optimization. In International conference on machine learning, pp. 1889–1897. find
  34. Wang, Y., Yang, L., Li, B., Tian, Y., Shen, K., and Wang, M. Revolutionizing reinforcement learning framework for diffusion large language models. In International Conference on Learning Representations (ICLR), 2026. find
  35. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. find
  36. Waskom, M. L. seaborn: statistical data visualization. Journal of Open Source Software, 6(60):3021, 2021. doi: 10.21105/joss.03021. URL https://doi.org/10. 21105/joss.03021. find
  37. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. DeepSeekMath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
  38. Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022.
  39. Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id=xcqSOfHt4g. ↗
  40. Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229–256, 1992. find
  41. Singhal, R., Horvitz, Z., Teehan, R., Ren, M., Yu, Z., McKeown, K., and Ranganath, R. A general framework for inference-time scaling and steering of diffusion models. In International Conference on Machine Learning (ICML), 2025. find
  42. Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., et al. Huggingface’s transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019.
  43. Song, Y., Zhang, Z., Luo, C., Gao, P., Xia, F., Luo, H., Li, Z., Yang, Y., Yu, H., Qu, X., et al. Seed diffusion: A large-scale diffusion language model with high-speed inference. arXiv preprint arXiv:2508.02193, 2025.
  44. Wu, C., Zhang, H., Xue, S., Liu, Z., Diao, S., Zhu, L., Luo, P., Han, S., and Xie, E. Fast-dLLM: Training-free acceleration of diffusion LLM by enabling KV cache and parallel decoding. In International Conference on Learning Representations (ICLR), 2026. find
  45. Su, X., Li, X., Uehara, M., Kim, S., Zhao, Y., Scalia, G., Hajiramezanali, E., Biancalani, T., Zhi, D., and Ji, S. Iterative distillation for reward-guided fine-tuning of diffusion models in biomolecular design. In International Conference on Learning Representations (ICLR), 2026. find
  46. Yadan, O. Hydra - a framework for elegantly configuring complex applications. Github, 2019. URL https://github.com/facebookresearch/hydra. ↗
  47. Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999. find
  48. Ye, J., Xie, Z., Zheng, L., Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. Dream 7B: Diffusion large language models. arXiv preprint arXiv:2508.15487, 2025.
  49. Tang, X., Dolga, R., Yoon, S., and Bogunovic, I. wd1: Weighted policy optimization for reasoning in diffusion language models. In International Conference on Learning Representations (ICLR), 2026. find
  50. Zhan, A. Simple policy gradients for reasoning with diffusion language models. In International Conference on Machine Learning (ICML), 2026. find
  51. Varin, C., Reid, N., and Firth, D. An overview of composite likelihood methods. Statistica Sinica, pp. 5–42, 2011. find
  52. Zhao, S., Gupta, D., Zheng, Q., and Grover, A. d1: Scaling reasoning in diffusion large language models via reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2025. von Werra, L., Belkada, Y., Tunstall, L., Beeching, E., Thrush, T., Lambert, N., Huang, S., Rasul,…
  53. Contents \#\# Introduction \#\# Background 2.1 Masked Diffusion Language Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.2 Any-order Autoregressive Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2.3 Reinforcement Learning with Verifi… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2025-d2-improving-reasoning-diffusion-language/