Paper: On the Implicit Bias of Adam (ICLR-2024)

Reviewer comment (Theory):
While the overall content is easily comprehensible, specific details are inaccessible unless one reviews the supplementary material. Specifically, many theoretical works provide a sketch of the proof in the main paper to explain the techniques newly considered/developed by the authors and clarify how they differ from existing techniques. This paper lacks such details.

Author response:
We wholeheartedly agree that including more details of the proof would be beneficial to the reader. To address this, we have introduced two additional levels of detail. Apart from the full proof in the Appendix, we have incorporated a 3-page derivation in Section SA-9, also in the Appendix. This derivation essentially forms the basis of the full proof, with the key difference being the characterization of all big-O bounds in precise terms. Additionally, following Theorem 3.1 in the main paper, we have added a concise sketch of this derivation.

Reviewer comment (Experiments):
The experimental setup in Section 6 is not self-contained and has elements that seem arbitrary, making cherry-picking possible.

Author response:
This is a very valid concern that we thank you for raising. We have therefore launched additional hyperparameter sweeps, and pictures pertaining to more learning rates have been added in the latest revision.

While excluding the spike regime might seem questionable, we believe it's justifiable in this case. Without explicit regularization and stochastic batching, the spike magnitudes are significant, often reverting the loss to values seen at the beginning of training. This can hardly go unrecognized by any researcher or engineer. We believe studying the spike regime itself is an important direction for research, but Adam in this regime is too unstable to be reasonably approximated by smooth ODE solutions.

To answer your question about the meaning of "sufficiently close," the threshold for entering the spike regime varies with the task and, crucially, the learning rate. A sufficiently small learning rate prevents spikes altogether. Of course, we do not recommend making $ \beta $ and $ \rho $ equal. For instance, with Resnet-50 on CIFAR-10, using full-batch Adam without explicit regularization and a learning rate of $10^{-4}$, $(\beta, \rho) = (0.9, 0.999)$ enters the spike regime, whereas $ (0.94, 0.999) $ does not. We opt for higher learning rates to accelerate training and make the bias term (linear in $h$) more prounounced.

Following your suggestion, we have added more discussion of this topic in the "Numerical experiments" section.

Reviewer comment (Experiments):
In Section 6, the authors refer to Ma et al. (2022) to conduct experiments for the stable oscillation regime when $\rho$ and $\beta$ are "sufficiently close". How do you define "sufficiently close"? Most ML practitioners are familiar with situations where $\rho$ (0.99~0.999) is larger than $\beta$ (0.9). As the authors suggest, is this a spike regime or a stable oscillation regime? There should be a discussion for such questions within Section 6, but the authors merely refer to Ma et al. (2022) without providing specific details.

Author response:
This is a very valid concern that we thank you for raising. We have therefore launched additional hyperparameter sweeps, and pictures pertaining to more learning rates have been added in the latest revision.

While excluding the spike regime might seem questionable, we believe it's justifiable in this case. Without explicit regularization and stochastic batching, the spike magnitudes are significant, often reverting the loss to values seen at the beginning of training. This can hardly go unrecognized by any researcher or engineer. We believe studying the spike regime itself is an important direction for research, but Adam in this regime is too unstable to be reasonably approximated by smooth ODE solutions.

To answer your question about the meaning of "sufficiently close," the threshold for entering the spike regime varies with the task and, crucially, the learning rate. A sufficiently small learning rate prevents spikes altogether. Of course, we do not recommend making $ \beta $ and $ \rho $ equal. For instance, with Resnet-50 on CIFAR-10, using full-batch Adam without explicit regularization and a learning rate of $10^{-4}$, $(\beta, \rho) = (0.9, 0.999)$ enters the spike regime, whereas $ (0.94, 0.999) $ does not. We opt for higher learning rates to accelerate training and make the bias term (linear in $h$) more prounounced.

Following your suggestion, we have added more discussion of this topic in the "Numerical experiments" section.

Reviewer comment (Theory):
The concept of ADAM's modified loss is timely. However, it seems that the resulting modified loss doesn't explain the differences between traditional ADAM and SGD. For example, as the authors mentioned, ADAM often provides worse generalization performance and sharper solutions than SGD. Yet, in NLP tasks using Transformers, ADAM significantly outperforms SGD [1,2].

Author response:
Thank you very much for pointing out our overstatement regarding Adam's generalization compared to SGD. We acknowledge that Adam does not universally underperform and have revised such assertions. Additionally, we've included a brief mention in the "Future directions" section, noting that Adam often excels in training transformers, as supported by referenced studies.

The paper does not focus on whether the minima targeted by Adam's implicit bias term are advantageous or detrimental for specific tasks. For instance, Andriushchenko et al. (2023) observed that, in training transformers, sharper minima often correlate with lower test errors. As our bias term is novel, further research is necessary to understand its impact on transformer training dynamics.

