Rank Collapse in Deep Transformers: Why Pure Attention Degenerates Doubly Exponentially and How Skip Connections Preserve Capacity

When the Transformer architecture was introduced in 2017 with the seminal paper "Attention Is All You Need", the central thesis was that recurrence and convolution could be completely discarded in favor of stacked self-attention mechanisms. However, theoretical analysis has shown that the title's premise is mathematically incomplete. Stacking pure self-attention layers in isolation does not produce an expressive deep model: it triggers a catastrophic failure mode known as rank collapse.

In a foundational study published at ICML 2021, researchers Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas demonstrated that without residual skip connections and multi-layer perceptron (MLP) blocks, the token representations in a deep self-attention network converge toward complete uniformity at a doubly exponential rate with respect to depth. As depth increases, all token vectors collapse to the exact same point in latent space, reducing the representation matrix to rank 1 and destroying all sequence information.

Understanding why pure attention fails geometrically, and how residual streams and feed-forward sub-layers prevent this collapse, is essential to understanding the structural design of modern foundation models.

The Geometry of Pure Self-Attention

To analyze how representations evolve across layers, consider an input matrix $X^{(0)} \in \mathbb{R}^{n \times d}$, representing a sequence of $n$ tokens embedded in a $d$-dimensional space. In a pure self-attention network without skip connections or feed-forward layers, the representation at layer $l+1$ is computed solely through self-attention:

$$P^{(l)} = \text{softmax}\left(\frac{X^{(l)} W_Q^{(l)} (X^{(l)} W_K^{(l)})^T}{\sqrt{d_k}}\right)$$

$$X^{(l+1)} = P^{(l)} X^{(l)} W_V^{(l)}$$

Here, $W_Q^{(l)}, W_K^{(l)} \in \mathbb{R}^{d \times d_k}$ and $W_V^{(l)} \in \mathbb{R}^{d \times d_v}$ denote the query, key, and value projection matrices at layer $l$.

The core vulnerability originates in the structure of the attention matrix $P^{(l)} \in \mathbb{R}^{n \times n}$. Because the softmax function normalizes each row independently, $P^{(l)}$ is a row-stochastic matrix: every entry satisfies $P_{i,j}^{(l)} \ge 0$, and the sum across each row is strictly equal to 1 ($\sum_{j=1}^n P_{i,j}^{(l)} = 1$).

Geometrically, multiplying the token matrix $X^{(l)}$ by a row-stochastic matrix $P^{(l)}$ computes a convex combination of the token vectors for each position:

$$x_i^{(l+1)} = \sum_{j=1}^n P_{i,j}^{(l)} x_j^{(l)} W_V^{(l)}$$

A convex combination of points always lies strictly within the convex hull of those points. Consequently, the convex hull of token representations shrinks at every successive layer:

$$\text{conv}(X^{(l+1)}) \subseteq \text{conv}(X^{(l)} W_V^{(l)})$$

This contraction mapping acts as a spatial smoothing operator. Each token moves closer to the weighted center of mass of all other tokens in the sequence.

The Doubly Exponential Decay Theorem

The degradation becomes fatal when self-attention is stacked recursively. In classical linear dynamic systems or Markov chains, convex averaging leads to exponential convergence toward a stationary state ($\mathcal{O}(\gamma^l)$ for some contraction factor $\gamma < 1$). Pure self-attention degenerates substantially faster: at a doubly exponential rate ($\mathcal{O}(\gamma^{2^l})$).

To formalize this, define the centroid vector $\bar{x}^{(l)} \in \mathbb{R}^d$ and the residual deviation matrix $R^{(l)} \in \mathbb{R}^{n \times d}$, which measures how far token representations deviate from total uniformity:

$$\bar{x}^{(l)} = \frac{1}{n} (X^{(l)})^T \mathbf{1}$$

$$R^{(l)} = X^{(l)} - \mathbf{1} (\bar{x}^{(l)})^T = \left(I - \frac{1}{n} \mathbf{1} \mathbf{1}^T\right) X^{(l)}$$

The rank of $X^{(l)}$ reflects the capacity of the model to distinguish distinct tokens. When $R^{(l)} = 0$, all tokens are identical, and $\text{rank}(X^{(l)}) = 1$.

In their 2021 proof, Dong, Cordonnier, and Loukas established that when weight matrices have bounded spectral norms, the residual norm $\|R^{(l)}\|$ obeys an upper bound governed by a double exponential:

$$\|R^{(l)}\| \le c \cdot \alpha^{2^l}$$

where $\alpha < 1$ and $c > 0$ are constants depending on the sequence length and projection norms.

