对 BLOOM 变压器中注意力头崩溃的外科修复
计算与语言
2026-03-11 v1
摘要
我们发现 BLOOM 系列变压器语言模型中存在一种系统性注意力崩溃病理，即 ALiBi 位置编码导致 31-44% 的注意力头几乎完全注意开头的序列标记。这种崩溃在四个模型规模（560M 至 7.1B 参数）间遵循可预测的模式，集中在 ALiBi 的斜率日程安排施加最陡距离惩罚的注意力头索引处。我们引入外科重新初始化：针对性的 Q/K/V 重新初始化，零化输出投影，并对除手术参数外的所有参数进行梯度掩码冻结。应用于 BLOOM-1b7 的单张消费级 GPU 上，该技术在两次迭代中恢复了 98.7% 的运行中注意力头容量（从 242 个恢复到 379 个）。对比 C4 训练数据的受控比较确认，重新初始化——而非语料库内容——驱动了恢复，揭示了两种不同的手术后现象：早期的全局功能再分配改善了模型，而晚期的局部退化在噪声训练信号下积累。一次实验中，重新初始化大多数健康的注意力头以及崩溃的注意力头，产生的模型在训练困惑度上暂时超过原始 BLOOM-1b7 25%（12.70 vs. 16.99），表明预训练的注意力配置是次优的局部极小值。我们随身携带的代码、检查点和诊断工具作为开源软件发布。
引用
@article{arxiv.2603.09616,
title = {Surgical Repair of Collapsed Attention Heads in ALiBi Transformers},
author = {Palmer Schallon},
journal= {arXiv preprint arXiv:2603.09616},
year = {2026}
}
备注
15 pages, 7 figures, 2 supplementary figures. Code: https://github.com/Palmerschallon/bloom-head-surgery Checkpoints: https://huggingface.co/TheNexus42/bloom-1b7-head-surgery