Research map / Architectures and efficiency / Attention and transformers
In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
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.
Empirically and theoretically dissects how randomly initialized one-layer multi-head softmax attention learns in-context linear regression, showing it converges to a near-debiased gradient descent predictor with sign-matching, homogeneous KQ scaling and zero-sum OV patterns, extended to non-isotropic and multi-task settings.
Problem
Understanding how multi-head softmax attention is trained to perform in-context linear regression, including the emergent KQ/OV weight structure, its learned predictor, and its statistical efficiency relative to single-head, linear attention, and Bayes-optimal estimators.
Approach
Shows that gradient training from random initialization produces diagonal KQ and last-entry OV patterns with sign-matched positive/negative heads and zero-sum OV; reduces the model to an approximate loss in (omega, mu) and proves the learned multi-head predictor approximates one-step debiased GD, nearly Bayes-optimal up to a constant factor.
Key results
- Two-head and larger multi-head softmax attention models outperform single-head attention and converge to nearly identical predictors that track vanilla and debiased GD errors on synthetic ICL linear regression.
- Theorem: the learned multi-head predictor approximates one-step debiased GD, and the ratio of its risk to the Bayes risk is bounded by 1 + sigma^{-2}{(1+xi sigma^2)(1+sigma^2+xi^{-1})}^{-1} in the proportional limit.
- Multi-head models generalize to sequence lengths longer than training (L=40), while linear attention requires explicit length normalization.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a second attention head | Single-head softmax attention · In-context linear regression | prediction error | Helped |
| Increase the number of heads from one to two or more | Softmax attention · In-context linear regression | prediction error | Helped |
| Increase the number of heads beyond two | Multi-head softmax attention · In-context linear regression | prediction error | No clear effect |
| Replace softmax attention with linear attention (ablation) | Transformer for in-context linear regression · Longer test sequences | length generalization | Hurt |
| Increase the test sequence length | Softmax attention · In-context linear regression | prediction error | Helped |
| Use multi-head softmax attention instead of the vanilla gradient descent predictor | In-context linear regression · In-context linear regression | prediction error | No clear effect |
| Use multi-head attention instead of the Bayes estimator | In-context linear regression · In-context linear regression | prediction error | Hurt |
| Increase the number of heads from one to four in the two-task setting | Multi-task in-context linear regression · Two-task linear regression (N=2) | task-specific predictor capability | Helped |
| Use three heads rather than assigning two dedicated heads to each task | Multi-task in-context linear regression · Two-task linear regression (N=2) | debiased gradient descent approximation | No clear effect |
Surprising findings
- Multi-head softmax attention (H>=2) learns a parametric GD-like predictor, whereas a single head learns a nonparametric kernel regressor.
- With H>=3 heads, dummy heads may appear with near-zero KQ and OV, depending on random initialization.
Limitations
- Analysis relies on a one-layer simplified reparameterization (diagonal KQ, last-entry OV) and an approximate loss valid under scaling conditions.
- Optimality is with respect to the approximate loss, not the exact ICL loss.
- Multi-task superposition regime (2<H<2N) is only studied in a toy case with H=3, N=2.
Experimental setup
Datasets: synthetic linear regression data generated on the fly
Benchmarks: synthetic in-context linear regression (ICL loss vs. sequence length)
Metrics: ICL prediction loss (MSE)
Baselines: single-head softmax attention linear attention vanilla GD predictor debiased GD predictor optimally tuned ridge (Bayes) estimator
Builds on: linear transformers implement preconditioned gradient descent softmax attention as kernel regressor in-context linear regression KQ and OV circuits
Closest papers by meaning
- Training Dynamics of In-Context Learning in Linear Attention (ICML 2025 · similarity 0.93)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.92)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.91)
- How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR 2024 · similarity 0.90)
- Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond (NeurIPS 2024 · similarity 0.89)
- Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency (NeurIPS 2025 · 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.