The Implicit Bias of Adam on Separable Data

We appreciate your support, and address your questions as follows.

Q1: Adam without $\epsilon$ sometimes does not converge. Contradiction with [1, 2]? Study a non-zero $\epsilon$ and let $\epsilon=0$ be a special case?

A1: Our goal is to study the implicit bias of Adam when the stability constant $\epsilon$ is negligible. The differences between our result and [1,2,3] that consider a non-negligible $\epsilon$ do not lead to any contradiction. Instead, these differences demonstrate that our work indeed provides novel and complementary insights into Adam. We address your detailed comments as follows:

• We emphasize that Adam without $\epsilon$ provably converges under our setting. Proving convergence even without $\epsilon$ is a contribution of our work.

• Our result does not contradict with [2]. By treating the correction term in the continuous-time approximation of Adam as a penalty, [2] gives an informal summary of the implicit bias of Adam without specifying the loss or model. In comparison, our work focuses on a specific setting, and provides a formal and rigorous analysis on the implicit bias of Adam with concrete convergence rates.

• Our theory does not contradict with the theories in [1,3]. [1,3] prove that Adam with a non-negligible $\epsilon$ has the implicit bias towards the maximum $\ell_2$-margin. The intuition is that, after sufficiently many iterations, $\nabla \mathcal{R}(w_t)$ and $v_t$ will be very close to zero (much smaller than $\epsilon$), and $\frac{m_t}{\sqrt{v_t} + \epsilon} \approx \frac{1}{\epsilon}m_t$, indicating that Adam with a non-negligible $\epsilon$ will eventually be similar to GD with momentum. In comparison, we focus on the setting where $\epsilon$ is negligible, and demonstrate that in this case, Adam will have a unique implicit bias towards the maximum $\ell_\infty$-margin. Therefore, our theoretical results do not contradict those in [1,3]. Instead, our paper and [1,3] can together provide a more comprehensive understanding of the implicit bias of Adam.

• Our experiments do not contradict with the experiments in [1,3]. In our simulations, we run different algorithms for $10^6$ iterations, and Adam can typically reach training losses below $10^{-8}$, demonstrating that it is reasonable to stop the training there. Our experiments show that within $10^6$ iterations, Adam with $\epsilon = 10^{-8}$ is indeed approaching the maximum $\ell_\infty$-margin. In comparison, in Figure 1 in [3], Adam and other algorithms are run for over $10^{14}$ iterations, where Adam with $\epsilon$ can eventually approach the maximum $\ell_2$-margin solution. Therefore, our paper focuses more on the practical stopping time, while [3] focuses more on the behavior of Adam after an extensively long training. Hence, our experiment results are complementary to [3], and there is no contradiction.

• “Start with a non-zero $\epsilon$ and let $\epsilon=0$ be a special case” could be an interesting future direction, and we will discuss it in the revision. However, given that existing works [1,2,3] have studied the case with a non-zero $\epsilon$, we believe that our current setting with $\epsilon = 0$ helps us demonstrate our claim in the cleanest setting.

Q2: Linear model + full-batch Adam is simple compared with [1]. Consider stochastic sampling noise or homogeneous models?

A2: Thanks for your comment. Again, our work aims to study a relatively understudied scenario where $\epsilon$ in Adam is negligible. This setting is complementary to the studies in [1,3], and therefore it is reasonable for us to start from the most classic setting to study implicit bias, which is solving linear logistic regression with full-batch algorithms [3,6].

To the best of our knowledge, existing works on the implicit bias of Adam or AdamW such as [1,3,6] mainly focus on the full batch setting. Extensions to stochastic Adam can possibly be done following the analysis in [5], where the implicit bias of full batch GD is extended to SGD. However, such an extension may be more challenging for Adam due to its complex form. This can be a future direction.

Extensions to homogeneous networks are also an important future direction. However, since our goal is to study the implicit bias of Adam with a negligible $\epsilon$, establishing concrete results in linear logistic regression is a reasonable starting point, and this clean and classic setting serves our purpose well in demonstrating the difference in the implicit bias of Adam when $\epsilon$ is negligible.

Q3: Line 210 claims the results are more accurate than [1], so add experiments on homogeneous networks?

