Research map / Large language models / Reasoning and chain of thought
KeyDiff: Key Similarity-Based KV Cache Eviction for Long-Context LLM Inference in Resource-Constrained Environments
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 KeyDiff, a training-free, attention-free KV cache eviction method that keeps keys with the lowest cosine similarity to a mean-key anchor, enabling block-wise prompt processing under strict memory budgets.
Problem
KV cache memory grows with prompt length, and existing attention-based eviction policies degrade under block-wise prompt processing (where only the current block and cache are visible) and require materializing attention matrices, which is incompatible with FlashAttention in resource-constrained set
Approach
Observes that keys with low pairwise cosine similarity (geometrically distinctive keys) tend to receive high attention; scores each cached key by cosine similarity to an anchor vector (mean of normalized keys) and evicts the most similar ones, giving an O(n) attention-free score with a theoretical link to attention weights and a subset-diversity optimization view.
Key results
- On LongBench with Llama-3.1-8B-Instruct, KeyDiff at 8K budget reaches 49.03 average versus 50.x-range non-evicting baseline (49.20), a gap under 0.2 points, and 48.51 at 6K.
- On LongBench with Llama-3.2-3B, KeyDiff at 8K reaches 45.78 average versus 45.87 for the full-context baseline, outperforming SnapKV (44.96) and TOVA (43.21) at the same budget.
- Time-to-first-token with FlashAttention is reported up to 30% lower than TOVA and SnapKV for Llama 3.2-3B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use KeyDiff instead of SnapKV for KV cache eviction | Llama3.1-8B · LongBench English subset, 2K–8K cache | Avg. accuracy: 44.33 vs 39.60; 47.66 vs 45.18; 48.51 vs 47.12; 49.03 vs 48.06 | Helped |
| Use KeyDiff instead of SnapKV for KV cache eviction | Llama3.2-3B · LongBench English subset, 2K–8K cache | Avg. accuracy: 38.14 vs 35.46; 43.67 vs 40.60; 45.27 vs 43.55; 45.78 vs 44.96 | Helped |
| Add a sliding window to KeyDiff | KeyDiff · LongBench coding tasks | Accuracy | Helped |
| Vary the anchor vector choice (ablation) | KeyDiff · Ablation benchmark results | Accuracy | No clear effect |
| Use cosine similarity as the distance metric (ablation) | KeyDiff · Ablation benchmark results | Accuracy | Helped |
| Use KeyDiff with FlashAttention instead of TOVA or SnapKV | Llama3.2-3B · Inference latency evaluation | End-to-end inference latency: up to 30% | Helped |
Surprising findings
- KeyDiff with a 2K budget on Math-500 slightly outperforms baselines in some configurations, and Llama-8B with KeyDiff at moderate budgets performs comparably to or slightly better than the eviction-free baseline.
- Attention-based baselines improve substantially when moved from per-token full-prompt eviction to block-wise processing only in some settings; Spearman correlation between negative key cosine similarity and attention averages about 0.94 across layers.
Limitations
- Designed and evaluated primarily for GQA attention used in Llama and Qwen; not validated on other attention variants such as Multi-Head Latent Attention.
- Qwen models show larger degradation on reasoning tasks, which the authors attribute to higher GQA compression.
- Theoretical lemmas rely on assumptions (bounded key norms, n to infinity limit) that are only empirically motivated.
Experimental setup
Datasets: LongBench Needle In a Haystack Math-500
Benchmarks: LongBench Needle In a Haystack Math-500
Metrics: accuracy F1 / task score time to first token
Baselines: H2O TOVA SnapKV StreamingLLM (sink attention)
Builds on: H2O TOVA SnapKV StreamingLLM FlashAttention
Closest papers by meaning
- Fast KVzip: Efficient and Accurate LLM Inference with Gated KV Eviction (NeurIPS 2026 · similarity 0.91)
- Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference (NeurIPS 2025 · similarity 0.90)
- SmallKV: Small Model Assisted Compensation of KV Cache Compression for Efficient LLM Inference (NeurIPS 2025 · similarity 0.90)
- Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference (ICML 2024 · similarity 0.89)
- SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.89)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 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.