Residual vector quantization for KV cache compression in large language model

📄 arXiv: 2410.15704v1 📥 PDF

作者: Ankur Kumar

分类: cs.LG

发布日期: 2024-10-21

💡 一句话要点

提出残差向量量化方法，用于压缩大语言模型中的KV缓存，实现高效推理。

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

关键词: 大语言模型 KV缓存 压缩 残差向量量化 模型量化

📋 核心要点

  1. 大语言模型解码时内存需求高，现有KV缓存压缩方法主要依赖标量量化，效果有限。
  2. 提出残差向量量化方法，通过缩放、分组和量化等步骤，有效压缩KV缓存，降低内存占用。
  3. 实验表明，残差深度为8时可恢复大部分性能，且非连续通道分组优于连续分组，压缩率达5.5倍。

📝 摘要（中文）

本文提出了一种基于残差向量量化（RVQ）的KV缓存压缩方法，用于降低大语言模型（LLM）解码过程中的内存需求。该方法借鉴了在高保真音频压缩中广泛应用的RVQ技术，并对其进行调整以压缩预训练LLM中任意键或值投影矩阵的输出。具体来说，该方法首先将向量按标准差缩放，然后将通道分成组，并使用相同的残差向量量化器量化每个组。码本通过指数移动平均学习，没有其他可学习的参数，包括向量量化设置中常用的输入和输出投影。实验结果表明，残差深度为8时，可以恢复未量化模型的大部分性能。此外，对于键矩阵的压缩，将非连续通道分组比连续通道分组效果更好。该方法还可以通过对LLM进行轻量级微调来进一步提升性能。总体而言，该技术与现有的量化方法相比具有竞争力，同时更加简单，并且可以实现5.5倍于半精度的压缩。

🔬 方法详解

问题定义：大语言模型（LLM）在解码过程中需要存储大量的Key和Value（KV）缓存，这导致了巨大的内存需求，尤其是在长序列生成时。现有的KV缓存压缩方法主要依赖于标量量化技术，这些方法在压缩率和性能之间往往难以取得平衡，并且可能引入显著的精度损失。因此，如何在保证模型性能的前提下，更有效地压缩KV缓存，是本文要解决的核心问题。

核心思路：本文的核心思路是将残差向量量化（RVQ）技术应用于KV缓存的压缩。RVQ是一种多阶段量化方法，通过逐级逼近原始向量，可以实现更高的压缩率和更好的重建质量。作者认为，KV缓存中的向量具有一定的结构性，可以通过RVQ进行有效压缩，从而降低内存需求。此外，作者还探索了不同的通道分组策略，并结合轻量级微调，进一步提升压缩性能。

技术框架：该方法主要包含以下几个步骤：1) 向量缩放：首先，将KV缓存中的向量按其标准差进行缩放，以减小向量的动态范围。2) 通道分组：然后，将向量的通道分成若干组，每组包含若干个通道。作者尝试了连续通道分组和非连续通道分组两种策略。3) 残差向量量化：对每个通道组使用相同的残差向量量化器进行量化。RVQ包含多个量化阶段，每个阶段都学习一个码本，并将输入向量量化到码本中最接近的码字。量化后的残差向量作为下一阶段的输入，直到达到预设的残差深度。4) 码本学习：码本通过指数移动平均（EMA）的方式进行学习，无需额外的可学习参数。

关键创新：该方法的主要创新点在于将残差向量量化技术应用于LLM的KV缓存压缩。与传统的标量量化方法相比，RVQ可以更好地捕捉向量之间的相关性，从而实现更高的压缩率和更好的性能。此外，作者还探索了不同的通道分组策略，并结合轻量级微调，进一步提升了压缩性能。该方法无需额外的可学习参数，实现简单，易于部署。

关键设计：在通道分组方面，作者发现对于键矩阵的压缩，非连续通道分组比连续通道分组效果更好。这可能是因为非连续通道之间可能存在更强的相关性。在残差深度方面，实验结果表明，残差深度为8时，可以恢复未量化模型的大部分性能。码本的学习率是EMA的关键参数，需要根据具体任务进行调整。此外，作者还发现，在量化后对LLM进行轻量级微调，可以进一步提升模型的性能。

📊 实验亮点

实验结果表明，该方法在压缩KV缓存的同时，能够保持较高的模型性能。具体来说，残差深度为8时，可以恢复未量化模型的大部分性能。与半精度相比，该方法可以实现5.5倍的压缩率。此外，通过对LLM进行轻量级微调，可以进一步提升模型的性能。该方法在多种LLM上进行了验证，证明了其通用性和有效性。

🎯 应用场景

该研究成果可广泛应用于各种需要部署大语言模型的场景，尤其是在资源受限的边缘设备或移动设备上。通过降低内存需求，可以支持更大规模的模型部署，提高推理速度，并降低能源消耗。此外，该技术还可以应用于云端推理服务，降低运营成本，提高服务质量。

📄 摘要（原文）

KV cache compression methods have mainly relied on scalar quantization techniques to reduce the memory requirements during decoding. In this work, we apply residual vector quantization, which has been widely used for high fidelity audio compression, to compress KV cache in large language models (LLM). We adapt the standard recipe with minimal changes to compress the output of any key or value projection matrix in a pretrained LLM: we scale the vector by its standard deviation, divide channels into groups and then quantize each group with the same residual vector quantizer. We learn the codebook using exponential moving average and there are no other learnable parameters including the input and output projections normally used in a vector quantization set up. We find that a residual depth of 8 recovers most of the performance of the unquantized model. We also find that grouping non-contiguous channels together works better than grouping contiguous channels for compressing key matrix and the method further benefits from a light weight finetuning of LLM together with the quantization. Overall, the proposed technique is competitive with existing quantization methods while being much simpler and results in 5.5x compression compared to half precision.