A3: We would like to clarify that around line 210, we do not compare our result with [1]. Instead, we are comparing with [3], which studies linear classification. We will revise our comments and highlight more on our purpose, which is to study a setting that is complementary to [1,3].

We have concerns that experiments of homogeneous models are out of the scope of our paper as we do not claim anything about homogeneous models. However, we have added some preliminary experiment results in the pdf rebuttal page.

[1] Wang, B., Meng, Q., Chen, W. and Liu, T.-Y. (2021). The implicit bias for adaptive optimization algorithms on homogeneous neural networks. ICML.

[2] Cattaneo, M.D., Klusowski, J.M. and Shigida, B. (2024). On the Implicit Bias of Adam. ICML.

[3] Wang, B., Meng, Q., Zhang, H., Sun, R., Chen, W., Ma, Z.-M. and Liu, T.-Y. (2022). Does momentum change the implicit regularization on separable data? NeurIPS.

[4] Xie, S. and Li, Z. (2024). Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization. ICML.

[5] Nacson, M. S., Srebro, N. and Soudry, D. (2019). Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. AISTATS.

[6] Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S. and Srebro, N. (2018). The implicit bias of gradient descent on separable data. JMLR.

I thank the authors for the detailed response.

I would like to clarify that, as I pointed out in my review, I understand that the current work considers Adam without the stability constant. My point of the contradiction between the current paper and previous works lies in that how letting $\epsilon \to 0$ changes the implicit bias from the $\ell_2$ solution to the drastically different $\ell_\infty$-margin solution, e.g., is the transition smooth or abrupt? There exists a gap.

The rest of the rebuttals addressed my other concerns.

We appreciate your positive comments. Your comments and questions are addressed as follows.

Q1: The paper does not present results for a fixed learning rate.

A1: When considering fixed learning rate $\eta_t=\eta$ for some small $\eta$, our analysis can imply that $\lim_{t\to \infty}\big|\min\frac{\langle w_{t}, y_i\cdot x_i\rangle}{\lVert w_t \lVert_{\infty}}-\frac{\langle w^*, y_i\cdot x_i\rangle}{\lVert w^*\lVert_{\infty}} \big|\leq O(\sqrt{\eta})$. We will add a comment in the revision.

Q2: The discussion after Corollary 4.7, comparing Adam and GD, should also compare Adam and GD with adaptive learning rates (see [1] and related work).

A2: Thanks for your suggestion and for pointing out the related work [1]. We will cite it and add more comprehensive comparisons and discussions in the revision.

Q3: (Minor) In Assumption 4.3, ‘non-increasing’ should be ‘decreasing’ or ‘diminishing’.

A3: Thanks for your suggestion. We will change "non-increasing" to "decreasing" in the revision. We will also add comments to clarify that we do not require the learning rates to be "strictly decreasing”.

Q4: Based on the discussion following Assumption 4.2, it might be better to state an assumption on the data and then show that the condition on the initialization holds as a Lemma.

A4: We propose such an assumption because it is the most general version of assumption and it can cover various different settings. For example, we expect that the current version of Assumption 4.2 covers the following two cases:

• Case 1 (fixed $w_0$, random data): consider an arbitrary fixed vector $w_0\in \mathbb{R}^d$. Then as long as the training data inputs $\mathbf{x}_1,\ldots, \mathbf{x}_n$ is sampled from any continuous and non-degenerate distribution, $\nabla \mathcal{R}(w_0)[k] \neq 0$ holds with probability $1$.

• Case 2 (fixed data, random $w_0$): consider any fixed training data inputs $\mathbf{x}_1,\ldots, \mathbf{x}_n$ satisfying that the matrix $[ \mathbf{x}_1,\ldots, \mathbf{x}_n ]$ has no all-zero row. Then as long as $w_0$ is initialized following a continuous and non-degenerate distribution, $\nabla \mathcal{R}(w_0)[k] \neq 0$ with probability $1$.

Therefore, we feel that the current version of Assumption 4.2 may be more general.

Besides, we would also like to clarify that assumptions on properties of initialization have been considered in various previous works studying implicit bias [2, 3, 4, 5]. In particular, [2] makes an assumption that is essentially the same as Assumption 4.2.

