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

2026-01-29

AI Summary

This work addresses the performance bottleneck in long-context inference caused by the substantial high-bandwidth memory (HBM) consumption of key-value (KV) caches. The authors propose StiefAttention, a post-training KV cache compression method that, for the first time, incorporates orthogonality constraints into low-rank approximation by optimizing end-to-end decoder output reconstruction error over the Stiefel manifold. The approach further enables adaptive per-layer rank allocation based on a user-specified error budget. Experimental results on Llama3-8B demonstrate that, at the same compression ratio, StiefAttention reduces C4 perplexity by 11.9 points and improves zero-shot MMLU accuracy by 5.4%, while achieving lower relative output error and higher cosine similarity compared to existing methods.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Data Compression
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• KV cache compression
• memory bottleneck
• decoder-layer reconstruction
• long-context inference
• post-training compression

Innovation

• KV cache compression
• Stiefel manifold
• orthonormal projection
• post-training optimization
• layer-wise rank allocation