キーバリューキャッシュからの効率的なエビクション学習
大規模言語モデルの推論効率化に向け、従来のヒューリスティック手法に依存しない、キーバリューキャッシュの効率的なエビクション手法の研究。
キーポイント
大規模言語モデル（LLM）推論時のメモリ負荷を軽減するため、強化学習を用いてKVキャッシュの削除（エビクション）を最適化する手法「KV Policy（KVP）」を提案
従来のヒューリスティック手法（直近性や過去の注意スコア）を超え、トークンの将来の有用性を予測してランキングする新たなパラダイムを確立
長文コンテキストベンチマーク（RULER）や対話ベンチマーク（OASST2-4k）で既存手法を大幅に上回る性能を実証し、トレーニング分布外のタスクや長いコンテキスト長にもゼロショットで汎化
基盤LLMや追加推論コストを変更せずに軽量なRLエージェントを導入可能で、実用的なスケーラビリティを備えている
影響分析・編集コメントを表示
影響分析
この研究は、LLM推論のボトルネックであるKVキャッシュ管理に強化学習を適用した画期的なアプローチであり、長文処理や対話システムの実用化を加速させる可能性がある。Appleの研究チームによる実証結果は、業界標準ベンチマークで優位性を示しており、今後の効率化技術の方向性に影響を与える重要な進展と言える。
編集コメント
LLMの実用化における最大の課題である推論コスト削減に、強化学習という新たな視点からアプローチした点が革新的。業界全体の効率化トレンドを後押しする重要な研究。
キー・バリューキャッシュからの学習によるエビクション
著者: Luca Moschella, Laura Manduchi, Ozan Sener
出版物を表示
大規模言語モデル（LLM）のサイズが増大するにつれ、オートリグレッシブなキー・バリュー（KV）キャッシュのメモリ要求が主な原因となり、効率的な推論は困難な課題となっている。既存のエビクション（追い出し）や圧縮手法はコストを削減するが、新近性や過去のアテンションスコアなどのヒューリスティックに依存しており、これらはトークンの将来的な有用性に対する間接的な代理指標でしかなく、計算上のオーバーヘッドも伴う。我々はKVキャッシュのエビクション問題を、将来のデコーディングに対する予測有用性に基づいてトークンをランク付けすることを学習する、強化学習（RL）問題として再定義する。この目的のために、我々はKVポリシー（KVP）を導入する。これは、キーとバリューベクトルのみを用いて事前計算された生成トレース上で訓練される、軽量なヘッドごとのRLエージェントのフレームワークである。各エージェントは、将来の有用性によって導かれる専門的なエビクションポリシーを学習する。この有用性は、すべてのキャッシュ予算にわたるランキングの品質を評価するものであり、基盤となるLLMの変更や追加の推論を必要としない。長文脈ベンチマークRULERとマルチターン対話ベンチマークOASST2-4kにおいて、2つの異なるモデルファミリーで評価した結果、KVPはベースライン手法を大幅に上回った。さらに、標準的な下流タスク（例：LongBench、BOOLQ、ARC）でのゼロショットテストは、KVPがその訓練分布を超えて、またより長い文脈長に対しても、よく汎化することを示している。これらの結果は、将来のトークンの有用性を予測することを学習することが、適応的なKVキャッシュ管理のための強力でスケーラブルなパラダイムであることを実証している。
関連する文献と最新情報。
EpiCache: 長い対話型質問応答のためのエピソード的KVキャッシュ管理
2025年9月23日研究分野 手法とアルゴリズム、研究分野 音声と自然言語処理
大規模言語モデル（LLM）の最近の進歩は文脈長を拡張し、アシスタントが長い履歴を維持して首尾一貫したパーソナライズされた応答を可能にしている。しかし、この能力はキー・バリュー（KV）キャッシングに依存しており、そのメモリ使用量は対話の長さに比例して線形に増加し、厳しいリソース制約下ではすぐに支配的になる。このオーバーヘッドを削減するための活発な研究分野の一つがKVキャッシュ圧縮であり、キャッシュサイズを制限しつつ…
KV-Runahead: 並列キー・バリューキャッシュ生成によるスケーラブルな因果的LLM推論
2024年5月14日研究分野 手法とアルゴリズム、研究分野 音声と自然言語処理会議 ICML
大規模言語モデル（LLM）の推論には、最初のトークンを出力するプロンプト（またはプリフィル）フェーズと、後続のトークンを生成するエクステンション（またはデコーディング）フェーズの2つのフェーズがある。本研究では、プロンプトフェーズを加速する効率的な並列化スキーム、KV-Runaheadを提案する。重要な観察は、キー・バリューキャッシュ（KVキャッシュ）のために、エクステンションフェーズがプロンプトフェーズよりも高速にトークンを生成するということである。したがって、KV-Runaheadは…
機械学習における機会を発見する。
私たちの機械学習研究は、日々新たな領域を切り拓いています。
原文を表示
Learning to Evict from Key-Value Cache
AuthorsLuca Moschella, Laura Manduchi, Ozan Sener
View publication
The growing size of Large Language Models (LLMs) makes efficient inference challenging, primarily due to the memory demands of the autoregressive Key-Value (KV) cache. Existing eviction or compression methods reduce cost but rely on heuristics, such as recency or past attention scores, which serve only as indirect proxies for a token’s future utility and introduce computational overhead. We reframe KV cache eviction as a reinforcement learning (RL) problem: learning to rank tokens by their predicted usefulness for future decoding. To this end, we introduce KV Policy (KVP), a framework of lightweight per-head RL agents trained on pre-computed generation traces using only key and value vectors. Each agent learns a specialized eviction policy guided by future utility, which evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two different model families on the long-context benchmark RULER and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms baselines. Furthermore, zero-shot tests on standard downstream tasks (e.g., LongBench, BOOLQ, ARC) indicate that KVP generalizes well beyond its training distribution and to longer context lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.
Related readings and updates.
EpiCache: Episodic KV Cache Management for Long Conversational Question Answering
September 23, 2025research area Methods and Algorithms, research area Speech and Natural Language Processing
Recent advances in large language models (LLMs) have extended context lengths, enabling assistants to sustain long histories for coherent, personalized responses. This ability, however, hinges on Key-Value (KV) caching, whose memory grows linearly with dialogue length and quickly dominates under strict resource constraints. An active line of research for reducing this overhead is KV cache compression, which seeks to limit cache size while…
KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
May 14, 2024research area Methods and Algorithms, research area Speech and Natural Language Processingconference ICML
Large Language Model or LLM inference has two phases, the prompt (or prefill) phase to output the first token and the extension (or decoding) phase to the generate subsequent tokens. In this work, we propose an efficient parallelization scheme, KV-Runahead to accelerate the prompt phase. The key observation is that the extension phase generates tokens faster than the prompt phase because of key-value cache (KV-cache). Hence, KV-Runahead…
Discover opportunities in Machine Learning.
Our research in machine learning breaks new ground every day.
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み