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

arXiv: https://arxiv.org/abs/2510.08525 · HTML: https://mlsanity.com/paper/2510.08525/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 55 of its 71 references that are in the map.
- First posted on arXiv 9 Oct 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Efficient LLM inference and serving / [Efficient reasoning and agent inference](https://mlsanity.com/inference/efficient-reasoning-and-agent-inference/)

## Builds on (most cited first, 30 of 55)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [Learning Transferable Architectures for Scalable Image Recognition](https://mlsanity.com/paper/1707.07012/) (2017)
- [Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement](https://mlsanity.com/paper/2409.12122/) (2024)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://mlsanity.com/paper/2402.02750/) (2024)
- [KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization](https://mlsanity.com/paper/2401.18079/) (2024)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [AMC: AutoML for Model Compression and Acceleration on Mobile Devices](https://mlsanity.com/paper/1802.03494/) (2018)
- [DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads](https://mlsanity.com/paper/2410.10819/) (2024)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)

## Built on by (most cited first, 10 of 10)

- [Fast KVzip: Efficient and Accurate LLM Inference with Gated KV Eviction](https://mlsanity.com/paper/2601.17668/) (2026)
- [Active Perception Agent for Omnimodal Audio-Video Understanding](https://mlsanity.com/paper/2512.23646/) (2025)
- [EarlyTom: Early Token Compression Completes Fast Video Understanding](https://mlsanity.com/paper/2605.30010/) (2026)
- [How Much Cache Does Reasoning Need? Depth-Cache Tradeoffs in KV-Compressed Transformers](https://mlsanity.com/paper/2604.17935/) (2026)
- [CompilerKV: Risk-Adaptive KV Compression via Offline Experience Compilation](https://mlsanity.com/paper/2602.08686/) (2026)
- [dMoE: dLLMs with Learnable Block Experts](https://mlsanity.com/paper/2605.30876/) (2026)
- [HARD-KV: Head-Adaptive Regularization for Decoding-time KV Compression](https://mlsanity.com/paper/2606.28831/) (2026)
- [Dynamic Thinking-Token Selection for Efficient Reasoning in Large Reasoning Models](https://mlsanity.com/paper/2601.18383/) (2026)
- [EchoPress: Query-Agnostic KV Cache Pruning via Virtual Context Reconstruction](https://mlsanity.com/paper/2610.00412/) (2026)
- [Spend Bytes on Breadth: Precision-Count Trade-offs for Decode-Time KV Compression in Long Chain-of-Thought Reasoning](https://mlsanity.com/paper/2610.05685/) (2026)

## Abstract (arXiv)

Reasoning large language models exhibit complex reasoning behaviors via extended chain-of-thought generation that are highly fragile to information loss during decoding, creating critical challenges for KV cache compression. Existing token-dropping methods directly disrupt reasoning chains by removing intermediate steps, while head-reallocation methods, designed for retrieval tasks, fail to preserve the heads essential for generative reasoning. However, no existing method can identify which attention heads genuinely maintain reasoning consistency and control generation termination. To address this, we propose RLKV, which uses reinforcement learning as a probe to discover which heads contribute to reasoning quality by directly optimizing their cache usage against actual generation outcomes. This discovery naturally leads to an efficient compression strategy: we allocate full KV cache to reasoning-critical heads while aggressively compressing others with constant-size KV cache. Experiments reveal that a fraction of heads proves essential for reasoning, enabling 20--60% cache reduction with near-lossless performance across diverse tasks and models, and up to 2.06x end-to-end speedup at 60% reduction.