[Paper Review] Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth

Yihe Dong, Jean-Baptiste Cordonnier+1 more|arXiv (Cornell University)|Mar 4, 2021

Neural Networks and ApplicationsComputer Science90 citations

TL;DR

The paper shows that pure self-attention networks (without skip connections or MLPs) cause outputs to converge to a rank-1 matrix doubly exponentially with depth; skip connections and MLPs counteract this collapse, analyzed via a path decomposition.

ABSTRACT

Attention-based architectures have become ubiquitous in machine learning, yet our understanding of the reasons for their effectiveness remains limited. This work proposes a new way to understand self-attention networks: we show that their output can be decomposed into a sum of smaller terms, each involving the operation of a sequence of attention heads across layers. Using this decomposition, we prove that self-attention possesses a strong inductive bias towards "token uniformity". Specifically, without skip connections or multi-layer perceptrons (MLPs), the output converges doubly exponentially to a rank-1 matrix. On the other hand, skip connections and MLPs stop the output from degeneration. Our experiments verify the identified convergence phenomena on different variants of standard transformer architectures.

Motivation & Objective

• Motivate a deeper understanding of why transformer-like architectures work beyond empirical success.
• Prove rank-collapse behavior of self-attention networks as depth increases.
• Introduce a path decomposition to analyze SANs as ensembles of shallow networks.
• Characterize how skip connections and MLPs counteract rank collapse.
• Verify theoretical results with experiments on common transformer variants.

Proposed method

• Decompose SAN outputs into a sum over paths through attention heads across layers.
• Show that each path converges to a rank-1 matrix with identical rows, leading to overall rank collapse.
• Derive convergence bounds (e.g., Theorem 2.2 and Theorem 2.3) for single-head and multi-head SANs.
• Analyze the role of architectural components (skip connections, MLPs, layer normalization) in mitigating rank collapse.
• Develop a path-based framework to interpret SANs as ensembles of shallow networks.
• Validate theory with experiments on architectures like BERT, ALBERT, and XLNet.

Experimental results

Research questions

• RQ1 Does pure self-attention cause rank collapse as depth increases?
• RQ2 How do architectural elements like skip connections, MLPs, and layer normalization influence rank collapse?
• RQ3 Can a path decomposition framework reveal the internal mechanisms of SANs and their inductive biases?
• RQ4 Do empirical results on standard transformers corroborate the theoretical predictions about rank convergence?
• RQ5 What are the practical implications for width-depth trade-offs and long-path utilization in transformers?

Key findings

• Self-attention networks without skip connections converge to a rank-1 output with identical rows, at a doubly exponential rate in depth.
• Skip connections dramatically diversify paths and prevent rank collapse in practice.
• MLPs slow down the convergence to rank-1 by increasing the Lipschitz constant, creating a tug-of-war with self-attention.
• Layer normalization does not mitigate rank collapse in this analysis.
• The path decomposition shows SANs behave like ensembles of shallow networks, with short paths carrying most expressiveness.
• Experiments on BERT, ALBERT, and XLNet confirm rapid rank collapse when skip connections are removed.