Surgical Repair of Collapsed Attention Heads in ALiBi Transformers

AI summary

Overview

• Research area: Natural language processing; mechanistic interpretability and targeted repair of transformer language models (attention head pathology in ALiBi-based models).
• Technical level: Advanced. The paper assumes familiarity with attention heads, Q/K/V projections, positional encodings, residual streams, and perplexity evaluation.
• Scope (one sentence): The paper diagnoses a systematic beginning-of-sequence (BOS) attention collapse in the BLOOM model family, attributes it to ALiBi's slope schedule, and shows that reinitializing the collapsed heads ("surgical reinitialization") recovers 98.7% of head capacity in BLOOM-1b7 and can even transiently improve the model beyond its stock perplexity.

What This Paper Is About

Prior work treats attention heads that dump most of their attention on the BOS token as redundant and prunes them. This paper argues the opposite for BLOOM models: those heads are not redundant, they are collapsed, and the collapse is a predictable side effect of ALiBi positional encoding rather than a sign that the capacity is unneeded. The goal is to repair the collapsed heads inside an existing checkpoint, on a single consumer GPU, instead of deleting them or preventing them from forming in the first place.

Key Contributions

1. Cross-scale diagnosis of a systematic pathology. A diagnostic tool classifies every head by BOS mass and attention entropy, revealing that 31–44% of heads are BOS-fixated across four BLOOM scales, concentrated in a contiguous band of upper head indices that tracks ALiBi's slope schedule.
2. Surgical reinitialization as a repair technique. Xavier Q/K/V reinitialization plus zeroed output projections plus gradient-masked freezing of all other parameters, applied in two passes to BLOOM-1b7, recovers 242 → 379 of 384 healthy heads (98.7%) — including three DEAD heads with near-zero entropy.
3. A controlled corpus comparison separating mechanism from content. Identical surgery on the C4 validation split versus a curated corpus produces identical head recovery (108/108 by epoch 3), establishing that reinitialization — not corpus content — drives recovery, while corpus content determines domain specialization.
4. Demonstration that healthy heads are also suboptimally configured. Extending surgery to the mostly-healthy H5 column produces a model that transiently beats stock BLOOM-1b7 by 25% on training perplexity (12.70 vs. 16.99), suggesting pretrained attention configurations are local minima rather than global optima.

Main Findings

• The collapse band tracks ALiBi slopes. Sick bands: BLOOM-560m (560M, 16 heads/layer) H9–H15 at 44%; BLOOM-1b7 (1.7B, 16) H9–H15 at 36%; BLOOM-3b (3B, 32) H20–H30 at 38%; BLOOM-7b1 (7.1B, 32) H21–H30 at 31%. The band always starts near the midpoint of the head index range and extends to the maximum index. ALiBi assigns slope m_h = 2^(−8(h+1)/H); for H15 in a 16-head model m15 ≈ 0.0039, the steepest distance penalty, which Appendix B lists as −0.39 nats at position 100.
• BOS mass is bimodal, making the classification threshold robust. Heads cluster near 0.0 or above 0.8, with few in between, so the 0.50 threshold (and 0.40 or 0.60) yields nearly identical classifications. Stock BLOOM-1b7 contains 242 healthy heads (63.0%), 136 BOS-sink (35.4%), 3 DEAD (0.8%), and 3 low-entropy (0.8%).
• Two-pass surgery recovers nearly all head capacity. Pass 1 targeted 108 heads in the H9–H15 band across layers 5–22 (302M parameters, 17.5% trainable); Pass 2 targeted 39 remaining collapsed heads outside the band (40M parameters, 2.3% trainable). Recovery went 242 → 341 (88.8%) → 379 (98.7%) healthy, with BOS-sink falling 136 → 39 → 1. One borderline head remains at L16-H4 (BOS mass 0.538). MLP dormancy is unchanged at 8.75%, confirming the surgery is attention-specific.
• Training perplexity improves while held-out perplexity worsens — a distribution shift, not capacity loss. Training PPL went 16.99 → 15.13 → 15.10, while held-out PPL on the 12 evaluation prompts went 21.45 → 27.52 → 28.76. On 50 held-out C4 texts (4,024 tokens), the C4-trained surgical model beat stock (29.30 vs. 32.42, a 9.6% improvement), while the curated model scored 74.11 on C4. Each model generalizes better inside its own training distribution.
• Reinitialization, not corpus content, is the mechanism. Curated and C4 corpora both woke 108/108 heads by epoch 3. The curated corpus reached lower training PPL (~15.4 vs. 20.80) in the same epochs; the C4 model never matched it and began overfitting (Table 4 lists C4 E15 training PPL at 36.31).
• Surgery causes global redistribution beyond the targeted heads. Untouched H0–H8 heads drifted: 76 of roughly 200 under curated E3 versus 56 under C4 E3 and 59 under C4 E15. The curated corpus caused more redistribution (mean |δ| 0.0426 vs. 0.0310) while producing better perplexity, ruling out the idea that it merely "protects" the model.
• Two distinct post-surgical phenomena. Phenomenon 1 is early, global, functional redistribution that improves the model (largely complete by epoch 3). Phenomenon 2 is late, local, pathological degradation: in-band frozen head drift roughly doubled from C4 E3 to E15 (5 → 10 of 18 drifters, mean |δ| 0.0399 → 0.0641) while outside-zone drift barely changed (56 → 59).
• Drift propagates along head-index columns. In the C4 baseline, the H15 column spread (mean |δ| 0.083 → 0.154) and H13 spread (0.046 → 0.070), while H9, H10, and H12 remained stable. The steepest-slope column is the most vulnerable to cascade effects.
• Healthy heads also improve when reinitialized. Adding the 18-head H5 column (14 healthy, 4 collapsed) brought the target total to 126 heads. All 18 returned with lower BOS mass — L14 dropped from a nominally healthy 0.483 to 0.023 (a 95% reduction), and even the healthiest head, L5 at 0.166, fell to 0.056.
• The extended model transiently beats stock. Training PPL: stock 16.99, band-only Pass 1 E3 15.13, band + H5 E1 13.90, and a sub-epoch minimum of 12.70 at step 42 of 140 (about 30% through epoch 1), after which it rises. Perplexity falls monotonically from 19.25 post-reinitialization, crossing the stock baseline by step 10 and rising above it after roughly one epoch.
• Generation remains coherent, with visible corpus imprinting. Across 50 prompts in five categories (150 completions total), all three conditions produced coherent text with no degeneration. The curated model showed a more abstract register but inserted HTML markup into outputs and sometimes answered French prompts with Python code; the C4 model resembled stock with slightly more verbosity.

