别这么Stief！在Stief尔流形上学习KV缓存低秩近似
机器学习
2026-01-30 v1
摘要
键-值（KV）缓存使自回归解码能够快速实现，但在长文本场景下成为高带宽内存（HBM）容量和带宽的主要瓶颈。常见的缓解措施是通过将每个头的矩阵投影到更低的秩来压缩缓存的键和值，只存储投影结果于HBM中。然而，现有的后训练方法通常使用基于SVD的代理目标来拟合这些投影，这可能不良好地反映softmax、值混合以及后续解码器层变换后的端到端重建。出于此种原因，我们引入StiefAttention，这是一种后训练的KV缓存压缩方法，通过直接最小化解码器层输出重建误差来学习“正交”投影基。StiefAttention此外为每个层预计算一个误差-秩轮廓，使可在用户指定的误差预算下进行灵活的层级秩分配。值得注意的是，在相同条件下，StiefAttention对Llama3-8B的C4困惑度提升11.9分，对0-shot MMLU准确率提升5.4%，相对于原始解码器输出实现更低的相对误差和更高的余弦相似度。
引用
@article{arxiv.2601.21686,
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 Yüzügüler and Lukas Cavigelli and Enrico Macii and Daniele Jahier Pagliari and Alessio Burrello},
journal= {arXiv preprint arXiv:2601.21686},
year = {2026}
}