Reference: Andriushchenko, M., Croce, F., Müller, M., Hein, M., \& Flammarion, N. (2023). A modern look at the relationship between sharpness and generalization. Proceedings of the 40th International Conference on Machine Learning (ICML'23), 202, Article 36, 840-902. JMLR.org.

Reviewer comment (Writing):
The writing of the paper seems subpar to me, and would benefit from being thoroughly proofread. In some locations it sounded very informal/colloquial, eg "which is ``eaten'' by the gradient".

Author response:
Thank you for your feedback on the style of our presentation. We have revised the sentence that ended with "which is `eaten' by the gradient" to ensure a more formal tone. We will also conduct a thorough proofreading of the entire document to identify instances where the writing is too informal.

Reviewer comment (Theory):
The analysis, though interesting, is also handwavy: see questions below.

Author response:
Thank you for your comment on the analysis. We included an "Informal summary" section at the beginning of our paper to make it more accessible and prevent readers from being overwhelmed by detailed mathematics. We would like to emphasize that the theorem and its proof are rigorous: the assumptions are precisely defined (as in Theorem 3.1 of the main paper), and the arguments are formal. Note that we do not use a common convention in pure mathematics to allow the value of constant $C$ to change from line to line (the statement and proof of Lemma SA-6.15 may illustrate this). In our latest revision, we have improved the connection between the theorem statement and the informal summary (see Remark 3.2).

Reviewer comment (Theory):
How do the authors characterize/show/test if "$\epsilon$ is very large compared to all squared gradient components"? What if it's smaller at the beginning, then becomes larger as the algorithm converges to an interpolating solution?

Author response:
You make a great point. In reality, for the majority of training, the situation tends to align more closely with the first extreme case, which we have now clarified in the revised version of our document. The purpose of outlining these extreme cases is to provide an intuitive understanding of the qualitative behavior. Indeed, at the onset of training, the value of $\varepsilon$ is likely to be smaller than the components. However, towards the end, particularly near an interpolating solution, $\varepsilon$ may become comparable to or even exceed these components. It is in this phase that Adam can essentially become gradient descent, regardless of other hyperparameter settings.

Reviewer comment (Theory):
How does penalizing the norm of the gradient lead to flat minima (Sec 4 discussion)? Since the gradient is 0 at optimum, don't $f$ and $f  + ||\nabla f||^2$ have the same set of minimizers? And doesn't this still hold when the 2 norm is replaced by any other norm?

Author response:
You are absolutely right in noting that the set of stationary points, where the gradient is zero, remains unchanged regardless of whether a norm is added or subtracted. However, this set can be broad, and varying parameter settings can lead different algorithms to converge to distinct points within this set. Figures 2 and 3 in our document illustrate this point, showing how Adam, with varying parameter settings, converges to different global minimizers. While adding a norm does not alter the location of the minima, it does impact the Hessian matrix. The Hessian plays a crucial role in determining the flatness of the landscape; changes to it affect this flatness. For instance, incorporating the 2-norm tends to penalize sharper regions in the optimization landscape, as highlighted in the works of Barrett \& Dherin (2021) and others.

Reference: Barrett, D., and Dherin, B. (2021). "Implicit Gradient Regularization." In International Conference on Learning Representations.

Reviewer comment (Evaluation):
Similarly, in the experiment, why is the perturbed 1 norm close to 0 at convergence? It seems the authors are performing early stopping, but that precisely means that implicit regularization is not happening, and that the model overfits.

Author response:
Thank you for the great point on the unusual nature of implicit regularization in discrete steps, such as in gradient descent. Contrary to typical cases, the implicit bias we observe here acts more like anti-regularization. You also make a great point about the perturbed 1-norm not appearing to be near zero in the figures. This discrepancy arises because, as noted in our "Notation" section, the perturbed norm is not technically a norm: at or near a stationary point, the square root of $ \varepsilon $ is multiplied by the dimension, which is in the tens of millions. In the final iterations, the gradient components themselves are indeed very small, making the perturbed norm close to its lower bound $ p \sqrt{\varepsilon} $, where $ p $ is the dimension. This might have given the impression of early stopping, but in reality, we are nearly fitting the training data exactly towards the end (with a training accuracy of 99.98%). We have made these points clearer in the latest revision. We also agree that one could interpret the model as overfitting due to the lack of both explicit and, as we argue, this particular form of implicit regularization, aligning with our analysis.

Reviewer comment (Presentation):
In the numerical illustrations, is it possible to display more than 3 curves/values of $h$ and $\beta$? In particular, for limiting values, why isn't the red cross attained if there is implicit bias?

Author response:
Thank you for the great question. In the latest revision, we have included more curves in numerical illustrations. In the toy problem section, attempting to widen the gap further at this learning rate causes instability in Adam, leading to oscillatory behavior around the global minima curve, a pattern similar to RMSProp's periodic solution for minimizing $x^2 / 2$ in Ma et al. (2022). Increasing the learning rate in Figure 3 causes lack of convergence. Even substantial explicit regularization doesn't lead Adam to the red cross without disrupting its trajectory. In scenarios without explicit regularization, akin to Barrett \& Dherin (2021) for GD and Ghosh et al. (2023) for momentum-GD (refer to Figure 1 in both), the red cross is not attained because of the relatively small magnitude of the implicit correction term.

Reference: Ghosh, A., Lyu, H., Zhang, X., and Wang, R. (2023). "Implicit Regularization in Heavy-Ball Momentum Accelerated Stochastic Gradient Descent." In The Eleventh International Conference on Learning Representations.

Reviewer comment (Evaluation):
The figures are not averaged across multiple runs to account for randomness.

Author response:
We haven't rerun the same experiments to smooth the curves, mainly because the sole source of randomness in our setup is the initialization seed. Our preliminary tests with various initialization seeds yielded visually indistinguishable results, suggesting that additional smoothing might not be effective. Moreover, adopting a Monte-Carlo approach for smoothing would significantly increase training time, likely by an order of magnitude. We have, however, added more experiments to demonstrate the variability in outcomes across different training sessions.

Reviewer comment (Theory):
Regarding the contribution summary, third bullet point — Why do the authors consider that Adam does not have an implicit bias, despite having a bias term in the backward analysis ODE? It seems to me that the meaning of "implicit regularization", e.g. in eq 1, is the same as "bias term" mentioned on page 2, but then the statement "Adam has no implicit regularization" is unclear.

Author response:
Thank you very much for pointing out the need for clearer wording here. We've revised the relevant section to convey that the implicit bias we observe typically functions as an "anti-regularization" for standard hyperparameters and during most of training. This bias essentially "penalizes" the negative norm, rather than the norm itself, as shown in Table 1. This clarification has been included in the latest revision.

Reviewer comment (Experiments):
Some transformer tasks might be helpful to see if we can see consistent behaviors in the $1-norm$ across different domains. If I'm not mistaken, Adam generalizes better than SGD in transformer related tasks which slightly contradicts the first conclusion in the discussion section.

Author response:
Thank you for highlighting our overly assertive statement regarding the generalization capabilities of adaptive gradient algorithms. A more accurate statement would be: "In the absence of carefully calibrated regularization, it has been observed that adaptive gradient methods may exhibit poorer generalization compared to non-adaptive optimizers." This amended statement aligns with the findings reported by Cohen et al., 2021, who also reference three other studies supporting this observation. We have amended our text accordingly.

We have launched additional experiments and added more evidence in the latest revision. We also plan to explore the training dynamics of transformers as a future research direction. Efforts are underway to incorporate a language translation task, although time constraints may limit our ability to do so.

Reference: Cohen, J. M., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z., Dahl, G. E., et al. (2022). "Adaptive Gradient Methods at the Edge of Stability." ArXiv:2207.14484.

Reviewer comment (Presentation):
Can the authors explain more about Figure 2 and Figure 3? I'm a bit confused about what these graphs are about and how we can see the change of the $1-norm$ from them.

Author response:
Thank you very much for pointing out the absence of axis labels in Figures 2 and 3. We have included them in the latest revision. To clarify, the horizontal axis represents $ \theta_1 $, while the vertical axis corresponds to $ \theta_2 $. The loss function, optimized by Adam, is defined as $ E(\theta_1, \theta_2) = (1.5 - 2\theta_1\theta_2)^2/2 $. The blue line in the figures marks the hyperbola of global minima of this loss, where $ \theta_1\theta_2 = 3/4 $. The gradient is expressed as $ (4\theta_1\theta_2 - 3) \times (\theta_2, \theta_1)^{T} $, which means its 1-norm is simply a factor of $
 4\theta_1\theta_2 - 3 $ times the 1-norm of the two-dimensional vector $ (\theta_2, \theta_1)^T $.

On the level sets of the form $4 \theta_1\theta_2 - 3 = c$, "lower 1-norm of the gradient" is equivalent to "lower 1-norm of $(\theta_2, \theta_1)^T$". The red cross in our illustrations marks the global minimum that corresponds to the lowest 1-norm of the parameter. Of course, on the entire hyperbola, the 1-norm of the gradient is zero. This is why we describe this point as a limit point of minimizers on the level sets (we clarify this in the latest revision). Thus, the more Adam drifts towards the red cross, the more it penalizes the 1-norm of the gradient.

Reviewer comment (Experiments):
Are the norms plotted in section 6 the norms in the final iterate of training?

Author response:
In the "Numerical experiments" section, we stopped training at near-perfect train accuracy. This decision was based on the clarity of the images at that stage, rendering further training unnecessary. Consequently, at the final iteration, the norms are close to their lower bound $p \sqrt{\varepsilon}$, limiting their informativeness. Therefore, this section primarily features plots of norms and accuracies mid-training. A more complete picture is offered by the loss and norm curves in Figure 6. We intend to include additional examples of these curves in the Appendix.

Reviewer comment (Experiments):
Is full batch required to observe the same behaviors of the norm and $\rho, \beta$ as in section 6? Can we do mini-batches instead?

Author response:
At present, our training utilizes full-batch Adam, due to our better understanding of its qualitative bias. Although our theorem introduces loss correction terms applicable to mini-batch Adam, we do not yet fully understand their implications. Gaining a deeper understanding of the qualitative behavior of these mini-batch correction terms represents a promising direction for future research.