💬 NLP

ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution

ForesightKVは、Golden Evictionアルゴリズムを教師あり学習および強化学習（GRPO）と組み合わせることで、最も重要なKVペアを予測・保持し、長い推論タスクにおける高い性能を維持しながらメモリコストを大幅に削減する、推論モデルのKVキャッシュ蒸発を最適化するための学習ベースのフレームワークです。

原著者： Zican Dong, Peiyu Liu, Junyi Li, Zhipeng Chen, Han Peng, Shuo Wang, Wayne Xin Zhao

公開日 2026-02-04

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

CC BY 4.0

原著者： Zican Dong, Peiyu Liu, Junyi Li, Zhipeng Chen, Han Peng, Shuo Wang, Wayne Xin Zhao

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

あなたは、非常に長く複雑なミステリーを解決しようとしている、天才的な探偵（AI）であると想像してください。これを行うために、あなたはデスクの上に膨大な手がかりのノート（KVキャッシュ）を置いておく必要があります。物語が進むにつれて、あなたのノートはどんどん厚くなっていきます。やがて、デスクは書類で散らかりすぎて、手が動かせなくなり、次の推論を行うために特定の手がかりを見つけ出すのに膨大な時間がかかるようになります。これが、大規模言語モデルにおける「メモリ・オーバーロード」の問題です。

従来、デスクを片付けるために、人々は単純なルールを使用してきました。「最も古い書類を捨てる」あるいは「インクが最も薄い書類を捨てる」といったルールです。しかし、論文ForesightKVは、これらのルールが愚かすぎると主張しています。時には、古い書類やインクが薄い書類が、物語の後半で謎を解くための最も重要な手がかりになることがあるからです。それを捨ててしまうと、探偵はミスを犯し、調査の残りの部分を台無しにしてしまいます。

以下に、ForesightKVがこの問題をどのように解決するかを、3つの簡単なステップで説明します。

1. 問題点：「ワンパス（一度きり）」の過ち

旅行の荷造りをしている場面を想像してください。標準的な方法では、「最後に詰めた10個のアイテムを残し、それ以外は捨てる」と言うかもしれません。しかし、もし最初に詰めたアイテム（地図）が、旅行全体において不可欠なものだったとしたらどうでしょう？たとえそれがスーツケースのずっと奥にあるとしてもです。

既存の手法は、単純なパターン（「最新のものを保持する」や「最も人気のあるものを保持する」など）に基づいて、どの手がかりを残すべきかを推測しようとします。論文では、複雑な推論（数学の問題など）において、手がかりには3種類の性格があることが分かりました。

グローバル・スター（世界的スター）： 何があっても常に重要。
ローカル・ネイバー（近隣住民）： 短い間だけ重要。
セマンティック・カメレオン（意味の変幻自在なカメレオン）： これが厄しい存在です。ある手がかりは、今は無意味に見えても、50ステップ後にはパズルの鍵となることがあります。単純なルールでは、これらの「カメレオン」を見逃してしまいます。

2. 解決策：「タイムトラベル」するコーチ

著者たちは、ForesightKVと呼ばれる新しいシステムを開発しました。硬直したルールブックを使う代わりに、彼らは未来を見通すことができるコーチとして機能する、小さくて賢いアシスタント（スコアリング・モデル）を訓練しました。

このコーチは、今現在の手がかりを見るだけではありません。将来、どの手がかりが最も重要になるかを予測することを学習します。これは、2段階のトレーニングプロセスを通じて行われます。

ステージ1：「黄金の基準」（教師あり学習）

まず、研究者たちは、何も捨てずにAIに数学の問題を解かせました。そして、AIの注意（アテンション）を観察し、「もし今、いくつかの手がかりを捨てなければならないとしたら、どれを捨てれば最終的な答えへのダメージが最小限で済むか？」と問いかけました。
彼らは、保持すべき完璧な手がかりのセットを見つけるために、**ゴールデン・エビクション（Golden Eviction）**と呼ばれる手法を用いました。そして、この「完璧な」意思決定プロセスを模倣するように、小さなアシスタントを教え込みました。それは、学生に解答集を見せて、「正しい答えの選び方を学びなさい」と言うようなものです。

ステージ2：「実世界シミュレーション」（強化学習）

