dKV-Cache：扩散语言模型的缓存

扩散语言模型 (DLMs) 被视为自回归语言模型的一个有前景的竞争对手。然而，扩散语言模型长期以来一直受限于缓慢的推理速度。一个核心挑战在于其非自回归架构和双向注意力排除了加速解码的关键值缓存。我们通过提出一种类似 KV 缓存的机制——延迟 KV 缓存 (delayed KV-Cache)——用于 DLM 的去噪过程，从而解决了这一瓶颈。我们的方法是基于观察到不同 token 在整个扩散过程中具有不同的表示动态性。因此，我们为关键和值状态提出了一个延迟和条件化的缓存策略。我们设计了两种互补的变体来逐步缓存关键和值：(1) dKV-Cache-Decode，它提供了几乎无损的加速，甚至改善了长序列上的性能，这表明现有的 DLM 可能在推理过程中未能充分利用上下文信息。(2) dKV-Cache-Greedy，它采用激进的缓存策略，生命周期较短，以性能下降为代价实现了更高的速度提升（具有二次时间复杂度）。最终，dKV-Cache 实现了 2-10 倍的推理加速，极大地缩小了自回归模型 (ARs) 和 DLM 之间的差距。我们在多个基准测试中评估了 dKV-Cache，它在通用语言理解、数学和代码生成基准测试中均提供了加速。实验表明，缓存也可以用于 DLM，甚至可以在不进行额外训练的情况下应用于现有的 DLM。