零样本压缩长文本缓存,让大模型高效处理超长内容
ZSMerge: Zero-Shot KV Cache Compression for Memory-Efficient Long-Context LLMs
- 按注意力头粒度动态分配内存,基于多维重要性评估精简缓存
- 通过补偿注意力评分保留关键信息,实现20:1压缩比且质量接近原版
- 无需重训练即可适配多种模型,适合需要长上下文的部署场景
键值(KV)缓存内存线性增长与注意力机制二次计算复杂度,是大语言模型在长文本处理中的主要瓶颈。现有优化方法通过剪枝或特征合并缓解问题,但常导致不可逆信息丢失或需昂贵参数重训。为此,我们提出ZSMerge,一种面向高效缓存管理的动态压缩框架,包含三项核心操作:(1) 基于头级别粒度的多维标记重要性指标进行细粒度内存分配;(2) 通过补偿注意力评分保留关键上下文的残差合并机制;(3) 兼容多种模型架构的零样本适应机制,无需重训练。在LLaMA2-7B上,该方法实现20:1的键值缓存压缩比(内存降至基线5%),维持相近生成质量,并在极端54k token上下文下提升三倍吞吐量,彻底避免内存溢出。代码已开源。
原文摘要 · Abstract (English)
The linear growth of key-value (KV) cache memory and quadratic computational in attention mechanisms complexity pose significant bottlenecks for large language models (LLMs) in long-context processing. While existing KV cache optimization methods address these challenges through token pruning or feature merging, they often incur irreversible information loss or require costly parameter retraining. To this end, we propose ZSMerge, a dynamic KV cache compression framework designed for efficient cache management, featuring three key operations: (1) fine-grained memory allocation guided by multi-dimensional token importance metrics at head-level granularity, (2) a residual merging mechanism that preserves critical context through compensated attention scoring, and (3) a zero-shot adaptation mechanism compatible with diverse LLM architectures without requiring retraining. ZSMerge significantly enhances memory efficiency and inference speed with negligible performance degradation across LLMs. When applied to LLaMA2-7B, it demonstrates a 20:1 compression ratio for key-value cache retention (reducing memory footprint to 5\% of baseline) while sustaining comparable generation quality, coupled with triple throughput gains at extreme 54k-token contexts that eliminate out-of-memory failures. The code is available at https://github.com/SusCom-Lab/ZSMerge.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。