让扩散语言模型学会何时固定生成的词元,提升生成质量与稳定性。
The Path Matters: Learning a Token-Commitment Policy for Diffusion Language Models

- 设计可学习的词元提交策略,替代人工规则决策
- 在问答、数学推理等任务上显著改善质量-速度权衡
- 无需重训即可跨设置通用,适合追求稳定生成的场景
扩散大语言模型通过并行优化多个词元位置实现快速生成,但这种并行性带来了隐藏的控制问题:每一步应将哪些候选词元提交到已解码序列中?我们称此为词元提交。现有固定生成器解码器多依赖手工设计的置信度规则或块级接受过滤器。本文提出将词元提交建模为可复用的轨迹状态策略,并引入TraceLock——一个轻量级插件控制器。由于理想提交时机不可得,TraceLock通过未来稳定性自监督:在解码步骤t,若某位置i的候选词元在完整解码后仍保持一致,则标记为稳定。该控制器对变长轨迹状态打分,决定哪些活跃候选应被提交。训练一次后,可直接部署于不同局部窗口宽度、生成长度和步数预算,无需重训或调参。在问答、数学推理和代码生成任务上的实验表明,TraceLock优于启发式与学习基线,在跨设置部署下表现尤为稳定。诊断分析显示其决策不能简化为标量置信度,表明冻结扩散语言模型存在超越置信度的可学习提交轨迹空间。代码开源:https://github.com/BobSun98/TraceLock。
原文摘要 · Abstract (English)
Diffusion large language models promise faster generation by refining many token positions in parallel, but this parallelism introduces a hidden control problem: which proposed tokens should be transferred into the partially decoded sequence at each step? We refer to this decision as token commitment. Existing frozen-generator decoders largely rely on hand-designed confidence rules or block-specific acceptance filters. We argue that token commitment can instead be learned as a reusable trace-state policy. We introduce TraceLock, a lightweight plug-in controller that instantiates this policy for a frozen diffusion language model. Since oracle commitment times are unavailable, TraceLock derives self-supervision from future stability: at decoding step t, a proposed token for position i is labeled stable if it matches the final token at position i after the full decoding trace completes. The controller scores variable-length trace states and decides which active token proposals should be committed to the partially decoded sequence. Once trained for a given frozen backbone, the controller can be deployed across local-window widths, generation lengths, and step budgets without retraining or per-setting calibration. Experiments on question answering, mathematical reasoning, and code generation show that TraceLock improves the quality-step tradeoff over heuristic and learned baselines, with particularly stable behavior under cross-setting deployment. Diagnostic analyses show that its decisions are not reducible to scalar confidence, suggesting that frozen diffusion language models expose a learnable space of commitment trajectories beyond confidence-based decoding. Code is available at https://github.com/BobSun98/TraceLock.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。