一致性扩散语言模型：推理速度提升高达14倍且不牺牲质量

摘要

1) 一句话总结 一致性扩散语言模型（CDLM）通过结合基于一致性的多Token最终化与分块KV缓存技术，解决了传统扩散模型的效率瓶颈，在数学和编程任务上实现了高达14.5倍的推理加速且不牺牲生成质量。

2) 核心要点

解决传统瓶颈：标准扩散语言模型（DLM）因全双向注意力无法使用标准KV缓存，且需大量细化步数来维持质量。CDLM通过训练后（post-training）方案同时解决了这两个问题。
分块因果掩码：CDLM在训练中采用分块因果掩码（block-wise causal mask）替代全双向注意力，使模型能够关注提示、已完成的块和当前块，从而实现精确的分块KV缓存。
三重复合训练目标：训练过程联合优化三个目标：蒸馏损失（指导新取消掩码位置的预测）、一致性损失（强制未取消掩码位置的块内时间一致性）以及辅助的掩码去噪损失（保持基础预测能力）。
高效推理机制：推理时采用分块自回归解码，重用已完成块的KV缓存；在当前块内，基于置信度阈值并行最终化多个Token，并支持遇到结束符时提前停止。
显著降低步数：CDLM将生成所需的细化步数大幅削减了4.1倍至7.7倍，且在大多数任务上准确率几乎没有下降（简单截断步数会导致准确率急剧下降）。
大幅延迟加速：在具体任务上，CDLM在GSM8K-CoT上实现了高达11.2倍的加速，在MBPP-Instruct上实现了高达14.5倍的加速，并达到了最高的每秒Token吞吐量。
硬件利用率平衡：系统分析显示，分块DLM的算术强度（AI）高于受内存带宽限制的自回归（AR）解码，低于受计算限制的原生DLM，在小批处理规模下达到了最佳的效率平衡点。
高扩展性：作为一种训练后方案，CDLM可应用于任何分块扩散模型，未来可通过更大、更强的DLM教师模型提取轨迹来训练中等规模的学生模型。

3) 风险与不足

解码动态差异与输出长度变化：由于CDLM是严格的块因果模型，不同任务可能会表现出不同的解码动态；在保持首选准确率（pass@1）质量的同时，模型可能会生成比原生模型更短的输出。

正文

作者：Minseo Kim, Chenfeng Xu, Coleman Richard Charles Hooper, Harman Singh, Ben Athiwaratkun, Ce Zhang, Kurt Keutzer, Amir Gholami 机构：首尔大学，加州大学伯克利分校，Together AI 日期：2026年2月19日

摘要我们介绍了一致性扩散语言模型（CDLM）。该模型通过将基于一致性的多Token最终化（multi-token finalization）与分块KV缓存（block-wise KV caching）相结合，加速了扩散语言模型的推理过程，在数学和编程任务上实现了高达14.5倍的延迟加速。

扩散语言模型（DLMs）正成为自回归（AR）语言模型极具潜力的替代方案。DLM并非一次生成一个Token，而是通过多个采样步骤迭代地细化部分掩码序列，逐步将完全掩码的序列转化为清晰的文本。这种细化过程带来了一个引人注目的优势：它支持并行生成，允许模型在每次迭代中最终化多个Token，从而可能实现比自回归解码更高的吞吐量。同时，它可以利用双向上下文来解锁文本填充和细化等新功能。

然而在实践中，标准DLM面临两个主要的效率瓶颈：

全双向注意力下的KV缓存不兼容：标准DLM通常使用双向（非因果）注意力机制，这要求在每个去噪步骤中对整个上下文重新计算注意力，导致推理成本高昂，且无法使用标准的KV缓存。
维持质量需要高细化步数：高质量的生成通常需要大量的去噪/细化步骤，通常与生成长度相当。简单地减少步数往往会导致生成质量急剧下降。

CDLM通过一种训练后（post-training）方案同时解决了这两个瓶颈，使得少步数推理变得可靠，同时实现了精确的分块KV缓存。

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

DLM的生成是在N个离散采样步骤上的迭代细化过程。它将 $t = 1$ 时的完全掩码序列转化为 $t = 0$ 时的清晰序列。在每一步中，模型根据当前噪声序列 $X_{t}$ 和提示 $c$ 预测清晰序列分布 $X_{0}$ ：

$p_{θ} (X_{0} ∣ X_{t}, c)$

一种常见的确定性实现是低置信度重掩码（low-confidence remasking）：模型贪婪地取消掩码Token（通常在块内），最终化置信度最高的掩码位置，同时保持其他位置被掩码。这形成了解码轨迹：

