HeadKV: Head-Aware KV-Cache Methods

• HeadKV is a family of head-aware key-value cache compression methods that allocate non-uniform budgets based on the empirical significance of individual attention heads.
• It is applied across long-context language, visual autoregressive, and audio-language models to balance memory usage with high inference quality.
• Empirical evaluations show near full performance recovery with substantial memory savings through techniques like retrieval, reasoning, and spectral smoothing.

Searching arXiv for papers on HeadKV and related KV-cache compression methods. I’ll look up the relevant arXiv entries to ground the article in the cited papers. HeadKV denotes a family of head-aware key-value cache compression methods that allocate memory and eviction decisions at the granularity of individual attention heads rather than treating all heads in a layer uniformly. Across long-context LLMs, grouped-query attention systems, visual autoregressive models, and audio-LLMs, these methods exploit the empirical heterogeneity of heads—such as retrieval, reasoning, semantic-context, locality-biased, global, or modality-specialized behavior—to improve the trade-off between cache size and inference quality under fixed memory budgets (Fu et al., 2024, Lin et al., 4 Aug 2025, Liang et al., 20 May 2026, Wang et al., 8 Apr 2026).

1. Scope, terminology, and the KV-cache bottleneck

In Transformer decoding, the KV cache stores historical key and value tensors so that new tokens can attend to prior context without recomputing the full history. In a multi-layer decoder with $L$ layers and $H$ heads, a per-head cache at decoding step $t$ is written as $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$, or equivalently as fourth-order tensors $K,V\in\mathbb{R}^{L\times H\times t\times d}$ (Lin et al., 2024). This cache is a major systems bottleneck because its storage and memory-transfer costs grow with sequence length.

Within this setting, “HeadKV” is not a single algorithm but a recurring design principle: allocate cache non-uniformly across heads because heads contribute unequally to downstream performance. The term is used explicitly for several different methods, and related work extends the same head-aware premise to cooperative allocation, low-rank compression, and modality-specific scoring.

Variant	Domain	Core mechanism
HeadKV / HeadKV-R2	Long-context LLM QA	Per-head importance from retrieval or retrieval+reasoning examples
CompressKV (a.k.a. HeadKV)	GQA-based LLMs	Semantic Retrieval Heads plus layer-adaptive budgets
HeadKV	AR image generation	Local/global head classification plus stratified eviction
HeatKV	Visual autoregressive modeling	Per-head, per-scale static pruning
AudioKV	Audio-LLMs	Audio-critical head scoring plus spectral smoothing

A recurring premise across these methods is that layer-level or uniform-budget compression is too coarse. In the language setting, HeadKV was introduced partly because prior methods such as PyramidKV and SnapKV allocate budgets at layer level and therefore may waste cache on unimportant heads or under-allocate critical ones (Fu et al., 2024). In visual generation, fixed per-head budgets were likewise identified as suboptimal because some heads are locality-biased while others capture broader contextual structure (Liang et al., 20 May 2026).

2. HeadKV in long-context LLMs

The long-context LLM variant of HeadKV formulates compression as a two-stage pipeline: estimate head importance offline, then allocate per-head budgets and evict tokens accordingly. In “Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning” (Fu et al., 2024), the retrieval-only score $S_h^{\text{ret}}$ checks whether the top-scoring token under head $h$ falls inside an inserted answer span, while HeadKV-R2 replaces this with a retrieval-plus-reasoning score that inspects top-$k$ attended tokens and weights those lying in the correct answer segment $A^2$. The stated motivation is that retrieval+reasoning examples yield a denser, more informative head-importance distribution.

Budget assignment is explicit. With base head budget $b$, inflation factor $H$0, and local window $H$1, the dynamic pool is

$H$2

and the per-head budget is

$H$3

At inference time, the method follows SnapKV-style attention pooling: during prefill, it aggregates attention from the last $H$4 tokens to all past tokens, ranks tokens by pooled attention score, and keeps the top $H$5 tokens for each head (Fu et al., 2024).

The implementation choices are deliberately conservative. The paper fixes $H$6 instruct tokens for all methods, tunes $H$7 in $H$8, and does not hard-zero heads; every head receives a nonzero basic budget (Fu et al., 2024). This is important because a common simplification is to interpret head-aware compression as pure head pruning, whereas the reported algorithm instead performs differentiated retention.

