Research map / Architectures and efficiency / Attention and transformers
Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes HeadKV-R2, a head-level KV cache compression method that allocates per-head cache budgets using retrieval-reasoning head importance scores estimated from constructed needle-in-a-haystack examples.
Problem
KV cache memory grows with input length, and existing KV compression methods allocate budgets at the layer level or ignore head-level importance for tasks needing both retrieval and reasoning.
Approach
Estimate each attention head's retrieval-and-reasoning importance using constructed retrieval-reasoning needle examples and attention on the full correct answer, then allocate a shared budget pool across heads proportionally to these scores, with SnapKV-style selection within each head.
Key results
- On Llama-3-8B-Instruct with KV size 128, HeadKV-R2 reaches 32.00 average F1 on LongBench QA versus 28.52 for Ada-SnapKV and 32.90 for FullKV.
- On Mistral-7B-Instruct with KV size 128, HeadKV-R2 averages 31.41 versus 28.09 for Ada-SnapKV on the same LongBench QA subset.
- On Reasoning-in-a-Haystack with KV size 128, HeadKV-R2 averages 56.84 on Llama-3-8B-Instruct versus 55.00 for Ada-SnapKV and 57.04 for FullKV.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| allocate cache using Retrieval-Reasoning heads rather than standard Retrieval heads (ablation) | HeadKV on Llama-3-8B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 30.38 -> 32.00 | Helped |
| use retrieval-reasoning examples and Retrieval-Reasoning head distribution instead of standard Retrieval heads (ablation) | HeadKV on Mistral-7B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 31.01 -> 31.41 | Helped |
| use the modified importance-score estimation over standard Retrieval Heads estimation (ablation) | HeadKV on Llama-3-8B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 30.38 -> 31.04 | Helped |
| use the modified importance-score estimation over standard Retrieval Heads estimation (ablation) | HeadKV on Mistral-7B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 31.01 -> 31.02 | No clear effect |
| allocate cache using Retrieval-Reasoning heads rather than standard Retrieval heads | HeadKV on Llama-3-8B-Instruct · Reasoning-in-a-Haystack, KV Size=128 | Avg. accuracy: 55.16 -> 56.84 | Helped |
| allocate cache using Retrieval-Reasoning heads rather than standard Retrieval heads | HeadKV on Mistral-7B-Instruct · Reasoning-in-a-Haystack, KV Size=128 | Avg. accuracy: 42.94 -> 43.09 | Helped |
| use HeadKV-R2 instead of Ada-SnapKV | KV cache compression on Llama-3-8B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 28.52 -> 32.00 | Helped |
| use HeadKV-R2 instead of Ada-SnapKV | KV cache compression on Mistral-7B-Instruct · LongBench Single-Doc and Multi-Doc QA, KV Size=128 | Avg.: 28.09 -> 31.41 | Helped |
| use HeadKV-R2 instead of Ada-SnapKV | KV cache compression on Llama-3-8B-Instruct · LooGLE Long Dependency QA, KV Size=128 | Avg.: 6.63 -> 7.16 | Helped |
| use HeadKV-R2 instead of Ada-SnapKV | KV cache compression on Mistral-7B-Instruct · LooGLE Long Dependency QA, KV Size=128 | Avg.: 8.38 -> 8.84 | Helped |
Surprising findings
- At KV size 1024, HeadKV-R2 on Llama-3-8B-Instruct (32.95) slightly exceeds the FullKV average (32.90) on the LongBench QA subset.
- Nearly 70% of heads receive zero importance score under the standard retrieval heads estimation, making the distribution too sparse for allocation.
Limitations
- Only one hyperparameter-sensitive beta and the importance distribution is task-static, computed from manually constructed examples rather than task-specific signals.
- Evaluated only on two 7-8B instruct models and QA-focused benchmarks.
- Gains on some datasets are small or negative (e.g., HeadKV-R2 is lower than baselines on NarrativeQA at KV 128 for Llama-3).
Experimental setup
Datasets: LongBench LooGLE Needle-in-a-Haystack (constructed retrieval-reasoning examples, 100 per model) bAbI / BABILong reasoning-in-a-haystack
Benchmarks: LongBench LooGLE Needle-in-a-Haystack Reasoning-in-a-Haystack (BABILong)
Metrics: F1 accuracy decoding latency peak memory usage
Baselines: SnapKV PyramidKV Ada-SnapKV HeadKV-R
Builds on: retrieval heads SnapKV PyramidKV Ada-KV
Closest papers by meaning
- RazorAttention: Efficient KV Cache Compression Through Retrieval Heads (ICLR 2025 · similarity 0.91)
- Which Heads Matter for Reasoning? RL-Guided KV Cache Compression (ICML 2026 · similarity 0.90)
- Semantic Integrity Matters: Benchmarking and Preserving High-Density Reasoning in KV Cache Compression (ICML 2026 · similarity 0.90)
- Training Free Exponential Context Extension via Cascading KV Cache (ICLR 2025 · similarity 0.90)
- KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction (NeurIPS 2025 · similarity 0.89)
- Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.