Global Convergence in Training Large-Scale Transformers

Understand

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.