Specialization of softmax attention heads: insights from the high-dimensional single-location model
Not on the map. Why?
Abstract (arXiv)
Multi-head attention enables transformer models to represent multiple attention patterns simultaneously. Empirically, head specialization emerges in distinct stages during training, while many heads remain redundant and learn similar representations. We propose a theoretical model capturing this phenomenon, based on the multi-index and single-location regression frameworks. In the first part, we analyze the training dynamics of multi-head softmax attention under SGD, revealing an initial unspecialized phase followed by a multi-stage specialization phase in which different heads sequentially align with latent signal directions. In the second part, we study the impact of attention activation functions on performance. We introduce the Bayes-softmax attention, which achieves optimal prediction performance in this setting.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 27 of its 36 references that are in the map. First posted on arXiv 4 Mar 2026.
Builds on
27 of 36 references.
References · 27
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
- 2020End-to-End Object Detection with Transformerscited by 7,384
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
- 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 994
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
- 2023Vision Transformers Need Registerscited by 844
- 2019Are Sixteen Heads Really Better than One?cited by 803
- 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 266
- 2018A mathematical theory of semantic development in deep neural networkscited by 182
- 2023SGD learning on neural networks: leap complexity and saddle-to-saddle dynamicscited by 142
- 2023Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMscited by 121
- 2020Online stochastic gradient descent on non-convex losses from high-dimensional inferencecited by 119
15 more
- 2019Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setupcited by 112
- 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 71
- 2022High-dimensional limit theorems for SGD: Effective dynamics and critical scalingcited by 69
- 2024LLM Circuit Analyses Are Consistent Across Training and Scalecited by 60
- 2018The committee machine: Computational to statistical gaps in learning a two-layers neural networkcited by 58
- 2023Saddle-to-Saddle Dynamics in Diagonal Linear Networkscited by 54
- 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
- 2024Loss Landscape Degeneracy and Stagewise Development in Transformerscited by 39
- 2024Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficientcited by 39
- 2024Attention layers provably solve single-location regressioncited by 15
- 2025Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholdscited by 12
- 2025Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networkscited by 9
- 2025Statistical Advantage of Softmax Attention: Insights from Single-Location Regressioncited by 6
- 2025High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classificationcited by 6
- 2024Online Learning and Information Exponents: On The Importance of Batch size, and Time/Complexity Tradeoffscited by 2
Built on by
Citing papers per year: 2026: 1.
1 paper.