mlsanityData to 14 Oct 2026

Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models

arXiv 2402.19449 ↗First posted 29 Feb 2024cs.LG cs.CL math.OCcited by 80

Not on the map. Why?

Abstract (arXiv)

Adam has been shown to outperform gradient descent on large language models by a larger margin than on other tasks, but it is unclear why. We show that a key factor in this performance gap is the heavy-tailed class imbalance found in language tasks. When trained with gradient descent, the loss of infrequent words decreases more slowly than the loss of frequent ones. This leads to a slow decrease on the average loss as most samples come from infrequent words. On the other hand, Adam and sign-based methods are less sensitive to this problem. To establish that this behavior is caused by class imbalance, we show empirically that it can be reproduced across architectures and data types, on language transformers, vision CNNs, and linear models. On a linear model with cross-entropy loss, we show that class imbalance leads to imbalanced, correlated gradients and Hessians that have been hypothesized to benefit Adam. We also prove that, in continuous time, gradient descent converges slowly on low-frequency classes while sign descent does not.

Key facts. Cited by 80 papers in the map (data through 14 Oct 2026). Builds on 50 of its 79 references that are in the map. First posted on arXiv 29 Feb 2024.

Builds on

50 of 79 references.

References · 50

  1. 2017Attention Is All You Needcited by 60,361
  2. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  3. 2014Adam: A Method for Stochastic Optimizationcited by 37,240
  4. 2020Language Models are Few-Shot Learnerscited by 32,419
  5. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  6. 2015Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shiftcited by 7,005
  7. 2016Layer Normalizationcited by 5,669
  8. 2020Training data-efficient image transformers & distillation through attentioncited by 4,586
  9. 2016Pointer Sentinel Mixture Modelscited by 3,096
  10. 2015Neural Machine Translation of Rare Words with Subword Unitscited by 2,918
  11. 2022Locating and Editing Factual Associations in GPTcited by 2,691
  12. 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,165
38 more
  1. 2020Prevalence of Neural Collapse during the terminal phase of deep learning trainingcited by 719
  2. 2019Large-Scale Long-Tailed Recognition in an Open Worldcited by 705
  3. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 655
  4. 2018Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidatescited by 639
  5. 2021How to train your ViT? Data, Augmentation, and Regularization in Vision Transformerscited by 569
  6. 2019Does Learning Require Memorization? A Short Tale about a Long Tailcited by 460
  7. 2019Why gradient clipping accelerates training: A theoretical justification for adaptivitycited by 341
  8. 2020An Investigation of Why Overparameterization Exacerbates Spurious Correlationscited by 310
  9. 2018What is the Effect of Importance Weighting in Deep Learning?cited by 251
  10. 2019Limitations of the Empirical Fisher Approximation for Natural Gradient Descentcited by 217
  11. 2020Understanding the Difficulty of Training Transformerscited by 210
  12. 2023Birth of a Transformer: A Memory Viewpointcited by 161
  13. 2022Scaling Laws and Interpretability of Learning from Repeated Datacited by 152
  14. 2022Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapsecited by 141
  15. 2017Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradientscited by 139
  16. 2023The Quantization Model of Neural Scalingcited by 138
  17. 2024Why Transformers Need Adam: A Hessian Perspectivecited by 124
  18. 2020Descending through a Crowded Valley - Benchmarking Deep Learning Optimizerscited by 119
  19. 2023Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Becited by 114
  20. 2018Convergence of Gradient Descent on Separable Datacited by 110
  21. 2022Better plain ViT baselines for ImageNet-1kcited by 98
  22. 2023Linear attention is (maybe) all you need (to understand transformer optimization)cited by 98
  23. 2022Robustness to Unbounded Smoothness of Generalized SignSGDcited by 81
  24. 2022Imbalance Trouble: Revisiting Neural-Collapse Geometrycited by 80
  25. 2023Tokenization and the Noiseless Channelcited by 59
  26. 2023Toward Understanding Why Adam Converges Faster Than SGD for Transformerscited by 53
  27. 2023Deep Transformers without Shortcuts: Modifying Self-attention for Faithful Signal Propagationcited by 45
  28. 2023Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stabilitycited by 42
  29. 2020The Geometry of Sign Gradient Descentcited by 41
  30. 2021Procrustean Training for Imbalanced Deep Learningcited by 24
  31. 2022How Does Adaptive Optimization Impact Local Neural Network Geometry?cited by 24
  32. 2021The Three Stages of Learning Dynamics in High-Dimensional Kernel Methodscited by 22
  33. 2022A Theoretical Analysis of the Learning Dynamics under Class Imbalancecited by 20
  34. 2024Learning Associative Memories with Gradient Descentcited by 16
  35. 2023Outliers with Opposing Signals Have an Outsized Effect on Neural Network Optimizationcited by 13
  36. 2014Zipf's law holds for phrases, not wordscited by 12
  37. 2021Frequency-aware SGD for Efficient Embedding Learning with Provable Benefitscited by 6
  38. 2017AdaBatch: Efficient Gradient Aggregation Rules for Sequential and Parallel Stochastic Gradient Methodscited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 14; 2025: 34; 2026: 32.

