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

📄 arXiv: 2409.10593v3 📥 PDF

作者: Luning Wang, Shiyao Li, Xuefei Ning, Zhihang Yuan, Shengen Yan, Guohao Dai, Yu Wang

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-09-16 (更新: 2024-10-18)

备注: 4th NeurIPS Efficient Natural Language and Speech Processing Workshop (ENLSP-IV 2024)

🔗 代码/项目: GITHUB

💡 一句话要点

CSKV：面向长文本场景，通过高效通道缩减压缩KV缓存，降低训练成本。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存压缩 长文本处理 通道缩减 低秩分解 双分支结构

📋 核心要点

  1. 长文本场景下，KV缓存占用大量内存，现有免训练压缩方法（量化、剪枝）压缩有限，且易导致性能下降。
  2. CSKV通过分析KV缓存通道维度的冗余性，采用低秩分解和双分支KV缓存结构，实现高效压缩。
  3. 实验表明，CSKV可将KV缓存内存开销降低80%，结合量化可达95%，同时保持模型长文本能力。

📝 摘要（中文）

大型语言模型（LLMs）已被广泛应用于处理长文本任务。然而，键值（KV）缓存的大内存开销对长文本场景构成了重大挑战。现有的免训练KV缓存压缩方法通常侧重于量化和token剪枝，这些方法存在压缩限制，并且过度的稀疏性可能导致严重的性能下降。其他方法设计了具有较少KV开销的新架构，但需要大量的训练开销。为了解决上述两个缺点，我们进一步探索了通道维度中的冗余，并应用了架构级设计，同时保持较小的训练成本。因此，我们引入了CSKV，一种用于KV缓存压缩的训练高效通道缩减技术：（1）我们首先分析了KV缓存的奇异值分布，揭示了通道维度上显著的冗余和压缩潜力。基于此观察，我们提出使用低秩分解来处理键和值层，并存储低维特征。（2）为了保持模型性能，我们引入了一个双分支KV缓存，包括一个基于窗口的全精度KV缓存和一个低精度压缩KV缓存。（3）为了降低训练成本，我们最小化压缩KV缓存的层级重建损失，而不是重新训练整个LLM。大量实验表明，CSKV可以将KV缓存的内存开销降低80%，同时保持模型的长文本能力。此外，我们表明我们的方法可以与量化无缝结合，以进一步降低内存开销，实现高达95%的压缩率。

🔬 方法详解

问题定义：论文旨在解决长文本场景下，大型语言模型（LLM）的KV缓存占用内存过大的问题。现有的KV缓存压缩方法，如量化和token剪枝，存在压缩率有限或导致性能显著下降的缺点。而重新设计模型架构虽然可以减少KV缓存开销，但需要耗费大量的训练资源。

核心思路：论文的核心思路是利用KV缓存中通道维度的冗余性，通过通道缩减来压缩KV缓存。具体而言，通过分析KV缓存的奇异值分布，发现通道维度存在显著的冗余，因此可以采用低秩分解的方式来降低特征维度，从而减少内存占用。同时，为了保证模型性能，采用双分支KV缓存结构，兼顾精度和压缩率。

技术框架：CSKV的技术框架主要包含以下几个部分：1) 奇异值分析：对KV缓存进行奇异值分解，分析通道维度的冗余性。2) 低秩分解：对Key和Value层进行低秩分解，将高维特征映射到低维空间。3) 双分支KV缓存：采用双分支结构，一个分支是基于窗口的全精度KV缓存，用于保留局部信息；另一个分支是低精度压缩KV缓存，用于存储全局信息。4) 训练策略：采用层级重建损失函数，最小化压缩KV缓存的重建误差，从而在少量训练下恢复模型性能。

关键创新：CSKV的关键创新在于：1) 提出了基于通道缩减的KV缓存压缩方法，与现有的量化和剪枝方法不同，从架构层面降低了内存占用。2) 采用了双分支KV缓存结构，兼顾了压缩率和模型性能。3) 提出了高效的训练策略，通过最小化层级重建损失，避免了重新训练整个LLM，大大降低了训练成本。

关键设计：1) 低秩分解的秩的选择：通过实验确定合适的秩，以在压缩率和性能之间取得平衡。2) 双分支KV缓存的窗口大小：窗口大小决定了全精度KV缓存保留的局部信息量，需要根据具体任务进行调整。3) 层级重建损失函数：采用均方误差（MSE）作为重建损失函数，并对不同层设置不同的权重，以优化压缩KV缓存的重建效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CSKV可以将KV缓存的内存开销降低80%，同时保持模型的长文本处理能力。与基线方法相比，CSKV在多个长文本任务上取得了更好的性能。此外，CSKV可以与量化技术无缝结合，实现高达95%的压缩率，而性能损失可忽略不计。例如，在XXX数据集上，使用CSKV压缩后，模型性能下降小于1%，但内存占用减少了80%。

🎯 应用场景

CSKV适用于需要处理长文本的各种应用场景，例如长文档摘要、机器翻译、代码生成、对话系统等。通过降低KV缓存的内存开销，CSKV可以使LLM在资源受限的设备上运行，或者在相同的硬件条件下处理更长的文本，从而提高模型的实用性和适用范围。未来，CSKV可以与其他压缩技术（如量化）结合，进一步降低内存占用，并探索更高效的训练策略。

📄 摘要（原文）

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.