RotateKV: Accurate and Robust 2-Bit KV Cache Quantization for LLMs via Outlier-Aware Adaptive Rotations

📄 arXiv: 2501.16383v2 📥 PDF

作者: Zunhai Su, Zhe Chen, Wang Shen, Hanyu Wei, Linge Li, Huangqi Yu, Kehong Yuan

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

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

💡 一句话要点

RotateKV：通过自适应旋转实现LLM的精确鲁棒的2比特KV缓存量化

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

关键词: KV缓存量化 低比特量化 大型语言模型 旋转变换 异常值感知

📋 核心要点

  1. 现有KV缓存量化方法在极低比特宽度下难以兼顾压缩率和鲁棒性，成为LLM部署的瓶颈。
  2. RotateKV通过异常值感知旋转、预RoPE分组头旋转和注意力Sink感知量化，实现高效的2比特KV缓存量化。
  3. 实验表明，RotateKV在多种任务上优于现有方法，显著降低内存占用并提升推理速度。

📝 摘要（中文）

Key-Value (KV) 缓存通过避免重复计算过去的KVs，从而促进高效的大型语言模型 (LLMs) 推理。随着批次大小和上下文长度的增加，过大的KV缓存成为显著的内存瓶颈，突显了高效压缩的必要性。现有的KV量化依赖于细粒度量化或保留大部分高比特宽度缓存，这两种方法都会降低压缩率，并且通常无法在极低的平均比特宽度下保持鲁棒性。在这项工作中，我们探索了旋转技术在2比特KV量化中的潜力，并提出了RotateKV，它通过以下创新实现了准确和鲁棒的性能：(i) 异常值感知旋转，它利用通道重排序来使旋转适应不同的通道级异常值分布，而不会牺牲快速Walsh-Hadamard变换 (FWHT) 的计算效率；(ii) 预RoPE分组头旋转，它减轻了旋转位置嵌入 (RoPE) 对所提出的异常值感知旋转的影响，并进一步平滑了跨头的异常值；(iii) 注意力Sink感知量化，它利用大量的激活来精确识别和保护注意力sink。RotateKV在使用LLaMA-2-13B的WikiText-2上实现了小于0.3的困惑度 (PPL) 降低，保持了强大的CoT推理和长上下文能力，在GSM8K上的降幅小于1.7%，优于现有的方法，即使在更低的平均比特宽度下也是如此。RotateKV还展示了峰值内存使用量减少了3.97倍，支持5.75倍更大的批次大小，并在解码阶段实现了2.32倍的加速。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）中KV缓存过大导致的内存瓶颈问题。现有的KV缓存量化方法，如细粒度量化或保留高比特宽度缓存，在极低比特宽度下无法同时保证压缩率和模型性能的鲁棒性，尤其是在长文本和复杂推理任务中表现不佳。

核心思路：论文的核心思路是利用旋转变换来优化KV缓存的量化过程，使其更适应数据中的异常值分布，并减轻旋转位置编码（RoPE）的影响。通过自适应的旋转，可以将KV缓存中的信息更紧凑地编码到低比特表示中，从而在保证模型性能的同时，显著降低内存占用。

技术框架：RotateKV的技术框架主要包含三个阶段：(1) 异常值感知旋转：对KV缓存的通道进行重排序，使旋转变换更适应通道级的异常值分布，同时保持计算效率。(2) 预RoPE分组头旋转：在应用旋转变换之前，先对RoPE进行处理，减轻其对旋转的影响，并平滑跨头的异常值。(3) 注意力Sink感知量化：利用大量的激活信息，精确识别并保护注意力机制中的重要部分（attention sinks），避免量化带来的性能损失。

关键创新：RotateKV的关键创新在于其旋转变换的自适应性和对RoPE的特殊处理。传统的旋转变换可能无法有效处理KV缓存中存在的异常值，而RotateKV通过通道重排序和预RoPE处理，使得旋转变换能够更好地适应数据的分布，从而提高量化精度和鲁棒性。此外，注意力Sink感知量化进一步提升了模型在关键信息上的保留能力。

关键设计：在异常值感知旋转中，论文采用了快速Walsh-Hadamard变换（FWHT）以保证计算效率。预RoPE分组头旋转的具体实现方式未知，但其目标是减少RoPE对后续旋转的影响。注意力Sink感知量化通过分析激活值来确定需要保护的注意力sink，并采用更高的比特宽度进行量化。具体的损失函数和网络结构细节在论文中可能有所描述，但摘要中未明确提及。

📊 实验亮点

RotateKV在LLaMA-2-13B模型上，使用2比特量化在WikiText-2数据集上实现了小于0.3的困惑度(PPL)下降。在GSM8K数据集上，CoT推理的性能下降小于1.7%。同时，峰值内存使用量减少了3.97倍，支持5.75倍更大的批次大小，并在解码阶段实现了2.32倍的加速，显著优于现有方法。

🎯 应用场景

RotateKV可应用于各种需要部署大型语言模型的场景，如移动设备、边缘计算和资源受限的服务器。通过降低内存占用和提高推理速度，RotateKV能够使LLM在更多平台上运行，并支持更大的批次大小和更长的上下文长度，从而提升用户体验和应用范围。

📄 摘要（原文）

Key-Value (KV) cache facilitates efficient large language models (LLMs) inference by avoiding recomputation of past KVs. As the batch size and context length increase, the oversized KV caches become a significant memory bottleneck, highlighting the need for efficient compression. Existing KV quantization rely on fine-grained quantization or the retention of a significant portion of high bit-widths caches, both of which compromise compression ratio and often fail to maintain robustness at extremely low average bit-widths. In this work, we explore the potential of rotation technique for 2-bit KV quantization and propose RotateKV, which achieves accurate and robust performance through the following innovations: (i) Outlier-Aware Rotation, which utilizes channel-reordering to adapt the rotations to varying channel-wise outlier distributions without sacrificing the computational efficiency of the fast Walsh-Hadamard transform (FWHT); (ii) Pre-RoPE Grouped-Head Rotation, which mitigates the impact of rotary position embedding (RoPE) on proposed outlier-aware rotation and further smooths outliers across heads; (iii) Attention-Sink-Aware Quantization, which leverages the massive activations to precisely identify and protect attention sinks. RotateKV achieves less than 0.3 perplexity (PPL) degradation with 2-bit quantization on WikiText-2 using LLaMA-2-13B, maintains strong CoT reasoning and long-context capabilities, with less than 1.7\% degradation on GSM8K, outperforming existing methods even at lower average bit-widths. RotateKV also showcases a 3.97x reduction in peak memory usage, supports 5.75x larger batch sizes, and achieves a 2.32x speedup in decoding stage.