Read the Trace, Steer the Path: Trajectory-Aware Reinforcement Learning for Diffusion Language Models

Abstract: Diffusion LLMs (dLLMs) generate responses by iteratively unmasking and revising many positions in parallel. This process leaves a rich denoising trace depicting which tokens become confident, which remain unstable, and when commitments form. Existing dLLM reinforcement learning methods use this signal only weakly. Flat rollouts are cheap, but assign a single outcome reward to the whole trajectory. Tree rollouts provide finer, verifiable training signals by branching partial trajectories and propagating leaf rewards upward, but are compute intensive. We ask whether the denoising trace itself can provide tree-like supervision without tree-level compute. We introduce CAPR (Cached-Amortized Path Refinement), a dLLM-RL algorithm that summarizes the denoising trace into a compact path state, uses cached trajectory states to generate cheap sibling continuations, and trains a block-level value head for local block-wise supervision. Under a block-wise unmasking schedule, CAPR records path-state and block-progress features, then redistributes the final outcome reward across blocks according to the tokens revealed in each block. This trains the value head to convert one sparse reward into block-level PPO weights. CAPR therefore recovers much of the granularity of tree search while avoiding full tree expansion, reducing rollout-generation cost to roughly 0.75x that of flat rollouts and 0.6x that of tree rollouts (under standard settings). Across 4x4 Sudoku, Countdown, GSM8K, and Math500, on dense and mixture-of-experts LLaDA backbones, CAPR sets a new state of the art for RL-tuned dLLMs at 256- and 512-token budgets. On Sudoku, it matches the strongest tree-structured baseline at less than one third of the per-step compute.

• The paper introduces CAPR, a trajectory-aware RL framework that refines denoising traces for localized reward assignment during language model fine-tuning.
• It leverages cached path states, selective branching, and a lightweight block critic to deliver state-of-the-art accuracy with reduced computational cost.
• Experimental results show CAPR outperforms both flat and tree-based baselines on reasoning benchmarks like Sudoku and GSM8K.

Trajectory-Aware RL for Diffusion LLMs via Cached-Amortized Path Refinement

Background and Motivation

Diffusion LLMs (dLLMs) use iterative denoising to generate completions, producing a trajectory of token predictions at each masked position across steps. This process captures rich trace information—including token-level confidence, entropy, and stability—which is unexploited by standard RL-based fine-tuning paradigms. Existing RL approaches for dLLMs exist at two extremes: "flat" rollouts that broadcast a sparse terminal reward across the whole trajectory (cheap but uninformative), and tree-based rollouts that compare branched continuations for segment-level attribution (granular but compute-intensive).

Figure 1: The denoising trace encodes evolving confidence and stability at each masked position, enabling token-wise discrimination between crucial and irrelevant decisions.

The CAPR Algorithm

The paper introduces CAPR (Cached-Amortized Path Refinement), a trajectory-aware RL framework that leverages the denoising trace for block-level reward assignment with significantly reduced compute, closing the gap between flat and tree-based techniques.

Mechanisms

1. Cache/Steer: At each denoising step, CAPR records per-position token confidence, entropy, and stability, encoding these into a dynamic path state vector. This path state is used both to steer the next denoising action (biasing toward stable positions, damping oscillating ones) and to cache branchable prefixes of the trajectory.
2. Branch/Prune: Rather than full tree expansion, CAPR branches sibling continuations at a fixed mid-trajectory step from the shared cached prefix using Gumbel perturbation. Siblings are pruned by local path-state quality, preserving only those with the strongest commitments, thus maintaining diversity without full combinatorial explosion.
3. Block Critic: CAPR maintains a local value head at block boundaries, learning to convert eventual terminal reward into per-block PPO weights. The block critic reads compact block-state embeddings and uses TD/GAE computed from a slowly-updated Polyak-averaged critic, thus producing token-attributed advantages that localize sparse rewards.

   Figure 2: Overview of CAPR: (a) Trace state is cached and used both to steer subsequent denoising, (b) Siblings are branched/pruned from mid-pathances and (c) a block critic generates per-block advantages for PPO from final rewards.

Self-distillation further repurposes solved rollouts with zero group variance into supervised targets for additional signal during RL fine-tuning.

Training and Inference Procedure