しかし、AIが実際に問題を解いている最中、AIの思考プロセスは変化するため、「完璧な」解答集が常に完璧であるとは限りません。
そこで、第2ステージはビデオゲームのシミュレーションのようになります。アシスタントにはこう命じられます。「さあ、いくつかの手がかりを捨ててごらん。もしAIが特定の種類の単語（例えば、間違いやすい数字や記号など）でミスをした場合、君にはペナルティを与える。もしAIが正しく解き続けた場合は、報酬を与える。」
これにより、アシスタントはさらに賢くなります。ある種の「退屈な」単語（低エントロピーのトークン）を保持しておくことが、後にAIが数値を幻覚（ハルシネーション）として出力してしまうのを防ぐために極めて重要であることを理解していくのです。

3. 結果：よりスマートで、より軽いデスク

論文では、これらを3つの異なるAIモデルを用いて、難解な数学ベンチマーク（AIME 2024および2025）でテストしました。

主張： ForesightKVは、フルサイズの重いノートを使用した場合と同じ精度でこれらの数学問題を解くことができますが、使用するメモリ空間は半分（50%）に抑えられます。
比喩： これは、バックパックが50%軽くなっているにもかかわらず、ゲームに勝つために必要なあらゆる手がかりを依然として記憶している状態のようなものです。
スピード： バックパックが軽いため、AIはより速く思考でき、より多くのユーザーを同時に処理できます（スループットの向上）。

まとめ

ForesightKVは、AIのメモリを管理するための新しい方法です。古いメモを盲目的に捨てるのではなく、長期的な解決策にとってどれが重要になるかを予測することを学習した、賢い訓練済みのアシスタントを使用します。「完璧な例」による学習フェードと、「実践による学習」というゲームフェーズを組み合わせることで、複雑な推論パズルを解くために必要な重要な詳細を忘れることなく、AIを高速かつ効率的に保ちます。

技術要約: ForesightKV

問題提起

大規模言語モデル（LLM）、特に推論モデルは、長い推論プロセスを生成する上で卓越した能力を示しています。しかし、シーケンス長が増加するにつれて、Key-Value（KV）キャッシュが線形に拡大し、メモリ消費量の増大と、メモリ帯域に制約されるデコーディングによるレイテンシの増加を招きます。既存のKVキャッシュ・エビクション（破棄）手法は、ヒューリスティックなルール（アテンションスコア、位置、特徴量など）に基づいて重要度の低いKVペアを破棄することでこれを軽減しようと試みていますが、異なるアテンションヘッド間の複雑で動的な依存関係を捉えることには失敗することが多々あります。具体的には、これらの学習を行わない（training-free）手法は、コンテキストの進化に応じてKVペアの重要性が動的に変化する「意味依存的（semantic-dependent）」なパターンへの対応に苦慮しています。さらに、重要なKVペアを破棄してしまうと、低エントロピー・トークン（決定論的な部分）において損失（loss）が急増し、それが後続の推論ステップを歪ませる事実誤認を引き起こす原因となります。

手法

著者らは、KVペアの長期的な貢献度を学習し、生成中に重要度の低いものを適応的に破棄するように設計された、学習ベースのフレームワークであるForesightKVを提案します。このフレームワークは固定のメモリ予算内で動作し、各KVペアの重要度を予測する軽量なスコアリングモデル（ $\pi_\theta$ ）を最適化するための、2段階の学習パラダイムを採用しています。極めて重要な点として、学習中もLLMのパラメータは凍結されたままです。更新されるのはスコアリングモデルのみです。

1. 動的なエビクション・プロセス

エビクションは定期的に（ $L$ トークンごと）トリガーされます。キャッシュサイズが $B + L$ （ $B$ は目標予算）に達したとき、システムは直近の $L$ 個のトークンを保持し、スコアリングモデルを使用して残りのキャッシュを $B - L$ まで削減します。スコアリングモデルは、キー（ $k_n$ ）、値（ $v_n$ ）、およびアテンション特徴量（ $a_n$ ）の結合を入力として受け取り、重要度スコア（ $\phi_n$ ）を予測します。エビクションの決定は、確率性を導入するために、Top-K選択の後に多項分布サンプリング（multinomial sampling）を通じて行われます。

2. ステージ1：「ゴールデン・エビクション」を用いた教師あり学習

スコアリングモデルを訓練するために、著者らは教師ラベルを構築するGolden Evictionアルゴリズムを導入しています。

