arXiv:2603.09616cs.CL2026-03

修复ALiBi模型中31%-44%注意力头失效问题,提升模型性能。

Surgical Repair of Collapsed Attention Heads in ALiBi Transformers

  • 针对ALiBi编码导致注意力集中于序列开头的问题,采用精准重初始化策略。
  • 在BLOOM-1b7上恢复98.7%注意力头功能,从242增至379个有效头。
  • 适用于研究或优化Transformer结构的开发者,尤其关注位置编码设计者。

我们发现BLOOM系列Transformer语言模型存在系统性注意力崩溃病理,其中ALiBi位置编码导致31%-44%的注意力头几乎完全关注序列起始标记。该现象在四个模型规模(560M至7.1B参数)中呈现可预测模式,集中在ALiBi斜率调度施加最陡峭距离惩罚的头索引处。我们提出手术式重初始化:对特定Q/K/V进行重初始化,输出投影设为零,并对非手术参数实施梯度掩码冻结。在单张消费级GPU上应用于BLOOM-1b7,仅两轮迭代即恢复98.7%的可用头容量(从242增至379/384)。与C4训练数据的对照实验表明,恢复效果源于重初始化而非语料内容;并揭示两种术后现象:早期全局功能再分配带来性能提升,晚期局部退化则在噪声训练信号下累积。扩展实验中,同时重初始化多数健康头与崩溃头,使模型训练困惑度短暂降至12.70(原为16.99),优于原始BLOOM-1b7的25%,表明预训练注意力配置可能处于次优局部极小值。代码、检查点及诊断工具已开源。

原文摘要 · Abstract (English)

We identify a systematic attention collapse pathology in the BLOOM family of transformer language models, where ALiBi positional encoding causes 31-44% of attention heads to attend almost entirely to the beginning-of-sequence token. The collapse follows a predictable pattern across four model scales (560M to 7.1B parameters), concentrating in head indices where ALiBi's slope schedule imposes the steepest distance penalties. We introduce surgical reinitialization: targeted Q/K/V reinitialization with zeroed output projections and gradient-masked freezing of all non-surgical parameters. Applied to BLOOM-1b7 on a single consumer GPU, the technique recovers 98.7% operational head capacity (242 to 379 of 384 heads) in two passes. A controlled comparison with C4 training data confirms that reinitialization -- not corpus content -- drives recovery, and reveals two distinct post-surgical phenomena: early global functional redistribution that improves the model, and late local degradation that accumulates under noisy training signal. An extended experiment reinitializing mostly-healthy heads alongside collapsed ones produces a model that transiently outperforms stock BLOOM-1b7 by 25% on training perplexity (12.70 vs. 16.99), suggesting that pretrained attention configurations are suboptimal local minima. Code, checkpoints, and diagnostic tools are released as open-source software.

注意力机制Transformer模型修复ALiBi

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。