80 papers.

Sort:

By area:

  1. 2024Adam-mini: Use Fewer Learning Rates To Gain Morecited by 137
  2. 2025Muon Outperforms Adam in Tail-End Associative Memory Learningcited by 46
  3. 2024Deconstructing What Makes a Good Optimizer for Language Modelscited by 41
  4. 2025Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wastefulcited by 41
  5. 2025In Search of Adam's Secret Saucecited by 40
  6. 2025Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?cited by 37
  7. 2024Does SGD really happen in tiny subspaces?cited by 36
  8. 2025Accelerating Neural Network Training: An Analysis of the AlgoPerf Competitioncited by 32
  9. 2025The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Trainingcited by 30
  10. 2025Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Lawcited by 24
  11. 2024Adam with model exponential moving average is effective for nonconvex optimizationcited by 23
  12. 2024Adam Exploits ℓ_∞-geometry of Loss Landscape via Coordinate-wise Adaptivitycited by 23
38 more
  1. 2024SWAN: SGD with Normalization and Whitening Enables Stateless LLM Trainingcited by 20
  2. 2025How Muon's Spectral Design Benefits Generalization: A Study on Imbalanced Datacited by 19
  3. 2024Revisiting Gradient Normalization and Clipping for Nonconvex SGD under Heavy-Tailed Noise: Necessity, Sufficiency, and Accelerationcited by 18
  4. 2025Is your batch size the problem? Revisiting the Adam-SGD gap in language modelingcited by 15
  5. 2024What Makes CLIP More Robust to Long-Tailed Pre-Training Data? A Controlled Study for Transferable Insightscited by 13
  6. 2025Towards Quantifying the Hessian Structure of Neural Networkscited by 13
  7. 2025The Rich and the Simple: On the Implicit Bias of Adam and SGDcited by 13
  8. 2025Memory-Efficient LLM Pretraining via Minimalist Optimizer Designcited by 13
  9. 2026Stable-DiffCoder: Pushing the Frontier of Code Diffusion Large Language Modelcited by 13
  10. 2026Mano: Restriking Manifold Optimization for LLM Trainingcited by 12
  11. 2024Understanding Adam Requires Better Rotation Dependent Assumptionscited by 11
  12. 2026Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillationcited by 11
  13. 2024Understanding and Minimising Outlier Features in Neural Network Trainingcited by 10
  14. 2025Gradient Multi-Normalization for Stateless and Scalable LLM Trainingcited by 9
  15. 2025GradPower: Powering Gradients for Faster Language Model Pre-Trainingcited by 8
  16. 2026Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectorycited by 8
  17. 2024Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformerscited by 7
  18. 2025Efficient Distributed Optimization under Heavy-Tailed Noisecited by 7
  19. 2025TROLL: Trust Regions improve Reinforcement Learning for Large Language Modelscited by 7
  20. 2025On Provable Benefits of Muon in Federated Learningcited by 7
  21. 2025Optimization Benchmark for Diffusion Models on Dynamical Systemscited by 7
  22. 2026On Surprising Effectiveness of Masking Updates in Adaptive Optimizerscited by 7
  23. 2026Why Muon Outperforms Adam: A Curvature Perspectivecited by 6
  24. 2025Unbiased and Sign Compression in Distributed Learning: Comparing Noise Resilience via SDEscited by 5
  25. 2025Provable Benefit of SignGD: A Minimal Model Under Heavy-Tailed Class Imbalancecited by 5
  26. 2024Adaptive Batch Size Schedules for Distributed Training of Language Models with Data and Model Parallelismcited by 4
  27. 2025On the Interaction of Batch Noise, Adaptivity, and Compression, under (L0,L1)-Smoothness: An SDE Approachcited by 4
  28. 2025Simple Convergence Proof of Adam From a Sign-like Descent Perspectivecited by 4
  29. 2025Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning Rulescited by 4
  30. 2026Towards Understanding Adam Convergence on Highly Degenerate Polynomialscited by 4
  31. 2025When Can You Get Away with Low Memory Adam?cited by 3
  32. 2025Stochastic Approximation with Block Coordinate Optimal Stepsizescited by 3
  33. 2026Muon Learns More Robust and Transferable Features than Adamcited by 3
  34. 2025Implicit Bias of Per-sample Adam on Separable Data: Departure from the Full-batch Regimecited by 2
  35. 2025Federated Stochastic Minimax Optimization under Heavy-Tailed Noisescited by 2
  36. 2026Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotationcited by 2
  37. 2026Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Ratescited by 2
  38. 2026Tight Complexity of Muon under Heavy-Tailed Noisecited by 2

50 most cited of 80 All 80 (JSON)