Transformer Global Convergence with Mean-Field Theory

1. Introduction: The Optimization Mystery of Large Transformers

Despite the widespread success of Transformer models across various domains—including natural language processing and computer vision—their optimization guarantees in large-scale settings remain poorly understood.¹

The Core Mystery: Why do gradient-based methods consistently succeed in training Transformers, despite the highly non-convex landscape of the training objective?

This phenomenon becomes particularly intriguing as model size increases: training algorithms typically converge globally, even when the loss landscape contains numerous local minima and saddle points. Understanding this theoretical puzzle is essential for:

• Algorithmic improvements: Designing better optimizers and learning rate schedules
• Architecture design: Understanding which components contribute to trainability
• Generalization theory: Connecting optimization dynamics to out-of-sample performance

1.1 Prior Work and Limitations

Previous theoretical analyses of deep network optimization have established global convergence guarantees for:

Model	Key Technique	Limitation
Two-layer NNs	Mean-field analysis	Limited to shallow architectures
Deep ResNets	Neural ODE / Skip connections	Requires homogeneity assumptions
NTK regime	Infinite width scaling	Excludes feature learning

Critical Gap: Existing tools for deep networks (particularly Lu et al., 2020) demand:

• Full homogeneity of the network function
• Global Lipschitz smoothness of gradients

These conditions are not satisfied by Transformer architectures, particularly due to:

• Softmax attention mechanism (not homogeneous)
• Layer normalization
• Complex interactions between attention and FFN blocks

1.2 Our Approach: Mean-Field Theory for Transformers

This paper bridges the gap between Transformer theory and practice by demonstrating global convergence of Transformer training via gradient flow in a large-scale model regime.¹

Key Innovation: Shift optimization analysis from parameter space to distributional dynamics in the Wasserstein metric, enabling:

1. Construction of the mean-field limit for Transformers
2. Rigorous approximation guarantees between discrete and continuous dynamics
3. Proof of global minimum convergence under mild assumptions

---

2. Mean-Field Limit Construction for Transformers

2.1 Transformer Model Architecture

Following common Transformer configurations, each block consists of two distinct layers:¹

Self-Attention Layer (with residual connection):

$${\text{Attn}}_{{\theta }_1,\dots ,{\theta }_M}(Z,\eta )=Z+\frac{\eta }{M}\sum _{j=1}^{M}f(Z,{\theta }^j)$$

Feed-Forward Layer (with residual connection):

$${\text{MLP}}_{w_1,\dots ,w_M}(Z,\eta )=Z+\frac{\eta }{M}\sum _{j=1}^{M}h(Z,w^j)$$

where:

• $Z\in {\mathbb{R}}^{D\times (N+1)}$ is the input sequence (matrix form)
• $M$ is the model width (number of attention heads/FFN units)
• $\eta$ is the residual step size
• $f(\cdot ,\theta )$ and $h(\cdot ,w)$ are the attention and FFN encoders

2.2 Deep Transformer Structure

For a Transformer with $L$ blocks (depth), the structure evolves as:

$$\begin{array}{rl}{\bar{T}}_{\Theta }(H,t+\Delta t/2)& ={\text{Attn}}_{{\theta }^{t,1},\dots ,{\theta }^{t,M}}\left({\bar{T}}_{\Theta }(H,t),\Delta t/2\right)\\ {\bar{T}}_{\Theta }(H,t+\Delta t)& ={\text{MLP}}_{w^{t,1},\dots ,w^{t,M}}\left({\bar{T}}_{\Theta }(H,t+\Delta t/2),\Delta t/2\right)\end{array}$$

where $\Delta t=1/L$ and $t=0,\Delta t,\dots ,(L-1)\Delta t$.

2.3 From Discrete to Continuous: Mean-Field Limit

Key Insight: As both width $M\to \infty$ and depth $L\to \infty$, we can interpret the Transformer as a continuous dynamical system where parameters follow a probability distribution.

The continuous Transformer $T_{\rho }(H,t)$ satisfies the ODE:

$${\dot{T}}_{\rho }(H,t)={\int }_{\theta ,w}\frac{f(T_{\rho }(H,t),\theta )+h(T_{\rho }(H,t),w)}{2}\rho (\theta ,w,t)d\theta dw$$

