Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models
Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang
Abstract
We propose TraceRL, a trajectory-aware reinforcement learning framework for diffusion language models (DLMs) that incorporates preferred inference trajectory into post-training, and is applicable across different architectures. Equipped with a diffusion-based value model that enhances training stability, we demonstrate improved reasoning performance on complex math and coding tasks. Besides, it can also be applied to adapt block-specific models to larger blocks, which improves sampling flexibility. Employing TraceRL, we derive a series of state-of-the-art diffusion language models, namely TraDo. Although smaller than 7B-scale AR models, TraDo-4B-Instruct still consistently outperforms them across complex math reasoning tasks. TraDo-8B-Instruct achieves relative accuracy improvements of 6.1% over Qwen2.5-7B-Instruct and 51.3% over Llama3.1-8B-Instruct on mathematical reasoning benchmarks. Through curriculum learning, we also derive the first long-CoT DLM, outperforming Qwen2.5-7B-Instruct on MATH500 with an 18.1% relative accuracy gain. To facilitate reproducible research and practical applications, we release a comprehensive open-source framework for building, training, and deploying diffusion LLMs across diverse architectures. The framework integrates accelerated KV-cache techniques and inference engines for both inference and reinforcement learning, and includes implementations of various supervised fine-tuning and RL methods for mathematics, coding, and general tasks.
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 14df2b4e-796d-4d07-bb76-5d57f3818b4fCited by top-tier papers31
- MMaDA: Multimodal Large Diffusion Language ModelsLing Yang, Ye Tian, Bowen Li, Xinchen Zhang et al.NeurIPS 2025 · 255 citations
- dParallel: Learnable Parallel Decoding for dLLMsZigeng Chen, Gongfan Fang, Xinyin Ma, Ruonan Yu et al.ICLR 2026 · 64 citations
- Co-Evolving LLM Coder and Unit Tester via Reinforcement LearningYinjie Wang, Ling Yang, Ye Tian, Ke Shen et al.NeurIPS 2025 · 56 citations
- SPG: Sandwiched Policy Gradient for Masked Diffusion Language ModelsChenyu Wang, Paria Rashidinejad, DiJia Andy Su, Song Jiang et al.ICLR 2026 · 45 citations
- d3LLM: Ultra-Fast Diffusion LLM using Pseudo-Trajectory DistillationYu-Yang Qian, Junda Su, Lanxiang Hu, Peiyuan Zhang et al.ICML 2026 · 33 citations
Builds on27
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 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
- Reinforcing the Diffusion Chain of Lateral Thought with Diffusion Language ModelsZemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li et al.NeurIPS 2025 · 55 citations
- VCORE: Variance-Controlled Optimization-based Reweighting for Chain-of-Thought SupervisionXuan Gong, Senmiao Wang, Hanbo Huang, Ruoyu Sun et al.ACL 2026
- AutoTool: Dynamic Tool Selection and Integration for Agentic ReasoningJiaru Zou, Ling Yang, Yunzhe Qi, Sirui Chen et al.ICML 2026 · 4 citations
- MetaAct-RL: Training Language Models for Reasoning Through Meta-Action-Based Reinforcement LearningZhiheng Xi, Yuhui Wang, Yiwen Ding, Guanyu Li et al.AAAI 2026
- LaDiR: Latent Diffusion Enhances LLMs for Text ReasoningHaoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Nicklas Majamaki et al.ICLR 2026 · 25 citations