Q5: The paper does not comment on how optimal the obtained rates in Corollary 4.7 are.

A5: Currently, we do not have any lower bounds on the margin convergence rates. Due to the complicated optimization dynamics, establishing such lower bounds may be very challenging, and therefore it is difficult to prove that our rate of convergence is optimal. We will add discussions in the revision, and will also mention in the revision that establishing lower bounds on the margin convergence rates is an interesting future work direction.

Despite not having matching lower bounds, our result demonstrates that the margin of the linear predictor converges in polynomial time for a general class of learning rates, which already significantly distinguishes Adam from gradient descent. Such polynomial convergence rates are also supported by our experiment results.

Q6: Why does considering the stability constant $\epsilon=0$ make the setting more challenging? The accumulated second-order moments would be non-zero.

A6: When studying Adam without $\epsilon$, it is true that under Assumption 4.2, the gradient coordinates at initialization are non-zero and therefore the accumulated second-order moments would be non-zero throughout training. However, please note that the impact of the initial gradient values decays exponentially fast during training: in a worst case where for a certain coordinate $k$, $\nabla \mathcal{R}(w_t)[k] = 0$ for all $t > 0$, we have $\mathbf{v}_t[k] = \beta_2 (1-\beta_2)^t \cdot \mathcal{R}(w_0)[k]^2$. Therefore, without a very careful analysis, the worst-case exponentially decaying $\mathbf{v}_t[k]$ may significantly affect the stability of the algorithm, and may even impose a $\exp( t )$ factor in the bounds of convergence rate, making the bounds vacuous.

In fact, many existing analyses of Adam and its variants rely on a non-zero stability constant $\epsilon$, having a factor of $1/\epsilon$ in their convergence bounds, e.g., [6] (see Theorem 4.3) and [7] (see Corollary 1 and Remark 2). In comparison, in our work, we implement a very careful analysis (also utilizing the relatively simple problem setting of linear logistic regression) to avoid having such factors in the bounds.

[1] Wang, M., Min, Z. and Wu, L. (2024). Achieving Margin Maximization Exponentially Fast via Progressive Norm Rescaling. International Conference on Machine Learning.

[2] Xie, S. and Li, Z. (2024). Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization. International Conference on Machine Learning.

[3] Lyu, K. and Li, J. (2020). Gradient Descent Maximizes the Margin of Homogeneous Neural Networks. International Conference on Learning Representations.

[4] Ji, Z. and Telgarsky, M. (2020). Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems.

[5] Wang, B., Meng, Q., Chen, W. and Liu, T.Y. (2021). The Implicit Bias for Adaptive Optimization Algorithms on Homogeneous Neural Networks. International Conference on Machine Learning.

[6] Zhou, D., Chen, J., Cao, Y., Yang, Z. and Gu, Q. (2024). On the Convergence of Adaptive Gradient Methods for Nonconvex Optimization. Transactions on Machine Learning Research.

[7] Huang, F., Li, J. and Huang, H. (2021). SUPER-ADAM: Faster and Universal Framework of Adaptive Gradients. Advances in Neural Information Processing Systems.

Thank you for your supportive comments! We address them in detail as follows:

Q1: Can the authors expand on how they arrive at the right side of inequality after line 292 using 6.1 ? Perhaps take me through the inequality step by step ?

A1: Thanks for your question. We would like to explain that there is a typo on line 292: "for all $t > t_2$" should be "for all $t \geq t_2$".

We now present the detailed derivation as follows. Our target is to derive equation after line 292, which is

$\mathcal{R}(w_{t+1}) \leq \Big(1-\frac{\gamma\eta_t}{2}\Big)\cdot \mathcal{R}(w_{t})\leq \mathcal{R}(w_{t})\cdot e^{-\frac{\gamma\eta_t}{2}}\leq \mathcal{R}(w_{t_2})\cdot e^{-\frac{\gamma\sum_{\tau=t_2}^{t}\eta_\tau}{2}}$. $\qquad$ (target)

To derive these results, we first recall that in equation (6.1), we have

$\mathcal{R}(w_{t+1})\leq \mathcal{R}(w_{t}) - \frac{\gamma\eta_t}{2}\mathcal{G}(w_t)$.

