🤖 AI

Gradient Heterogeneity Complements Hessian Heterogeneity in Transformer Optimization

This paper reveals that gradient heterogeneity, exacerbated by Post-LN architectures, degrades SGD convergence in Transformers but is effectively mitigated by coordinate-wise adaptive methods like Adam and SignSGD, which are less sensitive to gradient scale variations.

Original authors: Akiyoshi Tomihari, Issei Sato

Published 2026-08-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Akiyoshi Tomihari, Issei Sato

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast and rapidly evolving landscape of artificial intelligence, a specific type of computer model known as the Transformer has become the engine behind some of the most powerful language and image systems in existence. These models are trained by a process called optimization, where the computer repeatedly adjusts its internal settings to reduce errors, much like a hiker trying to find the lowest point in a foggy valley. For decades, the standard tool for this task was a method called stochastic gradient descent, which takes small, calculated steps based on the immediate slope of the terrain. However, as models grew larger and more complex, researchers found that this traditional tool often struggled, stumbling or moving too slowly to be practical. Consequently, the field largely shifted to a different tool called Adam, which adapts its step size dynamically and has proven far more effective at training these massive systems. Yet, a lingering mystery remained: why does Adam work so much better than the older method, and what is actually happening inside the model that makes the difference?

A team of researchers at The University of Tokyo has now peeled back the layers of this mystery, offering a clear explanation rooted in how the model's internal landscape is shaped. They discovered that the difficulty in training these models comes from a specific kind of unevenness, or heterogeneity, in the data the model processes. Imagine a landscape where some hills are incredibly steep and narrow while others are gentle and wide; the traditional method, which treats every direction equally, gets confused by this mix, often taking steps that are too large for the steep parts and too small for the gentle ones. The researchers found that this unevenness is not just a random quirk but a structural feature of how these models are built, particularly influenced by where certain stabilizing components are placed within the architecture.

The study reveals that the superior performance of the adaptive optimizer is not due to a magical ability to handle noise, as some previous theories suggested, but rather because it effectively ignores the scale of the steps in a way that the traditional method does not. By analyzing the mathematical behavior of these optimizers, the authors showed that the traditional method is highly sensitive to the variations in gradient strength across different parts of the model. When one part of the model requires a tiny adjustment and another needs a massive one, the traditional method tries to apply a single, uniform step size, which leads to inefficiency. In contrast, the adaptive method, and a simpler version of it called SignSGD, adjusts its approach by looking at the direction of the change rather than its magnitude, allowing it to navigate the jagged terrain of the model's landscape with much greater ease.

To prove this, the researchers conducted a rigorous theoretical analysis, deriving mathematical bounds that describe how many steps each method needs to reach a solution. Their work demonstrates that the traditional method's progress is slowed down significantly when the model exhibits high levels of this gradient unevenness, whereas the adaptive methods remain robust. They further traced the source of this unevenness back to the design of the Transformer itself, specifically the placement of a component called layer normalization. They found that when this component is placed after the main processing steps, it amplifies the differences in gradient strength, making the landscape even more treacherous for the traditional optimizer. When placed before the steps, the landscape is smoother, and the traditional method performs better, though the adaptive method still holds an advantage.

The team validated these theoretical insights with real-world experiments, fine-tuning models on both language and vision tasks. They observed that in scenarios where the gradient unevenness was high, the traditional optimizer struggled to converge, taking far longer to learn, while the adaptive methods and their sign-based counterparts trained efficiently. This confirmed that the key to the performance gap lies in how these optimizers handle the structural heterogeneity of the model. The findings suggest that the success of modern AI is not just about having more data or computing power, but about using the right mathematical tools to navigate the specific, uneven topography created by the architecture of the models themselves. By understanding this dynamic, researchers can better design training strategies and potentially develop new optimizers that are even more effective for the next generation of artificial intelligence.

Technical Summary: Gradient Heterogeneity Complements Hessian Heterogeneity in Transformer Optimization

Problem Statement

Transformers have achieved significant success, particularly in language models, yet they remain difficult to optimize using standard Stochastic Gradient Descent (SGD). In practice, training relies heavily on adaptive optimizers like Adam. While the performance gap between Adam and SGD is well-documented, the underlying mechanisms are not fully understood. Previous explanations have attributed this gap to stochastic gradient noise, heavy-tailed label distributions, or Hessian heterogeneity (differences in block-wise Hessian spectra). However, these factors do not fully account for the observed behavior, particularly in fine-tuning regimes with limited labeled data. This paper posits that the performance gap arises from the interaction between Hessian heterogeneity and a newly characterized phenomenon: gradient heterogeneity.

Methodology

The authors employ a theoretical analysis of optimization iteration complexity, comparing gradient-based methods (SGD) against sign-based methods (SignSGD), which serve as an analytically tractable proxy for Adam-like behavior.

Theoretical Framework

