通过键相似性优化缓存淘汰，实现长文本推理的高效低内存运行。
KeyDiff: Key Similarity-Based KV Cache Eviction for Long-Context LLM Inference in Resource-Constrained Environments
- 基于键向量相似性判断重要性，无需训练即可淘汰冗余缓存。
- 在8K缓存预算下，性能损失低于0.04%，比基线节省约23%缓存空间。
- 兼容FlashAttention等加速机制，适合资源受限环境下的长序列推理。
我们发现，大语言模型推理中几何上独特的键通常具有较高的注意力分数。基于此现象，提出无需训练的KeyDiff方法，仅依据键相似性进行KV缓存淘汰。与现有方法不同，KeyDiff可在严格资源约束下处理任意长提示，并高效生成响应。我们从理论上建立了键多样性与注意力分数之间的关联，表明KeyDiff能有效识别需保留的关键标记。值得注意的是，KeyDiff不依赖注意力分数，可兼容FlashAttention等优化机制。在严格内存限制下，对Llama和Qwen系列模型进行测试，在LongBench上，Llama 3.1-8B和Llama 3.2-3B使用8K缓存预算（约减少23% KV缓存）时，性能差距小于0.04%；Deepseek-R1-Distill-Llama-8B在Math500推理基准上接近基线表现，且相比其他淘汰方法，端到端推理延迟最高降低30%。
原文摘要 · Abstract (English)
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 ($\sim$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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。