Note that in the proof sketch, for simplicity we focus on the case of exponential loss, and for the exponential loss, it holds that $\ell(x) = -\ell'(x) =e^{-x}$. Therefore, by the definition of $\mathcal{R}(w_{t})$ and $\mathcal{G}(w_{t})$, we have $\mathcal{R}(w_{t}) = \mathcal{G}(w_{t})$. Replacing $\mathcal{G}(w_{t})$ by $\mathcal{R}(w_{t})$ in equation (6.1) then gives

$\mathcal{R}(w_{t+1})\leq \Big(1-\frac{\gamma\eta_t}{2}\Big)\cdot \mathcal{R}(w_{t})$.

This proves the first inequality in (target). To prove the second inequality in (target), we utilize the fact that $1-x\leq e^{-x}$ holds for all $x$. Based on this, we have

$1-\frac{\gamma\eta_t}{2} \leq e^{-\frac{\gamma\eta_t}{2}}$,

which further implies

$\Big(1-\frac{\gamma\eta_t}{2}\Big)\cdot \mathcal{R}(w_{t})\leq \mathcal{R}(w_{t})\cdot e^{-\frac{\gamma\eta_t}{2}}$.

This proves the second inequality in (target). So far, we have proved that

$\mathcal{R}(w_{t+1})\leq \mathcal{R}(w_{t})\cdot e^{-\frac{\gamma\eta_t}{2}}$

for all $t \geq t_2$. Applying this result recursively gives

$\mathcal{R}(w_{t+1})\leq \mathcal{R}(w_{t})\cdot e^{-\frac{\gamma\eta_t}{2}} \leq \mathcal{R}(w_{t-1})\cdot e^{-\frac{\gamma\eta_{t-1}+\gamma\eta_t}{2}}\leq \cdots \leq \mathcal{R}(w_{t_2})\cdot e^{-\frac{\gamma\sum_{\tau=t_2}^{t}\eta_\tau}{2}}$.

This proves the last inequality in (target), and our derivation is finished. We will clarify the derivation of this result in the revision.

Q2: Can the author provide some comments regarding the independence of convergence in the case of $a=2/3$ from $\rho$?

A2: We believe that the margin convergence rate in the case of $a=2/3$ does depend on $\rho$, as the corresponding bound in Corollary 4.7 is $O(\frac{d\cdot\log t + \log n + \log \mathcal{R}(w_0) + [\log(1/\rho)]^{1/3}}{t^{1/3}})$ for exponential loss, and $O(\frac{d\cdot\log t +nd+ n \mathcal{R}(w_0) + [\log(1/\rho)]^{1/3}}{t^{1/3}})$ for logistic loss. Since we are considering the margin convergence rate as $t$ goes to infinity, these convergence rates can be simplified into $O(\frac{d\cdot\log t}{t^{1/3}})$ for both exponential and logistic losses, because the term $d\cdot \log t$ can eventually dominate the other terms in the numerators. However, since $d\cdot \log t$ increases very slowly in $t$ and will require exponentially many iterations to dominate the other terms, in our result we still keep the other terms to make the bounds more concrete.

We guess that you are asking about why $\rho$ does not appear in the margin convergence rates for the case $a<2/3$. In fact, this is because we have applied the simplifications discussed above. Let us take exponential loss as an example. When $a<2/3$, for exponential loss, we can prove (see the bound below line 603 in Appendix B.2.) a margin convergence rate of the order

$O(\frac{d t^{1-3a/2} + d^{\frac{2(1-a)}{a}} + \log n + \log \mathcal{R}(w_0) + [\log(1/\rho)]^{1-a}}{t^{1-a}})$.

It is clear that the first term in the numerator $d t^{1-3a/2}$ is the only term that increases in $t$. Moreover, as $t$ increases, the term $d t^{1-3a/2}$ will dominate the other terms in the numerator in polynomial time, which leads to the following simplification:

$O(\frac{d t^{1-3a/2} + d^{\frac{2(1-a)}{a}} + \log n + \log \mathcal{R}(w_0) + [\log(1/\rho)]^{1-a}}{t^{1-a}}) = O(\frac{d t^{1-3a/2}}{t^{1-a}}) = O(\frac{d}{t^{a/2}})$,

