🤖 machine learning

CompilerKV: Risk-Adaptive KV Compression via Offline Experience Compilation

CompilerKV は、クロスプロンプトの規則性を活用してノイズの多いオンライン推定を効率的な $O(1)$ ルックアップに置き換えるオフラインコンパイルされた KV 保持ポリシーを導入し、既存のプレフィル専用圧縮手法と比較して、極端なメモリ制約下においても最先端のパフォーマンスと優れたスケーラビリティを達成する。

原著者： Ning Yang, Chengzhi Wang, Yibo Liu, Baoliang Tian, Haijun Zhang

公開日 2026-05-21

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Ning Yang, Chengzhi Wang, Yibo Liu, Baoliang Tian, Haijun Zhang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

非常に長い物語を後でその内容について質問に答えられるように記憶しようとしていると想像してください。あなたの脳（AI モデル）は、その物語の最も重要な部分を会話中に保持するために、「精神的なメモ帳」のような限られたスペース（KV キャッシュと呼ばれる）を持っています。

物語が短ければ、すべてを記憶できます。しかし、物語が 10 万語にもなると、そのメモ帳は溢れてしまいます。スペースを作るために、いくつかの部分を捨てなければなりません。問題は、一度何かを捨ててしまうと、二度と取り戻せないという点です。もし偶然にも、悪役が秘密の計画を明かす部分を捨ててしまった場合、物語全体が崩壊してしまいます。

従来の方法：その場しのぎの推測

従来の手法は、物語を「今、リアルタイムで」見て、何を保持するかを決定しようとしました。「この単語はアテンションスコアが高いから、重要に違いない！」あるいは「この文は複雑に見えるから、保持しよう！」といった具合です。

この論文の著者らは、これはノイズの多い単一のフレームを見て、映画全体を判断しようとするようなものだと主張しています。物語があまりに長く、データがごちゃごちゃしているため、単一のプロンプト（単一の物語）だけを見て判断すると、誤った推測をしてしまいます。華やかだが無意味な単語を保持し、静かだが決定的な手がかりを捨ててしまう可能性があります。

新しい方法：COMPILERKV（「試合前の戦略」）

著者らはCOMPILERKVを導入しました。その場しのぎで推測するのではなく、コーチが大きな試合前に試合映像を研究するように、事前に戦略を「コンパイル」します。

その仕組みを 3 つの簡単なステップに分解して説明します。

1. 「ノイズフィルター」（安定化ユーティリティ）

混雑した部屋で話を聞いていると想像してください。誰かが叫んでいる（一時的なスパイク）、あるいはマイク自体が他のものより大きい（スケールバイアス）といった状況です。

対策: COMPILERKV は、ただ一番大きな声で叫んでいる人を聞くだけではありません。ノイズを平滑化し、音量を正規化します。「この人は本当に重要なことを言っているのか、それとも単に大きい声を出しているだけなのか？」と問うのです。これにより、AI が一時的なノイズに騙されるのを防ぎます。

2. 「専門家マップ」（ヘッド異質性テーブル）

AI の内部には、多くの異なる「脳細胞」（アテンションヘッド）があります。あるものは事実を見つけるのが得意（図書館司書のような）、他のものは単なるおしゃべりのノイズです。

対策: 著者らはオフラインで数千の物語を研究し、特定の脳細胞は常に信頼でき、他のものは常にノイズであることに気づきました。彼らは永続的な「専門家マップ」を作成しました。
比喩: 全チームに意見を求めるのではなく、システムはこう判断します。「『図書館司書』細胞がその単語を重要だと言っているなら、たとえ『おしゃべり』細胞がそれを不要だと言っても、保持する。」これにより、信頼できる専門家に重要な情報を救うための「拒否権」を与えます。

3. 「リスクメーター」（リスク適応型閾値）

すべての物語が同じではありません。単純なもの（レシピ）もあれば、複雑で混乱を招くもの（ミステリー小説）もあります。

対策: システムは、現在の物語がどれほど「リスクが高いか」をチェックします。
- 低リスク（単純な物語）: 「攻撃的に行ける。テキストの 90% を捨てても大丈夫だ。」
- 高リスク（混乱を招く物語）: 「これは危険だ！慎重に行こう。万一のためにテキストの 95% を保持する。」
比喩: 旅行の荷造りのようなものです。天気が完璧だと分かれば、荷物は軽くします。嵐の警報が出れば、追加の装備を詰みます。COMPILERKV は、現在のプロンプトがどれほど「嵐のような状態」かによって、保持する量を自動的に調整します。

これが大きな意味を持つ理由

この論文は、その場しのぎの推測（オンラインヒューリスティック）ではなく、「試合前の研究」（オフラインコンパイル）を行うことで、はるかに優れた結果が得られると主張しています。

