TraceRL: Trajectory-Aware RL in Diffusion LMs

• TraceRL is a trajectory-aware reinforcement learning framework for diffusion language models that segments inference into discrete trace steps to enable fine-grained reward assignment.
• It employs a diffusion-based value model to compute low-variance advantage estimates, thereby stabilizing optimization and improving performance on reasoning and coding tasks.
• The framework adapts block-attention architectures and integrates curriculum learning to support long chain-of-thought reasoning, achieving significant accuracy gains on multiple benchmarks.

TraceRL is a trajectory-aware reinforcement learning framework for diffusion LLMs (DLMs), designed to incorporate preferred inference trajectories into post-training and to stabilize optimization through a diffusion-based value model. By aligning reinforcement learning policy updates with the generation process of DLMs, TraceRL delivers state-of-the-art performance on mathematical reasoning and coding benchmarks, enables flexible adaptation of block-attention architectures, and is disseminated via a comprehensive open-source framework for research and deployment.

1. Core Architecture and Mechanisms

TraceRL specializes in reinforcement learning for diffusion-based LLMs by emphasizing the alignment of rewards and policy updates with the actual inference trajectory. The framework decomposes generation into intermediate "trace steps," which correspond to groups of tokens produced at each decoding round. Rather than evaluating entire sequences post-hoc, rewards and policy gradients are computed along the trajectory, providing more granular control.

A diffusion-based value model estimates advantages and returns at both token and trajectory-step granularity. This mechanism aggregates rewards over the sequence, providing a variance-reducing baseline effect, functionally analogous to generalized advantage estimation (GAE) in autoregressive RL. The learning objective is formulated as:

$$J_{\text{policy}}(\theta_p) = \mathbb{E}\left[\sum_i \sum_t \sum_{ok} C_{eps}(r_k \cdot A_i / |\tau_i^{(s)}(t)|)\right] - \beta\, \mathrm{KL}[\pi_\theta / \pi_{\text{old}}]$$

where $C_{eps}$ is a clipped objective function, $A_i$ comprises token- and step-wise advantages, and $|\tau_i^{(s)}(t)|$ is the trajectory length aggregated by the shrinkage parameter $s$.

A "shrinkage parameter" ($s$) aggregates every $s$ adjacent steps of the trajectory, reducing the need for frequent value computation and lowering forward-pass overhead.

2. Trajectory-Awareness and Diffusion Value Modeling

The core innovation of TraceRL is the injection of trajectory awareness into both the reward signal and the policy optimization loop. By segmenting inference into discrete trace steps, the method captures the temporal evolution of model states and allows explicit reward assignment to partial sequences. The diffusion-based value model supplies stable and low-variance advantage estimates by smoothing returns over the trajectory.

This design ensures that training objectives are closely aligned with the actual left-to-right generation behavior of the DLM. The step- and token-wise baselining reduces instability and variance in RL updates, which is particularly critical given the stochasticity and sparsity of reward signals in reasoning tasks.

3. Performance and Benchmark Results

TraceRL yields substantial accuracy improvements on mathematical reasoning and coding tasks. Empirical results demonstrate the following:

• TraDo-4B-Instruct, trained with TraceRL and possessing only 4B parameters, consistently outperforms larger 7B-scale autoregressive (AR) models on complex math reasoning tasks.
• TraDo-8B-Instruct achieves a relative accuracy improvement of 6.1% over Qwen2.5-7B-Instruct and 51.3% over Llama3.1-8B-Instruct on mathematical benchmarks.
• The application of TraceRL in coding tasks, as evaluated on LiveCodeBench-V2 and LiveBench, results in superior performance compared to both baseline diffusion models and select AR LLMs.

Curriculum learning, combined with TraceRL, enables the creation of the first long chain-of-thought (long-CoT) diffusion LLM (TraDo-8B-Thinking), which attains an 18.1% relative accuracy boost on the MATH500 benchmark.

Model	Benchmark	Relative Accuracy Gain
TraDo-8B-Instruct	Qwen2.5-7B	6.1%
TraDo-8B-Instruct	Llama3.1-8B	51.3%
TraDo-8B-Thinking	MATH500	18.1% over Qwen2.5-7B-Instruct

These results indicate that trajectory-aware RL mechanisms and diffusion-based value modeling significantly amplify model performance on tasks requiring step-by-step reasoning.

4. Model Adaptation and Block Attention Flexibility

TraceRL supports the adaptation of block-attention DLMs to larger block sizes. For instance, a model trained with a fixed block length $B = 4$ is adapted through TraceRL to $B = 8$, expanding the model's inference window and sampling flexibility. This adaptation allows for more parallel token generation during dynamic sampling phases, with quantitative assessments showing retention of reasoning task performance and increased generation speed.

Such flexibility renders TraceRL well-suited for large-scale deployment scenarios where inference speed and efficiency must be balanced with reasoning ability.

5. Curriculum Learning and Chain-of-Thought Reasoning

TraceRL leverages a curriculum learning approach for training long-CoT DLMs. By progressively introducing more challenging chain-of-thought instances, the framework enables DLMs to generalize to extended multi-step reasoning without compromising inference efficiency. Empirical results for TraDo-8B-Thinking, based on this curriculum, demonstrate that trajectory-aware RL can robustly manage both short and long reasoning chains.

This capacity broadens the scope of practical applications, allowing for robust automated theorem proving, educational tutoring, and decision-making agents capable of multi-step reasoning.

6. Open-Source Framework and Technical Infrastructure

An open-source TraceRL framework is provided for reproducible research and industrial deployment (https://github.com/Gen-Verse/dLLM-RL). Key features include:

• Support for full-attention, block-attention, and adapted AR DLMs, allowing experimentation across architectural variants.
• Integration of accelerated KV-cache techniques and inference engines, with tunable window sizes for KV-cache acceleration.
• Implementation of post-training methods, including random masking, semi-autoregressive fine-tuning, and multiple RL algorithms (TraceRL, coupled RL, random-masking RL).
• End-to-end training and evaluation pipelines for mathematics, coding, and general tasks.

These technical resources are targeted toward practitioners seeking to build, train, and deploy diffusion LLMs under a unified trajectory-aware RL paradigm.

7. Practical Applications and Significance

TraceRL's improvements in training stability and inference trajectory control translate into enhanced practical capabilities for diverse reasoning-intensive domains. Applications include:

• Step-wise mathematical problem solving and education (e.g., via superior chain-of-thought reasoning on MATH500, AIME2024).
• High-veracity code generation and verification workflows (e.g., improvements documented on LiveCodeBench-V2, LiveBench).
• Flexible deployment for real-time systems needing rapid block-based sampling and accelerated inference, critical for interactive coding assistants, tutoring systems, and automated theorem provers.

By unifying trajectory awareness, diffusion-based value estimation, and reinforcement learning, TraceRL yields robust, highly adaptive diffusion LLMs that consistently outperform autoregressive baselines on reasoning-centric benchmarks and provide a reproducible toolkit for future research.