メカニズム: 完全な推論トレースに対して、アテンション行列をブロックに分割します。アルゴリズムは、ブロック内およびアテンションヘッドを横断したアテンションスコアを集計することで、「ブロックスコア」を計算します。
将来への意識（Future Awareness）: 各エビクションステップにおいて、アルゴリズムは各KVペアについて、すべての将来のステップにおける最大ブロックスコアを特定します。この「将来スコア」は、将来のアテンション計算への影響を最小限に抑えるための最適な保持戦略を表しています。
学習: スコアリングモデルは、任意の2つのKVペアに対する予測スコアが、Golden Evictionアルゴリズムから導出された将来スコアのランキングと一致するように、**ペアワイズ・ランキング損失（Pairwise Ranking Loss）**を用いて訓練されます。

3. ステージ2：強化学習（RL）

教師あり学習だけでは、推論中のエビクションによって生じる分布シフトを考慮できない可能性があります。このギャップを埋めるため、著者らはKVキャッシュのエビクションを**マルコフ決定過程（MDP）**として定式化し、GRPO（Group Relative Policy Optimization）アルアルゴリズムを適用しています。

状態（State）: 現在残っているKVキャッシュ。
行動（Action）: 保持するKVペアのサブセットを選択すること。
報酬（Reward）: 報酬関数は、低エントロピー・トークンにおける性能低下を抑制するように設計されています。具体的には、元々は低エントロピーであったが、エビクション後に損失が大幅に増加したトークンに焦点を当てます。報酬は、この特定のサブセットにおける損失増加の平均二乗誤差（MSE）の負の値として定義されます。
最適化: 方策モデルは、シーケンス全体にわたって累積報酬を最大化するように最適化され、決定論的な推論ステップに必要なクリティカルなコンテキストを効果的に保持することを学習します。

主な貢献

ForesightKV フレームワーク: 静的なヒューリスティック・ルールを超え、KVペアの長期的な貢献度を予測することを学習する、新しい2段階学習法（教師あり学習 + RL）を用いたフレームフロー。
Golden Eviction アルゴリズム: 将来のアテンション情報を用いて最適なエビクション・トレースを生成し、スコアリングモデルのための高品質な教師信号となる手法。
エビクションのための MDP 定式化: エビクションプロセスをMDPとしてモデル化し、GRPOを利用して、推論の正確性に不可欠な低エントロピー・トークンの損失スパイクの緩和を特に対象とする手法。
効率性: ベースとなるLLMのバックプロパゲーションやパラメータ更新を必要とせず、軽量なスコアリングモデルのみを訓練するフレームワーク。

実験結果

著者らは、AIME2024およびAIME2025の数学ベンチマークを用い、3つの推論モデル（Qwen3-1.7B, Qwen3-4B, DeepSeek-R1-Distill-Qwen-7B）でForesightKVを評価しました。

パフォーマンス: ForesightKVは、既存のベースライン（SnapKV, H2O, R-KV）を一貫して上回りました。特筆すべきは、半分のKVキャッシュ予算を使用しながら、フルキャッシュのベースラインと同等またはそれ以上の性能を達成したことです。例えば、Qwen3-4Bを用いたAIME2024において、ForesightKV（予算1K）は、R-KV（予算2K）よりも高い性能を示しました（54.5 vs 44.8 Pass@1）。
スループット: この手法はスループットを大幅に向上させました。32Kシーケンスにおいて1Kの予算を使用した場合、ForesightKVはフルKVキャッシュと比較して最大9.79倍のスループット向上を達成しました。
汎用性: 主に数学タスクで訓練されているにもかかわらず、ForesightKVはGPQA（科学）およびLiveCodeBench（コーディング）のベンチマークにおいても強い汎用性を示し、半分のキャッシュ予算を持つベースラインを上回りました。
アブレーション研究: 実験結果は、教師あり学習ステージ（初期の予測能力を提供）と強化学習ステージ（低エントロピー・トークンの保存のための方策を洗練）の両方が成功に不可欠であることを裏付けました。低エントロピー・トークンの損失増大に焦点を当てた特定の報酬設計が、最も効果的であることが判明しました。

意義と主張

本論文は、ForesightKVが、長文コンテキストの推論における適応的なKVキャッシュ管理のための有望な新しい方向性を提供すると主張しています。ルールベースのヒューリスティックから、KVペアの長期的な貢献を明示的にモデル化し、低エントロピー・トークンの保存をターゲットとする学習ベースのアプローチへと移行することで、メモリ効率と生成品質の優れたトレードオフを実現しています。著者らは、彼らのアプローチが特定の予算制約に過学習するのではなく、KVキャッシュのエントリの固有の重要性を捉えていることを強調しており、これにより異なる予算設定やタスク領域への優れた汎用性を実現しています。