PaperlayerView exact revision on arXiv

Source-linked AI summary

Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning

Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, Wen Xiao

arXiv:2410.19258v4cs.CLcs.AI

TL;DR

Longer inputs make KV-cache memory a growing challenge, while existing compression approaches do not allocate cache budgets across individual heads. HeadKV estimates head importance for retrieval and reasoning and allocates cache accordingly; experiments report stronger performance across contextual tasks, especially in low-cache settings.

  • Problem

    KV-cache memory grows rapidly with input length, and prior compression work has not varied cache size across individual heads.

    Source: §1 · p. 1

  • Method

    HeadKV estimates each attention head’s importance for retrieval and reasoning, then allocates larger KV-cache budgets to more important heads.

    Source: §1 · p. 1

  • Results

    Head-level compression outperforms strong baselines on contextual tasks, especially at 64 and 128 KV cache, while Retrieval-Reasoning allocation scores 32.95 versus the Full-KV cache average of 32.90 at 1024 KV cache.

    Source: §4.2 · p. 6

  • Takeaways & Limitations

    The reported results support allocating cache budgets across heads by importance to retain contextual retrieval and reasoning performance under compression.

    Source: §1 · p. 2, §5 · p. 10

    Dynamic per-head management can add synchronization and computational overhead in parallel execution, potentially reducing decoding speed in multi-GPU environments.

    Source: I LIMITATION · p. 20

Abstract

from arXiv · show

Key-Value (KV) caching is a common technique to enhance the computational efficiency of Large Language Models (LLMs), but its memory overhead grows rapidly with input length. Prior work has shown that not all tokens are equally important for text generation, proposing layer-level KV cache compression to selectively retain key information. Recognizing the distinct roles of attention heads in generation, we propose HeadKV, a head-level KV cache compression method, and HeadKV-R2, which leverages a novel contextual reasoning ability estimation for compression. Our approach operates at the level of individual heads, estimating their importance for contextual QA tasks that require both retrieval and reasoning capabilities. Extensive experiments across diverse benchmarks (LongBench, LooGLE), model architectures (e.g., Llama-3-8B-Instruct, Mistral-7B-Instruct), and long-context abilities tests demonstrate that our head-level KV cache compression significantly outperforms strong baselines, particularly in low-resource settings (KV size = 64 & 128). Notably, our method retains just 1.5% of the KV cache while achieving 97% of the performance of the full KV cache on the contextual question answering benchmark. Codes are available at https://github.com/FYYFU/HeadKV

1 INTRODUCTION

Longer LLM inputs make KV-cache memory a growing efficiency challenge. HeadKV allocates cache across individual heads using retrieval-and-reasoning importance, and experiments report gains over prior approaches across contextual tasks.

1 INTRODUCTION

  • KV-cache memory grows rapidly with input length, despite caching key and value states to reuse them during generation.The passage connects long inputs with increased memory use and latency from transformer self-attention.Source: §1 · p. 1
  • HeadKV allocates larger KV budgets to more important heads, addressing prior compression methods that had not varied cache size across individual heads.Importance estimation jointly considers retrieval and reasoning, which the paper identifies as necessary for contextual QA.Source: §1 · pp. 1–2
  • For contextual QA, HeadKV estimates head importance from attention scores while answering questions that require both retrieving context and reasoning over it.Within each head, it retains the most relevant KV entries using the strategy proposed in Li et al. (2024d).Source: §1 · p. 2
  • Head-level compression outperforms previous approaches on nearly all evaluated tasks and reduces memory use and decoding latency versus full KV.Evaluations include LongBench and LooGLE QA, needle-in-a-Haystack, and reasoning-in-a-Haystack tests.Source: §1 · p. 2

2 RELATED WORK

Prior work analyzes distinct attention-head functions and compresses KV caches, but the described methods retain layer-level constraints or full caches for key heads. The paper positions its approach as independent head-level allocation.

2.1 ATTENTION HEADS

  • Prior attention-head studies characterize specialized roles, including retrieval heads identified with Needle-in-a-Haystack tests, motivating compression that preserves retrieval and reasoning capabilities.The cited functions also include positional information, syntactic parsing, rare-word focus, and induction-head copying.Source: §2.1 · p. 3

2.2 KV CACHE COMPRESSION

  • Prior compression approaches reduce memory through varied strategies, including retaining initial tokens, selecting heavy hitters or recent-attention entries, and assigning smaller budgets to higher layers.The passage describes StreamingLLM, Heavy Hitter Oracle, attention-based selection, and PyramidKV.Source: §2.2 · p. 3
  • Earlier KV-cache compression methods use layer-level allocation; Feng et al. (2024) still use layer-level budgets, while Tang et al. (2024) retain full KV for key heads.The paper contrasts these constraints with allocating cache solely from head-level importance scores, independent of layer constraints.Source: §2.2 · p. 3

