テクニカル・サマリー: SelKV — パー・トークン単位のマージ・または・ドロップとアテンション補償による選択的KVキャッシュ・マージ

1. 問題提起

大規模言語モデル（LLM）は、自己回帰的な推論において、キー・バリュー（KV）キャッシュという重大なボトルネックに直面しています。コンテキスト長（$n$）が増加するにつれ、KVキャッシュのメモリ・フットプリントは線形（$O(L \cdot H \cdot d \cdot n)$）に増大し、ロングコンテキスト・アプリケーションにおけるバッチサイズ、スループット、およびデプロイコストを厳しく制限します。

既存の圧縮手法はこれを軽減しようと試みていますが、主に2つの制限事項に悩まされています：

1. 一律なマージ・または・ドロップの決定: 現在のアプローチは通常、バイナリ戦略を採用しています。つまり、すべての除外されたトークンをマージするか（例：KVMerger, WeightedKV）、あるいは完全にドロップするか（例：H2O, SnapKV）のいずれかです。この粒度の欠如により、除外されたトークンの値ベクトルがマージ対象のターゲットと類似していない場合、加重平均によって保持された意味表現が損なわれてしまいます。逆に、無差別なドロップは貴重な情報を失わせます。
2. アテンション・サグ（Attention Sag）: 複数のトークンが単一のキャッシュ・エントリにマージされると、そのエントリは単一の未マージ・トークンと同じソフトマックス・アテンション質量を受け取ることになります。これにより、「アテンション・サグ」が発生します。マージされた位置は、マージされたトークンの数（$\sim m$）に比例して系統的にアテンションが過小に割り当てられ、モデルが保持された情報を十分に活用できなくなります。

2. メソドロジー: SelKV

著者らは、これらの制限に対処するために設計された、トレーニング不要のデュアル・コンポーネント・フレームワークである SelKV（Selective KV Cache Merging）を提案しています。SelKVは、プリフィル（prefill）フェーズの後の既存のマージ・パイプラインに統合可能な、プラグ可能なモジュールとして機能します。

コア・コンポーネント

1. ソフト・コサイン・ゲート（一律な決定への対処）
バイナリのマージ・または・ドロップ決定の代わりに、SelKVは値ベクトルの類似性に基づく、連続的なパー・トークン・ゲーティング・メカニズムを導入しています。

• メカニズム: マージ・ターゲットにルーティングされる各除外トークンに対して、システムはそれらの値ベクトル間のコサイン類似度を計算します。
• ゲーティング関数: ソフト・ゲート $g = \max(\cos\_sim, 0)$ がマージの強度を調整します：
  • $g \approx 1$: 高い類似性 $\rightarrow$ フル・マージ。
  • $g = 0$: 非類似（直交）ベクトル $\rightarrow$ トークンをドロップ。
  • $0 < g < 1$: 中程度の類似性 $\rightarrow$ 部分的マージ。
• 実装: これには学習パラメータや閾値のチューニングは不要です。保持されたトークン $i$ のマージされた値は、元の値と、それぞれのゲートによってスケールされた除外トークンの寄与によるアテンション加重平均として計算されます。

2. アテンション補償（アテンション・サグへの対処）
マージされた位置の系統的なアテンション不足を修正するために、SelKVはデコーディング時のロジット・バイアスを適用します。

• メカニズム: 生のマージ数（高い圧縮率では不安定になる可能性がある）を使用する従来の手法とは異なり、SelKVはプリフィル統計からアテンション・レシオ（$R_{h,i}$）を導出します。この比率は、マージされた位置が元の状態に対して受けるべき総アテンション質量を推定します。
• 補正: デコーディング中にバイアス項 $\alpha \cdot \log(R_i)$ がアテンション・ロジットに加算されます。これにより、ソフトマックス分布が再校正され、マージされた位置が現在表現している情報に比例した適切なアテンション質量を受け取れるようになります。

パイプラインの概要

SelKVのパイプラインは、プリフィル後の6つのステージで進行します：

