wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models
Xiaohang Tang, Rares Dolga, Sangwoong Yoon, Ilija Bogunovic
Abstract
Improving the reasoning capabilities of diffusion-based large language models (dLLMs) through reinforcement learning (RL) remains an open problem. The intractability of dLLMs likelihood function necessitates approximating the current, old, and reference policy likelihoods at each policy optimization step. This reliance introduces additional computational overhead, and can lead to large variance and estimation error in RL objective -- particularly in computing the policy ratio for importance sampling. To mitigate these issues, we introduce wd1, a novel ratio-free policy optimization approach that reformulates the RL objective as a weighted log-likelihood, requiring only a single approximation for the current parametrized policy likelihood. We formally show that our proposed method can be interpreted as energy-guided discrete diffusion training combined with negative sample unlearning, thereby confirming its theoretical soundness. In experiments on LLaDA-8B model, wd1 outperforms diffusion-based GRPO (d1) while requiring lower computational cost, achieving up to a +59% improvement in accuracy. Furthermore, we extend wd1 to denoising-stepwise weighted policy optimization (wd1++), achieving state-of-the-art math performance of 44.2% on MATH500 and 84.5% on GSM8K with only 20 RL training steps.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 6c2abdce-7117-46fb-8d2f-94363be67195Cited by top-tier papers13
- Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy OptimizationYuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk et al.ICML 2026 · 13 citations
- Consolidating Reinforcement Learning for Multimodal Discrete Diffusion ModelsTianren Ma, Mu Zhang, Yibing Wang, Qixiang YeICLR 2026 · 10 citations
- Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion Large Language ModelsNianyi Lin, Jiajie Zhang, Lei Hou, Juanzi LiACL 2026 · 8 citations
- Breaking the Factorization Barrier in Diffusion Language ModelsIan Li, Zilei Shao, Benjie Wang, Rose Yu et al.ICML 2026 · 5 citations
- Lavida-R1: Advancing Reasoning for Unified Multimodal Diffusion Language ModelsShufan Li, Yuchen Zhu, Kangning Liu, Zhe Lin et al.ICML 2026 · 4 citations
Builds on30
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Let's Verify Step by StepHunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards et al.ICLR 2024 · 3,045 citations
- DAPO: An Open-Source LLM Reinforcement Learning System at ScaleQiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan et al.NeurIPS 2025 · 2,828 citations
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
Related papers
- Stabilizing Reinforcement Learning for Diffusion Language ModelsJianyuan Zhong, Wang Kaibo, Ding Ding, Zijin Feng et al.ICML 2026 · 3 citations
- Principled RL for Diffusion LLMs Emerges from a Sequence-Level PerspectiveJingyang Ou, Jiaqi Han, Minkai Xu, Shaoxuan Xu et al.ICLR 2026 · 33 citations
- d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement LearningSiyan Zhao, Devaansh Gupta, Qinqing Zheng, Aditya GroverNeurIPS 2025 · 191 citations
- Simple Policy Gradients for Reasoning with Diffusion Language ModelsAnthony ZhanICML 2026 · 4 citations
- The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language ModelsZanlin Ni, Shenzhi Wang, Yang Yue, Tianyu Yu et al.ICML 2026 · 4 citations
