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

2026-06-02

AI Summary

Existing reinforcement learning approaches for diffusion language models struggle to balance computational efficiency with fine-grained training signals. This work proposes CAPR, an algorithm that, for the first time, extracts tree-like supervision signals from denoising trajectories by compressing them into path states, caching inexpensive sibling continuations, and introducing a block-level value head. This head redistributes the final reward according to tokens revealed within each block, enabling efficient fine-grained reinforcement learning without explicit search trees. CAPR establishes new state-of-the-art results for diffusion-based language model reinforcement learning on benchmarks including Sudoku, Countdown, GSM8K, and Math500, achieving performance on par with the strongest tree-search baselines on Sudoku using less than one-third of the per-step computation.

Abstract

Diffusion large language models (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.

Problem

• diffusion language models
• reinforcement learning
• denoising trace
• trajectory-aware learning
• reward redistribution

Innovation

• diffusion language models
• reinforcement learning
• denoising trace
• block-wise supervision
• trajectory-aware