@inproceedings{
gao2024global,
title={Global Convergence in Training Large-Scale Transformers},
author={Cheng Gao and Yuan Cao and Zihao Li and Yihan He and Mengdi Wang and Han Liu and Jason Matthew Klusowski and Jianqing Fan},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=9wtlfRKwZS}
}
Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings are not well-understood. This paper rigorously analyzes the convergence properties of gradient flow in training Transformers with weight decay regularization. First, we construct the mean-field limit of large-scale Transformers, showing that as the model width and depth go to infinity, gradient flow converges to the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches a global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small. Our analysis is based on a series of novel mean-field techniques that adapt to Transformers. Compared with existing tools for deep networks (Lu et al., 2020) that demand homogeneity and global Lipschitz smoothness, we utilize a refined analysis assuming only $\textit{partial homogeneity}$ and $\textit{local Lipschitz smoothness}$. These new techniques may be of independent interest.
Global Convergence in Training Large-Scale Transformers
Cheng Gao Yuan Cao Zihao Li Yihan He Mengdi Wang Han Liu Jason Matthew Klusowski Jianqing Fan