# On the Optimization and Generalization of Multi-head Attention

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

## Key facts

- Cited by 37 papers in the map (data through 14 Oct 2026).
- Builds on 62 of its 69 references that are in the map.
- First posted on arXiv 19 Oct 2023.

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

- [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)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [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)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [A Structured Self-attentive Sentence Embedding](https://mlsanity.com/paper/1703.03130/) (2017)
- [Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://mlsanity.com/paper/1905.13210/) (2019)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [A Decomposable Attention Model for Natural Language Inference](https://mlsanity.com/paper/1606.01933/) (2016)
- [Towards moderate overparameterization: global convergence guarantees for training shallow neural networks](https://mlsanity.com/paper/1902.04674/) (2019)
- [Long Short-Term Memory-Networks for Machine Reading](https://mlsanity.com/paper/1601.06733/) (2016)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [On the linearity of large non-linear models: when and why the tangent kernel is constant](https://mlsanity.com/paper/2010.01092/) (2020)
- [Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks](https://mlsanity.com/paper/1909.12292/) (2019)

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

- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?](https://mlsanity.com/paper/2402.15607/) (2024)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (2024)
- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [Prompting a Pretrained Transformer Can Be a Universal Approximator](https://mlsanity.com/paper/2402.14753/) (2024)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [Superiority of Multi-Head Attention in In-Context Linear Regression](https://mlsanity.com/paper/2401.17426/) (2024)
- [Implicit Optimization Bias of Next-Token Prediction in Linear Models](https://mlsanity.com/paper/2402.18551/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [Next-token prediction capacity: general upper bounds and a lower bound for transformers](https://mlsanity.com/paper/2405.13718/) (2024)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [The Expressibility of Polynomial based Attention Scheme](https://mlsanity.com/paper/2310.20051/) (2023)
- [A Theoretical Insight into Attack and Defense of Gradient Leakage in Transformer](https://mlsanity.com/paper/2311.13624/) (2023)
- [One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space](https://mlsanity.com/paper/2311.14652/) (2023)
- [On the Power of Convolution Augmented Transformer](https://mlsanity.com/paper/2407.05591/) (2024)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability](https://mlsanity.com/paper/2405.16845/) (2024)
- [On the Benefits of Rank in Attention Layers](https://mlsanity.com/paper/2407.16153/) (2024)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [Parameter-Efficient Fine-Tuning for Continual Learning: A Neural Tangent Kernel Perspective](https://mlsanity.com/paper/2407.17120/) (2024)
- [Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers](https://mlsanity.com/paper/2502.20681/) (2025)
- [Optimal Rates for Generalization of Gradient Descent for Deep ReLU Classification](https://mlsanity.com/paper/2510.02779/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization](https://mlsanity.com/paper/2505.15009/) (2025)
- [Finite-Time Analysis of Gradient Descent for Shallow Transformers](https://mlsanity.com/paper/2601.16514/) (2026)
- [Stochastic Gradient Descent with Momentum is Algorithmically Stable](https://mlsanity.com/paper/2605.28517/) (2026)

## Abstract (arXiv)

The training and generalization dynamics of the Transformer's core mechanism, namely the Attention mechanism, remain under-explored. Besides, existing analyses primarily focus on single-head attention. Inspired by the demonstrated benefits of overparameterization when training fully-connected networks, we investigate the potential optimization and generalization advantages of using multiple attention heads. Towards this goal, we derive convergence and generalization guarantees for gradient-descent training of a single-layer multi-head self-attention model, under a suitable realizability condition on the data. We then establish primitive conditions on the initialization that ensure realizability holds. Finally, we demonstrate that these conditions are satisfied for a simple tokenized-mixture model. We expect the analysis can be extended to various data-model and architecture variations.