发表机构
University of Edinburgh; Chamber of Deputies, Brazil（爱丁堡大学; 巴西众议院）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文提出选择性注意力冻结（SAF），在训练中途用拟合均值替换低方差注意力头，实现静态定义，减少计算并加速训练，同时保持低困惑度损失并提升泛化能力。
AI 中文摘要
一些注意力头在不同输入上学习到相似的模式。重用这些模式可以通过避免重复的查询-键分数计算和softmax来降低训练成本。通过受控预训练对比，我们识别出选择性注意力冻结（Selective Attention Freezing, SAF），该方法选择注意力模式方差较低的头，并在训练中途用拟合的softmax后均值替换其注意力权重。我们用绝对位置和相对距离偏好来表示这些固定模式，将存储从序列长度的二次方降低到线性。一个融合内核重建这些模式，并同时执行普通注意力头和被替换的头。在匹配的训练token预算下，替换25%的注意力头在124M参数和4K上下文中使替换后的优化器更新速度提升1.056倍，困惑度增加0.77%。在1B参数和8K上下文中，包括通信在内的四块GPU上替换后的更新速度提升1.068倍，困惑度增加0.51%。生成的模型还加速了长输入微调和因果预填充。经过关联回忆适配后，124M模型在固定长度下对更多键值对的泛化能力优于普通注意力和两个剪枝对照组。
英文摘要
Some attention heads learn similar patterns across inputs. Reusing these patterns could reduce training cost by avoiding repeated query-key score computation and softmax. Through controlled pretraining comparisons, we identify Selective Attention Freezing (SAF), which selects heads with low attention-pattern variance and replaces their attention weights with fitted post-softmax means halfway through training. We represent these fixed patterns with absolute-position and relative-distance preferences, reducing storage from quadratic to linear in sequence length. A fused kernel reconstructs the patterns and executes ordinary-attention and replaced heads together. At matched training-token budgets, replacing 25% of attention heads gives 1.056x faster post-replacement optimiser updates at 124M parameters and 4K context, with a 0.77% perplexity increase. At 1B and 8K context, post-replacement updates are 1.068x faster on four GPUs including communication, with a 0.51% perplexity increase. The resulting models also accelerate long-input finetuning and causal prefill. After associative-recall adaptation, the 124M model with 25% replacement generalises to more key-value pairs at a fixed length better than ordinary attention and two pruning controls.
