提出注意力感知的压缩方法,实现5.8倍压缩下几乎无损精度。
KV Cache Compression Through the Lens of Transform Coding
- 基于信号处理理论,按注意力机制需求分配比特,而非均匀量化。
- 在多个基准测试中,实现约5.8倍压缩且精度几乎不降。
- 适合长序列推理场景,尤其对内存受限部署有显著价值。
键值(KV)缓存存储历史标记信息,是长序列推理中的主要内存瓶颈。现有量化方法通过低精度数据类型统一表示缓存,并设计量化方案以最小化缓存本身的重建误差,但未考虑误差在注意力机制中的传播。我们证明,在白噪声量化模型下,期望的注意力感知失真可分解为键与值的可加贡献,且在标记和通道间可分离。基于变换编码与反向水填法(经典信号处理与率失真理论工具),我们提出注意力感知变换编码(AATC),在校准集上分配比特以最小化注意力感知失真。在 Llama-3.1-8B-Instruct 和 Qwen-2.5-7B-Instruct 上,覆盖 LongBench、RULER、GSM8K、MMLU-Pro 与 MATH-500 多个数据集,本方法在约 5.8 倍压缩率下实现近似无损精度,而所有基线方法在至少一个设置中性能下降。
原文摘要 · Abstract (English)
The key-value (KV) cache stores information from past tokens and is a major memory bottleneck in long-context inference. Existing quantization methods address this bottleneck by representing the KV cache uniformly with lower-precision data types and designing quantization schemes to minimize reconstruction error in the cache itself, without accounting for how that error propagates through attention mechanisms. We prove that, under a white-noise quantization model, the expected attention-aware distortion decomposes into additive key and value contributions that factor across tokens and channels. Building on transform coding and reverse water-filling, which are classical tools from signal processing and rate-distortion theory, we introduce Attention-Aware Transform Coding (AATC), which allocates bits over a calibration set to minimize attention-aware distortion. On Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct, evaluated across LongBench, RULER, GSM8K, MMLU-Pro, and MATH-500, our method achieves near-lossless accuracy at approximately $5.8\times$ compression, whereas each baseline degrades in at least some settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。