mlsanityData to 14 Oct 2026

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

arXiv 2606.04396 ↗First posted 3 Jun 2026cs.CLcited by 0

On the map: Diffusion models → Discrete and text diffusion

Abstract (arXiv)

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.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Discrete and text diffusion, one of 14 directions in Diffusion models. Builds on 20 of its 35 references that are in the map. First posted on arXiv 3 Jun 2026.

Builds on

20 of 35 references.

Closest ancestors · 8

Same direction, newest first

  1. 2025Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization9 Oct
  2. 2025SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation7 Oct
  3. 2025Advancing Reasoning in Diffusion Language Models with Denoising Process Rewards2 Oct
  4. 2025Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step28 Sep
  5. 2025MDPO: Overcoming the Training-Inference Divide of Masked Diffusion Language Models18 Aug
  6. 2024Scaling up Masked Diffusion Models on Text24 Oct
  7. 2024Simplified and Generalized Masked Diffusion for Discrete Data6 Jun
  8. 2023Diffusion Language Models Can Perform Many Tasks with Scaling and Instruction-Finetuning23 Aug

Other references · 12

  1. 2017Proximal Policy Optimization Algorithmscited by 15,215
  2. 2021Training Verifiers to Solve Math Word Problemscited by 10,109
  3. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 9,379
  4. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,057
  5. 2023Let's Verify Step by Stepcited by 3,789
  6. 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 2,026
  7. 2025LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Modelscited by 285
  8. 2025Mercury: Ultra-Fast Language Models Based on Diffusioncited by 235
  9. 2024Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimizationcited by 78
  10. 2025Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspectivecited by 34
  11. 2025Entropy-Gated Branching for Efficient Test-Time Reasoningcited by 22
  12. 2025EntroPIC: Towards Stable Long-Term Training of LLMs via Entropy Stabilization with Proportional-Integral Controlcited by 13

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Reasoning in Diffusion Large Language Models is Concentrated in Dynamic Confusion Zones
  2. 2026Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models
  3. 2026Reinforcement Learning from Denoising Feedback
  4. 2025d-TreeRPO: Towards More Reliable Policy Optimization for Diffusion Language Models
  5. 2026SLIM-RL: Risk-Budgeted Random-Masking RL for Diffusion LLMs Without Trajectory Slicing
  6. 2025Learning Unmasking Policies for Diffusion Language Models
  7. 2026Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages
  8. 2026On Trajectory-Aware Training for Masked Diffusion Language Models