CAPR’s training objective combines block-weighted PPO loss, a KL reference constraint, a one-sided chi-square trust-region term to suppress upper-tail logit spikes, and a self-distillation loss. All trajectory-aware modules (branch/prune, block critic, path state steering) are training-only; at inference only the base denoising loop and learned LoRA dropout policy are active. This ensures computational efficiency at deployment.

Experimental Results

Benchmarks, Baselines, and Models

Evaluations were performed on four reasoning benchmarks—$4{\times}4$ Sudoku, Countdown, GSM8K, and Math500—using two state-of-the-art dLLM backbones: LLaDA-8B-Instruct (dense) and LLaDA-MoE-7B-A1B-Instruct (mixture-of-experts). Baselines encompass flat (Diffu-GRPO, VRPO, wd1), token-wise (SAPO, TraceRL), and tree-based (d-TreeRPO) RL algorithms.

Across all tasks and at both 256- and 512-token generation budgets, CAPR establishes new SOTA for RL-tuned dLLMs in accuracy while requiring less than one-third the per-step compute of the best tree-based method.

• On Sudoku, CAPR matches or marginally exceeds d-TreeRPO’s final accuracy (e.g., $93.0\%$ vs. $92.9\%$ at 256 tokens for LLaDA-8B) at only $165$ seconds per update versus $596$ seconds for d-TreeRPO.
• Gains over flat baselines are especially pronounced on tasks with highly sparse rewards (Sudoku, Countdown).
• On mathematical QA (GSM8K, Math500), where the base models already exhibit competence, CAPR continues to deliver consistent but smaller increases, avoiding the instability and slow learning of tree-based methods.

  Figure 3: Training reward curves for CAPR across eight runs, showing rapid and smooth reward ascent, particularly on tasks with sparse credit assignment.

  

  Figure 4: Ablation showing progressive gains via Branch/Prune, Block Critic, and self-distillation, with Block Critic driving the largest independent improvement.

  

  Figure 5: Peak Sudoku reward by ablation configuration confirms the complementary contribution of all components.

Robustness diagnostics (KL divergence, chi-square regularization, value-head MSE, gradients) confirm that the increased credit assignment granularity does not induce policy instability, with the CAPR policy consistently remaining within a tightly bounded trust region during learning.

Figure 6: Diagnostics of reference KL, chi-square regularization, and value-head MSE indicate stable training dynamics.

Figure 7: CAPR enables more efficient and higher-reward training compared to open-source and tree-based RL baselines across all four reasoning tasks.

Implications and Theoretical Perspective

The CAPR approach demonstrates that effective trajectory-aware reward assignment for dLLMs does not require combinatorial tree search. By caching and amortizing mid-trajectory prefixes, leveraging intra-path state statistics, and employing a lightweight block-critic, CAPR achieves localized reward propagation at lower computational cost. The chi-square trust region regularizer specifically suppresses rare token-level over-optimizations, a documented problem in sparse-reward RL for sequence models. Theoretical interpretations link CAPR to performance-difference bounds for trust-region RL and to amortized advantage estimation.

Pragmatically, trajectory-aware RL via CAPR enables scaling RL-fine-tuning of dLLMs to longer generation horizons and larger decoder budgets, filling an essential gap as masked diffusion LMs become ubiquitous (2606.04396).

Limitations and Outlook

• CAPR’s evaluation and design assume verifiable reward signals (e.g., programmatic verification, QA), and may require adaptation for subjective reward tasks (alignment, safety, instruction-following) with noisier signals.
• Generalization to other masked diffusion architectures will depend on the compatibility of the recorded trace statistics and may require hyperparameter retuning.
• The block-critic’s local value is learned online; for tasks demanding extremely long credit assignment chains, more global critics may be required.

Future directions include cross-task signal sharing, reward-model-based signal localization, and hybridization with multimodal or instruction-tuned objectives.

Conclusion

CAPR provides a trajectory-aware framework for reinforcement learning in masked diffusion LLMs that yields state-of-the-art accuracy and wall-clock efficiency by directly leveraging denoising traces. Through block-level path-state caching, selective branching/pruning, a fine-grained critic, and robust regularization, it resolves sparse-reward credit assignment without prohibitive compute. This architecture significantly improves task performance on constrained and symbolic reasoning while making RL-tuning of dLLMs more tractable at scale.