Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, Mengdi Wang

TL;DR

This work tackles the misalignment between diffusion language model post-training objectives and actual inference trajectories. It introduces TraceRL, a trajectory-aware RL framework that leverages a diffusion-based value model to stabilize training and optimizes over inference traces, applicable to both full- and block-attention models. The approach yields state-of-the-art diffusion models (TraDo) for math and coding reasoning, including a long-CoT diffusion model (TraDo-8B-Thinking), and demonstrates improved sampling speed via accelerated inference and block-size scaling. An open-source framework accompanies the method, enabling reproducible research and practical deployment across architectures. Overall, TraceRL advances diffusion-LM reasoning, sampling efficiency, and flexibility, with broad implications for scalable, explainable 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

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

TL;DR

This work tackles the misalignment between diffusion language model post-training objectives and actual inference trajectories. It introduces TraceRL, a trajectory-aware RL framework that leverages a diffusion-based value model to stabilize training and optimizes over inference traces, applicable to both full- and block-attention models. The approach yields state-of-the-art diffusion models (TraDo) for math and coding reasoning, including a long-CoT diffusion model (TraDo-8B-Thinking), and demonstrates improved sampling speed via accelerated inference and block-size scaling. An open-source framework accompanies the method, enabling reproducible research and practical deployment across architectures. Overall, TraceRL advances diffusion-LM reasoning, sampling efficiency, and flexibility, with broad implications for scalable, explainable 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

Paper Structure

This paper contains 47 sections, 1 theorem, 13 equations, 6 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Masked Diffusion Language Models with Full Attention
  4. Adapted From AR Models
  5. Block Diffusion
  6. Decoding Strategy and Accelerated Inference
  7. Mismatch Between Post-Training Objective and Inference Trajectory
  8. Semi-Autoregressive Fine-Tuning
  9. Aligning Post-training with Preferred Inference Traces
  10. Reinforcement Learning Naturally Leverages Inference Traces
  11. RL Training with Trajectory
  12. Accelerated Training with Shrinkage Parameter
  13. Value Model with Diffusion Modeling
  14. Sliced Training in Block Diffusion
  15. Experiments
  16. ...and 32 more sections

Key Result

Proposition 1

Let a trajectory $\tau$ be partitioned into trace steps $\tau(1),\dots,\tau(|\tau|)$, and let $t_j$ denote the unique step index with $j\in\tau(t_j)$. For token-wise rewards $r_j$ and token-wise values $V^{\text{old}}_j$, define the step-wise aggregates Let the step-wise return and GAE be given by Define token-wise quantities Then the following explicit expressions hold: For the boundary case

Figures (6)

  • Figure 1: Left: RL training dynamics with different methods, where our TraceRL achieves the best optimization. Right: Benchmark results on complex math reasoning tasks and LiveCodeBench-V2, all evaluated with KV-cache. TraDo series models outperform strong AR models even at smaller scales.
  • Figure 2: Semi-AR SFT for full attention and block attention model Overview. This example uses a block size of $B = 2$ and a sequence length of $L = 6$. Block diffusion models naturally adapt to semi-AR SFT efficiently, whereas full attention models require slicing the data by $B$.
  • Figure 3: TraceRL Overview (detailed pipeline in \ref{['algapp']}). This is an example for $s = 2$, $L = 6$ and $B = 3$. We aggregate every $s$ neighboring steps to perform trajectory-aware reinforcement learning. The integers in these squares represent the sequence of the policy’s inference process.
  • Figure 4: The TraceRL training curves for the 4B and 8B models on the math task. The red curve denotes the dynamic sampling accuracy, which achieves faster sampling speed, while the blue curve denotes the static sampling accuracy, which achieves higher accuracy. The 4B model is trained with a value model, whereas the 8B model is trained directly using $\mathcal{J}_{\text{policy}}$.
  • Figure 5: RL method ablations on block diffusion models and math RL tasks. The red and yellow curves represent TraceRL with and without a value model, respectively. The blue curve corresponds to training with a random masking objective within block, similar to the semi-autoregressive training approach. The green curve represents training with an additional complementary mask within block.
  • ...and 1 more figures

Theorems & Definitions (3)

  • Proposition 1: Token-wise return and advantage from step-wise recursions
  • proof
  • Remark 1: Special cases