通过注意力感知的变换优化KV缓存压缩,提升长文本生成速度。
Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms

- 基于校准统计推导出最优非正交变换,使注意力误差转为可量化均方误差
- 在每元素2比特条件下,恢复了标量量化损失的大部分长上下文精度
- 适用于高吞吐量大模型推理服务,尤其适合资源受限部署
长上下文大模型解码需在每步读取键值(KV)缓存,其加载时间超过注意力计算时间,导致吞吐受带宽限制。减少缓存大小可提升解码速度与服务容量。关键挑战在于:在保持注意力结果、重建成本低且每标记固定比特数的前提下压缩缓存。现有方法在每元素2比特时依赖正交变换,但多为数据无关或仅使用查询统计而未从失真准则推导变换;且多基于随机或哈达玛旋转,均衡各分量方差而非集中能量,同时采用固定宽度标量量化器,在低比特率下性能不佳。本文将KV缓存量化建模为变换编码问题,以注意力产物误差为失真度量。在高分辨率模型下,推导出键和值的闭式最优变换。结果显示,最优键变换非正交,满足广义帕塞瓦尔关系:注意力感知失真转化为变换域的均方误差。因此可在变换域直接应用均方误差最优向量量化器。为满足固定宽度布局要求,证明将系数分组为等体积块,可使等大小码本达到变率最优。在每元素2比特下,所提方法NOVA-KV在相当吞吐下恢复了标量量化损失的大部分长上下文检索准确率。
原文摘要 · Abstract (English)
Long-context LLM decoding reads the key-value (KV) cache at every step. Loading it takes longer than computing attention over it, so throughput is bandwidth-bound. Hence, reducing the cache size can raise both decoding speed and serving capacity. The challenge is to reduce cache size while preserving the attention products, keeping reconstruction cheap, and using a fixed per-token bit count. At two bits per element, the most competitive methods rely on orthogonal transforms. However, existing techniques are either data-oblivious or use the query statistics without deriving the transform from a distortion criterion. Moreover, they rely on transforms built on top of random or Hadamard rotations, which equalize variances across entries rather than compacting energy, and fixed-width scalar quantizers, which are suboptimal at low rates. In this paper, we formulate KV cache quantization as a transform coding problem in which distortion is the error in the attention products. We derive closed-form optimal transforms for keys and values from calibration statistics, under a high-resolution model. We show that the optimal key transform is not orthogonal and satisfies a generalized Parseval relation: the attention-aware distortion becomes mean-squared error (MSE) in the transform domain. Thus, we can use MSE-optimal vector quantizers applied directly to the transformed key coefficients. To meet the fixed-width layout requirement, we show that grouping coefficients into equal-volume partitions makes equal-size codebooks attain the variable-rate optimum under the same high-resolution model. At two bits per element, our method, termed NOVA-KV, recovers most of the long-context retrieval accuracy lost by scalar quantization methods at comparable throughput.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。