Which Heads Matter for Reasoning? RL-Guided KV Cache Compression
Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang
2025-10-13
Summary
This paper focuses on making large language models, specifically those good at complex reasoning, more efficient in terms of memory usage during operation.
What's the problem?
Large language models that can 'think through' problems step-by-step require a lot of memory to store information as they generate answers. Existing methods to reduce this memory footprint either hurt the model's reasoning ability by throwing away important information, or they compress the wrong parts of the model, leading to performance drops. The issue is that not all parts of the model are equally important for reasoning; some parts are crucial while others aren't, and current compression techniques don't distinguish between them.
What's the solution?
The researchers developed a new method called RLKV that uses a 'trial and error' learning process, similar to how you train a pet with rewards. This process figures out which parts of the model (specifically, 'heads' within the model) are most important for good reasoning. They then dedicate more memory to those crucial parts and compress the less important parts, resulting in a significant reduction in memory usage without sacrificing the model's reasoning skills.
Why it matters?
This work is important because it allows us to run powerful reasoning models on hardware with limited memory. By efficiently compressing these models, we can make them more accessible and practical for a wider range of applications, potentially enabling more complex AI systems to be deployed on devices with less processing power.
Abstract
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.