The Rich and the Simple: On the Implicit Bias of Adam and SGD

TL;DR

This paper compares the implicit biases of Adam and SGD in training neural networks, showing Adam's tendency to find more complex, optimal solutions leading to better generalization, especially under distribution shifts.

Contribution

It provides a theoretical and empirical analysis demonstrating Adam's resistance to simplicity bias compared to SGD, resulting in richer features and improved generalization.

Findings

Adam produces more complex decision boundaries than SGD.

Adam achieves higher test accuracy under distribution shifts.

Theoretical analysis confirms differences in implicit bias between Adam and SGD.

Abstract

Adam is the de facto optimization algorithm for several deep learning applications, but an understanding of its implicit bias and how it differs from other algorithms, particularly standard first-order methods such as (stochastic) gradient descent (GD), remains limited. In practice, neural networks (NNs) trained with SGD are known to exhibit simplicity bias -- a tendency to find simple solutions. In contrast, we show that Adam is more resistant to such simplicity bias. First, we investigate the differences in the implicit biases of Adam and GD when training two-layer ReLU NNs on a binary classification task with Gaussian data. We find that GD exhibits a simplicity bias, resulting in a linear decision boundary with a suboptimal margin, whereas Adam leads to much richer and more diverse features, producing a nonlinear boundary that is closer to the Bayes' optimal predictor. This richer decision boundary also allows Adam to achieve higher test accuracy both in-distribution and under certain distribution shifts. We theoretically prove these results by analyzing the population gradients. Next, to corroborate our theoretical findings, we present extensive empirical results showing that this property of Adam leads to superior generalization across various datasets with spurious correlations where NNs trained with SGD are known to show simplicity bias and do not generalize well under certain distributional shifts.