Search papers, labs, and topics across Lattice.
This paper introduces SAKI, a novel training-free key-value (KV) cache indexing method that prioritizes the preservation of attention scores over traditional key reconstruction quality. By deriving a covariance-weighted low-rank objective and optimizing for attention score distortion, SAKI significantly enhances top-k recall across multiple models, outperforming key PCA at every tested rank. Notably, SAKI achieves substantial improvements in recall metrics, with increases of up to 0.102 on LLaMA 3.1 8B and 0.064 on Qwen 2.5 7B, while also benefiting a majority of attention heads in deeper layers.
SAKI achieves up to 30% better recall than traditional key PCA methods by directly preserving attention scores, not just key variance.
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.