Definitions and Assumptions:
- Gradient Heterogeneity: Defined as the variation in gradient norms across parameter blocks $\{ \|\nabla L(\theta)\|_b \}$ .
- Hessian Heterogeneity: Defined as differences in block-wise Hessian spectra.
- Assumptions: The analysis assumes a fine-tuning setting where parameters remain within a localized region, the Hessian is Lipschitz continuous, and the Hessian matrix is approximately block-diagonal.
- Gradient-Hessian Correlation: The authors observe and utilize the empirical finding that large gradient norms in specific blocks often correlate with large Hessian operator norms in those same blocks.
Complexity Measures:
The paper introduces two weighted Hessian norms to bound the local curvature encountered during optimization:
- $\Lambda_G$ (Gradient-weighted): Weights Hessian block norms by the squared gradient norms of those blocks. This quantity controls the iteration complexity of gradient-based methods (SGD).
- $\Lambda_P$ (Parameter-weighted): Weights Hessian block norms by the parameter dimensions of those blocks. This quantity controls the iteration complexity of sign-based methods (SignSGD).
Derivation of Bounds:
Using steepest descent interpretations under different norms ( $\ell_2$ for SGD, $\ell_\infty$ for SignSGD), the authors derive upper bounds on iteration complexity for both deterministic and stochastic settings.
- For SGD, the complexity is proportional to $\Lambda_G$ . Under gradient-Hessian correlation, blocks with large gradients also have large curvature, causing $\Lambda_G$ to become large and slowing convergence.
- For SignSGD, the complexity is proportional to $\Lambda_P$ . Since sign-based updates use unit-magnitude directions, the update magnitude is independent of the gradient norm. Consequently, $\Lambda_P$ is less sensitive to the alignment of large gradients with high curvature, resulting in a smaller effective complexity bound.
Architectural Analysis:
The paper investigates the origin of gradient heterogeneity in Transformer architectures by analyzing the Jacobian structure of Layer Normalization (LN). It compares Post-LN (LN after residual connections) and Pre-LN (LN before residual connections). Theoretical derivation shows that in Post-LN, the LN Jacobian multiplies the entire residual branch, compounding input-dependent scale variations and amplifying gradient heterogeneity. In Pre-LN, the identity path bypasses the LN Jacobian, mitigating this effect.

Key Contributions

Theoretical Bounds: The authors derive upper bounds on iteration complexity for both gradient-based and sign-based methods in deterministic and stochastic settings. These bounds demonstrate that SGD is controlled by $\Lambda_G$ , while sign-based methods are controlled by $\Lambda_P$ .
Gradient Heterogeneity Characterization: The study identifies gradient heterogeneity as a critical factor complementing Hessian heterogeneity. It provides quantitative measures (Gini coefficients) and visualizations showing that Transformers exhibit significantly higher gradient heterogeneity than CNNs (e.g., ResNet).
Architectural Insight: The paper links gradient heterogeneity to architectural choices, specifically showing that Post-LN architectures exhibit pronounced heterogeneity due to the multiplicative nature of the LN Jacobian, whereas Pre-LN architectures are more stable.
Learning Rate Implications: The analysis suggests that SignSGD (and by extension, Adam) benefits from learning rates scaled by the $\ell_1$ -norm of the gradient, consistent with steepest descent under the $\ell_\infty$ -norm.

Experimental Results

The theoretical claims are validated through fine-tuning experiments on NLP (RoBERTa on SuperGLUE/GLUE) and Vision (ViT, ResNet on Flowers102, Aircraft) tasks.

Gradient Heterogeneity: Experiments confirm that Transformers (RoBERTa) exhibit substantially higher gradient heterogeneity (Gini coefficient $\approx 0.94$ ) compared to ResNet ( $\approx 0.43$ ).
Post-LN vs. Pre-LN: Models with Post-LN normalization show higher Gini coefficients (0.941) compared to Pre-LN (0.880), supporting the theoretical claim that Post-LN amplifies heterogeneity.
Optimizer Performance:
- SGD: Exhibits significantly slower training loss convergence on high-heterogeneity tasks (e.g., RoBERTa on RTE) compared to adaptive optimizers, even with gradient clipping and tuned learning rate schedules.
- SignSGD: Achieves performance comparable to Adam and converges reliably on both high-heterogeneity (RoBERTa) and lower-heterogeneity (ViT) tasks.
- Momentum: Experiments indicate that the performance gap is primarily driven by adaptive learning rates (coordinate-wise scaling) rather than momentum, as SignSGD without momentum still outperforms SGD.

Significance and Claims

The paper claims to provide a more complete theoretical explanation for the Adam-SGD gap in Transformer optimization by introducing gradient heterogeneity as a key factor. The authors argue that:

The superiority of Adam-like methods is not solely due to handling gradient noise or heavy-tailed distributions, but fundamentally stems from their sign-based nature (coordinate-wise updates) which mitigates the negative effects of gradient and Hessian heterogeneity.
The interaction between gradient heterogeneity and Hessian heterogeneity creates a scenario where gradient-based methods (SGD) face a much larger effective curvature ( $\Lambda_G$ ) than sign-based methods ( $\Lambda_P$ ).
Architectural choices, specifically the placement of Layer Normalization (Post-LN vs. Pre-LN), directly influence the degree of gradient heterogeneity, thereby affecting optimization difficulty.

The study concludes that understanding these heterogeneities is essential for explaining why Transformers are difficult to optimize with SGD and why adaptive, sign-based methods are effective. The findings suggest that block-wise normalization techniques could theoretically bridge part of this gap by reducing sensitivity to gradient heterogeneity, even without full coordinate-wise adaptivity.