Research map / Architectures and efficiency / Attention and transformers
On the Role of Attention Masks and LayerNorm in Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes rank collapse in transformer self-attention under general attention masks and LayerNorm, showing masks slow collapse while LayerNorm with orthogonal or proper value matrices can prevent collapse to rank one.
Problem
Rank collapse (token uniformity/oversmoothing) of token representations in multi-layer self-attention, which prior theory analyzed only for fully bidirectional attention and without LayerNorm (or with LayerNorm only under identity value matrices).
Approach
Treats self-attention as a discrete-time dynamical system with general masks described by directed graphs (quasi-strongly connected graphs and radius), proves exponential collapse for masked attention, and uses infinite products of nonnegative matrices and constructed counterexamples to show LayerNorm with suitable value matrices yields equilibria of any rank and prevents rank-one collapse.
Key results
- Pure masked self-attention with quasi-strongly connected masks shows exponential rank collapse, with rate (1-epsilon^r)^(t/r) that slows for larger graph radius r (local attention).
- With LayerNorm and orthogonal value matrices, collapse to a common point on the unit sphere still happens exponentially under quasi-strongly connected masks, but with non-negative initial cosine similarities.
- A causal-mask counterexample (N=2, d=2) and a general construction show LayerNorm dynamics admit at least 2^k equilibria of rank k for each k up to min(N,d), and non-rank-one convergence for a positive-measure set of inputs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use local or sliding-window attention instead of full attention | Self-attention networks · Wikipedia 128-token excerpts | token similarity μ(X) | Helped |
| Add LayerNorm | Randomly initialized self-attention networks · Wikipedia 128-token excerpts | token similarity μ(X) | Helped |
| Add LayerNorm alongside other components, including skip connections | Pretrained transformer models · Wikipedia 128-token excerpts | token representation stability | Helped |
| Add skip connections without the full transformer's other components | Self-attention networks · Deeper layers | token similarity μ(X) | Hurt |
| Add LayerNorm to counteract the instability from pure skip connections | Transformer components · Deeper layers | token representation stability | Helped |
| Use a smaller attention temperature d_QK | Pretrained self-attention networks · Different attention masks | initial rate of rank collapse | Helped |
| Use a smaller attention temperature d_QK | Pretrained self-attention networks · Global attention versus sparse masked attention | initial rate of rank collapse | Mixed |
| Use causal attention instead of full attention | LayerNorm self-attention with orthogonal value matrices · Rank collapse under attention masks | rate of rank collapse | Helped |
| Use LayerNorm with orthogonal value matrices | Self-attention dynamics · Strongly connected attention masks | token similarity μ(X) | Hurt |
| Use LayerNorm with properly chosen value matrices | Self-attention dynamics · A general class of input sequences | token rank | Helped |
| Add LayerNorm to pure self-attention | Self-attention networks · All attention masks in the numerical experiments | token similarity μ(X) | Mixed |
| Use self-attention with LayerNorm rather than without LayerNorm | Transformer token representations · Pretrained BERT, GPT2, T5, and ALBERT | token rank | Helped |
Surprising findings
- LayerNorm, previously hypothesized to play no role in rank collapse, can prevent rank collapse of self-attention dynamics with proper value matrices.
- Pretrained BERT, GPT2, T5, and ALBERT preserve full token rank across layers while tokens become anisotropic (concentrated in a narrow region).
Limitations
- Theoretical results for LayerNorm rank-collapse counterexamples are limited to specific constructions and small dimensions; general convergence to specific equilibria is not proven.
- Empirical validation uses only Wikipedia excerpts and pretrained encoder/decoder models, with no downstream task evaluation.
- Equilibria found are stable only in certain directions and have small regions of attraction.
Experimental setup
Datasets: Wikipedia excerpts (3000 samples of 128 tokens)
Metrics: mu(X) token collapse measure (Frobenius deviation from mean) minimal singular value average absolute pairwise cosine similarity
Baselines: pure self-attention without LayerNorm full bidirectional attention causal attention
Builds on: Dong et al. 2021 rank collapse of pure self-attention Geshkovski et al. continuous-time self-attention dynamics Wang et al. LayerNorm in MLPs isometry
Closest papers by meaning
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.92)
- The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training (ICML 2025 · similarity 0.90)
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.90)
- JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention (ICLR 2024 · similarity 0.89)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.89)
- On the Role of Hidden States of Modern Hopfield Network in Transformer (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.