PaperlayerView exact revision on arXiv

Source-linked AI summary

Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models

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

arXiv:2509.06949v1cs.CL

TL;DR

Diffusion language model post-training can misalign with the trajectories used during inference, especially across differing architectures. TraceRL trains on preferred trajectories with a diffusion-based value model, yielding strong math and coding results and new long-CoT models. The framework also supports flexible block sizes, accelerated inference, and reproducible development, while further process-reward exploration remains open.

  • Problem

    Post-training objectives can mismatch diffusion models’ inference behavior, while reinforcement learning remains underexplored for block diffusion and across architectures.

  • Method

    TraceRL aligns reinforcement learning with intermediate inference traces and uses a diffusion-based value model across diverse diffusion language model architectures.

  • Results

    TraceRL produces state-of-the-art diffusion language models across reasoning tasks, including TraDo models that surpass strong autoregressive baselines and a first long-CoT diffusion model.

  • Takeaways & Limitations

    TraceRL supports diffusion language model reasoning, larger block sizes, accelerated inference, and an integrated framework for building, training, and deployment.

  • Takeaways & Limitations

    Further exploration of process rewards for TraceRL optimization remains an important direction for future work.

Abstract

from arXiv · show

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

1. Introduction

Diffusion language models have advanced across coding and reasoning, but reinforcement learning remains insufficiently unified across architectures. TraceRL addresses this gap by aligning post-training with inference trajectories and producing strong reasoning models.

  • Diffusion language models have been explored for coding, fixed-format reasoning, and complex reasoning tasks across full- and block-attention architectures.
  • A unified and effective reinforcement learning framework spanning different diffusion language model architectures remains underexplored.
  • Existing post-training methods can mismatch optimal inference because random masking does not reflect language’s sequential and logical structure.
  • TraceRL yields 4B- and 8B-sized state-of-the-art diffusion instruction models that surpass strong autoregressive models on mathematical reasoning benchmarks.
  • 6.1% and 51.3% relative accuracy gains are reported for TraDo-8B-Instruct over Qwen2.5-7B-Instruct and Llama3.1-8B-Instruct, respectively.
  • TraceRL aligns training objectives with sampling trajectories, uses a diffusion-based value model for stability, and supports diverse architectures and applications.

2. Preliminaries

The preliminaries describe masked diffusion objectives, autoregressive-model adaptation, block diffusion, and iterative decoding. These components provide the architectural and inference context for trajectory-aware reinforcement learning.

  • 2.1. Masked Diffusion Language Models with Full Attention: Masked diffusion language models replace selected tokens with [MASK] tokens to form a partially masked sequence x_t from raw data x_0.
  • 2.1. Masked Diffusion Language Models with Full Attention: DLM training objectives can be derived from data likelihood through an evidence lower bound on log p_θ(x).
  • 2.2. Adapted From AR Models: Adapting from an autoregressive model can preserve AR capabilities while requiring only adaptation training, which is more efficient than training from scratch.
  • 2.3. Block Diffusion: Block diffusion combines autoregressive training efficiency with diffusion sampling efficiency and naturally supports KV-cache through block-wise generation.
  • 2.3. Block Diffusion: Fixed block size limits sampling flexibility because small blocks generate fewer tokens at once and can reduce speed relative to full-attention diffusion models.
  • 2.4. Decoding Strategy and Accelerated Inference: Iterative decoding progressively unmasks tokens using confidence-based selection, with static and threshold-based dynamic sampling as common strategies.
  • 2.4. Decoding Strategy and Accelerated Inference: Figure 2 contrasts efficient semi-autoregressive fine-tuning for block diffusion with data slicing by block size for full-attention models.

3. Mismatch Between Post-Training Objective and Inference Trajectory

