技術的概要：KV キャッシュ圧縮性を備えたトランスフォーマーの学習

1. 問題定義

長文脈言語モデル化は、キー・バリュー（KV）キャッシュによって次第にボトルネック化しており、そのメモリ使用量とデコード時のアクセスコストはプレフィックス長に比例して増大する。最近の研究では、トークンレベルの要約から最適化ベースの KV キャッシュ圧縮（勾配ベースのマッチングやアテンションマッチングなど）に至るまで、さまざまな文脈圧縮手法が提案されているが、これらのアプローチは固定された事前学習済みモデルに対する事後的介入として機能している。

これらの事後的手法の核心的な限界は、その有効性が本質的にモデル内部表現の圧縮性によって制約されている点にある。本論文は、圧縮性は入力シーケンスやタスクの属性のみならず、トランスフォーマー内の学習された表現の属性であると主張する。具体的には、異なるトランスフォーマー実装が同一の入力シーケンスに対して同一の次トークン分布を生成しつつも、内部構造は大きく異なり、その中には本質的に圧縮に抵抗するものも存在しうる。

2. 理論的動機

著者はKV 圧縮性の概念を形式化した。すなわち、あるトランスフォーマーが KV 圧縮性を有するとは、長い入力に対する KV キャッシュを、モデルの次トークン分布を有界誤差内で保持しつつ、より短いキャッシュへマッピングする圧縮ポリシーが存在することを意味する。

主要な理論的知見:

• 非圧縮性実装の存在: 本論文は、ほぼ任意のシーケンス・ツー・ベクトル関数に対して、本質的に非圧縮性（いかなる非自明な圧縮も定数誤差を伴う）なトランスフォーマー実装が存在しうることを証明するとともに、単一の KV ペアまで圧縮可能な高圧縮性実装も存在しうることを示した。
• 動機付けとなる例（ヒストグラム計算）: 著者は、トークン頻度を計算するヒストグラム関数の「自然な」トランスフォーマー実装は、圧縮中のスケーリング誤差を修正するために必要な位置情報を保持せずにトークン埋め込みを平均化することに依存しているため、非圧縮性となりうることを実証した。対照的に、位置情報と未正規化カウントを明示的に符号化するわずかに修正された構造化実装は、単一の KV ペアへの完全な圧縮を可能にする。
• 示唆: 標準的な学習目的（次トークン予測）は、本質的に圧縮可能な表現の出現を促進するものではなく、単純で非圧縮性の解決策を好む可能性がある。

3. 手法：KV 圧縮認識学習（KV-CAT）

理論的ギャップに対処するため、著者は圧縮可能な内部表現の出現を促進するように設計された継続的事前学習（CPT）手順である**KV 圧縮認識学習（KV-CAT）**を提案する。

学習メカニズム

KV-CAT は、学習時 KV 疎化ポリシーを導入することで標準的な学習ループを修正する。

1. 二重フォワードパス: 各学習シーケンスに対して、モデルは 2 回のパスを実行する。
   • 密パス: すべての KV スロットが有効な標準的なフォワードパス。
   • マスクパス: 学習されたルーターが KV スロットの一部をマスクする（例：50% のみ保持）フォワードパス。
2. ルーター機構: 層のグループ間に軽量な学習型ルーター（線形アテンションモジュールとして実装）が挿入される。これらは各トークンに対してスカラースコアを出力し、閾値処理によって二値マスクを生成する。これらのルーターはモデルパラメータと共同で学習される。
3. 学習目的: 損失関数は 3 つの項を組み合わせる。
   • 自己蒸留損失（$L_{mask}$）: マスク済みモデルの分布と密モデルの分布との間の KL 発散を最小化する。これにより、KV スロットがマスクされてもモデルが性能を維持することを強制する。
   • アンカー損失（$L_{anchor}$）: 密フォワードパスに適用される標準的な次トークン予測（NTP）損失であり、モデルの元の挙動を維持し、破滅的忘却を防ぐ。
   • 予算損失（$L_{budget}$）: ルーターが目標保持率（例：50%）を維持することを促す正則化項であり、モデルがすべてのトークン、あるいは何のトークンもマスクしないことを防ぐ。

