Understanding Transformer Optimization via Gradient Heterogeneity

TL;DR

This paper investigates why Adam outperforms SGD in training Transformers by analyzing gradient heterogeneity, revealing that Adam's normalization makes it less sensitive to this issue and explaining its empirical success.

Contribution

The study provides a theoretical framework linking gradient heterogeneity to optimizer performance in Transformers, highlighting the role of normalization placement and comparing sign-based methods to SGD.

Findings

Gradient heterogeneity degrades SGD convergence but less affects Adam.

Post-LN architectures exhibit higher gradient heterogeneity.

Experimental validation in NLP and vision confirms theoretical insights.

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 origin of gradient heterogeneity in Transformer architectures and show that it is strongly influenced by the placement of layer normalization, with Post-LN architectures exhibiting particularly pronounced heterogeneity. Experimental results from fine-tuning Transformers in both NLP and vision domains validate our theoretical analysis. Code is available at https://github.com/tom4649/gradient-heterogeneity.