The section shows that random-mask post-training can diverge from left-to-right inference, while semi-autoregressive and preferred-trace training better align optimization with actual decoding. Reinforcement learning offers a practical way to collect such traces during rollouts.

  • Fully random masking enables parallel decoding but mismatches inference behavior when decoding uses left-to-right block generation with KV-cache.
  • 3.1. Semi-Autoregressive Fine-Tuning: Semi-autoregressive fine-tuning trains later tokens conditioned on earlier context while retaining diffusion sampling efficiency.
  • 3.1. Semi-Autoregressive Fine-Tuning: For full-attention DLMs, semi-autoregressive training increases forward passes by approximately ⌈L/B⌉, yet outperforms equally costly independent random-mask repetitions.
  • 3.1. Semi-Autoregressive Fine-Tuning: Table 1 evaluates CoT-reasoning fine-tuning on MATH500 under non-CoT prompts using block-attention SDAR-4B-Chat and full-attention Dream-7B-Instruct.
  • 3.2. Aligning Post-training with Preferred Inference Traces: The model’s preferred inference traces achieve optimal performance over baselines at equal or lower computational load for both block- and full-attention structures.
  • 3.3. Reinforcement Learning Naturally Leverages Inference Traces: Collecting inference traces for fine-tuning requires substantial effort, whereas reinforcement learning naturally produces traces during rollouts.

4. RL Training with Trajectory

TraceRL performs reinforcement learning on diffusion-model inference trajectories rather than only final sequences, while shrinking trajectories and using a diffusion-based value model for efficient, stable training. For block diffusion, sliced block-attention training extends this approach to policy and value models with greater parallelism.

  • Trajectory-aware reinforcement learning: TraceRL rewards or penalizes intermediate inference traces according to verifiable response rewards, enabling trajectory-aware reinforcement learning across DLM architectures.The trajectory records tokens decoded at each inference step, and process-level rewards can also be incorporated.
  • Accelerated Training with Shrinkage Parameter: Aggregating every s neighboring inference steps reduces the number of training forward passes while preserving a trajectory-based policy objective.The shrinkage parameter produces a shorter trajectory with length ⌈|τ_i|/s⌉ and reduces training computation complexity by a factor of s.
  • Diffusion-based value model: A diffusion-based value model estimates step-wise values and token-wise advantages, providing a variance-reducing baseline for policy updates.The value network conditions on each trace-step prefix, predicts values for tokens generated at that step, and uses a stop-gradient frozen baseline when constructing advantages.
  • Diffusion-based value model: TraceRL maps token-wise rewards to step-level returns and trains the value network with clipped regression while updating the policy using token advantages.Step-wise rewards are mapped back to tokens for return construction, and the explicit reward and advantage forms are given in the paper’s proposition.
  • Sliced Training in Block Diffusion: For block diffusion, traces are organized into blocks and sliced so each slice is forwarded once with block attention, enabling parallel policy and value-model training.The maximum slice size is bounded by ⌈B/s⌉, and the formulation is described as more efficient than full-attention training.

5. Experiments

TraceRL is evaluated across block- and full-attention diffusion models on mathematics and coding, with static and dynamic sampling analyses. The experiments report strong reasoning performance, improved optimization, sampling acceleration, and applications to larger block sizes.

  • Experimental Setup: Experiments cover mathematics and coding using full-attention Dream models and block-attention SDAR models with block size 4.Evaluation includes GSM8K, MATH500, AIME, LiveCodeBench-V2, and LiveBench.
  • Instruction Models: TraceRL-trained TraDo instruction models achieve state-of-the-art reasoning performance among current diffusion models, with TraDo-4B-Instruct outperforming Qwen2.5-7B-Instruct across all math tasks.Both dynamic, faster sampling and static, more accurate sampling improve significantly.
  • Long-CoT Model: 85.8% accuracy on MATH500 is achieved by TraDo-8B-Thinking, the first long-CoT diffusion language model derived using long-CoT SFT and TraceRL.The model demonstrates strong reasoning capabilities across benchmarks.
  • Training Dynamics: 5.4% static and 4.2% dynamic MATH500 improvements are obtained by TraDo-4B-Instruct, while TraDo-8B-Instruct improves 4.2% static and 4.8% dynamic.TraDo-4B-Instruct surpasses Qwen2.5-7B-Instruct after optimization.
  • RL Ablations: TraceRL achieves the best optimization performance with and without a value model and remains superior under both dynamic and static evaluation in block diffusion.The comparison supports optimizing over the preferred inference trace, including within a small block.
  • Coding Tasks: 25.0% accuracy on LiveCodeBench-V2 establishes a new state of the art among open-source full-attention diffusion language models.TraceRL converges faster and achieves the best performance compared with the baselines.
  • Sampling Analysis: 15.4% speedup on MATH500 is achieved for dynamic sampling after TraceRL optimization.The analysis also reports increased average response length on complex math reasoning tasks.

