Which Heads Matter for Reasoning? RL-Guided KV Cache Compression
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.
68ecd411b0aa732b88f31f43
Wenjie Du
claimed_verified
2025-10-13T11:05:02.293Z
6624ce8248e016b5ea4ba952
https://cdn-avatars.huggingface.co/v1/production/uploads/6624ce8248e016b5ea4ba952/WbpwxXABS5gomQJtSiiCo.jpeg
wenjiedu
user
Kurt232
68ecd411b0aa732b88f31f44
Li Jiang
68ecd411b0aa732b88f31f45
Keda Tao
68ecd411b0aa732b88f31f46
Xue Liu
68ecd411b0aa732b88f31f47
Huan Wang
Which Heads Matter for Reasoning? RL-Guided KV Cache Compression
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.
A reinforcement learning framework identifies and prioritizes critical attention heads for efficient KV cache compression in large language models, maintaining reasoning quality with reduced overhead.
chain-of-thought generation KV cache token-dropping head-reallocating attention heads functional heterogeneity RLKV reinforcement learning cache usage reasoning quality KV compression