mlsanityData to 7 Oct 2026

Training Dynamics of In-Context Learning in Linear Attention

arXiv 2501.16265 ↗First posted 27 Jan 2025cs.LGcited by 44

Not on the map. Why?

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.

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

Builds on

81 of 98 references.

References · 81

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2022In-context Learning and Induction Headscited by 930
  4. 2023Progress measures for grokking via mechanistic interpretabilitycited by 822
  5. 2013Exact solutions to the nonlinear dynamics of learning in deep linear neural networkscited by 748
  6. 2022Transformers learn in-context by gradient descentcited by 714
  7. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  8. 2018On Lazy Training in Differentiable Programmingcited by 671
  9. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  10. 2021Linear Transformers Are Secretly Fast Weight Programmerscited by 471
  11. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  12. 2022Data Distributional Properties Drive Emergent In-Context Learning in Transformerscited by 303
38 more
  1. 2020Kernel and Rich Regimes in Overparametrized Modelscited by 302
  2. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  3. 2018On the Optimization of Deep Networks: Implicit Acceleration by Overparameterizationcited by 280
  4. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
  5. 2018Gradient descent aligns the layers of deep linear networkscited by 182
  6. 2018A mathematical theory of semantic development in deep neural networkscited by 182
  7. 2018A Convergence Analysis of Gradient Descent for Deep Linear Neural Networkscited by 180
  8. 2018Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balancedcited by 172
  9. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  10. 2023Sudden Drops in the Loss: Syntax Acquisition, Phase Transitions, and Simplicity Bias in MLMscited by 121
  11. 2023The mechanistic basis of data dependence and abrupt learning in an in-context classification taskcited by 115
  12. 2024The Evolution of Statistical Induction Heads: In-Context Learning Markov Chainscited by 115
  13. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 114
  14. 2023In-Context Convergence of Transformerscited by 112
  15. 2021Neural Networks as Kernel Learners: The Silent Alignment Effectcited by 106
  16. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  17. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 97
  18. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 93
  19. 2022Vision Transformers provably learn spatial structurecited by 92
  20. 2024What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formationcited by 86
  21. 2022Training Trajectories of Language Models Across Scalescited by 80
  22. 2023The Transient Nature of Emergent In-Context Learning in Transformerscited by 72
  23. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 71
  24. 2023Transformers as Support Vector Machinescited by 67
  25. 2023Max-Margin Token Selection in Attention Mechanismcited by 57
  26. 2024Asymptotic theory of in-context learning by linear attentioncited by 54
  27. 2023Transformers learn through gradual rank increasecited by 52
  28. 2024Competition Dynamics Shape Algorithmic Phases of In-Context Learningcited by 52
  29. 2024LoRA Training in the NTK Regime has No Spurious Local Minimacited by 48
  30. 2023Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regressioncited by 47
  31. 2024Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformerscited by 46
  32. 2024Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscapecited by 43
  33. 2024Dynamic metastability in the self-attention modelcited by 42
  34. 2024Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic taskscited by 40
  35. 2024Infinite Limits of Multi-head Transformer Dynamicscited by 38
  36. 2024From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformerscited by 34
  37. 2024The broader spectrum of in-context learningcited by 33
  38. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31

50 most cited of 81 All 81 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 22; 2026: 22.

44 papers.

Sort:

By area:

  1. 2025The emergence of sparse attention: impact of data distribution and benefits of repetitioncited by 25
  2. 2026There Will Be a Scientific Theory of Deep Learningcited by 21
  3. 2025Strategy Coopetition Explains the Emergence and Transience of In-Context Learningcited by 17
  4. 2025Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network Architecturescited by 14
  5. 2025Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networkscited by 11
  6. 2026Understanding LoRA as Knowledge Memory: An Empirical Analysiscited by 11
  7. 2025Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformerscited by 10
  8. 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamicscited by 9
  9. 2026Sequential Group Composition: A Window into the Mechanics of Deep Learningcited by 7
  10. 2026To Use or not to Use Muon: How Simplicity Bias in Optimizers Matterscited by 7
  11. 2025Statistical Advantage of Softmax Attention: Insights from Single-Location Regressioncited by 6
  12. 2025Weight-Space Linear Recurrent Neural Networkscited by 5
32 more
  1. 2025On the Convergence of Gradient Descent on Learning Transformers with Residual Connectionscited by 5
  2. 2025What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformerscited by 5
  3. 2025Single-Head Attention in High Dimensions: A Theory of Generalization, Weights Spectra, and Scaling Lawscited by 5
  4. 2025In-Context Occam's Razor: How Transformers Prefer Simpler Hypotheses on the Flycited by 4
  5. 2025Exact Learning Dynamics of In-Context Learning in Linear Transformers and Its Application to Non-Linear Transformerscited by 3
  6. 2025On the Emergence of Induction Heads for In-Context Learningcited by 3
  7. 2025Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspectivecited by 3
  8. 2025Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descentcited by 2
  9. 2026Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridgecited by 2
  10. 2026Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learningcited by 2
  11. 2026Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledgecited by 2
  12. 2026Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizercited by 2
  13. 2025Multi-Layer Attention is the Amplifier of Demonstration Effectivenesscited by 1
  14. 2025Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regressioncited by 1
  15. 2025Learning Linear Regression with Low-Rank Tasks in-Contextcited by 1
  16. 2025Understanding the Staged Dynamics of Transformers in Learning Latent Structurecited by 1
  17. 2026Incremental Learning of Sparse Attention Patterns in Transformerscited by 1
  18. 2026Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Modelscited by 1
  19. 2026Specialization of softmax attention heads: insights from the high-dimensional single-location modelcited by 1
  20. 2026Conservation Laws for Modern Neural Architecturescited by 1
  21. 2025Decomposing Prediction Mechanisms for In-Context Recallcited by 0
  22. 2025What Neuroscience Can Teach AI About Learning in Continuously Changing Environmentscited by 0
  23. 2026Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMscited by 0
  24. 2026Learning to Adapt: In-Context Learning Beyond Stationaritycited by 0
  25. 2026A Theory of Saddle Escape in Deep Nonlinear Networkscited by 0
  26. 2026Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regressioncited by 0
  27. 2026The Stability of Singular Distribution: A Spectral Perspective on the Two-Phase Dynamics of Language Model Pre-trainingcited by 0
  28. 2026An Asymptotic Theory of Chain-of-Thought in In-Context Learningcited by 0
  29. 2026The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspacescited by 0
  30. 2026When Does Depth Matter For In-Context Learning? Adaptive Inference in Deep Transformerscited by 0
  31. 2026Sequential Correlations Change In-Context Learning: Effective Context Length and Architectural Mismatchcited by 0
  32. 2026Syntax vs. Semantics: How Transformers Learn Deep Dependenciescited by 0