Technical Summary: Sink vs. Diagonal Patterns as Mechanisms for Attention Switch and Oversmoothing Prevention

Problem Statement

The paper addresses the theoretical and geometric understanding of attention sinks—a phenomenon where specific tokens (often the beginning-of-sequence, or BOS) receive disproportionate attention mass from other tokens. While attention sinks are empirically ubiquitous in pretrained Large Language Models (LLMs) across various modalities, their precise geometric representation, their necessity for preventing oversmoothing (the convergence of token representations to a single vector), and their relationship to alternative mechanisms like diagonal attention (self-attention) remain incompletely understood.

Existing literature offers conflicting or incomplete explanations: some suggest sinks prevent oversmoothing, while others link them to an "active-dormant" mechanism (attention switch). However, prior work often fails to:

1. Specify the geometric conditions under which sinks can be represented.
2. Rigorously prove that sinks are the only mechanism for a "hard attention switch" (where output is zero).
3. Quantify the trade-offs between sinks and diagonal patterns in terms of representational cost.
4. Account for the joint effects of skip connections, layer normalization, and value-output transformations on oversmoothing.

Methodology

The authors employ a combination of geometric analysis, theoretical proofs, and empirical verification on pretrained transformers (including GPT-2, LLaMA-3, Gemma, and Mistral).

1. Geometric Representation Analysis: The paper derives necessary and sufficient conditions for a group of tokens to attend exclusively to a sink token. This involves analyzing the alignment of embedding vectors in the query-key space.
2. Oversmoothing Analysis: The authors model a single attention step with skip connections and layer normalization. They derive a theoretical approximation for the expected pairwise cosine similarity of token representations as a function of attention density (parameterized by $\lambda$, interpolating between identity and uniform attention). This analysis incorporates the value-output transformation matrix $W$.
3. Equivalence Proofs: The paper proves that under mild geometric assumptions (specifically regarding the rank of input embeddings and their span), a "hard attention switch" (zero output) is mathematically equivalent to an attention sink.
4. Cost Comparison: The authors introduce a "soft attention switch" (allowing diagonal patterns where tokens attend to themselves) and compare the $\ell_2$ (Frobenius) and nuclear norm costs of representing sinks versus diagonal patterns. They analyze these costs under specific input-output distributions, including a generalized bigram-backcopy task.
5. Empirical Validation: The theoretical predictions are validated against empirical data from multiple models and datasets (TinyStories, WikiText, CodeSearchNet, etc.), measuring cosine similarities, attention uniformity, and rank distributions.

Key Contributions and Results

1. Geometric Conditions for Sink Representation

The paper establishes that for an attention layer to represent a sink, the context embeddings must lie in a specific geometric configuration relative to the sink embedding. Specifically, the set of vectors $\{z_0 - z_i\}$ (where $z_0$ is the sink and $z_i$ are context tokens) and the set of context embeddings $\{z_j\}$ must lie in half-spaces.

• Result: Empirical analysis of pretrained models (GPT-2, Llama 3.1, Gemma) confirms that in middle layers, the cosine similarity between BOS and all other tokens is non-positive, satisfying the condition that allows any attention head to produce a sink pattern.

2. Oversmoothing and Attention Density

The authors analyze whether sparse attention is strictly necessary to prevent oversmoothing.

• Theory: They prove that skip connections and value-output transformations ($W$) can theoretically counteract the mixing effect of uniform attention. However, this depends on specific conditions involving the trace of matrix products ($\beta \text{tr}(BW)$).
• Empirical Finding: In practice, these counteracting conditions are rarely met. In most heads of trained LLMs, increasing attention density (moving toward uniform attention) leads to an increase in token cosine similarity (oversmoothing). This confirms that sparse attention patterns are practically necessary to prevent oversmoothing in current models.

3. Equivalence of Sinks and Hard Attention Switch

The paper defines a hard attention switch as a mechanism where the output of an attention head on a fraction of tokens is identically zero.

• Theorem: Under mild geometric assumptions (verified in pretrained LLMs), a head implements a hard attention switch if and only if it implements an attention sink.
• Implication: Tokens cannot achieve a zero output by attending to a convex combination of past tokens that sums to zero, nor by attending to themselves with zero value. The only geometrically viable solution to "switch off" a head is to attend fully to the BOS token (the sink).

4. Sink vs. Diagonal Patterns (Soft Attention Switch)

The authors introduce a soft attention switch, which prohibits inter-token communication but allows tokens to attend to themselves (diagonal attention). This prevents oversmoothing without requiring a sink.

• Cost Analysis: The paper derives a quantitative comparison of the representational costs (nuclear norm of $W_{QK}$) for sinks versus diagonal patterns.
  • Sink Cost: Scales with $\kappa \sqrt{T}$ (low-rank structure).
  • Diagonal Cost: Scales with $\kappa T$ (full-rank structure).
• Result: Sinks are provably cheaper to represent than diagonal patterns across a wide range of hyperparameters and context lengths.
• Trade-off: While diagonal attention is cheaper in specific high-dimensional, low-rank scenarios, the theoretical analysis and empirical evidence (Figure 3) suggest that for typical backcopy tasks and generic input-output mappings, sinks incur lower parameter costs. This explains the prevalence of sinks over diagonal patterns in pretrained transformers.

Significance and Claims

The paper claims to close several gaps in the literature regarding attention mechanisms:

1. Geometric Clarity: It provides the first rigorous geometric conditions explaining how and when transformers can represent sinks.
2. Necessity of Sinks: It proves that sinks are not just a common artifact but are theoretically necessary for implementing a hard attention switch under general conditions.
3. Oversmoothing Mechanism: It clarifies that while architectural components could theoretically prevent oversmoothing with dense attention, trained models rely on sparse patterns (sinks) because the theoretical counter-mechanisms are rarely active in practice.
4. Cost Efficiency: It establishes that the prevalence of sinks over diagonal patterns is driven by representational efficiency. Sinks offer a lower-cost solution for switching off token communication, which is a critical function for preventing oversmoothing and enabling specific tasks like copy-paste mechanisms.

The authors conclude that the introduction of diagonal patterns as a theoretical alternative and the cost analysis of sinks versus diagonals provide a unified explanation for why attention layers in pretrained models often act like MLPs (when communication is not needed) and why sinks are the dominant mechanism for achieving this state.