Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913171585892352 |
|---|---|
| author | Benfenati, Luca Risso, Matteo Vannozzi, Andrea Yüzügüler, Ahmet Caner Cavigelli, Lukas Macii, Enrico Pagliari, Daniele Jahier Burrello, Alessio |
| author_facet | Benfenati, Luca Risso, Matteo Vannozzi, Andrea Yüzügüler, Ahmet Caner Cavigelli, Lukas Macii, Enrico Pagliari, Daniele Jahier Burrello, Alessio |
| contents | 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 matrices 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 constructs layer-wise error-rank profiles over candidate ranks, enabling sequential rank allocation under a user-specified KV cache budget. Notably, on Llama3-8B under the same conditions, StiefAttention outperforms EigenAttention by $4.2$ points on C4 perplexity and $8.9$ points on 0-shot MMLU accuracy at iso-compression, yielding lower relative error and higher cosine similarity with respect to the original decoder-layer outputs. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_21686 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold Benfenati, Luca Risso, Matteo Vannozzi, Andrea Yüzügüler, Ahmet Caner Cavigelli, Lukas Macii, Enrico Pagliari, Daniele Jahier Burrello, Alessio Machine Learning 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 matrices 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 constructs layer-wise error-rank profiles over candidate ranks, enabling sequential rank allocation under a user-specified KV cache budget. Notably, on Llama3-8B under the same conditions, StiefAttention outperforms EigenAttention by $4.2$ points on C4 perplexity and $8.9$ points on 0-shot MMLU accuracy at iso-compression, yielding lower relative error and higher cosine similarity with respect to the original decoder-layer outputs. |
| title | Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2601.21686 |