# Superiority of Multi-Head Attention in In-Context Linear Regression

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 41 references that are in the map.
- First posted on arXiv 30 Jan 2024.

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

- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Language Models are Multilingual Chain-of-Thought Reasoners](https://mlsanity.com/paper/2210.03057/) (2022)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [MetaICL: Learning to Learn In Context](https://mlsanity.com/paper/2110.15943/) (2021)
- [Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought](https://mlsanity.com/paper/2210.01240/) (2022)
- [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)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers](https://mlsanity.com/paper/2212.10559/) (2022)
- [Active Example Selection for In-Context Learning](https://mlsanity.com/paper/2211.04486/) (2022)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Noisy Channel Language Model Prompting for Few-Shot Text Classification](https://mlsanity.com/paper/2108.04106/) (2021)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Overcoming a Theoretical Limitation of Self-Attention](https://mlsanity.com/paper/2202.12172/) (2022)
- [Uncovering mesa-optimization algorithms in Transformers](https://mlsanity.com/paper/2309.05858/) (2023)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Improving In-Context Few-Shot Learning via Self-Supervised Training](https://mlsanity.com/paper/2205.01703/) (2022)

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

- [How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?](https://mlsanity.com/paper/2402.15607/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (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)
- [A Theoretical Understanding of Chain-of-Thought: Coherent Reasoning and Error-Aware Demonstration](https://mlsanity.com/paper/2410.16540/) (2024)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data](https://mlsanity.com/paper/2402.00743/) (2024)
- [On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability](https://mlsanity.com/paper/2405.16845/) (2024)
- [From Unstructured Data to In-Context Learning: Exploring What Tasks Can Be Learned and When](https://mlsanity.com/paper/2406.00131/) (2024)
- [The Initialization Determines Whether In-Context Learning Is Gradient Descent](https://mlsanity.com/paper/2512.04268/) (2025)
- [Impact of Positional Encoding: Clean and Adversarial Rademacher Complexity for Transformers under In-Context Regression](https://mlsanity.com/paper/2512.09275/) (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)
- [Toward generalizable learning of all (linear) first-order methods via memory augmented Transformers](https://mlsanity.com/paper/2410.07263/) (2024)
- [TabGRU: An Enhanced Design for Urban Rainfall Intensity Estimation Using Commercial Microwave Links](https://mlsanity.com/paper/2512.02465/) (2025)
- [Taming Visual Neglect: A Variational Information Bottleneck Framework for Adaptive Attention in Multimodal In-Context Learning](https://mlsanity.com/paper/2608.23570/) (2026)

## Abstract (arXiv)

We present a theoretical analysis of the performance of transformer with softmax attention in in-context learning with linear regression tasks. While the existing literature predominantly focuses on the convergence of transformers with single-/multi-head attention, our research centers on comparing their performance. We conduct an exact theoretical analysis to demonstrate that multi-head attention with a substantial embedding dimension performs better than single-head attention. When the number of in-context examples D increases, the prediction loss using single-/multi-head attention is in O(1/D), and the one for multi-head attention has a smaller multiplicative constant. In addition to the simplest data distribution setting, we consider more scenarios, e.g., noisy labels, local examples, correlated features, and prior knowledge. We observe that, in general, multi-head attention is preferred over single-head attention. Our results verify the effectiveness of the design of multi-head attention in the transformer architecture.