Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models
Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang
2025-09-09
Summary
This paper introduces TraceRL, a new method for improving diffusion language models (DLMs) – which are a type of AI that generates text – by teaching them to follow better 'paths' to arrive at correct answers. They then used this method to create a series of new models called TraDo, which perform very well on challenging tasks like math and coding.
What's the problem?
Large language models, even the advanced diffusion-based ones, often struggle with complex reasoning tasks like solving math problems or writing code. They can get stuck in inefficient or incorrect thought processes, leading to wrong answers. Existing methods for improving these models don't always address how the model *arrives* at its answer, focusing more on the final result. Also, adapting smaller, efficient models to handle more complex tasks is difficult.
What's the solution?
The researchers developed TraceRL, which is a reinforcement learning technique that rewards the model for choosing good 'trajectories' or steps during the text generation process. Think of it like guiding the model to think through a problem in a helpful way. They also incorporated a 'value model' based on diffusion to make the training process more stable. They applied this to create TraDo models of different sizes, and also used a technique called curriculum learning to train a model specifically for long, complex reasoning chains. Finally, they released all their code and models publicly.
Why it matters?
This work is important because it shows a way to significantly improve the reasoning abilities of diffusion language models, even when those models are relatively small. The TraDo models outperform much larger models on math problems, meaning you can get better performance without needing massive computing resources. The open-source release of the framework and models allows other researchers and developers to build upon this work and create even more powerful AI systems.
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