用去噪轨迹生成细粒度奖励,让扩散语言模型强化学习更高效。
Read the Trace, Steer the Path: Trajectory-Aware Reinforcement Learning for Diffusion Language Models

- 利用去噪过程中的动态信息生成局部奖励信号,替代昂贵的树状搜索。
- 在4个任务上达到新SOTA,256-512token预算下性能领先。
- 计算开销仅为平铺回溯的0.75倍,树搜索的60%,适合资源受限场景。
扩散大语言模型(dLLMs)通过迭代并行解掩码与修正多个位置生成文本,留下丰富的去噪轨迹,反映哪些词元变得确定、哪些仍不稳定、何时形成承诺。现有dLLM强化学习方法对此信号利用不足:平铺回溯成本低但仅分配单一结果奖励;树状回溯能提供更精细可验证信号,但计算开销大。本文提出CAPR(缓存-分摊路径精炼)算法,将去噪轨迹压缩为紧凑路径状态,利用缓存的轨迹状态生成低成本兄弟延续,并训练块级价值头实现局部块级监督。在块级解掩码调度下,CAPR记录路径状态与块进度特征,根据每块揭示的词元重新分配最终奖励,训练价值头将稀疏奖励转化为块级PPO权重。该方法在保持树状搜索大部分精度的同时,避免全树展开,回溯生成成本降至平铺回溯的0.75倍、树回溯的0.6倍(标准设置下)。在4×4数独、倒计时、GSM8K和Math500任务上,使用密集型与混合专家架构的LLaDA模型,CAPR在256-512词元预算下刷新dLLM强化学习新纪录。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。