Unraveling the Gradient Descent Dynamics of Transformers
This work addresses a foundational problem in machine learning by providing theoretical insights into Transformer training dynamics, which is incremental as it builds on existing knowledge to clarify optimization behaviors.
The study tackled the lack of theoretical understanding of Transformer optimization dynamics by analyzing which architectures allow Gradient Descent to achieve guaranteed convergence and under what conditions rapid convergence occurs, finding that with appropriate initialization, GD can train Transformers to global optimal solutions, especially with large input dimensions, though Softmax attention may lead to suboptimal local solutions while Gaussian attention performs more favorably.
While the Transformer architecture has achieved remarkable success across various domains, a thorough theoretical foundation explaining its optimization dynamics is yet to be fully developed. In this study, we aim to bridge this understanding gap by answering the following two core questions: (1) Which types of Transformer architectures allow Gradient Descent (GD) to achieve guaranteed convergence? and (2) Under what initial conditions and architectural specifics does the Transformer achieve rapid convergence during training? By analyzing the loss landscape of a single Transformer layer using Softmax and Gaussian attention kernels, our work provides concrete answers to these questions. Our findings demonstrate that, with appropriate weight initialization, GD can train a Transformer model (with either kernel type) to achieve a global optimal solution, especially when the input embedding dimension is large. Nonetheless, certain scenarios highlight potential pitfalls: training a Transformer using the Softmax attention kernel may sometimes lead to suboptimal local solutions. In contrast, the Gaussian attention kernel exhibits a much favorable behavior. Our empirical study further validate the theoretical findings.