Lossless KV Cache Compression to 2%

📄 arXiv: 2410.15252v1 📥 PDF

作者: Zhen Yang, J. N. Han, Kan Wu, Ruobing Xie, An Wang, Xingwu Sun, Zhanhui Kang

分类: cs.CL, cs.AI

发布日期: 2024-10-20

💡 一句话要点

提出CLLA架构，将KV缓存无损压缩至2%，显著提升LLM推理效率

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

关键词: KV缓存压缩 大型语言模型 推理加速 跨层注意力 量化

📋 核心要点

  1. 大型语言模型推理加速依赖KV缓存，但其内存需求日益增长，成为高效部署的瓶颈。
  2. 论文提出跨层潜在注意力（CLLA）架构，融合多种压缩技术，旨在大幅降低KV缓存大小。
  3. 实验表明，CLLA能将KV缓存压缩至原始大小的2%以下，并在多数任务上保持无损性能。

📝 摘要（中文）

大型语言模型（LLM）彻底改变了众多领域的数据处理方式，尤其是在处理扩展上下文推理方面的能力备受认可。为了加速推理过程，维护键值（KV）缓存至关重要。然而，对KV缓存内存日益增长的需求给高效实现带来了重大障碍。本研究提出了一种新颖的架构，即跨层潜在注意力（CLLA），旨在将KV缓存压缩到原始大小的2%以下，同时保持相当的性能水平。CLLA将KV缓存压缩的多个方面，包括注意力头/维度缩减、层共享和量化技术，整合到一个统一的框架中。大量的实验表明，CLLA在利用极小KV缓存的同时，在大多数任务上实现了无损性能，标志着在实际KV缓存压缩方面取得了重大进展。

🔬 方法详解

问题定义：大型语言模型在推理过程中需要维护一个KV缓存来存储先前层的激活值，以便在后续的token生成中使用。这个KV缓存的大小随着上下文长度的增加而线性增长，导致内存占用过高，成为部署和扩展LLM的瓶颈。现有的KV缓存压缩方法通常会引入性能损失，或者压缩率不够高。

核心思路：论文的核心思路是通过跨层共享信息、减少注意力头/维度以及使用量化技术，在尽可能不损失性能的前提下，大幅降低KV缓存的大小。CLLA架构旨在将多种压缩技术整合到一个统一的框架中，实现更高的压缩率和更好的性能。

技术框架：CLLA架构包含以下几个主要模块：1) 注意力头/维度缩减：通过减少每个注意力头或维度的数量来降低KV缓存的大小。2) 层共享：在不同的Transformer层之间共享KV缓存，减少冗余信息。3) 量化：使用低精度表示（如INT8或INT4）来存储KV缓存，进一步降低内存占用。CLLA将这些模块集成到一个统一的框架中，并通过联合优化来提高整体性能。

关键创新：CLLA的关键创新在于其跨层潜在注意力机制，该机制允许不同层之间共享和传递信息，从而减少了KV缓存中的冗余信息。此外，CLLA还通过联合优化注意力头/维度缩减、层共享和量化等多种压缩技术，实现了更高的压缩率和更好的性能。与现有方法相比，CLLA能够在保持无损性能的同时，将KV缓存压缩到更小的尺寸。

关键设计：CLLA的具体实现细节包括：1) 使用可学习的线性变换来减少注意力头/维度的数量。2) 设计了一种层共享策略，允许相邻层之间共享KV缓存。3) 使用混合精度量化技术，对不同的KV缓存部分使用不同的量化精度。4) 使用知识蒸馏技术来弥补压缩带来的性能损失。具体的参数设置和损失函数需要在实际应用中进行调整和优化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CLLA架构能够将KV缓存压缩到原始大小的2%以下，同时在大多数任务上保持无损性能。例如，在某些文本生成任务上，CLLA的性能与未压缩的模型相当，甚至略有提升。与现有的KV缓存压缩方法相比，CLLA在压缩率和性能方面都取得了显著的提升。

🎯 应用场景

CLLA架构在资源受限的设备上部署大型语言模型具有广泛的应用前景，例如移动设备、边缘计算设备等。该技术可以降低LLM的内存占用，使其能够在这些设备上运行，从而实现更广泛的AI应用。此外，CLLA还可以用于加速LLM的推理速度，提高用户体验。未来，CLLA有望成为LLM部署和应用的关键技术之一。

📄 摘要（原文）

Large language models have revolutionized data processing in numerous domains, with their ability to handle extended context reasoning receiving notable recognition. To speed up inference, maintaining a key-value (KV) cache memory is essential. Nonetheless, the growing demands for KV cache memory create significant hurdles for efficient implementation. This work introduces a novel architecture, Cross-Layer Latent Attention (CLLA), aimed at compressing the KV cache to less than 2% of its original size while maintaining comparable performance levels. CLLA integrates multiple aspects of KV cache compression, including attention head/dimension reduction, layer sharing, and quantization techniques, into a cohesive framework. Our extensive experiments demonstrate that CLLA achieves lossless performance on most tasks while utilizing minimal KV cache, marking a significant advancement in practical KV cache compression.