# How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?

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

## Key facts

- Cited by 35 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 74 references that are in the map.
- First posted on arXiv 23 Feb 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [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)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [Learning Structured Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1608.03665/) (2016)
- [ThiNet: A Filter Level Pruning Method for Deep Neural Network Compression](https://mlsanity.com/paper/1707.06342/) (2017)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [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)
- [Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning](https://mlsanity.com/paper/2012.09816/) (2020)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://mlsanity.com/paper/1905.13210/) (2019)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [WANLI: Worker and AI Collaboration for Natural Language Inference Dataset Creation](https://mlsanity.com/paper/2201.05955/) (2022)
- [Self-Adaptive In-Context Learning: An Information Compression Perspective for In-Context Example Selection and Ordering](https://mlsanity.com/paper/2212.10375/) (2022)
- [An Improved Analysis of Training Over-parameterized Deep Neural Networks](https://mlsanity.com/paper/1906.04688/) (2019)

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

- [From Words to Numbers: Your Large Language Model Is Secretly A Capable Regressor When Given In-Context Examples](https://mlsanity.com/paper/2404.07544/) (2024)
- [When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers](https://mlsanity.com/paper/2504.10957/) (2025)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (2024)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [Evaluating the Prompt Steerability of Large Language Models](https://mlsanity.com/paper/2411.12405/) (2024)
- [Can In-context Learning Really Generalize to Out-of-distribution Tasks?](https://mlsanity.com/paper/2410.09695/) (2024)
- [A Provably Effective Method for Pruning Experts in Fine-tuned Sparse Mixture-of-Experts](https://mlsanity.com/paper/2405.16646/) (2024)
- [What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encoding](https://mlsanity.com/paper/2406.01977/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)
- [Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers](https://mlsanity.com/paper/2405.15926/) (2024)
- [Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis](https://mlsanity.com/paper/2410.09605/) (2024)
- [Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysis](https://mlsanity.com/paper/2406.17167/) (2024)
- [Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective](https://mlsanity.com/paper/2505.14808/) (2025)
- [Provable In-Context Learning of Nonlinear Regression with Transformers](https://mlsanity.com/paper/2507.20443/) (2025)
- [Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent](https://mlsanity.com/paper/2510.10425/) (2025)
- [Differentially Private Kernel Density Estimation](https://mlsanity.com/paper/2409.01688/) (2024)
- [AERO: Entropy-Guided Framework for Private LLM Inference](https://mlsanity.com/paper/2410.13060/) (2024)
- [Merging Smarter, Generalizing Better: Enhancing Model Merging on OOD Data](https://mlsanity.com/paper/2506.09093/) (2025)
- [Asymptotic Study of In-context Learning with Random Transformers through Equivalent Models](https://mlsanity.com/paper/2509.15152/) (2025)
- [Understanding Generalization in Transformers: Error Bounds and Training Dynamics Under Benign and Harmful Overfitting](https://mlsanity.com/paper/2502.12508/) (2025)
- [How Can Mamba Learn In Context with Outliers and Generalize Provably?](https://mlsanity.com/paper/2510.00399/) (2025)
- [Transformers Provably Learn Directed Acyclic Graphs via Kernel-Guided Mutual Information](https://mlsanity.com/paper/2510.25542/) (2025)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning](https://mlsanity.com/paper/2510.12026/) (2025)
- [How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs](https://mlsanity.com/paper/2510.25753/) (2025)
- [HIFICL: High-Fidelity In-Context Learning for Multimodal Tasks](https://mlsanity.com/paper/2603.12760/) (2026)
- [The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces](https://mlsanity.com/paper/2606.06814/) (2026)
- [In-Context Learning with Hypothesis-Class Guidance](https://mlsanity.com/paper/2502.19787/) (2025)
- [To Theoretically Understand Transformer-Based In-Context Learning for Optimizing CSMA](https://mlsanity.com/paper/2508.09146/) (2025)
- [Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions](https://mlsanity.com/paper/2510.18638/) (2025)

## Abstract (arXiv)

Transformer-based large language models have displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning by simply augmenting the query with some input-output examples from that task. Despite the empirical success, the mechanics of how to train a Transformer to achieve ICL and the corresponding ICL capacity is mostly elusive due to the technical challenges of analyzing the nonconvex training problems resulting from the nonlinear self-attention and nonlinear activation in Transformers. To the best of our knowledge, this paper provides the first theoretical analysis of the training dynamics of Transformers with nonlinear self-attention and nonlinear MLP, together with the ICL generalization capability of the resulting model. Focusing on a group of binary classification tasks, we train Transformers using data from a subset of these tasks and quantify the impact of various factors on the ICL generalization performance on the remaining unseen tasks with and without data distribution shifts. We also analyze how different components in the learned Transformers contribute to the ICL performance. Furthermore, we provide the first theoretical analysis of how model pruning affects ICL performance and prove that proper magnitude-based pruning can have a minimal impact on ICL while reducing inference costs. These theoretical findings are justified through numerical experiments.