Cache Me If You Must: Adaptive Key-Value Quantization for Large Language Models

📄 arXiv: 2501.19392v4 📥 PDF

作者: Alina Shutova, Vladimir Malinovskii, Vage Egiazarian, Denis Kuznedelev, Denis Mazur, Nikita Surkov, Ivan Ermakov, Dan Alistarh

分类: cs.LG

发布日期: 2025-01-31 (更新: 2025-02-28)

备注: Preprint, under review

💡 一句话要点

AQUA-KV：自适应键值量化，提升大语言模型KV缓存压缩率并保持精度

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

关键词: 键值缓存 量化 大语言模型 模型压缩 自适应量化

📋 核心要点

  1. 大语言模型KV缓存占用大量内存，现有压缩方法在提高压缩率的同时通常会降低模型精度。
  2. AQUA-KV利用层间Key和Value的依赖关系，通过紧凑的适配器预测信息，并对无法预测的信息进行最优压缩。
  3. 实验表明，AQUA-KV在Llama 3.2模型上实现了接近无损的推理，同时显著提高了压缩率。

📝 摘要（中文）

大型语言模型(LLM)的实际高效部署依赖于键值(KV)缓存，以处理和生成长输出，从而减少重复计算的需求。对于大型上下文，KV缓存会占用数十GB的设备内存，因为它们存储每个token和每一层的向量表示。最近的研究表明，可以通过量化、剪枝或合并来压缩缓存的向量，但这些技术通常会以牺牲质量为代价来提高压缩率。本文旨在通过利用两个观察结果来改进键和值的压缩：1)不同层之间键和值之间的内在依赖性，以及2)内部网络状态的高压缩机制。我们提出了AQUA-KV，一种用于键值缓存的自适应量化方法，它依赖于紧凑的适配器来利用键和值之间现有的依赖关系，并旨在“最优地”压缩无法预测的信息。AQUA-KV显著提高了压缩率，同时保持了最先进的LLM系列的高精度。在Llama 3.2 LLM上，我们以每个值2-2.5比特的近乎无损的推理实现了低于1%的困惑度和LongBench分数相对误差。AQUA-KV是一次性的、简单的和高效的：即使对于70B模型，它也可以在单个GPU上在1-6小时内完成校准。

🔬 方法详解

问题定义：大语言模型在处理长文本时，需要缓存每一层每个token的Key和Value向量，导致KV缓存占用大量内存。现有的量化、剪枝等压缩方法虽然可以降低内存占用，但往往会牺牲模型的生成质量，例如困惑度升高、长文本评估指标下降。

核心思路：AQUA-KV的核心思路是利用不同层之间Key和Value向量的内在依赖关系，通过学习一个小型适配器网络来预测这些向量，从而只需要存储预测残差，降低存储空间。对于无法预测的部分，则采用自适应量化方法进行压缩。

技术框架：AQUA-KV包含以下主要步骤：1) 训练小型适配器网络，用于预测每一层的Key和Value向量。适配器的输入是前几层的Key和Value向量，输出是当前层的预测向量。2) 计算预测向量与原始向量之间的残差。3) 对残差向量进行自适应量化，根据残差的分布动态调整量化参数，以最小化量化误差。4) 存储量化后的残差向量和适配器网络的参数。在推理时，首先使用适配器网络预测Key和Value向量，然后将预测向量与量化后的残差向量相加，得到最终的Key和Value向量。

关键创新：AQUA-KV的关键创新在于：1) 利用层间依赖关系，通过适配器网络预测Key和Value向量，减少需要存储的信息量。2) 采用自适应量化方法，根据残差的分布动态调整量化参数，实现更高的压缩率，同时保持精度。3) AQUA-KV 是一种 one-shot 方法，只需要少量数据进行校准，无需重新训练整个模型。

关键设计：适配器网络采用小型MLP结构，输入是前几层的Key和Value向量，输出是当前层的预测向量。损失函数采用均方误差，用于最小化预测向量与原始向量之间的差异。自适应量化采用分位数量化方法，根据残差的分布将残差映射到不同的量化级别。量化级别的数量和边界由残差的分布决定。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

AQUA-KV在Llama 3.2系列模型上取得了显著的性能提升。在2-2.5比特/值的压缩率下，困惑度的相对误差低于1%，LongBench的评分误差也低于1%。与现有方法相比，AQUA-KV在保持精度的同时，实现了更高的压缩率，并且校准过程简单高效，即使对于70B模型，也只需要1-6小时。

🎯 应用场景

AQUA-KV可应用于各种需要高效部署大型语言模型的场景，例如移动设备、边缘计算设备和资源受限的服务器。通过降低KV缓存的内存占用，AQUA-KV可以使这些设备能够运行更大规模的模型，或者在相同的硬件资源下处理更长的上下文，从而提升用户体验。

📄 摘要（原文）

Efficient real-world deployments of large language models (LLMs) rely on Key-Value (KV) caching for processing and generating long outputs, reducing the need for repetitive computation. For large contexts, Key-Value caches can take up tens of gigabytes of device memory, as they store vector representations for each token and layer. Recent work has shown that the cached vectors can be compressed through quantization, pruning or merging, but these techniques often compromise quality towards higher compression rates. In this work, we aim to improve Key & Value compression by exploiting two observations: 1) the inherent dependencies between keys and values across different layers, and 2) high-compression mechanisms for internal network states. We propose AQUA-KV, an adaptive quantization for Key-Value caches that relies on compact adapters to exploit existing dependencies between Keys and Values, and aims to "optimally" compress the information that cannot be predicted. AQUA-KV significantly improves compression rates, while maintaining high accuracy on state-of-the-art LLM families. On Llama 3.2 LLMs, we achieve near-lossless inference at 2-2.5 bits per value with under $1\%$ relative error in perplexity and LongBench scores. AQUA-KV is one-shot, simple, and efficient: it can be calibrated on a single GPU within 1-6 hours, even for 70B models.