PaperlayerView exact revision on arXiv

Source-linked AI summary

Why Transformers Need Adam: A Hessian Perspective

Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, Zhi-Quan Luo

arXiv:2402.16788v4cs.LGcs.AI

TL;DR

The paper addresses why SGD performs worse than Adam on Transformers. It examines Hessian spectra across parameter blocks and reports that block heterogeneity is associated with worse SGD performance across the examined settings. Its initial theory points to the use of one learning rate across blocks as a limitation that block-specific rates can mitigate.

  • Problem

    Why SGD underperforms Adam on Transformers remains unclear, despite Transformers’ reliance on Adam and SGD’s strong performance on CNNs.

    Source: §1 · p. 1

  • Method

    The paper examines blockwise Hessian spectra and compares optimizer behavior across Transformers, CNNs, MLPs, and quadratic problems, with initial theoretical analysis.

    Source: §1 · p. 2

  • Results

    Across examined Transformers, CNNs, and MLPs, SGD performs worse than Adam on problems with block heterogeneity but similarly on problems without it.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The initial analysis indicates that one learning rate across blocks limits gradient descent on heterogeneous problems, and different rates across blocks can mitigate this limitation.

    Source: §1 · p. 2

    The asymptotic Adam analysis requires β2 = 1; the authors report a non-convergence issue for β2 < 1.

    Source: §4.2 · p. 9

Abstract

from arXiv · show

SGD performs worse than Adam by a significant margin on Transformers, but the reason remains unclear. In this work, we provide an explanation through the lens of Hessian: (i) Transformers are "heterogeneous": the Hessian spectrum across parameter blocks vary dramatically, a phenomenon we call "block heterogeneity"; (ii) Heterogeneity hampers SGD: SGD performs worse than Adam on problems with block heterogeneity. To validate (i) and (ii), we check various Transformers, CNNs, MLPs, and quadratic problems, and find that SGD can perform on par with Adam on problems without block heterogeneity, but performs worse than Adam when the heterogeneity exists. Our initial theoretical analysis indicates that SGD performs worse because it applies one single learning rate to all blocks, which cannot handle the heterogeneity among blocks. This limitation could be ameliorated if we use coordinate-wise learning rates, as designed in Adam.

1 Introduction

The paper asks why SGD underperforms Adam on Transformers and investigates the gap through Hessian spectra. It proposes block heterogeneity across parameter blocks as an important explanation, while noting that full Hessian spectra alone do not distinguish Transformers from CNNs.

1 Introduction

  • Full Hessian spectra are often similar for CNNs and Transformers despite their different optimizer behavior, motivating a finer-grained investigation.The authors report that the full spectra did not reveal critical features associated with the performance gap.Source: §1 · p. 1, Figure 1 · p. 2
  • CNNs stack similar convolution blocks, whereas Transformers combine disparate parameter blocks, motivating analysis of each block’s Hessian spectrum.The paper calls the spectrum of each parameter block the blockwise Hessian spectrum.Source: §1 · p. 2
  • Transformers exhibit block heterogeneity, and across examined Transformers, CNNs, and MLPs, SGD performs worse than Adam with heterogeneity but similarly without it.The paper also reports that gradient descent underperforms Adam on heterogeneous quadratic problems, with comparable performance otherwise.Source: §1 · p. 2
  • The paper’s initial theoretical analysis attributes slower gradient descent on heterogeneous quadratic problems to one learning rate across blocks, a limitation mitigated by block-specific rates.The authors emphasize that block heterogeneity is at least one important cause of the SGD–Adam gap, not necessarily the only cause.Source: §1 · p. 2

2 Problem Settings and Initial Attempts

The section describes the models and Hessian-spectrum estimation setup, finds that full spectra do not explain the optimizer gap, and motivates blockwise spectra as a more informative view. A near-block-diagonal Hessian structure supports focusing on principal blocks.

2.1 Problem Settings

  • The study estimates full and blockwise Hessian spectra, using Stochastic Lanczos Quadrature to approximate spectra that are intractable to compute and store directly.CNN experiments use ResNet18 and VGG16 on ImageNet; Transformer experiments cover GPT2, ViT-base, BERT, and GPT2-nano on the listed tasks.Source: §2.1 · p. 3

