提出StiefAttention方法，优化长文本生成中的键值缓存压缩效率
Don't be so Stief! Learning KV Cache low-rank approximation over the Stiefel manifold
- 在Stiefel流形上学习正交投影基，直接最小化解码层输出误差
- 在相同压缩率下，比EigenAttention在C4困惑度上提升4.2点，MMLU准确率提升8.9点
- 支持按层分配秩，适配用户指定的缓存预算，适合大模型推理优化
键值（KV）缓存可加速自回归解码，但在长上下文场景下成为高带宽内存（HBM）容量与带宽的主要瓶颈。常见缓解方法是将每头矩阵投影到低秩，仅存储投影结果。但现有后训练方法通常使用SVD类代理目标拟合投影，可能无法反映经过softmax、值混合及后续解码层变换后的端到端重构效果。为此，我们提出StiefAttention，一种后训练的KV缓存压缩方法，通过直接最小化解码层输出重构误差来学习正交投影基。StiefAttention还构建了各层候选秩下的误差-秩曲线，支持在用户指定的缓存预算下进行分层秩分配。值得注意的是，在相同条件下，Llama3-8B模型上，StiefAttention在C4困惑度上比EigenAttention提升4.2点，0-shot MMLU准确率提升8.9点，且相对误差更低，与原始解码层输出的余弦相似度更高。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。