TraceRL Framework

• TraceRL is a trajectory-aware reinforcement learning framework that recasts diffusion language model decoding as a Markov decision process to assign credit at the token level.
• It implements curriculum-based block-size scaling by progressively doubling tokens unmasked per batch, enhancing parallelism without sacrificing performance.
• TraceRL improves model stability and reasoning in math and coding domains by integrating PPO-based policy optimization with temporal-difference learning.

The TraceRL framework is a trajectory-aware reinforcement learning (RL) methodology designed for post-training and optimizing diffusion LLMs (DLMs) and masked diffusion LLMs (MDMs) using multi-step inference traces. By recasting DLM/MDM decoding as a Markov decision process (MDP), TraceRL enables step-level credit assignment to all tokens unmasked throughout the generation process—significantly improving the stability and reasoning capacity of block-based and diffusion-based LLMs in mathematical and coding domains. TraceRL has facilitated the development of high-performing models such as the TraDo series and enabled progressive block-size scaling curricula as exemplified by the T$^\star$ approach.

1. Markov Decision Process Formulation

TraceRL formalizes DLM/MDM decoding as an MDP, where the state, action, and reward are defined over the denoising trajectory rather than terminal outputs (Wang et al., 8 Sep 2025, Xia et al., 16 Jan 2026):

• State at diffusion step $t$ is $s_t = (\tau_{<t}, Q)$, where $Q$ is the fixed prompt and $\tau_{<t}$ denotes all position–token pairs finalized prior to $t$.
• Action at step $t$ involves selecting $\tau(t)$: a set of masked positions to finalize. Each action $o \in \tau(t)$ corresponds to revealing a specific token, with policy $\pi_\theta(o \mid \tau_{<t}, Q)$ describing the distribution across all such positions.
• Transition occurs as tokens are revealed; although parallel unmasking is used for efficiency, TraceRL conceptually treats each token's unveiling as a unique action for granular credit assignment.
• Reward is generally a verifiable, sequence-level scalar (e.g., correctness of solution), broadcast to all tokens unless token-level metrics or heuristics are available. Sequence-level rewards are distributed back along the trajectory using temporal-difference (TD) learning or Generalized Advantage Estimation (GAE).

This trajectory-centric formulation aligns the RL optimization objective tightly with the model’s inference mechanism, allowing improved policy learning within the actual operational constraints of masked diffusion inference.

2. Policy Objective and Value Estimation

The TraceRL framework extends Proximal Policy Optimization (PPO) to optimize rewards over entire denoising trajectories (Xia et al., 16 Jan 2026, Wang et al., 8 Sep 2025). For a trajectory $\tau$, the surrogate policy loss to be minimized is

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta_{\mathrm{old}}}} \left[ \sum_{t=1}^{T} \frac{1}{|\tau(t)|} \sum_{o\in\tau(t)} C_\epsilon \bigl(\rho_t(o),A(o)\bigr) \right] - \beta\, \mathrm{KL}\!\bigl(\pi_\theta\|\pi_{\theta_{\mathrm{old}}}\bigr)$$

where $C_\epsilon$ is the PPO clipping function and $\rho_t(o)$ is the policy ratio. Step-level advantages $A(o)$ are computed by first aggregating value predictions over each denoising step, applying GAE/TD to obtain a step-advantage $\hat{A}_t$, and then distributing $\hat{A}_t$ to every token in $\tau(t)$.

The value function $V_\phi$ is trained to regress from denoising step prefixes to clipped TD-return estimates, producing low-variance step- and token-level baselines that further stabilize policy updates.

3. Progressive Block-Size Scaling and Curriculum Learning

TraceRL underpins curricula that progressively scale the block-size ($B$) used for parallel token unmasking in DLM/MDM inference. T$^\star$ (Xia et al., 16 Jan 2026) is a prominent curriculum that begins with a small, autoregressive-initialized block size (e.g., $B_0 = 4$) and iteratively doubles $B$ through staged epochs.

At each curriculum stage:

• 50% of batches use block-aligned TraceRL rollouts, unmasking tokens in contiguous $B$-sized blocks.
• The other 50% use TraceRL on block-shifted data (offset by $\Delta = B/2$), exposing the model to cross-block dependencies.
• After each stage, adjacent blocks are merged, and $B \leftarrow 2B$.
• The process continues until a target maximum block size $\hat{B}$ is reached (practically $B=32$; larger sizes degrade stability).

