Title: KeyDiff: Key Similarity-Based KV Cache Eviction for Long-Context LLM Inference in Resource-Constrained Environments
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models

LLM

Key-Value Cache

KV Cache

KV Cache Eviction

Token Eviction

Efficient LLM Inference
Abstract: We demonstrate that geometrically distinctive keys during LLM inference tend to have high attention scores. Based on the phenomenon we propose KeyDiff, a training-free KV cache eviction method based solely on key similarity. Unlike other KV cache eviction methods, KeyDiff can process arbitrarily long prompts within strict resource constraints and efficiently generate responses.
We provide a theoretical basis for KeyDiff by relating key diversity with attention scores. These results imply  KeyDiff can efficiently identify the most important tokens to retain. Notably KeyDiff does not rely on attention scores, allowing the use of optimized attention mechanisms like FlashAttention. Under a strict memory allowance, we demonstrate the effectiveness of KeyDiff for the Llama and Qwen model families by observing a performance gap of less than 0.04\% with 8K cache budget (~23\% KV cache reduction) from the non-evicting baseline on LongBench for Llama 3.1-8B and Llama 3.2-3B. We also observe near baseline performance for Deepseek-R1-Distill-Llama-8B on the Math500 reasoning benchmark and decrease end-to-end inference latency by up to 30\% compared to the other token-eviction methods.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes KEYDIFF, a training-free policy for evicting KV-cache entries when running large language models under limited memory and resource budgets. The key insight is that keys that are geometrically distinct—those with low cosine similarity to the rest of the cache—tend to receive higher attention weights, so preferentially retaining these keys preserves model performance. Algorithmically, the method computes the mean of all current keys as an anchor, scores each key by its negative cosine similarity to this anchor, and keeps the top-N least-similar keys, yielding an efficient O(N) procedure. Extensive experiments demonstrate that KEYDIFF outperforms existing baselines in both latency and memory consumption while maintaining accuracy.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The idea is novel. Requires only cosine distance; independent of attention implementation, so it can work with FlashAttention.\n2. Clear proof on the insight.\n3. Benchmarks cover essential aspects,\n4. Linear time solution; easy to apply\n\nWeakness:\n1. Limited architecture coverage; no evidence to show its effectiveness on MoE and encoder-decoder.\n2. Need to test on resource constrained environment.\n3. Need to show the anchor sensitivity at low budget.\n\nQUESTIONS:\n1. For the anchor update schedule, do you recompute the anchor after each eviction? What is the incremental overhead?\n\n2. What will KEYDIFF perform if we use quantized key? Does cosine similarity remain reliable?\n\n3. In chat mode, how does sliding-window KEYDIFF choose the window size?\n\n4. Could KEYDIFF combine with RAG? In this case, the cache contains both prompt and retrieved chunks.\n\n5. How does KEYDIFF perform on mobile device, since it is a resource contained environment.\n\nLIMITATIONS:\n1. Only decoder-only models are tested.\n2. Benchmarks assume ample GPU bandwidth; edge-class DRAM and cache is not tested\n3. KEYDIFF ranks keys on a per-layer basis but not per head; if one head relies on a subset of highly self-similar keys, those keys might be evicted.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-20"}

comments: SUMMARY:
To address the challenges of the existing attention-based KV cache eviction methods, this work proposes a KV cache evictino method based solely on key similarity, named KeyDiff. The proposed KeyDiff method minimizes the pairwise cosine similarity among keys, thereby maximizing the diversity of the compressed KV cache. Extensive experiments validate that KeyDiff achieves significant superiority in both algorithm accuracy and inference effectiveness.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The motivation of this work is demonstrated clearly and is easy to follow.
2. The experiments are rich, which demonstrate the effectiveness and efficiency of the proposed method.

Weaknesses:
1. The complete evaluation result combining different KV eviction methods with the 'sliding window' strategy is lacking. The KeyDiff with 'sliding window' can achieve better performance in reasoning and code tasks, while the effects of other methods (TOVA, Sink, SnapKV) combined with 'slicing window' are unclear. 
2. The discussion on the ‘Attention-Based Token Eviction Challenges’ in Section 2.4 is more like a general challenge associated with block-wise processing strategies. How does replacing attention with cosine similarity specifically address this challenge?
3. The referencing of experimental figures and tables is poorly organized, compromising the article's coherence and readability.

