arXiv:2409.10593cs.LGcs.AI2024-09被引 2

通过通道压缩降低大模型长文本缓存内存占用,训练成本低且效果好。

CSKV: Training-Efficient Channel Shrinking for KV Cache in Long-Context Scenarios

  • 利用低秩分解压缩键值通道,挖掘通道维度冗余。
  • 80%内存压缩率下保持长文本处理能力,结合量化可达95%压缩比。
  • 仅需少量微调,适合部署在资源受限的长文本场景。

大语言模型在处理长文本任务时面临键值(KV)缓存内存开销大的挑战。现有无需训练的压缩方法多依赖量化和删减令牌,存在压缩上限,过度稀疏会严重损害性能;而新架构虽降低缓存开销,但需大量训练成本。为此,本文探索通道维度冗余,提出一种训练高效的通道压缩方法CSKV:(1) 分析KV缓存奇异值分布,发现通道维度存在显著冗余,提出对键值层进行低秩分解并存储低维特征;(2) 设计双分支缓存结构,包含基于窗口的全精度缓存与低精度压缩缓存,以保留模型性能;(3) 仅最小化压缩缓存的逐层重建损失,避免重训整个模型。大量实验表明,CSKV可将KV缓存内存减少80%,同时保持长文本能力;还可无缝融合量化技术,实现最高达95%的压缩比。代码已开源。

原文摘要 · Abstract (English)

Large Language Models (LLMs) have been widely adopted to process long-context tasks. However, the large memory overhead of the key-value (KV) cache poses significant challenges in long-context scenarios. Existing training-free KV cache compression methods typically focus on quantization and token pruning, which have compression limits, and excessive sparsity can lead to severe performance degradation. Other methods design new architectures with less KV overhead but require significant training overhead. To address the above two drawbacks, we further explore the redundancy in the channel dimension and apply an architecture-level design with minor training costs. Therefore, we introduce CSKV, a training-efficient Channel Shrinking technique for KV cache compression: (1) We first analyze the singular value distribution of the KV cache, revealing significant redundancy and compression potential along the channel dimension. Based on this observation, we propose using low-rank decomposition for key and value layers and storing the low-dimension features. (2) To preserve model performance, we introduce a bi-branch KV cache, including a window-based full-precision KV cache and a low-precision compressed KV cache. (3) To reduce the training costs, we minimize the layer-wise reconstruction loss for the compressed KV cache instead of retraining the entire LLMs. Extensive experiments show that CSKV can reduce the memory overhead of the KV cache by 80% while maintaining the model's long-context capability. Moreover, we show that our method can be seamlessly combined with quantization to further reduce the memory overhead, achieving a compression ratio of up to 95%. Code is available at https://github.com/wln20/CSKV.

KV缓存通道压缩长文本低秩分解

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