Table of Contents
Fetching ...

Understanding Why Adam Outperforms SGD: Gradient Heterogeneity in Transformers

Akiyoshi Tomihari, Issei Sato

Abstract

Transformers are challenging to optimize with SGD and typically require adaptive optimizers such as Adam. However, the reasons behind the superior performance of Adam over SGD remain unclear. In this study, we investigate the optimization of transformers by focusing on gradient heterogeneity, defined as the disparity in gradient norms among parameters. Our analysis shows that gradient heterogeneity hinders gradient-based optimization, including SGD, while sign-based optimization, a simplified variant of Adam, is less affected. We further examine gradient heterogeneity in transformers and show that it is influenced by the placement of layer normalization. Experimental results from fine-tuning transformers in both NLP and vision domains validate our theoretical analyses. This study provides insights into the optimization challenges of transformers and offers guidance for designing future optimization algorithms. Code is available at https://github.com/tom4649/gradient-heterogeneity.

Understanding Why Adam Outperforms SGD: Gradient Heterogeneity in Transformers

Abstract

Transformers are challenging to optimize with SGD and typically require adaptive optimizers such as Adam. However, the reasons behind the superior performance of Adam over SGD remain unclear. In this study, we investigate the optimization of transformers by focusing on gradient heterogeneity, defined as the disparity in gradient norms among parameters. Our analysis shows that gradient heterogeneity hinders gradient-based optimization, including SGD, while sign-based optimization, a simplified variant of Adam, is less affected. We further examine gradient heterogeneity in transformers and show that it is influenced by the placement of layer normalization. Experimental results from fine-tuning transformers in both NLP and vision domains validate our theoretical analyses. This study provides insights into the optimization challenges of transformers and offers guidance for designing future optimization algorithms. Code is available at https://github.com/tom4649/gradient-heterogeneity.

Paper Structure

This paper contains 90 sections, 7 theorems, 100 equations, 75 figures, 14 tables.

Key Result

Theorem 4.7

Assume $\delta_{D}< \min(\Lambda_{G},\Lambda_{P})/3$. Then, the iteration complexities in deterministic settings are bounded as follows. For the gradient-based sequence, suppose that $\varepsilon < \frac{\Lambda_G^2}{\rho_H \sqrt{P}}$ holds and that learning rate at time $t$ satisfies $\eta_t = \zet For the sign-based sequence, suppose that $\varepsilon < \frac{\Lambda_P^2}{\rho_H \sqrt{P}}$ holds

Figures (75)

  • Figure 1: Correlation between the gradient norm and the maximum Hessian eigenvalue. Each point represents the mean value for a parameter block (pre-trained RoBERTa on RTE).
  • Figure 2: Correlation between the full-batch gradient and gradient error. Each point represents the absolute values of a coordinate (pre-trained RoBERTa on RTE).
  • Figure 3: RoBERTa on RTE
  • Figure 4: ResNet18 on Flowers102
  • Figure 6: RoBERTa
  • ...and 70 more figures

Theorems & Definitions (20)

  • Definition 4.4: Gradient heterogeneity
  • Definition 4.5: Weighted Hessian complexity
  • Definition 4.6: Iteration complexity
  • Theorem 4.7: Deterministic setting
  • Theorem 4.9: Stochastic setting
  • Lemma C.1
  • proof
  • Lemma C.2
  • proof
  • proof : Proof of gradient-based sequence
  • ...and 10 more