发现大模型注意力头可线性预测,节省显存。
Linear Predictability of Attention Heads in Large Language Models
- 用少数参考头的线性组合重构其他注意力头。
- 在多个模型上重建准确率超0.76,部分达0.85以上。
- 仅缓存参考头能减半显存,且重构键比值更稳健。
大型语言模型推理日益受制于键值缓存,但注意力头激活的细粒度结构仍不明确。研究发现预训练Transformer普遍存在跨头线性结构:对特定词元,某注意力头的查询、键、值向量通常可由同层少量同伴头的线性组合重建。在Llama-3.1-8B、Falcon3-10B、OLMo-2-7B和Qwen3-32B上,仅需2-5个参考头即可高保真重建目标头(如在C4数据集上键的均方决定系数R²约0.76,五参考时;在GSM8K上常超过0.85)。该可预测性为学习所得而非架构设计:随机初始化时几乎不存在,随OLMo-2训练过程迅速提升,且理论下界显示初始阶段线性预测误差极高。我们进一步发现此现象与键投影子空间的层内对齐增强相关。最后,利用这一冗余,仅缓存参考头的键值状态,通过轻量线性映射实时重建其余头,实现2倍键值缓存压缩,模型依赖的精度损失为4.5-5.5个百分点(在五个基准测试上),且重构键比重构值影响更小。
原文摘要 · Abstract (English)
Large language model (LLM) inference is increasingly bottlenecked by the Key-Value (KV) cache, yet the fine-grained structure of attention-head activations remains poorly understood. We show that pretrained Transformers exhibit a pervasive inter-head linear structure: for a given token, the Query, Key, and Value (QKV) vectors of an attention head can often be reconstructed as a linear combination of a small number of peer heads, typically within the same layer. Across Llama-3.1-8B, Falcon3-10B, OLMo-2-7B, and Qwen3-32B, just 2-5 reference heads recover many target heads with high fidelity (e.g., mean R^2 approx 0.76 for Keys on C4 with five references, and frequently R^2 > 0.85 on GSM8K). This predictability is learned rather than architectural: it is largely absent at random initialization, rises rapidly during pretraining as we track through OLMo-2 checkpoints, and is supported by a theoretical lower bound showing high mean-squared error for linear prediction at initialization. We further connect this emergence to increasing intra-layer alignment of Key projection subspaces. Finally, we exploit this redundancy for efficiency by caching only reference-head KV states and reconstructing the remaining heads on the fly via lightweight linear maps, achieving 2x KV-cache reduction with model-dependent accuracy trade-offs (4.5-5.5 percentage point average drop on Falcon3-10B and Qwen3-32B across five benchmarks, and larger drops on Llama-3.1-8B), and we find that reconstructing Keys is substantially less harmful than reconstructing Values.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。