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

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.

We respond to specific comments below and will incorporate these clarifications in the paper.

W2: Test accuracy gap for the synthetic setting is not as significant

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

W3, Q8: Presentation

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

Q1: Conditions where gradient flow beats signGD?

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

Q2: Initialization in Theorem 6

Thank 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}))$.

Q3: Reconciling with Zou et al. which shows that Adam generalizes better than SGD for CNNs

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

Q4: Similarity between the predictors learned by Adam vs signGD in the toy setting

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

About the MNIST with spurious feature Experiment (Section 4.2)

Q6: Train loss/accuracy comparison for SGD and Adam

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

Q5: “Agreement with a linear model” as a complexity proxy

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

Q7: Margin distribution on test set

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

Thanks a lot for the time and effort to review our work. We appreciate your positive endorsement and support.

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.

We respond to specific comments below and will incorporate these clarifications in the paper.

W1: Concern about Theorem 1

Thank you for catching this. We will update the proof of Theorem 1 with the following additional steps that address your concern.

The 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}\||}$ for some constant $C’>0$.

Using 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$.

Next, 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 $cos\theta_{k,t}-cos\theta_{k,0}\geq Cm\log(1+t)$, which diverges as $t\to\infty$, leading to a contradiction as $|cos(\cdot)|\leq 1$.

Hence, $sin\theta_{k,t}\to 0$, and thus $cos\theta_{k,t} \to sign(a_k)$.

W2: Concern about the toy setting

Thank you for the thoughtful comments.

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

• [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.

W3: Related work

Thank you for the pointers to related work, we will add these in the updated version of the paper.

We hope that these clarifications help address your concerns and that you would consider increasing your score.

While dealing with an unbounded objective, optima do not exist in the classical sense, only infima. And yes, the linear loss artificially creates additional solutions, since all divergent trajectories are equally good. However, if we look only at the direction, there exists only one optimal direction. In this direction, the derivative of the objective w.r.t. the angle is zero (i.e. it is optimal). What you showed is that the direction of gradient flow converges to this direction.

The best way to see that these results are due to this linear loss, is to replace it with an exponential. And as I stated in the initial review, under exponential loss, gradient flow will have the advantage. Additionally, even with linear loss, your proofs are limited to very specific settings (that have some kind of symmetry). If adaptive methods really had an advantage, it should have been easy to generalize your proofs, for example, to separable data. However, I cannot see an easy way to do it.

For clarification, I do believe that adaptive methods have an advantage, but the paper hinges on a problematic setting described above, which depicts the wrong image (the wrong reason why these methods can be superior).

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.

We respond to specific comments below and will incorporate these clarifications in the paper.

W2, Q1: Intuition on Adam’s behaviour vs GD

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

W3: Regarding Section 4.2 and Code

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

W1: Fixed second layer weights in the analysis of two-layer ReLU NN

We fix the second‑layer weights to isolate and make analyzable the hidden‑layer dynamics. This is a standard practice, e.g., [1-4].

Q2: Does Theorem 1 hold for any initialization?

Yes, 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).

Thanks a lot for the time and effort to review our work. We appreciate your positive endorsement and support.

References:

[1] Frei et al., “Benign Overfitting without Linearity: Neural Network Classifiers Trained by Gradient Descent for Noisy Linear Data,” COLT 2022.

[2] Cao et al., “Benign Overfitting in Two-layer Convolutional Neural Networks,” NeurIPS 2022.

[3] Frei et al., “Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data,” ICLR 2023.

[4] Kou et al., “Implicit Bias of Gradient Descent for Two-layer ReLU and Leaky ReLU Networks on Nearly-Orthogonal Data,” NeurIPS 2023.

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.

Q: Do the results hold under specific conditions such as particular final weight configurations and neuron initializations?

Thanks for the question about the scope of our theoretical results.

Theorem 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).

Importantly, 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.

We hope that this clarification helps address your concern and that you would consider increasing your score.