The double exponential rate arises from a positive feedback loop between two interacting components:

1. Spatial Averaging (The Value Step): Multiplying by $P^{(l)}$ averages token positions, reducing the pairwise distances between token representations by a factor proportional to the non-uniformity of $P^{(l)}$.
2. Logit Homogenization (The Attention Step): The attention matrix $P^{(l)}$ is itself computed from pairwise inner products: $X^{(l)} W_Q (X^{(l)} W_K)^T$. When tokens move closer together ($X^{(l)} \approx \mathbf{1} (\bar{x}^{(l)})^T$), the variance of the attention logits shrinks toward zero.
3. Uniform Softmax Feedback: Applying softmax to nearly identical logits forces the attention distribution toward the uniform matrix: $P^{(l)} \to \frac{1}{n} \mathbf{1} \mathbf{1}^T$.
4. Total Collapse: When $P^{(l)}$ approaches the uniform distribution, the subsequent layer computes a uniform arithmetic mean over all tokens, collapsing the residual variance to zero in a single transition.

Because spatial contraction accelerates logit homogenization, and logit homogenization accelerates spatial contraction, the collapse compounds exponentially at every layer step. By layer 6 to 8, a pure self-attention network retains virtually zero token distinction regardless of sequence length.

Why Skip Connections Halt Rank Collapse

The standard Transformer architecture avoids rank collapse primarily through residual skip connections:

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

Residual connections fundamentally alter the algebraic topology of the computation graph. As demonstrated by Andreas Veit et al. (2016) and extended to self-attention by Dong et al., a deep network with residual connections can be unrolled into an ensemble of $2^L$ distinct paths of varying lengths:

$$X^{(L)} = X^{(0)} + \sum_{i} F_i(X^{(0)}) + \sum_{i,j} F_j(F_i(X^{(0)})) + \dots$$

This formulation provides two structural guarantees against rank collapse:

1. Identity Gradient Highway: The unperturbed identity path $X^{(0)}$ runs directly from input to output. The representation at layer $L$ retains an explicit linear component of the original high-rank embedding matrix, establishing a hard lower bound on matrix rank.
2. Binomial Path Distribution: The distribution of effective path depths across the $2^L$ computational sub-paths follows a binomial distribution centered at $L/2$. Because shallow paths do not undergo sufficient recursive attention multiplications to reach the collapse threshold, they deliver high-rank, diverse representations directly to deeper layers.

The Role of Feed-Forward Networks (MLPs)

While skip connections prevent complete rank-1 collapse, residual attention alone can still suffer from directional alignment and anisotropy, where representations become concentrated in a narrow cone of latent space. Feed-forward sub-layers (MLPs) provide the complementary countermeasure:

$$\text{FFN}(x) = W_2 \cdot \sigma(W_1 x + b_1) + b_2$$

Modern architectures replace standard activations with gated variants like SwiGLU:

$$\text{SwiGLU}(x) = W_2 \cdot (\text{swish}(W_1 x) \otimes W_3 x)$$

MLPs counteract representation collapse through three distinct mechanisms:

• Pointwise Non-Linearity: Self-attention performs linear combinations of token vectors across the sequence dimension. MLPs apply non-linear activations ($\sigma$) coordinate-wise across the feature dimension. This non-linear mapping projects co-linear or contracted tokens back into higher-dimensional curved manifolds, breaking token clustering.
• Subspace Expansion: The intermediate hidden dimension of the MLP ($d_{ff}$) is typically expanded to $4d$ (in standard Transformers) or $\frac{8}{3}d$ (in SwiGLU architectures). Expanding into a higher-dimensional space allows the network to untangle representations that have approached the same subspace during attention aggregation.
• Independent Channel Mixing: While attention mixes information across the sequence dimension ($n$), the MLP mixes information exclusively across the embedding dimension ($d$). Alternating between sequence-mixing and channel-mixing prevents either operator from driving the state space to degenerate attractors.

Architectural Implications for Modern Foundation Models

The mathematical reality of rank collapse dictates several foundational design decisions in frontier LLM architectures:

• Layer Normalization Placement: Pre-LayerNorm (Pre-LN) and RMSNorm place normalization inside the residual branch ($x + \text{Attn}(\text{RMSNorm}(x))$) rather than around the sum. This ensures the residual highway remains strictly linear and unscaled, preserving representation rank across hundreds of layers.
• QK-Normalization: Models such as Gemma 2 and Qwen 2.5 incorporate QK-Norm ($Q = \text{RMSNorm}(Q)$, $K = \text{RMSNorm}(K)$ before the dot product) to prevent logit scale explosion. This maintains stable attention entropy and prevents sudden drops into degenerate uniform or delta distributions.
• Parameter Allocation (Attention vs. MLP): Standard scaling recipes allocate approximately one-third of model parameters to self-attention mechanisms and two-thirds to feed-forward networks. This ratio balances the routing capacity of attention with the non-linear feature transformation capacity of MLPs.

