どのヘッドが推論に重要か？RLガイドによるKVキャッシュ圧縮

Which Heads Matter for Reasoning? RL-Guided KV Cache Compression

October 9, 2025

著者: Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang

cs.AI

要旨

大規模言語モデルは、拡張された連鎖思考生成を通じて複雑な推論行動を示し、デコードフェーズにおいて前例のないKey-Value（KV）キャッシュのオーバーヘッドを生み出します。既存のKVキャッシュ圧縮手法は、推論モデルにおいて十分な性能を発揮しません。トークン削除手法は重要な情報を破棄することで推論の整合性を損ない、ヘッド再割り当て手法は検索タスク向けに設計されているため、推論に重要なヘッドを誤って圧縮し、圧縮率が上がるにつれて性能が大幅に低下します。我々は、推論モデルにおいてKVヘッドが機能的な異質性を示すと仮説を立てました。つまり、一部のヘッドは連鎖思考の一貫性に不可欠であり、他のヘッドは圧縮可能であると考えます。この洞察を検証し活用するために、我々はRLKVという新しい推論重要ヘッド識別フレームワークを提案します。RLKVは、各ヘッドのキャッシュ使用量と推論品質の関係を直接最適化するために強化学習を利用します。RLKVはトレーニング中に実際に生成されたサンプルから報酬を生成するため、推論行動に関連するヘッドを自然に特定します。その後、これらのヘッドには完全なKVキャッシュを割り当て、他のヘッドには圧縮された定数KVキャッシュを適用して効率的な推論を実現します。実験の結果、推論に不可欠なアテンションヘッドはごく一部であることが明らかになり、我々のKV圧縮手法はベースライン手法を上回りながら、20～50%のキャッシュ削減を達成し、非圧縮結果とほぼ同等の性能を維持することができました。

English

Reasoning large language models exhibit complex reasoning behaviors through the extended chain-of-thought generation, creating unprecedented Key-Value (KV) cache overhead during the decoding phase. Existing KV cache compression methods underperform on reasoning models: token-dropping methods break reasoning integrity by discarding critical information, while head-reallocating methods mistakenly compress reasoning-critical heads since they are designed for retrieval tasks, resulting in significant performance degradation as compression rates increase. We hypothesize that KV heads exhibit functional heterogeneity in reasoning models-some heads are critical for chain-of-thought consistency while others are compressible. To validate and exploit this insight, we propose RLKV, a novel reasoning-critical head identification framework, which uses reinforcement learning to directly optimize the relationship between each head's cache usage and reasoning quality. As RLKV produces rewards from actual generated samples during training, it naturally identifies heads relevant to reasoning behaviors. We then allocate full KV cache to these heads while applying compressed constant KV cache to others for efficient inference. Our experiments reveal that only a small fraction of attention heads is essential for reasoning, enabling our KV compression approach to outperform baseline methods while achieving 20-50% cache reduction with near lossless performance compared to uncompressed results.