用可交换向量量化压缩注意力缓存，1比特也能保持高精度。
CommVQ: Commutative Vector Quantization for KV Cache Compression
- 设计与旋转位置编码兼容的量化码本，支持快速解码。
- 2比特量化下内存减少87.5%，1比特时仍保持低精度损失。
- 适合在单卡上运行超长上下文的大模型，如LLaMA-3.1 8B。
大型语言模型在长上下文应用中面临显存瓶颈，尤其是键值（KV）缓存占用过大。本文提出共性向量量化（CommVQ），通过轻量级编码器和可交换码本对KV缓存进行加性量化，解码仅需矩阵乘法。设计码本与旋转位置编码（RoPE）共性，并使用期望最大化（EM）算法训练，使解码可无缝融入自注意力机制。实验表明，在长上下文基准和GSM8K测试中，2比特量化将FP16 KV缓存压缩87.5%，优于现有方法；更实现1比特量化下极小精度损失，使LLaMA-3.1 8B模型在单张RTX 4090 GPU上支持128K上下文长度。代码已开源。
原文摘要 · Abstract (English)
Large Language Models (LLMs) are increasingly used in applications requiring long context lengths, but the key-value (KV) cache often becomes a memory bottleneck on GPUs as context grows. To address this, we propose Commutative Vector Quantization (CommVQ) to significantly reduce memory usage for long-context LLM inference. We first introduce additive quantization with a lightweight encoder and codebook to compress the KV cache, which can be decoded via simple matrix multiplication. To further reduce computational costs during decoding, we design the codebook to be commutative with Rotary Position Embedding (RoPE) and train it using an Expectation-Maximization (EM) algorithm. This enables efficient integration of decoding into the self-attention mechanism. Our approach achieves high accuracy with additive quantization and low overhead via the RoPE-commutative codebook. Experiments on long-context benchmarks and GSM8K show that our method reduces FP16 KV cache size by 87.5% with 2-bit quantization, while outperforming state-of-the-art KV cache quantization methods. Notably, it enables 1-bit KV cache quantization with minimal accuracy loss, allowing a LLaMA-3.1 8B model to run with a 128K context length on a single RTX 4090 GPU. The source code is available at: https://github.com/UMass-Embodied-AGI/CommVQ.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。