# On the Role of Attention Masks and LayerNorm in Transformers

arXiv: https://arxiv.org/abs/2405.18781 · HTML: https://mlsanity.com/paper/2405.18781/

## Key facts

- Cited by 48 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 41 references that are in the map.
- First posted on arXiv 29 May 2024.

## Builds on (most cited first, 30 of 35)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Representation Degeneration Problem in Training Natural Language Generation Models](https://mlsanity.com/paper/1907.12009/) (2019)
- [Neighborhood Attention Transformer](https://mlsanity.com/paper/2204.07143/) (2022)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse](https://mlsanity.com/paper/2206.03126/) (2022)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Fast Attention Requires Bounded Entries](https://mlsanity.com/paper/2302.13214/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Scaling Transformer to 1M tokens and beyond with RMT](https://mlsanity.com/paper/2304.11062/) (2023)
- [O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformers](https://mlsanity.com/paper/2006.04862/) (2020)
- [Revisiting Over-smoothing in BERT from the Perspective of Graph](https://mlsanity.com/paper/2202.08625/) (2022)
- [Anisotropy Is Inherent to Self-Attention in Transformers](https://mlsanity.com/paper/2401.12143/) (2024)
- [Demystifying Oversmoothing in Attention-Based Graph Neural Networks](https://mlsanity.com/paper/2305.16102/) (2023)

## Built on by (most cited first, 30 of 48)

- [Why do LLMs attend to the first token?](https://mlsanity.com/paper/2504.02732/) (2025)
- [On the Emergence of Position Bias in Transformers](https://mlsanity.com/paper/2502.01951/) (2025)
- [Attention Sinks and Compression Valleys in LLMs are Two Sides of the Same Coin](https://mlsanity.com/paper/2510.06477/) (2025)
- [Dynamic metastability in the self-attention model](https://mlsanity.com/paper/2410.06833/) (2024)
- [Hallucination Detection in LLMs Using Spectral Features of Attention Maps](https://mlsanity.com/paper/2502.17598/) (2025)
- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (2024)
- [Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing](https://mlsanity.com/paper/2501.00658/) (2024)
- [The Asymptotic Behavior of Attention in Transformers](https://mlsanity.com/paper/2412.02682/) (2024)
- [Pay Attention to Attention Distribution: A New Local Lipschitz Bound for Transformers](https://mlsanity.com/paper/2507.07814/) (2025)
- [All You Need is One: Capsule Prompt Tuning with a Single Vector](https://mlsanity.com/paper/2510.16670/) (2025)
- [Optimizing Rank for High-Fidelity Implicit Neural Representations](https://mlsanity.com/paper/2512.14366/) (2025)
- [Transformer Normalisation Layers and the Independence of Semantic Subspaces](https://mlsanity.com/paper/2406.17837/) (2024)
- [Attention-Only Transformers via Unrolled Subspace Denoising](https://mlsanity.com/paper/2506.03790/) (2025)
- [Rethinking Layer-wise Model Merging through Chain of Merges](https://mlsanity.com/paper/2508.21421/) (2025)
- [DoPE: Denoising Rotary Position Embedding](https://mlsanity.com/paper/2511.09146/) (2025)
- [Lambda-Skip Connections: the architectural component that prevents Rank Collapse](https://mlsanity.com/paper/2410.10609/) (2024)
- [Revisiting [CLS] and Patch Token Interaction in Vision Transformers](https://mlsanity.com/paper/2602.08626/) (2026)
- [Attention Drift: What Autoregressive Speculative Decoding Models Learn](https://mlsanity.com/paper/2605.09992/) (2026)
- [ReLU's Revival: On the Entropic Overload in Normalization-Free Large Language Models](https://mlsanity.com/paper/2410.09637/) (2024)
- [AERO: Entropy-Guided Framework for Private LLM Inference](https://mlsanity.com/paper/2410.13060/) (2024)
- [NerVE: Nonlinear Eigenspectrum Dynamics in LLM Feed-Forward Networks](https://mlsanity.com/paper/2603.06922/) (2026)
- [MODA: MOdular Duplex Attention for Multimodal Perception, Cognition, and Emotion Understanding](https://mlsanity.com/paper/2507.04635/) (2025)
- [Stochastic Parroting in Temporal Attention -- Regulating the Diagonal Sink](https://mlsanity.com/paper/2602.10956/) (2026)
- [How Data Shapes RoPE Frequency Usage: From Positional Scale Matching to Length Generalization](https://mlsanity.com/paper/2607.07678/) (2026)
- [Transforming Rank: How Architecture Navigates the Spectral Pathologies of Depth](https://mlsanity.com/paper/2607.14018/) (2026)
- [Provable optimal transport with transformers: The essence of depth and prompt engineering](https://mlsanity.com/paper/2410.19931/) (2024)
- [TTrace: Lightweight Error Checking and Diagnosis for Distributed Training](https://mlsanity.com/paper/2506.09280/) (2025)
- [Understanding the Failure Modes of Transformers through the Lens of Graph Neural Networks](https://mlsanity.com/paper/2512.09182/) (2025)
- [Memory as Resonance: A Biomimetic Architecture for Infinite Context Memory on Ergodic Phonetic Manifolds](https://mlsanity.com/paper/2512.20245/) (2025)
- [Rank, Head-Channel Non-Identifiability, and Symmetry Breaking: A Precise Analysis of Representational Collapse in Transformers](https://mlsanity.com/paper/2604.23681/) (2026)

## Abstract (arXiv)

Self-attention is the key mechanism of transformers, which are the essential building blocks of modern foundation models. Recent studies have shown that pure self-attention suffers from an increasing degree of rank collapse as depth increases, limiting model expressivity and further utilization of model depth. The existing literature on rank collapse, however, has mostly overlooked other critical components in transformers that may alleviate the rank collapse issue. In this paper, we provide a general analysis of rank collapse under self-attention, taking into account the effects of attention masks and layer normalization (LayerNorm). In particular, we find that although pure masked attention still suffers from exponential collapse to a rank one subspace, sparse or local masked attention can provably slow down the collapse rate. In the case of self-attention with LayerNorm, we first show that for certain classes of value matrices, collapse to a rank one subspace still happens exponentially. However, through construction of nontrivial counterexamples, we then establish that with proper choice of value matrices, a general class of sequences may not converge to a rank one subspace, and the self-attention dynamics with LayerNorm can simultaneously possess a rich set of equilibria with any possible rank between one and full. Our result refutes the previous hypothesis that LayerNorm plays no role in the rank collapse of self-attention and suggests that self-attention with LayerNorm constitutes a much more expressive, versatile nonlinear dynamical system than what was originally thought.