arXivDaily arXiv每日学术速递 周一至周五更新
​通知:联邦节假日调休:哥伦布日 / 原住民日(Columbus Day / Indigenous Peoples' Day),今日(13日)无更新。
arXiv 2503.10714cs.CLcs.AI

ZSMerge:面向内存高效长上下文大语言模型的零样本KV缓存压缩方法

ZSMerge: Zero-Shot KV Cache Compression for Memory-Efficient Long-Context LLMs

Information Hub, The Hong Kong University of Science and Technology (Guangzhou)(香港科学与技术大学信息中心);School of Data Science, The Chinese University of Hong Kong, Shenzhen(香港中文大学(深圳)数据科学学院)

Xin Liu, Xudong Wang, Pei Liu, Guoming Tang

AI总结:

针对大语言模型长上下文处理中KV缓存内存增长与注意力计算的瓶颈,该研究提出无需重训练的ZSMerge动态压缩框架,在大幅降低内存占用、提升推理速度的同时维持了生成质量。

摘要:

键值(KV)缓存内存的线性增长以及注意力机制的二次计算复杂度,给大语言模型(LLM)的长上下文处理带来了显著瓶颈。现有KV缓存优化方法虽通过令牌剪枝或特征合并应对这些挑战,但往往会造成不可逆的信息损失,或需要成本高昂的参数重训练。为此,我们提出ZSMerge——一种专为高效缓存管理设计的动态KV缓存压缩框架,具备三项核心操作:(1)以头级粒度、由多维度令牌重要性指标引导的细粒度内存分配;(2)通过补偿注意力评分保留关键上下文的残差合并机制;(3)无需重训练即可适配多种LLM架构的零样本适配机制。ZSMerge在各类LLM上的性能下降可忽略不计,同时显著提升了内存效率与推理速度。将其应用于LLaMA2-7B时,该方法实现了20:1的KV缓存保留压缩比(将内存占用降至基线的5%),同时维持了相当的生成质量;在5.4万令牌的极端上下文场景下,吞吐量提升至原来的3倍,消除了内存不足故障。代码可在https://github.com/SusCom-Lab/ZSMerge获取。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