移植性: 彼らが作成した「専門家マップ」は、異なる AI モデル間でも機能します。異なる種類の脳に適用できる普遍的なルールブックのようなものです。
メモリ節約: 元のテキストのわずか 1.5% のメモリ予算で AI を稼働させながら、それでも質問に正しく答えることができます。
長文への対応: 大規模なコンテキスト（最大 12 万語）でテストされた際、他の手法が崩壊する中、COMPILERKV は強さを保ちました。例えば、「干し草の山の中の針」（巨大なテキストから特定の事実を見つける）テストでは、COMPILERKV は 89% の確率で針を見つけましたが、次点の手法では 42% しか見つけられませんでした。

結論

COMPILERKVは、一瞬で完璧な決定を下そうとするのをやめます。代わりに、何を保持するかを決定するために、事前に計算されたリスクを考慮した戦略を使用します。これは、荒々しい賭けをするギャンブラーと、長年の経験に基づいて最善の手をすでに計算済みのチェスのグランドマスターの違いです。その結果、メモリ不足になったり重要な部分を忘れたりすることなく、膨大な量のテキストを記憶できる AI が実現します。

技術サマリー：COMPILERKV

1. 問題定義

大規模言語モデル（LLM）は、コンテキスト長の増加に伴うキー・バリュー（KV）キャッシュの線形成長により、メモリボトルネックに直面している。プリフィル専用 KV 圧縮は、プリフィル段階の末尾でトークンの部分集合を選択し、デコードパス全体で固定する、展開に優しい解決策として登場した。このアプローチは、ランタイムメモリ使用量を固定コストに変換し、デコード時の介入をゼロにする。

しかし、本論文は既存のプリフィル専用手法（SnapKV、PyramidKV、HeadKV など）に根本的な統計的ミスマッチが存在することを特定している。これらの手法は、単一のノイズを含むプロンプトからの統計量を用いて、ヘッドごとの信頼性（どのアテンションヘッドが検索に重要か）およびプロンプトレベルの圧縮感度（特定のプロンプトがどの程度圧縮可能か）という修正信号をオンラインで推定しようとしている。著者らは、これが誤った統計単位であると主張する：

E1（単一プロンプトノイズ）: プロンプトごとのアテンションスコアは重尾分布を示し、スケールバイアスがかかる。大きな値は必ずしも意味的な重要性を保証するものではなく、小さな推定誤差が不可逆的な選択エラーを引き起こす。
E2（ヘッドの異質性）: アテンションヘッドは機能的に特化している。一部は「検索に不可欠」である一方、他は一貫してノイズを含む。単純な集約は、ノイズを含むヘッドがスコアを支配することを許容する。
E3（プロンプトレベルのリスク分散）: 異なるプロンプトは異なる圧縮限界を持つ。固定された保持ルールは、高エントロピーまたは高パープレキシティのプロンプトでは失敗し、壊滅的な情報損失を招く。

既存の手法は、これらの信号を各プロンプトに対してオンラインで再推定しており、これらの信号が単一プロンプト内の安定性ではなく、高いクロスプロンプト規則性を示すという事実を活用できていない。

2. 手法：COMPILERKV

著者らは、COMPILERKVを提案する。これは、校正コーパスから修正テーブルをオフラインでコンパイルする KV 保持ポリシーであり、オンラインで推定するものではない。核心的な洞察は、ヘッドの信頼性ランキングおよびリスク特徴から安全な閾値へのマッピングが、一度学習して再利用可能なアーキテクチャ事前分布として安定しているという点である。

このフレームワークは、プリフィル段階中に動作する 3 つの結合コンポーネントで構成される：

A. 安定化されたトークン有用性（E1 への対応）

学習なしで一時的なアテンションスパイクとヘッド間スケールバイアスを除去するため、パラメータフリーの有用性スコアを計算する：

ウィンドウ集約アテンション: プロンプト尾部の観測ウィンドウ $\Omega$ 上でアテンションを集約することで、一時的なスパイクを平滑化する。
ヘッド相対正規化: ヘッドの平均値ノルムでトークンスコアを正規化し、スケールバイアスを除去する。
数式: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$ 。ここで、 $\alpha_t$ はウィンドウ質量、 $\rho$ は正規化された値ノルムである。

B. ヘッド異質性テーブル（E2 への対応）

すべてのヘッドを均等に扱うか、オンラインで重要性を再推定するのではなく、COMPILERKV は連続的な信頼性重みの凍結テーブル $W_{head}[l, h]$ をコンパイルする。

オフラインコンパイル: 校正コーパス上で**保守的 Q ラーニング（CQL）**を通じて学習される。状態は $(l, h, B_l)$ であり、行動は離散化された信頼性重みである。
メカニズム: 推論中、トークンの有用性は、そのヘッドの信頼性によって再スケーリングされる。重要なのは、集約に重み付き最大プーリング（ $\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$ ）を使用することである。これにより、「拒否権」が検索に不可欠なヘッドに与えられる。単一の重要なヘッドがトークンに強くアテンションする場合、ノイズを含むヘッドからの合意に関わらず、そのトークンは保持される。

C. リスク適応型閾値ゲート（E3 への対応）

圧縮限界が異なるプロンプトを処理するため、リスク適応型閾値ゲートがプロンプトレベルのリスク特徴を保持閾値にマッピングする。