1. 重要度スコアリング: トークンは、アテンション・ウェイト単独ではなく、（VATPに従い）アテンション・ウェイトと値ベクトルの大きさの積に基づいてスコアリングされます。
2. トークン選択: 層ごと（MHAの場合はヘッドごと、GQAの場合は結合による）に、上位 $m_s$ 個のトークンを保持するトークンと共に、固定された直近のウィンドウが選択されます。
3. マージ・ターゲット・ルーティング: 除外されたトークンは、プリフィル時のアテンション・パターンに基づき、位置バケット内の保持されたトークンへとルーティングされます。
4. 選択的マージ: ソフト・コサイン・ゲートが、各除外トークンのマージ強度を決定します。
5. アテンション補償: ロジット・バイアスが計算され、適用されます。
6. RoPE 再配置: キーは連続した位置に再埋め込みされます（位置の歪みを避けるため、ロングコンテキスト評価では無効化されています）。

3. 主な貢献

• パー・トークン単位の粒度: 意味的類似性に基づいてマージ強度を連続的に変調することを可能にするソフト・コサイン・ゲートの導入により、従来の手法にある「全か無か」の欠陥を排除しました。
• 安定したアテンション補償: 生のカウントベースの手法よりも、特に重度の圧縮下でより堅牢にアテンション・サグを修正する、アテンション・レシオに基づくロジット・バイアスを導入しました。
• トレーニング不要の統合: 本フレームワークはファインチューニングを必要とせず、既存のワンショット圧縮パイプラインに統合可能です。
• アーキテクチャへの適応性: 本手法は、MHAおよびGQAアーキテクチャの両方をサポートしており、GQAモデルにおけるトークン選択のユニオン（結合）に関する特定の処理を備えています。

4. 実験結果

フレーム本フレームワークは、LongBench（16の英語データセット）において、3つのモデル（LongChat-7B (MHA), LLaMA-3.1-8B (GQA), Gemma-2-9B (GQA)）を用いて評価されました。なお、KVキャッシュの**25%**のみを保持しています。

• ベースラインとの比較: SelKVは、代表的なワンショット・ベースライン（SnapKV, LOOK-M, PyramidKV）を一貫して上回りました。
  • GQAモデルにおいて、SelKVはフル・キャッシュ・ベースラインに最も近い性能を達成しました（Gemma-2で $-0.67$、LLaMA-3.1で $-0.73$ の平均差）。
  • MHAモデル（LongChat）では、特定の31.5kトークンの設定においてPyramidKVがわずかに優れた性能を示しましたが、SelKVも競争力のある性能を維持しました。
• マルチドキュメントQA: 特筆すべきことに、SelKVはいくつかのマルチドキュメントQAタスク（例：LLaMA-3.1におけるHotpotQAおよびMuSiQue）において、フル・キャッシュ・ベースラインを上回りました。著者らは、これを選択的マージが、重要な証拠を保持しつつ、注意をそらすコンテキストを抑制する「暗黙的なアテンション・フィルター」として機能している証拠であると解釈しています。
• 効率性: コンテキスト長が100kトークンの場合、圧縮されたキャッシュはフル・キャッシュ・ベースラインと比較して3.3倍のデコーディング速度向上を実現しました。
• 堅牢性: 本手法は、保持率10%〜90%の様々な圧縮比、および要約、フューショット学習、コード補完を含む様々なタスクカテゴリにおいて、堅牢性を示しました。

5. 意義と主張

本論文は、SelKVを、品質を犠牲にすることなく、ロングコンテキストLLMの推論を現実的なものにするための重要な一歩として位置付けています。著者らは以下のことを主張しています：

• 選択的マージは一律な戦略よりも優れている: 表現の類似性に基づいてマージ強度を変調することで、SelKVは無差別なマージによる意味的破壊を回避します。
• アテンション補償は不可欠である: アテンション・サグを補正しなければ、マージされたトークンは十分に活用されません。SelKVのレシオベースの補正は、保持された情報が効果的にアクセスされることを保証します。
• 暗黙的フィルタリング: 複雑なマルチドキュメントQAにおいてSelKVがフル・キャッシュ・ベースラインを上回る能力は、圧縮プロセス自体がノイズをフィルタリングすることでモデルの性能を高め得ることを示唆しており、これは「圧縮は常に能力を低下させる」という仮説に挑戦する発見です。

著者らは、SelKVが、生成品質をほぼ損なうことなく、特にGQAベースのアーキテクチャにおいて、メモリ制約のある環境でのLLMデプロイのための実用的かつパラメータフリーなソリューションを提供すると結論付けています。