SAKI通过优化注意力得分直接提升键值缓存效率,显著减少推理误差。
SAKI: Score-Aware Low-Rank Key Indexing with Random-Matrix Noise Correction for KV Retrieval

- 基于协方差加权的奇异值分解,直接保留注意力得分而非键重建质量
- 在32阶压缩下,使LLaMA 3.1和Qwen 2.5的顶部64召回率分别提升至0.799和0.850
- 无需训练,仅用512个校准令牌即可识别真实信号,适合深层模型优化
现有低秩键值缓存方法要么保留模型权重,要么保留键的方差,均未直接反映推理中使用的注意力得分。本文推导了秩为r的键压缩所引起的期望注意力得分失真,并证明其等价于协方差加权的低秩目标。在边界条件下,控制该失真可提升前k召回率。最优秩r解具有闭式非对称分解,由协方差加权查询-键算子的SVD得到。这启发我们提出SAKI:一种无需训练的键值缓存索引,直接保留注意力得分。在LLaMA 3.1 8B、Qwen 2.5 7B、Mistral 7B v0.1和Llama 3.2 3B上,SAKI在所有测试秩下均优于键主成分分析(PCA)。在秩32时,其将PCA的剩余顶部64召回误差减少13%至30%,包括在LLaMA 3.1 8B上从0.748提升至0.799,以及在Qwen 2.5 7B上从0.786提升至0.850。它改善了每模型68%至89%的注意力头,深层层收益最大。预测得分均方误差降低与实测结果高度一致,皮尔逊相关系数达0.997;消融实验表明增益源于优化注意力得分目标,而非协方差加权本身。对评分算子的分析进一步解释为何仅保留权重、不变子空间或键重建的方法可能次优。SAKI利用随机矩阵理论分离真实协方差信号与自相关采样噪声,在仅512个校准令牌下匹配PCA性能,并在PCA无法检测可靠信号处仍具价值。
原文摘要 · Abstract (English)
Existing low rank KV cache methods preserve either model weights or key variance, neither of which directly reflects the attention scores used during inference. We derive the expected attention score distortion caused by rank r key compression and show that it yields a covariance weighted low rank objective. Under a margin condition, controlling this distortion also improves top k recall. The optimal rank r solution has a closed form asymmetric factorization obtained from the SVD of the covariance weighted query key operator. This motivates SAKI, a training free KV cache index that directly preserves attention scores rather than key reconstruction quality. Across LLaMA 3.1 8B, Qwen 2.5 7B, Mistral 7B v0.1, and Llama 3.2 3B, SAKI outperforms key PCA at every tested rank. At rank 32, it removes 13 to 30 percent of PCA's remaining top 64 recall error, including improvements from 0.748 to 0.799 on LLaMA 3.1 8B and from 0.786 to 0.850 on Qwen 2.5 7B. It improves 68 to 89 percent of attention heads per model, with the largest gains in deeper layers. Predicted score MSE reductions closely match empirical measurements, with a Pearson correlation of 0.997, while ablation studies confirm that the gains arise from optimizing the attention score objective rather than covariance weighting alone. Analysis of the scoring operator further explains why weight only, invariant subspace, and key reconstruction methods can be suboptimal. SAKI uses random-matrix theory to separate genuine covariance signal from autocorrelated sampling noise, matching PCA with only 512 calibration tokens and adding value exactly where PCA sees no reliable signal.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。