テクニカル・サマリー: KVzap - 高速、適応型、かつ忠実なKVキャッシュ・プルーニング

1. 問題提起

トランスフォーマー型言語モデルのコンテキスト長が数万から数十万トークンへと拡大するにつれ、Key-Value (KV) キャッシュが効率的な大規模言語モデル (LLM) 推論における支配的なボトルネックとなっています。中間アテンション状態を保存するKVキャッシュは、膨大なGPUメモリを消費し（例：65Bモデルで128kのシーケンス長の場合335 GB）、Time-to-First-Token (TTFT) を増大させ、デコーディングのスループットを低下させます。

アーキテクチャの変更（例：Grouped Query Attention、Multi-head Latent Attention）やスパース・アテンション機構は、ヘッド ($H$)、次元 ($D$)、レイヤー ($L$) の軸に沿った圧縮には対処してきましたが、シーケンス長軸 ($T$) は、広く採用されている手法では依然としてほとんど対処されていません。既存のアドホックなプルーニング手法は、プロダクション環境へのデプロイメントに不可欠な以下の4つの基準を満たせないことが多いです。

1. 高速かつ軽量であること: オーバーヘッドが無視できる程度であること。
2. フェーズに依存しないこと: プリフィリング（長いコンテキスト）とデコーディング（推論）の両方に適用可能であること。
3. 最適化フレンドリーであること: FlashAttention2やPagedAttentionのようなカーネルと互換性があること。
4. 忠実であること: 精度の低下が最小限であること。

KVzipのような現在の最先端手法は、高い圧縮率を実現しますが、拡張されたプロンプトに対して計算コストの高い二重フォワードパスを必要とするため、オンライン推論やデコーディング・タスクには適していません。

2. 手法

著者らは、プリフィリングとデコーディングの両方で動作する、KVzipの高速かつ入力適応型の近似手法であるKVzapを提案しています。この手法は、3つのコアコンポーネントで構成されています。

2.1. 高度化されたスコアリング・ポリシー (KVzip+)

著者らはまず、KVzipのスコアリングメカニズムを洗練させました。KVzipは、「コンテキスト再構成」に基づいてトークンのスコアを算出します。これは、プロンプトが繰り返された際にそのトークンが受ける最大のアテンション重みを測定するものです。著者らはこれに対し、Devotoら (2025) に着想を得た正規化項を組み込むことで、手法を強化しました。これにより、トークンのバリューベクトルの残差ストリームへの寄与を考慮できるようになります。
KVzip+ スコア ($s^+_i$) は以下のように定義されます：
$$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$$
ここで、$a_{ji}$ はアテンション重数、$W_O$ は出力射影行列、$v_i$ はバリューベクトル、$h_j$ は隠れ状態を表します。

2.2. サロゲート・モデリング

KVzipの禁止的なコストである二重フォワードパスを回避するため、KVzapは、モデルの入力隠れ状態 ($h$) から直接対数スコア ($\log(s^+)$) を予測する軽量なサロゲート・モデル（線形層または2層のMLP）を学習します。

• 入力: 隠れ状態 $h_t \in \mathbb{R}^{D_h}$。
• 出力: 各KVヘッドに対するスコア $\mathbb{R}^H$。
• 学習: モデルは、Nemotron-Pretraining-Datasetからサンプリングされた120万組の (隠れ状態, 対数スコア) ペアを用いて学習されます。MLPバリアントは、隠れ幅 $D_h/8$ とGELU活性化関数を使用します。
• 性能: サロゲートモデルは、KVzip+ オラクルとの平均 $R^2$ 相関において0.60–0.80を達成しており、高コストなスコアが隠れ状態から正確に近似可能であることを示しています。

2.3. 適応型プルーニングとエビクション

固定の予算（例：正確に50%のトークンを保持する）を強制するKVzipとは異なり、KVzapは閾値ベースのプルーニングを採用しています。

• 閾値設定: 予測スコアが固定の閾値 $\tau$ を下回るトークンは破棄されます。これにより、入力プロンプトの情報密度に応じて圧縮率を動的に調整できます。
• スライディング・ウィンドウ: ローカルなコンテキストを保持するため、直近の $w=128$ トークンはStreamingLLMの手法に従って常に保持されます。
• 実装: プルーニングはアテンション操作の後に実行されます。このプロセスは計算効率が高く、オーバーヘッドは極めて低いです（MLPで1.1%、線形モデルで0.02%に限定）。

3. 主な貢献

1. KVzip+: 隠れ状態の正規化を取り入れた、KVキャッシュの重要性に関する改良されたスコアリング指標。
2. サロゲート近似: 軽量なモデルを用いて、KVzip+ スコアが隠れ状態から効果的に近似できることを実証し、二重フォワードパスの必要性を排除したこと。
3. KVzap: これらのサロゲートを閾値ベースのエビクションとスライディング・ウィンドウと共に適用し、速度、フェーズ非依存性、最適化フレンドリー、および忠実性の4つの実用的な基準をすべて満たしながら、最先端の性能を達成する新しいプルーニング技術。

4. 実験結果

実験は、Qwen3-8B、Llama-3.1-8B-Instruct、および Qwen3-32B を対象に、3つのベンチマーク（RULER（長文脈検索/集約）、LongBench（マルチタスク長文脈）、AIME25（数学的推論））を用いて行われました。

• 圧縮率 vs 精度: KVzapは、無視できる程度の精度低下で 2–4倍のKVキャッシュ圧縮 を実現します。
  • RULER 4k において、KVzapはKVzip+ オラクルと同等以上の性能を示し、他の15のプルーニング手法（Expected AttentionやDuo Attentionを含む）を上回りました。
  • LongBench において、KVzapは2–3倍の圧縮までほぼ完璧な精度を維持します。
  • AIME25 において、KVzap-MLPはKVキャッシュの50%以上を破棄した場合でも、推論精度を維持します。
• 適応性: 閾値ベースのアプローチにより、タスクの難易度に基づいて圧縮率を自動的に調整します。例えば、合成データであるRULERよりも、より密度の高い実世界のLongBenchデータに対して、より高い圧縮率を実現しています。
• 効率性: この手法は、計算およびメモリのオーバーヘッドをほとんど追加しません。デコーディング中、追加のFLOPsは、本来であればメモリ帯域幅の制限によって停止しているはずのGPUサイクルを活用します。

5. 意義と主張

本論文は、KVzapが学術的なプルーニング研究と実際の推論エンジンの間の溝を埋めるものであると主張しています。その意義は以下の通りです：

• 実用性: アーキテクチャの変更を必要とせず、速度を犠牲にすることなく最先端の圧縮を実現する初めての手法であり、主要な推論エンジン（vLLM, SGLeng, TRT-LLM）との互換性があります。
• 忠実性: LLMはKVキャッシュを完全には活用しておらず、未使用のペアをベースモデルのエンドツーエンドの再学習なしに、隠れ状態から特定できることを実証しました。
• スケーラビリティ: 8Bから32Bパラメータに及ぶモデル間で検証されており、将来のスパース・アテンション・アーキテクチャとの互換性も考慮されています。

著者らは限界についても認めており、KVzapはエンドツーエンドの学習目的ではなく事後的な追加要素であること、また、最適化されたカーネル（例：PagedAttention）を通じて圧縮率を実際の実行時間の高速化に変換するにはさらなるエンジニアリングが必要であることを指摘しています。しかし、シンプルさ、高い圧縮率、そして堅牢な性能の組み合わせにより、KVzapはプロダクション環境へのデプロイメントの有力な候補であると位置付けています。