💬 NLP

Just on Time: Token-Level Early Stopping for Diffusion Language Models

本文介绍了一种针对扩散语言模型的无需训练的、基于标记层级的早停方法，该方法通过轻量级预测信号动态识别并确定稳定的标记，在显著降低计算成本的同时，保持了多种任务中的生成质量。

原作者： Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

发布于 2026-08-04

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试解决一个巨大的拼图游戏，但你不是看盒子上印着的图案，而是面对一盒空白的灰色方块。为了完成它，你必须猜测每个方块应该是什颜色，然后观察整个画面，发现有些猜测错了，于是重新涂抹它们。你重复这个“猜测并修正”的循环数百次，直到画面最终看起来正确为止。这就是一种被称为**扩散语言模型（Diffusion Language Model）**的新型人工智能编写文本的方式。不同于以往那种像打字员一样一次写一个词的旧型 AI，这些模型从一张白纸开始，同时对整个句子进行精炼，就像画家在画布上慢慢添加细节一样。

问题在于，这种绘画过程极其缓慢。尽管 AI 通常很快就能搞定句子的简单部分（比如“The”或“cat”），但它仍然会反复涂抹这些部分，以确保万无一失。这就像一位厨师尝了尝汤，意识到味道已经完美了，却还是在端上桌之前又搅拌、品尝了一个小时。这浪费了大量的计算能力和时间。科学家们提出的问题是：我们如何告诉 AI，“嘿，这一部分你已经做对了，别再管它了，去处理那些难的部分吧”？

这就是一种名为 JoT（Just on Time，即时）的新方法。把 JoT 想象成站在 AI 画家身旁的一位聪明的监工。它不会同时命令整个团队停止绘画，而是观察画布上的每一个单词。一旦 AI 对某个特定单词非常有信心——比如它 99% 确定这个词是“cat”——JoT 就会大喊：“冻结这个词！不要再碰它了！”并将 AI 的注意力转移到下一个模糊、不确定的单词上。

开发 JoT 的研究人员发现，这种“词元级提前停止”（token-level early stopping）效果惊人。他们在两个强大的 AI 模型 Dream-7B 和 LLaDA-8B 上进行了测试，涵盖了四个挑战：解决数学题、回答常识题、补全句子以及编写代码。结果令人瞩目。在一项名为 GSM8K 的数学推理测试中，JoT 让 AI 的速度提升了 5.5 倍，而得分仅略微下降（约 2.3 分）。在编程挑战 HumanEval 中，提速效果更为显著，达到了近 20 倍（19.6倍），且 AI 依然能答对几乎所有的题目。

该论文反对那种让所有单词同时停止的“一刀切”方法。相反，JoT 使用了一个巧妙的技巧：它会降低那些紧邻已完成单词的单词的置信度门槛。如果 AI 已经搞定了句子的开头，它就可以更早地信任下一个词。这使得 AI 能够将精力仅集中在真正令人困惑的部分，跳过对简单部分的冗余工作。

虽然这种方法速度极快，但作者也谨慎地指出，它并非魔法。他们测量到 AI 仍会犯一些错误，主要是在遇到棘手的数学步骤并过早锁定错误数字时。然而，这些错误很少见，且这种权衡明显有利于速度。论文表明，通过让每个单词在各自的完美时刻“退出”过程，我们可以让这些强大但缓慢的 AI 模型在日常使用中变得更加实用，而无需重新训练或改变其核心设计。这是一个简单的、无需训练的修复方案，它让 AI 工作得更聪明，而不是更辛苦。

技术摘要：扩散语言模型的即时停止 (Just on Time, JoT)

问题陈述

扩散语言模型 (DLMs) 通过将一个全掩码序列迭代细化为清晰的输出进行文本生成。虽然这种方法能够实现并行 Token 预测和双向上下文整合，但它存在解码效率低下的问题。标准的 DLMs 需要对整个序列进行固定数量的细化步骤 ( $T$ )，即使许多 Token 在最终步骤之前就已经达到了稳定且高置信度的预测。这导致了不必要的计算开销，因为模型仍在继续细化那些已经收敛的 Token。现有的早停方法通常采用全局标准或依赖于特定任务的适配，无法动态地解决单次生成中不同位置之间存在的收敛速率差异。

方法论：JoT (Just on Time)

作者提出了 JoT，一种无需训练的、针对每个 Token 的早停机制，旨在识别序列中每个位置独立的收敛情况。JoT 允许不同的 Token 在不同的步骤中“退出”迭代细化过程，从而将计算集中在需要的地方。

核心组件

