arXiv:2412.08521cs.CL2024-12被引 5

提出新方法压缩大模型推理中的键值缓存,显著提升长文本处理效率。

EMS: Adaptive Evict-then-Merge Strategy for Head-wise KV Cache Compression Based on Global-Local Importance

  • 用全局与局部注意力融合评分精准识别重要令牌
  • 自适应合并策略在极端压缩比下仍保持低困惑度
  • 适合需要高效长文本推理的部署场景

随着大语言模型发展,长上下文高质量、快速处理需求日益增长。键值缓存通过存储已生成的键值令牌,有效减少推理中的重复计算。然而,内存开销成为关键瓶颈,因此高效压缩键值缓存备受关注。现有方法多从重要性识别与压缩策略设计两方面入手,但常因累积注意力分数或位置编码影响,导致重要令牌分布偏差,并忽略不同注意力头间的稀疏性与冗余性,难以在头级别保留有效信息。为此,本文提出EMS,通过全局-局部联合得分更准确评估令牌重要性,设计自适应统一的‘剔除-合并’压缩框架,兼顾各头间稀疏性与冗余性。同时引入零类机制实现头级并行压缩,提升效率。大量实验表明,即使在极端压缩比下,EMS仍达当前最优性能:在LongBench上,四个模型平均提升超1.28分,256缓存预算下困惑度最低;在Needle-in-a-Haystack任务中,缓存预算低于上下文长度2%时,仍保持95%检索准确率。

原文摘要 · Abstract (English)

As large language models (LLMs) continue to advance, the demand for higher quality and faster processing of long contexts across various applications is growing. KV cache is widely adopted as it stores previously generated key and value tokens, effectively reducing redundant computations during inference. However, as memory overhead becomes a significant concern, efficient compression of KV cache has gained increasing attention. Most existing methods perform compression from two perspectives: identifying important tokens and designing compression strategies. However, these approaches often produce biased distributions of important tokens due to the influence of accumulated attention scores or positional encoding. Furthermore, they overlook the sparsity and redundancy across different heads, which leads to difficulties in preserving the most effective information at the head level. To this end, we propose EMS to overcome these limitations, while achieving better KV cache compression under extreme compression ratios. Specifically, we introduce a Global-Local score that combines accumulated attention scores from both global and local KV tokens to better identify the token importance. For the compression strategy, we design an adaptive and unified Evict-then-Merge framework that accounts for the sparsity and redundancy of KV tokens across different heads. Additionally, we implement the head-wise parallel compression through a zero-class mechanism to enhance efficiency. Extensive experiments demonstrate our SOTA performance even under extreme compression ratios. EMS consistently achieves the lowest perplexity, improves scores by over 1.28 points across four LLMs on LongBench under a 256 cache budget, and preserves 95% retrieval accuracy with a cache budget less than 2% of the context length in the Needle-in-a-Haystack task.

KV缓存压缩长文本推理注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。