Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models
Yinjie Wang Princeton University Ling Yang Bowen Li Princeton University Ye Tian Princeton University Ke Shen Mengdi Wang Princeton University
Abstract
Code: https://github.com/Gen-Verse/dLLM-RL Models: TraDo-4B/8B 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 comprehe
中文速览
扩散语言模型(Diffusion Language Model, DLM)虽然拥有并行生成速度快、双向注意力一致性好等优势,但现有的强化学习后训练方法在给序列打分时采用随机掩码,与模型实际推理时从左到右逐步揭示词语的真实轨迹严重脱节,导致优化效果不佳。为此,研究者提出了 TraceRL——一种轨迹感知强化学习框架,它直接利用模型推理时生成的中间轨迹来计算策略梯度,并配合一个扩散式价值模型来降低训练方差、提升稳定性,同时通过"收缩参数"合并相邻步骤以大幅减少前向传播次数,兼顾效率与精度。基于 TraceRL 训练出的 TraDo-4B/8B 系列模型在数学推理基准上分别比同规模的 Qwen2.5-7B-Instruct 高出 6.1%、比 Llama3.1-8B-Instruct 高出 51.3%,并首次实现了支持长链式推理(long Chain-of-Thought)的扩散语言模型。这项工作证明扩散语言模型同样可以通过强化学习达到甚至超越自回归大模型的复杂推理能力,并配套开源了完整的训练与部署框架,为该方向的后续研究提供了坚实基础。
原文 arXiv:2509.06949;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2509.06949v1