2.2 Full Hessian Spectrum Is Not Informative Enough

  • Full Hessian spectra do not reveal critical phenomena linked to the optimizer gap: CNNs and Transformers show similar eigenvalue dispersion, shape, and training evolution.The spectra begin approximately symmetric around zero; during training, most negative eigenvalues disappear and the shape develops a bulk and outliers.Source: §2.2 · p. 3

2.3 Motivations of Investigating Blockwise Hessian Spectra

  • Transformer blocks differ in design and are stacked non-sequentially, while CNNs repeatedly stack similar convolution blocks; prior near-block-diagonal findings suggest principal blocks retain much Hessian information.This motivates examining principal-block spectra to distinguish the model families.Source: §2.3 · pp. 3–4
  • Blockwise Hessian spectra carry more information than full spectra for distinguishing CNNs and Transformers; near-block-diagonal structure appears at 1% step and persists during training, with each block representing an output neuron.Parameter blocks follow the default PyTorch partition and need not correspond to whole layers.Source: §2.1 · p. 3, §2.3 · p. 4

3 Main Results

Transformers show greater blockwise Hessian-spectrum differences than CNNs, and SGD tends to lag Adam when this heterogeneity is present. Reduced heterogeneity after training or fine-tuning accompanies a narrower optimizer gap, while JS0 is proposed as a pre-training indicator of that gap.

3.1 Transformers Exhibit Block Heterogeneity in Hessian, while CNNs Do Not

  • Transformers have different blockwise Hessian spectra, whereas CNN block spectra are similar, a contrast termed block heterogeneity versus block homogeneity.The comparison uses initial Hessian spectra; the Figure 4 JS distances are significantly smaller for CNNs than Transformers.Source: Figure 3 · p. 4, §3.1 · p. 4, Figure 4 · p. 5

3.2 SGD Performs Worse than Adam on Various Tasks with Block Heterogeneity

  • Across the checked tasks, SGD performs worse than Adam when block heterogeneity exists, including a scaled MLP where the gap grows with heterogeneity and MLP-mixer.The MLP-mixer also exhibits initial Hessian heterogeneity; the pattern does not depend on Transformers or attention mechanisms.Source: §3.2 · pp. 5–6

3.3 Reduced Block Heterogeneity in Pre-trained Transformers

  • After training or fine-tuning, block heterogeneity tends to reduce; on fine-tuned GPT2, SGD remains slower than Adam but reaches similar loss at convergence.The paper reports this reduction for pretrained GPT2 and along ViT-base training.Source: §3.3 · p. 6, Figure 6 · p. 6
  • ViT-base heterogeneity declines during training, and SGD performs better when initialized at a less-heterogeneous checkpoint, such as 50% of total training steps.The authors hypothesize that attenuation of Hessian heterogeneity after training may be common, leaving detailed investigation for future work.Source: §3.3 · p. 6

3.4 Implication on Choosing SGD or Adam

  • JS0, the average initial JS distance between block spectra, is optimizer-independent, and the SGD–Adam performance gap grows as JS0 increases.The authors present JS0 as a potential pre-training indicator of whether SGD may underperform Adam.Source: §3.4 · pp. 6–7, Figure 7 · p. 7, Table 1 · p. 7
  • Original SQL computation takes comparable time to one training run, while simplified SQL preserves the reported message and takes under 0.001 seconds for ResNet18.The authors state that the simplified method reduces computation time while still detecting Hessian heterogeneity.Source: §3.4 · p. 7

4 Case Study of Quadratic Models and Preliminary Theory

Quadratic experiments show that block heterogeneity distinguishes cases where GD lags Adam from cases where they perform similarly, despite matched Hessian condition numbers. Initial theory attributes this gap to shared versus block-specific learning rates, while noting important bounds and convergence qualifications.

4.1 Experimental Observations

  • The study models a strongly convex quadratic with positive-definite block-diagonal Hessian and compares four Hessian types, all with condition number 5000.The cases include Transformer-type and CNN-type spectra, plus simplified heterogeneous and homogeneous spectra.Source: §4 · p. 7, §4.1 · pp. 7–8
  • With heterogeneous blocks, GD largely lags Adam, while with homogeneous blocks GD is on par, even though all four Hessians have condition number 5000.Cases 3 and 4 also share all eigenvalues; the reported difference comes from their blockwise spectra.Source: §4.1 · pp. 8–9, Figure 8 · p. 8

