![[2502.00213] Understanding Transformer Optimization via Gradient Heterogeneity](https://arxiv.org/static/browse/0.3.4/images/arxiv-logo-fb.png){kind=logo}
[2502.00213] Understanding Transformer Optimization via Gradient Heterogeneity
Summary
This paper explores the optimization challenges of Transformer models, focusing on gradient heterogeneity and its impact on convergence when using stochastic gradient descent (SGD) versus adaptive optimizers like Adam.
Why It Matters
Understanding the optimization dynamics of Transformer models is crucial for improving their training efficiency and performance. This study sheds light on why adaptive methods like Adam outperform SGD, providing insights that could lead to better optimization strategies in machine learning applications.
Key Takeaways
- Gradient heterogeneity affects the convergence of SGD in Transformer models.
- Adam's coordinate-wise normalization allows it to perform better than SGD by relying on gradient signs.
- Layer normalization placement significantly influences gradient heterogeneity in Transformer architectures.
- The study provides theoretical bounds on iteration complexity for learning rate scaling in SignSGD.
- Experimental validation across NLP and vision domains supports the theoretical findings.
Computer Science > Machine Learning arXiv:2502.00213 (cs) [Submitted on 31 Jan 2025 (v1), last revised 18 Feb 2026 (this version, v4)] Title:Understanding Transformer Optimization via Gradient Heterogeneity Authors:Akiyoshi Tomihari, Issei Sato View a PDF of the paper titled Understanding Transformer Optimization via Gradient Heterogeneity, by Akiyoshi Tomihari and 1 other authors View PDF HTML (experimental) Abstract:Transformers are difficult to optimize with stochastic gradient descent (SGD) and largely rely on adaptive optimizers such as Adam. Despite their empirical success, the reasons behind Adam's superior performance over SGD remain poorly understood. In this study, we analyze the optimization of Transformer models through the lens of \emph{gradient heterogeneity}, defined as the variation in gradient norms across parameter blocks. We provide a theoretical analysis showing that gradient heterogeneity, together with Hessian heterogeneity, degrades the convergence of gradient-based methods such as SGD, while sign-based methods are substantially less sensitive to this effect. Adam's coordinate-wise normalization makes its update directions depend mainly on gradient signs, so Adam can be interpreted as a soft variant of SignSGD. Our analysis uses the fact that SGD and SignSGD follow steepest descent directions under different norms, and derives upper bounds on the iteration complexity with implications for learning rate scaling in SignSGD. We further investigate the o...
