技術概要：アテンションマッチングによる高速 KV コンパクト化

1. 問題定義

大規模言語モデル（LLM）を長文脈に拡張する現状は、キー・バリュー（KV）キャッシュのメモリ要件によってボトルネックとなっています。自己回帰型トランスフォーマーアーキテクチャでは、KV キャッシュはすべての過去のトークンのキーとバリューを保持する必要があり、長期的な設定（例：マルチセッション対話、エージェント型コーディング）では、リクエストあたりギガバイト規模に達することがあります。

KV キャッシュサイズを削減する既存のアプローチは、概ね 2 つのカテゴリに分類され、どちらも重大な限界を有しています。

• トークンスペース手法: トークンエビクション（H2O、SnapKV）、トークンマージ（KVMerger）、要約化などの技術は、高い圧縮率で急速に性能が劣化します。特に要約化は損失が大きく、正確な情報抽出を必要とするタスクにおける下流性能を損ないます。
• 潜在空間手法: Cartridges（Eyuboglu ら、2025）などの最近の研究は、完全な文脈の性能に一致するように、潜在空間でコンパクトな KV キャッシュを訓練できることを示しました。しかし、このアプローチはすべての文脈に対してエンドツーエンドの勾配ベース最適化（プレフィックスチューニング）に依存しており、極めて遅く、高コストです（文脈あたり数 GPU 時間を要する）。

未解決の問題は、文脈ごとの勾配最適化の計算コストを伴わずに、下流モデルの挙動を保持する高速かつ高品質な文脈コンパクト化を潜在空間で達成することです。

2. 手法：アテンションマッチング

著者は、アテンションマッチングに基づく一連の手法を提案します。これは、参照クエリのセットに対して、元のキャッシュ $(K, V)$ のアテンション出力とアテンション質量を再現するコンパクトな KV キャッシュ $(C_k, C_v)$ を構築します。

核心的な目的

目標は、サイズ $t \ll T$ のコンパクトなキャッシュ $(C_k, C_v)$ を見つけることで、任意のクエリ $q$ および任意の追加された固定/将来のトークン $(K_{fixed}, V_{fixed})$ に対して、アテンション出力がほぼ保持されるようにすることです。著者は、連結ブロックに対するアテンションの分解を利用し、参照クエリ $Q_{ref}$ のセットに対して以下の 2 つの量を一致させるだけで十分であることを示しています。

1. ローカルアテンション出力: バリューの重み付き和。
2. アテンション質量: ソフトマックスの分母である、正規化されていないアテンションロジットの和。

コンパクト化されたキャッシュはトークン数が少なく（$t < T$）、したがって（例えば $q=0$ の場合、質量は $T$ 対 $t$）、自然には総アテンション質量を一致できないという不一致に対処するため、この手法はトークンごとのスカラーバイアス $\beta$ を導入します。最適化のターゲットは以下の通りです。
$$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$$
$$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$$

アルゴリズムパイプライン

このフレームワークは、コンパクト化時に勾配降下を回避し、効率的な閉形式の部分問題に分解します。

1. 参照クエリ生成 ($Q_{ref}$):

   • リピートプレフィル: 文脈自体（またはその反復版）に対してプレフィルパスを実行し、クエリアクティベーションを抽出します。
   • セルフスタディ: 文脈を条件とした合成相互作用を生成し、クエリ分布を広げます。
   • オンポリシー: 層を順次コンパクト化し、部分的にコンパクト化されたモデルからクエリを抽出することで、分布シフトを低減します。

2. キー選択 ($C_k$):

   • 最高アテンションキー: 参照クエリ全体で最も高いアテンションスコアを受け取る元のキーを選択します（RMS 経由で集約）。
   • 直交マッチングパーシュート (OMP): アテンション質量マッチングの目的を最もよく満たすようにキーを貪欲に選択し、その後、重みをフィットさせるために非負最小二乗法（NNLS）を適用します。

3. バイアスとバリューフィット ($\beta, C_v$):

   • $\beta$ のフィット: 元のキーのアテンション質量に一致させるために、非負最小二乗法（NNLS）で解きます。
   • $C_v$ のフィット: 選択されたキーとフィットされたバイアスを前提として、アテンション出力に一致させるために、通常最小二乗法（OLS）で解きます。

非一様コンパクト化