Empirically, the method is evaluated on LongBench and LooGLE with Llama-3-8B-Instruct and Mistral-7B-Instruct under KV budgets $H$9 plus Full KV. At KV$t$0, HeadKV-R2 reaches average F1 values of $t$1 on Llama-3-8B-Instruct and $t$2 on Mistral-7B-Instruct, compared with $t$3 and $t$4 for FullKV, leading to the paper’s summary that it recovers about $t$5 of full-KV performance while retaining about $t$6 of tokens. The same paper reports that peak memory is reduced by about $t$7 at KV$t$8 versus FullKV, and that decoding latency is identical to other compressed-KV methods (Fu et al., 2024).

The paper also frames a broader methodological claim: contextual QA requires both retrieval and reasoning capability, so head importance should be estimated on examples that exercise both. The superiority of HeadKV-R2 over HeadKV-R and HeadKV-ER in the ablation table is consistent with that interpretation (Fu et al., 2024).

3. Semantic Retrieval Heads and layer-adaptive HeadKV

A second language-model line uses the name HeadKV through CompressKV, which is described as “CompressKV (a.k.a. HeadKV)” (Lin et al., 4 Aug 2025). Its central claim is that token eviction should not be driven by all heads in grouped-query attention systems. Instead, it first identifies, for each layer $t$9, a subset $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$0 of “Semantic Retrieval Heads.”

The offline head score is the layer-specific

$K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$1

where $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$2 is the correct answer span and $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$3 is the attention weight assigned to cached position $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$4 when generating token $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$5 (Lin et al., 4 Aug 2025). Heads are sorted by this score, and the top-$K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$6 heads form $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$7.

The retention rule then operates at layer level but is driven only by the selected heads. After prefilling a prompt of length $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$8, token importance for position $K_{\ell,h},V_{\ell,h}\in\mathbb{R}^{t\times d}$9 in layer $K,V\in\mathbb{R}^{L\times H\times t\times d}$0 is

$K,V\in\mathbb{R}^{L\times H\times t\times d}$1

computed over a short observation window, often the last few decoder steps. Given a per-layer budget $K,V\in\mathbb{R}^{L\times H\times t\times d}$2, the algorithm retains the top $K,V\in\mathbb{R}^{L\times H\times t\times d}$3 positions ranked by $K,V\in\mathbb{R}^{L\times H\times t\times d}$4, and for all other positions evicts $K,V\in\mathbb{R}^{L\times H\times t\times d}$5 for all heads $K,V\in\mathbb{R}^{L\times H\times t\times d}$6 in that layer (Lin et al., 4 Aug 2025).

Budgeting across layers is itself adaptive. CompressKV simulates extreme compression of each layer to a tiny cache, such as $K,V\in\mathbb{R}^{L\times H\times t\times d}$7 tokens, compares compressed and full attention-block outputs, and derives a normalized per-layer error $K,V\in\mathbb{R}^{L\times H\times t\times d}$8. Starting from lower bounds $K,V\in\mathbb{R}^{L\times H\times t\times d}$9, it distributes the remaining budget proportionally to $S_h^{\text{ret}}$0, subject to lower and upper bounds $S_h^{\text{ret}}$1 and a final fix-up loop ensuring $S_h^{\text{ret}}$2 (Lin et al., 4 Aug 2025).

This formulation directly challenges a common assumption in earlier heuristic eviction schemes: that aggregating over all heads is an adequate proxy for token importance. The paper reports that the selected semantic heads not only attend to initial and final tokens, as “streaming heads” do, but also collectively activate on answer spans and surrounding semantic context; masking these semantic heads causes a much larger performance drop than masking traditional copy-and-paste heads (Lin et al., 4 Aug 2025).

On LongBench, the reported results are unusually aggressive. With only $S_h^{\text{ret}}$3 of the full KV cache, CompressKV retains more than $S_h^{\text{ret}}$4 of question-answering accuracy; with $S_h^{\text{ret}}$5 cache, it matches over $S_h^{\text{ret}}$6 of full-cache average performance. On Needle-in-a-Haystack, the paper reports $S_h^{\text{ret}}$7 of baseline retrieval accuracy at $S_h^{\text{ret}}$8 of full cache, described as $S_h^{\text{ret}}$9 entries versus about $h$0K full (Lin et al., 4 Aug 2025).

