论文

SAKI:用于 KV 检索的具有随机矩阵噪声校正的分数感知低秩键索引

SAKI: Score-Aware Low-Rank Key Indexing with Random-Matrix Noise Correction for KV Retrieval

模型推理KV Cache

摘要

现有的低秩 KV 缓存方法保留模型权重或关键方差,这两者都不能直接反映推理过程中使用的注意力分数。我们推导了由秩 r 键压缩引起的预期注意力分数失真,并表明它产生了协方差加权的低秩目标。在裕度条件下,控制这种失真也可以提高 top k 召回率。最优秩 r 解具有从协方差加权查询键运算符的 SVD 获得的封闭形式非对称分解。这激发了 SAKI,一种 无需训练 KV 缓存索引,它直接保留注意力分数而不是关键重建质量。在 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%,其中更深层次的收益最大。预测得分 MSE 降低与经验测量非常匹配,皮尔逊相关性为 0.997,而消融研究证实,收益来自于优化注意力得分目标,而不是单独的协方差加权。对评分算子的分析进一步解释了为什么仅权重、不变子空间和关键重建方法可能不是最优的。 SAKI 使用随机矩阵理论将真正的协方差信号与自相关采样噪声分开,仅将 PCA 与 512 个校准标记相匹配,并在 PCA 没有看到可靠信号的地方精确地添加值。