[Paper] 恰逢其时：Token-Level Early Stopping for Diffusion Language Models

Source: arXiv - 2602.11133v1

概述

扩散语言模型（DLM）通过反复“去噪”噪声令牌序列，直至生成连贯的输出。虽然功能强大，但这种迭代过程常常低效：许多令牌在仅几步后就已收敛到最终形态，却仍被模型一直更新到最后一步扩散。论文 “Just on Time: Token‑Level Early Stopping for Diffusion Language Models” 提出了一种 无需训练、基于令牌的提前停止 机制，能够在每个令牌收敛时即时检测并冻结它，从而在不牺牲质量的前提下降低总的扩散步数。

关键贡献

• Token‑level convergence detection：引入轻量级、仅用于推理的信号，以逐位置决定何时令牌“足够稳定”而停止更新。
• Training‑free approach：该方法可直接在任何预训练的扩散语言模型上使用；无需额外微调或辅助损失。
• Adaptive per‑token freezing：使每个令牌可以在不同的扩散步数停止，从而形成动态调度，而非固定的全局步数。
• State‑of‑the‑art efficiency：在多个基准（数学推理、开放域问答、科学理解）上，该技术将平均扩散步数降低 30‑55 %，同时保持 BLEU/ROUGE/Exact‑Match 分数在完整步数基线的 0.2 % 以内。
• Broad applicability：已在开源（如 DiffuSeq、Diffusion‑GPT）和商业扩散语言模型上验证，表明该方法与模型无关。

方法论

1. 信号提取 – 在每个扩散步骤中，模型已经为每个 token 生成了词汇表上的概率分布。作者为每个位置计算两个低成本统计量：

   • 预测熵（模型对该 token 的不确定程度）。
   • 局部一致性得分（当前预测与周围上下文的一致性，通过浅层注意力掩码衡量）。

2. 收敛准则 – 当一个 token 的熵同时低于预设阈值 且 其一致性得分高于第二个阈值时，该 token 被标记为“就绪”。这些阈值只需设定一次（例如通过小规模验证扫描），随后在所有下游任务中保持不变。

3. 动态冻结 – 一旦 token 满足准则，其嵌入即被冻结：后续的扩散步骤会跳过该位置的去噪计算，从而有效降低每一步的工作负载。其余“未稳定”的 token 继续被细化。

4. 实现细节 – 早停逻辑作为模型前向传播的薄层包装加入，几乎不增加开销（< 2 % 的总推理时间）。无需修改扩散调度、损失函数或模型架构。

结果与发现

关键要点

• 效率：平均推理延迟随扩散步数的减少成比例下降（在 V100 GPU 上约提升 40%）。
• 质量保持：在所有任务中，标准指标的下降在统计上并不显著，证明早停不会截断有用的细化。
• 鲁棒性：相同的阈值在各领域均表现良好，表明这些信号具有广泛的可靠性。

实际影响

• 更快的 LLM‑as‑a‑service：提供商可以以更低的每请求 GPU 时间提供基于扩散的模型，从而实现成本节约和更高的吞吐量。
• 边缘部署：减少的步数使得在资源受限的硬件上运行扩散 LM 成为可能（例如移动 GPU、边缘 TPU），而完整步推理在这些设备上几乎不可行。
• 混合流水线：开发者可以将提前停止与其他加速技巧（例如 classifier‑free guidance scaling、量化）结合，以实现更大的延迟降低。
• 动态质量‑vs‑速度控制：通过在推理时调整熵/一致性阈值，用户可以在需求下以略微降低质量换取更快的响应。
• 工具支持：作者发布了一个轻量级 Python 库（diffuse‑early‑stop），可接入现有的扩散‑LM API（Hugging Face、Diffusers），降低集成门槛。

限制与未来工作

• 阈值敏感性：虽然作者报告了良好的默认值，但在极端领域（例如诗歌生成）可能需要特定领域的调优。
• 非单调收敛：在罕见情况下，被视为“稳定”的标记可能因长程依赖而随后翻转，可能损害连贯性；当前方法不会重新激活已冻结的标记。
• 对超大词表的可扩展性：熵计算随词表大小线性增长；对于词表超过 100k 的模型，开销可能变得显著，暗示需要近似熵估计器。
• 未来方向：论文暗示通过小型元网络学习自适应阈值，探索多模态扩散（文本+图像）提前停止，并与基于强化学习的解码策略集成。

底线：基于标记级的提前停止提供了一种务实、即插即用的方式，使扩散语言模型 显著更快，且不牺牲其在复杂推理任务中吸引人的质量。对于构建 AI 驱动产品的开发者而言，这是一种低成本的优化，可在现有硬件上解锁新用例。

作者

• Zahar Kohut
• Severyn Shykula
• Dmytro Khamula
• Mykola Vysotskyi
• Taras Rumezhak
• Volodymyr Karpiv

论文信息

• arXiv ID: 2602.11133v1
• 分类: cs.LG, cs.CL
• 发布时间: 2026年2月11日
• PDF: 下载 PDF