技術要約：アテンション行列を用いないエピファニー（直感）を考慮したKVキャッシュ・エビクション

問題提起

推論モデル（例：DeepSeek-R1）が数万トークンに及ぶ思考の連鎖（Chain of Thought）を生成する際、Key-Value（KV）キャッシュは線形に増大し、デプロイメントにおける主要なメモリボトルネックとなる。現在のKVキャッシュ・エビクション戦略は、通常、アテンション重みに基づいてトークンをランク付けする。このアプローチには、以下の2つの決定的な限界がある：

1. ノイズの多い重要度のプロキシ（代理指標）： 長い推論のトレースにおいて、アテンション重みはトークンの重要性を示す信頼できる指標ではない。「アテンション・シンク（Attention Sinks）」は内容に関わらず重みを吸収し、フィラー（埋め草）トークンは生成中に重気を引くものの、後に参照されることは稀である。
2. アーキテクチャ上の不適合性： アテンション重みを読み取るには、$n \times n$ のアテンション行列を実体化させる必要がある。これにより「イガー（eager）」カーネルの使用を強制され、FlashAttentionのような融合（fused）カーネルが使用できなくなる。その結果、メモリ使用量は二次関数的（$O(n^2)$）にスケールし、典型的な推論トレースよりもはるかに短いコンテキスト長（例：8192トークン）でメモリ不足（OOM）エラーを引き起こす。

これらを軽減しようとする既存の手法（ThinKVやRaaSなど）は依然としてアテンション由来の信号に依存しており、他の手法（LagKVなど）は汎用的ではあるが、推論に特化した最適化を欠いている。

手法：EPIKV

著者らは、アテンション重みの代わりにモデルの内部表現の変化に基づいてトークンをスコアリングする、トレーニングフリーのエビクション・ポリシーであるEPIKV（Epiphany-Aware KV Cache）を提案する。

1. エピファニー・スコア（Epiphany Score）

アテンションの代わりに、EPIKVは連続するデコードステップ間における残差ストリーム（隠れ状態）の変化を測定する。レイヤー $l$ におけるトークン $t$ に対し、信号はL2ノルムの差として定義される：
$$g_l(t) = \|h_l(t) - h_l(t-1)\|_2$$
大きな変化は「エピファニー・トークン」、すなわち、モデルが探索的な推論から収束的な推論へと移行したり、ステップを完了したり、あるいは洞察を確定させたりする転換点を示している。

2. 2バンド・レイヤー構造（Two-Band Layer Anatomy）

反事実的遮蔽実験（counterfactual occlusion experiments）を通じて、著者らは32レイヤーの推論モデルにおける特定のレイヤー構造を特定した：

• バンドA（レイヤー7–13）： ここでの隠れ状態の変化は、トークンの重要性と正の相関を持つ。これらのレイヤーは、既知の事実検索や特徴ルーティングが行われる部位と一致する。
• バンドB（レイヤー18–25）： ここでの隠れ状態の変化は、重要性と負の相関を持つ。これらの上層中間レイヤーは、流暢で驚きの少ない継続においても活性化するため、ここでの大きな変化は予測可能で、捨ててもよいトークンであることを示唆する。

最終的なスコアは、これらのバンドを組み合わせたものである：
$$s(t) = \bar{g}_{10}(t) - \bar{g}_{21}(t)$$
ここで、$\bar{g}$ は因果的ウィンドウにおける移動平均を表す。

3. 時間的トレンドの補正

生の信号は、単調な位置的トレンド（内容ではなく位置を追跡してしまう現象）を示す。これを補正するために、著者らは因果的ローリングzスコアを適用する：
$$z_l(t) = \frac{g_l(t) - \mu_l(t)}{\sigma_l(t) + \epsilon}$$
最終的なエビクション・スコアは $z_{10}(t) - z_{21}(t)$ である。この正規化により位置バイアスが除去され、スコアが絶対的な大きさではなく、局所的なコンテンツの偏差を反映するようにする。