4. Cooperative allocation and low-rank relatives

HeadKV-style work also sits alongside two neighboring directions: cooperative head-level budget allocation and feature-dimension compression. These papers do not always use the exact name “HeadKV,” but they sharpen the same underlying question: which parts of the KV cache matter, and how should a fixed budget be assigned?

CoKV argues that evaluating heads independently can misestimate their true effect because heads contribute cooperatively during inference. It formulates the extra-budget allocation problem

$h$1

casts heads or head-groups as players in a cooperative game, and scores them with Shapley or Sliced Shapley values. After discarding the $h$2 heads with the smallest values, it allocates

$h$3

The paper reports state-of-the-art LongBench performance on Llama-3-8B-Instruct and Mistral-7B, $h$4 of FullKV on Llama-3-8B at $h$5 tokens, under $h$6 of FullKV end-to-end prefill+decode time on Mistral-7B, and about $h$7 peak GPU memory reduction at $h$8K tokens (Sun et al., 21 Feb 2025).

MatryoshkaKV addresses a different axis: the feature dimension of cached keys and values. Rather than distributing token budgets, it projects each head’s cache from $h$9 to a lower rank $k$0 with trainable orthogonal matrices. The paper starts from PCA and explicitly notes its failure mode: when $k$1 falls below about $k$2, PCA projections suffer a rapid accuracy drop even though the keys or values can still be reconstructed with small squared error. MatryoshkaKV replaces static PCA with a distillation objective over trainable orthogonal matrices parameterized by Cayley transforms and uses a nested schedule with ranks in $k$3. It then performs adaptive search over layers and heads under a global budget, reporting more than $k$4 performance with an average KV compression rate of $k$5, up to $k$6 in extreme scenarios, and about $k$7 gain from head-wise adaptive allocation over uniform truncation (Lin et al., 2024).

ReCalKV is another feature-dimension method, but it explicitly incorporates head structure. For keys, it uses Head-wise Similarity-aware Reordering (HSR), computing pairwise CKA similarities among heads, greedily grouping similar heads, and applying grouped SVD after reordering the columns of $k$8. For values, it uses Offline Calibration and Matrix Fusion (OCMF), updating low-rank factors with closed-form offline calibration and fusing the calibrated factor into $k$9 so that no explicit reconstruction is required online. On LLaMA-2-7B at $A^2$0 compression, the paper reports WikiText2 perplexity $A^2$1 versus $A^2$2 for Palu, zero-shot average $A^2$3 versus $A^2$4, and LongBench average $A^2$5 versus $A^2$6 (Yan et al., 30 May 2025).

A plausible implication is that head-aware token retention and head-aware feature compression are complementary rather than mutually exclusive. The papers do not present a unified system, but together they show that “head awareness” can govern both which cached tokens are kept and how each head’s cached vectors are represented.

5. Visual autoregressive HeadKV variants

In autoregressive image generation, HeadKV refers to a distinct but structurally similar idea: classify or rank heads by attention behavior, then assign asymmetric cache budgets. The 2026 paper “Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation” defines two head types. A head is “local” if its minimal backward span $A^2$7 satisfying

$A^2$8

falls below a small window size $A^2$9; otherwise it is “global.” The default threshold is $b$0 (Liang et al., 20 May 2026).

Budgeting follows directly from this classification. Local heads keep only conditional tokens plus a fixed sliding window, $b$1, while global heads keep conditional tokens plus a fraction $b$2 of the historical cache and the same window, $b$3. Historical tokens for global heads are compressed by Stratified Token Eviction, which partitions history into long-range and near-range subsets, scores both with average attention over a query buffer of length $b$4, and evicts the lowest-scored $b$5 tokens from each subset (Liang et al., 20 May 2026).

The stated advantage is that head types can be identified early in generation and then reused, requiring no additional training or dataset-level statistics. The paper reports $b$6–$b$7 memory savings and $b$8–$b$9 speedups relative to full caching, with specific results such as LlamaGen-XL at $H$00 using $H$01 GB, $H$02 s latency, and $H$03 it/s on an NVIDIA A6000, and Janus-Pro-7B at $H$04 using $H$05 GB, $H$06 s, and $H$07 it/s (Liang et al., 20 May 2026).