置信度指标：
JoT 通过前两个预测概率 ( $p_1$ 和 $p_2$ ) 的比例来量化每个掩码位置 $i$ 的预测确定性：
$r_i = \frac{p_1}{p_2 + \epsilon}$
该指标源自未缩放的 Logits，因此对采样温度具有不变性。高比例表示对顶级候选者的决定性偏好，而接近 1 的比例则表示不确定性。
空间调制 (Spatial Modulation)：
考虑到位于已解析（未掩码）上下文附近的 Token 能从更丰富的局部信息和更低的有效噪声中受益，JoT 引入了一种空间自适应阈值。
- 几何核根据位置 $i$ 与窗口 $D$ 内未掩码 Token 的接近程度计算权重 $w_i$ 。
- 该权重被归一化为一个软化因子 $\phi_i \in [0, 1]$ 。
- 位置 $i$ 的接受阈值 $\tau_i$ 在最大阈值 $\tau_{max}$ 和最小阈值 $\tau_{min}$ 之间进行插值：
  $\tau_i = \tau_{max} - (\tau_{max} - \tau_{min}) \cdot \phi_i$
- 靠近已解析上下文的 Token 会获得较低的阈值（更容易完成），而处于掩码区域中心的 Token 则保留更严格的阈值。
早退决策 (Early-Exit Decision)：
在每一步中，如果掩码 Token 的置信度比例 $r_i$ 超过其自适应阈值 $\tau_i$ ，则该 Token 被视为完成（通过 argmax 或采样选择）并从掩码位置集合中移除。该过程持续进行，直到所有 Token 都被完成或达到最大步数。

理论分析

论文通过形式化分析（命题 1）探讨了“领先对反转”（leading-pair reversals）的概率——即在更多上下文被揭示时，一个已提交的 Token 随后被第二候选者超越的情况。在模型预测符合真实条件分布（校准）的假设下，此类反转的概率受限于 $2/(1+\tau)$ 。作者指出，高阈值下的残余误差可能源于校准失败，而非停止机制本身。

主要贡献

JoT 框架： 一种无需训练的、基于每个 Token 的早停方法，通过根据与已解析上下文的空间邻近性来调整接受阈值。
全面的评估： 在两个最先进的模型（Dream-7B 和 LLaDA-8B）以及四个多样化的基准测试（GSM8K（数学）、MMLU（知识）、HellaSwag（常识）和 HumanEval（代码））上进行了实验。
消融研究： 详细分析了隔离阈值选择和空间调制的效应，为超参数配置提供了实践指导。

实验结果

作者将 JoT 与全量解码以及其他早停基准方法（Prophet、KLASS、Fast-dLLM）进行了对比评估。

速度-质量权衡：
- Dream-7B： JoT 在 GSM8K 上实现了高达 5.54× 的加速，在 HumanEval 上实现了 19.60× 的加速。准确率下降幅度较小（例如，GSM8K 下降 2.3 点，HumanEval 下降 0.6 点），在大多数基准测试中，性能保持在全量解码的约 3 个百分点以内。
- LLaDA-8B： JoT 在所有基准测试中均将准确率维持在全量解码的 3.1 个百分点以内，同时提供了具有竞争力的加速效果（例如，GSM8K 上为 3.75×）。
对比： JoT 通常优于 Prophet（后者遭受了更大的准确率下降），并且在多个任务上提供了比 KLASS 和 Fast-dLLM 更好的速度-准确率帕累托前沿（Pareto frontiers）。
可组合性： JoT 成功地与 Fast-dLLM 的 DualCache（减少每步延迟）相结合，证明了步数减少与单步优化是正交的，并且可以结合使用。
失效模式： 对 500 个 GSM8K 样本的分析显示，错误主要发生在模型在关键推理分支处过早提交 Token 时（例如，选择了 "× 2" 而不是 "/ 2"）。这些错误并非新型失效模式，而是以与全量解码不同的速率发生。

重要性与主张

论文声称，JoT 为扩散语言模型的解码低效问题提供了一个引人注目的解决方案，且无需重新训练或特定任务的微调。通过在 Token 表现出足够置信度时动态地将其完成，JoT 大幅减少了扩散步骤的总数。

作者将 JoT 定位为一种在 Dream-7B 上实现了最佳速度-质量权衡，并在 LLaDA-8B 上保持竞争力的方法。他们强调，该方法的有效之处在于它允许模型跳过对“显而易见”的上下文 Token 进行冗余细化，同时保留处理复杂推理阶段所需的计算量。论文结论认为，尽管最优阈值可能因任务而异，但该方法在超参数方面的鲁棒性使其在多样化的应用场景中具有实用价值。