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

arXiv: https://arxiv.org/abs/2402.19449 · HTML: https://mlsanity.com/paper/2402.19449/

## 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 (most cited first, 30 of 50)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates](https://mlsanity.com/paper/1804.10959/) (2018)
- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://mlsanity.com/paper/2106.10270/) (2021)
- [Does Learning Require Memorization? A Short Tale about a Long Tail](https://mlsanity.com/paper/1906.05271/) (2019)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Scaling Laws and Interpretability of Learning from Repeated Data](https://mlsanity.com/paper/2205.10487/) (2022)
- [Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse](https://mlsanity.com/paper/2206.03126/) (2022)
- [Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients](https://mlsanity.com/paper/1705.07774/) (2017)
- [The Quantization Model of Neural Scaling](https://mlsanity.com/paper/2303.13506/) (2023)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)

## Built on by (most cited first, 30 of 80)

- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Muon Outperforms Adam in Tail-End Associative Memory Learning](https://mlsanity.com/paper/2509.26030/) (2025)
- [Deconstructing What Makes a Good Optimizer for Language Models](https://mlsanity.com/paper/2407.07972/) (2024)
- [Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful](https://mlsanity.com/paper/2507.07101/) (2025)
- [In Search of Adam's Secret Sauce](https://mlsanity.com/paper/2505.21829/) (2025)
- [Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?](https://mlsanity.com/paper/2511.00674/) (2025)
- [Does SGD really happen in tiny subspaces?](https://mlsanity.com/paper/2405.16002/) (2024)
- [Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition](https://mlsanity.com/paper/2502.15015/) (2025)
- [The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training](https://mlsanity.com/paper/2502.19002/) (2025)
- [Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Law](https://mlsanity.com/paper/2505.19227/) (2025)
- [Adam with model exponential moving average is effective for nonconvex optimization](https://mlsanity.com/paper/2405.18199/) (2024)
- [Adam Exploits ℓ_∞-geometry of Loss Landscape via Coordinate-wise Adaptivity](https://mlsanity.com/paper/2410.08198/) (2024)
- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [How Muon's Spectral Design Benefits Generalization: A Study on Imbalanced Data](https://mlsanity.com/paper/2510.22980/) (2025)
- [Revisiting Gradient Normalization and Clipping for Nonconvex SGD under Heavy-Tailed Noise: Necessity, Sufficiency, and Acceleration](https://mlsanity.com/paper/2410.16561/) (2024)
- [Is your batch size the problem? Revisiting the Adam-SGD gap in language modeling](https://mlsanity.com/paper/2506.12543/) (2025)
- [What Makes CLIP More Robust to Long-Tailed Pre-Training Data? A Controlled Study for Transferable Insights](https://mlsanity.com/paper/2405.21070/) (2024)
- [Towards Quantifying the Hessian Structure of Neural Networks](https://mlsanity.com/paper/2505.02809/) (2025)
- [The Rich and the Simple: On the Implicit Bias of Adam and SGD](https://mlsanity.com/paper/2505.24022/) (2025)
- [Memory-Efficient LLM Pretraining via Minimalist Optimizer Design](https://mlsanity.com/paper/2506.16659/) (2025)
- [Stable-DiffCoder: Pushing the Frontier of Code Diffusion Large Language Model](https://mlsanity.com/paper/2601.15892/) (2026)
- [Mano: Restriking Manifold Optimization for LLM Training](https://mlsanity.com/paper/2601.23000/) (2026)
- [Understanding Adam Requires Better Rotation Dependent Assumptions](https://mlsanity.com/paper/2410.19964/) (2024)
- [Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillation](https://mlsanity.com/paper/2605.09253/) (2026)
- [Understanding and Minimising Outlier Features in Neural Network Training](https://mlsanity.com/paper/2405.19279/) (2024)
- [Gradient Multi-Normalization for Stateless and Scalable LLM Training](https://mlsanity.com/paper/2502.06742/) (2025)
- [GradPower: Powering Gradients for Faster Language Model Pre-Training](https://mlsanity.com/paper/2505.24275/) (2025)
- [Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory](https://mlsanity.com/paper/2602.09314/) (2026)
- [Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers](https://mlsanity.com/paper/2411.00999/) (2024)
- [Efficient Distributed Optimization under Heavy-Tailed Noise](https://mlsanity.com/paper/2502.04164/) (2025)

## 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.