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

TL;DR

KV-cache memory bottlenecks in long-context LLMs motivate post-training compression. StiefAttention learns orthonormal projection bases on the Stiefel manifold by directly minimizing decoder-layer output reconstruction error, rather than proxy KV-reconstruction objectives. It trains lightweight predictors from activation statistics to produce layer-specific bases, with layer-wise rank surfaces and deployment-time rank-allocation policies. Empirically, on Llama-3-8B under iso-memory budgets, StiefAttention outperforms EigenAttention in end-to-end metrics (e.g., C4 perplexity and MMLU accuracy) and yields better alignment of compressed layer outputs with the originals, especially in early layers. This approach offers a tunable accuracy-memory tradeoff with practical deployment benefits for large-scale autoregressive decoding.

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 \emph{orthonormal} projection bases by directly minimizing \emph{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. Noteworthy, on Llama3-8B under the same conditions, StiefAttention outperforms EigenAttention by $11.9$ points on C4 perplexity and $5.4\%$ on 0-shot MMLU accuracy at iso-compression, yielding lower relative error and higher cosine similarity with respect to the original decoder-layer outputs.

Figures (6)

• Figure 1: Overview of StiefAttention. From lightweight activation statistics $\mu_{K}, \mu_{V}$ and $\sigma_{K}^2, \sigma_{V}^2$, we learn orthonormal projection bases $P_K, P_V$ which are optimized to minimize decoder-layer output error.
• Figure 2: Zero-shot accuracy--memory tradeoff under KV-cache compression on HellaSwag, PIQA, and MMLU. Each point corresponds to an error budget $\epsilon$ that determines layer-wise ranks; the x-axis reports the resulting KV-cache ratio (lower is better).
• Figure 3: Perplexity--memory tradeoff under KV cache compression on WikiText-2 and C4. Each point corresponds to an error budget $\epsilon$ that determines layer-wise ranks; the x-axis reports the resulting KV cache ratio (lower is better).
• Figure 4: Effect of rank-allocation policies.
• Figure 5: Layer-wise rank profiles under Weighted Pareto rank-allocation policy across all error budgets $\epsilon$