# Provably learning a multi-head attention layer

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 80 references that are in the map.
- First posted on arXiv 6 Feb 2024.

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

- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Textbooks Are All You Need](https://mlsanity.com/paper/2306.11644/) (2023)
- [Textbooks Are All You Need II: phi-1.5 technical report](https://mlsanity.com/paper/2309.05463/) (2023)
- [TinyStories: How Small Can Language Models Be and Still Speak Coherent English?](https://mlsanity.com/paper/2305.07759/) (2023)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Theoretical Limitations of Self-Attention in Neural Sequence Models](https://mlsanity.com/paper/1906.06755/) (2019)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Saturated Transformers are Constant-Depth Threshold Circuits](https://mlsanity.com/paper/2106.16213/) (2021)
- [Convergence Analysis of Two-layer Neural Networks with ReLU Activation](https://mlsanity.com/paper/1705.09886/) (2017)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [Understanding and Improving Transformer From a Multi-Particle Dynamic System Point of View](https://mlsanity.com/paper/1906.02762/) (2019)
- [Self-Attention Networks Can Process Bounded Hierarchical Languages](https://mlsanity.com/paper/2105.11115/) (2021)
- [Statistically Meaningful Approximation: a Case Study on Approximating Turing Machines with Transformers](https://mlsanity.com/paper/2107.13163/) (2021)
- [Recovery Guarantees for One-hidden-layer Neural Networks](https://mlsanity.com/paper/1706.03175/) (2017)
- [Formal Language Recognition by Hard Attention Transformers: Perspectives from Circuit Complexity](https://mlsanity.com/paper/2204.06618/) (2022)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [Learning One-hidden-layer Neural Networks with Landscape Design](https://mlsanity.com/paper/1711.00501/) (2017)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)

## Built on by (most cited first, 21 of 21)

- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [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)
- [Measure-to-measure interpolation using Transformers](https://mlsanity.com/paper/2411.04551/) (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)
- [Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis](https://mlsanity.com/paper/2410.09605/) (2024)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [Transformers Meet In-Context Learning: A Universal Approximation Theory](https://mlsanity.com/paper/2506.05200/) (2025)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [Transformers Provably Learn Sparse XOR with Polylogarithmic Parameters](https://mlsanity.com/paper/2502.07553/) (2025)
- [Provably Learning from Modern Language Models via Low Logit Rank](https://mlsanity.com/paper/2512.09892/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Provably Learning Multi-Head Attention with Queries](https://mlsanity.com/paper/2608.03294/) (2026)
- [Provably Learning Attention with Queries](https://mlsanity.com/paper/2601.16873/) (2026)
- [Convergent Stochastic Training of Multi-Headed Attention and Understanding LoRA](https://mlsanity.com/paper/2605.07959/) (2026)

## Abstract (arXiv)

The multi-head attention layer is one of the key components of the transformer architecture that sets it apart from traditional feed-forward models. Given a sequence length k, attention matrices Theta1,ldots,Thetam∈ℝd× d, and projection matrices W1,ldots,Wm∈ℝd× d, the corresponding multi-head attention layer F: ℝk× d→ ℝk× d transforms length-k sequences of d-dimensional tokens X∈ℝk× d via F(X) triangleq summi=1 softmax(XThetaiX^top)XWi. In this work, we initiate the study of provably learning a multi-head attention layer from random examples and give the first nontrivial upper and lower bounds for this problem: - Provided Wi, Thetai satisfy certain non-degeneracy conditions, we give a (dk)O(m3)-time algorithm that learns F to small error given random labeled examples drawn uniformly from ± 1k× d. - We prove computational lower bounds showing that in the worst case, exponential dependence on m is unavoidable. We focus on Boolean X to mimic the discrete nature of tokens in large language models, though our techniques naturally extend to standard continuous settings, e.g. Gaussian. Our algorithm, which is centered around using examples to sculpt a convex body containing the unknown parameters, is a significant departure from existing provable algorithms for learning feedforward networks, which predominantly exploit algebraic and rotation invariance properties of the Gaussian distribution. In contrast, our analysis is more flexible as it primarily relies on various upper and lower tail bounds for the input distribution and "slices" thereof.