4. FlashAttention との互換性

極めて重要な点として、EPIKVは以下のもののみを使用してスコアを計算する：

1. キャッシュされたキーおよび値（すでに高帯域メモリ内にある）。
2. 標準的なフォワードパスによって公開される、レイヤーごとの隠れ状態。
   EPIKVは output_attentions=True を要求することも、アテンション行列を実体化させることも決してない。これにより、EPIKVは標準的な FlashAttention インファレンス・スタック内で動作し、イガー・カーネルへのフォールバックを強制することなく実行できる。

主な貢献

1. 2バンド構造の特定： 中間レイヤー（7–13）における隠れ状態の変化は重要性と正の相関を持ち、上層中間レイヤー（18–25）における変化は負の相関を持つことを発見した。これはアテンションベースの信号よりも優れた性能を示す。
2. 位置非依存のスコアリング： 因果的ローリングzスコアを用いることで、生の隠れ状態の分散に内在する位置的トレンドを除去し、高品質なエビクション性能を回復できることを実証した。
3. アテンションなしでの性能パリティ： EPIKVは、アテンション行列を一度も実体化させることなく、展開可能なキャッシュ予算において、最強のアテンションベースのベースライン（ThinKV, H2O）と同等またはそれを上回る性能をMATH-500およびAIME-2024で達成した。
4. エンジニアリング効率： 同等の予算において、EPIKVはアテンションベースのエビクションよりも最大 2.8倍高速 に動作し、アテンションベースのスコアリングが制限するメモリの壁を回避する。

結果

実験は、DeepSeek-R1-Distill-LLaMA-8B を用い、MATH-500 および AIME-2024 データセットに対して行われた。

• 精度：
  • MATH-500において、4096トークンのキャッシュを用いた場合、EPIKVは 72% の精度を達成し、ThinKV (71%) と同等であり、H2O (67%) を上回った。
  • AIME-2024において、8192トークンのキャッシュを用いた場合、ラグ正規化されたEPIKVのバリアントは 37% に達し、最強のアテンションベースのベースライン (33%) を上回った。
• 速度：
  • AIME-2024 (8192トークン) において、EPIKVは問題を 440秒 で処理した。これに対し、最速のアテンションベースのベースライン（ThinKV）は 721秒、RaaSは 1239秒 であった。これは 1.6倍から2.8倍の高速化 を意味する。
• メモリのスケーラビリティ：
  • アテンションベースのスコアリングは、$O(n^2)$ のアテンション行列により、約8192トークンで80 GBのA100 GPUを使い果たすが、EPIKV（FlashAttention経由）は、わずか48 GBのメモリ使用量で 65,536トークン までスケール可能である（実現可能なコンテキスト長が16倍増加）。
• ベースラインの失敗モード： H2Oのようなアテンションベースの手法は、推論のトレースにおいて崩壊し、タイトな予算（例：1024トークン）において、緩やかに劣化するのではなく、空の出力や構造化されていないテキストを生成する。

意義と主張

本論文は、EPIKVが 標準的なFlashAttentionサービング・スタックでのデプロイを可能にする と主張している。アテンション行列から重要度スコアリングを切り離すことで、以下のことが実現される：

• アテンション重みを実体化させるアーキテクチャ上のコストを排除する。
• アテンション重みがノイズとなる推論トレースにおいて、より堅牢な信号を提供する。
• 既存のハードウェア上で、より長いコンテキストウィンドウと高いバッチサイズを可能にする。

著者らは、この2バンド構造は競技数学に特有のものであり、タスクの難易度（例：GSM8Kではより早いレイヤーへ移動する）によって変化することを明記しており、これは信号が固定されたレイヤーインデックスではなく、推論の定着（consolidation）を追跡していることを示唆している。また、検証用リソースとして、反事実的重要度ラベルを公開している。