Surgical Repair of Collapsed Attention Heads in ALiBi Transformers
Organizations: Independent researcher
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.
Figures & tables
| Model | Params | Heads/Layer | Sick Band | Sick % | Layers Affected |
|---|---|---|---|---|---|
| BLOOM-560m | 560M | 16 | H9–H15 | 44% | 56–94% per index |
| BLOOM-1b7 | 1.7B | 16 | H9–H15 | 36% | 56–94% per index |
| BLOOM-3b | 3B | 32 | H20–H30 | 38% | 63–94% per index |
| BLOOM-7b1 | 7.1B | 32 | H21–H30 | 31% | 66–94% per index |
| Metric | Stock | Pass 1 (E3) | Pass 2 (E1) |
|---|---|---|---|
| Healthy | 242 (63.0%) | 341 (88.8%) | 379 (98.7%) |
| BOS-sink | 136 | 39 | 1 |
| DEAD | 3 | 0 | 0 |
| Low-entropy | 3 | 4 | 4 |
| Training PPL | 16.99 | 15.13 | 15.10 |
| Held-out PPL | 21.45 | 27.52 | 28.76 |
| Model | C4 Validation PPL | Eval Prompt PPL |
|---|---|---|
| Stock BLOOM-1b7 | 32.42 | 21.45 |
| C4 Surgical (E3) | 29.30 | 29.59 |
| Curated Surgical (Pass 2) | 74.11 | 28.76 |
| Metric | Curated (E3) | C4 (E3) | C4 (E15) |
|---|---|---|---|
| Heads woken | 108/108 | 108/108 | 108/108 |
| Training PPL | 15.4 | 20.80 | 36.31 |
| Held-out PPL | 27.52 | 29.59 | — |
| Metric | Curated E3 | C4 E3 | C4 E15 |
|---|---|---|---|
| Drifting heads (of 200) | 76 | 56 | 59 |
| Mean | 0.0426 | 0.0310 | 0.0398 |
| Worst case | L23H14: +0.444 | L23H14: +0.364 | L19H5: +0.353 |
| Metric | Curated E3 | C4 E3 | C4 E15 |
|---|---|---|---|
| Drifting heads (of 18) | 9 | 5 | 10 |
| Mean | 0.0990 | 0.0399 | 0.0641 |
| Column | Frozen Heads | E3 Mean | E15 Mean | Trend |
|---|---|---|---|---|
| H15 | 4 | 0.083 | 0.154 | Spreading |
| H13 | 2 | 0.046 | 0.070 | Spreading |
| H9 | 4 | 0.034 | 0.051 | Stable |
| H10 | 3 | 0.046 | 0.046 | Stable |
| H12 | 3 | 0.007 | 0.019 | Stable |
| Layer | Stock BOS Mass | Stock Status | Post-Surgery BOS Mass | Change |
| L5 | 0.166 | Healthy | 0.056 | |
| L7 | 0.317 | Healthy | 0.034 | |
| L10 | 0.374 | Healthy | 0.040 | |
| L11 | 0.517 | Collapsed | 0.079 | |
| L12 | 0.666 | Collapsed | 0.027 | |
| L14 | 0.483 | Healthy | 0.023 |
| Condition | Training PPL |
|---|---|
| Stock BLOOM-1b7 | 16.99 |
| Band-only surgery (Pass 1 E3) | 15.13 |
| Band + H5 surgery (E1) | 13.90 |
| Band + H5 sub-epoch best (step 42/140) | 12.70 |
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
| Head | Slope | Distance Penalty at pos 10 | At pos 100 |
|---|---|---|---|
| H0 | 0.7071 | ||
| H4 | 0.1768 | ||
| H8 | 0.0442 | ||
| H12 | 0.0110 | ||
| H15 | 0.0039 |
Attention Head Analysis
Attention head analysis studies how individual Transformer heads route information, support computations, and affect model behavior.
See all 81 papers