# Training Infinitely Deep and Wide Transformers

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 48 references that are in the map.
- First posted on arXiv 17 May 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [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)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [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)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (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)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [On the linearity of large non-linear models: when and why the tangent kernel is constant](https://mlsanity.com/paper/2010.01092/) (2020)
- [A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth](https://mlsanity.com/paper/2003.05508/) (2020)
- [Transformers are Universal In-context Learners](https://mlsanity.com/paper/2408.01367/) (2024)
- [Convergence of gradient descent for deep neural networks](https://mlsanity.com/paper/2203.16462/) (2022)
- [Scaling ResNets in the Large-depth Regime](https://mlsanity.com/paper/2206.06929/) (2022)
- [On the Convergence of Encoder-only Shallow Transformers](https://mlsanity.com/paper/2311.01575/) (2023)
- [On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field regime](https://mlsanity.com/paper/2110.02926/) (2021)
- [Unraveling the Gradient Descent Dynamics of Transformers](https://mlsanity.com/paper/2411.07538/) (2024)
- [Understanding the training of infinitely deep and wide ResNets with Conditional Optimal Transport](https://mlsanity.com/paper/2403.12887/) (2024)
- [Control to flocking of the kinetic Cucker-Smale model](https://mlsanity.com/paper/1411.4687/) (2014)
- [On the Convergence of Gradient Descent on Learning Transformers with Residual Connections](https://mlsanity.com/paper/2506.05249/) (2025)
- [Global Convergence in Training Large-Scale Transformers](https://mlsanity.com/paper/2410.23610/) (2024)
- [Local Conditions for Global Convergence of Gradient Flows and Proximal Point Sequences in Metric Spaces](https://mlsanity.com/paper/2304.05239/) (2023)

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

- [Convergence Guarantees of Gradient Descent for Neural Networks via Generalized Lipschitz Smoothness](https://mlsanity.com/paper/2608.11479/) (2026)

## Abstract (arXiv)

Transformers have become the dominant architecture in modern machine learning, yet the theoretical understanding of their training dynamics remains limited. This paper develops a rigorous mathematical framework for analyzing gradient-based training of transformers in the mean-field regime, where both the depth (number of layers) and width (number of attention heads) tend to infinity. While ResNet training can be understood as controlling a neural ODE, transformer training corresponds to controlling a neural PDE, due to the coupling of multiple token distributions through the attention mechanism. Our mean-field model features two types of measure representations: token distributions evolving through layers and attention parameters at each layer. We establish well-posedness of the forward pass through infinitely deep transformers, characterizing token evolution via flow maps that satisfy ODEs in function spaces. Using adjoint sensitivity analysis, we derive an explicit formula for the conditional Wasserstein gradient of the training risk, involving adjoint variables governed by backward ODEs. We prove the existence and uniqueness of gradient flow curves in the conditional Wasserstein metric space, establishing a rigorous foundation for gradient-based transformer training. A key technical contribution is providing necessary and sufficient conditions for injectivity of the Neural Tangent Kernel (NTK) for attention mechanisms: we show that NTK injectivity is equivalent to linear independence of log-sum-exp functions modulo affine functions, a condition satisfied by diverse token distributions, including discrete distributions, uniform distributions, and Gaussian mixtures. Under this NTK injectivity assumption, we prove that gradient flow converges to global minima when the initial loss is sufficiently small, eliminating spurious local minima from the optimization landscape.