# Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality

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

## Key facts

- Cited by 74 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 29 Feb 2024.

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

- [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)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (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)
- [How Well Can Transformers Emulate In-context Newton's Method?](https://mlsanity.com/paper/2403.03183/) (2024)
- [The emergence of sparse attention: impact of data distribution and benefits of repetition](https://mlsanity.com/paper/2505.17863/) (2025)
- [Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot](https://mlsanity.com/paper/2406.06893/) (2024)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [One-Layer Transformer Provably Learns One-Nearest Neighbor In Context](https://mlsanity.com/paper/2411.10830/) (2024)
- [Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time](https://mlsanity.com/paper/2510.01098/) (2025)
- [Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond](https://mlsanity.com/paper/2405.03251/) (2024)
- [Understanding the Language Model to Solve the Symbolic Multi-Step Reasoning Problem from the Perspective of Buffer Mechanism](https://mlsanity.com/paper/2405.15302/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)
- [Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations](https://mlsanity.com/paper/2406.11828/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [A Theoretical Understanding of Chain-of-Thought: Coherent Reasoning and Error-Aware Demonstration](https://mlsanity.com/paper/2410.16540/) (2024)
- [Adversarial Training Can Provably Improve Robustness: Theoretical Analysis of Feature Learning Process Under Structured Data](https://mlsanity.com/paper/2410.08503/) (2024)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics](https://mlsanity.com/paper/2510.06954/) (2025)
- [Amuro and Char: Analyzing the Relationship between Pre-Training and Fine-Tuning of Large Language Models](https://mlsanity.com/paper/2408.06663/) (2024)
- [Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis](https://mlsanity.com/paper/2410.09605/) (2024)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [How Transformers Get Rich: Approximation and Dynamics Analysis](https://mlsanity.com/paper/2410.11474/) (2024)
- [Federated In-Context Learning: Iterative Refinement for Improved Answer Quality](https://mlsanity.com/paper/2506.07440/) (2025)

## Abstract (arXiv)

We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We establish the global convergence of gradient flow under suitable choices of initialization. In addition, we prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics, where each attention head focuses on solving a single task of the multi-task model. Specifically, we prove that the gradient flow dynamics can be split into three phases -- a warm-up phase where the loss decreases rather slowly and the attention heads gradually build up their inclination towards individual tasks, an emergence phase where each head selects a single task and the loss rapidly decreases, and a convergence phase where the attention parameters converge to a limit. Furthermore, we prove the optimality of gradient flow in the sense that the limiting model learned by gradient flow is on par with the best possible multi-head softmax attention model up to a constant factor. Our analysis also delineates a strict separation in terms of the prediction accuracy of ICL between single-head and multi-head attention models. The key technique for our convergence analysis is to map the gradient flow dynamics in the parameter space to a set of ordinary differential equations in the spectral domain, where the relative magnitudes of the semi-singular values of the attention weights determines task allocation. To our best knowledge, our work provides the first convergence result for the multi-head softmax attention model.