拡散型大規模言語モデルにおけるKVキャッシュにはAttentionが全てである

要旨

本研究では、拡散型大規模言語モデル（DLMs）におけるキー・バリュー（KV）キャッシュを適応的に再計算し、予測精度を最大化しながらデコード遅延を最小化する方法を検討する。従来の手法では、すべてのデノイジングステップとレイヤーにおいて、すべてのトークンに対してQKVを再計算していたが、特に浅いレイヤーではKV状態がほとんど変化しないため、大幅な冗長性が生じていた。我々は以下の3つの観察を行った：(1) 遠く離れた{bf MASK}トークンは主に長さバイアスとして機能し、アクティブな予測ウィンドウを超えてブロック単位でキャッシュ可能である、(2) KVの動的変化は深さとともに増加するため、深いレイヤーから選択的にリフレッシュを行うことで十分である、(3) 最も注目されるトークンはKVのドリフトが最小であり、他のトークンのキャッシュ変更に対する保守的な下限を提供する。これらの観察に基づき、我々は{bf Elastic-Cache}を提案する。これは、トレーニング不要でアーキテクチャに依存しない戦略であり、{いつ}リフレッシュするか（最も注目されるトークンに対する注意認識ドリフトテストにより決定）と{どこで}リフレッシュするか（選択されたレイヤー以降を再計算し、浅いレイヤーのキャッシュとウィンドウ外のMASKキャッシュを再利用する深さ認識スケジュールにより決定）を共同で決定する。固定周期スキームとは異なり、Elastic-Cacheは拡散型LLMに対して適応的かつレイヤー認識型のキャッシュ更新を行い、冗長な計算を削減し、生成品質の低下をほとんど伴わずにデコードを加速する。LLaDA-Instruct、LLaDA-1.5、およびLLaDA-Vを用いた数学的推論とコード生成タスクにおける実験では、一貫した高速化が確認された：GSM8K（256トークン）で8.7倍、長いシーケンスで45.1倍、HumanEvalで4.8倍の高速化を達成し、ベースラインよりも高い精度を維持した。我々の手法は、既存の信頼度ベースのアプローチよりも大幅に高いスループット（GSM8Kで6.8倍）を実現しつつ、生成品質を維持し、拡散型LLMの実用的な展開を可能にする。

English

This work studies how to adaptively recompute key-value (KV) caches for diffusion large language models (DLMs) to maximize prediction accuracy while minimizing decoding latency. Prior methods' decoders recompute QKV for all tokens at every denoising step and layer, despite KV states changing little across most steps, especially in shallow layers, leading to substantial redundancy. We make three observations: (1) distant {bf MASK} tokens primarily act as a length-bias and can be cached block-wise beyond the active prediction window; (2) KV dynamics increase with depth, suggesting that selective refresh starting from deeper layers is sufficient; and (3) the most-attended token exhibits the smallest KV drift, providing a conservative lower bound on cache change for other tokens. Building on these, we propose {bf Elastic-Cache}, a training-free, architecture-agnostic strategy that jointly decides {when} to refresh (via an attention-aware drift test on the most-attended token) and {where} to refresh (via a depth-aware schedule that recomputes from a chosen layer onward while reusing shallow-layer caches and off-window MASK caches). Unlike fixed-period schemes, Elastic-Cache performs adaptive, layer-aware cache updates for diffusion LLMs, reducing redundant computation and accelerating decoding with negligible loss in generation quality. Experiments on LLaDA-Instruct, LLaDA-1.5, and LLaDA-V across mathematical reasoning and code generation tasks demonstrate consistent speedups: 8.7times on GSM8K (256 tokens), 45.1times on longer sequences, and 4.8times on HumanEval, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput (6.8times on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.

拡散型大規模言語モデルにおけるKVキャッシュにはAttentionが全てである

Attention Is All You Need for KV Cache in Diffusion LLMs

要旨

Support