3 METHOD

The method estimates head importance, allocates KV cache budgets across heads, and selects retained cache entries. Its importance estimation targets contextual QA by accounting for both retrieval and reasoning.

3 METHOD

  • The method has three components: identify important heads, allocate cache budgets across heads, and determine which key and value vectors to retain within each head.Source: §3 · p. 3

3.1 HEAD-LEVEL IMPORTANCE SCORE ESTIMATION

  • Retrieval-Reasoning examples add a reasoning step and competing answers, and score heads using attention to the entire correct answer rather than only an exact-match top-attention token.The original retrieval-only distribution is too sparse for allocation, with nearly 70% of heads receiving zero scores, and does not account for contextual reasoning.Source: §3.1 · pp. 4–5

3.2 HEAD-LEVEL KV CACHE ALLOCATION

  • Each head receives a basic KV budget plus a share of a pooled dynamic budget proportional to its normalized importance score, while the last α instruction tokens are preserved.The hyperparameter β controls the pool size.Source: §3.2 · p. 5

3.3 KV CACHE SELECTION

  • Local windows formed from the last α instruction tokens guide selection by pooling their attention scores to remaining tokens and retaining higher-scoring entries.This attention-based selection follows budget allocation.Source: §3.3 · p. 5, §3.2 · p. 5

4 EXPERIMENTS AND ANALYSIS

The experiments compare head-level KV cache allocation with strong baselines across contextual QA benchmarks and long-context tests. HeadKV-R2 performs strongly, especially at small cache sizes, while maintaining comparable decoding efficiency and reducing memory use relative to Full KV.

4.1 EXPERIMENT SETTINGS

  • The experiments cover Single-Doc QA and Multi-Doc QA from LongBench and Long Dependency QA from LooGLE, using Llama-3-8B-Instruct and Mistral-7B-Instruct.Baselines retain the same number of KV entries; the proposed method uses SnapKV for per-head selection and sets the local window to α = 8.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • Head-level compression outperforms baselines, especially at KV size 128, and reaches 32.95 versus FullKV's 32.90 on Llama-3-8B-Instruct at KV size 1024.Table 1 compares LongBench and LooGLE results for Llama-3-8B-Instruct and Mistral-7B-Instruct.Source: §4.2 · p. 6, Table 1 · p. 7
  • Figure 3 reports average accuracy across six LongBench datasets for KV cache sizes 64, 128, 256, 512, and 1024.The average input length is 8,683 tokens, and size 64 retains 0.7% of total tokens.Source: Figure 3 · p. 7

4.3 RETRIEVAL-REASONING HEADS

  • HeadKV-R2 combines retrieval-reasoning examples with modified importance estimation and consistently outperforms the retrieval-based alternatives.The ablation also finds that Enhanced-Retrieval improves slightly, but it still relies on retrieval examples and lacks full consideration of contextual reasoning.Source: §4.3 · p. 8, Table 2 · p. 8

4.4 LONG-CONTEXT RETRIEVAL AND REASONING

  • At KV size 128, head-level compression outperforms strong baselines on Needle-in-a-Haystack retrieval and Reasoning-in-a-Haystack tasks.The reasoning test inserts multiple needles and averages results across QA1–QA5 for each length; the authors note that Needle-in-a-Haystack does not assess contextual reasoning.Source: §4.4 · pp. 8–9, Table 3 · p. 9

4.5 MEMORY & LATENCY

  • Decoding latency matches other compression methods, while peak memory is comparable to compression baselines and lower than Full KV.The latency comparison includes prefilling and decoding; the reported peak-memory results average three trials.Source: §4.5 · p. 10, Figure 6 · p. 10

5 CONCLUSION AND FUTURE WORK

The conclusion presents HeadKV-R2 as head-level cache allocation guided by importance scores that account for retrieval and reasoning. The authors identify exploration of additional head types and task-specific score estimation as future directions.

5 CONCLUSION AND FUTURE WORK

  • HeadKV-R2 allocates cache budgets across attention heads using importance scores estimated from examples involving both retrieval and reasoning.The method assigns larger budgets to more important heads and smaller budgets to less important ones.Source: §5 · p. 10
  • Future work could explore attention heads tied to in-context learning or truthfulness and task-specific score estimation using gradients.The authors suggest such task-specific allocation could tailor compression across applications.Source: §5 · p. 10

A DETAILE RESULTS

The appendix points to detailed results for the different KV cache sizes shown in Figure 3.

A DETAILE RESULTS

  • Table 4 provides detailed results for KV cache sizes 64, 128, 256, 512, and 1024 on Llama-3-8B-Instruct and Mistral-7B-Instruct.The appendix identifies Table 4 as the detailed results corresponding to Figure 3.Source: Table 4 · p. 15, A DETAILE RESULTS · p. 15

