Year: 2024
Type: conference-paper
Abstract: 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.
DOI: 10.52202/079017-0921
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: CPS: Medium: Collaborative Research: Provably Safe and Robust Multi-Agent Reinforcement Learning with Applications in Urban Air Mobility, CAREER: Statistical Learning with Recursive Partitioning: Algorithms, Accuracy, and Applications, Deep Learning and Random Forests for High-Dimensional Regression, G4720003262, G5429545050, G6671297155, DMS/NIGMS 2: Collaborative Research: Developing Statistical Learning Methods for Revealing the Molecular Signatures of Microvascular Changes in Neural Injury, Interface of Statistical Learning and Optimal Decisions, HDR TRIPODS: Data Science Principles of the Human-Machine Convergence, G7712150002, RI: Medium: Provable Reinforcement Learning with Function Approximation and Neural Networks, G8876996369
OpenAlex ID: W4415798130