技術概要：xKV - 整合された特異ベクトル抽出によるクロスレイヤー KV キャッシュ圧縮

1. 問題定義

長文脈大規模言語モデル（LLM）は、大規模情報検索やコードベースのデバッグなどの高度なアプリケーションを可能にしますが、重要なボトルネックに直面しています。それは推論中のキー・バリュー（KV）キャッシュのメモリ消費量です。文脈長が数十万、あるいは数百万トークンにまで拡張するにつれ、KV キャッシュのメモリフットプリントは許容できなくなるほど膨大となり、並列推論リクエストやスループットを深刻に制限しています。

既存の緩和策は一般的に 2 つのカテゴリに分類されます。

• レイヤー内手法: 量子化、トークン排除、単一レイヤー低ランク分解などの技術は、各レイヤーの KV キャッシュを独立して圧縮します。ある程度効果的ですが、レイヤー間にある潜在的な冗長性を活用できていません。
• クロスレイヤー手法: クロスレイヤーアテンション（CLA）や YOCO などのアプローチは、高価なアーキテクチャ変更と事前学習を必要とします。一方、MiniCache のような手法は、トークンごとのコサイン類似度に基づいて隣接レイヤーを事後にマージすることに依存しています。しかし、実証分析により、レイヤー間のトークンごとの類似度は、精度の大幅な低下なしに堅牢な圧縮を支援するにはしばしば弱すぎるものであることが示されています。

2. 手法：xKV

著者は、モデルの再学習やアーキテクチャ変更を必要とせず、レイヤー間の深い構造的冗長性を活用する、学習後のプラグアンドプレイ圧縮手法であるxKVを提案します。

2.1. 核心的な発見：整合された特異ベクトル

本論文は、トークンごとのコサイン類似度に基づいてレイヤーが区別されるという仮定に異議を唱えています。**Centered Kernel Alignment（CKA）**を用いることで、著者は個々のトークン埋め込みは分岐する可能性があっても、支配的特異ベクトル（トークン空間における主要な変動を定義する基底）は隣接レイヤー間で高度に整合していることを実証しています。この高い CKA スコアは、ポイントごとの類似度が低くても、KV キャッシュ部分空間の幾何学的構造が共有されていることを示しています。

2.2. クロスレイヤー因数分解（CLF）

この整合性を利用し、xKV はモデルの $N$ レイヤーをサイズ $W$ の連続するウィンドウに分割します。各ウィンドウに対して、クロスレイヤー因数分解を実行します。

1. 連結: ウィンドウ内の全レイヤーの KV キャッシュを水平方向に連結します。
2. SVD: 連結されたキャッシュに対してランク $r$ の特異値分解（SVD）を計算します。
3. 分解: 結果として、ウィンドウ用の共有トークン基底（$A_k$）とレイヤー固有の再構成行列（$B_\ell$）が得られます。
4. 保存: $W$ 個のフルランクキャッシュを保存する代わりに、システムは 1 つの共有基底と $W$ 個のより小さな再構成行列を保存します。これにより、メモリ複雑度が $O(W \cdot L \cdot d)$ から $O(L \cdot r + W \cdot r \cdot d)$ に削減されます。

2.3. 選択的再構成（SR）

デコードフェーズ中にすべてのトークンを再構成することによる計算ボトルネック（シーケンス長 $L$ に比例して増加）に対処するため、xKV は**選択的再構成（SR）**を導入します。

• メカニズム: 密な再構成の代わりに、SR はランドマーク誘導チャンク選択（ShadowKV に類似）を用いて、現在のクエリと最も関連性の高いトークンのスパースな部分集合（$S$）を特定します。
• 実行: 共有基底とレイヤー固有の行列を使用して、これらの重要なトークンに対応する行のみを再構成します。
• 利点: これにより、再構成コストをシーケンス長に対して一定に抑え、エンドツーエンドの高速化を可能にします。

3. 主要な貢献

1. クロスレイヤー冗長性の発見: 本論文は、KV キャッシュがレイヤー間で高度に整合した支配的特異ベクトルを共有していることを明らかにしました。これは、低いトークンごとの類似度によって隠蔽されていた冗長性の形態です。
2. xKV アルゴリズム: クロスレイヤー因数分解（CLF）を通じて共有基底を抽出する学習不要な手法であり、$W$ レイヤー（$W > 2$）間での効果的な情報共有を可能にします。
3. 選択的再構成: 推論時にクエリに関連するトークンのみを再構成する最適化であり、密な再構成に伴う線形計算オーバーヘッドを排除します。
4. プラグアンドプレイ互換性: この手法は微調整やアーキテクチャ変更を必要とせず、既存の事前学習済みモデル（例：Llama-3.1、Qwen2.5、DeepSeek-V2）に適用可能です。

4. 実験結果

著者は、RULER、LongBench、Multi-turn NIAH ベンチマークを用いて、Llama-3.1-8B-Instruct、Qwen2.5-7B/14B-Instruct、DeepSeek-Coder-V2-Lite における xKV を評価しました。

• 圧縮率と精度:
  • xKV は、フルアテンションのベースラインと比較して、最小限の精度低下（平均で≤3%）で最大8 倍の KV キャッシュ圧縮を達成します。
  • Llama-3.1-8B において、xKV は 8 倍圧縮で RULER 上で**88.50%**の平均精度を達成し、MiniCache（1.3 倍圧縮で 45.04%）を大幅に上回り、SnapKV や PyramidKV などのトークン排除手法と同等かそれ以上の性能を示しました。
  • マルチターン NIAH タスクにおいて、xKV は会話のターンを通じて安定した性能を維持しますが、トークン排除手法は最初のターン後に急激な低下を招きます。
• 効率性とスループット:
  • メモリ: xKV は、同じ GPU メモリ予算下でより大きなバッチサイズを可能にします。
  • レイテンシ: 密な再構成により、アテンションレイテンシが最大3.6 倍削減されます。
  • スループット: 選択的再構成（xKV-SR）と組み合わせることで、NVIDIA A100 GPU においてフルアテンションベースラインに対して最大4.23 倍のエンドツーエンド高速化を達成します。
  • 比較: xKV-SR は、同程度の精度レベルにおいてスループットで ShadowKV を30% 上回り、Keys と Values の両方を GPU 上に保持することで、PCIe 転送のボトルネックを回避しています。
• 堅牢性: この手法は、DeepSeek-V2 における Grouped-Query Attention（GQA）や Multi-head Latent Attention（MLA）を含む、異なるモデルアーキテクチャにおいても効果的です。

5. 意義と主張

本論文は、xKV が長文脈 LLM 推論のメモリフットプリントとレイテンシを同時に削減する実用的なプラグアンドプレイアプローチを提供すると主張しています。トークンごとの類似度から特異ベクトルの整合性へと焦点を移すことで、xKV は従来の手法が見逃していた未活用のレイヤー間冗長性を解き放ちます。

著者は、彼らのアプローチが以下の点で優れていることを強調しています。

• アーキテクチャ変更に必要な事前学習の高コストを回避する。
• 類似度ベースのマージ（MiniCache）の精度制限を克服する。
• 選択的再構成を通じて、密な低ランク再構成の計算ボトルネックを解決する。
• ShadowKV のようなオフロードベースの手法を制約する PCIe バンド幅の制限を回避し、圧縮されたキャッシュの完全なオンデバイス（GPU）推論を可能にする。

この研究は、効率的な長文脈推論に向けた基礎的なステップとして提示され、採用を促進するためにコードは公開されています。