B DATASET DETAILS

The appendix describes the question-answering datasets used to evaluate retrieval and contextual reasoning.

B DATASET DETAILS

C LONGBENCH RESULTS

Head-level KV cache compression outperforms the other LongBench baselines, particularly on summarization and Code datasets across Llama-3-8B-Instruct and Mistral-7B-Instruct. The Mistral-7B-Instruct figure also reports HeadKV-R2 outperforming other strong baselines on the Needle-in-a-Haystack test.

C LONGBENCH RESULTS

  • Head-level KV cache compression outperforms other LongBench baselines, particularly on summarization and Code datasets across Llama-3-8B-Instruct and Mistral-7B-Instruct.Table 6 reports results with KV cache=128.Source: Table 6 · p. 16, C LONGBENCH RESULTS · p. 16
  • The enhanced head distribution also covers important retrieval heads on Mistral-7B-Instruct.Source: Figure 7 · p. 16
  • HeadKV-R2 still outperforms other strong baselines on the Mistral-7B-Instruct Needle-in-a-Haystack test.The figure caption says these results are consistent with those on Llama-3-8B-Instruct.Source: Figure 8 · p. 17

D REASONING-IN-A-HAYSTACK DETAILS

The section describes the Reasoning-in-a-Haystack evaluation and how Retrieval-Reasoning examples are constructed to estimate head importance. The resulting distribution guides head-level KV budget allocation, with implementation details provided in pseudo-code.

D REASONING-IN-A-HAYSTACK DETAILS

  • The test inserts sentences from curated “Input” needles at varied positions in haystacks; five tasks and different context lengths are evaluated.The dataset is curated by Kuratov et al. (2024), and its examples are shown in Figure 12.Source: D REASONING-IN-A-HAYSTACK DETAILS · p. 17
  • The proposed method achieves the best average score across Reasoning-in-a-Haystack methods and surpasses Full KV cache on Llama-3-8B-Instruct QA3, the most challenging task.The authors attribute this to distributions better capturing semantically relevant and important content.Source: D REASONING-IN-A-HAYSTACK DETAILS · p. 17

E CONSTRUCT RETRIEVAL-REASONING EXAMPLES

  • Retrieval-Reasoning examples add reasoning paths and incorrect-answer distractors to retrieval examples, aiming to identify heads attending to correct answers despite similar distractors.The question is reversed to create two examples; 100 examples per model span five haystack lengths and ten question depths per length.Source: E CONSTRUCT RETRIEVAL-REASONING EXAMPLES · pp. 17–18
  • The estimation focuses on correct answers so the maximum head importance score remains 1, matching the Needle-in-a-Haystack output and the Retrieval Heads setup.The resulting scores guide head-level KV budget allocation; heads focused on incorrect answers are ignored within the shared global budget pool.Source: E CONSTRUCT RETRIEVAL-REASONING EXAMPLES · p. 18

F PSEUDO CODE

  • Pseudo-code derives per-head capacities from normalized importance scores and selects cached keys and values using SnapKV attention scores.Selected cache entries are combined with the observation window; the static importance distribution means budget determination runs once at initialization.Source: F PSEUDO CODE · pp. 18–19

G HYPER-PARAMETER ANALYSIS

The hyper-parameter analysis reports that Head-R2 performs better with smaller β and that both head-based strategies outperform Ada-SnapKV across tested β values. The discussion contrasts KV cache compression with context compression and notes potential parallel-execution overhead.

G HYPER-PARAMETER ANALYSIS

  • β is the method’s only introduced hyper-parameter and controls shared global budget pool size; tested values range from 1.005 to 10.The reported choice is based on average performance rather than selecting a β separately for each dataset.Source: G HYPER-PARAMETER ANALYSIS · pp. 19–20
  • Head-R2 performs better with smaller β, while Head-R and Head-R2 outperform Ada-SnapKV across tested β values on both models.A smaller β corresponds to a larger shared budget pool and greater reliance on the importance-score distribution.Source: G HYPER-PARAMETER ANALYSIS · p. 20

H DISCUSSION

  • KV cache compression targets the cache itself, whereas context compression targets the original input or representations derived from it.The discussion characterizes current KV cache methods as simpler and more computationally efficient than context-compression methods because they avoid extra trained components and recomputing key and value matrices.Source: H DISCUSSION · p. 20

I LIMITATION

Head-level KV cache compression may add overhead from dynamically managing individual heads, particularly during parallel execution, creating a potential trade-off with decoding speed in multi-GPU settings.

I LIMITATION

  • Head-level operations may add synchronization and computation costs in multi-GPU execution, potentially reducing decoding speed.The authors identify mitigating parallel overhead while maintaining practical deployment benefits as an open challenge.Source: I LIMITATION · p. 20

Related papers

Loading 2410.19258v4…