Paper 2509.06949
Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models
- Published
- Sep 2025
- Research lab
- Independent
- Citations
- 75
- GitHub
- 516 stars
01 In brief
Summary
The paper introduces TraceRL, a trajectory-aware reinforcement learning framework for diffusion language models (DLMs) that aligns training objectives with inference trajectories, improving reasoning performance on math and coding tasks.
It includes a diffusion-based value model for training stability and supports both full-attention and block-attention architectures.
The authors release TraDo models (4B and 8B) that outperform autoregressive baselines, with TraDo-8B-Instruct achieving 6.1% relative accuracy improvement over Qwen2.5-7B-Instruct and 51.3% over Llama3.1-8B-Instruct on math benchmarks.
They also develop the first long-CoT DLM, TraDo-8B-Thinking, with an 18.1% relative gain on MATH500 over Qwen2.5-7B-Instruct.
TraceRL uses a shrinkage parameter for efficient training, a diffusion-based value model for variance reduction, and sliced training for block diffusion.
The framework supports various model architectures, accelerated inference via KV-cache, and multiple post-training methods.
Experiments show TraceRL outperforms existing RL methods, accelerates dynamic sampling, and enables block size scaling.
The open-source framework is released for building, training, and deploying diffusion LLMs.
02 From the paper
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. Code and Models: https://github.com/Gen-Verse/dLLM-RL