# Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 23 references that are in the map.
- First posted on arXiv 2 Mar 2026.

## Builds on (most cited first, 17 of 17)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [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)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (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)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [Low-rank Solutions of Linear Matrix Equations via Procrustes Flow](https://mlsanity.com/paper/1507.03566/) (2015)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [Formal Algorithms for Transformers](https://mlsanity.com/paper/2207.09238/) (2022)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)

## Built on by (most cited first, 5 of 5)

- [Echo: KV-Cache-Free Associative Recall with Spectral Koopman Operators](https://mlsanity.com/paper/2605.06997/) (2026)
- [Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention](https://mlsanity.com/paper/2605.21292/) (2026)
- [Reachability and asymptotics of Gaussian Transformer dynamics](https://mlsanity.com/paper/2606.07600/) (2026)
- [Stochastic approximation in non-markovian environments revisited](https://mlsanity.com/paper/2603.21091/) (2026)
- [Learning Decision-Stump Thresholds in Context: Dynamics of Softmax Attention](https://mlsanity.com/paper/2610.07074/) (2026)

## Abstract (arXiv)

We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.