The Rich and the Simple: On the Implicit Bias of Adam and SGD
Authors: Bhavya Vasudeva
Jung Whan Lee
Vatsal Sharan
Mahdi Soltanolkotabi
Venue: NeurIPS 2025 poster ()
Keywords: adam
sgd
implicit bias
simplicity bias
Primary area: optimization

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.

The Rich and the Simple: On the Implicit Bias of Adam and SGD
Bhavya Vasudeva
Jung Whan Lee Vatsal Sharan
Mahdi Soltanolkotabi
Department of Computer Science University of Southern California
Bovasudev,jjee7// vsharan, soltanol]eusc. wowaaaannooooaaaollleaaaa.. edu
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 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 simplicity bias, resulting in linear decision boundary with a suboptimal margin, whereas Adam leads to much richer and more diverse features, producing nonlinear boundary that is closer 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.
Introduction
Adaptive optimization algorithms, particularly Adam [1], have become ubiquitous in training deep neural networks (NNs) due their faster convergence rates and better performance, particularly on large language models (LLMs), as compared to (stochastic) gradient descent (SGD) [2]. Despite its widespread use, the theoretical understanding of how Adam works and when/why outperforms (S)GD remains limited.
Modern NNs heavily overparameterized and thus the training landscape has numerous global optima. As a result, different training algorithms may exhibit preferences or biases towards different global optima a.k.a. implicil bias. There is extensive prior work the implicit bias of GD [3, 4, 5, 6), for both linear and nonlinear models (see Section 5 for a detailed discussion of related work). However, there is limited work investigating the implicit bias of Adam. Recently. Zhang al. [7] showed that for linear logistic regression with separable data, Adam iterates directionally converge to the minimum lor norm solution, contrast to GD which converges to minimum fn-nrrm [3]. This difference between the implicit bias of Adam vs GD in simple linear settings motivates the central question of this paper:
What the implicit bias of Adam for nonlinear models such as NNs, and how does differ from the implicit bias of (SJGD?
Given the popularity of Adam, surprisingly little is known about the implicit bias of Adam for training NNs. A notable exception is the recent work Tsilivis et al. [8], which characterized the
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Bayes
5G0
Adam
Accuscy 184)
30
0.8
80-
70-
50
Watarb/ds
Carib4
MutiNn
CheCCommetts NIT- Doctean feats.
Figure I: Illustration of the synthetic dataset considered in this work, and comparison of the Bayes' optimal predictor with the decision boundaries of two-layer NNs trained with Adam and GD.
Figure 2: Training with Adam leads better perfor- mance across different test set metrics on six bench- mark datasets with spurious correlations, as compared to SGD. See Section 4 for details.
late-stage implicit bias of family of steepest descent algorithms, homogeneous NNs, in terms of maximizing an angorithm-dependee geometric margin (e for signGD or Adam without momentum) and convergence to KKT points. However, the stationary points of the respective margin-maximizzan may not be unique. Additionally, this characterization does not relate the implicit bias of the algorithm with properties of the learned solution, such as the type of features learned, or the complexity of the decision boundary, which in turn impact its generalization. In particular, in many applications, NNs trained with SGD are known exhibit simplicity bias, that is, they learn simple solutions [9]. which can lead to suboptimal generalization. For example, SGD-trained two-layer NNs rely on now-dimensional projections of the data to make predictions [10]. This simplicity bias can be particularly detrimental in the presence of spurious features where it is simpler for NNs trained with SGD to utilize them to achieve zero training error." This leads to the question: Does training with Adam lead to solutions that are resistant to this simplicity bias?
In this paper, we answer these questions in two ways. Theoretically, show that two-layer ReLU NNs trained on a Gaussian mixture data setting (see Fig. 1) with SGD exhibit simplicity bias while training with Adam leads richer feature learning. Empirically, we demonstrate that training with Adam can lead to better performance as compared to SGD on various benchmark datasets with spurious features (see Fig. 2). Our main contributions
follows.
We identify a simple yet informative setting with mixture of Gaussians where GD and Adam exhibit different implicit biases (see Fig. 1). The Bayes' optimal predictor in this setting is nonlinear (piecewise linear), and we show both theoretically and empirically- - that while GD exhibits a simplicity bias resulting in a linear predictor, Adam encourages reliance on richer features leading to a nonlinear decision boundary. which is closer to the Bayes' optimal predictor. We theoretically prove this difference in the implicit bias by analyzing the population gradients and updates of GD and Adam without momentum (signGD). We also show that this leads to better test accuracy in distribution well as across some distribution shifts. Additionally, to theoretically understand the behaviour of Adam with momentum, we analyze simpler setting where the variance of the Gaussians approaches 0, in the infinite width limit. We show that the decision boundaries learned with signGD and Adam are more nonlinear (and closer to the Bayes' optimal predictor) than the one learned with GD.
We conduct extensive experiments various datasets and show that Adam leads to richer features that are more robust compared to simpler features learned via SGD, allowing Adam achieve higher test accuracy both in distribution and under certain distribution shifts. First, we consider an MNIST-based task with a colored patch as a spurious feature and show that compared to SGD, training with Adam leads to more nonlinear decision boundary. larger margins overall, and has better generalization on a test set with flipped correlation. Next, we show that Adam achieves better worst-group accuracy on four benchmark datasets (Waterbirds, CelebA, MultiNLI, and CivilComments) for subgroup robustness, and better decoded worst- group accuracy on the Dominoes or MNIST-CIFAR dataset [9], with images from CIFAR and MNIST classes as the complex/core and simple/spurious features, respectively. Finally, we study the Boolean features
"We remark that simplicity bias may not always be detrimental; for instance, it can be beneficial for in di-iistribbuion generalization (see Appendix C for further discussion). Our focus is characterizing and contrasting the implicit bias of Adam vs (S)GD in terms of rich simple feature learning, not advocating for one always be better.

---

dataset proposed in Qiu et al. [11]. and show that training with Adam leads to better core feature learning as compared SGD.
2 Setup
We consider a two-layer bomogeneous neural network with fixed final layer and ReLU activation, defined as f( (W::)) x)) a' a W æ), where æ IR denotes the input, W R denotes the trainable parameters, a {=1}m are the final layer weights, and (--) := max(0,-) -) is the ReLU activation. Let S = (( y,) denote the set of train samples, where the label y é {-1.1}. The model trained to minimize the empirical risk T.())= = 1/2 Σ (W a.)), where denotes a decreasing loss function. We consider two loss functions, namely logistic loss, where {()) == log(1 exp(2)) and correlation or linear loss, where {(z) := 2, for é R. We focus on the following two update rules.
Gradient Descent. The updates for GD with step-size 7 > at iteration 2 are written as W41 W.- n,G,, where G2 := V WWI[W,) each row of which is written as:
L.W (W,;,, Σ
Σ o (x;) =1
where tr denotes (-v ((W, for convenience, and d'(2) := 1[ 0], for = R.
Adam.
The update rule for the Adam optimizer [1] is as follows:
W+1 W. 17,, o (V, +111))0112/
where M. M++
B, (1 (1- (.) is the bias-corrected first-moment estimate,
1- 85-1
1-
55-1
and V.- V--1
1
82 V. (1 (1- G.) is the bias-corrected second (raw) moment
1 1-8
1-
1-
estimate. c is the numerical precision parameter, which is set as for the theoretical results. Also, and (.)% denote the Hadamard product and power, respectively, and Mo and Vo are initialized as zeroes. Note that we can write
Mr STan and V. Er 3J G o GGT ΣT-OBT
Σt 35
At each optimization step, the descent direction is different from the gradient direction because of the entry-wise division with the second (raw) moment. Further, the first update step exactly matches the update of signGD, which uses the sign of the gradient zign(G.) instead of directly using the gradient G, for the update. This is because Mo Go and Vo Go o Go, and hence
(Mo @ v ).,-
[Co),,) sign((G))..). Similarly, when the parameters 81 and B2 are set as the Adam updates the as signGD for every t 0.
(Ga).,,
Dataset. Our synthetic dataset designed to investigate the impact of feature diversity on the implicit biases of GD and Adam in NN training. It models two classes with differing feature distributions to emulate real-world scenarios where feature complexity may vary between classes. See Fig. for illustration of the dataset. Concretely, each sample (x, y) is generated as follows:
y Unif((+1)), E~ Unif([+1])
(I)
31 ~ N مفرح
1IT 2 22 ~ (E o2)) Ty a 03), Vj (3....,dd)
The first two dimensions contain information about the label while the rest are noisy. Our dataset construction is inspired by the synthetic "slabs" dataset introduced by Shah et al. [9]. While their approach utilizes slab features to represent non- non-linearly separable components, we consider Gaussian features instead. This modification enhances the realism of the synthetic data and facilitates more nuanced analysis of the NN training dynamics.
We first write the Bayes' optimal predictor for this dataset when using only the signal dimensions as follows.
Proposition (Bayes' Optimal Predictor). The optimal predictor for the data in Eq. (1) with d= 2 is:
(μ + p3)a) + 2 U222 = xi-x2 +
-² log (0.5 + expp
2u2z2
20
3

---

Initialization
Epoch 500
Epoch 1000
Epoch 10300
Epoch 20000
Epoch 500
Epoch 1000
Epoch 10000
Epoch 20000
24
Figure 3: Evolution of the decision boundary (top row) and the neurons (bottom row) over time, for GD (left) and Adam (right) with learning rates 0.1 and 10 over 20 200 epochs of training width 100 NN with small initialization (the neurons are colored based the quadrant they were initialized in) using population gradients (the samples are plotted for illustration purposes) on the Gaussian data setting (Eq. (1)) with μ = 0.3, ~ 2. 0.1. GD leads a linear decision boundary, with neurons mostly aligned with the directions [ OJT, while Adam (with 31 = 32 = 0.9999) leads to a non-linear decision boundary, with neurons aligned with three main directions []', [1, 1]" [1, -1]", which
is closer to the Bayes' optimal predictor.
Since consider homogeneous NNs (no bias parameter), make the following assumption the data generating process to make the setting realizable, ie.,, ensure that the Bayes' optimal predictor passes through the origin.
Assumption 1 (Realizabilitty) Let J 12
pi t ux
2 1.
For realizability, 141
(NW- - 4 and μ3 2/2 (nw+ + 1
Here, x denotes the degree anisotropy of the clusters, and corresponds to the slopes of the two linear components of the optimal predictor.
3 Theoretical Results
In this section, we aim to theoretically analyze GD and Adam and the differences in the learned solution arising from the update rules. For clarity of exposition we focus on simple setting. Specifically, in this section, we consider the infinite sample limit, fixed outer layer weights, d 2, and truining with correlation or linear loss. As we will see later Section 4, these algorithms learn different solutions even when these assumptions are relaxed.
3.1 Gaussian Data
We first obtain closed-form of the population gradient for the Gaussian cluster data as follows. The proof is included in Appendix A.
Proposition (Population Gradient). Consider the data Eq. (1) with og = When using correlation loss, the population gradient for neuron w is written as:
V.. V.. = -aE w x 2 0]yz]
= - (((A)T)))) + a A) 24 (AHT, + (o(AU),w) s(AūIw)) 20 (0(Ap1)12) ()) )w))
where A = # w p4 ==
1-- 2n O, O,..., 00' i = [. -1 - 0, .j......' pio =
[W++1
w
[1,,0,0,.... 0J" and Φ and &o denote the normal PDF and CDF, respectively.
Here, fa., fa.. denote the (normalized) mean vectors of the clusters with label 1 and ján denotes the mean of the cluster with label -1.
Before theoretically analyzing the training dynamics of GD and Adam, compare them empirically. Fig. 3 shows the evolution of the decision boundary and the neurons for two-layer NNs trained with GD and Adam using the population gradients in Proposition 2, a function of the training epochs. We observe that for GD (left), all neurons align in direction [=1,0] O and the learned decision boundary is linear, ie., the model only relies the first dimension to make predictions. On the other hand, for Adam (right), neurons converge along three different directions, O] 7
1 [1, [1, -1]T, and
the learned decision boundary is piece- wise linear, i.e., the model uses both the signal dimensions.

---

We will now prove these results for the two optimizers in Theorem 1 and Theorem 2. First, we leverage the gradient expression in Proposition 2 to show that gradient descent with infinitesimal step size, ie., gradient flow (GF) exhibits simplicity bias and learns a linear predictor.
Theorem 1. (Informal) Consider the data in Eq. (1), neurons initialized such that ag =1 with probability 0.5, d 2, 02 By J, and w > C1 and > Cz. where C1,02 are constants. Let = := lim Wh.t and W1o00 == w0.00
for k [m]] Then, the solution learned by gradient flow 1p..
is: 00 0& [T
The proof is included in Appendix A. The main step in the proof is using the relation cos 0x. w. w. where w" := [1, 0J', and showing that
doos On.e aj
whe (1 (1-200 Wk, 1 ) " C (sin 8x .)² ag k. 20 ak [wx,,l
dt
where C > O is a constant, for every O. To show this, we use the fact that uōk, = -V we. I(W,) for gradient flow, and use the population gradient from Proposition 2. Further, using Proposition 2 and showing that W S at + 1) for some constant <> O, we prove by contradiction that sin 0x,, - as + oo, and hence cos 0x.t 1.
Next, we analyze Adam with 31 = f32 = 0 (signGD), and show that it learns both features resulting in nonlinear predictor.
Theorem 2. (Informal) Consider the data in Eq. (1), neurons initialized such that a =1 with probability 0.5, small initialization scale, d = 2, Oz = σy a. w> and e1 $ L S 02, where c, 0, C2 constants. Let Wk,00 == lim What and Wk [w for k E [m]] Let Gu denote the
direction of wg.0- Then, the solution learned by signGD is:
[1,1J
0g > O, sin 00 > 0,
V2 [1 [1.0]T [-
ag > 0, sin 0,0 < 0,
0k
ag > 0, sin 000 0,
a < O.
The proof is included in Appendix A. At high level, we leverage Proposition show that each t > 0, the gradient update in each of the two dimensions is nonzero and signGD updates are in the direction (sign(()), sign(a, sin 0x Notably, when ag > 0, neurons are either in 1 [1, or
v/
in J [11 -13 direction at each iteration. However, when ag 0, the neurons point [11,
and /2 =1]" in alternating iterations, leading convergence in the [-1,0]* direction.
These results characterize the direction in which each neuron converges asymptotically. For GF, all neurons are in the same direction, with exactly half the neurons in [1,0]" and [-1, 0j" directions, which leads
linear predictor. In contrast, for signGD, there is fraction of neurons aligned in the directions J [1, 1J" and / [1, -1] which leads to a piece-wise linear decision boundary,
These results explain the behaviour considers GF, the continuous time version of GD, it is still predictive of the behaviour we observe for diserete-timee GD with small step-size. We also observe similar behaviour for Adam vs GD in Fig. 1, where we consider finite samples (see Section 4 for details).
observed in Fig. 3. We note that while our theoretical result
Analyzing this setting allows us conceptually understand how Adam (without momentum) operates and leads to rich feature learning, while GD exhibits simplicity bias. Importantly, we make no assumptions regarding the initialization direction of the neural network parameters, ensuring that any differences observed between Adam and GD arise solely from the inherent characteristics of the optimization algorithms themselves. Since analyze the population setting, these results are in the under- palddrr-pramttrrizedd regime and don't 't require any lower bound on the network width m to ensure overparameterizatton..
Next, we show that under some conditions on the distribution parameters, the piece- wise linear predictor learned by Adam (without momentum) obtains a strictly lower test error than the linear predictor learned by GD. For simplicity, we assume that m + oo for the following result, so that we can write the predictor learned by Adam in a piece-wise linear form which is symmetric (with respect to the first dimension), using p(sin 0x.0 0) = p(sin 0g.0 < 0) 0.5. However, even with finite m, these probabilities concentrate well and we can expect the following still hold.
5

---

w.p. Theorem 3. (Informal) Consider the data Eq. () with d and w @(1), -.. [ 1].
P 2 ee/EW, where is constant. Consider two predictors,
By
Linear: - = sign(x1):
Piece-wise Linear: i/ = sign (31 + x2)
if 20,
sign x2)
if <
Then, it holds that E(O' + v) E(U + y) O.
The proof is included in Appendix A. Note that we consider isotropic distributions (k 1) for Theorems I and 2, whereas the above result on the test error applies is : 1]. This shows that training with Adam can provably lead to better test accuracy both in-distribution and across certain distribution shifts.
In the next section, we consider a simplified setting to investigate the effect of setting B1, 1 for Adam, which closer to the setting used in practice, where $1 0.9, 32 0.999.
3.2 Toy Data Setting
We consider a simple yet informative setting where = and 0g ay 0, which we refer to as the toy data setting. Specifically, the samples generated as follows:
~ Unif((= c Unif((=1)). T1 = 2/2 (yw-1)
T2 6
2 n.
This setting allows us to characterize the full trajectory of each neuron for the three algorithms. namely GD, signGD, and Adam (31 $2 2s 1). We now state our main result.
Theorem 4. (Informal) Consider the toy data in Eq. (2), neurons initialized small scale, and C1 < C2. where C1, C2 are constants, Let Wk lim
We. weit and W1 [w...ol .fr k [m]
and p:= lan 21
Then, for m + 00, the solutions learned by GD, signGD, and Adam are shown Table 1, where is a constant e [0.72, 1], the probabilities are over the neurons, and the sign of the first element of Wko00 is the same as signn( a.).
GD
Adam (3 = B2 = 0) or signGD
Adam (8, = Bg 1)
[1,0]'
w.p. p
[1.0]T
w.p.
[1,0]
w.p. P
[-,,]]
w.p.
[]
w.p.
[-1,0]
w.p. 1
v2 [1,1]]
w.p. $
W1,00
J 1, 2w 1,
w.p. v2 1J
w.p.
v2
[1,-1] w.p. $
2
[1,-1]
w.p.
[s, w.p.
[s,-1]] w.p.
Table Solutions learned by different algorithms on the toy dataset (see Theorem 4).
Note that < 0.5, which implies (from Table 1) that the fraction of neurons in the direction [1 larger for GD compared to Adam. Consequently, the decision boundary learned by Adam is more non-linear. The proof mainly relies analyzing the updates of each algorithm, so we defer to Appendix A. Fig. shows the decision boundaries learned by the three algorithms, mentioned in Table 1. The predictor learned by Adam is more non- linear and closer to the Bayes' optimal predictor.
-3.
We note that the difference between the predictors learned by GD and Adam is more significant in the Gaussian setting, as compared to the toy dataset. The main reason is that in the toy dataset, there is a larger region where the gradients are the [1, 0JT direction, which makes the decision boundary for signGD more linear, well larger region where the neurons are only active for one of μ4 μ_, which makes the decision boundary for GD more non-linear.
Figure 4: Comparison of the Bayes' opti- mal predictor and the predictors learned by two-layer NNs trained with GD, Adam (3) = 32 s 0) or signGD, and Adam with 31 8g 1 on the toy dataset (Gaussian dataset with o - 0).

---

4
Experimental Results
In this section, we present experimental results across synthetic and woal-world datasets showing that GD exhibits simplicity bias while Adam promotes rich feature learning.
4.1 Gaussian Dataset
We consider the Gaussian data in Eq. (1) in this section, focusing on the finite sample setting, which is closer to practice as we train with the binary cross-entropy loss and consider Adam with momentum parameters 31 = 32 = 0.9999, although the results generalize to other values as well. We consider a small initialization scale and fix the outer layer weights. Specifically, wy -N Vi for & é [m], where a is a small constant.
and ax Vm
Fig. compares the decision boundaries learned by Adam and GD in the finite sample setting, with the Bayes' optimal predictor (for the population version of this setting). We set n = 5000, m = 1000, = 0.3, w = 2, Jz = 0.2, By = 0.15, a 0.001, and use learning rates 0.1 and 10 for GD and Adam, respectively. These results are similar to the population setting and show that the difference in the implicit bias of Adam and GD is quite robust to the choice of the training setting, For comparable train loss, the test accuracy of Adam is 0.32% more than that of GD in this case. We also find that reducing μ increases the accuracy gap: repeating the same experiment with μ 0.25 leads to a gap of 0.595%. These results also generalize to settings where d 2: with m = 500, d 20 and μ = 0.25, the gap is 0.203%. See Appendix B for additional results.
4.2 MNIST Dataset with Spurious Feature
In this section, we conduct an experiment to provide additional support for our theoretical results. We construct binary classification task using 14 14 MNIST images, where inject a 2 2 colored patch at the top left corner of the image (to model the simple feature). One class comprises digit '0' images with red patch, while the other has digits '1' and '2' with green patch. We train two-layer NN using SGD or Adam. For test samples, we flip the patch color to check the model's reliance the simple feature. We also train a linear model on this task, and measure agreement between its predictions on the test set with those of the NN trained with Adam/SGD. This serves as the measure of the complexity of the learned decision boundary. The results are shown in Table 2. We observe that the NN trained with Adam relies more on the digit features and generalizes better. It also has complex/nonlinear decision boundary, as the agreement with the linear predictor is lower as compared to SGD. In addition, we plot the distribution of the train set margin for the two NNs as shown in Fig. 5. We observe that the margin for Adam is generally larger compared SGD. These results support our theoretical findings that SGD learns simpler features compared to Adam.
Test Accuracy
Agreement w/ Lineur Model
COFT
SGD
66.6
95.5
Adam
86
84
Table 2: Comparison of test accuracy and agreement with linear model for a two-layer NN trained on MNIST with spurious correlation.
40
12A 150
26
K0
Figure 5: Distribution of margins of training-set samples from the MNIST dataset with spurious correlation for a two-layer NN trained using SGD (left) and Adam (right).
4.3 Dominoes Dataset
The Dominoes dataset [9, 12]. specifically MNIST-CIFAR, is a binary image classification task. It contains images where the top half shows an MNIST digit [13] from classes (0, 1). while the bottom half shows 1I [14] images from classes (automobile, truck). The MNIST portion corresponds to simple or spurious features, which are 95% correlated with the label, while the CIFAR portion is the complex or core feature and is fully correlated. See Fig. for example images from the dataset and Appendix B for further details.
Figure 6: Example images from class -1 and from the MNIST-CIFAR dataset.
We define groups based on the labels predicted by the core or the spurious feature. The minority groups correspond to images where the core and spurious features
7

---

disagree, and the majority groups correspond to images where they agree. The groups in the test set are balanced. The model can rely on the spurious feature to attain good train performance but can only generalize well on the balanced test set if it learns to use the core feature.
In Table 3, report the average original worst-group accuracy on the balanced test set, for ResNet- 18 model (also see Appendix B, where we include results with ResNet-34 model and 99% spurious correlation). In addition, we report the average core-only worst-group accuracy on a test set where the spurious top half of the image is removed and replaced by black image. This measures how much of the core features have been learned by the model. Lastly, we also report the average decoded worst-group accuracy, obtained by retraining the last layer of the model using logistic regression to fit a group balanced validation dataset, and then evaluating on the original test set. This gives better evaluation for how much of the core features have been learned in the latent representation. We find that training with Adam leads to a significant gain across all three metrics as compared to training with SGD.
We also remark that these results (as well as those in Fig. 2) challenge the widely held consensus that SGD generally performs/generafinn
better than Adam on image data [15]. We use ResNet-based models across the three image datasets and Adam leads to better worst-group accuracy across these cases. These results show that while SGD could be better for in-distribution generalization, Adam can be better for generalization under distribution shifts because promotes richer feature learning.
Optimizer
Original Ace.
Core-Only Acc. Decoded Acc.
Optimizer
Test Acc.
Deculed Core Corr.
Decoded Spurious Com:
SGD 0.81
1.66.1779
T1.0Loas 84.66_0.15
SGD
89.58...z
0.51.o.t
0,78...
Adam
14.17_315
20.63
Adam
97.87.0.0 00.00
0.87.0.30
0.36 Do
Table 3: Training with Adam leads to bet-
Table 4: Training with Adam leads to better test
ter worst-group accuracy on the original and core-only test sets, and after decoding, on the Dominoes (MNIST-CIFAR) dataset (95% spu- rious correlation), as compared to SGD, for ResNet- 18 model.
accuracy, decoded core, and decoded spurious corre- lations on the Boolean features dataset, as compared to SGD, for three-layer NN.
4.4 Subgroup Robustness Datasets
We consider four benchmark subgroup robustness datasets, namely Waterbirds [16], CelebA [17], MultiNLI [18] and CivilComments [19, 20]. Each of these contains core feature that is fully correlated with the label, and spurious feature that is simpler but has a lower correlation. Waterbirds contains images of different birds on various backgrounds. The task is to classify whether the bird is a landbird or a waterbird. The background of the image is either land or water, and is spuriously correlated with the target label. CelebA consists of images of celebrity faces. The task whether the hair color is blonde or not blonde, and the gender of the celebrity male or female is the spurious feature. MultiNLI contains sentence pairs and the task relates to the first, out of three classes: entailment, neutral, and contradiction. The spurious features are negation words which are often but not always associated with contradiction. The CivilComments dataset consists of online comments and the task is to predict whether the comment is toxic or non-toxic. Toxicity is spuriously correlated with the mention of various demographic attributes in the comments, based on gender, race, or religion.
classify
predict how the second sentence
Prior work has shown that simplicity bias can be detrimental worst-group test performance in the presence of spurious features [21]. The standard practice is to use SGD for image datasets and Adam(W) for language datasets. However, since Adam promotes richer feature learning, it should be more robust to spurious correlations across all datasets. Hence, we compare the performance of SGD and Adam on these datasets, when fine-tuning pretrained BERT bert basee uuuuuucceeee sed model [22] on the language datasets, and an ImageNet-pretrained pretrained ResNet-50 [23] model on the image datasets. To ensure fair comparison, we sweep optimizer-sepeinn hyperparameters, namely learning rate, momentum, and weight decay (see Appendix B for details). Fig. 2 and Table 13 in Appendix B show the worst-group and average (group-balanced) accuracies on these datasets when training with SGD or Adam, based on the best worst-group validation accuracy. We see that training with Adam leads to significantly better worst-group accuracy as well as (slightly) better average accuracy compared to training with SGD.
8

---

4.5 Boolean Features Dataset
In this section, we consider the synthetic Boolean features dataset proposed by Qiu et al. [11] to study feature learning under spurious correlation. The dataset is designed to model the presence of two types of features: a set of complex core features with dimension d, that are fully correlated with the label, and set of simple spurious features with dimension de that have correlation strength A e [o, 1] with the label. The rest of the du = d - de dg features are uncorrelated with the label. For our experiments, the core and spurious features are modeled as staircase functions with degrees dg and de < de, respectively. Degree d threshold staircase functions for a Boolean input æ (-1,+1)* defined as:
[1 if x1 111++11122 + 21 ..d. 20
staircase (x) == (-1 otherwise.
We train a three-layer NN using SGD and Adam. We consider d = 50, d. 8, = 1, = 0.9. See Appendix B for formal description of the dataset and further details of the experimental setting. To measure feature learning, we used the decoded core and spurious correlations [11. 24], which measure the extent to which the model has effectively learned the core and spurious features. The decoded core correlation is measured by retraining the last layer of the model using logistic regression to fit the core function and evaluating its correlations with fc on the uniform distribution, E [fc(2)) siggnssnn()) where f is the model. The decoded spurious correlation is measured similarly by retraining on the spurious function and measuring the correlation with fs on the test set.
We report the results in Table 4 evaluated at the lowest comparable training loss achieved by both optimizers. See Appendix B for the training curves and additional results. Adam records significantly higher average test accuracy and decoded core correlation, as well as lower decoded spurious correlation. This suggests that Adam's superior performance on the test set can be attributed to richer feature learning as encourages the utilization of the core features and forgetting or down- weeghtiggg the spurious features. In contrast, SGD relies heavily on the simple spurious feature.
Related Work
In this section, we discuss related work the implicit bias of GD, simplicity bias of NNs trained with GD, implicit bias of Adam and adaptive algorithms, and comparison of Adam and (S)GD in various settings.
Implicit Bias of GD. Since the pionecring studies that identified the implicit bias of linear classifiers on separable datasets [25], extensive research has been conducted on the implicit bias of gradient- based methods for linear models, NNs, and even self-attention models. Wang et al. [26] shows that GD with momentum exhibits the same implicit bias for linear models trained on separable data as vanilla GD. Nacson et al. [27]. Ji and Telgarsky [28], Ji et al. [29] demonstrate fast convergence (in direction) of GD-based approaches with adaptive step-sizes the &a max-margin predictor. It has also been shown that multilayer perceptrons (MLPs) trained with exponentially tailed loss functions on classification tasks, GD gradient flow converge in direction the KKT points of the max- margin problem in both finite [30, 31] and indtnite-width [32] networks. Additionally, Phuong and Lampert [33], Frei et al. [34], Kou al. [35] analyze the implicit bias of ReLU and Leaky- R networks trained with GD on orthogonal data, while Mulayoff et al. [36] investigate convergence to stable minima. Other studies focus on the implicit bias to minimize rank regression tasks using squared loss [37, 38, 39]. The recent survey Vardi [40] includes comprehensive review of related work. More recently, Tarzanagh et al. [41, 42), Vasudeva et al. [43] studied single-head prompt and self-attention models with fixed linear decoder and characterized the implicit bias of attention weights trained with GD to converge to the solution of hard-margin SVM problem.
Simplicity Bias of NNs Trained with GD. Kalimeris et al. [44] empirically demonstrate that NNs trained with SGD first learn to make predictions that are highly correlated with those of the best possible linear predictor for the task, and only later start to use more complex features to achieve further performance improvement. Shah et al. [9] created synthetic datasets and show that in the presence of 'simple" and 'complex' features (linearly separable vs non-lineare non-linearly separable), (two- layer) NNs trained with SGD rely heavily on 'simple' features even when they have equal or even slightly worse predictive power than the 'complex' features. They also show that using SGD leads to learning small-margin and ieature-imooverissee classifiers, instead of large-margin and feature-dense classifiers, even on convergence, which contrasts with Kalimeris et al. [44].

---

Implicit Bias of Adam and Other Adaptive Algorithms. Wang et al. [45] show that homogeneous NNs trained with RMSprop or signGD converge to KKT point of the & max-margin problem. similar to GD, while AdaGrad has a different implicit bias. Zhang et al. [7] show that linear models trained on separable data with Adam converge to the lo max-margin solution. Recently Fan et al. [46] characterized the implicit bias of steepest descent algorithms for multiclass linearly separable data. Xie and Li [47] analyze loss minimization with Adam W and show that under some conditions, it converges to KKT point of the lo -norm constrained loss minimization.
Adam vs (S)GD. Zhou et al. [48] show that SGD converges to flatter minima while Adam converges to sharper minima. Andriushchenko et al. [49] show that flatter minima can correlate with better in-distribution generalization but may not be predictive of even be negatively correlated with generalization under distribution shifts. Zou et al. [50] study an image-inspired dataset and show that CNNs trained with GD can generalize better than Adam. Ma al. [51] show that adding noise to lower or higher frequency components of the data can lead to lower or higher robustness of Adam compared to GD. Kunstner al. [52] show that the reason why Adam outperforms SGD on language data is because the performance of SGD deteriorates under heavy-tailed class imbalance, i.e., when minority classes constitute significant part of the data, whereas Adam less sensitive and performs better. In contrast to their focus on multiple classes and training performance, our work focuses on generalization in a binary classification setting. Several works [53, 54, 55] also study why Adam outperforms SGD on attention models or Transformers.
6 Conclusion
In this work, we investigate the implicit bias of Adam and contrast it with (S)GD. NNs trained with SGD exhibit simplicity bias, whereas we find that training with Adam leads to richer feature learning, making the model more robust to spurious features and certain distribution shifts. We note that richer feature learning may not always be desirable; for instance, simplicity bias can be beneficial for better in-distribution generalization. However, it's important to characterize and contrast the implicit bias of Adam vs GD)GD in this context. To get a principled understanding, we identify a synthetic data setting with Gaussian clusters and theoretically show that two-layer ReLU NNs trained with GD Adam on this task learn different solutions. GD exhibits simplicity bias and learns a linear predictor with a suboptimal margin, while Adam leads to richer feature learning and learns a nonlinear predictor that is closer to the Bayes' optimal predictor. Through theoretical and empirical results, our work adds to the conceptual understanding of how Adam works and poses important directions for future work, such as studying the implicit bias of Adam for other architectures, and the effect of weight decay on simple vs rich feature learning study the implicit bias of AdamW.
Acknowledgments
BV thanks Puneesh Deora for extensive discussions and feedback on the manuscript, and Surbhi Goel for helpful discussions. The authors acknowledge use of USC CARC's Discovery cluster, and thank the anonymous reviewers for their helpful comments. VS was supported by NSF CAREER Award CCF-2239265 and an Amazon Research Award. MS was supported by the Packard Fellowship in Science and Engineering, Sloan Research Fellowship in Mathematics, an NSF CAREER Award #1846369, NSF-CIF awards # 1813877 and #2008443, NSF SLES award #2417075, and NIH DP2LM014564-00 This work was done in part while BV, VS and MS were visiting the Simons Institute for the Theory of Computing at UC Berkeley.
References
[1] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), fass/1412.690,, 2015. 1,3
[2] J. Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Surinder Kumar, and Suvrit Sra. Why adam beats sgd for attention models. ArXiv, abs/1912.03199 2019.
URL https: //api. semanticcchoolll org/ CorppsID ( 208858389. 1
[3] Daniel Soudry, Elad Hoffer, and Nathan Srebro. The implicit bias of gradient descent on
separable data. Journal of Machine Learning Research, 19, 10 2017. 1
10

---

[4] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 1832-1841, PMLR, 10-15 Jul 2018. UtI.https://proceddii mlr.preeee wwwpsssss////eeee /ttp///gaaaaaraaeeeeeeeea 18a. html.
[5] Jingfeng Wu, Difan Zou, Vladimir Braverman, and Quanquan Gu. Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate. In International Conference on Learning Representations, 2021. URL https: httppp:////ppereeeeee //tp///eep//eeeee/////ee/e wwtrrriwwwweeeeeeeeeeeee net/forum? id 3X64RLg2Y60.
[6] Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-econd Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1772-1798. URL btttps://proceet
PMLR, 25-28 Jun 2019.
ngs mlr wwwpppressssss ji19a. html.
[7] Chenyang Zhang, Difan Zou, and Yuan Cao. The implicit bias of adam on separable data, 2024.
URL
arxiv. org/abs ha//ab//200 /a0es/246..00 10650. 1, 10
[8] Nikolaos Tsilivis, Gal Vardi, and Julia Kempe. Flavors of margin: Implicit bias of steepest de- scent in homogeneous neural networks, 2024. URL https: org/ org/abs//r abs/2410. 22069.
1,
[9] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 9573-9585. Curran Associates, Inc., 2020. 2, 3, 7, 9, 26, 30
[10] Depen Morwani, jatin batra, Prateek Jain, and Praneeth Netrapalli. Simplicity bias in 1-hidden layer neural networks. In Shirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: / http://pppeeviieeeeeeee
net
2
[11] GuanWen Qiu, Da Kuang, and Surbhi Goel. Complexity matters: Feature learning in the presence of spurious correlations. In Forty-fors International Conference on Machine Learning,
2024. URL https: /ttp////peeeeeieeeeeeeeee wwt/ffouuue irrrruuiidddttuudgBiii 3, 9, 26
[12] Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy. Agree
to disagree: Diversity through disagreement for better transferability, 2022. URL https: / //arxiv. org/ ab2/ 04414. 7, 26
[13] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):227-2222 1998. doi: 10. 109/5.76679. 7, 26
[14] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. 7, 26
[15] Ashia C. Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods machine learning. Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17. page 4151-4161, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 97815 10860964
[16] Shiori Sagawa*, Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020. URL
https://openreviie www.tee//oouue werpummid=dyyyoooo ditrrumnddddyy 8, 26
[17] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) December 2015,
8
[18] Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pages 1112-1122. Association for Computational Linguistics, 2018 URL pettp://acliee
httg//anhoooog//le--oo 11011.11111111111111111111111111 8
[19] Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion
Proceedings of The 2019 World Wide Web Conference, 2019. 8
11

---

[20] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, al. Wilds: A benchmark of in-the- wild distribution shifts. In International Conference on Machine Learning, pages 5637-5664. PMLR, 2021.
[21] Bhavya Vasudeva, Kameron Shahabi, and Vatsal Sharan. Mitigating simplicity bias in deep learn- ing for improved OOD generalization and robustness. Transactions on Machine Learning Re- search, 2024. ISSN 2835-8856. URL https: httt:p:://ppereevvie
8, 30
[22] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10. N0.18653/v1/N19/ 1423.
URL https: / aacccaaa.aaa.....c lallaannannloggoooooooooc wwweooooggoooooooooooooooo orgg N19 1423. 8, 27
[23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016. doi: 10.1109/CV.
8, 26
[24] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations, 2023. URL https://arxiv org/abst abs/ 02937. 9, 26
[25] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-288
2018.
[26] Bohan Wang, Qi Meng. Huishuai Zhang, Ruoyu Sun, Wei Chen, Zhi-Ming Ma, and Tie-Yan Liu. Does momentum change the implicit regularization on separable data? In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22 Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 978171387 1088. 9
[27] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019.
[28] Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via primal-dual analysis. In Algorithmic Learning Theory, pages 772-804.70 PMLR, 2021. 9
[29] Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration.
In International Conference on Machine Learning, pages 4860-4869, PMLR, 2021.
[30] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Ad-
vances in Neural Information Processing Systems, 33:17176-17 186, 2020. 9
[31] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020. 9
[32] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference Learning Theory, pages 1305-1338.
PMLR, 2020. 9
[33] Mary Phuong and Christoph H Lampert. separable data. In International Conference on Learning Representaaionss, 2021. URL https: //openreview
The inductive bias of ReLU networks on orthogonally
d=krz770xU9Z..
[34] Spencer Frei, Gal Vardi, Peter L Bartlett, Nathan Srebro, and Wei Hu, Implicit bias in leaky relu networks trained on high-dimensional data. arXiv preprint 2000022200.07082, 2022. 9
[35] Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit bias of gradient descent for two-layer relu and leaky relu networks on nearly-orthogonal data, 2023. 9
[36] Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: a view from function space. In Proceedings of the 35th International Conference on Neural Information Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845393. 9
12

---

[37] Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference on Learning Theory, pages 4224 4258. PMLR, 2021. 9
[38] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems, volume 32, 2019. 9
[39] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, 2021.
[40] Gal Vardi. On the implicit bias in deep-learning algorithms, 2022.
[41] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023, 9
[42] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transform-
ers as support vector machines, 2023. 9
[43] Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast conver- gence rates for self-attention, 2024. URL
wwtp//////aa.///////////// 05738. 9
[44] Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd neural networks learns functions of increasing complexity. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
9
[45] Bohan Wang. Qi Meng, Wei Chen, and Tie- Yan Liu. The implicit bias for adaptive optimiza- tion algorithms on homogeneous neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pages 10849-10858. PMLR, 18-24 Jul 2021. URL https: https://proee /ttps//////oceeiggsseeee mlr ppwss/v139/ Vatss//v33//waag211 html. 10
[46] Chen Fan, Mark Schmidt, and Christos Thrampoulidis. Implicit bias of normalized steepest descent in multiclass classification: Sign descent, spectral descent, and adam, 2025. URL
ptttps://arxer org/ abs/2502.04664, 10
[47] Shuo Xie and Zhiyuan Li. Implicit bias of AdamW: -orm constrained optimization. In Forty-first International Conference on Machine Learning, 2024. URL hnttps://oppnreeeee
net net/forum? i CecCkXXkk06.J.... 10
[48] Pan Zhou, Jiashi Feng. Chao Ma. Caiming Xiong, Steven Hoi, and E. Weinan. Towards theoret- ically understanding why sgd generalizes better than adam in deep learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 1711879546...... 10
[49] Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flammarion. A modern look at the relationship between sharpness and generalization. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
10
[50] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. In The Eleventh International Conference on Learning Representations, 2023. URL http:::/t/////t /openreview met/forrtt
id=1UYpN14qjTF. 10
[51] Avery Ma, Yangchen Pan, and Amir massoud Farahmand. Understanding the robustness difference between stochastic gradient descent and adaptive gradient methods. Transactions on Machine Learning Research, 2023. ISSN 1835-8856. URL htttttt////peereeieeee .nt//
Featured Certification. 10
[52] Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy- tailed class imbalance and why adam outperforms gradient descent on language models. In The Tiihth Annual Conference on Neural Information Processing Systems, 2024. URL wwt./feeefeeee formmrrmddiee (drrrm6m 6aV80c. 10
https://openreet
[53] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020, Curran Associates Inc. ISBN 9781713829546. 10
13

---

[54] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi- Quan Luo. Why transformers need adam: A hessian perspective. In The Thirty-eightt Annual Conference on Neural Information Processing Systems, 2024, htt.httts:///ppeeeeees X6rqEpbnj3. 10
www. met/forummiee id=
[55] Yan Pan. Toward understanding why adam converges faster than sgd for transformers. ArXiv, abs/2306.00204, 2023, https://apppp semanticscholes http:oorg//ooooooeeeee httg/ooppuuDD::54228330
10
[56] John Burkardt. The truncated normal distribution. https: /ttpss::/ppppeeeeeee sc. fsu. edu/ wwapenattiinnttttteeteeoee wwweeeeettaiionsttunaaaaaaa wtruucaaeudnormntt wo..ooomm............... 2023. Department of Scientific
Computing, Florida State University. 16
[57] Evan Z. Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6781 PMLR, 18 24 Jul 2021. URL https://:/ htttpp::///roooe //ppocee//o/cc/ngs.. mlr. peress/vv39/110
html. 27
[58] Devansh Arpit, Stanistaw Jastrzębski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization deep networks. In International conference on machine learning,
pages 233-242. PMLR, 2017. 30
[59] Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-fumme funnttiorrrunnioinnnnnn map biased towards simple functions. International Conference on Learning Representations, 2019. 30
[60] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5301- 5310. PMLR, 09-15 Jun 2019. URL htthstp://. htttps://ppooceedinng mlr www.mmwww......eeemmmmmmmmeeeeee htt.s//////raaaaaa11100
html. 30
[61] David Barrett and Benoit Dherin. Implicit gradient regularization. In International Con- ference on Learning Representations, 2021. URL /ttpp//ps///ooerrreree wwt.:/ntttrrr/.////
3q5IqUrkcF. 30
[62] Robert Geirhos, Henn-Herikk Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning deep neural networks. Nature Machine Intelligence, 2(1 2(66)665,,7,, 2020. 30
14

---

Appendix
A Omitted Proofs
15
A.I Gaussian Data
15
A.2 Toy Data
20
A.3 Auxiliary Results
24
B Additional Experiments and Details of Settings
25
B.1
Gaussian Dataset
26
B.2
MNIST Dataset with Spurious Correlation.
26
B.3
Dominoes Dataset
26
B.4
Subgroup Robustness Datasets
26
B.5 Boolean Features Dataset
28
C Discussion
30
A Omitted Proofs
The proof for Proposition as follows.
Proof. The optimal predictor can be found by solving for the following:
2
1
exp (#I-w.)²
20
(wa-Ne)²² + exp
20
(#1-2:)?
20
(#2+
20
= exp
20
20
x3
Simplification yields
0.5 + exp
Rugzz = exp
(Hi+ta)er
12222 N
-
of
20
2
Taking log on both sides and rearranging, we get:
(#1+pa)&2
Maza 1i-r3 2o
23 20
log 00.5 exp (-2ugr 2
of
of
+
For isotropic Gaussians, simplifies to (M1 + (a))a1 + 1222
-i-u3-u6 u
o² log (0.5 (1 exp (- 2/1222
Under realizability, we get
w + 82 log (0.5 (1 exp expeeppp2p00 2umz
(-203) )))
A.1
Gaussian Data
can prove Proposition 2 as follows.
Proof. The population gradient can be simplified as follows.
E w x 2 O]yæ]] [(x)w 0, = |u
- x2y 1,e= --1
0ly=
[Pr[w | 1,6= 1) Pr| 0
4
[(x|w" x 0, 10yy=1,e- 11 2 x 0ly -1 0, -1)).
15

---

The conditional expectation E(x' w x' 0) can be simplified as follows. Let μz' E(x'). Since we can write x' " x' w
w| x' w 2 x'. we have
E( a' 2 0) [ a' 0) E(z'Iw<x' 20
= [(w x w x' 2 0) L(x())
w F( x' w x' 2 0) E E(x') EE (2())
w
= 1 u
||| E(w ww x' 2 0).
Using result on the mean of truneated normal distribution from Burkardt [56], and that for a given w, w x' is Gaussian random variable, we have,
o( Haw
E(w's a'|w' '' 2 20 pww Jw
Rw
where == w p', Fw 1. Then, we have
Hu
E(x'lw' x' 2 20) 1l
Ow
Using the above, we can write the population gradient -aE (w' x 2 0lyx]] as:
-0.25a (1+, + (u,, ()) + (n_ oI' (H, w) 2po + ol' i))
where p. Pr[wTat x Oly 1,e 1] Φ0 س p Pr[w Oly 11,=-11 -1] Φ 4Lاي
Po = Pr[wTax æ 2 Oly = -1] = $b ( ) and (1z, w)) 1-
(using the facts
that for any 2, x-=)) ((=) and 1 ((-2) (=))) Simplifying the expression then finishes the proof.
Next, we state the full version of Theorem and prove as follows.
Theorem 5. (Full version of Theorem 1.) Consider the data in Eq. (1), neurons initialized such that
ag +1 with probability 0.5, d == 2, 02 = oy G, w 2 and Xo := 2 2 0.8. Let Wk, := lim4--00 unk. and 0k,00 wx-- f for k [m]. Then, the solution learned by gradient flow in the infinite sample setting with correlation loss is:
00 = ax[1,
Proof. For neuron j E [m]] let Bj.e denote the angle between wy and the r-axis at iteration 2 0. We drop the subscript for convenience. Let waD [1,0]T. Then, cos 0₂ w w We want to see if 6g tends with time. Specifically, given 0 [-x, w7. we want to show that dens 0. We dt
have:
deos0;
ais WED a [ (T w. WEDD in?
a²o
((W2-1) ( w.) (Vii[ + 2A
Wi.1 2 205 ( m)))
a|w:l w²+1
Twel
wt.1
A( w: w, (AAi (.) w. 24h (2( ( + (AAūI -20(N1
a²od (wf-1)
- ())))) 1 (ApIIw.)) wh+1 |w.1"
11.1 (Ø (QAi) w.) $ (App' w.)) + iwi1
- 15 (App)
A|wal
wh+1
|.m|"
a'adsin'00 4|
0.
(لبيم) (.(AiI ((A/I
2n cos प..
( (b (.)) & u. 200
sin 0,
The first and third terms are always positive, so the sign depends on the second term. We note that the derivative is 0 when w,2 0, i.e. 0, = 0. This indicates that once 0, becomes 0, it remains 0.
16

---

Also, using the mean value theorem, we can write:
A((W² M((w>-1)am 2wm)),))
+ M(w2-1)m 10
= ((() Diwsin 8x (w*+1)
Iwo (w=+1)
Iwil (w=++1
for some ce
A((w 1(w,,1) -www,.) A((w² 1)w, 1 2
Clearly, o(c) ((A). The second term is
[0: (wx((
wrd (w++1)
lower bounded by
(w*+ o( X) cos dr. We now consider two cases:
8dw2
Case 1: 0g é - or 0, [m/2, In this case, cos 0g 0, so the second term, and hence the derivative, positive.
Case 2: 0g E [- x /2]: In this case, cos 0, > 0, so the second term is negative, and we have to compare its magnitude to the other terms. Using & (AAxi, wn;) (>p w.) 1 and png wi 11 have:
dcus Og
4² sin 01
(w²-1)
8w A
cos (,)(()() - (-1
4| ww
Since cos 0₂ 1 and ((-1) O, the RHS is positive when
~²-1
4w 2 (()).
Let E ((A0, w) = dw
w²-1
A0 Ao
w²+1)
20
dE
(Nu
x + (Ao )
w²+1 2w
2 0,
dba
2w
when Ag
2u w
The RHS here is a decreasing function of w for w 2 2. The condition becomes
Ag 0.8. dE
(0A0 2++
2 3:سس1 (x>o - 0.
2
Since E' is an increasing function of both w and X, and we can numerically verify that E 2) 0, the result is true for all and do 0.8.
This shows that for neuron wk. a₁ dros dt
2 C (sin for some constant C' > 0. [wr.el
Next, using Proposition 2, we can show that the gradients are bounded and consequently, the iterate deos 9 de
norm is upper bounded as ||wx,,|| c(t + 1), for some constant 0. This gives as C" (s)n for some constant C" ++1
Next, consider ox. and suppose cos Oz. stayed below some I, 1 for all t. Then, (sin 02,e))2
1 L 0, deasla.
(L+t) C1(1-13). Integrating both sides, get cos 0n, cos fos0k C'(1 _²²) log(1 + t). which diverges as t + co, leading to a contradiction as cos(-) 1. The case where aş follows similarly.
Hence, as + co, sin 0n, + O. and thus cos cos + sign(ax).
Next, we state the full version of Theorem and prove follows.
Theorem 6. (Full version of Theorem 2.) Consider the data in Eq. (1), neurons initialized such that
ag = +1 with probability 0.5, and sup; Wk,oo| 7/2, = 2, 0g By = a, w 2 2 and 0.8 s " s 1.5. I for k [m]. Let 6g denote the direction of wko Then, the solution learned by signGD in the infinite sample setting with correlation loss is:
Let Wy = lim +
wh.t and Wk == w0.
V2 [1,
0g > 0, sin 0 k.o 0 > 0,
V [1,-
0g 0, sin 0 R < 0,
[1.0]T
> 0, sin 00.0 = 0,
[-1.0]
ag < 0.
Proof. For signGD, we can analyze the gradient expression for any wr
vwI(W)= --ar ( ( ((ApI
24 (1
(d(A))0)) + ((Nū[ 2 w)
17

---

Specifically, the gradient can be in the direction [=1,0]* only when [(W) = 0. We have:
(0.1]VoI(WW) = ad
wit1
2)
1 (b w) sin x) + o (Ap__w) 2Φ 2Φ 22((5)))
0a sin 0
2nd w*+1
<(c) c( + (Ap 2Φ (1)z
where CE [Apl w. Apr w0. Consider the expression in the parenthesis. Assuming
w- 1, have: 2n
2 2 + 1 ω{ ابرهة 2()) ((c) + ( (Ap) @(Apl 26( 20(A)))) 2 ((( ((
22 (2) 1) )( (()) 0,
whenever S 1.5 (we can check this numerically, and use the fact that ((2) is a decreasing function of 2 2 0).
Thus, the gradient only in the [31,0] direction when sin 0, i.e., when is in that direction.
Next, we can check there are neurons in the [0, =1] direction. We have:
[(())
22 - (*(ApI ((Ai[m)) 2A + cos ( d(Ap 20(2)
The expression in the parenthesis is positive as long _ 0.5A 2 0.4, or T.6w
3-2-1 2 L. Let
1.6w E(Aa,w) w) = Ao 3--=1 We can show that it is an increasing function of both Ag and w. Since E(0.8,) 2) 0, the result holds for all w 2 2 and Ao 2 0.8.
Based on these calculations, the updates are along [+1,+1] directions, depending on the sign of and sin Specifically, we have four cases as shown in Table 5 for O, and 0.11 any L.
sig sign(sin +ve
+ve sagn(a)
sign(-[o,1]p 1]v I(W))
esin 0,.1)
+ve -ve
-ve
-ve
-ve +ve
-ve -ve
+ve
+ve
Table 5: Different cases for 0g and 0 in the analysis of signGD.
Using the small initialization condition, the first iterate is dominated by the update direction: WK,1
Wk.0 1/80, where 80 sign(V Zign(VT.(w...))) so for each coordinate i, [w,,, 2 n/2 n/2,, hence Wk,1 sign-aligned with So- Consequently, the next update follows from Table and the argument proceeds by recursion.
This shows that whenever > O, the updates for neurons in the first/second or third/fourth quadrant are along [1, 1]T or [1,-1]* respectively. However, when a < O, the updates for neurons in the first/second or third/fourth quadrants alternate between [-1,+1]] and [-1, F1] As a result, even iterations, these neurons are close to the [] direction (but may not be exactly aligned due to the initialization). However, in the limit oo, these neurons converge in this direction.
We state the full version of Theorem 3 and prove it below.
Theorem 7. (Full version of Theorem 3.) Consider the data in Eq. (1) with d 2 and w 12].
[1) 1}. 2 0.8./w. Consider two predictors,
oy
Linear: = sign((x1)
s sign (3 x2) sign(3x x22)
32 2 0, 32 < O.
Piece-wise Linear: T
Then, it holds that E(j' +y) E(ÿ # v) O.
18

---

sign ax1 Bagg(aa11
+ bxz)
322 2 0,
Proof. Linear: أ sign (x0)) + bro).. Piece-wise Linear: j'
sign axt1 bbxz)
I2 < 0.
( ((((m~-))
+ + ---0)
( a 6 nw- : -
a
nwnt #1
E(j # v) Φ
1/2 +0
Va
Bg Vatko E2
Vxov
When 1,b= = 0, we get:
w now- w
M afmww 1
E(Ú + v) = 1000
+ ob
w
20, v
20
Considering the non-isotropic case, when 3, b 1, we have:
E(g'= # v) < 1 Φb
3mw2-3-2 2ad
+
3p
kw
V9 loy
2
20
= E(g' +)) E(Ú + v) 1 Φb
3nm -- 2
+b
kw +1
2 1b (-
V9N+ldv
2
2
VKo,
2
3nm
+ 1
u
Kiw
3u
nw=+1 +1
=: 0.555 (oo,
VBK+Iay
2
Vkoy
2w
Vxou v
2
where Ao == +. We can analyze the first derivatives: 5."
dE dw
An (- 3(ww0+1) (6 2nd
(-
VBe+i
An
3NW2-3+24
2w
+ (ww)+1)
Vic
O da mw-=1)
2w
30 2-1)
VR
of 3dg mwitd))
2w
dolmm o( do -w==1)
3/R
Aa 8w
(3kw" Tauuw--+220 342w)2
(now -199 3(mw2-1)
A2 nwn +1) (xw)-1)²)
2w 2w2 /n
Vx
v9x+1
exp
9
(mw)++1)
exp
8w3
Aolnu> <(
de VR
xw²=1)
3/4
(rooo -193) (M -3+2w 2- (9rce
3(mW2-1 (mo*+1)
A2 " +2+5me
2nd 2=²VE VR
20
v9x+1
exp
Smw (9)c+1
exp
2wth
This is positive when the following two conditions hold:
(m(3kw>² - 3+ 200² (9x () - 1) = x ( + 9 4w - 18kw> - 12w + 12kiw") ))- - (9h 1) + 1
= (12w" w²) + "Go'" 12w) 1> O,
when -(3w-6)+ (3w- (12-))))+ -(3w-0)+ vB -w4 24w 36 0 since S 12. K S 1 implies //w-))))33000) (12w2-w*)
(12w#-w*)
12w3 w Gw² - 12w - 12 O. Since do 2 0.8 Kww, and w [2, 12], we have:
3/R expl
(122
)ox(0 -1(66)-122)-1)) 12))-11)
3(ow2
2w" 1 2+ Bmww)) 2n
Sno (9N+1)
(Mw²+1) exp
1
exp
0.8"79
145 exp
0.8854 O.
VIO
80
2
Next, we compute the derivative wrt Ag.
dE aho 2
1 (- 3mw2-3-2 Φ1
An
3NW -3+2w 2w
(nw2-1)
0i da )w²=1) 3lw3++1
oi 3dg men
VaN+I
V9x+I
VK 2n
VR
(1w2-1) 200 21 do xut yic
1333mu73322 exp (- [-10(" (120" 8w x(9+t))
-w -w==-=""(((( )n(( p -12w)-1)
3( +1)
Xb(222 ( +2+ Sw?
(xw exp
Since Rww 1,
1 VR(3ku2 -3+2w) exp
(1w))1111
((w0++1)
((Eー-1)) exp AG(222 #425mm= 2wr
ViK+I(mu2-1)
2 v10
24-1+ exp( w--1
0.8279 80
3 w-- exp(
0.8254 1
v10 9 exp(
0.8279 ) 3 exp( 0.89
> 0.
Next, we compute the derivative wrt K.
dE dn (- =ow""" 6( An V
3nw 2 -342
(Rw)+1)
di da اwاس-) 31m=~~11 (1 Vos
3d
mew ))
dx Iu (9x+I))99NTI
2w
RVR
Vx
2w
>0(m=0+11 Si mor-1
(1 / 2 +ow 22243)
-23(n" (12m" "nw -12
Awni v 2n
(9x=1) / 9n- ( exp
Ow xw9+1)
3(mm²-1) (mw+-1) exp
W 8
The expression in the parenthesis is lower bounded as:
2-w33
exp(- 80
0.8270 3 145
exp(
0.8254
57(144) 21
exp(-
0.8"79
3143
exp(
0.8"54
> 0.
109/2
2
T04/2
145
80
145
2
0
19

---

A.2 Toy Data
We can write the Bayes' optimal predictor in the toy setting as follows.
We consider three datapoints (æ, y): ([-Ha,0] ",' -1), [[2",,2]".11 .1) and (//1,=42]1, :1), where P1, H2, Mg > O. The optimal predictor can be found by solving the following:
min( (1n(V/11-10 - 11)) (x2 12) (01 - 11)2 ((1-111)
(12 12) V(01 nz) 2 + xz.
Solving this gives piecewise linear function:
(N1 ps) M232 -1+13-13
22
11 2
(N pz)21 P2
22 S
In the realizable setting, this is:
WT + 22 0,
22 0
22 S 0. WTT 12 0,
We now state the full version of Theorem followed by the proof.
Theorem (Full version Theorem 4), Consider ||wx| s Bw
nM
((3w I)(w² 1) 4w?)
(4w² -ww--ww- (w² A Sw (2w + 11 and + v3 w² < 2,2, Let WK
lims-coo for
Ilima
tan w²-1 2x
neuron & E [m] and :== Adam as shown Table 6,"
Then, for m m oo, the solutions learned by GD, signGD, and
GD
Adam (3) 82 0) or signGD
Adam (3, = B3 1)
[1,0]'
w.p. P
[1,0]T
w.p. P
[1,0]' [-
w.p.
[-1,0]
w.p. 3
[1,
w.p.
w.p.
v2 [1,
w.p. 8
WF.00
2 1, 2w
w.p.
[1.1]
w.p.
v2 [1,-1]
w.p. B
2 1,
w.p.
[1=1]"
w.p.
[8 11 72
w.p.
W.p. 8
Table 6: Solutions learned by GD, signGD and Adam (see Theorem 4).
where is constant between 0.72 and 1. In each case, the sign of the first element of Wko is the same as sign(ax).
-. -.-2. Define three sets S1, S2, S3 as S : (= S: : 02. S: 0). ( (2eS:32- 0).
Proof. Let 2 == (x, v). and 3 := 22
First iteration. We first analyze the gradients at the first iteration. Consider different cases where "ho0 0 depending on different samples Table 7 lists the population gradients depending
which samples contribute to the gradient. See Fig. 7 for an illustration. Note that = tan Ma
tan
20 and h 0 tan 2-11 T.
2w
20

---

Sct s.L. w. 20
Pop. Gradient Es-pl yoeja
Prob. of such we
Sgu 8q
82
S.
5 + pal
Iuc
8,08
9
Jun
Syu 8
$ [e 2pp,. -pol
8g
Figure 7: An illustration of the toy dataset and set such that w depending on the region of initializa- tion.
Table 7: Population gradients and corresponding proba- bilities depending on the region of initialization of the neurons.
Using the population gradients in Table 7, the updates for the different algorithms, are written as:
w.p.
2
x
R-1
w.p.
x
w
ag 4
w
w.p.
tan
2w
GD:
Wk,1
wk.0 +
# 1
W
w.p.
tan
2w
2u
2 9-11
3x
tan 2
2w
w.p.
tan 2-1
33 +1
2w
11
w.p.
2 2w
w.p. tan w--1
[1,0]T
200
3 -1
SignGD/Adam: Wk,1 wk.0 d&77
[1,1]T w.p. 1
26 tan
[1,-1]T w.p. 2
T
Second iteration. Next, these updates to analyze the second iteration. Tables 8 to 10 include the updates at the second iteration for GD, signGD and Adam, respectively, where we use the conditions and the (small) initialization scale. Specifically, the small initialization scale helps ensure that the gradient and the corresponding updated neuron are in the same region (in terms of which samples contribute the gradient for the next iteration). Using the condition on w, the updates in rows 5 and 7 of Table remain in the direction of the points Z₂ and ž3. respectively, whereas those in rows and 11 get along the direction of [1,0]*
Based on the updates in Table 8, we can write the GD iterate at any time t > as:
tan 2
w.p.
2x
2
w.p.
Whit Wk.1 + nu((-1)
tan 11 w
W
w.p.
205
2
2-1
w.p.
25
21

---

[wxx1 Wk.0 wk.o)/ (np)
Prob.
Set S s.t. W x 0
4sign (22) [ x-D
x
0
y 1
w-
[1-11 0]"
tan 11
2r. 2
S2u U.S.
[
-1
Si
0)
2m
tan -1 -
0J
2m 2
S2 u Sa
tan -1 -
2
Si
2 -
-tan11 2
2
S2
2w
As
2
T tan 2w
w
2
S1
2w
4m
π 2 tan -
-1
Sa
20
4m
2
1 2
T ttan" 2
S
2
4m
T
2
2w
S2 U3
3w
T tan"
2
1
w
2w
SI
2w
4m
π tan 2w
2
S2u S3
2
4m: -
2
tan
S
0)"
20
4m
Table 8: Population gradients at the second iteration for GD.
Wk.1
Prob.
41E2-D
1
11 2: 2
0
0yyar)
1/p p
Wk.2
ta --11
0[1,0]"
2w
[ _0TT
0[,,0]T
-][
tan
2w
[+ OOT
x
1]
T:
2
ttn
w
2
2-2- _1]T
1[1,1]
2 ttnn" -1 w--1
1]]
2
[
-0]1,00
2x
v|1,-1]"
2 -1]
7/1.-1]
2x w
-n 11]
tan""
اقف _OT
- [1.0] OT
2g Table 9: Population gradients at the second iteration for SignGD (Adam, 81 = B3 = 0).
Based on the updates in Table 9, we can write the signGD iterate at any time as:
tan -1 w
[1,0]T
w.p.
2
tau --1
-[1,0]
w.p.
w 1 tan
[1,1] w.p.
2
Wk,t W + mt
2T
W 2 2w
[1,-1]]
w.p. 1
tan
w 2
1[1
w.p-
tan
2w
-1/]1
w.p. 4
Lan 2 2
22

---

Wk1
wk.
Prob.
4L
w
x
2 0
Wk.2 1
o[1,0]
- 0J'
tan -1
2
7 "1,00"
2x tan -
0]
-nl 72 (w== ( 0]"
2
2x
tan 3
o[1,0]"
(1 00'
n[ v (w== () 0]'
2n
2x
tan w -
-1[1,0]'
-n[1,0]"
32 tan
n[1.1]'
2
2w
n|1,1]"
2
2 -tan" -
-"[1,11"
w 1
1/l (w² (u²-1)/2+((A+)) 1]"
2w
1x
V((=² 11/22
TI 2 -tan" w
-11,-11"
A 1
1|1,-
2w
1x
20
-1],,-11
2 tan 2
0
2w
(1)22) +(W++1)² 1]
(1²-1)/2
2w
4x
n[1.1]"
B +1
11
2 tan 11
2w
n6 72
1]
2w
1x
2
taan w
-1[
2w 1]
2
007
2,5
(w=+1)² 1]
4x
/2
v((3=2+1)
7[4,-1]"
3w 11T
tan 2M
11 /2
1u
2w
4x
2
--1
3w 1
2 tan
2.35
2w
4x
V2 ((3wx++)//2)
Table 10: Population gradients at the second iteration for Adam, B1 = B2 1.
Based on the updates in Table 10, we can write the Adam iterate at any time as follows:
tan 1
[1.0]
w.p.
2T
w 1+ (T- 1) 21x
tan 2
VF
Vise 1)=+ (-- مF
0
w.p.
2xr
tan-1 -
VF
1
(
J++1+ 1ا
11)
w.p.
2m
2w
tan 2
[-1,0]
w.p.
2
2
tan-1 w²-1
[1.1]
w.p. /
4
(w 1|(2
tan 2w
1
w.p. 5
(1-
1Ke = w0.0 + 11 Σ
tan 2-1
[1. 11
w.p.
4 =
1
(w (-2-))))11)
tan 2w
w.p.
tan 2-11
20 (r- (1(1)-
w.p.
4x
tan-1 -
(3w 11 (7-1)(
w.p.
((w
4T
tan 1 W
VF 1
3w 1)(W
w.p.
2
2-11
(3w (322=11)/----11(2+111
tan 2w
w.p. B
4/
23

---

t + oo iterations. Based on the analysis above, we can compute lim
#k..
for each algorithm.
For GD, have:
tan w
w.p.
2x 2n
w.p.
lim
tan
w.p.
3 -1
w
2w
w.p.
tan
1n
20
For Adam with 8 0 or signGD, we have:
[1,0] w.p.
tan 2
[-,,0] w.p.
lim 1-000
What n
tan 10
[1 w.p.
2w
2
2x
2-1
[[,--1] w.p. 2x 2w
tan
For Adam with 3 1, using the results in Appendix A.3, we have:
--1
[1,0]'
w.p.
2T 2u
21
[-mi,
w.p.
2T
[m2,
0]
w.p.
tan
2ST 2-11
[-1,0]
w.p.
tan 2u2
2m
lim
Wh.r
[1
w.p. 8
tan
-1
4m P-1
26
[-
[1] w.p.
tan
2u
4
W -1
[-m3, 0]
w.p. 1
4
2m W
[m4, 1]
w.p.
4m
tan 2 -1
[m4, w.p. 5
[-mss, 0] w.p. 1 tan 2-11
2u
where ,,,...,, my are constants that satisfy 0.935 $ 1, 0.923 S m2 $ 0.84 ma s 1, 0.72 < my 1, 0.98 my $ 1. Taking my and normalizing each direction then finishes the proof.
A.3 Auxiliary Results
Lemma Given a constant r O and function fr(x)
/z(x-1-15) where x: 2 1, it holds that (x) 2 0 when x 2 rr Further, when x e IN, the minima occurs at either x = 1+ or = 2 [r]r],
g-1+r
and it holds that:
min
br|r [(++))
1-irler
$ff(z) 1
The result can be obtained by examining the derivative of Sz fx(x) with respect x, so we omit the proof.
24

---

Further, given "1 w²-1
12:= w²-1
w²+1
1/r1, 7=:= 0.5r1 71
3w²+1 w
0.5 F's
0.5 w++1 and 1t/5
holds that:
0.4472 S <
r < 2.236, 0.2236 & rs 0.5, 3 5.4721, 1.2236 & rs 1.5.
Alternately, for specific value of
compute these exactly. For instance, when w f1 = 0.6, r2 = 1.6667, r3 0.3, r4 4.3333, ng 1.3.
Also, can simplify the lower bound on (x) as follows:
0<r<l
v (2(1+)2) 2tr
1<r<2
)3(24r2)
3+F
<<<<2/3
c(r) min
In|tr ((4+=1)))))))) ((++))(+))))))))))
1+|r|-r
V4(3+r2) 4+r
</3<r<4
V5(14n2) S+r
<<r<5
Vo(5+r2)
6+r
5<r<6
VT(G+23)
We use this to obtain the exact lower bounds for the aforementioned intervals:
0.9354,
0.4472 5 < 1,
0.9238,
mi)
0.8443,
0,2236 5 r 0.5.,
0.7240,
$Srs5..4211.
0.9802,
1.2236 s r 1.5.
When = 2, c(r1) s 0.9713, c(r2) s 0.9681, c(rs) ~ 0.8803, c(ra) 0.7808, c(rs) 0.9936.
Lemma 2. The sum f(x) := Σf J satisfies 2 25 fxx) s -1
Proof.
To establish the bounds for ((x), we can compare the sum to the corresponding integral. We
have:
f(x) Σ 1 J
1
Vn dn 2/ 1-
f(x) Σ s1+ f
1
dn 1 2-27--1
Combining both inequalities and using Using the fact that vx + +12 2 V finishes the proof.
Additional Experiments and Details of Settings
All experiments in this section were performed on an internal cluster with NVIDIA V100 and P100 GPUs with 32GB memory each.
We list the licenses under which various datasets used in this work were released as follows. All the datasets are publicly available. MNIST is released under the CC BY-SA 3.0 license. CIFAR-10 released under the MIT license. The code to gen- erate the Waterbirds dataset is released under the MIT license. The creators of the CelebA dataset encourage its use for non-commercial research purposes only, but do not mention license name. MultiNIII is released under the ODC-By license. CivilComments is released under the CC BY-NC 4.0 license.
Adam
Figure 8: Comparison of the decision bound- aries learned by GD and Adam with the Bayes' optimal predictor (dashed green) across three tings of our synthetic setup (details in the text).
25

---

B.1 Gaussian Dataset
Additional Results. In this section, we discuss additional results for the Gaussian dataset. Across all settings, we consider full-batch GD with a learning rate of 0.1 and Adam with a learning rate of 10 and 31 = 82 = 0.9999. We set m 500, a 0.01, w 2 for the results in this section. In the finite sample setting, we consider d = 100 and (T y xx) = (0.2, 0.15. 0.01). When n. = 5000, μ = 0.2, Adam achieves 0.618% better test accuracy than GD. When n = 30000, μ = 0.25, the gap is 0.521%. In the population setting, we consider d = 10, μ 0.3, (y: σ.) (0.1,0.1, 0.01) and Adam achieves 0.292% better test accuracy than GD. The decision boundaries learned by GD and Adam in these three cases shown in Fig. 8(a), (b), and (c), respectively. In Fig. 3, the test accuracy of Adam is 0.55% more than that of GD.
Effect of Stochasticity. We repeated the experiment in Fig. 1 with batch size 50 to see the effect of stochasticity/minibbaae
training. The results are shown in Fig. 9. We observe that the decision boundary learned by Adam is more nonlinear than SGD and closer to Bayes' optimal predictor. However, the decision boundary learned by mini-batch Adam is less nonlinear as compared to full- batch Adam.
Figure 9: Comparison of the deci- sion boundaries learned by GD and Adam with the Bayes' optimal predic- tor (dashed green) with batch size 50.
B.2 MNIST Dataset with Spurious Correlation.
Training Details. We train two-layer NN with width
64 using Adam or SGD with learning rates 10-3 and 0.1, respectively, using batch size 64. The network is initialized as follows. Hidden layer weights are initialized with small-scale random initialization, specifically from Gaussian distribution with dimension. Final layer weights are initialized as + #1/ ///m with equal probability and kept fixed, where m is the hidden dimension. We train using BCE loss till the loss reaches 5e 3.
0.001/Vd, where d is the input
B.3 Dominoes Dataset
Dataset. The Dominoes dataset [9, 12] is composed of images where the top half of the image shows an MNIST digit [13] from class {0, 1). while the bottom half shows an image from other image datasets such as MNIST, Fashion-MNIST or CIFAR-10. In our case, we use CIFAR-10 10 [14] images from classes (automobile, truck). which the MNIST-CIFAR dataset. Fig. in Qiu et al. [11] demonstrates how the MNIST-CIFAR dataset is generated as well as more example images.
Training Details. We train both
ResNet-18 and ResNet-34 model for to 500 epochs or until convergence using batch size of 32, weight decay 10-5, initial learning rates 10-3 for SGD and 10 for Adam, and a cosine annealing learning rate scheduler. We average results across 5 random seeds. The groups in the test set are balanced.
Additional Results. Following Kirichenko al. [24], we generate datasets with spurious correlation strengths of 99% and 95% between the spurious features and the target label, while the core features are fully predictive of the label. We report the final average worst-group accuracies for each correlation strength and optimizer in Table 1l and Table 12 for ResNet-18 and ResNet-34, respectively. In both cases, training with Adam leads significantly better metries.
Method
99% comrlatiin
95% correlation
Original Acc.
Core-Oaly Acc. Decoded Acc.
Original Ace.
Core-Oaly Acc.
Decoded Acc.
SGD
0.00_a.00
0.D0
42..0
1.8l
1.00_.77
71.04...8 84.66.0.18
Adam
0.20
0.41 Lo.50
7137
14.17,A15
20.63
Table 11: Worst- group accuracies for original accuracy, core-only accuracy, and decoded accuracy for a 18 model trained using SGD and Adam.
B.4 Subgroup Robustness Datasets
Training Details. Following Sagawa* et al. [16], for the image datasets, we use the Pytorch torchvision implementation of ResNet50 [23] which is pre-trained on the ImageNet dataset, and
26

---

Method
Original Acc.
Core-Only Acc.
90% correlation
Decoded Acc.
Original Acc.
95% correlation
Core-Only Acc.
Decoded Acc.
SGD
0. 0...0
0.00.0.00
47.60
0.24
_..0 18.31_5.39
69.2613.94 82.68, on
0.0
3,32,2.97
69.11
9.60,419
Table 12:
Worst-group accuracies for original accuracy, core-only accuracy, and decoded accuracy for ResNet-34 model trained using SGD and Adam.
for the language-basee based datasets, we use the Hugging Face pytorch-trans implemen- tation of the pre-trained BERT gooch/hyperparamete setting with the highest worst-group accuracy on the validation set.
model [22]. We report test results for the
Main Results. We use batch size 512 for Waterbirds and MultiNL and 1024 for CelebA and CivilComments. The results averaged over five independent runs for the image datasets and four independent runs for the language datasets. We fine-ture until convergence, which takes 7 epochs on the language datasets, epochs on CelebA and 100 epochs Waterbirds. Fig. 10 shows the hyperparameters (learning rate, weight decay and momentum parameters) considered for the two optimizers across these datasets and the worst-group test accuracies the last the-tuning epoch. Fig. 2 and Table 13 show the final results. For Adam, we find that lower values of 31, 32 (compared to the default (0.9, 0.999)) are generally better: the best values for each of them are 10 0.5, across the four datasets.
SGD
& m
Adam A
Adart 8 00W
718000
(10.
SrouD That A2C
Waterbirds
2000 TICk
005
150
6J5
100
P
6
SGD
Adam
B 10
Adom -8-00 0.5
Adam A-0.99 0.939
1001 STOID Teet AcE
VIOT STOUD TOnk AZC
CelebA
0.25
0:50
1.62
033
0.55
2-5
1.59
600
325
0.5
0.75
1.00
0.39
0.25
0,50
6.5
1.50
SGD
Adam A 10
Adam A
0.5
Adam Bg 0,999
01311
0991
crleds
(18.00
MuitiNUI
5GD
Adam = 10-1
Adam By = 0.5
Adam B 01999
CivilComments
Womngne Kep
MOreo Thst Aest
322
0.50
1.50
8.30
4.50
9.54
0.75
140
020
8.30
3.75
1.00
020
030
6.50
475
Figure 10: Hyperparameter sweep for SGD and Adam optimizers on four benchmark datasets for subgroup robustness, showing worst-group test accuracy at the last fine-tuning epoch. We sweep the momentum parameters (3 for SGD, 31,32 for Adam), learning rate (1)), and weight decay (A).
Results with Smaller Batch Sizes. Here, following Liu et al. [57], we use batch size of 128 for the image datasets and 32 for the language datasets. The results are averaged over four independent runs, and presented in Table 14. We used the default momentum values for Adam in this setting, but tuned the learning rate and weight decay.
27

---

Waterbirds
CelobA
MulaNII
CivilC ommmnts
Optimizer
Average
Worst-group
Average
Worst-group
Average
Worst-group
Average
Worst-group
SGD
16.9.1.09
35.33.
93. 7..227
67.56
81.12.0.22
05.92.2.8
86.08,0.18
62.89.1.24
Adam
87.91
76.04_204
90. 19.70
77.22.8.22
81.99.0.32
71.181274
85.7320.58
66.25_3.m
Table 13: Comparison of average and worst-group test accuracy on four benchmark datasets for subgroup robustness, using larger batch sizes: Adam outperforms SGD. See Section 4 and Appendix B for details.
For the image datasets, we use SGD with momentum 0.9 and set the learning rate 10-3 for Waterbirds and 10 for CelebA and the weight decay as 10 for both datasets. For Adam, tried learning rates {10-5, 10=4) and weight decays {10-6, 10-4) for both datsets. We use learning rate and weight decay of 10-5 and 10-6 for Waterbirds, and 10 and 10 for CelebA for the final results.
For the language datasets, we tried the following settings. For Adam, we tried learning rates {10-5, 2. 1 for both datasets. For SGD, we set the learning rate as 10 and tried momentum values {0, 0.9}. For both optimizers, weight decay as 0 across all cases and use a learning rate of 10 for Adam for both datasets. The SGD-mmoeentmm set as 0.9 for MultiNLI and 0 for CivilComments.
tried weight decays {0, 10-$). For the final results, we set
Consistent with the results using larger batch size, we find that even with smaller batch sizes, Adam attains better worst-group test accuracy and comparable average test accuracy compared to SGD, as shown in Table 14. However, we observe that the gains are smaller for image datasets and larger for language datasets in this setting.
Optimizer
Waterbirds
CelebA
MultiNLI
CivilComments
Average
Worst-group
Average
Worst-group
Average
Worst-group
Average
Worst-
SGD
85. 05.74.0.34
71 21.95
93. 73..32
67. 67.64_..00
80. 17.0.7
54.38
72, 12.99
43. 13.72.4.22
Adam
86.33,1.08
78.44.267
94.10.1.11
68.19.2.6
81.78,0.16
67.21
83.71,0.9
69.64.211
Table 14: Comparison of average and worst-group test accuracy on four benchmark datasets for sub- group robustness, using smaller batch sizes: Adam outperforms SGD. See Section 4 and Appendix B for details.
B.5 Boolean Features Dataset
Dataset. Formally, let de, d E N be the number of core and spurious features respectively, d, e N be the number of independent features that are independent of the label, and let d' de d, d,, be the total dimension of the vector. For some x ε (-1,+13d denote I. E [-1,+11d H, E +11,+1)d, and I. e {11,+1) to be the coordinates of x that correspond to the core, spurious, and independent features respectively. Let À [0, 1] be the strength of the spurious correlation. Define Boolean functions:
: (+1,-111d [+1,-1}, fs: (+1, (+1.-1}
Next, we define the distributions D Dsame. Danr and Dx. Define D Dsame to be the uniform distribution over the set of points in {-1, +1) where the and spurious features agree:
Damme Unif (( (-1,+13 fo(x.) fxx)))
Similarly define Dair as the uniform distribution over {-1, +1]d where the core and spurious features disagree:
Dun == Unif (fx f-1,+11) : fe(x.) fa(x))))
Lastly, define Dx as the distribution where with probability X, a sample is drawn from D and otherwise sample is drawn from Dair-
Details of Hyperparameters. We use a width 20 NN with Leaky ReLU activation with 0.01 negative slope and train with cross-entropy loss. We use 10000 training samples and 5000 test samples drawn from data distribution where d 50, d, 41, training for 80000 epochs. We consider normal initialization with μ 0 and o V Jeatures For each optimizer, we did sweep across various learning rates and picked the best performing learning rate, The results are reported after averaging across 5 random seeds.
28

---

Additional Results. In Table 15, we consider 4 initialization schemes for the model parameters: (1) the default PyTorch linear initialization, which is a uniform initialization on (-Vk, Vk) where k in features' and (2) normal initialization with = O and a= av featuress with a = 10-3, 10 & 1. These results in line with those in Table 4 and show that Adam leads to richer feature learning. We report in Table 16 the final test accuracies, decoded core, and decoded spurious correlations at the end of training for each optimizer, highlighting that Adam maintains its superior performance throughout training. We also report in Table 17 the average test accuracies, decoded core, and decoded spurious correlations at lowest comparable training loss for SGD and Adam, where we use 9 core features instead of 8, so that de = 9, dg = 1, du 40. We use the same correlation strength = 0.9 and uniform initialization.
Mathod
Text
Unioon DCC
DSC
(.01)
DSC
Nomnal
DCC
Text
DCC
DSC
98.61,212
0.7
0.54..33
05.00.
0.95 0o
0.77
0.77at.33
19.58,1.02
0.78. .7
0.51.am 0.87
Adam
0,84.
0.42.cmm
13.035
0.89.o.e
97.77 29.0
0,87
0.42.e09
Table 15: Test accuracy, decoded core, and decoded spurious correlations averaged across 5 random seeds for SGD and Adam at lowest comparable training loss. Higher test accuracy and decoded core correlation are better. Lower decoded spurious correlation is better.
Method
Test
Uniform Ini
DCC
DSC
Test ace
Normal Inkt (n-
DCC
0.01)
DSC
Test
Nornal init to
DCC
0.31
DSC
Test aoe
Normal iin (G- 1)
DCC
DSC
SGD Adam
95. 26.54.m
0. 0.75
0.10,
96
0.75
0.05yn
96, 12
0.81
0.76,033
80.54,1.0
v.5
90 20.03.0077
0.94.
0,31,19
99 28,
2.96
0 0.
99.22
0. 0.98
0.36.
22
0. 0.96
0 0.20
Table 16: Final test accuracy, decoded core, and decoded spurious correlations averaged across 5 random seeds for SGD and Adam at the end of training.
Training Dynamics. We report in Figures 11 and 12 the test accuracy, decoded core, and decoded spuri- ous correlations throughout training with de = 8, d, 1, x 0.9, for 5 different seeds with normal initial- ization (a 1) for SGD and Adam respectively. From these dynamics, we observe that SGD learned the spurious feature early in training and gradually learns some of the core feature but still retains the spurious feature information. In contrast, when training with Adam, we see the spurious feature is learned early in training, but forgotten quickly as training progresses to instead learn the core feature.
Unitorm init
Method
Test acc
DCC
DSC
SGD
83.47 94
0.23 os
0.74_0.02
Adam
94.22,
0.70...06
0.51 .0.05
Table 17: Test accuracy, decoded core, and decoded spurious correlations on the Boolean features dataset for SGD and Adam at lowest comparable training loss. de 9, d, 1, d,, 40.
Text Accunwyy
Descolal Core Condation
Decoded Spurieus Comn atton
0.
Figure 11: Training curves for SGD with normal weight initialization (a 1) for random seeds.
Text Accusac
Descded Core Commlation
Decoded Spur Comeeation
100
0.4
58
0
5top
330
35C
20k
42x
Figure 12: Training curves for Adam with normal weight initialization (a= 1) for 5 random seeds.
29

---

C Discussion
In this section, we discuss some limitations of our work.
First, we note that the theoretical results hold under some assumptions, such as the population setting and with linear loss function, which are quite different from practice. However, given that this the first result of its kind it is natural to start with some simplifying assumptions. We would like to note that there are no prior works analyzing the training dynamics and implicit bias of signGD or Adam for NNs. The only exception we are aware of is Tsilivis et al. [8], which analyzes steepest descent algorithms (including signGDD for homogeneous NNs trained with an exponentially-taiie function. This paper focuses the late stage of training, and assumes separable data, which does not apply in our setting. Additionally, in their setting (for instance, using uniform disc distributions instead of Gaussian to ensure separability in the population setting), it can be shown that both the linear and the nonlinear predictors are KKT points the max-&2 and max- margin problems. Consequently, cannot distinguish between the implicit biases of GD and signGD in their framework. Therefore, it is natural to make some additional assumptions to make the analysis more tractable and take a step towards understanding and contrasting the implicit biases of Adam/signGD and GD. Generalizing these results to broader settings is an interesting direction for future work.
loss
Second, in this work, we focus on settings where simplicity bias in NNs hurts generalization. However, this is not always the case. There large body of work showing that simplicity bias DL is helpful and can explain good in-distribution generalization [58, 59, 60, 61]. Simultaneously, when the goal to ensure good performance under certain distribution shifts, such as OOD generalization robustness to spurious features, simplicity bias has been shown to be detrimental in such cases [62, 9, 21]. In this work, we focus on the latter setting to showcase the benefit of richer feature learning encouraged by Adam. However, in general, either one of richer simpler feature learning may be desirable depending on the problem.
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See Appendix C.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach only tested on few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken low lighting. Or a speech-to-text system might not used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
31

---

Assumptions are stated in the paper and full theorem statements and proofs are included in Appendix A.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
If the contribution
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soures models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
32

---

Answer: [No]
Justification: We use open source data and models, and standard optimization algorithms. All details of experimental settings are included Appendix
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr
public/guides seniddesCooeeuubmie aessonPPoicyy for more details.
While we encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /nniss... cccpc/////cc.c cte/public/ae gtwggg//ecoedduiiiieeeeeee iss assoonoliccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
33

---

It should be clear whether the emror bar is the standard deviation or the standard error of the mean.
is OK report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https http///eerrpp..pee www.liic/thhcsuuieeeee.
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special considera- tion due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of Machine Learning general. is not tied to specific application. There are many potential societal consequences of this work, but we don't feel any of them must be specifically discussed here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance). fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
34

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: See Appendix B.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswi thcode. com com/dataeet has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license the derived asset (if it has changed) should be provided.
35

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
36

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
for Please refer to our LLM policy (https: pettts://neurre ceseurrippscc/ooofeeeeeeeeeee what should or should not be described.
37


Reviews:
[{"review_id": "a3sAb2Py16", "reviewer": "Reviewer_xUiG", "summary": "This paper demonstrates that, in comparison to Adam, stochastic gradient descent exhibits a simplicity bias when training two-layer ReLU neural networks  on a binary classification task with Gaussian data. The empirical analysis also highlights how SGD preferentially converges to simpler models.", "questions": "I did not thoroughly examine the proof, but a critical question arises: does the conclusion hold for all two-layer ReLU neural networks on a binary classification task with Gaussian data, or only under specific conditions such as particular final weight configurations and neuron initializations? This distinction is crucial because it determines whether the observed simplicity bias is an isolated phenomenon or a universal property of such models. Clarifying this would strengthen the paper’s contribution by addressing whether the result is broadly applicable or contingent on specific architectural or initialization choices. I will change my score according to the authors response.", "limitations": "see questions", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThe paper introduces a compelling theoretical framework by leveraging the Bayes optimal predictor as a benchmark since there are no prior works analyzing the training dynamics and implicit bias of signGD or Adam for NNs. The derivation of signGD as a first step in the analysis is particularly novel and insightful.\n\nWeaknesses:\n\nThe theoretical guarantees seems to rely on assumptions about initialization and final weights. The analysis of Adam with momentum remains incomplete (but it's still a good start)", "was_revised": true, "final_justification": "The authors have explained about the limitation and I think it's a interesting work.", "rebuttal": "Thank you for your time and helpful feedback. We are encouraged that you find our theoretical framework to contrast the implicit bias of Adam and GD in NNs compelling, and the analysis novel and insightful.\n\n**Q: Do the results hold under specific conditions such as particular final weight configurations and neuron initializations?**\n\nThanks for the question about the scope of our theoretical results.\n\nTheorem 1 and Theorem 2 hold for any random (small scale) initialization of the hidden weights ($w_k$). There is no constraint on the directions in which the hidden weights are initialized. This is in accordance with standard initialization conditions used in the deep learning theory literature as well as the Pytorch default Kaiming initialization (upto a constant). The only condition is that the output weights $a_k$​ are symmetric ($a_k=\\pm 1$ w.p. 0.5). \n\nImportantly, we **do not** assume any condition on the final weight configuration. The limiting directions (as stated in the final equations in Theorems 1 and 2) arise from the training dynamics themselves. These equations present the limiting directions each neuron converges to, depending on how it was (randomly) initialized. We also show how the neurons evolve over time in practice in Fig. 3 (second row), which matches the theoretical results. We will further clarify this in the paper.\n\nWe hope that this clarification helps address your concern and that you would consider increasing your score."}, {"review_id": "THyt4SzYUu", "reviewer": "Reviewer_uF13", "summary": "The paper compares the implicit bias of Adam to that of Gradient Descent (GD) and shows that in some settings GD converges to \"simple\" solution, whereas Adam is able to capture richer features. The paper contains theoretical results as well as a range of experimental results. While the theoretical results are derived for one specific type of input and for a simple model, the experimental results explore a number of different datasets and models.", "questions": "See above.\n* Can you give a longer more extensive intuitive explanation for the effects observed? Or is this difficult\n* Does Theorem 1 hold regardless of initialization?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n* The implicit bias of Adam is far less well studied than the implicit bias of GD. So it is great that this paper adds to the investigation of Adam.\n* The experiments are done using simple, but insightful, clean settings. They are also relatively broad and add coherently to the story.\n* While the derivations of some of the theoretical results are a bit brute force (which may be necessary, I am not claiming to know a simpler way), I found the paper quite well written in general.\n\nWeakness:\n* As common for these kinds of things, the theoretical results are quite special. The model is a two-layer ReLU network, but the second layer weights are not even trainable.\n* There could be even more of an intuitive explanation or intuition why Adam behaves the way it does.\n* Pleas provide code. While I agree that it seems that the code should be pretty standard, it would allow for much easier replication. Also, certain details are difficult to find in the paper. For example, for the MNIST Dataset with Spurious Features, you say you train a two-layer neural network, but what is the width?", "was_revised": true, "final_justification": "I have taken note of some concerns raised in another reviews regarding the loss function used and, just in general, the setting that is studied formally here is very specialized. However, for me the inclusion of the experiments partly make up for those weaknesses and so I would like to keep my positive score.", "rebuttal": "Thank you for your time, careful reading and helpful feedback. We are encouraged that you appreciate the paper’s contribution to the study of implicit bias of Adam, and find the experimental settings simple yet insightful, clean and broad, and the paper well-written overall.\n\nWe respond to specific comments below and will incorporate these clarifications in the paper.\n\n**W2, Q1: Intuition on Adam’s behaviour vs GD**\n\nThanks for the question. In our setting, different latent features have different effective scales. GD treats all coordinates uniformly, so the largest‑margin (“simplest”) feature dominates. Adam adaptively rescales coordinates by their second‑moment estimates, approximately normalizing these scale differences; this allows smaller / higher‑frequency features to be learned rather than being suppressed by the dominant one. \n\n\n**W3: Regarding Section 4.2 and Code**\n \nWe used MLP width 64 for the experiment in Section 4.2 (MNIST dataset with spurious feature). As we cannot share external links during the rebuttal, we will release the code for the experiment upon acceptance, as suggested.\n\n**W1: Fixed second layer weights in the analysis of two-layer ReLU NN**\n\nWe fix the second‑layer weights to isolate and make analyzable the hidden‑layer dynamics. This is a standard practice, e.g., [1-4].\n\n**Q2: Does Theorem 1 hold for any initialization?**\n\nYes, Theorem 1 holds for any random initialization of the hidden weights ($w_k$). The only condition is that the output weights $a_k$​ are symmetric ($a_k=\\pm 1$ w.p. 0.5). \n\n\nThanks a lot for the time and effort to review our work. We appreciate your positive endorsement and support.\n\n\n**References:**\n\n[1] Frei et al., “Benign Overfitting without Linearity: Neural Network Classifiers Trained by Gradient Descent for Noisy Linear Data,” COLT 2022.\n\n[2] Cao et al., “Benign Overfitting in Two-layer Convolutional Neural Networks,” NeurIPS 2022.\n\n[3] Frei et al., “Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data,” ICLR 2023.\n\n[4] Kou et al., “Implicit Bias of Gradient Descent for Two-layer ReLU and Leaky ReLU Networks on Nearly-Orthogonal Data,” NeurIPS 2023."}, {"review_id": "H1iDf7e39T", "reviewer": "Reviewer_w3en", "summary": "A bulk of papers have shown that (S)GD has a bias towards simple solutions in different machine learning tasks. Many view this as a positive property, as it aligns with the Minimum Description Length principle and introduces implicit regularization. This paper comes to undermine this thinking by showing that SGD's solutions are too simple compared to Adam's. The main focus is to demonstrate that the implicit bias of SGD can be counterproductive in some simple settings.", "questions": "None", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 1, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Presentation and writing are quite good. The authors did an excellent job of conveying their ideas\n- This paper comes to challenge a long-standing convention, and I commend the authors for it.\n\nWeaknesses:\n- I am afraid that there might be some flaws in one of the proofs. The proof of Theorem 1 relies on the fact that $ \\\\frac{d}{dt} \\\\cos(\\\\theta\\_t ) >0 $ for all $ t \\\\geq 0 $. From here, the authors conclude that $ \\\\lim\\_{t \\\\to \\\\infty} \\\\cos(\\\\theta\\_t ) = 1 $, which is not true. For example  $ \\\\theta\\_t = \\\\frac{\\\\pi}{6}(1+e^{-t}) $ satisfies\n$$ \\\\frac{d}{dt} \\\\cos(\\\\theta\\_t ) = \\\\frac{\\\\pi}{6}e^{-t} \\\\sin \\left(\\\\frac{\\\\pi}{6} \\left(1+e^{-t} \\right)  \\right) > 0 \\qquad \\\\forall t \\\\geq 0$$\nbut $ \\\\lim\\_{t \\\\to \\\\infty} \\\\cos(\\\\theta\\_t ) = \\\\frac{\\sqrt{3}}{2} $. This can render the Gaussian-Mixture section irrelevant. (Resolved in rebuttal, and the score increased)\n\n- The toy data example is less significant for the following reasons:\n1) Both Adam and GD classify 100% correctly for this problem (see Fig. 4), so there is no generalization or other benefit in using Adam.\n2) If we use exponential loss instead of linear, then we end up with a separable data+homogeneous network+exponential loss. According to [40], GD will converge to the maximum margin solution, which in this case coincides with the Bayes' optimal predictor (please correct me if I am wrong). So this result depends only on the unusual choice of linear loss function, whereas with the exponential loss, GD has a clear advantage.\n\nRelated work:\nThere are additional papers that you should add to your related work section. For example\n- The implicit bias of minima stability: A view from function space\n- The Implicit Bias of Minima Stability in Multivariate Shallow ReLU Networks\n- On linear stability of SGD and input-smoothness of neural networks", "was_revised": true, "final_justification": "After diving more into the details of the proof, I am now convinced, as I said in my initial review, that this choice of linear loss function is problematic. First, the objective function (risk) is not bounded. This raises the question of how we can measure the effectiveness of the algorithms. Since this is a homogeneous network in a classification task, we can look at convergence in direction, as the authors did. In this regard, we see that gradient flow actually converges in direction to the \"optimal\" solution. Namely, it does exactly what it is supposed to do. If it doesn't converge to a good separator, it is only because of the bad choice of the objective. In contrast, the adaptive methods don't converge to the optimal direction at all. However, in this carefully designed setting, two wrongs make it right for Adam. Overall, I can't see a real phenomenon here, sorry. This is the main reason why I don't support acceptance.", "rebuttal": "Thank you for your time, careful reading and helpful comments to improve our work. We are encouraged by your appreciation of the paper’s writing and presentation, and of challenging a long-standing convention.  \n\nWe respond to specific comments below and will incorporate these clarifications in the paper.\n\n**W1: Concern about Theorem 1**\n\nThank you for catching this. We will update the proof of Theorem 1 with the following additional steps that address your concern.\n\nThe current proof shows that for neuron $w_k$, $a_k\\tfrac{d(cos\\theta_{k,t})}{dt} \\geq C’ \\tfrac{sin^2(\\theta_{k,t})}{\\||w_{k,t}\\||}$ \nfor some constant $C’>0$. \n\nUsing Prop. 1, we can show that the gradients are bounded and consequently, the iterate norm is upper bounded as $\\||w_{k,t}\\||\\leq c(t+1)$, for some constant $c>0$. This gives $a_k\\tfrac{d(cos\\theta_{k,t})}{dt} \\geq C \\tfrac{sin^2(\\theta_{k,t})}{t+1}$ for some constant $C>0$.\n\nNext, consider $a_k=1$, and suppose $cos⁡\\theta_{k,t}​$ stayed below some $L<1$ for all $t$. Then, $sin⁡^2\\theta_{k,t}\\geq 1−L^2=:m>0$, so $\\tfrac{d cos\\theta_{k,t}}{dt}\\geq\\tfrac{Cm}{(1+t)}$. Integrating both sides, we get \n$cos\\theta_{k,t}-cos\\theta_{k,0}\\geq Cm\\log(1+t)$, \nwhich diverges as $t\\to\\infty$, leading to a contradiction as $|cos(\\cdot)|\\leq 1$. \n\nHence, $sin\\theta_{k,t}\\to 0$, and thus $cos\\theta_{k,t} \\to sign(a_k)$.\n\n\n**W2: Concern about the toy setting**\n\nThank you for the thoughtful comments.\n- We agree that both Adam and GD reach 100% training accuracy on this toy distribution. Our goal here is mechanistic: the linear loss lets us compute the population gradients in closed form and see that Adam’s adaptive rescaling uses both features while GD collapses to the dominant (“simpler”) one. This difference shows up in margin: Adam attains a larger margin, so under mild shifts (e.g., additive noise to the first coordinate), Adam generalizes better. \n\n- [40] (Lyu & Li, 2020) proves that for separable, homogeneous networks with exponentially‑tailed loss, GD/flow converges in direction to a KKT point of the parameter‑space $\\ell_2$ max‑margin problem, not necessarily the global max‑margin solution. It can be shown that in the toy setup, the linear predictor, the Bayes’ optimal predictor and other predictors with decision boundaries between these predictors, all are KKT points of the $\\ell_2$ max-margin problem. Thus, switching to exponential/logistic loss does not necessarily imply GD converges to the Bayes optimal predictor in the toy setup. \n\n\n**W3: Related work**\n\nThank you for the pointers to related work, we will add these in the updated version of the paper.\n\n\nWe hope that these clarifications help address your concerns and that you would consider increasing your score."}, {"review_id": "wTqMXpFz8z", "reviewer": "Reviewer_mTkP", "summary": "This paper presents a theoretical and empirical investigation of the simplicity biases of SGD and Adam, giving evidence that SGD can rely on simple, spurious features in a way that may cause reduce out-of-distribution generalization, while Adam tends to learn \"better\" features. Theoretically, it is proven that when training a two-layer NN (fixed second layer) on a binary classification problem with two-dimensional data coming from a particular gaussian mixture, gradient flow learns an appoximately linear predictor, while signGD learns a non-linear predictor that achieves a lower test loss under certain additional conditions on the data. In a further simplified setting, there are analogous results for GD, signGD, and Adam, although there is no comparison of the test losses in this case. Empirically, it is demonstrated with a large number of synthetic/real datasets that Adam is more robust to spurious correlations than SGD and tends to learn the \"core\" features instead of spurious features.", "questions": "1. Theorem 3 shows conditions under which the test error of signGD is smaller than that of gradient flow. Are there similar conditions under which the opposite conclusion holds?\n2. I don't understand the condition on initialization in the statement of Theorem 6, which says that $\\|w_{k,0}\\| \\rightarrow 0$. Can you define this more precisely?  Also, how is this condition used in the proof?\n3. How do you reconcile your results with those of Zou et al, who show that GD can generalize better than Adam when training CNNs?\n4. In Figure 4, the decision boundaries for Adam and SignGD are super similar. Why is this? I think the paper would be improved by pointing this out and adding some explanation, since in practice we expect Adam to be very different from SignGD.\n5. In Table 2, why is agreement with a linear model a good measure of complexity? It seems that you could have low agreement with a baseline linear model while having a very simple decision boundary.  For example, if the baseline linear model has parameter $\\theta$, then another linear model with parameter $-\\theta$ will have 0 agreement with the baseline linear model.\n6. For the experiments in Section 4.2, what were the training loss/accuracy for SGD and Adam? Is it possible that both algorithms have a similar generalization gap but Adam has lower training error?\n7. In the setting of Figure 5, did you measure the margins on test set samples? Does it also hold for test data that Adam tends to have larger margins?\n8. Do you have any theoretical results that don't use $d=2$? If not, I recommend improving the clarity of the presentation by changing the dataset definition in Equation (1) to remove all dimensions besides two.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The problem considered is very important, and the paper provides a new perspective on the generalization of Adam vs SGD. As pointed out in the paper, it seems a conventional wisdom that SGD generalizes better than Adam, so it is interesting to see evidence of the opposite conclusion. These are the first theoretical results (that I'm aware of) showing a generalization improvement for Adam.\n2. The paper contains a lot of empirical results, which gives confidence that the theoretical conclusions are not specific to the synthetic data model. In particular, the experiments with ResNets and CIFAR/MNIST data in Section 4.3 are pretty convincing.\n\nWeaknesses:\n1. The data model used for the theory is somewhat artificial and the data is always two-dimensional, which is a big limitation in my eyes. Still, I think it is okay for a first step in this direction.\n2. It is not clear whether the complexity of the decision boundary is actually a meaningful way to compare the algorithms. The main conclusion of the theory is that gradient flow learns an approximately linear predictor, while signGD learns a non-linear predictor, and experiments back this up (i.e. Figure 1). But the gap in terms of test accuracy is not as significant, both in theory and practice. Theorem 3 says that signGD's test error is smaller than that of gradient flow under certain conditions, but it is not clear how much smaller or whether we should expect the opposite conclusion to hold under other natural conditions. Similarly, in Figure 1, it is true that Adam's decision boundary is more non-linear than that of SGD, but visually it appears a comparable distance from the Bayes optimal predictor, so is Adam's predictor necessarily better?\n3. Theorems 1 and Theorem 2 show a difference between gradient flow and signGD, which is not the same as GD vs signGD. It seems possible to me that the difference in learned features could be due to the continuous time dynamics compared to the discrete time dynamics, and this point is not acknowledged in the paper. I think that the paper will be more transparent if the authors use the name \"gradient flow\" instead of GD throughout the paper when referring to the results of Theorem 1.", "was_revised": true, "final_justification": "I have kept my score the same. The authors have addressed many of my small questions/concerns, and I maintain my positive opinion of the paper, although my confidence is still somewhat low, especially after seeing the proof issues raised by reviewer w3en. I have not thoroughly studied the proofs, and I am getting the impression that there are some details of the theory that could be ironed out. In particular, the comparison of SGD and Adam in Section 3.1 actually compares gradient flow (continuous time) with Adam (discrete time), and it's unclear whether the difference in algorithm behavior could be attributed to this difference rather than the their different learning rules. Also, the full statement of Theorem 2 (Theorem 6 in Appendix C.1) seems a bit informal in how it describes the initialization; I'm not sure whether it is fully rigorous, and the rewording suggested by the authors in their rebuttal does not fix the issue. Still, I appreciate this paper's contribution to our understanding of SGD vs. Adam.", "rebuttal": "Thank you for your time, careful reading and detailed helpful feedback. We are encouraged that you find the problem very important, our perspective on the generalization of Adam vs GD novel and interesting, and appreciate our comprehensive empirical results.\n\nWe respond to specific comments below and will incorporate these clarifications in the paper.\n\n**W2: Test accuracy gap for the synthetic setting is not as significant**\n\nWe agree that, on the synthetic Gaussian task, absolute accuracy gaps are modest (between 0.2 to 0.6%, depending on the setting). By design, this is a controlled setting to isolate the difference between (S)GD and Adam, and the results serve as corroboration for the theory. The practical impact appears on the larger-scale benchmarks where Adam’s advantage is more significant.\n\n\n**W3, Q8: Presentation**\n\nThank you for the suggestions, we will incorporate them in the updated version. We will use gradient flow instead of GD to refer to Theorem 1, as suggested. Regarding the dataset definition in Equation (1), our theoretical results are for $d=2$, but we report some empirical results for $d>2$. We will further emphasize that for theory, $d=2$, as suggested. \n\n\n**Q1: Conditions where gradient flow beats signGD?**\n \nThank you for the question. As discussed in Appendix B, there can be cases where the solution learned by GD is better. For the synthetic setting, the solution learned by gradient flow can have better test accuracy when $\\sigma_y$ is much larger than $\\sigma_x$. We will add an example in the updated version to demonstrate this concretely. More broadly speaking our goal is not to say one algorithm Adam is better but to provide an explanation why it might be better in some instances. However, there is a kind of no free lunch phenomenon at play here and one can always find settings where one training algorithm is better than another.\n\n\n**Q2: Initialization in Theorem 6**\n\nThank you for the question. We will state a more precise version of the initialization condition as follows. There exists $\\alpha\\to 0$ such that $sup_⁡k\\|w_{k,0}\\|\\leq \\alpha <<\\eta$. Initializing with a small norm ensures that the first and subsequent iterates are dominated by the update directions rather than by the (arbitrary) initial offset, i.e. $w_{k,1}\\approx −\\eta sign(\\nabla L(w_{k,0}))$. \n\n\n**Q3: Reconciling with Zou et al. which shows that Adam generalizes better than SGD for CNNs**\n\nZou et al. analyze in‑distribution generalization, while the focus of our work is on generalization under distribution shifts. For in-distribution generalization, the simplicity bias exhibited by (S)GD can indeed be beneficial. However, under distribution shifts, such as in the presence of spurious correlations, this simplicity bias can be detrimental and the richer feature learning enabled by Adam can be helpful. \n\n**Q4: Similarity between the predictors learned by Adam vs signGD in the toy setting**\n\nGood observation! In the toy setting, the decision boundary of Adam is slightly more nonlinear than signGD, and the difference is smaller than what we may see in practice. Following the reviewer’s suggestion on highlighting the difference between signGD and Adam, in the updated version of the paper, we will add experimental results for signGD on the group robustness datasets (Section 4.4) to more clearly showcase and discuss the effect of $\\beta_1$ and $\\beta_2$ on the implicit bias of Adam.\n\n\n**About the MNIST with spurious feature Experiment (Section 4.2)**\n\n**Q6: Train loss/accuracy comparison for SGD and Adam**\n\nFor both predictors, the train loss is approximately 0.005 (as mentioned in Appendix D.2), and train accuracy is 99.98%. Thus, Adam has better test accuracy and generalization gap compared to SGD.\n\n**Q5: “Agreement with a linear model” as a complexity proxy**\n\nThanks for the question. The linear model here serves as the max‑margin linear classifier on the training data (obtained via ERM on the same data, leading to near-perfect train set accuracy). A linear alternative that shows noticeably lower agreement with the max‑margin linear classifier (such as the one suggested by the reviewer) would have a smaller margin or worse accuracy on the train set. In contrast, we see that the NNs trained with SGD or Adam get near-perfect train accuracy. To summarize, a higher capacity model that also attains near‑zero training error, would either have high agreement with this (max-margin) linear classifier and hence, would be close to linear, or have noticeably lower agreement (without sacrificing accuracy), which is achieved by exploiting nonlinear structure. \n\n**Q7: Margin distribution on test set**\n\nThanks for the suggestion. We checked the margin on the test set, as suggested, and observed a similar trend: Adam has an overall larger margin compared to SGD. To be precise, the mean $\\pm$ std dev for the margins across the test set is $0.36\\pm 1.22$ for SGD and $1.69\\pm 1.75$ for Adam.  \n\n\nThanks a lot for the time and effort to review our work. We appreciate your positive endorsement and support."}]

Author rebuttal:

Meta-review: This paper provides a theoretical and empirical comparison of the implicit biases of SGD and Adam, particularly in settings with structured data and linear losses. The authors demonstrate that Adam tends to prefer nonlinear predictors that approximte Bayes' rule, while SGD prefers linear predictors, offering valuable insights into how optimization algorithms affect generalization.

Even though the setting is rather restrictive (e.g., linear loss, structured data distribution), this work clearly illustrates the distinct implicit biases of GD and Adam, which is an important contribution. The analysis is rigorous and well-structured, and the observations are novel and relevant to the broader understanding of optimization in deep learning.
Decision: Accept (poster)