This block-scaling loop enables higher-parallelism decoding with minimal performance loss, and prepares models for efficient, scalable reasoning or generation.

4. Empirical Findings and Alternative Decoding Schedules

TraceRL-trained models, including T$^\star$, do not revert to strict left-to-right (canonical) decoding as block size increases. Instead, they adopt a non-canonical schedule ($\hat{S}$) that preserves high local monotonicity while trading off maximal parallelism against sequential dependencies. The LocalStrict metric quantifies the closeness to purely monotonic scheduling:

$\mathrm{LocalStrict} = \frac{1}{n} \sum_{k=1}^{n} \mathbbm{1}\!\left[\pi_k = \min_{j \geq k} \pi_j\right]$

Empirically, models trained via T$^\star$ reach $\mathrm{LocalStrict} \approx 0.80$–$0.85$ across block sizes while matching or exceeding the math reasoning performance of standard autoregressive or AR-initialized models—for example, improving MATH500 Pass@3 from 55.9% to 63.4% at $B=8$ (Xia et al., 16 Jan 2026), and achieving +6.1% to +51.3% relative gains on MATH500 for TraDo-8B Instruct over Qwen2.5-7B and Llama3.1-8B (Wang et al., 8 Sep 2025).

5. Implementation and Practical Guidance

TraceRL and T$^\star$ can be instantiated with the following empirically validated configurations (Xia et al., 16 Jan 2026, Wang et al., 8 Sep 2025):

Hyperparameter	Typical Value	Usage
Policy LR	$1\times10^{-6}$	AdamW optimizer for policy updates
Value LR	$5\times10^{-6}$	Value optimizer (if learned)
KL Penalty ($\beta$)	0.01	Regularize divergence from base
PPO Clipping ($\epsilon$)	0.1–0.2	Prevent extreme ratio updates
Discount Factor ($\gamma$)	$1.0$	Undiscounted, sparse rewards
GAE Parameter ($\lambda$)	0.95 or 1.0	Value-advantage estimation
Batch Size	128 tasks × 16–32 rollouts	Parallel RL sampling

For curriculum execution:

1. Initialize from an AR-trained small-block checkpoint.
2. Alternate TraceRL rollouts on aligned and shifted blocks.
3. Double the block size stagewise, carrying forward trained policy parameters.
4. Maintain static or dynamic sampling, batch slicing, and parallel inference as per system specifications.

The TraceRL framework and its reference implementation are available at https://github.com/Gen-Verse/dLLM-RL (Wang et al., 8 Sep 2025).

6. Applications and Comparative Analysis

TraceRL’s trajectory-based credit assignment, diffusion-based value modeling, and block-scaling curricula enable:

• Adaptation of DLMs/MDMs to larger block sizes without significant performance loss.
• Efficient long chain-of-thought (Long-CoT) reasoning through staged curriculum learning.
• Outperformance of standard AR RL baselines and random-masking PPO in mathematical and coding evaluation benchmarks.

TraceRL’s empirically demonstrated variance reduction and efficient credit assignment (via grouped macro-steps and value baselines) allow models like TraDo-8B-Instruct to achieve substantial empirical improvements over comparably sized AR models (Wang et al., 8 Sep 2025). Algorithms based on TraceRL stabilize learning and generalization over long contexts and complex reasoning traces, supporting advanced deployments in mathematics and code generation.

7. Reproducibility and System Architecture

A comprehensive open-source framework supports TraceRL post-training, including policy/value optimizer modules, rollout engines for block/full attention models, and auxiliary tools for evaluation and deployment (Wang et al., 8 Sep 2025). Key system features:

• Accelerated attention engines (e.g., JetEngine for block diffusion).
• Configurable RL/exploration setups (masking, curriculum, block-size adaptation).
• Docker-based build environments, checkpoint availability, and full pipeline evaluation scripts.

All major components—prompt templates, RL routines, ablation infrastructure, logging, and reproducibility settings—are documented for direct extension and comparative evaluation. Pretrained TraceRL-based models are uploaded under the Gen-Verse organization on HuggingFace.

---

The TraceRL framework constitutes a significant trajectory-aware RL methodology underpinning curriculum-based scaling and fine-tuning in diffusion LLMs, with proven empirical gains and a robust open-source foundation for continued research (Xia et al., 16 Jan 2026, Wang et al., 8 Sep 2025).