Research map / Optimization / Implicit bias and generalization
The Rich and the Simple: On the Implicit Bias of Adam and SGD
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.
Shows theoretically and empirically that Adam (without momentum or with momentum) resists the simplicity bias of GD on two-layer ReLU nets, learning richer nonlinear decision boundaries that generalize better under spurious correlations.
Problem
Understanding the implicit bias of Adam for nonlinear neural networks and how it differs from (S)GD, in particular whether Adam avoids the simplicity bias that makes SGD-trained networks rely on spurious simple features.
Approach
Analyzes population gradients of two-layer homogeneous ReLU nets on a Gaussian mixture with a simple and a complex feature: gradient flow converges to a linear predictor, while signGD (Adam with beta=0) and Adam with beta near 1 learn multiple feature directions yielding a piecewise-linear boundary closer to Bayes-optimal.
Key results
- On Dominoes (MNIST-CIFAR, ResNet-18, 95% spurious correlation), Adam reaches 14.17 original worst-group accuracy vs 0.81 for SGD, and 84.66 decoded vs 71.04 for SGD.
- On the Boolean features dataset (three-layer NN), Adam gets 97.87 test accuracy vs 89.58 for SGD, with higher decoded core correlation (0.87 vs 0.51).
- On the Gaussian data with finite samples, Adam's test accuracy is 0.32% higher than GD's (mu=0.3), and 0.595% higher with mu=0.25.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Adam instead of SGD | ResNet-18 · Dominoes (MNIST-CIFAR) balanced test | original worst-group accuracy: 0.81_{\pm 0.38} -> 14.17_{\pm 3.15} | Helped |
| Use Adam instead of SGD | ResNet-18 · Dominoes (MNIST-CIFAR) core-only test | core-only worst-group accuracy: 1.66_{\pm 1.79} -> 20.63_{\pm 5.75} | Helped |
| Use Adam instead of SGD | ResNet-18 · Dominoes (MNIST-CIFAR) balanced test | decoded worst-group accuracy: 71.04_{\pm 0.63} -> 84.66_{\pm 0.18} | Helped |
| Use Adam instead of SGD | Three-layer NN · Boolean features dataset | test accuracy: 89.58_{\pm 1.92} -> 97.87_{\pm 0.69} | Helped |
| Use Adam instead of SGD | Three-layer NN · Boolean features dataset | decoded core correlation: 0.51_{\pm 0.08} -> 0.87_{\pm 0.03} | Helped |
| Use Adam instead of SGD | Three-layer NN · Boolean features dataset | decoded spurious correlation: 0.78_{\pm 0.08} -> 0.36_{\pm 0.06} | Helped |
| Use Adam instead of GD | Two-layer NN on finite-sample Gaussian data · Gaussian dataset | test accuracy: +0.32% | Helped |
| Use Adam instead of GD | Two-layer NN on finite-sample Gaussian data with μ=0.25 · Gaussian dataset | test accuracy: 0.595% | Helped |
| Use Adam instead of GD | Two-layer NN with m=500, d=20, and μ=0.25 · Gaussian dataset | test accuracy: 0.203% | Helped |
Surprising findings
- Adam outperforms SGD on image datasets with spurious correlations, challenging the consensus that SGD generalizes better than Adam on image data.
- Theory shows SGD-trained (gradient flow) two-layer nets learn only a linear predictor, while signGD learns multiple feature directions.
Limitations
- Theory covers gradient flow, signGD (no momentum), and a toy limit for Adam with beta near 1; the finite-sample Adam with momentum case is only empirical.
- Theoretical analysis restricted to d=2, fixed outer layer weights, and correlation loss in the infinite-sample setting.
- Simplicity bias can be beneficial for in-distribution generalization, so richer features are not always better.
Experimental setup
Datasets: Gaussian mixture synthetic data MNIST with colored patch Dominoes (MNIST-CIFAR) Waterbirds CelebA MultiNLI CivilComments Boolean features dataset
Benchmarks: Waterbirds CelebA MultiNLI CivilComments Dominoes (MNIST-CIFAR) Boolean features dataset MNIST with spurious patch
Metrics: test accuracy worst-group accuracy average accuracy decoded core accuracy decoded spurious correlation margin distribution test error on Gaussian mixture
Baselines: SGD GD signGD Linear model
Builds on: Adam signGD Shah et al. simplicity bias (slabs dataset) Gunasekar et al. implicit bias of GD Wang et al. implicit bias of Adam on linear models
Closest papers by meaning
- The Implicit Bias of Adam on Separable Data (NeurIPS 2024 · similarity 0.91)
- A Comprehensive Framework for Analyzing the Convergence of Adam: Bridging the Gap with SGD (ICML 2025 · similarity 0.89)
- Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks (NeurIPS 2025 · similarity 0.89)
- Simplicity Bias in Overparameterized Machine Learning (AAAI 2024 · similarity 0.89)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.89)
- Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wild (CVPR 2025 · similarity 0.88)
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.