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

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

## Key facts

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

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

- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [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)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [An Introduction to Matrix Concentration Inequalities](https://mlsanity.com/paper/1501.01571/) (2015)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decomposition](https://mlsanity.com/paper/1503.02101/) (2015)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [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)
- [Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel](https://mlsanity.com/paper/1810.05369/) (2018)
- [No Spurious Local Minima in Nonconvex Low Rank Problems: A Unified Geometric Analysis](https://mlsanity.com/paper/1704.00708/) (2017)
- [Matrix Completion has No Spurious Local Minimum](https://mlsanity.com/paper/1605.07272/) (2016)
- [What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization](https://mlsanity.com/paper/2305.19420/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (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)
- [Towards a Mathematical Understanding of Neural Network-Based Machine Learning: what we know and what we don't](https://mlsanity.com/paper/2009.10713/) (2020)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [A Priori Estimates of the Population Risk for Two-layer Neural Networks](https://mlsanity.com/paper/1810.06397/) (2018)
- [How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations](https://mlsanity.com/paper/2310.10616/) (2023)
- [Gradient Descent Can Take Exponential Time to Escape Saddle Points](https://mlsanity.com/paper/1705.10412/) (2017)
- [Representation formulas and pointwise properties for Barron functions](https://mlsanity.com/paper/2006.05982/) (2020)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)

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

- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [In-Context Learning with Representations: Contextual Generalization of Trained Transformers](https://mlsanity.com/paper/2408.10147/) (2024)
- [Transformers are Minimax Optimal Nonparametric In-Context Learners](https://mlsanity.com/paper/2408.12186/) (2024)
- [Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought](https://mlsanity.com/paper/2502.21212/) (2025)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency](https://mlsanity.com/paper/2410.05459/) (2024)
- [Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot](https://mlsanity.com/paper/2406.06893/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Beyond Induction Heads: In-Context Meta Learning Induces Multi-Phase Circuit Emergence](https://mlsanity.com/paper/2505.16694/) (2025)
- [Towards Understanding the Universality of Transformers for Next-Token Prediction](https://mlsanity.com/paper/2410.03011/) (2024)
- [Federated In-Context Learning: Iterative Refinement for Improved Answer Quality](https://mlsanity.com/paper/2506.07440/) (2025)
- [In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning](https://mlsanity.com/paper/2409.12293/) (2024)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [Transformers Meet In-Context Learning: A Universal Approximation Theory](https://mlsanity.com/paper/2506.05200/) (2025)
- [On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability](https://mlsanity.com/paper/2405.16845/) (2024)
- [Provably Neural Active Learning Succeeds via Prioritizing Perplexing Samples](https://mlsanity.com/paper/2406.03944/) (2024)
- [Global Convergence in Training Large-Scale Transformers](https://mlsanity.com/paper/2410.23610/) (2024)
- [Inclusive KL Gradient Flows: Otto-Wasserstein, Fisher-Rao-Gaussian, and Local-Estimator Dynamics](https://mlsanity.com/paper/2411.00214/) (2024)
- [In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separation](https://mlsanity.com/paper/2502.08136/) (2025)
- [ICL CIPHERS: Quantifying "Learning" in In-Context Learning via Substitution Ciphers](https://mlsanity.com/paper/2504.19395/) (2025)
- [Minimalist Softmax Attention Provably Learns Constrained Boolean Functions](https://mlsanity.com/paper/2505.19531/) (2025)
- [Provable Benefit of Curriculum in Transformer Tree-Reasoning Post-Training](https://mlsanity.com/paper/2511.07372/) (2025)
- [Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers](https://mlsanity.com/paper/2502.20681/) (2025)
- [Asymptotic Study of In-context Learning with Random Transformers through Equivalent Models](https://mlsanity.com/paper/2509.15152/) (2025)
- [How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability](https://mlsanity.com/paper/2601.19208/) (2026)
- [Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points](https://mlsanity.com/paper/2509.16974/) (2025)
- [How Can Mamba Learn In Context with Outliers and Generalize Provably?](https://mlsanity.com/paper/2510.00399/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)

## 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.