Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold

Abstract

Key--value (KV) caching enables fast autoregressive decoding but at long contexts becomes a dominant bottleneck in High Bandwidth Memory (HBM) capacity and bandwidth. A common mitigation is to compress cached keys and values by projecting per-head matrixes to a lower rank, storing only the projections in the HBM. However, existing post-training approaches typically fit these projections using SVD-style proxy objectives, which may poorly reflect end-to-end reconstruction after softmax, value mixing, and subsequent decoder-layer transformations. For these reasons, we introduce StiefAttention, a post-training KV-cache compression method that learns orthonormal projection bases by directly minimizing decoder-layer output reconstruction error. StiefAttention additionally precomputes, for each layer, an error-rank profile over candidate ranks, enabling flexible layer-wise rank allocation under a user-specified error budget.

arXiv PDF

Cite this work

@misc{dontbesostieflearningkvcachelowrankapproximationoverthestiefelmanifold,
  title = {Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold},
  author = {Luca Benfenati and Matteo Risso and Andrea Vannozzi and Ahmet Caner Yuzuguler and Lukas Cavigelli and Enrico Macii and Daniele Jahier Pagliari and Alessio Burrello},
  year = {2026},
  note = {arXiv},
  url = {https://arxiv.org/abs/2601.21686}
}

Download