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

TL;DR

The paper tackles the memory bottleneck of extended chain-of-thought in reasoning LLMs by proposing RLKV, a framework that identifies reasoning-critical KV heads and preserves their full KV cache, while compressing others. It introduces a mixed-attention scheme with per-head gating, trained via reinforcement learning with sparsity pressure and stabilizing techniques (self-distillation sampling and adaptive penalty weighting) to discover which heads are essential for reasoning. RLKV leverages Group Relative Policy Optimization to optimize gating parameters $\boldsymbol{\alpha}$, using rewards from verifiable reasoning quality and an $L1$ penalty to promote sparsity. Empirical results across two reasoning models and four benchmarks show 20–50% KV-cache reductions with near lossless performance, and analyses confirm that reasoning heads are functionally distinct and more crucial than retrieval heads. The work advances memory-efficient inference for reasoning LLMs and offers insights into the heterogeneous roles of attention heads in complex reasoning tasks, with practical implications for deploying reasoning models on memory-constrained hardware.

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.

Figures (12)

• Figure 1: Motivation.Left: Existing KV cache compression methods underperform on reasoning models. Token-dropping and head-reallocation methods maintain relatively stable performance on Llama-3.1-8B-Inst but drop substantially on Llama-3.1-8B-R1, due to the 8x longer generation sequences in the reasoning models (MBPP results shown). Right: Failure modes: Token-dropping methods degenerate to repetitive behavior due to dropping critical tokens, while head-reallocation methods generate unnecessary steps, suggesting reasoning process degradation. See Appendix \ref{['appendix:motivation']} for complete results.
• Figure 2: Overview of RLKV: Our method proposes to utilize RL to identify reasoning heads. The RL pipeline naturally captures reasoning behaviors, since it samples the current model's generations to produce reward signals. The reward function evaluates the samples to assess reasoning quality. We employ $L \times H$ learnable gating adapters to mix full attention and local attention for each head, quantifying each head's reliance on full versus local KV cache access. We apply an L1 penalty to encourage adapter sparsity, while RL optimizes the adapters to preserve reasoning behaviors. After training, we identify reasoning heads with high adapter values and allocate full KV cache to them while applying compressed KV cache to others for efficient inference.
• Figure 3: Gating adapter distribution after RLKV training on two models, which both are GQA architecture.
• Figure 4: The conflict of sparse reward versus dense penalty leads to training collapse without our stabilization techniques. As adapters become sparse (decreasing average), model performance degrades (dropping reward), creating a vicious cycle where dense L1 penalties dominate increasingly sparse rewards.
• Figure 5: Performance comparison of RLKV against KV cache compression baselines across reasoning benchmarks. We evaluate RLKV (Ours) and existing methods on two reasoning models (Llama-3.1-8B-R1 and Qwen-2.5-7B-R1) across four benchmarks (GSM8K, MATH, AIME24, MBPP) at sparsity levels of 0.2, 0.4, 0.6, and 0.8. RLKV consistently outperforms all baselines across different sparsity levels, demonstrating particularly strong advantages at high sparsity levels (0.4 or 0.6) where competing methods suffer significant performance degradation. Complete numerical results are provided in Appendix \ref{['appendix:full_results']}.