Numerical Demonstration: Pure Attention vs. Residual Architectures

The following Python script simulates the singular value spectrum and effective rank across 20 layers for pure self-attention compared to a standard Transformer block with residual connections and non-linear MLPs.

import numpy as np

def compute_effective_rank(X):
    """
    Computes the continuous entropy-based effective rank (Roy & Vetterli, 2007)
    from the singular values of token representation matrix X.
    """
    # Center matrix across sequence
    X_centered = X - np.mean(X, axis=0, keepdims=True)
    _, s, _ = np.linalg.svd(X_centered, full_matrices=False)
    s_norm = s / np.sum(s)
    s_norm = s_norm[s_norm > 1e-12]
    entropy = -np.sum(s_norm * np.log(s_norm))
    return np.exp(entropy)

def softmax(Z):
    exp_Z = np.exp(Z - np.max(Z, axis=-1, keepdims=True))
    return exp_Z / np.sum(exp_Z, axis=-1, keepdims=True)

# Simulation parameters
np.random.seed(42)
seq_len = 64
d_model = 128
n_layers = 20

# Initialize random sequence embeddings
X_pure = np.random.randn(seq_len, d_model)
X_trans = X_pure.copy()

print(f"Layer 0 Initial Effective Rank: {compute_effective_rank(X_pure):.4f}")
print("-" * 60)
print(f"{'Layer':<6} | {'Pure Attention Rank':<22} | {'Transformer Rank (Res+MLP)':<25}")
print("-" * 60)

for layer in range(1, n_layers + 1):
    # Weight matrices with unit spectral scaling
    W_q = np.random.randn(d_model, d_model) / np.sqrt(d_model)
    W_k = np.random.randn(d_model, d_model) / np.sqrt(d_model)
    W_v = np.random.randn(d_model, d_model) / np.sqrt(d_model)
    W_1 = np.random.randn(d_model, d_model * 4) / np.sqrt(d_model)
    W_2 = np.random.randn(d_model * 4, d_model) / np.sqrt(d_model * 4)

    # 1. Pure Attention Path
    Q_p = X_pure @ W_q
    K_p = X_pure @ W_k
    V_p = X_pure @ W_v
    attn_p = softmax((Q_p @ K_p.T) / np.sqrt(d_model))
    X_pure = attn_p @ V_p

    # 2. Full Transformer Block Path (Pre-LN + Residual + MLP)
    # RMSNorm helper
    norm = lambda x: x / (np.sqrt(np.mean(x**2, axis=-1, keepdims=True)) + 1e-6)
    
    # Attention sub-layer with skip connection
    X_norm = norm(X_trans)
    Q_t = X_norm @ W_q
    K_t = X_norm @ W_k
    V_t = X_norm @ W_v
    attn_t = softmax((Q_t @ K_t.T) / np.sqrt(d_model))
    X_trans = X_trans + (attn_t @ V_t)

    # MLP sub-layer with skip connection (GELU non-linearity)
    X_mlp_in = norm(X_trans)
    gelu = lambda z: 0.5 * z * (1 + np.tanh(np.sqrt(2 / np.pi) * (z + 0.044715 * z**3)))
    mlp_out = gelu(X_mlp_in @ W_1) @ W_2
    X_trans = X_trans + mlp_out

    if layer in [1, 2, 3, 4, 5, 8, 12, 16, 20]:
        r_pure = compute_effective_rank(X_pure)
        r_trans = compute_effective_rank(X_trans)
        print(f"{layer:<6} | {r_pure:<22.4f} | {r_trans:<25.4f}")

Running this simulation demonstrates the collapse trajectory: by layer 4, pure attention rank drops precipitously toward 1.0 (complete token co-linearity), whereas the residual and MLP-equipped Transformer maintains stable effective dimensionality across arbitrary depth.

Self-attention provides dynamic token routing, but residual streams and feed-forward networks provide the geometric stability that prevents deep models from collapsing into trivial representations.

Sources

• Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth (arXiv:2103.03404)
• Residual Networks Behave Like Ensembles of Relatively Shallow Networks (arXiv:1605.06431)
• GLU Variants Improve Transformer (arXiv:2002.05202)
• The Effective Rank: A Measure of Effective Dimensionality (Roy & Vetterli, 2007)