where $\rho (\theta ,w,t)$ is the probability distribution of parameters at “time” $t$.

Interpretation: Each encoder $f$ or $h$ is conceptualized as a particle, and $\rho$ describes the distribution of these particles in parameter space.

2.4 Gradient Flow on Parameter Distribution

The training objective with ${\ell }_2$ regularization is:

$$\mathcal{Q}(\rho )=\mathcal{R}(\rho )+\frac{\lambda }{2}{\int }_0^1{\int }_{\theta ,w}\left(\Vert \theta {\Vert }_2^2+\Vert w{\Vert }_2^2\right)\rho (\theta ,w,t)d\theta dwdt$$

where the population risk is:

$$\mathcal{R}(\rho )={\mathbb{E}}_{\mu }\left[\frac{1}{2}{\left(\text{Read}[T_{\rho }(H,1)]-y(H)\right)}^2\right]$$

The Wasserstein gradient flow of this functional is given by the McKean-Vlasov PDE:

$$\frac{\partial \rho (\tau )}{\partial \tau }=\text{div}\left(\rho \cdot \nabla \frac{\delta \mathcal{Q}}{\delta \rho }\right)$$

This PDE governs how the parameter distribution evolves during training.

---

3. Wasserstein Gradient Flow Representation

3.1 Mathematical Framework

The Wasserstein space ${\mathcal{P}}_2({\mathbb{R}}^d)$ is the space of probability measures with finite second moment, equipped with the Wasserstein-2 distance:

$$W_2(\mu ,\nu )={\left(\underset{\gamma \in \Gamma (\mu ,\nu )}{\inf }\int \Vert x-y{\Vert }^{2}d\gamma (x,y)\right)}^{1/2}$$

where $\Gamma (\mu ,\nu )$ is the set of couplings between $\mu$ and $\nu$.

Why Wasserstein geometry? Unlike $L^2$ geometry on function spaces, Wasserstein geometry:

• Respects the nonlinear structure of probability distributions
• Enables analysis of propagation of chaos (convergence of particle systems)
• Provides a natural metric for gradient flows on probability measures

3.2 Functional Gradient Derivation

The functional derivative of $\mathcal{Q}$ with respect to $\rho$ is:

$$\frac{\delta \mathcal{Q}}{\delta \rho }(\theta ,w,t)={\mathbb{E}}_{\mu }\left[\text{Tr}\left({\frac{f(T_{\rho }(H,t),\theta )+h(T_{\rho }(H,t),w)}{2}}^{\top }{p}_{\rho }(H,t)\right)\right]+\frac{\lambda }{2}\left(\Vert \theta {\Vert }_2^2+\Vert w{\Vert }_2^2\right)$$

where $p_{\rho }(H,t)={\frac{\partial \mathcal{R}}{\partial T_{\rho }(H,t)}}^{\top }$ is the adjoint variable.

3.3 The Gradient Flow PDE

Explicitly, the Wasserstein gradient flow satisfies:

$$\frac{d{\rho }^{(\tau )}(\theta ,w,t)}{d\tau }={\text{div}}_{\theta }\left({\rho }^{(\tau )}{G}_f(\theta ,{\rho }^{(\tau )},t)\right)+{\text{div}}_w\left({\rho }^{(\tau )}{G}_h(w,{\rho }^{(\tau )},t)\right)$$

with gradient functions:

$$G_f(\theta ,\rho ,t)=\frac{1}{2}{\mathbb{E}}_{\mu }\left[{\nabla }_{\theta }\text{Tr}(f(T_{\rho }(H,t),\theta {)}^{\top }{p}_{\rho }(H,t))\right]+\lambda \theta$$ $$G_h(w,\rho ,t)=\frac{1}{2}{\mathbb{E}}_{\mu }\left[{\nabla }_w\text{Tr}(h(T_{\rho }(H,t),w{)}^{\top }{p}_{\rho }(H,t))\right]+\lambda w$$

Interpretation: Particles $\theta$ and $w$ flow “downhill” according to the gradient of the training objective, with the distribution $\rho$ transporting mass through this vector field.

3.4 Well-Posedness of the Gradient Flow

Proposition (Existence and Uniqueness): Under mild assumptions, there exists a unique solution $\{{\rho }^{(\tau )}{\}}_{\tau \ge 0}$ to the gradient flow equation, satisfying:

1. Bounded support: ${\rho }^{(\tau )}$ concentrates on $\left\{(\theta ,w):\Vert \theta {\Vert }^2+\Vert w{\Vert }^2\le R_{\tau }^2\right\}\times [0,1]$
2. Normalized mass: ${\int }_{\theta ,w}{\rho }^{(\tau )}(\theta ,w,t)d\theta dw=1$ for all $t$
3. Regularization effect: The $\lambda$ parameter controls parameter growth

Key insight: The $\lambda >0$ regularization is essential for well-posedness—it stabilizes the optimization by controlling both maximum and average parameter norms.

---

4. Partial Homogeneity and Local Lipschitz Smoothness

4.1 Beyond Full Homogeneity

Previous mean-field analyses of deep networks required full homogeneity:

$$g(T,c\theta )=c\cdot g(T,\theta )\text{for all }c>0$$

This property enabled certain technical arguments but excludes:

• Softmax attention (exponential in parameters)
• Sigmoid activations
• Complex interactions in multi-head attention

4.2 Partial Homogeneity Assumption

This paper introduces partial homogeneity that applies to only a subset of parameters:¹

Assumption 4 (Partial 1-Homogeneity): There exists a partition $\alpha =({\alpha }_1,{\alpha }_2)$ such that:

$$g(T,c{\alpha }_1,{\alpha }_2)=c\cdot g(T,{\alpha }_1,{\alpha }_2)\text{for all }c\in \mathbb{R}$$

Interpretation: Only a subset of parameters (${\alpha }_1$) scales the output homogeneously, while the remaining parameters (${\alpha }_2$) can have arbitrary dependence.

Example: In self-attention, the key-query projection matrices might exhibit partial homogeneity while the value projection does not.

4.3 Local Lipschitz Smoothness

Instead of global Lipschitz continuity of gradients (too restrictive for Transformers), we assume local Lipschitz smoothness in expectation:

Assumption 3 (Locally Lipschitz Continuous Gradient in Expectation): For any $L_T>0$ and $L_T$-Lipschitz continuous functions $T_1(H)$ and $T_2(H)$:

$${\mathbb{E}}_{\mu }\Vert {\nabla }_{\theta }f(T_1,\theta {)}_{:,i}-{\nabla }_{\theta }f(T_2,\theta {)}_{:,i}{\Vert }_2\le {\varphi }_{PT}(\Vert \theta \Vert ,K_T,L_T)\underset{H}{\sup }\Vert T_1-T_2{\Vert }_{2\text{-col}}$$

where ${\varphi }_{PT}$ is a continuous, monotonically increasing function.

Key relaxation: This assumption:

• Accommodates ReLU activations (where derivatives exist almost everywhere)
• Is satisfied by softmax attention with standard initialization
• Holds “in expectation” rather than pointwise

4.4 Comparison with Prior Work

Property	Prior Work (ResNets)	This Work (Transformers)
Homogeneity	Full homogeneity required	Partial homogeneity sufficient
Lipschitz constant	Global, uniform	Local, expectation-based
Activation functions	Smooth required	ReLU, softmax, sigmoid allowed
Analysis scope	Single encoder per block	Two distinct encoders (Attn + FFN)

---

5. Global Minimum Convergence Proof

5.1 Main Theorem: Gradient Flow Approximation

Theorem 3.1 (Gradient flow approximation of discretization): Define the empirical distribution:

$$\hat{\rho }(\tau ):=\frac{1}{ML}\sum _t\sum _{j=1}^M{\delta }_{({\theta }_{t,j}^{(\tau )},w_{t,j}^{(\tau )},t)}$$

Under Assumptions 1-3, the empirical distribution weakly converges to the Wasserstein gradient flow solution $\rho (\tau )$ almost surely as $L\to \infty$, $M/\log L\to \infty$.

Moreover, for any fixed $\tau >0$ and $\delta >0$, with probability at least $1-3\exp (-\delta )$:

$$\underset{s\in [0,\tau ]}{\sup }\left|\text{Read}[{\bar{T}}_{\Theta }(s)(H,t)]-\text{Read}[T_{\rho }(s)(H,t)]\right|\le C\left(L^{-1}+\frac{\sqrt{\delta }+\log (L+1)}{\sqrt{M}}\right)$$