リスク信号: 観測ウィンドウ上で計算されるアテンションエントロピー（構造的リスク）と局所パープレキシティ（意味的リスク）を使用する。
オフラインコンパイル: 離散化されたリスクビン $(b_{ent}, b_{ppl})$ を保持閾値 $\tau^{(l)}$ にマッピングするルックアップテーブル（LUT）が、CQL 経由でコンパイルされる。
保守的バイアス: CQL の目的関数は、稀な高リスク状態（高エントロピー/高パープレキシティ）における積極的な閾値にペナルティを課し、オンラインヒューリスティックが通常失敗する領域でのより安全な保持にテーブルをバイアスする。
弾性選択: 最終的な選択は、閾値 $\tau^{(l)}$ 以上の有用性を持つすべてのトークンを保持する。候補セットが予算 $B_l$ を超える場合、トップ $B_l$ トークンにクリップされる。下回る場合は、すべての候補が保持される（弾性的な過少保持）。

D. 理論的根拠

本論文は、近似誤差が破棄された尾部質量によって制御されることを示す形式的分析を提供している（命題 4.1）。また、有限テーブルの校正誤差は、単一プロンプト内のアテンション行内の独立性を仮定するのではなく、プロンプトレベル（独立したプロンプト間）で制御されると主張しており、これにより分散を低減する。

3. 主要な貢献

問題の再定義: プリフィル専用圧縮における修正の統計単位は、単一プロンプトではなくクロスプロンプトスケールであるべきであることを特定。
オフラインコンパイル: ヘッドの信頼性とリスク適応型閾値をオフラインでコンパイルする手法を導入し、それらをポータブルなアーキテクチャ事前分布に変換。
結合ポリシー: ヘッドの信頼性重みとリスク閾値を、独立したモジュールではなく、同じ下流の報酬に対して較正する共同最適化を提案。
ポータビリティ: コンパイルされたテーブルが、分離したコーパス間、さらには異なるモデルファミリー間（例：LLaMA から Mistral へ）でも、最小限の性能低下で転移することを実証。

4. 実験結果

この手法は、LongBench、RULER、および Needle-in-a-Haystack（NIAH）ベンチマークにおいて、4 つの 7B–8B バックボーン（LLaMA-3、Mistral、Qwen2、InternLM）で評価された。

LongBench（512 トークン予算）: COMPILERKV は 4 つのバックボーンすべてで圧縮 SOTAを達成し、最強のベースライン（例：ChunkKV、HeadKV）を平均**+1.67 ポイント**（95% CI [+1.08, +2.37]）上回った。一部のバックボーンでは、完全 KV（FullKV）の性能をわずかに上回っており、著者らはこれは情報獲得ではなくノイズ除去によるものだと帰因している。
圧迫レジーム:
- RULER（128k コンテキスト）: 固定の 512/32k キャッシュ比率において、COMPILERKV は約 73 のスコアを維持し、SnapKV（約 38）や KVZip（約 66.5）を大幅に上回り、FullKV（約 79）の下に留まる。
- NIAH（32k）: 0.89 の精度を達成し、SnapKV の 0.42 を上回る。
- バッチ処理: 32k 入力、バッチサイズ 16 において、COMPILERKV は実行可能であるのに対し、FullKV はメモリ不足（OOM）に陥る。
転移性:
- クロスコーパス: ヘッドの信頼性ランキングは、分離した校正コーパス間で高いスピアマン相関（ $\bar{\rho} \approx 0.90$ ）を示す。
- クロスモデル: モデル間（例：LLaMA-3 から Mistral へ）の直接テーブル転送は、ターゲット固有のコンパイルと比較して、わずか0.4–0.8 LongBench ポイントの損失しか生じない。
効率性: オンラインオーバーヘッドは最小限である。レイヤーあたりの $O(1)$ 回のテーブルルックアップと、候補プール上の単一のソートのみで済む。32k コンテキストにおいて、FullKV と比較して最大2.45 倍のデコードスループットを達成。

5. 意義と主張

本論文は、COMPILERKV が、保持ポリシーの計算における統計的ミスマッチを修正することで、固定予算における FullKV とのギャップを埋めると主張している。修正信号の計算を、オンラインでノイズが多く単一プロンプトに基づく推定から、オフラインでコンパイルされ、クロスプロンプト的な事前分布へとシフトさせることで、この手法はノイズや異質性に対する選択プロセスを安定化させる。

著者らは、新奇性が単なる「ヘッド認識」（HeadKV などの先行手法も有する）ではなく、連続的なヘッド信頼性とプロンプトリスク閾値の結合されたオフラインコンパイルにあることを強調している。これにより、ヘッドの信頼性はチューニングのトリックから、一度コンパイルされ、 $O(1)$ のオンラインオーバーヘッドでモデルやコンテキスト全体に展開可能な再利用可能な科学的対象へと変容する。この研究は、プリフィル専用圧縮にとって「正しい」信号は、入力プロンプトの一時的な特徴ではなく、モデルのアーキテクチャ的特性であることを示唆している。