4.2 Initial Theoretical Results

  • The full Hessian condition number alone does not distinguish the four cases; the analysis therefore uses finer-grained blockwise spectra.The paper says κ is the same across cases and identifies blockwise sub-matrix spectra as necessary for distinguishing GD and Adam.Source: §4.2 · p. 9
  • For asymptotic analysis, Adam is considered with β2 = 1 because constant-learning-rate Adam with β2 < 1 has a non-convergence issue.The text reports that iterates can reach near-optimal solutions and then bounce back.Source: §4.2 · p. 9
  • Across 1,000 trials in heterogeneous Case 3 with standard Gaussian random initialization, r·max_l κ_l ≤ 1,000 with high probability, about 5× below κ = 5,000, suggesting Adam could be 5× faster than GD.Source: §4.2 · p. 10
  • Tightening Adam’s complexity bound is difficult without extra structure on each H_l, and the authors say it is unclear whether relevant structure holds in Transformers.In the tested Case 3, the diagonal preconditioner increased the reported block condition numbers.Source: §4.2 · p. 10
  • The Adam complexity coefficient is independent of the global condition number κ, whereas the stated tight GD complexity is ˜O(κ).The discussion attributes the potential improvement to the block-diagonal Hessian structure and heterogeneous blockwise spectra.Source: §4.2 · pp. 9–10
  • The theory indicates that single-learning-rate methods like GD can lag coordinate-wise methods like Adam on problems with block heterogeneity.The proposed mechanism is that one learning rate cannot handle variation among blocks; Adam’s preconditioner chooses different rates across block sub-matrices.Source: §4.2 · pp. 9–10

5 Conclusion

The conclusion links Transformer underperformance of SGD to block heterogeneity in the Hessian. It reports numerical checks across several model types and initial theoretical support for that link.

5 Conclusion

  • The work explores why SGD performs worse than Adam when training Transformers.It approaches the question through the Hessian.Source: §5 · p. 10
  • The authors establish block heterogeneity in the Hessian and link it to the performance gap between Adam and SGD.Source: §5 · p. 10
  • The authors numerically verify their claim on Transformers, CNNs, MLPs, and quadratic problems, and provide initial theory.Source: §5 · p. 10

Broader Impacts

The broader-impacts statement says the work may help the community understand large AI model training, while noting a potential risk if AI models are used illegally.

Broader Impacts

  • The authors say the work can help the community better understand large AI model training.Source: Broader Impacts · p. 11
  • The authors identify potential harm if AI models are used for illegal purposes.Source: Broader Impacts · p. 11
  • The paper frames its research question as why SGD performs worse than Adam for Transformer training.Source: Broader Impacts · p. 11

A Related Works

Related work offers several explanations and analyses of optimizer behavior, but the paper emphasizes Hessian block heterogeneity as an additional distinction between CNNs and Transformers. It positions its contribution alongside work on adaptive optimizers, Hessian spectra, and Transformer training difficulties.

On the unsatisfactory performance of SGD on Transformers

  • Prior work attributed SGD’s Transformer gap to heavy-tailed language-task noise, but reports of the gap in full-batch settings [46] and on Vision Transformers [92] suggest architecture may matter.Source: A Related Works · p. 18

Understanding of Adam

  • Prior analyses connect Adam’s advantage to changing learning rates over time [84] [19] [51], while this work identifies handling block heterogeneity over space as another benefit.Source: A Related Works · p. 18
  • Unlike diagonal-Hessian accounts of Adam, this work finds blockwise eigenvalues useful where Hessians are closer to block-diagonal than purely diagonal.The authors note that it is unclear whether diagonal structure commonly holds in real problems.Source: A Related Works · p. 18

Hessian Spectrum Analysis

  • Prior Hessian-spectrum studies focused on MLPs and CNNs, whereas this work examines Transformers as well as optimizer impacts.Source: A Related Works · p. 18

On the difficulties of Transformer training

  • Existing Transformer-training remedies target gradient vanishing or explosion, while this work investigates Hessian eigenvalues and their relation to the SGD–Adam gap.The authors describe this analysis as a first step and leave finer-grained work for the future.Source: A Related Works · p. 19

