Rank Collapse in Transformer Architectures: Why Pure Attention Decays Doubly Exponentially and How Residuals Restore Expressivity

The foundational premise of the Transformer architecture is that self-attention acts as a dynamic routing engine, computing pairwise affinities across sequence positions to contextualize representations. However, theoretical analysis reveals an underlying instability: stripped of its auxiliary components, a pure self-attention network rapidly destroys its own expressive capacity.

Without residual connections and non-linear feed-forward networks (MLPs), stacked self-attention layers suffer from severe rank collapse. As proven by Dong et al. (ICML 2021), the token representation matrix in a pure self-attention network converges toward a rank-1 matrix doubly exponentially with respect to network depth. Under this regime, every token in the sequence converges to an identical vector, eliminating all positional and semantic distinctions within a handful of layers.

Understanding why self-attention causes this hyper-accelerated degeneration—and how skip connections and non-linearities counteract it—provides fundamental insight into why modern large language models require specific architectural scaffolding to remain trainable at scale.

The Geometry of Self-Attention: Row-Stochastic Contraction

To analyze rank collapse, consider an input sequence matrix $X \in \mathbb{R}^{n \times d}$, where $n$ is the sequence length and $d$ is the hidden dimension. In a standard multi-head or single-head self-attention layer, queries and keys compute an affinity matrix passed through a row-wise softmax operator:

$$P = \text{softmax}\left(\frac{X W_Q W_K^\top X^\top}{\sqrt{d_k}}\right) \in \mathbb{R}^{n \times n}$$

$$\text{Attn}(X) = P X W_V W_O$$

The defining mathematical characteristic of the attention matrix $P$ is that it is strictly row-stochastic: all entries are non-negative, and each row sums exactly to one ($\sum_{j=1}^n P_{i,j} = 1$ for all $i$).

When $X$ is multiplied on the left by a row-stochastic matrix $P$, the $i$-th row of the resulting matrix is a convex combination (weighted average) of all row vectors in $X$:

$$(P X)_i = \sum_{j=1}^n P_{i,j} X_j$$

In convex geometry, any weighted average of a set of points lies strictly within their convex hull. Consequently, applying a row-stochastic matrix is a non-expansive contraction mapping that pulls all token vectors toward the geometric center of the sequence representation.

In classical Markov chains and linear dynamical systems, multiplying by fixed transition matrices leads to geometric mixing, where states converge to a stationary distribution at a singly exponential rate:

$$\|X^{(l)} - \mathbf{1} x^\top\| \le \lambda^l \|X^{(0)} - \mathbf{1} x^\top\|$$

where $\lambda < 1$ corresponds to the subdominant eigenvalue of the transition matrix and $\mathbf{1} x^\top$ represents the rank-1 stationary consensus matrix. However, self-attention exhibits a far more destructive convergence dynamic.

Why Self-Attention Collapses Doubly Exponentially

The fundamental reason self-attention degenerates faster than classical linear systems is that the transition matrix $P(X)$ is data-dependent and parameterized by the tokens themselves.

To quantify token diversity across layers, Dong et al. decomposed the token matrix $X$ into a rank-1 sequence average and a residual diversity component $R$:

$$X = \mathbf{1} x^\top + R$$

where $\mathbf{1} \in \mathbb{R}^n$ is the vector of ones, $x = \frac{1}{n} X^\top \mathbf{1} \in \mathbb{R}^d$ is the centroid representation, and $R \in \mathbb{R}^{n \times d}$ is the residual matrix satisfying $\mathbf{1}^\top R = 0$. The Frobenius norm $\|R\|_F$ directly measures the total variance and distinctiveness among tokens.

When expanding the attention logits with this decomposition, the rank-1 mean component factors out due to the shift-invariance property of the softmax function ($\text{softmax}(M + \mathbf{1} c^\top) = \text{softmax}(M)$ for any vector $c$):

$$X W_Q W_K^\top X^\top = (\mathbf{1} x^\top + R) W_Q W_K^\top (\mathbf{1} x^\top + R)^\top$$

$$P = \text{softmax}\left(\frac{R W_Q W_K^\top R^\top + \mathbf{1} r^\top}{\sqrt{d_k}}\right) = \text{softmax}\left(\frac{R W_Q W_K^\top R^\top}{\sqrt{d_k}}\right)$$

This formulation exposes the positive feedback loop driving collapse:

