自回归语言模型长期以来是文本生成的主流范式,逐 token 预测的计算方式虽然逻辑清晰,但在长序列生成时面临严重的内存带宽瓶颈和顺序依赖延迟。扩散语言模型(Diffusion Language Models, DLMs)的出现为这一困境提供了另一种思路:不再逐 token 生成,而是通过多步迭代逐步去除噪声,最终从完全掩码的序列恢复出目标文本。这种并行生成机制理论上具备更高的吞吐量潜力,然而实际部署中仍有两大致命缺陷阻碍其广泛应用。首先,全双向注意力机制在每一步去噪时都需要重新计算完整上下文的注意力矩阵,导致标准 KV 缓存技术失效,计算成本随序列长度平方增长;其次,高质量生成往往需要与目标长度相当的去噪步数,步数减少会导致输出质量急剧下降。这两个问题严重制约了扩散语言模型的推理效率,使其难以在延迟敏感的生产环境中落地。
一致性扩散语言模型(Consistency Diffusion Language Models, CDLM)针对上述两项瓶颈提出了系统性的解决方案,其核心思想是将图像生成领域的一致性模型训练范式迁移到语言模型,并通过块级因果注意力机制实现精确的 KV 缓存。该方法在 7B 参数规模的 Dream 模型上实现了 4.1 倍至 7.7 倍的步数压缩,在 GSM8K 数学推理任务中将延迟降低至原来的 1/11.2,在 MBPP 代码生成任务中更是达到了 14.5 倍的加速效果。与传统扩散模型需要数十甚至上百步迭代不同,CDLM 能够在少数几步内完成高质量的 token 定稿,显著提升了每秒处理的 token 数量。
CDLM 的训练流程包含四个关键阶段。首先是轨迹采集阶段,模型在特定领域的提示词上运行推理,记录每一步的部分精化序列、隐藏状态缓冲区以及对应的标准答案。在实际训练中采用了块大小为 32、生成长度 256 的配置,总计进行 256 步迭代,这一保守设置能够产生高质量的蒸馏轨迹。其次是块级因果掩码的设计,训练阶段使用完整的双向注意力,而在推理阶段切换为块级因果掩码,使其能够关注提示词、已完成的前序块以及当前解码块,这种设计使得模型能够在保持局部双向上下文的同时支持精确的 KV 缓存复用。第三个阶段是三目标联合训练,蒸馏损失针对新解锁的 token 位置,将学生模型的预测分布与教师模型在轨迹中存储的隐藏状态所对应的预测分布对齐;一致性损失针对仍处于掩码状态的 token 位置,强制学生在中间状态和块完成状态下的预测保持一致;辅助去噪损失则使用随机掩码的标准去噪目标,保持模型的通用掩码 token 预测能力。最后是推理阶段,模型以块级自回归方式解码,利用已缓存的提示词和所有已完成块的 KV 值,在当前块内部采用置信度阈值并行定稿,当块内出现结束符时立即停止。
从系统层面的算术强度分析来看,CDLM 处于自回归模型与全注意力扩散模型之间的最佳平衡点。自回归解码在小批量时严重受限于内存带宽,算术强度接近 1;全双向注意力的扩散模型即使在单样本时也高度 compute-bound,因为每步都需要处理完整序列的注意力计算;而块级扩散模型通过块内并行处理在相似内存开销下实现了更高的算术强度,同时避免了全注意力的二次方计算量。这一特性解释了为何 CDLM 在小批量推理场景下能够提供显著优于传统方案的效率提升。
在工程实践中,CDLM 的部署需要关注几个关键参数。块大小 B 决定了每一步并行处理的 token 数量,较大的块能够提高算术强度但会降低 KV 缓存的精细度,建议在 16 至 32 范围内根据目标硬件进行调优。生成长度 Lg 与总步数 N 的比例直接影响推理延迟和质量保留,对于需要长输出的任务可以适当增加步数预算,但对于短响应场景可以激进地压缩至原始步数的四分之一左右。置信度阈值决定了块内并行定稿的激进程度,较低的阈值会更快地锁定 token 但可能引入错误传播,建议通过验证集上的通过率指标进行校准。此外,早停策略应当基于当前块内是否出现结束符而非全局序列长度,这能够在保持输出完整性的同时避免不必要的计算浪费。
CDLM 的意义不仅在于单点加速,更在于它验证了一条将图像生成领域的一致性模型思想迁移到语言模型的可行路径。随着更强大的扩散语言模型骨干网络不断涌现,CDLM 这类后训练优化方案的价值将进一步放大,未来可以期待从更大规模的教师模型采集轨迹并蒸馏出更高效的学生模型,从而在保持生成质量的前提下持续压缩推理成本。
资料来源:Together AI 博客《Consistency diffusion language models: Up to 14x faster inference without sacrificing quality》