一致性扩散语言模型：推理速度提升14倍，质量不减

摘要

我们提出了一致性扩散语言模型（CDLM），通过将基于一致性的多令牌最终化与分块KV缓存相结合，加速扩散语言模型推理，在数学和编程任务上实现了高达14.5倍的延迟加速。

扩散语言模型（DLM）正成为自回归语言模型（AR LM）的一种有前景的替代方案。DLM不是一次生成一个令牌，而是在多个采样步骤中迭代地细化部分掩码序列，逐步将完全掩码的序列转换为干净的文本。这种细化过程创造了一个引人注目的机会：它支持并行生成，允许模型在每次迭代中最终确定多个令牌，从而可能实现比自回归解码更高的吞吐量。同时，它可以利用双向上下文来解锁新的能力，如文本填充和细化。

然而，在实践中，标准DLM存在两个主要效率低下的问题。[1]

完全双向注意力下的KV缓存不兼容。标准DLM通常使用双向（非因果）注意力，这需要在每个去噪步骤中重新计算整个上下文的注意力，使得推理成本高昂，并阻止了标准KV缓存的使用。

保持质量所需的高细化步数。高质量生成通常需要许多去噪/细化步骤，通常与生成长度相当。简单地减少步数往往会急剧降低质量。

CDLM通过一种训练后方案针对这两个瓶颈，使少步推理变得可靠，同时实现精确的分块KV缓存。

预备知识：扩散语言模型中的推理

DLM生成是在N个离散采样步骤上的迭代细化。它将时间t=1处的完全掩码序列转换为t=0处的干净序列。在每个步骤中，模型根据当前噪声序列xt和提示c预测干净序列分布x0：

$p_{\theta}(\mathbf{X}_0 \mid \mathbf{X}_t, c)$

一种常见的确定性实例化是低置信度重新掩码：模型贪婪地取消掩码令牌（通常在块内），最终确定最高置信度的掩码位置，同时保持其他位置掩码。这导致解码轨迹：

$\mathcal{T}_{\mathbf{x}} = \left(\mathbf{x}_{t_0}, \mathbf{x}_{t_1}, \ldots, \mathbf{x}_{t_N}\right), \quad t_k = 1 - \frac{k}{N}$

它记录了部分细化序列逐步演变的过程。该轨迹成为CDLM训练的核心对象。

CDLM训练

1) 轨迹收集

我们通过使用DLM在特定领域提示上运行推理来离线收集轨迹。对于每个提示x，我们记录令牌级解码轨迹T_x，一个紧凑的隐藏状态缓冲区H_x，其中包含令牌最终化时刻的最后一层隐藏状态，以及真实文本ŷ。具体来说，我们采用分块解码，生成长度L_g = 256，块大小B = 32，总步数N = L_g（即，在当前块内每步恰好最终确定一个令牌）。这种保守的设置产生了更高质量的轨迹用于蒸馏。

2) 块因果学生模型和注意力掩码

在轨迹提取过程中，我们使用全双向注意力掩码。相比之下，在训练CDLM时，我们采用块级因果掩码，该掩码关注提示、先前完成的块以及当前解码块。这种设计使模型能够从全双向切换到块扩散模型（如[2]），从而实现对已完成块的精确块级KV缓存。

3) 训练目标

CDLM联合最小化三个目标：

(i) 蒸馏损失（新解掩码位置）

对于在中间状态y及其块完成y*之间新解掩码的位置，我们将学生的预测分布与教师从存储的隐藏状态获得的重构分布进行匹配。

直觉：该目标作为主要锚点，教导学生在块级因果约束下完成块内的多个标记。

(ii) 一致性损失（仍掩码位置）

我们通过将学生在状态y的预测与自身在更知情状态y*下对仍掩码位置的预测对齐（使用停止梯度目标），来强制块内时间一致性。

直觉：该目标鼓励沿解码轨迹的稳定多步转换。

(iii) 辅助DLM掩码去噪损失

我们包含一个标准掩码去噪目标，应用于随机掩码的真实文本。

直觉：该目标保留了模型的一般掩码标记预测能力，并有助于保持推理行为，特别是在数学任务上。

4) 推理

在推理时，CDLM以块级自回归方式解码，重用提示和所有先前完成块的KV缓存。在每个块内，我们应用置信度阈值并行完成。[3] 一旦当前块中出现文本结束标记，我们还采用早停。

我们有意避免引入额外超参数的启发式方法（例如，具有任务相关设置的块间并行），而是专注于基于精确KV缓存和可靠步骤减少的稳健默认解码流程。

主要结果：CDLM–Dream

我们看到：

• CDLM–Dream在基准测试中实现了最大的步骤减少，将细化步骤减少了约4.1倍至7.7倍，同时在大多数任务上精度变化很小。
• 这些步骤减少转化为显著的延迟改进：在GSM8K-CoT上高达11.2倍，在MBPP-Instruct上高达14.5倍。
• CDLM通常达到最高的每秒标记数吞吐量，但有一个值得注意的细微差别：由于CDLM严格块级因果，可能产生较短的输出，同时保持pass@1质量。

有效的步骤减少：为什么训练很重要

简单截断步骤数会导致明显的精度下降，而CDLM在相似的步骤预算下保持质量（并且由于缓存，延迟大约减半）。这突出了一个核心点：稳定的多标记细化并非免费；它需要强制轨迹一致行为的训练。

系统级分析：为何块级 DLM 处于最佳平衡点

为了理解硬件利用率，我们分析了随着批次大小增加时的算术强度（AI，每字节移动的 FLOPs），比较了：AR 解码、普通（全注意力）DLM、块级 DLM（CDLM，B∈{4,16,32}）。

关键解读：

• AR 解码在小批次大小时强烈受限于内存（bs=1 时 AI 接近 1），随着批次增大，由于权重加载的摊销，AI 会提升。
• 普通 DLM 即使在 bs=1 时也受限于计算，因为全双向注意力每一步都处理整个序列，导致饱和。
• 块级 DLM（CDLM）处于中间状态：由于块内并行性（在相似内存流量下处理 B 个 token），AI 高于 AR，但低于普通 DLM，通常是小批次设置下的平衡工作点。

总体而言，该分析解释了为何类似 CDLM 的块级扩散可以在小批次大小时实现高效：它利用并行性摊销内存访问，同时仍处于受益于实际扩展的状态。

讨论与结论

表现力与效率

DLM 中的全双向注意力需要在每个去噪步骤重新计算 O(L^2) 注意力，使得推理高度计算密集。CDLM 实现了精确的 KV 缓存，同时保留每个块内的双向上下文，保持了局部精化能力（例如，在当前块内进行填充）。

随更强 DLM 骨干扩展

CDLM 是一种训练后方案，可应用于任何块扩散模型，其优势应随着更强 DLM 的出现而增长。一个有前景的方向是从更大、更强的 DLM 教师收集轨迹，并使用 CDLM 训练中等规模的学生模型。

结论

我们提出了 CDLM，一种基于训练的加速方案，将一致性建模引入 DLM。通过强制块内时间一致性并微调块级因果学生模型，CDLM 减少了精化步骤并实现了精确的 KV 缓存。在数学和编程任务上，CDLM 实现了更快的推理、更少的步骤、更低的延迟和更高的吞吐量，同时保持了有竞争力的准确性。

[1] Beyond Next-Token Prediction: A Performance Characterization of Diffusion versus Autoregressive Language Models

[2] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

[3] Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding