From Projection Geometry to Attention Fidelity: Explaining Cross-Layer KV Prediction
Tinuade Adeleke ⋅ William Stella ⋅ Arya Srikar Reddy ⋅ Ruizhe Li
Abstract
The key-value (KV) cache is a major memory cost in large language model (LLM) inference. Although several compression methods exploit similarities between KV activations in adjacent layers, the underlying sources of this predictability and its relationship to attention behavior remain unclear. We deconstruct the origins of cross-layer predictability by analyzing the two stages of the key-generation pipeline that act on the residual stream: the key projection and Rotary Positional Embedding (RoPE). A ladder of linear probes separates the contribution of each stage and shows that adjacent key projections read from substantially overlapping input subspaces. Across five model configurations, this projection overlap correlates with cross-layer key predictability ($r=0.54$--$0.75$, permutation $p<10^{-4}$), while the learned cross-layer maps themselves remain geometrically complex and do not reduce to simple rotations even in later layers. Norm-controlled patching experiments then show that prediction errors aligned with the leading query directions disrupt attention roughly $5\times$ more than equal-norm errors in the complementary subspace. Finally, in an equal-payload low-rank codec, a query-weighted objective accepts higher key MSE but reduces attention JS divergence by up to 24\%. These results show that KV-cache compression should optimize attention-relevant error, not Euclidean reconstruction error alone.
Chat is not available.
Successful Page Load