6. Open-source Framework

The paper releases an open-source framework for training and deploying large diffusion language models. It supports multiple architectures, accelerated inference, and several supervised fine-tuning and reinforcement learning methods.

  • Architecture Support: The released framework supports pretrained and adapted full-attention models alongside block diffusion models.Examples include LLaDA, MMaDA, Dream, DiffuCoder, SDAR, and TraDo.
  • Inference: The framework extends KV-cache acceleration for full-attention diffusion models and uses JetEngine to accelerate block-model inference.The full-attention extension adds a tunable further-horizon window size for each forward pass.
  • Training Methods: Supported training methods include fully random masking, semi-autoregressive fine-tuning, TraceRL, coupled RL, and random masking RL.Block models can optionally use a diffusion-based value model, and all methods support multi-node training.

7. Related Work

Related work establishes diffusion language models as scalable architectures with bidirectional attention and parallel decoding, while prior RL methods mainly target full-attention models. The paper positions TraceRL against unresolved alignment and cross-architecture challenges.

  • Reasoning Enhancement: Chain-of-thought prompting, self-improvement, and reinforcement learning have been used to improve language-model reasoning and complex-task performance.The cited work spans prompting, iterative refinement, and RL approaches.
  • Efficient Reasoning: Long-CoT models have strong reasoning power but prohibitively slow inference, motivating faster inference methods.Diffusion language models are presented as a promising direction because of parallel decoding and bidirectional mechanisms.
  • Diffusion RL: Prior full-attention diffusion RL methods use random masking or complementary rollout objectives, while reinforcement learning for block diffusion remains unexplored.The prior methods include PPO optimization with random masks and doubled effective training data through complementary counterparts.
  • Diffusion Language Models: Diffusion language models have been developed through continuous latent diffusion, token-space transition processes, and masked diffusion architectures.Masked diffusion models have emerged as scalable large-language-model architectures.
  • Diffusion Language Models: Masked diffusion models combine bidirectional attention for global consistency with parallel decoding for inference acceleration.These properties distinguish them from left-to-right autoregressive transformers in the cited related work.

8. Conclusion

The paper presents TraceRL as a reinforcement learning method for diffusion language models with diverse architectures, alongside a released framework for research and deployment. It reports benefits for inference acceleration, block-size scaling, and reasoning-oriented model development.

  • TraceRL is introduced as a reinforcement learning method for diffusion language models with diverse architectures.
  • Extensive experiments support TraceRL’s effectiveness across different reinforcement learning tasks, producing three state-of-the-art diffusion language models.
  • The method is reported to benefit inference acceleration and block-size scaling, motivating further exploration.
  • The authors identify combining diffusion models’ accelerated inference with strong reasoning ability as an important research avenue.
  • The released framework supports building, training, and deploying diffusion language models across diverse architectures, with accelerated inference integrated into reinforcement learning pipelines.

A. Theoretical Results

This section states token-wise return and advantage results derived from step-wise trajectory recursions, then describes the demonstration data and prompt templates used for evaluation and reinforcement learning.

  • Theoretical Results: The theoretical section introduces token-wise return and advantage quantities for trajectories partitioned into trace steps.
  • Theoretical Results: The stated expressions follow by reindexing the trace-step sums and collecting value terms into a telescoping series.
  • Theoretical Results: For λ = 0, the derivation reduces to the stated one-step temporal-difference form.
  • Theoretical Results: The special case (γ, λ) = (1, 1) corresponds to an undiscounted Monte Carlo return and advantage with a token-wise baseline.
  • Theoretical Results: The special case (γ, λ) = (1, 0) gives an undiscounted return and a one-step TD(0) advantage using the trace-level baseline.
  • Demonstration Setup: The demonstration uses chain-of-thought responses from Qwen2.5-32B-Instruct on 2,000 OpenR1-MATH tasks filtered against Qwen2.5-7B-Instruct solvability.
  • Prompt Templates: The same prompt template is used for reinforcement learning and evaluation, while TraDo and SDAR share templates and long-CoT mode uses a thinking prompt.