推論

評価時にはルーターは無効化され、モデルは標準的な密トランスフォーマーとして動作する。その後、この学習済みモデルの KV キャッシュに対して、事後圧縮手法（アテンションマッチングや勾配ベース最適化など）が適用される。

4. 実証的評価

著者は、FineWeb-Eduデータセットを用いた継続的事前学習により、QWEN2.5モデル（0.5B および 1.5B パラメータ）上で KV-CAT を評価した。彼らは、最先端の事後圧縮手法（アテンションマッチングと勾配ベース最適化）を用いて、KV-CAT チェックポイントをベースモデルと比較し、以下の 3 つの軸で評価を行った。

1. 非圧縮性能（Q1）:

   • KV-CAT モデルは、HellaSwag、WinoGrande、PIQA などの標準的な多肢選択ベンチマークにおいて、ベースモデルと同等の性能を維持し、精度の差は無視できるレベルであった。これは、学習手順が標準的な能力と圧縮性の間でトレードオフを行っていないことを確認するものである。

2. 固定予算下での圧縮品質（Q2）:

   • サフィックスパープレキシティ: 固定された最適化予算の下で、KV-CAT モデルはベースモデルよりも一貫して優れたサフィックスパープレキシティ保持を達成した。パープレキシティ保持の改善は最大で3.21 倍に達した。
   • 最適化速度: 勾配ベース圧縮を使用する場合、KV-CAT モデルはベースモデルと同じ圧縮品質を達成するために、最大で5 倍少ない最適化ステップを必要とした。

3. ダウンストリームタスク（Q3 & Q4）:

   • 検索（干し草の山の中の針）: KV-CAT は、圧縮された文脈からの完全一致検索精度を向上させ、特定の保持率（例：30-50% の保持率）において最大**68%**の改善をもたらした。
   • 長文脈 QA: LongBench v2 タスクにおいて、KV-CAT モデルは文脈が圧縮された場合に精度が大幅に向上（最大39%）し、さまざまなサブドメイン（法務、学術、エージェント履歴など）においてベースモデルを上回った。

5. 主要な貢献

本論文は以下の貢献を主張する。

1. 圧縮性の特性評価: 入力シーケンスだけでなく、学習されたトランスフォーマー表現の属性として KV 圧縮性を定義すること。
2. 理論的証明: ほぼすべてのシーケンス・ツー・ベクトル関数が、圧縮可能な実装と非圧縮可能な実装の両方を許容することを示すこと。
3. KV-CAT フレームワーク: 学習時疎化と自己蒸留を通じて明示的に圧縮可能な表現を促進する学習手順を提案すること。
4. 実証的検証: KV-CAT が、非圧縮性能を劣化させることなく、検索、パープレキシティ、長文脈 QA における事後圧縮手法の品質と予算のトレードオフを改善することを示すこと。

6. 意義と主張

本論文は、KV-CAT を事後圧縮アルゴリズムの代替としてではなく、モデルをそれらに対してより適合しやすくする補完的なアプローチとして位置づけている。その意義は、最適化対象を圧縮アルゴリズム自体からモデルの学習ダイナミクスへとシフトさせる点にある。学習中にモデルを圧縮制約に曝すことで、著者はモデルが情報損失に対して構造的に頑健な表現を学習し、それによってメモリと計算コストを削減しつつより効果的な長文脈推論を可能にすると主張する。

著者は限界について謙虚であり、この手法が追加のルーターと二重パスによる学習オーバーヘッドと実装の複雑さを導入すること、ならびに大規模モデルや異なるドメインへの一般化にはさらなる調査が必要であることを指摘している。