arXiv 2603.044272026-03-28cs.LGcs.AI

细键，全值：通过低维注意力选择减少KV缓存

Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection

Sapient Intelligence；Department of Computing Science, University of Alberta（阿尔伯塔大学计算机科学系）

Hengshuai Yao, Xing Chen, Ahmed Murtadha, Guan Wang

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出通过低维注意力选择减少KV缓存，展示选择仅需O(logN)维度即可区分N个相关token类别，相比传统Transformer注意力节省参数和训练时间。

摘要：

标准Transformer注意力使用相同的维度性来处理查询、键和值，但这些组件扮演不同的角色：查询和键产生标量注意力权重（选择），而值承载丰富的表示（值传输）。我们证明选择仅需O(logN)维度来区分N个相关token类别（例如，语法角色、语义簇、位置模式）——远少于值传输所需。我们引入了因子键，利用这种不对称性，使任何预训练模型的KV缓存物理上缩小，而无需重新训练——不同于组查询注意力（GQA）和多头潜在注意力（MLA），这些必须在预训练前设计到架构中。我们将每个键投影W_K≈A_{d×r}B_{r×d}通过截断奇异值分解（SVD）分解，设置W_K' = A作为新的键投影，产生紧凑的r维键用于缓存，并将B^T吸收进查询投影（W_Q' = W_Q B^T）零成本——因为查询从不缓存。在7B规模上，使用r = d/4（d是模型维度）从头训练，与全注意力的困惑度（9.24 vs 9.25 PPL在20B tokens后，平均超过两个种子）匹配，同时使用12%更少的参数和训练8%更快。对于现有模型，SVD后接QK微调（3个epoch，少于1%的预训练数据）在GPT-2和Mistral-7B上实现75%的键缓存节省，质量成本约为2%。该方法可与GQA和量化结合，达到最多16倍的键缓存压缩。对于一个7B模型服务128K上下文，因子键每用户节省25GB的KV缓存，使相同硬件上可支持大约60%更多的并发用户。

英文摘要：

Standard Transformer attention uses identical dimensionality for queries, keys, and values, yet these components serve different roles: queries and keys produce scalar attention weights (selection), while values carry rich representations (value transfer). We show that selection requires only $O(\log N)$ dimensions to distinguish among $N$ relevant token categories (e.g., syntactic roles, semantic clusters, positional patterns) -- far fewer than value transfer needs. We introduce factored keys, which exploit this asymmetry to physically shrink the KV cache of any pretrained model without retraining from scratch -- unlike Grouped-Query Attention (GQA) and Multi-Head Latent Attention (MLA), which must be designed into the architecture before pretraining. We factorize each key projection $W_K \approx A_{d \times r} B_{r \times d}$ via truncated singular value decomposition (SVD) (where $r$ is the chosen compression dimension), set $W_K' = A$ as the new key projection producing compact $r$-dimensional keys for the cache, and absorb $B^\top$ into the query projection ($W_Q' = W_Q B^\top$) at zero cost -- since queries are never cached. At the 7B scale, training from scratch with $r = d/4$ (where $d$ is the model dimension) matches full-attention perplexity ($9.24$ vs $9.25$ PPL after 20B tokens, mean over two seeds) while using 12% fewer parameters and training 8% faster. For existing models, SVD followed by QK fine-tuning (3 epochs, less than 1% of pretraining data) achieves 75% key cache savings at roughly 2% quality cost on both GPT-2 and Mistral-7B. The approach composes with GQA and quantization for up to $16\times$ combined key cache compression. For a 7B model serving a 128K context, factored keys save 25 GB of KV cache per user, enabling roughly 60% more concurrent users on identical hardware.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。

AI总结：

摘要：

英文摘要：

相关论文