Research map / Optimization / Implicit bias and generalization
Global Convergence in Training Large-Scale Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that gradient flow on a deep, wide Transformer with weight decay converges to a Wasserstein gradient flow as width and depth grow, and reaches a near-global minimum when the regularization is small.
Problem
Optimization guarantees for gradient-based training of large-scale Transformers are poorly understood; prior mean-field results for deep networks require full homogeneity and global Lipschitz smoothness, which Transformers do not satisfy.
Approach
Builds a mean-field limit of residual Transformers with alternating self-attention and feed-forward encoders (averaged in the continuous limit), and proves gradient-flow convergence to a Wasserstein gradient flow using only partial 1-homogeneity and local Lipschitz smoothness in expectation.
Key results
- Discrete Transformer gradient flow is shown to be uniformly close to the Wasserstein gradient flow, with error C(L^{-1}+sqrt((delta+log(L+1))/M)) for width M and depth L.
- Global convergence result: with small weight decay lambda, the risk of gradient flow approaches within epsilon of zero as L and M/log L grow (Corollary: lambda <= C_lambda epsilon suffices).
- Appendix experiments on Vision Transformers on CIFAR-10 reportedly reach near-zero training loss and 100% training accuracy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase width and depth toward the large-scale limit and train with gradient flow | Transformer · training risk | risk: \epsilon + C_1\Big(L^{-1}+\sqrt{\frac{\delta+\log(L+1)}{M}}\Big)+C_2\lambda | Helped |
| use positive $l_2$ weight decay regularization | Transformer gradient flow · optimization dynamics | well-posedness | Helped |
Limitations
- Convergence results are asymptotic and give no explicit rate.
- Global convergence relies on assumptions (long-time boundedness, support separation) justified only at a high level.
- Universal approximation of self-attention layers is assumed or left open rather than proven in general.
Experimental setup
Datasets: CIFAR-10
Benchmarks: CIFAR-10
Metrics: training loss training accuracy
Builds on: mean-field analysis of overparameterized networks neural ODE view of ResNets global convergence via homogeneity (Chizat and Bach)
Closest papers by meaning
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.89)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.88)
- Scaling Laws for Sparsely-Connected Foundation Models (ICLR 2024 · similarity 0.87)
- Learning Theory of Transformers: Local-to-Global Approximation via Softmax Partition of Unity (NeurIPS 2026 · similarity 0.87)
- Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data (NeurIPS 2024 · similarity 0.86)
- Infinite Limits of Multi-head Transformer Dynamics (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.