Gradient differences across layers help explain why adaptive optimizers beat SGD for Transformers

Training large Transformer models with plain stochastic gradient descent (SGD) is notoriously inefficient compared to adaptive methods like Adam, yet the underlying reasons have remained murky. A new analysis published on arXiv digs into this puzzle by examining how variations in gradient strength across different parameter blocks affect optimization during fine-tuning.

The problem with heterogeneous gradients#

The research introduces the concept of gradient heterogeneity—the spread in gradient norms when you compare one group of parameters to another. The key finding is that when gradients vary significantly in magnitude across network layers, standard SGD struggles to converge efficiently. The same applies to the Hessian (the curvature of the loss landscape), which also tends to be uneven. Together, these two sources of variation compound the optimization challenge.

However, not all optimizers suffer equally. Sign-based methods, which rely on the direction of gradients rather than their raw magnitudes, prove far more resilient to this kind of heterogeneity. This insight motivated the researchers to use SignSGD as an analytically tractable stand-in for Adam-like behavior, since both methods perform coordinate-wise updates that dampen the impact of gradient scale variations.

Architecture shapes gradient behavior#

The analysis reveals that layer normalization placement profoundly influences how gradients scatter across the network. Post-LN architectures—where normalization is applied after the main transformation in each block—show especially pronounced gradient heterogeneity, while alternative placements (like Pre-LN) distribute gradients more evenly.

Theoretical insights and learning rate implications#

By analyzing how SGD and SignSGD follow steepest descent in different geometric spaces, the researchers derived convergence bounds that inform how learning rates should be scaled when using sign-based methods. The theory carries practical implications for tuning optimizers on Transformer models.

Validation across domains#

The theoretical predictions held up in experiments spanning both natural language processing and computer vision fine-tuning tasks, grounding the abstract analysis in real-world training scenarios.