# How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 52 of its 57 references that are in the map.
- First posted on arXiv 8 Aug 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Pruning Filters for Efficient ConvNets](https://mlsanity.com/paper/1608.08710/) (2016)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (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)
- [Benign Overfitting in Linear Regression](https://mlsanity.com/paper/1906.11300/) (2019)
- [ShortGPT: Layers in Large Language Models are More Redundant Than You Expect](https://mlsanity.com/paper/2403.03853/) (2024)
- [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)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Physics of Language Models: Part 3.1, Knowledge Storage and Extraction](https://mlsanity.com/paper/2309.14316/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [The Unreasonable Ineffectiveness of the Deeper Layers](https://mlsanity.com/paper/2403.17887/) (2024)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Thinking Like Transformers](https://mlsanity.com/paper/2106.06981/) (2021)
- [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)
- [Benign overfitting in ridge regression](https://mlsanity.com/paper/2009.14286/) (2020)
- [Physics of Language Models: Part 3.2, Knowledge Manipulation](https://mlsanity.com/paper/2309.14402/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)

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

- [Towards Understanding Fine-Tuning Mechanisms of LLMs via Circuit Analysis](https://mlsanity.com/paper/2502.11812/) (2025)
- [Short-length Adversarial Training Helps LLMs Defend Long-length Jailbreak Attacks: Theoretical and Empirical Evidence](https://mlsanity.com/paper/2502.04204/) (2025)
- [Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B](https://mlsanity.com/paper/2504.00132/) (2025)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [Transformers versus the EM Algorithm in Multi-class Clustering](https://mlsanity.com/paper/2502.06007/) (2025)
- [On the Robustness of Transformers against Context Hijacking for Linear Classification](https://mlsanity.com/paper/2502.15609/) (2025)
- [Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads](https://mlsanity.com/paper/2510.16807/) (2025)
- [Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression](https://mlsanity.com/paper/2508.07571/) (2025)
- [The Initialization Determines Whether In-Context Learning Is Gradient Descent](https://mlsanity.com/paper/2512.04268/) (2025)
- [Provable and Practical In-Context Policy Optimization for Self-Improvement](https://mlsanity.com/paper/2603.01335/) (2026)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [Toward generalizable learning of all (linear) first-order methods via memory augmented Transformers](https://mlsanity.com/paper/2410.07263/) (2024)
- [First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training](https://mlsanity.com/paper/2510.14614/) (2025)
- [Adaptive Mean Estimation by In-Context Learning: A Gradient-Flow Analysis](https://mlsanity.com/paper/2610.07804/) (2026)

## Abstract (arXiv)

Despite the remarkable success of transformer-based models in various real-world tasks, their underlying mechanisms remain poorly understood. Recent studies have suggested that transformers can implement gradient descent as an in-context learner for linear regression problems and have developed various theoretical analyses accordingly. However, these works mostly focus on the expressive power of transformers by designing specific parameter constructions, lacking a comprehensive understanding of their inherent working mechanisms post-training. In this study, we consider a sparse linear regression problem and investigate how a trained multi-head transformer performs in-context learning. We experimentally discover that the utilization of multi-heads exhibits different patterns across layers: multiple heads are utilized and essential in the first layer, while usually only a single head is sufficient for subsequent layers. We provide a theoretical explanation for this observation: the first layer preprocesses the context data, and the following layers execute simple optimization steps based on the preprocessed context. Moreover, we demonstrate that such a preprocess-then-optimize algorithm can significantly outperform naive gradient descent and ridge regression algorithms. Further experimental results support our explanations. Our findings offer insights into the benefits of multi-head attention and contribute to understanding the more intricate mechanisms hidden within trained transformers.