著者は、異なるアテンションヘッドがコンパクト化に対して異なる感度を示すことを観察しました。これに基づき、異なるヘッドと層に異なる保持率を割り当てる非一様コンパクト化スケジュールを提案します。このスケジュールは、ヘッドごとの感度曲線に基づく貪欲な交換アルゴリズムを用いてモデルごとに 1 回事前計算され、異なる文脈に対して堅牢に適用可能となります。

チャンク化されたコンパクト化

非常に長い文脈（例：60k 以上のトークン）をサポートするため、この手法は入力の一続きのチャンクに対して独立してコンパクト化を適用します。グローバルな位置意味論とチャンク間相互作用を保持するために、デフォルトではKV ベースのチャンキングアプローチ（完全な文脈のプレフィル、KV 状態のスライス、コンパクト化、マージ）が使用されます。

3. 主要な貢献

• アテンションマッチング定式化: 出力尤度ではなく、アテンション挙動の保持を直接最適化する、潜在空間コンパクト化のための新規目的関数。これは効率的な線形代数問題に自然に分解されます。
• 閉形式解: コンパクト化中のエンドツーエンドの勾配ベース最適化の必要性を排除し、最小二乗法および NNLS ソルバーに置き換えます。
• 速度 - 品質のパレートフロンティア: このアプローチは、単一 GPU で数秒以内に最大 50 倍の圧縮率を達成し、Cartridges よりも桁違いに高速でありながら、同等またはそれ以上の下流精度を維持します。
• 非一様スケジューリング: コンパクト化に対して最も敏感なアテンションヘッドに容量を優先する、モデル固有の事前計算済み予算配分戦略。
• バイアス機構: アテンション質量の不一致を修正するためのトークンごとのスカラーバイアスの導入により、削減されたキーのセットで正確な質量マッチングを可能にします。

4. 結果

この手法は、QuALITY（長文書 QA）、LongHealth（高密度臨床記録）、QASPER（生成 QA）、およびLongBench v2で評価されました。

• 性能対速度: QuALITY において 50 倍の圧縮率で、アテンションマッチング手法はパレートフロンティアを描き、トークンエビクションのベースライン（H2O+、KVzip）や要約化を上回ります。Cartridges と同等またはそれ以上の精度を達成しつつ、2 桁高速（数時間対数分）です。
• 高圧縮: 極端な比率（最大 100 倍）においても、これらの手法は重要な下流精度を保持します（例：LongHealth で約 60% 対、要約化ベースラインの約 30%）。
• 汎化: このアプローチは生成タスク（QASPER）や検索集約的なベンチマーク（RULER）に汎化し、最速のバリアント（AM-HighestAttnKeys-fast）は RULER において最先端の KVzip を上回りました。
• 効率性: Gemma-3-12B における 60k トークン文脈のコンパクト化は、キー選択戦略（OMP 対 Highest Attention）に応じて約 100〜600 秒を要し、これは主にクエリ生成とキー選択のステップによって支配されます。
• オンラインコンパクト化: AIME に関する予備結果は、軌道中の繰り返しによるコンパクト化が推論性能を保持し、推論深度を物理的メモリ制限から切り離すことを示しています。

5. 意義と主張

本論文は、アテンションマッチングが、高速かつ損失の大きいトークンスペース手法と、低速かつ高品質な潜在空間手法の間のギャップを埋める、実用的かつ効率的な KV キャッシュコンパクト化のプリミティブを提供すると主張しています。

• 実用性: 文脈ごとの勾配降下を回避することで、遅延と計算コストが制約となる実世界の実装において、高忠実度なコンパクト化を可能にします。
• システム統合: このアプローチは既存の推論エンジン（例：FlashAttention、可変長 KV パッキング）と互換性があり、クリティカルパス外（分離されたコンパクト化）で実行可能です。
• 限界: 著者は、この手法がCartridges よりも大幅に高速である一方で、単純なヒューリスティックよりは遅いことを認めています。さらに、極端な圧縮率（例：LongHealth における 100 倍）では、元のキーに制限されないより広い表現空間を探索できるため、Cartridges などの勾配ベース手法が依然として優れている可能性があります。
• 将来の方向性: 論文は、クエリ生成の最適化、部分集合選択を超えたコンパクトキーの直接最適化、およびオンライン・長期的エージェント設定への推論エンジンへのコンパクト化の統合に関する将来の作業を提案しています。

要約すると、この研究は、アテンション挙動を直接マッチングすることで高速かつ高品質な文脈コンパクト化が達成可能であることを実証し、長文脈 LLM におけるメモリボトルネックに対するスケーラブルな解決策を提供しています。