dKV-Cache: 拡散言語モデルのためのキャッシュ

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

著者: Xinyin Ma, Runpeng Yu, Gongfan Fang, Xinchao Wang

cs.AI

要旨

Diffusion Language Models（DLMs）は、自己回帰型言語モデルにとって有望な競合相手と見なされてきました。しかし、Diffusion言語モデルは長い間、推論速度の遅さに制約されてきました。その核心的な課題は、非自己回帰型アーキテクチャと双方向アテンションが、デコードを加速するキー・バリューキャッシュを妨げていることです。我々はこのボトルネックを解決するため、DLMのノイズ除去プロセスに対して、遅延型KVキャッシュ（delayed KV-Cache）と呼ばれるKVキャッシュに似たメカニズムを提案します。このアプローチは、異なるトークンが拡散プロセス全体で異なる表現ダイナミクスを持つという観察に基づいています。それに応じて、キーとバリューの状態に対して、遅延的かつ条件付きのキャッシュ戦略を提案します。我々は、キーとバリューを段階的にキャッシュするための2つの補完的なバリアントを設計しました：(1) dKV-Cache-Decodeは、ほぼロスレスな加速を提供し、長いシーケンスでは性能さえ向上させ、既存のDLMが推論中に文脈情報を十分に活用していない可能性を示唆しています。(2) dKV-Cache-Greedyは、寿命を短縮した積極的なキャッシュを行い、二次時間計算量でより高い高速化を実現しますが、一部の性能低下を伴います。最終的に、dKV-Cacheは推論において2～10倍の高速化を達成し、AR（自己回帰型モデル）とDLMのギャップを大幅に縮めました。我々はdKV-Cacheを複数のベンチマークで評価し、一般的な言語理解、数学、コード生成のベンチマーク全体で加速を実現しました。実験結果は、キャッシュがDLMでも使用可能であり、現在のDLMからトレーニング不要で適用できることを示しています。

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.