按层不确定度动态分配缓存预算，显著降低长文本推理内存占用
ZigZagkv: Dynamic KV Cache Compression for Long-context Modeling based on Layer Uncertainty
- 基于注意力与隐藏状态的层间不确定性，动态分配每层缓存预算
- 仅需全量缓存20%内存，性能近乎无损
- 适合需要长上下文推理的部署场景，如文档生成、代码补全
大语言模型（LLMs）已成为研究热点。为加速推理，将计算缓存存储于内存已成为标准做法。然而，随着推理长度增加，不断增长的键值（KV）缓存可能导致内存不足。现有方法通过压缩KV缓存缓解此问题，通常在所有层中保留关键令牌以减少信息损失，并为每层分配统一缓存预算。但本文观察到，不同层和模型所需的最小缓存预算大小存在差异，这源于注意力机制和隐藏状态输出的层间变化。基于此，提出一种简单有效的动态压缩方法，利用层不确定性自适应分配每层缓存预算。实验表明，该方法在保持近似无损性能的前提下，将KV缓存内存使用降至全量推理的约20%。
原文摘要 · Abstract (English)
Large Language models (LLMs) have become a research hotspot. To accelerate the inference of LLMs, storing computed caches in memory has become the standard technique. However, as the inference length increases, growing KV caches might lead to out-of-memory issues. Many existing methods address this issue through KV cache compression, primarily by preserving key tokens throughout all layers to reduce information loss. Most of them allocate a uniform budget size for each layer to retain. However, we observe that the minimum budget sizes needed to retain essential information vary across layers and models based on the perspectives of attention and hidden state output. Building on this observation, this paper proposes a simple yet effective KV cache compression method that leverages layer uncertainty to allocate budget size for each layer. Experimental results show that the proposed method can reduce memory usage of the KV caches to only $\sim$20\% when compared to Full KV inference while achieving nearly lossless performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。