SAKI: Score-Aware Low-Rank Key Indexing for Long-Context KV Retrieval

2026-08-04

AI Summary

Existing low-rank KV cache methods fail to directly optimize attention scores, limiting retrieval accuracy. This work proposes a training-free, covariance-weighted low-rank indexing approach that, for the first time, formulates attention score distortion as a covariance-weighted low-rank approximation problem. We derive a closed-form asymmetric factorization solution, revealing the suboptimality of conventional methods. Our method performs covariance-weighted decomposition of the query-key operator via SVD and explicitly optimizes Top-k recall. Evaluated across multiple mainstream large language models, it substantially outperforms Key PCA: at rank 32, Top-64 recall error is reduced by 13%–30%, with over 70% of attention heads benefiting. Predicted and empirical mean squared errors exhibit strong agreement (r = 0.997).

Abstract

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.

Problem

• attention scores
• low-rank compression
• KV cache
• score distortion
• long-context retrieval

Innovation

• attention score preservation
• low-rank KV cache
• covariance-weighted SVD
• training-free indexing
• long-context retrieval