# Training Dynamics of In-Context Learning in Linear Attention

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

## Key facts

- Cited by 47 papers in the map (data through 14 Oct 2026).
- Builds on 81 of its 98 references that are in the map.
- First posted on arXiv 27 Jan 2025.

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

- [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)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [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)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization](https://mlsanity.com/paper/1802.06509/) (2018)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [A mathematical theory of semantic development in deep neural networks](https://mlsanity.com/paper/1810.10531/) (2018)
- [A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks](https://mlsanity.com/paper/1810.02281/) (2018)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [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)
- [Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMs](https://mlsanity.com/paper/2309.07311/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [The mechanistic basis of data dependence and abrupt learning in an in-context classification task](https://mlsanity.com/paper/2312.03002/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Neural Networks as Kernel Learners: The Silent Alignment Effect](https://mlsanity.com/paper/2111.00034/) (2021)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (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)

## Built on by (most cited first, 30 of 47)

- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [The emergence of sparse attention: impact of data distribution and benefits of repetition](https://mlsanity.com/paper/2505.17863/) (2025)
- [Strategy Coopetition Explains the Emergence and Transience of In-Context Learning](https://mlsanity.com/paper/2503.05631/) (2025)
- [Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network Architectures](https://mlsanity.com/paper/2512.20607/) (2025)
- [Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks](https://mlsanity.com/paper/2506.06489/) (2025)
- [Understanding LoRA as Knowledge Memory: An Empirical Analysis](https://mlsanity.com/paper/2603.01097/) (2026)
- [Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers](https://mlsanity.com/paper/2506.10887/) (2025)
- [From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics](https://mlsanity.com/paper/2510.06954/) (2025)
- [Sequential Group Composition: A Window into the Mechanics of Deep Learning](https://mlsanity.com/paper/2602.03655/) (2026)
- [To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters](https://mlsanity.com/paper/2603.00742/) (2026)
- [Statistical Advantage of Softmax Attention: Insights from Single-Location Regression](https://mlsanity.com/paper/2509.21936/) (2025)
- [Single-Head Attention in High Dimensions: A Theory of Generalization, Weights Spectra, and Scaling Laws](https://mlsanity.com/paper/2509.24914/) (2025)
- [What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers](https://mlsanity.com/paper/2506.13688/) (2025)
- [On the Convergence of Gradient Descent on Learning Transformers with Residual Connections](https://mlsanity.com/paper/2506.05249/) (2025)
- [Weight-Space Linear Recurrent Neural Networks](https://mlsanity.com/paper/2506.01153/) (2025)
- [In-Context Occam's Razor: How Transformers Prefer Simpler Hypotheses on the Fly](https://mlsanity.com/paper/2506.19351/) (2025)
- [On the Emergence of Induction Heads for In-Context Learning](https://mlsanity.com/paper/2511.01033/) (2025)
- [Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent](https://mlsanity.com/paper/2510.10425/) (2025)
- [Exact Learning Dynamics of In-Context Learning in Linear Transformers and Its Application to Non-Linear Transformers](https://mlsanity.com/paper/2504.12916/) (2025)
- [Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspective](https://mlsanity.com/paper/2512.11784/) (2025)
- [Specialization of softmax attention heads: insights from the high-dimensional single-location model](https://mlsanity.com/paper/2603.03993/) (2026)
- [Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledge](https://mlsanity.com/paper/2603.20969/) (2026)
- [Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer](https://mlsanity.com/paper/2605.05176/) (2026)
- [Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression](https://mlsanity.com/paper/2509.23779/) (2025)
- [Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge](https://mlsanity.com/paper/2602.02470/) (2026)
- [Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learning](https://mlsanity.com/paper/2602.04872/) (2026)
- [Incremental Learning of Sparse Attention Patterns in Transformers](https://mlsanity.com/paper/2602.19143/) (2026)
- [Learning to Adapt: In-Context Learning Beyond Stationarity](https://mlsanity.com/paper/2604.10946/) (2026)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [Learning Linear Regression with Low-Rank Tasks in-Context](https://mlsanity.com/paper/2510.04548/) (2025)

## Abstract (arXiv)

While attention-based models have demonstrated the remarkable ability of in-context learning (ICL), the theoretical understanding of how these models acquired this ability through gradient descent training is still preliminary. Towards answering this question, we study the gradient descent dynamics of multi-head linear self-attention trained for in-context linear regression. We examine two parametrizations of linear self-attention: one with the key and query weights merged as a single matrix (common in theoretical studies), and one with separate key and query matrices (closer to practical settings). For the merged parametrization, we show that the training dynamics has two fixed points and the loss trajectory exhibits a single, abrupt drop. We derive an analytical time-course solution for a certain class of datasets and initialization. For the separate parametrization, we show that the training dynamics has exponentially many fixed points and the loss exhibits saddle-to-saddle dynamics, which we reduce to scalar ordinary differential equations. During training, the model implements principal component regression in context with the number of principal components increasing over training time. Overall, we provide a theoretical description of how ICL abilities evolve during gradient descent training of linear attention, revealing abrupt acquisition or progressive improvements depending on how the key and query are parametrized.