Research map / Optimization / Optimizers
Why Transformers Need Adam: A Hessian Perspective
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Explains why SGD underperforms Adam on Transformers via "block heterogeneity" in the blockwise Hessian spectrum, validated on Transformers, CNNs, MLPs and quadratic problems, with a theory that a single learning rate (GD) lags coordinate-wise Adam under heterogeneity.
Problem
Why SGD performs significantly worse than Adam on Transformers while performing on par with Adam on CNNs; the full Hessian spectrum does not reveal the difference.
Approach
Shows that Transformers exhibit block heterogeneity: the Hessian spectra of different parameter blocks (Q, K, V, MLP, embedding) differ substantially, unlike CNNs. SGD uses one learning rate for all blocks and so struggles under heterogeneity, while Adam's coordinate-wise learning rates handle it. Proposes JS^0, the average JS distance among blockwise initial Hessian spectra, as a predictor of SGD
Key results
- SGD performs significantly worse than Adam on Transformers (GPT2, ViT-base, BERT, GPT2-nano) while performing on par on ResNet18 and VGG16.
- JS^0 values: ResNet18 0.10, VGG16 0.09, versus BERT 53.38, GPT2 83.23, ViT-base 286.41; the SGD-Adam gap grows with JS^0.
- On quadratic problems with condition number 5000, GD lags far behind Adam in heterogeneous-block cases (Case 1 and 3) but matches it in homogeneous cases (Case 2 and 4); theory shows Adam complexity O~(r max kappa_l) vs GD O~(kappa).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SGD instead of Adam on Transformer training | Transformer models · Transformer tasks | training performance | Hurt |
| Use SGD instead of Adam on CNN training | CNN models · ImageNet | training performance | No clear effect |
| Increase layer-scaling heterogeneity | 4-layer MLP · MNIST | test accuracy | Hurt |
| Pretrain GPT2 before fine-tuning | GPT2 · Alpaca Eval | converged loss | Helped |
| Initialize ViT-base from the checkpoint at 50% training steps rather than from scratch | ViT-base with SGD · ImageNet | training performance | Helped |
| Use a single learning rate across blocks rather than coordinate-wise learning rates | GD and Adam on quadratic problems · Block-heterogeneous quadratic problems | convergence performance | Hurt |
| Use a single learning rate across blocks rather than coordinate-wise learning rates | GD and Adam on quadratic problems · Block-homogeneous quadratic problems | convergence performance | No clear effect |
| Set Adam β2 below 1 rather than β2 = 1 with a constant learning rate | Adam on quadratic problems · Strongly convex quadratic | convergence | Hurt |
| Reduce SQL num_v from 10 to 1 (ablation) | Simplified SQL · BERT and ResNet18 initialization | JS^0 estimation | No clear effect |
| Reduce Lanczos steps from 100 to 10 (ablation) | Simplified SQL · BERT and ResNet18 initialization | heterogeneity detection | No clear effect |
| Sample 50% of blocks and use batch size 32 instead of batch size 1024 (ablation) | Simplified SQL · BERT and ResNet18 initialization | heterogeneity detection | No clear effect |
Surprising findings
- The full Hessian spectrum looks similar for CNNs and Transformers, so it does not explain the SGD-Adam gap; the blockwise spectra are needed.
- Diagonal (Adam) preconditioning can increase the condition number of individual Hessian blocks, so the theoretical advantage comes from coefficient independence from the global condition number rather than local improvement.
Limitations
- Authors state block heterogeneity is not the only cause of the SGD-Adam gap, only one important cause.
- Theory is initial and restricted to quadratic problems with block-diagonal Hessian and beta2=1 Adam variant; constant-beta2 Adam does not converge in the analysis.
- JS^0 is a predictor from observed correlation across a small set of models rather than a validated general criterion.
Experimental setup
Datasets: ImageNet OpenWebText Cornell Movie-Dialogs Corpus MNIST Alpaca Eval
Benchmarks: ImageNet OpenWebText Cornell Movie-Dialogs Corpus MNIST
Metrics: test accuracy training loss JS distance among blockwise Hessian spectra (JS^0)
Baselines: SGD with momentum AdamW gradient descent
Builds on: Stochastic Lanczos Quadrature Adam AdamW
Closest papers by meaning
- Adam Exploits $\ell_\infty$-geometry of Loss Landscape via Coordinate-wise Adaptivity (ICLR 2025 · similarity 0.87)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.86)
- On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent (ICLR 2025 · similarity 0.86)
- Understanding Adam Requires Better Rotation Dependent Assumptions (NeurIPS 2025 · similarity 0.85)
- The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training (ICML 2025 · similarity 0.85)
- A Comprehensive Framework for Analyzing the Convergence of Adam: Bridging the Gap with SGD (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.