Surgical Repair of Collapsed Attention Heads in ALiBi Transformers

TL;DR

This paper identifies a systematic attention collapse in ALiBi transformers, introduces a surgical reinitialization method to recover attention head functionality, and demonstrates significant performance improvements through targeted reinitialization.

Contribution

The paper presents a novel surgical reinitialization technique to fix collapsed attention heads in ALiBi transformers, improving model capacity and performance.

Findings

Reinitialization recovers 98.7% of attention heads in BLOOM-1b7.

Reinitializing healthy heads alongside collapsed ones improves perplexity by 25%.

Attention collapse follows a predictable pattern across model scales.

Abstract

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.