1. As tokens undergo convex averaging, the residual variance $\|R\|$ decreases.
2. Because the logits scale with $R W_Q W_K^\top R^\top$, smaller $\|R\|$ forces the attention logits toward zero.
3. In the zero-logit limit, the softmax output $P$ approaches the uniform matrix $\frac{1}{n} \mathbf{1} \mathbf{1}^\top$.
4. A uniform attention matrix assigns equal weight to every token, maximizing over-smoothing and collapsing $\|R\|$ in the subsequent layer by another factor of $\|R\|^2$.

[Token Representations X^(l)] ──> Convex Averaging P(X)X ──> Reduced Residual ||R||
            ▲                                                        │
            │                                                        ▼
            └──── Uniform Attention P ≈ (1/n)11^T ◄── Quadratic Logit Shrinkage R W_QK R^T

Mathematically, this feedback loop establishes a cubic contraction recurrence between successive layers:

$$\|R^{(l+1)}\|_F \le C \cdot \|R^{(l)}\|_F^3$$

for a constant $C$ determined by weight norms and sequence parameters. Solving this cubic recurrence demonstrates that the residual norm decays according to:

$$\|R^{(l)}\|_F \le \frac{1}{\sqrt{C}} \left( \sqrt{C} \|R^{(0)}\|_F \right)^{3^l}$$

This is a doubly exponential convergence rate with base 3. In a 12-layer pure self-attention network without skip connections, token diversity $\|R\|_F$ drops to machine precision within 5 to 7 layers, leaving every token identical and reducing the effective rank of the sequence representation to exactly 1.

Counteracting Forces: How Transformers Prevent Degeneration

Standard Transformer architectures do not suffer from catastrophic rank collapse in practice because they combine self-attention with two structural counter-mechanisms: residual connections (skip connections) and non-linear feed-forward networks (MLPs).

1. Residual Connections as Information Highways

The standard Transformer layer introduces an identity skip connection around the self-attention block:

$$X^{(l+1)} = X^{(l)} + \text{Attn}(X^{(l)})$$

Through the lens of ensemble path expansion (Veit et al., 2016), an $L$-layer Transformer can be represented as an unrolled summation over $2^L$ distinct paths of varying length. The shortest paths—those composed primarily of identity skip connections—bypass the contractive attention operators entirely.

Dong et al. proved that adding skip connections establishes a strictly positive lower bound on the residual norm:

$$\|R^{(l)}\|_F \ge (1 - \gamma)^l \|R^{(0)}\|_F > 0$$

where $\gamma < 1$ is an upper bound on the contraction introduced by attention. By guaranteeing that unmixed, full-rank token representations propagate directly from early layers to the network head, skip connections prevent asymptotic rank decay.

2. Non-Linear MLPs Break Convex Hulls

While self-attention mixes representations across tokens (row-wise operations), feed-forward networks apply non-linear transformations independently to each token (column-wise operations):

$$\text{MLP}(X) = \sigma(X W_1 + b_1) W_2 + b_2$$

Because the activation function $\sigma$ (such as GELU or SwiGLU) is non-linear, it breaks the affine invariance of convex combinations. The MLP expands, rotates, and disperses the token vectors across high-dimensional space, pushing points outward from the sequence centroid.

Dong et al. showed that introducing MLPs without skip connections slows the rate of rank collapse from doubly exponential ($3^l$) to singly exponential ($\lambda^l$). When both residual connections and MLPs are present, rank collapse is suppressed entirely.

• Pure Self-Attention (SAN): Converges doubly exponentially ($\mathcal{O}(c^{3^l})$) toward rank 1, resulting in total representation collapse within 5 to 7 layers.
• Attention + MLPs (No Skips): Slows convergence to singly exponential ($\mathcal{O}(\lambda^l)$), but still suffers severe representation degradation at depth.
• Attention + Skip Connections: Establishes a strictly positive lower bound on representation variance, preserving full rank across arbitrary depth.
• Full Transformer (Attn + MLP + Skips): Maintains maximal rank and high feature expressivity through complementary row-wise routing and column-wise non-linear dispersion.

Signal Propagation: Rank Collapse vs. Entropy Collapse

While skip connections prevent rank collapse, balancing deep Transformer architectures introduces a secondary optimization failure mode. As demonstrated by Noci et al. (NeurIPS 2022) in their study of signal propagation, Transformer initialization must navigate between two distinct collapse regimes:

