arXiv:2505.15781cs.CL2025-05NeurIPS被引 131

为扩散语言模型设计缓存机制,实现2-10倍推理加速

dKV-Cache: The Cache for Diffusion Language Models

  • 提出延迟条件缓存策略,分步缓存键值状态
  • 实现2-10倍推理速度提升,长序列性能反而更好
  • 无需训练即可适配现有模型,适合高效部署场景

扩散语言模型(DLMs)被视为自回归模型的有力竞争者,但长期受制于推理缓慢。其非自回归架构与双向注意力机制导致无法使用传统键值缓存加速解码。本文提出一种针对去噪过程的延迟键值缓存机制——dKV-Cache。基于不同词元在扩散过程中表征动态差异的观察,设计了两种互补变体:dKV-Cache-Decode 实现近乎无损加速,甚至在长序列上提升性能,表明现有DLM在推理中可能未充分利用上下文信息;dKV-Cache-Greedy 采用激进缓存策略,寿命缩短,实现二次方时间复杂度下的更高加速,代价是轻微性能下降。实验表明,dKV-Cache 在多个基准测试中均实现加速,涵盖通用语言理解、数学推理与代码生成任务。结果证明,该缓存可直接应用于现有DLM,且无需额外训练。

原文摘要 · Abstract (English)

Diffusion Language Models (DLMs) have been seen as a promising competitor for autoregressive language models. However, diffusion language models have long been constrained by slow inference. A core challenge is that their non-autoregressive architecture and bidirectional attention preclude the key-value cache that accelerates decoding. We address this bottleneck by proposing a KV-cache-like mechanism, delayed KV-Cache, for the denoising process of DLMs. Our approach is motivated by the observation that different tokens have distinct representation dynamics throughout the diffusion process. Accordingly, we propose a delayed and conditioned caching strategy for key and value states. We design two complementary variants to cache key and value step-by-step: (1) dKV-Cache-Decode, which provides almost lossless acceleration, and even improves performance on long sequences, suggesting that existing DLMs may under-utilise contextual information during inference. (2) dKV-Cache-Greedy, which has aggressive caching with reduced lifespan, achieving higher speed-ups with quadratic time complexity at the cost of some performance degradation. dKV-Cache, in final, achieves from 2-10x speedup in inference, largely narrowing the gap between ARs and DLMs. We evaluate our dKV-Cache on several benchmarks, delivering acceleration across general language understanding, mathematical, and code-generation benchmarks. Experiments demonstrate that cache can also be used in DLMs, even in a training-free manner from current DLMs.

扩散模型推理加速缓存机制语言模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。