Interpretation: Large-scale discrete Transformers can be arbitrarily well approximated by their mean-field limit, with approximation error decaying as $O(1/\sqrt{M}+1/L)$.

5.2 Global Convergence Theorem

Theorem 4.1 (Global convergence up to $\lambda$): Suppose the Wasserstein gradient flow $\rho (\tau )$ weakly converges to ${\rho }_{\infty }$, and the following conditions hold:

1. Bounded support: $\rho (\tau )$ concentrates on a bounded region for large $\tau$
2. Separation property: The support of ${\rho }_{\infty }$ at some depth $t^{\ast }$ spans a set that separates inner and outer parameter regions

Then, for any $ϵ>0$, there exists ${\tau }_0>0$ such that:

$$\underset{\tau \ge {\tau }_0}{\sup }\mathcal{R}(\Theta (\tau ))\le ϵ+C_1\left(L^{-1}+\frac{\sqrt{\delta }+\log (L+1)}{\sqrt{M}}\right)+C_2\lambda$$

Key implications:

1. As $M\to \infty$ and $L\to \infty$, the risk $\mathcal{R}(\Theta (\tau ))$ approaches zero
2. The additional term $C_2\lambda$ arises from regularization
3. By choosing sufficiently small $\lambda >0$, we achieve arbitrarily small training risk

5.3 Proof Sketch

The proof proceeds in three main steps:

Step 1: Establishing Continuity of Functional Gradient

Show that the functional gradient $\frac{\delta \mathcal{Q}}{\delta \rho }{|}_{{\rho }_{\infty }}$ remains constant if the derivative is constant over a region. This requires careful analysis of the Transformer dynamics.

Step 2: Bounding the Energy at Fixed Points

Derive the key bound for $\mathcal{Q}({\rho }_{\infty })\propto \lambda$ by analyzing the landscape of the functional energy through its derivatives.

Step 3: Finite-Time Risk Approximation

Show that the finite-time risk can approach this bound. Using Theorem 3.1, demonstrate that $\mathcal{Q}(\Theta ({\tau }_0))$ becomes sufficiently small, and since $\mathcal{Q}(\rho (\tau ))$ is non-increasing, it remains small for all $\tau \ge {\tau }_0$.

5.4 Practical Corollary

Corollary 4.1: For any fixed $\delta >0$ and $ϵ>0$, there exist constants ${\tau }_0,L_0,K_0>0$ such that:

$$\underset{\tau \ge {\tau }_0,L>L_0,M/\log L>K_0}{\sup }\mathcal{R}(\Theta (\tau ))\le (1/2+C_2{C}_{\lambda })ϵ$$

whenever $\lambda \le C_{\lambda }ϵ$.

Practical meaning: With sufficiently large width and depth, gradient flow training of Transformers guarantees convergence to near-zero training loss.

---

6. Novel Mean-Field Techniques for Transformers

6.1 Technical Contributions

This paper develops several novel techniques that extend mean-field theory to Transformer architectures:¹

6.1.1 Uniform Error Control

Previous works analyzed error at specific time points. This work achieves uniform error control over any finite time interval $[0,\tau ]$:

$$\underset{s\in [0,\tau ]}{\sup }\Vert {\bar{T}}_{\Theta }(s)-T_{\rho }(s){\Vert }_F\le C(M,L)$$

This enables continuous monitoring of maximum error across the training trajectory.

6.1.2 Two-Encoder Analysis

Unlike ResNet models with a single encoder per block, Transformers use two distinct encoders ($f$ for attention, $h$ for FFN) that alternate. The analysis:

• Treats each encoder separately with appropriate regularity conditions
• Unifies them through the average $(f+h)/2$ in the continuous limit
• Provides rigorous validation of the “ensemble of paths” concept

6.1.3 Propagation of Chaos Framework

Extended the classical propagation of chaos theory to the Transformer setting:

• Particle systems with non-i.i.d. initializations at different depths
• Uniform bounds on particle differences using Grönwall’s inequality
• Concentration estimates for empirical distributions

6.2 Novel Assumptions and Their Justification

Assumption	Novel Element	Justification
2	Column-wise norms for sequential data	Natural for token-based inputs
3	Local Lipschitz in expectation	Accommodates ReLU/softmax
4	Partial homogeneity	Enables softmax/sigmoid activations

