The Structural Origin of Attention Sink: From Variance Discrepancy to Multi-Stage Collapse

TL;DR

Why do Large Language Models (LLMs) obsess over the first token? This paper reveals that the "attention sink" is not a design choice but a structural failure. It starts with a simple variance discrepancy in causal attention, escalates through FFN "super neurons," and results in a "dimension disparity" that locks the model's focus. The authors propose a simple fix—Head-wise RMSNorm—which not only kills the sink but also makes training faster and more stable.

Background: The Mystery of the First Token

In decoder-only Transformers, the first token consistently receives an absurdly high attention score, even when it is just a BOS (Beginning of Sentence) tag or random noise. While this has been exploited for KV cache streaming, it also correlates with pathological behaviors like representation collapse and activation outliers. This research moves beyond functional "explanations" to find the actual hardware-level "why."

The Problem: A Chain Reaction of Variance

The authors identify a three-stage "causal chain" that creates the sink:

1. Causal Mining: Because of the causal mask, the first token only attends to itself. Subsequent tokens average values from multiple prior tokens. This prevents the first token from "averaging out," leaving it as a high-variance outlier.
2. Super Neuron Activation: FFN layers contain "super neurons" with massive weight norms. The high-variance first token selectively triggers these neurons, generating massive activation spikes.
3. Structural Locking: These spikes are channeled through sparse down-projections, creating "Dimension Disparity." The representation collapses into a single direction, which then "locks" the Query-Key dot products in the next layer, forcing the attention to sink.

Figure 1: The full propagation chain from value aggregation variance to attention sink formation.

Methodology: Intervening in the Machine

To prove this theory, the authors performed "surgical" interventions:

• Mask Intervention: By forcing a middle token to only attend to itself, they could create a new attention sink at any position.
• Variance Amplification: Simply scaling the variance of a token (without scaling its absolute norm) was enough to turn it into a sink.

The Solution: Head-wise RMSNorm

Instead of replacing Softmax (which can be unstable), the authors propose adding a normalization step immediately after value aggregation: ${\hat{\mathbf{o}}}_t^{(h)}=\frac{{\mathbf{o}}_t^{(h)}}{\mathrm{RMS}({\mathbf{o}}_t^{(h)})}\odot \mathbf{\lambda }$ This ensures that every head and every position contributes a standardized variance to the residual stream, preventing any single token from becoming a structural outlier.

Experimental Evidence: Faster Convergence & Better Rank

The authors tested their method on a Llama-2 architecture trained from scratch. The results were striking:

• Convergence: Models with Head-wise RMSNorm reached lower validation loss faster than the baseline.
• Manifold Health: The "Effective Rank" of the hidden states stayed high, whereas the baseline suffered from "Shallow Layer Collapse"—a sudden drop in representation diversity.

Figure 2: Head-wise RMSNorm maintains a significantly higher effective rank across layers, preventing manifold collapse.

Critical Insights

This work represents a shift from describing LLM artifacts to engineering them away. It proves that the attention sink is a structural byproduct of the asymmetry in causal masking. By restoring statistical parity across the sequence, we don't just fix a weird attention pattern—we improve the fundamental optimization landscape of the Transformer.

Limitations

The study was primarily conducted on 152M-parameter models. While the authors validated the presence of these mechanisms in Llama-2-7B and Llama-3-8B, the full training benefit of Head-wise RMSNorm at the 70B+ scale remains to be seen.

Conclusion

The attention sink is a "canary in the coal mine" for variance imbalance. By addressing the root structural cause rather than the functional symptom, we can build LLMs that are more representationally dense and easier to train.