TACG轨迹感知提交门控提升扩散语言模型解码效率

背景

扩散语言模型（Diffusion Language Models，DLLMs）通过逐步对掩码位置进行去噪来生成文本，在此过程中会暴露出一系列预测分布（即轨迹），而非单一的快照信念。现有的大多数解码器却忽略了这一轨迹，仅从当前快照提交 token，从而将置信度与提交就绪性混为一谈：在上下文不完整时，一个瞬时的 top-1 峰值可能被锁定，而那些在跨步骤中具有一致支持的候选 token 则被推迟。这种机制导致提交决策不稳定，既可能过早锁定错误候选，也可能延迟正确的 token。

为了解决这一问题，作者提出了 Trajectory-Aware Commit Gating（TACG），一种无需训练的门控解码器，它利用轨迹感知信号来决定当前提议是否准备好被提交。

核心内容

TACG 是一种无需额外训练的门控级别解码器，其核心思想是将 token 的身份锚定到基础后验分布（base posterior），仅利用轨迹感知信号来决定是否提交当前候选，从而分离置信度评估与提交决策。

TACG 由两个关键组件构成：

1. Temporal Implicit Logits Guidance (TILG)
   TILG 保持过去 logits 的指数移动平均作为自参考，并在自然参数空间（natural-parameter space）中将当前 logits 与该参考进行对比。这一对比能够捕捉 token 在时间上的稳定性：如果当前 logits 与历史平均显著偏离，意味着该候选尚不可靠，应推迟提交。

2. History Gate (HG)
   HG 强制要求在提交之前，当前提议必须保持短时间内的持久性（short-term proposal persistence）。它作为一个硬性门控，确保候选至少在连续几步中都是 top-1，从而过滤掉瞬时的噪声峰值。

此外，TACG 还包含一个带有上限的额外促进预算（capped extra-promotion budget），用于在满足稳定性约束后适度加速提交。这些组件共同构成一个稳定性约束的提交规则，不需要辅助网络或额外的前向传播。

作者在三个扩散语言模型上进行了评估：LLaDA、Dream 和 LLaDA2-Mini，使用的基准包括代码生成（HumanEval、MBPP）和数学推理（GSM8K、MATH500）。实验结果表明，TACG 通常能够提升或保持准确率，同时减少所需去噪步数，并提高每步前向传播生成的 token 数量（Tokens Per Forward, TPF）。代码已公开在 GitHub 上。

关键要点

• 问题本质：现有解码器仅基于当前快照的置信度提交 token，忽略了扩散过程中的轨迹信息，导致提交时机不当。
• 方法特点：TACG 是一种无需训练的解码器，不依赖任何辅助网络或额外前向传播，直接通过轻量级门控机制改善提交决策。
• 核心组件：
  • TILG：通过指数移动平均构建自参考，对比当前 logits 与历史趋势，在自然参数空间中进行对比。
  • HG：强制候选 token 在连续多步中保持 top-1 后才允许提交。
  • 上限促进预算：在稳定性约束下适度加速提交，避免过度延迟。
• 评估设置：在 LLaDA、Dream、LLaDA2-Mini 三个模型上，使用 HumanEval、MBPP、GSM8K、MATH500 四个基准。
• 主要结果：在保持或提升准确率的同时，减少去噪步数，提高每步生成的 token 数量（TPF）。
• 开源情况：代码公开，便于复现和进一步研究。

意义与影响

TACG 为扩散语言模型的解码提供了一种全新的视角：通过显式利用去噪轨迹中的时间一致性信息，而非仅仅依赖单步预测，这使得提交决策更加稳健。该方法无需额外训练或修改模型架构，具有很强的实用性和迁移性，可以直接应用于现有的 DLLM。

从效率角度看，TACG 能够减少去噪步数并提升每步输出 token 数，这意味着在相同计算预算下可以生成更长的文本，或加速推理过程。这对于扩散模型在实际部署中的竞争力至关重要——扩散模型虽然生成质量高，但推理速度常被诟病，TACG 提供了一条轻量级优化路径。

此外，该工作强调了“提交就绪性”与“置信度”的分离，这一思路可能启发后续研究在其他序列决策任务（如多步推理、强化学习中的动作选择）中引入类似的轨迹感知门控机制。公开的代码也为社区验证和扩展提供了基础。