6.3 Verification for Concrete Architectures

The paper provides explicit verification of assumptions for specific Transformer configurations, demonstrating that:

• Feed-forward layers with ReLU satisfy the universal kernel property
• Self-attention layers can serve as universal approximators under certain conditions
• The partial homogeneity condition holds for standard architectural choices

---

7. Connections to Prior Work

7.1 Neural Network Mean-Field Theory

The foundation builds on the seminal work of Mei, Montanari, and Nguyen (2018) on two-layer networks:¹

$$\text{NTK limit}\stackrel{\text{mean-field}}{\to }\text{Wasserstein gradient flow}$$

Key developments:

• Chizat & Bach (2018): Established global convergence for overparameterized models using optimal transport
• Lu et al. (2020): Extended to deep ResNets with skip connections
• Ding et al. (2021, 2022): Proved overparameterization guarantees for deep ResNets

7.2 ResNet vs Transformer Analysis

Aspect	ResNet Analysis	Transformer Analysis
Structure	Single encoder per block	Two distinct encoders
Homogeneity	Full homogeneity required	Partial homogeneity sufficient
Analysis method	ODE discretization	PDE/ODE hybrid
Skip connections	Identity shortcuts	Learned attention patterns

7.3 Neural ODE Connection

The Transformer ODE perspective connects to Neural ODE theory:

• ResNets: $\dot{x}=f(x)$ with identity-like skip connections
• Transformers: $\dot{T}={\mathbb{E}}_{\rho }[f(T,\theta )+h(T,w)]$ with stochastic averaging

This provides a unifying framework for understanding deep network training dynamics.

7.4 In-Context Learning Connections

Recent work on in-context learning (ICL) provides complementary perspectives:

• Ahn et al., 2023: Showed Transformers can perform ICL via linear regression approximation
• Kim & Suzuki, 2024: Analyzed mean-field dynamics for in-context feature learning
• This work: Provides optimization-theoretic foundation for these phenomena

---

8. Future Directions and Open Problems

8.1 Theoretical Extensions

1. Direct gradient descent analysis: Current results use continuous gradient flow; discrete-time analysis with step sizes is needed
2. Generalization bounds: Connect optimization convergence to finite-sample generalization
3. Self-attention as universal kernel: Rigorous conditions for attention’s approximation capacity

8.2 Practical Implications

1. Initialization schemes: Theory suggests optimal scaling for width/depth tradeoffs
2. Regularization tuning:指导选择合适的 $\lambda$ 以平衡收敛速度和泛化
3. Architecture design: Principles for choosing attention/FFN ratios

8.3 Open Questions

• Can we remove the partial homogeneity assumption entirely?
• What are the necessary and sufficient conditions for global convergence?
• How does the theory extend to mixture of experts and sparse Transformers?

---

9. Summary

This paper establishes the first rigorous global convergence theory for large-scale Transformer training using mean-field methods.¹

Key Contributions:

1. Mean-field limit construction: Showed that as width and depth go to infinity, Transformers converge to a Wasserstein gradient flow described by a PDE

2. Novel technical assumptions: Introduced partial homogeneity and local Lipschitz smoothness—weaker conditions that accommodate real Transformer architectures

3. Two main theorems:

   • Theorem 3.1: Close approximation between discrete gradient flow and continuous limit
   • Theorem 4.1: Global convergence to near-zero training loss

4. Practical implications: Demonstrated that basic gradient flow can successfully navigate the complex non-convex landscape to find optimal solutions

Significance: These results provide a theoretical foundation for understanding why Transformers train so successfully in practice, and open new avenues for optimization theory of modern deep learning architectures.

---

References

---

Related Topics

• Transformer as Differential Equation — Continuous modeling of Transformers
• Neural Tangent Kernel Theory — Infinite-width neural network analysis
• Unified Gradient Flow Convergence — General convergence theory for neural networks
• Backpropagation as Gradient Flow — Optimization dynamics foundations
• ResNet as Dynamical System — Neural ODE perspective on skip connections

Footnotes

1. Gao, C., Cao, Y., Li, Z., He, Y., Wang, M., Liu, H., Klusowski, J. M., & Fan, J. (2024). Global Convergence in Training Large-Scale Transformers. NeurIPS 2024. arXiv:2410.23610 ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