B More Results and Discussions

Appendix results show that careful learning-rate tuning does not close the SGD–Adam gap on Transformers, while comparisons across CNNs, MLPs, and training stages relate optimizer performance to block heterogeneity. Further examples qualify what the theory implies about Adam under homogeneous versus heterogeneous Hessians, and the appendix describes a more efficient heterogeneity measure and optimizer updates.

B More Results and Discussions

C.1 Preliminaries on Optimizers

  • The appendix specifies SGD momentum and Adam’s first- and second-moment updates, including elementwise operations and default β settings.Source: §C.1 · p. 24

C.2 Preliminaries on the Stochastic Lanczos Quadrature Method

The Stochastic Lanczos Quadrature method (SLQ) estimates Hessian spectra by turning trace estimation into averages of quadratic forms for random vectors. It approximates those forms through Gaussian quadrature, with Lanczos steps providing the matrix needed to construct the quadrature rules.

C.2 Preliminaries on the Stochastic Lanczos Quadrature Method

  • SLQ formulates Hessian spectrum approximation as trace estimation and estimates the trace using averages of quadratic forms for normalized random vectors.The random vectors may be Rademacher or Gaussian; the implementation uses Gaussian vectors.Source: §C.2 · pp. 24–25
  • The method reduces each quadratic form to an integral over a measure tied to Hessian eigenpairs, which are unknown and make direct integration difficult.Gaussian quadrature approximates the integral with nodes and weights, and is exact for polynomial f of degree less than 2m.Source: §C.2 · p. 25
  • Finding Gaussian quadrature rules is difficult because the orthogonal polynomials depend on an unknown measure; obtaining the Jacobi matrix bypasses the need to find them directly.The Jacobi matrix is sufficient to obtain the quadrature rules.Source: §C.2 · p. 26
  • With the measure induced by the Hessian and a unit vector, m Lanczos steps produce the Jacobi matrix for the associated orthogonal polynomials.The matrix’s eigenpairs provide Gaussian quadrature nodes and weights; these approximate each quadratic form, which is then averaged to estimate the trace.Source: §C.2 · p. 26

D.1 Implementation Details on SLQ and Training Configurations

The implementation estimates Hessian spectra with SLQ using repeated Hessian-vector products, while controlling randomness and selecting block coverage according to model size. At the largest scale, full-data Hessian-vector products make a complete SLQ run unaffordable, so the Hessian is approximated using the largest feasible batch.

D.1 Implementation Details on SLQ and Training Configurations

  • SLQ queries the network only for Hessian-vector products; runs use m = 100 and nv = 10, with settings reported to yield estimation error below 10^-14.Randomness in forward and backward passes is removed, and attention blocks are implemented naively because Flash Attention does not support Hessian-vector products.Source: §D.1 · p. 28
  • At GPT2 scale, one Hessian-vector product over 9B tokens takes about 9 GPU days, making the at-least-9,000-day complete SLQ run unaffordable.The study approximates the Hessian using the largest possible batch size under GPU bandwidth and time constraints.Source: §D.1 · p. 28
  • For ResNet18, VGG16, BERT, and GPT2-nano, block spectra are computed for all eligible blocks except BatchNorm or LayerNorm layers, with GPT2-nano sampled at even indices.The respective block exclusions and sampling choices differ by model.Source: §D.1 · p. 28
  • For ViT-base and GPT2, large model size limits blockwise SLQ to selected embedding, output, attention, and MLP blocks.The selected blocks are spaced across model depth; GPT2 spectra estimation takes about 7 days on one A100 GPU.Source: §D.1 · pp. 28–29
  • Training uses default codebase configurations and learning rates grid-searched for SGD and Adam under the same budget.Vision tasks use cosine decay; SFT uses nanoGPT, with GPT2 pretraining on OpenwebText before fine-tuning on an Alpaca Eval subset.Source: §D.1 · p. 29

D.2 Ablation Studies of SQL on a Small Tranformer

On a 75k-parameter GPT model, the SQL method accurately estimates the true Hessian spectrum. Figure 16 compares the true spectrum with the estimated curve.

D.2 Ablation Studies of SQL on a Small Tranformer

  • On a 75k-parameter GPT model, SQL produces an accurate estimate of the true Hessian spectrum.Figure 16 compares the true spectrum with the estimated spectrum curve.Source: §D.2 · p. 29, Figure 16 · p. 29