$T_{x} = (x_{t_{0}}, x_{t_{1}}, \dots, x_{t_{N}}), t_{k} = 1 - \frac{k}{N}$

该轨迹记录了部分细化的序列是如何逐步演变的。这一轨迹构成了CDLM训练的核心对象。

CDLM 训练机制

1. 轨迹收集 我们通过在特定领域的提示上运行DLM推理来离线收集轨迹。对于每个提示 $x$ ，我们记录Token级别的解码轨迹 $T_{x}$ 、包含Token最终化时刻最后一层隐藏状态的紧凑隐藏状态缓冲区 $H_{x}$ ，以及真实文本 $\overset{y}{^}$ 。具体而言，我们采用分块解码，生成长度为 $L_{g} = 256$ ，块大小为 $B = 32$ ，总步数 $N = L_{g}$ （即在当前块内每步精确最终化一个Token）。这种保守的设置能为知识蒸馏提供更高质量的轨迹。

2. 块因果学生模型与注意力掩码 在轨迹提取期间，我们使用全双向注意力掩码。相比之下，在训练CDLM时，我们采用分块因果掩码（block-wise causal mask），使其关注提示、先前完成的块以及当前解码的块。这种设计使模型能够从全双向模型转换为块扩散模型，从而为已最终化的块实现精确的分块KV缓存。

3. 训练目标 CDLM联合最小化三个目标：

蒸馏损失（针对新取消掩码的位置）：对于在中间状态 $y$ 及其块完成状态 $y^{*}$ 之间新取消掩码的位置，我们将学生模型的预测分布与从存储的隐藏状态中获得的教师模型重建分布进行匹配。直觉上，该目标作为主要锚点，教会学生模型在块因果约束下最终化块内的多个Token。
一致性损失（针对仍被掩码的位置）：我们通过对齐学生模型在状态 $y$ 的预测与其在信息更丰富的状态 $y^{*}$ 下对仍被掩码位置的预测（使用停止梯度目标），来强制实现块内的时间一致性。直觉上，该目标鼓励沿解码轨迹实现稳定的多步过渡。
辅助DLM掩码去噪损失：我们包含了一个应用于随机掩码真实文本的标准掩码去噪目标。直觉上，该目标保留了模型一般的掩码Token预测能力，并有助于保持推理行为，特别是在数学任务上。

4. 推理 在推理阶段，CDLM以分块自回归的方式进行解码，重用提示和所有先前已最终化块的KV缓存。在每个块内，我们应用基于置信度阈值的并行最终化。一旦当前块中出现文本结束（end-of-text）Token，我们也会采用提前停止策略。我们刻意避免引入额外超参数的启发式方法（例如依赖任务的块间并行设置），而是专注于基于精确KV缓存和可靠步数减少的稳健默认解码流程。

主要结果：CDLM–Dream

我们的主要发现包括：

CDLM-Dream在各项基准测试中实现了最大幅度的步数减少，将细化步数削减了约4.1倍至7.7倍，且在大多数任务上准确率变化微小。
这些步数的减少转化为显著的延迟改善：在GSM8K-CoT上加速高达11.2倍，在MBPP-Instruct上加速高达14.5倍。
CDLM通常能达到最高的每秒Token吞吐量。值得注意的是，不同任务可能表现出不同的解码动态，因为CDLM是严格的块因果模型，在保持pass@1质量的同时可能会生成更短的输出。

有效的步数减少：为什么训练至关重要

简单地截断步数会导致准确率显著下降，而CDLM在相似的步数预算下保持了生成质量（并且得益于缓存，延迟降低了约一半）。这突显了核心观点：稳定的多Token细化并非没有代价，它需要通过训练来强制实现与轨迹一致的行为。

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

为了理解硬件利用率，我们分析了随着批处理大小（batch size）的增加，算术强度（AI，即每移动一个字节的FLOPs），并对比了：自回归（AR）解码、原生（全注意力）DLM、以及块大小 $B \in {4, 16, 32}$ 的分块DLM（CDLM）。

核心解读：

AR解码在小批处理规模下严重受限于内存带宽（当bs=1时，AI接近1），随着批处理规模增加，由于权重加载的摊销，性能得以扩展。
原生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实现了更快的推理、更少的步数、更低的延迟和更高的吞吐量，同时保持了极具竞争力的准确率。

参考文献

Beyond Next-Token Prediction: A Performance Characterization of Diffusion versus Autoregressive Language Models
Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding

wsl-docs

探索