mlsanityData to 7 Oct 2026

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

arXiv 2401.17426 ↗First posted 30 Jan 2024cs.LG cs.AI stat.MLcited by 17

Not on the map. Why?

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.

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

Builds on

38 of 41 references.

References · 38

  1. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  2. 2022Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?cited by 1,413
  3. 2021What Makes Good In-Context Examples for GPT-3?cited by 1,184
  4. 2021Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivitycited by 1,122
  5. 2022Transformers learn in-context by gradient descentcited by 714
  6. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  7. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  8. 2022Language Models are Multilingual Chain-of-Thought Reasonerscited by 518
  9. 2021Linear Transformers Are Secretly Fast Weight Programmerscited by 471
  10. 2021MetaICL: Learning to Learn In Contextcited by 450
  11. 2022Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thoughtcited by 334
  12. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
26 more
  1. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  2. 2019Theoretical Limitations of Self-Attention in Neural Sequence Modelscited by 286
  3. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
  4. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 220
  5. 2022Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizerscited by 204
  6. 2022Active Example Selection for In-Context Learningcited by 195
  7. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  8. 2021Noisy Channel Language Model Prompting for Few-Shot Text Classificationcited by 163
  9. 2023Pretraining task diversity and the emergence of non-Bayesian in-context learning for regressioncited by 138
  10. 2023In-Context Convergence of Transformerscited by 112
  11. 2022Overcoming a Theoretical Limitation of Self-Attentioncited by 110
  12. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 97
  13. 2023Uncovering mesa-optimization algorithms in Transformerscited by 94
  14. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 93
  15. 2023In-Context Learning through the Bayesian Prismcited by 78
  16. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  17. 2023Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Contextcited by 63
  18. 2022Improving In-Context Few-Shot Learning via Self-Supervised Trainingcited by 56
  19. 2023Attention Scheme Inspired Softmax Regressioncited by 49
  20. 2023On the Role of Attention in Prompt-tuningcited by 49
  21. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 47
  22. 2023Memorization Capacity of Multi-Head Attention in Transformerscited by 39
  23. 2023On the Optimization and Generalization of Multi-head Attentioncited by 37
  24. 2023Sequence Length Independent Norm-Based Generalization Bounds for Transformerscited by 28
  25. 2023Why Softmax Attention Outperforms Linear Attentioncited by 15
  26. 2020Repulsive Attention: Rethinking Multi-head Attention as Bayesian Inferencecited by 13

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 10; 2025: 5; 2026: 2.

17 papers.

Sort:

By area:

  1. 2024How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?cited by 35
  2. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysiscited by 14
  3. 2024How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regressioncited by 13
  4. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  5. 2024A Theoretical Understanding of Chain-of-Thought: Coherent Reasoning and Error-Aware Demonstrationcited by 13
  6. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  7. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  8. 2024Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Datacited by 5
  9. 2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5
  10. 2024From Unstructured Data to In-Context Learning: Exploring What Tasks Can Be Learned and Whencited by 3
  11. 2025The Initialization Determines Whether In-Context Learning Is Gradient Descentcited by 3
  12. 2025Impact of Positional Encoding: Clean and Adversarial Rademacher Complexity for Transformers under In-Context Regressioncited by 2
5 more
  1. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  2. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalizationcited by 1
  3. 2024Toward generalizable learning of all (linear) first-order methods via memory augmented Transformerscited by 0
  4. 2025TabGRU: An Enhanced Design for Urban Rainfall Intensity Estimation Using Commercial Microwave Linkscited by 0
  5. 2026Taming Visual Neglect: A Variational Information Bottleneck Framework for Adaptive Attention in Multimodal In-Context Learningcited by 0