论文
在查询所关注的地方花费位:具有注意力保留变换的 KV 缓存矢量量化
Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms
摘要
长上下文 LLM 解码在每一步都会读取键值 (KV) 缓存。加载它比计算它花费的时间更长,因此吞吐量受到带宽限制。因此,减小缓存大小可以提高解码速度和服务容量。面临的挑战是减少缓存大小,同时保留注意力产品,保持重建成本低廉,并使用固定的每个词元位数。每个元素两位,最具竞争力的方法依赖于正交变换。然而,现有技术要么忽略数据,要么使用查询统计数据而不从失真标准导出变换。此外,它们依赖于建立在随机或哈达玛旋转之上的变换,它均衡条目之间的方差而不是压缩能量,以及固定宽度标量量化器,这在低速率下不是最优的。在本文中,我们将 KV 缓存量化表述为变换编码问题,其中失真是注意力产品中的误差。我们在高分辨率模型下从校准统计数据中导出键和值的封闭形式最优变换。我们证明最佳密钥变换不是正交的,并且满足广义的 Parseval 关系:注意感知失真变成变换域中的均方误差(MSE)。因此,我们可以使用直接应用于变换后的关键系数的 MSE 最优矢量量化器。为了满足固定宽度布局要求,我们证明将系数分组为等体积分区使得相同大小的码本在相同的高分辨率模型下获得可变速率最优。在每个元素两位的情况下,我们的方法(称为 NOVA-KV)在可比较的吞吐量下恢复了标量量化方法损失的大部分长上下文检索精度。