This gives the final bound in Corollary 4.7 for $a<2/3$, which does not depend on $\rho$.

Q3: Is there some intuition with regards to the boundaries and case on $a$?

A3: Again, we explain it for the case of exponential loss. We note that Corollary 4.7 is derived based on Theorem 4.5, in which the upper bound for margin convergence is

$O(\frac{\sum_{\tau=0}^{t_0-1}\eta_\tau + d\sum_{\tau=t_0}^{t-1}\eta_\tau^{\frac{3}{2}}}{\sum_{\tau=0}^{t-1}\eta_\tau})$.

Therefore, to give more concrete convergence rates when $\eta_t$ is specifically set as $(t+2)^{-a}$, we need to calculate $\sum_{\tau=0}^t \eta_\tau^{3/2}$ and $\sum_{\tau=0}^t \eta_\tau$ for different values of $a$. By properties of the series $\{ (t+2)^{-a} \}$, this calculation can be separated into 4 cases as follows:

1. When $0<a<\frac{2}{3}$, $\sum_{\tau=0}^t \eta_\tau^{3/2}=\sum_{\tau=0}^t (\tau+2)^{-3a/2} = \Theta(t^{1-3a/2})$ and $\sum_{\tau=0}^t \eta_\tau=\sum_{\tau=0}^t \tau^{-a} = \Theta(t^{1-a})$.

2. When $a=\frac{2}{3}$, $\sum_{\tau=0}^t \eta_\tau^{3/2}=\sum_{\tau=0}^t \tau^{-3a/2} = \Theta(\log t)$ and $\sum_{\tau=0}^t \eta_\tau=\sum_{\tau=0}^t \tau^{-a} = \Theta(t^{1-a})$.

3. When $\frac{2}{3}<a<1$, $\sum_{\tau=0}^t \eta_\tau^{3/2}=\sum_{\tau=0}^t \tau^{-3a/2} = \Theta(1)$ and $\sum_{\tau=0}^t \eta_\tau=\sum_{\tau=0}^t \tau^{-a} = \Theta(t^{1-a})$.

4. When $a=1$, $\sum_{\tau=0}^t \eta_\tau^{3/2}=\sum_{\tau=0}^t \tau^{-3a/2} = \Theta(1)$ and $\sum_{\tau=0}^t \eta_\tau=\sum_{\tau=0}^t \tau^{-a} = \Theta(\log t)$.

The calculations above lead to the boundary cases in Corollary 4.7.

Thank you for your constructive feedback! We address your comments as follows:

Q1: The paper does not provide an intuition why Adam and GD have different implicit biases. Relation to SignGD?

A1: Thanks for your suggestion. Several recent works have discussed that Adam and SignGD are closely related [2], and the implicit bias of SignGD shown by [1] can indeed provide valuable insights into the implicit bias of Adam. Some of our lemmas are also motivated by this intuition (Lemma 6.1 and Lemma A.3). We will add more discussions on this intuition in the revision. However, we would also like to clarify that, despite the clear intuition, our proof for Adam is not a simple combination of existing techniques. For example, the proofs of Lemma 6.1 and Lemma A.3 are non-trivial, and in Lemma A.3, we have implemented a particularly careful analysis which enables us to cover general learning rates and general values of $\beta_1,\beta_2$.

Q2: The authors claim that Corollary 4.7 is derived under the worst case and this is why in experiments we often observe margins converging faster than the bounds. This statement lacks supporting evidence. The paper should prove that the rate of convergence is tight.

A2: Currently, we do not have any lower bounds on the margin convergence rates. Due to the complicated optimization dynamics, establishing such lower bounds may be very challenging, and therefore it is difficult to prove that our rate of convergence is optimal. We will make it clear in the revision. We will also mention in the revision that establishing lower bounds on the margin convergence rates is an interesting future work direction.

Despite not having matching lower bounds, our result demonstrates that the margin of the linear predictor converges in polynomial time for a general class of learning rates, which already significantly distinguishes Adam from gradient descent. Such polynomial convergence rates are also supported by our experiment results.

Q3: The term "convergence" is used unclearly: objective convergence vs normalized margin convergence.

