arXiv cs.LG

SAKI-RoPE: Score-Aware Low-Rank Key Indexing with RoPE-Aware Covariance and Random-Matrix Noise Analysis

· 1小时前
AI 摘要

arXiv:2608.03228v4 Announce Type: replace Abstract: Existing low-rank KV-cache indexes preserve model weights or key variance, not the attention scores that determine retrieval. We derive the expected score distortion of key-side rank-r compression, a two-sided covariance-weighted low-rank problem that uses both query and key statistics, and solve it exactly: the optimum is an optimal rank-r linear score map, computed in closed form from two d x d covariances per head, and its predicted per-head score-MSE reduction matches measurement (Pearson r = 0.997). The resulting training-free index, SAKI, beats key-PCA at every rank on LLaMA-3.1-8B, Qwen2.5-7B, Mistral-7B and Llama-3.2-3B, at every context length from 4K to 64K, and in every calibration-domain pair (0.748 -> 0.799 recall@64 at r = 32 on LLaMA-3.1-8B). Because RoPE rotates keys by their distance from the query, we make the closed form RoPE-aware by averaging the query covariance over rotations; this SAKI-RoPE, our main method, raises recall to 0.861, beats PCA on every head of all four models, and needs only 28-39% of PCA's rank to match it at r = 32 across the four models. Random-matrix theory of the covariance estimation noise shows that 512 calibration tokens suffice and that clipping the Marchenko-Pastur noise bulk helps SAKI when calibration is short. We make no perplexity claim: with the attention sink pinned, perplexity differences between indexes are small and depend on model and text, and we show that any low-rank index calibrated on a long prefix drops the attention-sink token, which collapses 32K perplexity unless the sink is pinned.

查看原文(arXiv cs.LG)