mlsanityData to 7 Oct 2026

Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape

arXiv 2402.01258 ↗First posted 2 Feb 2024stat.ML cs.LGcited by 43

Not on the map. Why?

Abstract (arXiv)

Large language models based on the Transformer architecture have demonstrated impressive capabilities to learn in context. However, existing theoretical studies on how this phenomenon arises are limited to the dynamics of a single layer of attention trained on linear regression tasks. In this paper, we study the optimization of a Transformer consisting of a fully connected layer followed by a linear attention layer. The MLP acts as a common nonlinear representation or feature map, greatly enhancing the power of in-context learning. We prove in the mean-field and two-timescale limit that the infinite-dimensional loss landscape for the distribution of parameters, while highly nonconvex, becomes quite benign. We also analyze the second-order stability of mean-field dynamics and show that Wasserstein gradient flow almost always avoids saddle points. Furthermore, we establish novel methods for obtaining concrete improvement rates both away from and near critical points. This represents the first saddle point analysis of mean-field dynamics in general and the techniques are of independent interest.

Key facts. Cited by 43 papers in the map (data through 7 Oct 2026). Builds on 42 of its 67 references that are in the map. First posted on arXiv 2 Feb 2024.

Builds on

42 of 67 references.

References · 42

  1. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  2. 2022Transformers learn in-context by gradient descentcited by 714
  3. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  4. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  5. 2023Gated Linear Attention Transformers with Hardware-Efficient Trainingcited by 490
  6. 2018A Mean Field View of the Landscape of Two-Layers Neural Networkscited by 463
  7. 2018On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transportcited by 422
  8. 2015An Introduction to Matrix Concentration Inequalitiescited by 391
  9. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  10. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  11. 2017How to Escape Saddle Points Efficientlycited by 284
  12. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
30 more
  1. 2015Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decompositioncited by 241
  2. 2019Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernelcited by 211
  3. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  4. 2018Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernelcited by 131
  5. 2017No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysiscited by 127
  6. 2016Matrix Completion has No Spurious Local Minimumcited by 124
  7. 2023What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalizationcited by 112
  8. 2023In-Context Convergence of Transformerscited by 112
  9. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 97
  10. 2020Towards a Mathematical Understanding of Neural Network-Based Machine Learning: what we know and what we don'tcited by 94
  11. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 93
  12. 2023Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretrainingcited by 82
  13. 2018A Priori Estimates of the Population Risk for Two-layer Neural Networkscited by 76
  14. 2017Gradient Descent Can Take Exponential Time to Escape Saddle Pointscited by 74
  15. 2023How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representationscited by 72
  16. 2020Representation formulas and pointwise properties for Barron functionscited by 62
  17. 2023Learning time-scales in two-layers neural networkscited by 57
  18. 2022Convex Analysis of the Mean Field Langevin Dynamicscited by 53
  19. 2023On Learning Gaussian Multi-index Models with Gradient Flowcited by 53
  20. 2016Risk Bounds for High-dimensional Ridge Function Combinations Including Neural Networkscited by 26
  21. 2022How Does Adaptive Optimization Impact Local Neural Network Geometry?cited by 24
  22. 2022Uniform-in-time propagation of chaos for mean field Langevin dynamicscited by 23
  23. 2023Understanding Emergent In-Context Learning from a Kernel Regression Perspectivecited by 23
  24. 2019Global convergence of neuron birth-death dynamicscited by 22
  25. 2020Complexity Measures for Neural Networks with General Activation Functions Using Path-based Normscited by 19
  26. 2023Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reductioncited by 19
  27. 2023Leveraging the two timescale regime to demonstrate convergence of neural networkscited by 18
  28. 2017First-order Methods Almost Always Avoid Saddle Pointscited by 14
  29. 2023Symmetric Mean-field Langevin Dynamics for Distributional Minimax Problemscited by 11
  30. 2022Sequential Attention for Feature Selectioncited by 9

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 15; 2025: 18; 2026: 10.

43 papers.

Sort:

By area:

  1. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  2. 2024In-Context Learning with Representations: Contextual Generalization of Trained Transformerscited by 36
  3. 2024Transformers are Minimax Optimal Nonparametric In-Context Learnerscited by 36
  4. 2025Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thoughtcited by 32
  5. 2024Pretrained transformer efficiently learns low-dimensional target functions in-contextcited by 29
  6. 2024From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiencycited by 28
  7. 2024Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannotcited by 21
  8. 2024On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixturescited by 15
  9. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  10. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  11. 2025Beyond Induction Heads: In-Context Meta Learning Induces Multi-Phase Circuit Emergencecited by 10
  12. 2024Towards Understanding the Universality of Transformers for Next-Token Predictioncited by 9
31 more
  1. 2025Federated In-Context Learning: Iterative Refinement for Improved Answer Qualitycited by 7
  2. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  3. 2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5
  4. 2024Provably Neural Active Learning Succeeds via Prioritizing Perplexing Samplescited by 5
  5. 2024In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learningcited by 5
  6. 2024Global Convergence in Training Large-Scale Transformerscited by 5
  7. 2024Inclusive KL Gradient Flows: Otto-Wasserstein, Fisher-Rao-Gaussian, and Local-Estimator Dynamicscited by 4
  8. 2025ICL CIPHERS: Quantifying "Learning" in In-Context Learning via Substitution Cipherscited by 4
  9. 2025Minimalist Softmax Attention Provably Learns Constrained Boolean Functionscited by 4
  10. 2025Transformers Meet In-Context Learning: A Universal Approximation Theorycited by 4
  11. 2025Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Trainingcited by 4
  12. 2025Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformerscited by 3
  13. 2025Asymptotic Study of In-context Learning with Random Transformers through Equivalent Modelscited by 3
  14. 2026How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretabilitycited by 3
  15. 2025How Can Mamba Learn In Context with Outliers and Generalize Provably?cited by 2
  16. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  17. 2024Probability Distribution Learning and Its Application in Deep Learningcited by 1
  18. 2025In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separationcited by 1
  19. 2025Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Pointscited by 1
  20. 2025Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learningcited by 1
  21. 2025How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPscited by 1
  22. 2026Bayesian Optimality of In-Context Learning with Selective State Spacescited by 1
  23. 2026Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Modelscited by 1
  24. 2026On the global convergence of gradient flow for wide shallow models beyond homogeneous nonlinearitiescited by 1
  25. 2025In-Context Learning with Hypothesis-Class Guidancecited by 0
  26. 2025A Framework for Quantifying How Pre-Training and Context Benefit In-Context Learningcited by 0
  27. 2026Linearized Attention Cannot Enter the Kernel Regime at Any Practical Widthcited by 0
  28. 2026Reconciling In-Context and In-Weight Learning via Dual Representation Space Encodingcited by 0
  29. 2026Demonstrations, CoT, and Prompting: A Theoretical Analysis of ICLcited by 0
  30. 2026Handover of In-Context Learning State Across Session Boundariescited by 0
  31. 2026Relative Generalization Invariance of LLM Pretrainingcited by 0