通过低维注意力选择压缩键值缓存，大幅节省显存且无需重训练。
Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection
- 将键向量分解为低维表示，仅保留必要注意力信息。
- 7B模型压缩4倍后仍保持9.24的困惑度，训练快8%、参数少12%。
- 适用于已有模型，可与分组查询、量化技术叠加使用。
标准Transformer中查询、键、值维度相同，但功能不同：查询与键用于生成注意力权重（选择），值用于传递丰富表征（值传递）。我们发现，选择只需$O(\log N)$维即可区分$N$个相关标记类别（如句法角色、语义簇、位置模式），远低于值传递所需维度。本文提出因子化键（factored keys），利用这一不对称性，在不重新训练的前提下物理压缩任意预训练模型的KV缓存——区别于GQA和MLA需在预训练前设计。通过截断奇异值分解（SVD）将键投影$W_K \approx A_{d \times r} B_{r \times d}$，取$W_K' = A$作为新键投影，生成$r$维紧凑键存入缓存，并将$B^\top$吸收进查询投影$W_Q' = W_Q B^\top$（零成本，因查询不缓存）。7B模型在$ r = d/4 $时，训练从头达到9.24对9.25的困惑度（200亿词，两种子均值），参数减少12%，训练加速8%。对现有模型，经SVD加3轮QK微调（<1%预训练数据）即可实现75%键缓存节省，质量损失约2%（适用于GPT-2和Mistral-7B）。该方法可与GQA和量化结合，实现最高16倍键缓存压缩。对128K上下文的7B模型，每用户节省25GB缓存，相同硬件下支持约60%更多并发用户。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。