Research map / Optimization / Implicit bias and generalization
Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models
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.
Argues that heavy-tailed class imbalance in language data is a key cause of the performance gap between Adam and (S)GD, showing slow GD progress on rare classes across vision CNNs, ResNets, ViTs, linear models and transformers, and proving GD is slow vs sign descent on a simple imbalanced problem.
Problem
Why Adam outperforms gradient descent by a much larger margin on language models than on vision tasks; what property of the problem makes SGD slow and Adam faster.
Approach
Heavy-tailed class imbalance (Zipf-like next-token frequencies) makes GD progress slowly on low-frequency classes, which dominate the average loss; sign-based/per-coordinate normalized updates are insensitive to class frequency. Gradient and Hessian become correlated across class blocks through an assignment mechanism, which explains why diagonal preconditioning helps.</247
Key results
- On a CNN with barcoded heavy-tailed MNIST, GD makes almost no progress on the ~10k low-frequency classes while Adam reduces loss on both groups.
- Theorem: on a simple imbalanced setting, gradient flow loss for class k decays as Θ(1/(π_k t)) while continuous-time sign descent decays as Θ(e^{-ct}), independent of class frequency.
- Upweighting low-frequency classes by 1/sqrt(π_k) lets SGD consistently outperform plain SGD on the tested problems.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train with Adam instead of GD under heavy-tailed class imbalance | CNN · MNIST+Barcoded | training loss on low-frequency classes | Helped |
| Introduce heavy-tailed class imbalance instead of balanced classes | CNN trained on MNIST · MNIST+Barcoded vs. MNIST | training loss | Hurt |
| Train with Adam instead of SGD under heavy-tailed class imbalance | ResNet18 · Heavy Tailed ImageNet | training loss | Helped |
| Introduce heavy-tailed class imbalance instead of uniform class frequencies | ResNet18 trained on ImageNet · Heavy Tailed ImageNet vs. Small ImageNet | training loss | Hurt |
| Train with Adam instead of SGD under heavy-tailed class imbalance | SimpleViT · Heavy Tailed ImageNet | training loss | Helped |
| Introduce heavy-tailed class imbalance instead of balanced classes | Vision Transformer trained on ImageNet · ImageNet | training loss | Hurt |
| Use full-batch deterministic updates rather than stochastic updates | GD vs. Adam on TinyPTB · TinyPTB | training loss on low-frequency classes | No clear effect |
| Use sign descent instead of gradient descent | small transformer last layer · TinyPTB | training loss on low-frequency classes | Helped |
| Upweight low-frequency classes with weights 1/√π_k | SGD · heavy-tailed imbalance tasks | original unweighted training loss | Helped |
| Upweight low-frequency classes with weights 1/π_k | SGD · some heavy-tailed imbalance tasks | original unweighted training loss | No clear effect |
| Use gradient flow instead of continuous-time sign descent | simple imbalanced setting · simple imbalanced setting | class loss convergence rate: Θ(e^{-ct}) vs. Θ(1/{π_k t}) | Hurt |
| Use inputs sampled from approximately orthogonal zero-mean data rather than aligned positive-mean data | linear softmax regression · Random Heavy-Tailed Labels | training loss on low-frequency classes | Helped |
Surprising findings
- Class imbalance alone is not sufficient to make GD slow; orthogonal or zero-mean inputs can remove the gap.
- Gradient-Hessian correlation across class blocks holds only during training and does not hold globally; the opposite pattern exists at other parameters.
Limitations
- Focuses on training loss/optimization rather than generalization; validation results only partially hold.
- Theory is for a simplified orthogonal-input setting in continuous time and biased toward sign descent.
- Reweighting by 1/pi_k can hurt the original objective; does not capture stochasticity interactions fully.
Experimental setup
Datasets: WikiText-103 WikiText-2 PTB TinyPTB MNIST ImageNet Random Heavy-Tailed Labels Barcoded MNIST Heavy Tailed ImageNet
Benchmarks: WikiText-103 training loss PTB training loss Barcoded MNIST training loss Heavy Tailed ImageNet training loss
Metrics: training loss overall training loss per class-frequency group validation loss
Baselines: SGD gradient descent normalized GD sign descent
Builds on: Adam sign descent normalized gradient descent
Closest papers by meaning
- Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf’s Law (NeurIPS 2025 · similarity 0.88)
- Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization (ICML 2024 · similarity 0.86)
- Imbalances in Neurosymbolic Learning: Characterization and Mitigating Strategies (NeurIPS 2025 · similarity 0.86)
- In Search of Adam’s Secret Sauce (NeurIPS 2025 · similarity 0.86)
- Clipping Improves Adam-Norm and AdaGrad-Norm when the Noise Is Heavy-Tailed (ICML 2025 · similarity 0.86)
- Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks (NeurIPS 2025 · similarity 0.86)
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.