QUESTIONS:
See above.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-21

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a training-free KV-cache eviction policy that scores tokens purely by how geometrically distinctive their keys are (low mean cosine similarity to other cached keys), which shows advantages in the chunk-prefill setting.\n\nThis observation that distinctive keys tend to attract high attention regardless of the current query lets KeyDiff reduce peak memory of prefilling with higher accuracy, which is critical for serving/edge deployments where the peak memory of a full-prompt prefill would exceed RAM budgets.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength:**\n- Well-motivated setting: Chunk prefill under a hard KV budget is important in both serving and memory-constrained on-device LLMs inference.\n- Simple but novel insight: Linking low pairwise cosine similarity to high long-range importance is intuitive yet, to my knowledge, new and is backed by both visualization and a neat geometric bound.\n\n**Weaknesses:**\n- NIAH still lags the no-evict baseline: Figure 6 shows recall tapering for deeper needles; although KeyDiff beats other evictors, it is still materially below full-cache accuracy for long docs.\n- Missing harder retrieval tasks: The paper lacks harder benchmarks, such as RULER stress very rare tokens, which raises the concern about whether the property always holds.\n\nQUESTIONS:\n- How does KeyDiff perform on RULER or similarly harder long-context datasets?\n- The authors mention extending KeyDiff to MLA. In MLA, all heads share the same low-rank token embedding E, with different Keys computed through different K projections. Dropping a token means discarding E and thus removing that token’s keys from every head, even if it is distinctive for only some heads. How will your method handle this head-coupling constraint?\n- The authors' “distinct keys ⇒ high future attention” assumption is intriguing, but will it hold on tasks where most tokens are already mutually dissimilar? When there are many key scores as “distinct,” KeyDiff may have no reliable ranking signal and could evict truly important tokens. Have you evaluated such scenarios, and how does the method safeguard against this failure mode?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLong-context inference has become increasingly common. Due to resource constraints, the prefill must be chunked into small segments, which requires sparse attention algorithms to make decisions without access to the full context. This paper proposes a KV cache eviction policy based on the cosine similarity of Key vectors. By identifying Key vectors that significantly deviate from the mean, KeyDiff effectively retains important tokens. Evaluation shows that KeyDiff achieves high accuracy while also improving performance.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- Detailed evaluation across different models and datasets.\n- The proposed method is easy to implement and introduces limited overhead.\n- KeyDiff supports chunked prefill and GQA.\n\n## Weaknesses\n- More evidence is needed to demonstrate the correlation between Key cosine similarity and the attention score.\n\nQUESTIONS:\nThanks for submitting to NeurIPS. This paper claims that key similarity can determine which keys are important, which is simple yet effective. I enjoyed reading the paper and would advocate for it if the proposed solution is widely applicable. I appreciate the effort to provide an intuitive explanation in the appendix. However, I kindly ask the authors to provide some additional evidence to further strengthen the argument:\n\n1. **Correlation Analysis**  \n   I would like to see the correlation between Key vector cosine similarity and attention scores. I understand that figures cannot be submitted during the rebuttal, so a table showing cosine similarity and attention scores sampled from request prompts would be helpful. Please include the Spearman rank correlation between cosine similarity and attention scores.\n\n2. **Partial Cache Attention Scores**  \n   What percentage of the total attention score is retained when only a subset of the KV cache (e.g., 4K tokens) is used compared to full attention?\n\n3. **Performance on Retrieval-Critical Tasks**  \n   I would like to see the performance of KeyDiff on tasks that require retrieving all keys to get the correct answer — for example, the Counting Star dataset. Due to time constraints, evaluating one model (e.g., LLaMA3.1 8B) is sufficient.\n\n4. **RULER Benchmark**  \n   Could you include results on the RULER benchmark, which is considered more challenging than LongBench?\n\n5. **NIAH Recall Saturation**  \n   For the NIAH dataset, while the token budget is 6K, none of the policies achieve 100% recall even when the token limit is 1000 or 4222. Is this due to model limitations? What about larger models?\n\n6. **PCA Analysis Interpretation**  \n   The paper shows multiple PCA analysis figures. While it is true that KeyDiff selects tokens more uniformly, it is unclear why this leads to better performance. Could you elaborate on how to interpret these graphs?\n\n7. **Key Vector Norm Variations**  \n   While the paper claims that Key vectors have similar norms, the beginning-of-text token tends to be an outlier with a smaller norm. Does KeyDiff account for this?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-29"}