B.3. Evaluation Details

The evaluation specifies model-dependent decoding, response, temperature, block-size, and horizon settings, evaluates AIME2024 more repeatedly than other datasets, and uses KV-cache acceleration throughout.

  • Model Settings: LLaDA uses block size 32, response limit 1024, temperature 0.1, and further horizon size 128.
  • Model Settings: Dream uses temperature 0.1 and further horizon size 128, with response limits of 1600 for mathematics and 1024 for coding.
  • Model Settings: SDAR and TraDo retain block size 4 and response limit 2000, while TraDo-8B-Thinking uses a 30,000-token evaluation limit and dynamic sampling only.
  • Evaluation Protocol: AIME2024 problems are evaluated 20 times each, whereas all other test datasets are evaluated three times with average accuracy reported.
  • Inference Acceleration: KV-cache is used in all evaluations, with fast-dllm adapted for full-attention models and jetengine used for block diffusion models.
  • Additional Settings: A separate Dream setting uses temperature 0.8, block size 32, horizon size 128, response limit 1024, and static one-token-per-step decoding.

B.5. Training Details

Training combines architecture-specific fine-tuning settings with TraceRL optimization, trajectory-based rollout processing, verifiable rewards, and separate policy and value updates. The section also documents compute comparisons and a padding-token stability constraint.

  • Optimization Settings: Full-attention fine-tuning uses learning rate 1 × 10^-6, while adapted Dream models can use autoregressive cold-start training at 1 × 10^-5.
  • Optimization Settings: RL training uses learning rate 1 × 10^-6, β = 0.01, ε = 0.2, and a default k = 3 estimator for KL.
  • Rewards and Data: Math reinforcement learning uses binary verifiable rewards and retains tasks with accuracy between 0.2 and 0.8 for training.
  • Efficiency Comparison: TraceRL uses shrinkage s = 8 and approximately 47.5 forward passes per data point in the reported Figure 6(a) comparison.
  • Training Stability: Large n_pad values can enlarge pad-token logits and terminate inference prematurely, while n_pad = 0 may prevent long-CoT SFT outputs from stopping.
  • TraceRL Procedure: The TraceRL algorithm freezes the old policy for PPO ratios and KL calculations before performing policy and value optimization steps.
  • TraceRL Procedure: TraceRL rollouts generate trajectory traces, obtain verifiable or optional process-level rewards, construct trace-based returns and advantages, and update policy and value objectives.
  • Reported Statistics: Table 5 reports average response lengths for TraDo-8B-Thinking across different benchmarks.

C.2. Example of Long-CoT outputs

The example uses algebra to show that sums of three consecutive integers are multiples of 3, then checks candidate cubes to conclude that 27 is the smallest positive solution.

  • Algebraic setup: Three consecutive integers can be represented as x, x+1, and x+2, whose sum is 3x + 3 = 3(x + 1).This establishes that any such sum must be divisible by 3.
  • Divisibility condition: Because the sum must be divisible by 3, any matching perfect cube must also be divisible by 3.The reasoning later expresses the cube as k^3 and concludes that k must be divisible by 3.
  • Candidate checking: The candidate cubes 1 and 8 do not satisfy the divisibility condition, while 27 does.The example explicitly checks the smaller cubes before accepting 27.
  • Conclusion: The response concludes that 27 is the smallest positive perfect cube expressible as the sum of three consecutive integers.It also notes that allowing negative integers does not produce a smaller positive cube.
Loading 2509.06949v1…