HeatKV applies head tuning to visual autoregressive models that generate multiple scales. Rather than classifying heads as local or global, it aggregates per-head, per-scale attention into $H$08, defines a head-level Cumulative Attention Score (CAS) and a Scale-dependent CAS (S-CAS), and ranks head-scale pairs by importance. Using only $H$09 calibration prompts, it builds a static pruning schedule under budget fraction $H$10, with a formula for the number of head-scales to prune at scale $H$11:

$H$12

It then applies an early-pruning procedure that the paper states is optimal in the binary case under mild dominance conditions (Cederlund et al., 14 May 2026).

On Infinity-2B, the reported results at equal KV memory show clear gains over ScaleKV. At $H$13 budget, both methods use $H$14 MB, but HeatKV reports FID $H$15 versus $H$16, LPIPS $H$17 versus $H$18, and PSNR $H$19 versus $H$20. At $H$21 budget, HeatKV reports FID $H$22 versus $H$23. The paper further states that at $H$24, HeatKV matches or surpasses ScaleKV at $H$25 on all fidelity metrics, describing this as a $H$26 effective compression improvement (Cederlund et al., 14 May 2026).

These two visual variants illustrate that HeadKV is not tied to one scoring primitive. One method uses online behavioral classification with no offline statistics (Liang et al., 20 May 2026); the other uses offline calibration and a static per-scale schedule (Cederlund et al., 14 May 2026). What they share is the rejection of uniform budgets across heads.

6. Audio specialization, limitations, and recurring themes

AudioKV extends the head-aware KV premise to large audio-LLMs. It identifies “audio-critical” heads offline by semantic-acoustic alignment on ASR data. Using ASR anchors with confidence threshold $H$27, it maps recognized words to audio-token spans, records top-$H$28 attended positions for each decoding step, computes a per-step hit ratio

$H$29

and averages it over ASR steps to obtain $H$30 (Wang et al., 8 Apr 2026). Total budget $H$31 is then split proportionally across heads, or equivalently as a window-plus-baseline-plus-score decomposition.

Its second distinctive component is Spectral Score Smoothing (SSS), which treats per-head importance as a temporal signal, applies RFFT, chooses a cutoff by cumulative spectral energy, filters high frequencies, and mixes the smoothed and raw signals with coefficient $H$32. The stated rationale is to suppress high-frequency noise and recover smooth global trends in importance scores (Wang et al., 8 Apr 2026).

The reported evaluation spans five models—Qwen2.5-Omni-3B, Qwen2.5-Omni-7B, Qwen3-Omni-30B-A3B, Gemma-3n-E2B, and Gemma-3n-E4B—across ASR, speech translation, and audio question answering. The abstract states that at a $H$33 compression ratio, AudioKV maintains near-full accuracy on Qwen3-Omni-30B with only a $H$34 drop, while traditional methods suffer catastrophic degradation and repetition. The detailed results also report end-to-end slowdown under $H$35 and KV-cache memory reduction up to $H$36–$H$37 on NVIDIA A100 (Wang et al., 8 Apr 2026).

Several limitations recur across the HeadKV family. First, importance estimation may be task-specific. AudioKV explicitly notes that scores derived from ASR may need recalibration for other downstream tasks such as audio question answering (Wang et al., 8 Apr 2026). Second, independent head ranking may miss synergy, which is precisely the critique motivating CoKV’s Shapley-based allocation (Sun et al., 21 Feb 2025). Third, the best compression axis may depend on the setting: token eviction and budget allocation dominate many HeadKV methods, whereas MatryoshkaKV and ReCalKV demonstrate that feature-dimension compression can also preserve performance under high compression (Lin et al., 2024, Yan et al., 30 May 2025).

This suggests that HeadKV is best understood as a broader architectural principle: KV compression should follow the functional heterogeneity of attention heads rather than impose uniform retention rules. In LLMs that heterogeneity is expressed through retrieval and reasoning (Fu et al., 2024), semantic context (Lin et al., 4 Aug 2025), or cooperation among heads (Sun et al., 21 Feb 2025); in visual models through locality, global structure, and scale dependence (Liang et al., 20 May 2026, Cederlund et al., 14 May 2026); and in audio-LLMs through semantic-acoustic grounding and temporal continuity (Wang et al., 8 Apr 2026).