新研究解释 KV-cache 压缩为何有效:注意力空间复杂度的新下界
Attention under Bounded Key-Query Similarity: Space Complexity and Transformer Anisotropy
做 KV-cache 压缩的人可以看看,这篇用理论证明了为什么连均匀采样这种简单方法在真实模型上就好用。
arXiv 论文研究在 key/query 余弦相似度被限制在零附近小区间时的注意力近似空间复杂度。作者给出了近似空间最优的算法和对应的下界证明,下界构造中 key/query 相似度分布恰好呈现真实 Transformer 已观测到的各向异性。论文还证明在常见参数范围内,简单的均匀采样算法已接近最优,解释了它在真实数据集上的效果。这项工作为 KV-cache 压缩方法提供了理论基础。
Attention under Bounded Key-Query Similarity: Space Complexity and Transformer Anisotropy
A growing body of work on KV-cache compression aims to reduce the memory requirements of autoregressive generation. These methods compress the stored context while approximately preserving attention computation, enabling generation with reduced memory consumption. Almost tight theoretical results on the worst case space complexity of attention have recently been obtained, but they are achieved by instances where queries can be perfectly aligned with keys. This directly contradicts the observed anisotropy of the key/query geometry in real transformers. In particular, it has been shown empirically that key/query cosine similarities are small even for most correlated pairs, i.e., near-perfect alignment of queries with keys does not happen in practice. In this paper, we study the space complexity of attention approximation under the assumption that key/query cosine similarities belong to a small interval around zero, i.e., a query cannot align too much with any key. We give a nearly space optimal algorithm and a corresponding space lower bound. Surprisingly, the distribution of key/query cosine similarities in our lower bound exhibits the anisotropy property empirically established and studied in the literature. Additionally, we show that the simple uniform sampling algorithm is nearly optimal for a typical range of parameters, supporting its surprising effectiveness on real datasets.