D.3 Implementation Details on Figure 2

The Figure 2 implementation uses a synthetic binary-classification dataset and a one-hidden-layer Tanh network. The MLP experiments vary layer-output scaling to study different degrees of heterogeneity.

D.3 Implementation Details on Figure 2

  • The Figure 2 setup uses 100 binary-classification data points and a one-hidden-layer Tanh network with input size 64 and width 8.The data are generated around random class centers with noise scaled by 0.5; Adam training for 1000 iterations at learning rate 1 × 10^-4 reaches 100% classification accuracy.Source: §D.3 · p. 29

D.4 Implementation Details on the MLP experiments in Figure 5

  • The Figure 5 MLP experiments scale each layer’s output by c from 1 to 15 and compare best test accuracy after one epoch for SGD and Adam.The model is a four-layer ReLU MLP trained on MNIST; learning rates are grid-searched from 1e-4 to 1e-1.Source: §D.4 · p. 29

E.1 Proof of Proof of Proposition 1

The proof determines the optimal learning rate by considering three ranges of η; in the middle range, it minimizes a convex objective.

E.1 Proof of Proof of Proposition 1

  • When η ≤ 1/L, the optimal choice is η=1/L; when η ≥ 1/µ, it is η=1/µ.In each outer range, both terms in the objective decrease toward or increase away from the stated boundary choice.Source: §E.1 · p. 30
  • The middle-range objective is convex because its second derivative is nonnegative throughout [1/L, 1/µ].The proof uses this convexity to find the optimizer by setting the first derivative to zero.Source: §E.1 · p. 30
  • For 1/L ≤ η ≤ 1/µ, solving g′(η)=0 gives the stated interior optimizer.The proof writes g_t(η)=µ(ηL−1)^t+L(1−ηµ)^t and identifies a solution for all t.Source: §E.1 · p. 30

E.2 Proof of Theorem 1

The appendix proves Theorem 1 by establishing nonzero block quantities under random initialization and bounding convergence through cost-to-go and per-step improvement. It also records experimental reproducibility details and the absence of error bars for costly large-Transformer runs.

E.2 Proof of Theorem 1

  • Random initialization makes each relevant block quantity nonzero with probability one, after the proof shows the zero-gradient-coordinate sets have measure zero.The proof uses the continuous initialization distribution and the fact that each Hessian row has a nonzero entry.Source: §E.2 · p. 30
  • The proof assumes w*=0 without loss of generality by translating the quadratic, and applies the Adam update to the full Hessian and each block.The translation changes coordinates to z=w−w*; the proof then states the corresponding blockwise update relations.Source: §E.2 · pp. 30–31
  • The proof characterizes convergence using cost-to-go and per-step improvement, expresses both through block quantities, and concludes by summing the bounds over blocks.It defines transformed Hessian spectra for Adam and for each block before completing the bound.Source: §E.2 · p. 31

3. Theory Assumptions and Proofs

  • The paper states that theoretical assumptions are in Section 4.2 and the complete proof is in Appendix E.The claims checklist says the abstract and introduction accurately reflect the experimental results.Source: §3 · pp. 33–34, §1 · p. 33

4. Experimental Result Reproducibility

  • The authors report that Appendix D describes experimental details and compute resources, and that supplementary materials include code.These disclosures are presented as supporting reproduction of the experiments.Source: §5 · p. 34, §6 · p. 35, §8 · p. 36

7. Experiment Statistical Significance

  • Error bars are not provided because repeating experiments on large Transformers is too expensive for multiple runs.This limits the reported statistical-significance information for those experiments.Source: §7 · p. 35

9. Code Of Ethics

  • The checklist reports that the work conforms to the NeurIPS Code of Ethics and discusses positive and negative societal impacts.It also states that no special safeguards are needed because the paper poses no such risks.Source: §9 · p. 36, §10 · p. 36, §11 · p. 37

12. Licenses for existing assets

  • The authors report that existing assets are credited, no new assets are released, and the work involves no human-subject research.The checklist marks the asset-license and human-subject items as addressed or not applicable.Source: §12 · p. 37, §13 · p. 37, §14 · p. 38
Loading 2402.16788v4…