mlsanityData to 7 Oct 2026

On the Optimization and Generalization of Multi-head Attention

arXiv 2310.12680 ↗First posted 19 Oct 2023cs.LG math.OC stat.MLcited by 37

Not on the map. Why?

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.

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

Builds on

62 of 69 references.

References · 62

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,456
  4. 2020Language Models are Few-Shot Learnerscited by 32,199
  5. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  6. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  7. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  8. 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
  9. 2020Training data-efficient image transformers & distillation through attentioncited by 4,560
  10. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
  11. 2018A Convergence Theory for Deep Learning via Over-Parameterizationcited by 784
  12. 2022Transformers learn in-context by gradient descentcited by 714
38 more
  1. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  2. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
  3. 2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 572
  4. 2019Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networkscited by 562
  5. 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 411
  6. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 404
  7. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  8. 2017A Structured Self-attentive Sentence Embeddingcited by 318
  9. 2019Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networkscited by 221
  10. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 220
  11. 2016A Decomposable Attention Model for Natural Language Inferencecited by 214
  12. 2019Towards moderate overparameterization: global convergence guarantees for training shallow neural networkscited by 176
  13. 2016Long Short-Term Memory-Networks for Machine Readingcited by 161
  14. 2023Birth of a Transformer: A Memory Viewpointcited by 154
  15. 2021Inductive Biases and Variable Creation in Self-Attention Mechanismscited by 144
  16. 2023Representational Strengths and Limitations of Transformerscited by 132
  17. 2020On the linearity of large non-linear models: when and why the tangent kernel is constantcited by 129
  18. 2019Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networkscited by 120
  19. 2020Fine-Grained Analysis of Stability and Generalization for Stochastic Gradient Descentcited by 118
  20. 2020Infinite attention: NNGP and NTK for deep attention networkscited by 117
  21. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  22. 2022Vision Transformers provably learn spatial structurecited by 92
  23. 2020Tight Bounds on the Smallest Eigenvalue of the Neural Tangent Kernel for Deep ReLU Networkscited by 87
  24. 2019How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?cited by 84
  25. 2020O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformerscited by 82
  26. 2018Risk and parameter convergence of logistic regressioncited by 81
  27. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  28. 2023Transformers as Support Vector Machinescited by 67
  29. 2020Global Convergence of Deep Networks with One Wide Layer Followed by Pyramidal Topologycited by 66
  30. 2023On the Role of Attention in Prompt-tuningcited by 49
  31. 2021A Local Convergence Theory for Mildly Over-Parameterized Two-Layer Neural Networkcited by 47
  32. 2021On the Expressive Power of Self-Attention Matricescited by 44
  33. 2023Memorization Capacity of Multi-Head Attention in Transformerscited by 39
  34. 2013Margins, Shrinkage, and Boostingcited by 35
  35. 2020The Effects of Mild Over-parameterization on the Optimization Landscape of Shallow ReLU Neural Networkscited by 33
  36. 2022Unraveling Attention via Convex Duality: Analysis and Interpretations of Vision Transformerscited by 30
  37. 2020Gradient Methods Never Overfit On Separable Datacited by 29
  38. 2022Stability vs Implicit Bias of Gradient Methods on Separable Data and Beyondcited by 29

50 most cited of 62 All 62 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 3; 2024: 19; 2025: 9; 2026: 6.

37 papers.

Sort:

By area:

  1. 2024Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMscited by 54
  2. 2024In-Context Learning with Representations: Contextual Generalization of Trained Transformerscited by 36
  3. 2024How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?cited by 35
  4. 2024Mechanics of Next Token Prediction with Self-Attentioncited by 34
  5. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31
  6. 2024Prompting a Pretrained Transformer Can Be a Universal Approximatorcited by 25
  7. 2024Superiority of Multi-Head Attention in In-Context Linear Regressioncited by 17
  8. 2024Implicit Regularization of Gradient Flow on One-Layer Softmax Attentioncited by 16
  9. 2024Implicit Optimization Bias of Next-Token Prediction in Linear Modelscited by 14
  10. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysiscited by 14
  11. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  12. 2024Next-token prediction capacity: general upper bounds and a lower bound for transformerscited by 11
25 more
  1. 2025Gating is Weighting: Understanding Gated Linear Attention through In-context Learningcited by 11
  2. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  3. 2024Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selectioncited by 9
  4. 2023The Expressibility of Polynomial based Attention Schemecited by 7
  5. 2023A Theoretical Insight into Attack and Defense of Gradient Leakage in Transformercited by 6
  6. 2023One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Spacecited by 6
  7. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  8. 2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5
  9. 2024On the Power of Convolution Augmented Transformercited by 5
  10. 2024Learning Linear Attention in Polynomial Timecited by 5
  11. 2024On the Benefits of Rank in Attention Layerscited by 4
  12. 2025Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformerscited by 3
  13. 2025Optimal Rates for Generalization of Gradient Descent for Deep ReLU Classificationcited by 3
  14. 2024Parameter-Efficient Fine-Tuning for Continual Learning: A Neural Tangent Kernel Perspectivecited by 2
  15. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  16. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalizationcited by 1
  17. 2025Stability-based Generalization Analysis of Randomized Coordinate Descent for Pairwise Learningcited by 0
  18. 2025Generalization and Optimization of SGD with Lookaheadcited by 0
  19. 2025Minimax Rates for Learning Pairwise Interactions in Attention-Style Modelscited by 0
  20. 2025Vocabulary In-Context Learning in Transformers: Benefits of Positional Encodingcited by 0
  21. 2026Finite-Time Analysis of Gradient Descent for Shallow Transformerscited by 0
  22. 2026Towards Initialization-dependent and Non-vacuous Generalization Bounds for Overparameterized Shallow Neural Networkscited by 0
  23. 2026Learning Theory of the SVRG: Generalization and Convergence Analysiscited by 0
  24. 2026Stochastic Gradient Descent with Momentum is Algorithmically Stablecited by 0
  25. 2026Faster Query-Key Learning Sharpens Attention in Self-Attention Modelscited by 0