On the Implicit Bias of Adam

Thank you so much for your comments. Each point has been addressed in our response below.

Weakness 1

Thank you for pointing this out. We will add more intuition in the introduction.

Weakness 2

Thank you for the great suggestions. We have incorporated a brief outline of the derivation and included a comment discussing the relationship between the informal summary and Theorem 3.1 (see Remark 3.2). Thank you for identifying the typo; it has been corrected in the revised version.

Question 1

Figure 1 illustrates what kind of function of the gradient is getting "penalized" by taking dimension $p = 1$. The bias term in Eq. (4), ignoring the $h / 2$ coefficient, is the derivative with respect to (one-dimensional) $\theta$ of the function $F(E'(\theta))$, where $F(x) := \int_0^x \bigl[ \frac{1 + \beta}{1 - \beta} - \frac{1 + \rho}{1 - \rho} + \frac{1 + \rho}{1 - \rho} \cdot \frac{\varepsilon}{y^2 + \varepsilon}\bigr] d \sqrt{\varepsilon + y^2}$ and $E'(\theta)$ is the derivative (= gradient) of the loss. Specifically, $2/h \cdot \mathrm{bias} = \frac{d}{d \theta} F(E'(\theta))$ represents what is being penalized, akin to how gradient descent penalizes the loss. For small $\varepsilon$, the penalty is effectively plus or minus the 1-norm of the gradient, as shown in the left picture of Figure 1, adjusted by a positive coefficient. As $\varepsilon$ increases, this term gradually shifts to resemble more closely the 2-norm of the gradient, again modified by a positive coefficient, reflecting Adam's convergence towards gradient descent.

Question 2

We believe that in practice $\varepsilon$ is small and, during most of training, the first extreme case describes the situation better. We do, however, also have an understanding of what happens if $\varepsilon$ is not small. In such cases, where $\rho > \beta$, the anti-regularization effect is less pronounced compared to regular Adam. This leads to more stable training, potentially improving generalization, particularly in situations where SGD is superior. Of course, this is contingent on adjusting the effective learning rate and $h$, as detailed in Example 2.3. Our hypothesis is that increasing $\varepsilon$ and modifying the learning rate can bring training dynamics closer to momentum-GD, in line with Choi et al. (2019).

Conversely, in language model training, where researchers sometimes fit the data almost exactly and continue training (a phase associated with "grokking"), the situation may hover between the two extremes, leaning towards the second. It is an interesting regime to study in the future, and perhaps the qualitative behavior of our bias term can contribute some insights.

Reference: Choi, D., et al. (2019). "On Empirical Comparisons of Optimizers for Deep Learning." ArXiv:1910.05446.

Thank you so much for your comprehensive response. We address each of the points you raised below.

Weakness 1

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.

Weakness 2

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.

Weakness 3

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.

We highly appreciate your feedback. We have carefully considered and responded to each specific point below.

Weakness 1

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.

Weakness 2

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

Question 1

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.

Question 2

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.

Question 3

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.

(continued below)

Question 4

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.

Question 5

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.

Question 6

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.

We greatly appreciate your response. Below, we have addressed each point in detail.

Weakness 1 and question 1

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.

Question 2

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.

Question 3

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.

Weakness 2

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.