1. Rank Collapse (Token Uniformity): Occurs when skip connection strength is insufficient relative to attention layer depth, causing pairwise cosine similarity $\frac{\langle x_i, x_j \rangle}{\|x_i\| \|x_j\|} \to 1$.
2. Entropy Collapse (Attention Saturation): Occurs when query-key product variance $\|W_Q W_K^\top\|$ is excessively large. Under large logit variance, the softmax function approaches an indicator (one-hot argmax) function, concentrating the entire attention distribution onto a single token (such as the initial sequence delimiter or an attention sink). In this regime, gradients for all non-attended tokens vanish.

       [Low Query-Key Variance]                   [High Query-Key Variance]
         Uniform Softmax Matrix                     Peaked / One-Hot Softmax
                   │                                          │
                   ▼                                          ▼
            Rank Collapse                              Entropy Collapse
         (Token Uniformity)                         (Vanishing Gradients)
                   \                                          /
                    \─── [Trainable Balanced Window] ────────/

Signal propagation theory demonstrates that maintaining stable forward and backward signal flow in 100+ layer models requires bounding both token similarity and attention entropy. This dynamic explains why modern frontier models implement:

• Pre-Layer Normalization and RMSNorm: Normalizing inputs prior to attention blocks ensures input variances do not scale exponentially with depth.
• QK-Normalization (QK-Norm): Explicitly applying $L_2$ or RMS normalization to query and key vectors (Dehghani et al., 2023) prevents attention logits from growing with depth, eliminating entropy collapse.
• Residual Stream Scaling ($\alpha$-Scaling / DeepNorm): Down-scaling residual branches by factors proportional to $1/\sqrt{2L}$ (Wang et al., 2022) preserves signal variance across extreme depths.

Over-Smoothing in Vision and Graph Transformers

The implications of rank collapse extend across all architectures utilizing self-attention:

• Vision Transformers (ViTs): In pure vision transformers applied to dense patch grids, over-smoothing manifests as patch representation collapse in deeper layers (Gong et al., 2021). Without proper architectural constraints, background patches and fine-grained visual features merge into a single global background vector.
• Graph Transformers and GNNs: Standard Graph Neural Networks suffer from graph over-smoothing due to repeated Laplacian matrix multiplications. Because dense self-attention is mathematically equivalent to a fully connected graph with dynamic edge weights, graph transformers exhibit accelerated rank collapse unless augmented with virtual nodes, structural encodings, or strong residual scaling.

Key Takeaways

• Pure self-attention is an aggressive over-smoother: Because attention matrices are row-stochastic, multiplying by $P$ performs convex averaging that contracts representations toward the sequence centroid.
• Collapse occurs doubly exponentially: Because attention weights depend quadratically on token differences, reducing token variance flattens attention logits, accelerating contraction at a cubic rate ($3^l$).
• Skip connections provide the essential lower bound: Identity residual connections carry uncompressed, full-rank token representations directly through the network, preventing asymptotic rank decay.
• MLPs break convex hulls: Non-linear activations operate per-token to disperse vectors across high-dimensional space, disrupting the contractive geometric forces of attention.
• Deep scaling requires balancing rank and entropy: Stable training requires preventing both over-smoothing (rank collapse) and over-peaking (entropy collapse) through techniques like RMSNorm, QK-Norm, and residual stream scaling.

Sources

• Dong, Y., Cordonnier, J. B., & Loukas, A. (2021). Attention is not all you need: pure attention loses rank doubly exponentially with depth. International Conference on Machine Learning (ICML).
• Noci, L., Anagnostidis, S., Biggio, L., Orvieto, A., Singh, S. P., & Lucchi, A. (2022). Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse. Advances in Neural Information Processing Systems (NeurIPS).
• Bhojanapalli, S., Chakrabarti, A., Glasner, D., Kang, D., Passos, A., & Veit, A. (2020). Low-Rank Bottleneck in Multi-head Attention Models. International Conference on Machine Learning (ICML).
• Veit, A., Wilber, M. J., & Belongie, S. (2016). Residual Networks Behave Like Ensembles of Relatively Shallow Networks. Advances in Neural Information Processing Systems (NeurIPS).
• Wang, H., Ma, S., Dong, L., Huang, S., Wang, W., & Wei, F. (2022). DeepNet: Scaling Transformers to 1,000 Layers. arXiv preprint arXiv:2203.00555.
• Dehghani, M., et al. (2023). Scaling Vision Transformers to 22 Billion Parameters. International Conference on Machine Learning (ICML).
• Gong, C., et al. (2021). Improve Vision Transformers Training by Preventing Rank Collapse. arXiv preprint arXiv:2103.17239.