A3: Thanks for pointing out the unclear expressions. We will revise them.

Q4: on page 6, line 183, $\ell_2$-margin convergence speed for GD should be $O(1/\log t)$.

A4: Thanks for pointing it out. We will clarify it in the revision.

Q5: (page 1, line 25) Typo: reply on -> rely on

A5: Thanks for pointing out this typo. We will fix it.

Q6: Does Theorem 4.5 imply that Adam reduces loss faster than GD? The authors should provide a detailed comparison.

A6: Thanks for your suggestion. You are right that our result shows an $O(e^{-\frac{\gamma t^{1-a}}{4(1-a)}})$ convergence rate of the training loss for Adam in logistic regression when the data are linearly separable, and this is faster than the $O(1/t)$ convergence rate of gradient descent. In the revision, we will comment on it, and give a detailed comparison of the convergence rates of loss between Adam and (S)GD with/without momentum.

Q7: Is $\beta_1\leq\beta_2$ necessary? What happens if we use Adam with $\beta_1>\beta_2$?

A7: Yes, this condition $\beta_1\leq\beta_2$ is necessary in our analysis. Such a condition ensures the stability of Adam. Under this condition, we can show that in each iteration of Adam, the update in each coordinate is bounded by a constant (Lemma 6.5).

On the other hand, without this condition, there exist extreme cases where Adam is unstable. For example, consider $\beta_2=0$, $\beta_1 > 0$, and suppose that there exists a coordinate index $k$ such that at iteration $t = 1$, the corresponding coordinate of the gradient is very close to zero, satisfying $0 \leq \nabla \mathcal{R}(w_1)[k] \leq \beta_1(1-\beta_1)\rho / 10000$. Then, by Assumption 4.2, we can see that $\mathbf{m}_1[k] \geq \beta_1(1-\beta_1)\rho$, and hence $\mathbf{m}_1[k] / \sqrt{\mathbf{v}_1[k]} \geq \frac{ \beta_1(1-\beta_1)\rho }{ \beta_1(1-\beta_1)\rho / 10000 } = 10000$. Clearly, $\mathbf{m}_1[k] / \sqrt{\mathbf{v}_1[k]}$ can be arbitrarily large when $\nabla \mathcal{R}(w_1)[k]$ tends to zero. This implies that when $\beta_2=0$, $\beta_1 > 0$, there are cases where Adam is unstable. For general $\beta_1,\beta_2$ with $\beta_1 > \beta_2$, there are also certain extreme scenarios where Adam is unstable.

We would also like to point out that the condition $\beta_1\leq\beta_2$ is a common condition considered in recent works studying Adam [3, 4, 5]. Besides, this condition aligns well with practice, where popular choices are ($\beta_1=0.9, \beta_2=0.99$), or ($\beta_1=0.9, \beta_2=0.999$). We will add more explanations in the revision.

Q8: Assumption 4.4 seems non-standard. Is it a necessary condition? Why is such an assumption needed?

A8: Assumption 4.4 is necessary for our proof. Our proof of Lemma 6.1 relies on this assumption (line 482-line 484, line 485-line 488). As we commented below Assumption 4.4 and proved in Lemma C.1, this assumption is quite mild: it holds for fixed (small enough) learning rate, or decaying learning rates $\eta_t = (t+2)^{-a}$ with $a \in (0, 1]$.

[1] Gunasekar, S., Lee, J., Soudry, D. and Srebro, N. (2018). Characterizing implicit bias in terms of optimization geometry. International Conference on Machine Learning.

[2] Balles, L., Pedregosa, F., and Roux, N. L. (2020). The geometry of sign gradient descent. arXiv preprint arXiv:2002.08056.

[3] Xie, S. and Li, Z. (2024). Implicit Bias of AdamW: $\ell_\infty$-Norm Constrained Optimization. International Conference on Machine Learning.

[4] Hong, Y. and Lin, J. (2024). On Convergence of Adam for Stochastic Optimization under Relaxed Assumptions. arXiv preprint arXiv:2402.03982.

[5] Zou, D., Cao, Y., Li, Y. and Gu, Q. (2023). Understanding the generalization of adam in learning neural networks with proper regularization. International Conference on Learning Representations.