Methodology in Plain English

The authors first build a diagnostic that runs text through a BLOOM model and, for every attention head, measures two things: how much of its attention lands on position 0 (BOS mass) and how spread out its attention is (Shannon entropy). Heads are labeled Healthy, BOS-sink, or DEAD using fixed thresholds.

For repair, they pick the collapsed heads and do four things: give the head's Q, K, and V projection matrices fresh random weights (Xavier normal), set the head's output projection to zero so it contributes nothing to the residual stream at first, freeze every other parameter in the model using gradient masks, and then train only the surgical parameters. They note that training the collapsed weights with gradients alone, without reinitializing, produced zero recovery after 15 epochs — the BOS-sink state is a sharp local minimum.

Training runs on a single NVIDIA RTX 5070 Ti with 16GB of VRAM in bfloat16, using AdamW at a learning rate of 5×10^−5 with no weight decay, linear warmup to cosine decay, batch size 1 with gradient accumulation over 8 steps, gradient clipping at max norm 1.0, gradient checkpointing enabled via model.enable_input_require_grads(), and a sequence length of 512 tokens. They perform surgery in two passes: the main H9–H15 band first, then the remaining collapsed heads outside that band.

Why This Matters

The paper reframes BOS-sink heads from waste to be pruned into dormant capacity that can be recovered. Since pruning is irreversible while reinitialization costs minutes of training on a consumer GPU, the practical calculus changes for anyone working with existing BLOOM checkpoints. The redistribution findings also carry a warning for the wider field: attention heads are connected through a shared residual stream, so any targeted parameter modification (LoRA, full fine-tuning, adapters) may alter untouched components, and the character of that alteration may depend on training data.

Real-world applications:

• Reviving open-weight checkpoints. Recovering head capacity in a deployed BLOOM checkpoint on a single consumer GPU instead of retraining or discarding the model.
• Rethinking compression pipelines. Attention-head pruning workflows could recover capacity before removal rather than permanently discarding it.
• Model auditing and diagnostics. The released diagnostic tool classifies head health in any BLOOM model, giving a fast read on whether a checkpoint has collapsed attention.
• Informing fine-tuning and adapter practice. The measurement framework for detecting global redistribution through the residual stream applies to any targeted modification technique.

Industry relevance centers on the fact that the technique is cheap, requires no architecture changes, and applies to an existing public model family — while the caution about corpus imprinting in generation output is directly relevant to anyone selecting fine-tuning data.

Future Directions

• Scale up the surgery. The full repair technique is demonstrated only on BLOOM-1b7; the paper confirms the pathology diagnostically at larger scales but does not attempt surgery on BLOOM-3b, 7.1B, or 176B.
• Isolate what corpus properties drive redistribution quality. The curated and C4 corpora differ simultaneously in domain, vocabulary, structure, repetition, and language distribution; a controlled ablation varying one property at a time at fixed token count is needed.
• Test whether the H5 result generalizes. The healthy-head experiment uses a single seed and a single column. Whether the improvement holds at other head indices, or under full-model reinitialization, is unknown — H5 may be special because of its moderate position in the ALiBi slope schedule.
• Make the improvement sustained. The H5 perplexity gain persists for only about one epoch before the small corpus is memorized; sustained access to the better attention configuration would require a larger training corpus, not a different technique.
• Extend beyond ALiBi. BOS-sink collapse occurs in other architectures, but the specific upper-index band pattern is a consequence of ALiBi's slope schedule, and whether the surgical method transfers is untested.

Target Audience

Researchers and engineers working on transformer interpretability, attention-head pruning, and model compression; practitioners who fine-tune or repair open-weight language models such as the BLOOM family; and anyone studying how positional encoding choices shape learned attention. Readers need working knowledge of attention mechanics and training loops, since the paper reports hyperparameters, perplexity, and head-level diagnostics without tutorial scaffolding.

Authors’ 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.