Source-linked AI summary
Beyond Mean Attention: Diversity-Aware, Layer-Wise Scoring for KV Cache Eviction
Tianfang Xie, Wei Zhu
TL;DR
Mean-attention KV eviction overlooks query disagreement and redundancy among selected tokens. The paper adds dispersion and selected-token similarity to a layer-wise score, then compares global and depth-varying settings. At B=64, a global setting improves 13 of 16 datasets, while passage retrieval shows a transferable depth profile.
Problem
SnapKV and PyramidKV rank tokens by mean attention over a small observation window, which overlooks query disagreement and redundancy among selected tokens.
Source: §1 · p. 1Method
The paper scores tokens using mean attention, attention dispersion, and similarity to the selected set, comparing global coefficients with searched depth profiles.
Source: §1 · p. 1, §3.2 · p. 2Results
At B=64, Global improves over PyramidKV on 13 of 16 datasets (+1.13 macro points); its macro gains are +0.73 at B=32 and +0.20 at B=128.
Source: §4.2 · p. 3Takeaways & Limitations
A global diversification constant is a robust default at small cache budgets, while passage retrieval supports a transferable task-specific depth profile.
Source: §4.2 · p. 3, §5 · p. 4On 12 HotpotQA prompts, greedy admission takes 222 ms per prompt versus 32 ms for top-k, increasing end-to-end latency by about 14%.
Source: §4.3 · p. 4
Abstract
from arXiv · showhide
KV cache eviction methods such as SnapKV and PyramidKV rank tokens solely by mean attention over a small observation window. We study a unified score, $\mu_i+\lambda_1\sigma_i+\lambda_2\mathrm{corr}(i,S)$, adding attention dispersion across window queries and redundancy relative to selected tokens. For $\lambda_2<0$, the score penalizes similarity to selected tokens as in maximal marginal relevance (MMR), without extra forward passes. To test whether this relevance-diversity balance should vary with depth, we compare fixed global coefficients with three-segment and quadratic profiles. Only these depth profiles are searched on a development split under a $\sinh$ reparameterization. On all 16 English LongBench datasets with Mistral-7B at a budget of 64 entries per layer, a single global diversification constant improves 13 of 16 datasets (macro +1.1); the gain holds at budget 32 and narrows at 128. Per-dataset search finds no detectable layer structure on most datasets; on passage retrieval it finds a large one: a mid-layer sign flip that rewards similarity and is worth +9.6 over the baseline at budget 64 and, without re-tuning, +13.2 over the global constant at budget 128. Ablations attribute the gain to the redundancy term; replaying every accepted search state on the held-out test set separates genuine structure from tuning noise.
1. INTRODUCTION
Mean-attention eviction addresses the growing memory footprint of long-context KV caches, but overlooks query disagreement and redundancy among selected tokens. The paper proposes a diversified, layer-wise scoring framework and evaluates global and depth-varying settings across LongBench.
1. INTRODUCTION
- The KV cache grows linearly with context length and often exceeds model weights, motivating token eviction under a small per-layer budget.SnapKV [1] and PyramidKV [2] rank historical tokens by mean attention from a recent observation window.Source: §1 · p. 1
- Mean attention misses query-to-query dispersion and can select near-duplicate tokens, evicting complementary evidence at tight budgets.The redundancy problem is one that maximal marginal relevance (MMR) addresses in information retrieval.Source: §1 · p. 1
- The proposed score combines mean attention, attention dispersion, and similarity to the selected set through layer-wise coefficients.The framework generalizes attention-based scoring under a fixed budget allocation and adds no extra forward passes.Source: §1 · p. 1
- At B=64, a fixed global setting improves 13 of 16 English LongBench datasets without per-dataset search.The study also reports a large transferable task-specific structure and ablations of the score terms.Source: §1 · p. 1
2. RELATED WORK
Related methods compress KV caches through attention-based eviction or redundancy-aware selection. This work distinguishes itself by conditioning redundancy on retained tokens and testing depth-dependent profiles across budgets.
2. RELATED WORK
- SnapKV [1] and PyramidKV [2] rank tokens independently using mean-attention variants; PyramidKV also decreases budgets with depth.SnapKV pools observation-window attention at the end of prefill.Source: §2 · p. 1
- R-KV, MixKV, and TwinKV [11] address redundancy through key similarity, per-head diversity, and repair swaps, respectively.Their approaches target decoding-time compression, vision-language caches, and retained near-duplicate keys, respectively.Source: §2 · pp. 1–2
- This work uses greedy MMR conditioned on retained tokens, searches redundancy weight by depth, and tests profile transfer across budgets.The related work also situates this trade-off alongside determinantal point processes and mean–variance portfolio theory.Source: §2 · p. 2
3. METHOD
The method extends mean-attention cache ranking with attention dispersion and similarity to already selected tokens, then compares compact depth-varying coefficient profiles. Tokens are admitted greedily under the cache budget, while profile searches use development scores to decide accepted moves.
3.1. Preliminaries
- SnapKV/PyramidKV rank tokens by mean attention from the last 8 queries and retain the top B, while PyramidKV decreases budgets with depth but preserves the uniform allocation’s total budget.The scoring pipeline retains the baselines’ max-pooling.Source: §3.1 · p. 2
3.2. Diversity-Aware Scoring
- The unified score adds query-to-query attention dispersion and greatest similarity to a retained token; negative λ2 penalizes similarity, while positive λ2 rewards it.The base score is normalized per head before greedy selection; the process uses existing attention and keys without extra forward passes.Source: §3.2 · p. 2
3.3. Layer-Wise Parameterization
- To limit overfitting from tuning 2L coefficients on 100 development questions, the method compares global, three-segment, and quadratic depth profiles and accepts search moves using development scores only.The quadratic profile is mapped to coefficients through sinh; Global uses λ1=0 and λ2=−0.15 and is not searched.Source: §3.3 · pp. 2–3
4. EXPERIMENTS
On Mistral-7B across 16 English LongBench datasets, fixed global diversification improves 13 of 16 datasets at B=64, while depth-specific search finds its strongest benefit on passage retrieval.
4.1. Experimental Setup
- The study evaluates Mistral-7B-Instruct-v0.2 on all 16 English LongBench datasets, usually with 64 entries per layer and a shared total cache budget.B=64 is approximately 1% of a typical prompt.Source: §4.1 · p. 3
4.2. Main results
- On passage retrieval, Segments scores 70.17, exceeding PyramidKV by 9.57 points and Global by 3.36 points.The B=64-tuned profile scores 63.03 at B=32 and 81.70 at B=128; at B=128 it is 13.40 above PyramidKV and 13.16 above Global.Source: §4.2 · p. 3
- At B=64, Global raises the LongBench macro score from 33.31 to 34.44 (+1.13 points) and beats PyramidKV on 13 of 16 datasets.The comparison uses Mistral-7B and the official LongBench scorer.Source: Table 1 · p. 3, §4.2 · p. 3
- Across B=32, 64, and 128, Global improves the macro score over PyramidKV by 0.73, 1.13, and 0.20 points, while Seg. gains 1.01, 1.84, and 0.60.Profiles deployed at B=64 transfer unchanged to the other budgets.Source: Table 2 · p. 3, §4.2 · p. 3
4.3. Ablation studies and further analysis
- For passage retrieval, λ2 turns positive in middle and deep layers; replay shows the sign flip yields a large test gain, and the profile beats PyramidKV and Global at every budget.The accepted λ1 updates do not improve the test score.Source: §4.3 · pp. 3–4, Figure 1 · p. 4
- Searching λ2 alone reaches 34.57 and improves 13 datasets, versus 33.50 and 8 datasets for λ1 alone; searching both reaches 35.15.Without mean attention, the fixed rule scores 33.16, below the 33.31 baseline, and drops passage retrieval to 57.41.Source: Table 3 · p. 4, §4.3 · p. 4
- Three-segment profiles score 35.15, compared with 35.03 for eight segments and 34.75 for two; Curve scores 34.33.The reported scores put segment count within the noise floor of the three-segment result.Source: §4.3 · p. 4
5. CONCLUSION
Diversified, layer-wise scoring extends mean-attention KV eviction without extra forward passes. A global diversification constant is a robust default at small budgets, while passage retrieval benefits from a layer-wise profile.
5. CONCLUSION
- A single global diversification constant is a robust default at small cache budgets, while passage retrieval benefits from a layer-wise profile whose gain survives budget changes.The conclusion reports both the broad default and the task-specific exception.Source: §5 · p. 4
- Ablations place the benefit in the redundancy term computed on cached keys.A trajectory-replay diagnostic makes the findings verifiable.Source: §5 · p. 4
- Instance-adaptive coefficients are identified as a next step.Source: §5 · p. 4
7. COMPLIANCE WITH ETHICAL STANDARDS
The study uses the publicly available LongBench benchmark and collected no new data or human or animal subjects; ethical approval was therefore not required.