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

• URL: http://arxiv.org/abs/2606.04396v1
• Date: Wed, 03 Jun 2026 03:22:54 GMT
• Title: Read the Trace, Steer the Path: Trajectory-Aware Reinforcement Learning for Diffusion Language Models
• Abstract summary: We introduce CAPR (Cached-Amortized Path Refinement), a dLLM-RL algorithm that summarizes the denoising trace into a compact path state.<n>It reduces rollout-generation cost to roughly 0.75x that of flat rollouts and 0.6x that of tree rollouts.<n>Across 4x4 Sudoku, Countdown, GSM8K, and Math500, it matches the strongest tree-structured baseline at less than one third of the per-step compute.
• Score: 26.87342465042021
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Fork Where the Model Changes Its Mind: Belief-Shift Branching for Tree-Structured Reinforcement Learning [62.24941741019998]
  We formalize fork placement as locating the emphss of the chain's value curve, where the expected outcome turns.<n>Three instantiations, none needing step-level supervision, span access levels: a black-box probe, a logit-lens depth profile, and a learned activation direction.
  arXiv  Detail & Related papers  (2026-09-10T04:06:45Z)
• HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees [10.936836651826722]
  We present HARTS (Hybrid-Attention RL over Tree Structures)<n>HarTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model.
  arXiv  Detail & Related papers  (2026-08-28T10:18:54Z)
• Trajectory-Level Speculative Decoding for Diffusion Language Models [16.97093186372498]
  Diffusion-based language models (dLLMs) enable parallel token generation through iterative denoising, but existing decoding strategies collapse to single-token generation under low confidence.<n>We develop a trajectory-level speculative framework that constructs draft denoising trajectories via confidence-stratified tree exploration.<n>Our framework reduces denoising by 30-40% and increases tokens-per-step from 2.6 to 4.3, achieving 7-14x speedup over vanilla dLLMs and 1.3x over Fast-dLLM with less than 1% accuracy change across reasoning and code benchmarks.
  arXiv  Detail & Related papers  (2026-08-27T09:42:20Z)
• Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing [65.59464371673623]
  Local Branch Routing (LBR) is a token-level test-time scaling framework.<n>It expands a small local lookahead tree, forwards all sampled branches through the language model, and uses a lightweight router to select the depth-1 subtree to commit.<n>LBR allows each token decision to use evidence beyond the root next-token distribution while avoiding full solution-level search.
  arXiv  Detail & Related papers  (2026-06-24T03:42:44Z)
• Breaking chains with trees: Deep learning with $\mathcal{O}(\log N)$ parallel time complexity [0.749500254646884]
  We propose Hierarchical Block-Local Learning (HBLL), a framework that decomposes deep neural networks into hierarchically linked blocks trained using local learning objectives.<n>HBLL is the first algorithm that is able to train deep neural networks in $mathcalO(log N)$ parallel time complexity, where $N$ is the number of network layers.<n>We show that HBLL implicitly defines a family of neuralworks corresponding to different hierarchical paths, enabling flexible inference with different effective numbers of layers.
  arXiv  Detail & Related papers  (2026-06-19T14:44:07Z)
• TreePS-RAG: Tree-based Process Supervision for Reinforcement Learning in Agentic RAG [71.06073770344732]
  Agentic retrieval-augmented generation (RAG) formulates question answering as a multi-step interaction between reasoning and information retrieval.<n>We present TreePS-RAG, an online, tree-based RL framework for agentic RAG that enables step-wise credit assignment while retaining outcome-only rewards.
  arXiv  Detail & Related papers  (2026-01-11T14:07:30Z)
• Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers [36.26426380985327]
  Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost limits scaling to long token sequences.<n>Recent Top-K sparse attention approaches reduce the computation of DiTs by compressing tokens into block-wise representation.<n>We introduce Log-linear Sparse Attention (LLSA), a trainable sparse attention mechanism for extremely long token sequences.
  arXiv  Detail & Related papers  (2025-12-18T14:53:12Z)
• Lookahead Tree-Based Rollouts for Enhanced Trajectory-Level Exploration in Reinforcement Learning with Verifiable Rewards [48.321707628011005]
  Lookahead Tree-Based Rollouts (LATR) is a novel rollout strategy designed to explicitly promote trajectory-level diversity.<n>LATR accelerates policy learning by 131% on average and improves final pass@1 performance by 4.2%.
  arXiv  Detail & Related papers  (2025-10-28T11:12:02Z)
• TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling [65.46347858249295]
  TreePO is a self-guided rollout algorithm that views sequence generation as a tree-structured searching process.<n>TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity.
  arXiv  Detail & Related papers  (2025-08-24T16:52:37Z)
• R-Stitch: Dynamic Trajectory Stitching for Efficient Reasoning [80.104336426172]
  Chain-of-thought (CoT) enhances problem-solving ability of large language models.<n>CoT incurs substantial inference cost due to long autoregressive trajectories.<n>We introduce R-Stitch, a training-free hybrid decoding framework.
  arXiv  Detail & Related papers  (2025-07-23T08:14:36Z)
• Recursive Speculative Decoding: Accelerating LLM Inference via Sampling Without Replacement [11.91629418177851]
  Speculative decoding is an inference-accel method for large language models. Recent works have advanced this method by establishing a draft-token tree. We present Recursive Speculative Decoding (RSD), a novel tree-based method that samples draft tokens without replacement.
  arXiv  Detail & Related papers  (2024-02-21T22:57:49Z)
• Provably Efficient Offline Reinforcement Learning with Trajectory-Wise Reward [66.81579829897392]
  We propose a novel offline reinforcement learning algorithm called Pessimistic vAlue iteRaTion with rEward Decomposition (PARTED) PARTED decomposes the trajectory return into per-step proxy rewards via least-squares-based reward redistribution, and then performs pessimistic value based on the learned proxy reward. To the best of our knowledge, PARTED is the first offline RL algorithm that is provably efficient in general MDP with trajectory-wise reward.
  arXiv  Detail & Related papers  (2022-06-13T19:11:22Z)
• Reinforcement Learning for Branch-and-Bound Optimisation using Retrospective Trajectories [72.15369769265398]
  Machine learning has emerged as a promising paradigm for branching. We propose retro branching; a simple yet effective approach to RL for branching. We outperform the current state-of-the-art RL branching algorithm by 3-5x and come within 20% of the best IL method's performance on MILPs with 500 constraints and 1000 variables.
  arXiv  Detail & Related papers  (2022-05-28T06:08:07Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.