Research map / Architectures and efficiency / Attention and transformers
How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies how a trained multi-head transformer performs in-context sparse linear regression, finds multi-head use in the first layer and a single dominant head in later layers, and proposes and theoretically analyzes a preprocess-then-optimize mechanism.
Problem
Trained transformers' in-context learning mechanisms, especially the role of multiple heads across layers, are poorly understood; prior theory mostly relies on hand-designed constructions for linear regression.
Approach
Empirically, the first attention layer uses all heads significantly while later layers rely on one head; this motivates a preprocess-then-optimize hypothesis where layer one reweights features by estimated feature-label correlations and later layers run gradient descent on the preprocessed data, which provably achieves lower excess risk than GD or ridge on raw features.
Key results
- In trained models, the first layer's head importance is spread evenly across heads, while later layers have one dominant head (e.g., pruned models with all first-layer heads match the original).
- Theory shows the preprocess-then-optimize algorithm achieves Õ(s/n) excess risk under d≲n, versus Ω(d/n) for OLS, ridge, and GD.
- Transformers with more heads (h=4,8) yield lower excess risk than h=1 for few-shot in-context examples, and outperform OLS and ridge.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the number of attention heads (from h=1 to h=4 or h=8) | Transformer · Sparse linear regression | excess risk | Helped |
| Mask individual heads in the first layer (ablation) | Trained multi-head transformer · Sparse linear regression | relative excess-risk increase after masking | Mixed |
| Mask individual heads in subsequent layers (ablation) | Trained multi-head transformer · Sparse linear regression | relative excess-risk increase after masking | Mixed |
| Keep all first-layer heads and only the most important head in subsequent layers (ablation) | Trained multi-head transformer · Sparse linear regression | excess risk | No clear effect |
| Use multiple heads (h=4,8) rather than a single head to produce first-layer representations for gradient-descent probing | Transformer · Sparse linear regression | excess risk after gradient descent on probed representations | Helped |
| Preprocess the data before gradient descent | Gradient descent on sparse linear regression examples · Sparse linear regression | excess risk | Helped |
| Use multi-head attention in the first layer and retain one head in subsequent layers, rather than single-head attention throughout | Transformer · Sparse linear regression | excess risk | Helped |
Surprising findings
- Only the first layer needs multiple heads; later layers effectively use a single head, with pruned multi-head-first-layer models matching the original performance.
Limitations
- Only sparse linear regression with synthetic Gaussian data is studied.
- Analysis uses linear attention-only transformers, neglecting MLPs and layernorm.
- Training dynamics of multi-layer transformers are not analyzed.
Experimental setup
Datasets: synthetic sparse linear regression data (d=16, s=4, Gaussian features)
Benchmarks: synthetic sparse linear regression
Metrics: excess risk head importance weight probing excess risk
Baselines: ordinary least squares ridge regression lasso gradient descent on raw inputs
Builds on: gradient descent emulation by transformers linear attention transformers ridge regression Lasso
Closest papers by meaning
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.92)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.91)
- How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR 2024 · similarity 0.90)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.90)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.90)
- One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.