On the Implicit Bias of Adam

Gradient descent (GD) can be seen as a numerical method solving the ordinary differential equation (ODE) , where is the loss function and is its gradient. Starting at , it creates a sequence of guesses , which lie close to the solution trajectory governed by the aforementioned ODE. Since the step size is finite, one could search for a modified differential equation such that is exactly zero, or at least closer to zero than , that is, all the guesses of the descent lie exactly on the new solution curve or closer compared to the original curve. This approach to analysing properties of a numerical method is sometimes called backward error analysis in the numerical integration literature (see Chapter IX in Ernst Hairer & Wanner (2006) and references therein).

Barrett & Dherin (2021) used this idea for full-batch GD and found that the modified loss function makes the trajectory of the solution to approximate the sequence one order of better than the original ODE, where is the Euclidean norm. In related work, Miyagawa (2022) obtained the correction term for full-batch GD up to any chosen order, also studying the global error (uniform in the iteration number) as opposed to the local (one-step) error.

The analysis was later extended to mini-batch GD in Smith et al. (2021). Assume that the training set is split into batches of size and there are batches per epoch (so the training set size is ). The cost function is rewritten as with mini-batch costs denoted . It was obtained in that work that after one epoch, the mean iterate of the algorithm, averaged over all possible shuffles of the batch indices, is close to the solution to , where the modified loss is given by .

Modified equations have also been derived for GD with heavy-ball momentum , where is the momentum parameter. In the full-batch setting, it turns out that for large enough it is close to the continuous trajectory solving

Versions of this general result were proven in Farazmand (2020), Kovachki & Stuart (2021), Ghosh et al. (2023) under different assumptions. The focus of the latter work is the closest to ours since they interpret the correction term as implicit regularization. Their main theorem also provides the analysis for the general mini-batch case.

In another recent work, Zhao et al. (2022) introduce a regularization term to the loss function as a way to ensure finding flatter minima, improving generalization. The only difference between their term and the first-order correction coming from backward error analysis (up to a coefficient) is that the norm is not squared and regularization is applied on a per-batch basis.

The application of backward error analysis for approximating the discrete dynamics of adaptive algorithms such as RMSProp (Tieleman et al., 2012) and Adam (Kingma & Ba, 2015) is currently missing in the literature. Barrett & Dherin (2021) note that “it would be interesting to use backward error analysis to calculate the modified loss and implicit regularization for other widely used optimizers such as momentum, Adam and RMSprop”. Smith et al. (2021) reiterate that they “anticipate that backward error analysis could also be used to clarify the role of finite learning rates in adaptive optimizers like Adam”. Ghosh et al. (2023) agree that “RMSProp … and Adam …, albeit being powerful alternatives to SGD with faster convergence rates, are far from well-understood in the aspect of implicit regularization”. In a similar context, in Appendix G to Miyagawa (2022) it is mentioned that “its [Adam’s] counter term and discretization error are open questions”.

This work fills the gap by conducting backward error analysis for (mini-batch, and full-batch as a special case) Adam and RMSProp. Our main contributions are listed below.

• •

  In Theorem 3.1, we provide a global second-order in continuous piecewise ODE approximation to Adam in the general mini-batch setting. (A similar result for RMSProp is moved to Appendix C.) For the full-batch special case, it was shown in prior work Ma et al. (2022) that the continuous-time limit of both these algorithms is a (perturbed by the numerical stability parameter ) signGD flow component-wise; we make this more precise by finding a linear in correction term on the right.

• •

  We analyze the full-batch case in the context of regularization (see the summary in Section 2). In contrast to the case of GD, where the two-norm of the loss gradient is implicitly penalized, Adam typically anti-penalizes the perturbed one-norm of the loss gradient (i. e., penalizes the negative norm), as specified in (5). Thus, the implicit bias of Adam that we identify serves as anti-regularization (except for the unusual case , large or very late at training).

• •

  We provide numerical evidence consistent with our theoretical results by training various vision models on CIFAR-10 using full-batch Adam. In particular, we observe that the stronger the implicit anti-regularization effect predicted by our theory, the worse the generalization. This pattern holds across different architectures: ResNets, simple convolutional neural networks (CNNs) and Vision Transformers. Thus, we propose a novel possible explanation for often-reported poor generalization of adaptive gradient algorithms. The code used for training the models is available at https://github.com/borshigida/implicit-bias-of-adam.

We are ready to describe our theoretical result (Theorem 3.1 below) in the full-batch special case. Assume is the loss, whose partial derivatives up to the fourth order are bounded. Let be iterations of Adam as defined in Definition 1.1. We find an ODE whose solution trajectory is -close to , meaning that for any time horizon there is a constant such that for any step size we have (for between and ). The ODE is written the following way (up to terms that rapidly go to zero as grows): for the component number

with initial conditions for all , where the correction term is

Depending on hyperparameters and the training stage, the correction term can take two extreme forms listed below. The reality is in between, but typically much closer to the first case.

• •

  If is small compared to all components of , i. e. , which is usually the case during most of the training, then we can write

  (4)

  For small , the perturbed one-norm is indistinguishable from the usual one-norm, and for it is penalized (in much the same way as the squared two-norm is implicitly penalized in the case of GD), but for the typical case its decrease is actually hindered by this term (so the bias is anti-regularization). The ODE in (2) approximately becomes

  (5)

• •

  If is large compared to all gradient components, i. e. (which may happen during the late learning stage, or if non-standard hyperparameter values are chosen), the fraction in (3) with in the numerator approaches one, the dependence on cancels out, and

  (6)

  In other words, becomes up to an additive constant, giving

  The form of the ODE in this case is

  (7)

These two extreme cases are summarized in Table 1. In Figure 1, we use the one-dimensional () case to illustrate what kind of term is being implicitly penalized.

Usually, is chosen to be small, and during most of the training Adam is much better described by the first extreme case. It is clear from (5) that, if , the correction term provides the opposite of regularization, in contrast to (1). The larger compared to , the stronger the anti-regularization effect is.

This finding may partially explain why adaptive gradient methods have been reported to generalize worse than non-adaptive ones (Chen et al., 2018; Wilson et al., 2017), as it offers a previously unknown perspective on why they are biased towards “higher-curvature” regions and find “sharper” minima. Indeed, note that standard (non-adaptive) -sharpness at can be defined by for some radius . This or similar definitions have been considered often in literature, see, e. g., Andriushchenko et al. (2023), Foret et al. (2021). Replacing the difference of the losses with its first-order approximation under the maximum (Foret et al., 2021; Ghosh et al., 2023)

we see that Adam typically anti-penalizes the approximation of -sharpness. Although the connection between sharpness and generalization is not clear-cut (Andriushchenko et al., 2023), our empirical results (Section 5) are consistent with this theory.

This overview also applies to RMSProp by setting ; see Theorem C.4 for the formal result.

Example 2.1(Backward Error Analysis for GD with Heavy-ball Momentum).

We only make one assumption, which is standard in the literature: the loss for each mini-batch is 4 times continuously differentiable, and partial derivatives of up to order 4 are bounded, i. e. there is a positive constant such that for in the region of interest

Theorem 3.1.

Remark 3.2.

Derivation sketch.

We now analyze the effect of the first-order term for Adam in the same model as Barrett & Dherin (2021) and Ghosh et al. (2023) have studied. Namely, assume the parameter is 2-dimensional, and the loss is given by . The loss is minimized on the hyperbola . We graph the trajectories of Adam in this case: the left part of Figure 2 shows that increasing forces the trajectory to the region with smaller , and increasing does the opposite. The right part shows that increasing the learning rate moves Adam towards the region with smaller if (just like in the case of GD, except the norm is different if is small compared to gradient components), and does the opposite if . All these observations are exactly what Theorem 3.1 predicts.

As a first sanity check, we train a relatively small fully-connected neural network with around parameters on the first 10,000 images of MNIST with full-batch Adam for 100 epochs and plot the value , i. e. the maximal weight difference between the Adam iteration and the piecewise ODE solution.¹¹1Since it makes little sense to numerically solve an ODE by further discretization, is estimated using the iteration (11) with ignored. Strictly speaking, this is not the trajectory obtained by the final backward error analysis step but rather the step immediately preceding it (after removing long-term memory but before converting the iteration to an ODE). We see in Figure 3 that even on this very large time horizon the trajectories are close in infinity-norm.

Further, we offer some preliminary empirical evidence that Adam (anti-)penalizes the perturbed one-norm of the gradients, as discussed in Section 2.

Ma et al. (2022) divide training regimes of Adam into three categories: the spike regime when is much larger than , in which the training loss curve contains very large spikes and the training is obviously unstable; the (stable) oscillation regime when is sufficiently close to , in which the loss curve contains fast and small oscillations; the divergence regime when is much larger than , in which Adam diverges. We exclude the last regime. In the spike regime, the loss spikes to large values at irregular intervals. This has also been observed in the context of large transformers, and mitigation strategies have been proposed in Chowdhery et al. (2022) and Molybog et al. (2023). Since it is unlikely that an unstable Adam trajectory can be meaningfully approximated by a smooth ODE solution, we reduce the incidence of large spikes by only considering and that are not too far apart, which is what Ma et al. (2022) recommend to do in practice.

We train Resnet-50, CNNs and Vision Transformers (Dosovitskiy et al., 2020) on the CIFAR-10 dataset with full-batch Adam. In this section, we provide the results for Resnet-50; the pictures for CNNs and Transformers are similar and are given in Section H.4. Figure 4 shows that in the stable oscillation regime increasing appears to increase the perturbed one-norm (consistent with our analysis: the smaller , the more this “norm” is penalized) and decrease the test accuracy. Figure 5 shows that increasing appears to decrease the perturbed one-norm (consistent with our analysis: the larger , the more this norm is penalized) and increase the test accuracy. The picture confirms the finding in Ghosh et al. (2023) (for the momentum parameter in momentum GD).

Figure 6 shows the graphs of as functions of the epoch number. The “norm” decreases, then rises again, and then decreases further until it flatlines.²²2Note that the perturbed one-norm cannot be near-zero at the end of training because it is bounded from below by . Throughout most of the training, the larger the smaller the “norm”. The “hills” of the “norm” curves are higher with smaller and larger . This is consistent with our analysis because the larger compared to , the more is prevented from falling by the correction term.

As far as we know, the assumption similar to (8) is explicitly or implicitly present in all previous work on backward error analysis of gradient-based machine learning algorithms. (Recently, Beneventano (2023) weakened this assumption for SGD without replacement, but their focus is somewhat different.) There is evidence that large-batch algorithms often operate near or at the edge of stability (Cohen et al., 2021, 2022), in which the largest eigenvalue of the hessian can be large, making it unclear whether the higher-order partial derivatives can safely be assumed bounded near optimality. In addition, as Smith et al. (2021) point out, in the mini-batch setting backward error analysis can be more accurate. We leave a qualitative analysis of the behavior of first-order terms in Theorem 3.1 in the mini-batch case as a future direction.

Relatedly, Adam does not always generalize worse than SGD: for transformers, Adam often outperforms (Zhang et al., 2020; Kumar et al., 2022). Moreover, for NLP tasks a long time can be spent training close to an interpolating solution. Our analysis suggests that in the latter regime the anti-regularization effect disappears, which does indeed confirm the finding that generalization can be better. However, we believe this explanation is not complete, and more work is needed to connect the implicit bias to the training dynamics of transformers.

In addition, the constant in Theorem 3.1 goes to infinity as goes to zero. Theoretically, our proof does not exclude the case where for very small the trajectory of the piecewise ODE is only close to the Adam trajectory for small, suboptimal learning rates, at least at later stages of learning. (For the initial learning period, this is not a problem.) It appears to also be true of Proposition 1 in Ma et al. (2022) (zeroth-order approximation by sign-GD). This is especially noticeable in the large-spike regime of training (see Section 5) which, despite being obviously unstable, can still lead to acceptable test errors. It would be worthwhile to investigate this regime in detail.

We extend our special thanks to Boris Hanin, Sam Smith, and the anonymous reviewers for their insightful comments and suggestions that greatly enhanced this work. We are also grateful to Jianqing Fan, Pier Beneventano, and Rae Yu for engaging and productive discussions. Cattaneo gratefully acknowledges financial support from the National Science Foundation through DMS-2210561 and SES-2241575. Klusowski gratefully acknowledges financial support from the National Science Foundation through CAREER DMS-2239448. Additionally, we acknowledge the Princeton Research Computing resources, coordinated by the Princeton Institute for Computational Science and Engineering (PICSciE) and the Office of Information Technology’s Research Computing.

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

The appendix provide some omitted details and proofs.

We consider two algorithms: RMSProp and Adam, and two versions of each algorithm (with the numerical stability parameter inside and outside of the square root in the denominator). This means there are four main theorems: Theorem B.4, Theorem C.4, Theorem D.4 and Theorem E.4, each residing in the section completely devoted to one algorithm. The simple induction argument taken from Ghosh et al. (2023), essentially the same for each of these theorems, is based on an auxiliary result whose corresponding versions are Theorem B.3, Theorem C.3, Theorem D.3 and Theorem E.3. The proof of this result is also elementary but long, and it is done by a series of lemmas in Appendix F and Appendix G. Out of these four, we only prove Theorem B.3 since the other three results are proven in the same way with obvious changes.

Appendix H contains some details about the numerical experiments.

Definition B.1.

Assumption B.2.

Theorem B.3(RMSProp with outside: local error bound).

The proof of Theorem B.3 is conceptually simple but very technical, and we delay it until Appendix G. For now assuming it as given and combining it with a simple induction argument gives a global error bound which follows.

Theorem B.4(RMSProp with outside: global error bound).

Proof.

Definition C.1.

Assumption C.2.

Theorem C.3(RMSProp with inside: local error bound).

The argument is the same as for Theorem B.3.

Theorem C.4(RMSProp with inside: global error bound).

The argument is the same as for Theorem B.4.

Definition D.1.

Assumption D.2.

Theorem D.3(Adam with outside: local error bound).

The argument is the same as for Theorem B.3.

Theorem D.4(Adam with outside: global error bound).

Proof.

Definition E.1.

Assumption E.2.

Theorem E.3(Adam with inside: local error bound).

The argument is the same as for Theorem B.3.

Theorem E.4(Adam with inside: global error bound).

The argument is the same as for Theorem D.4.

Our first goal is to argue that the first derivative of is uniformly bounded in absolute value. To achieve this, we just need to bound all the terms on the right-hand side of the ODE (13).

Lemma F.1.

Proof of Lemma F.1.

Lemma F.2.

Proof.

Our next goal is to argue that the second derivative of is bounded in absolute value. For this, we need to bound the first derivatives of all the three additive terms on the right-hand side of (13).

Lemma F.3.

Proof of Lemma F.3.

Lemma F.4.

Proof.

Finally, we need to argue that the third derivative of is bounded in absolute value. To achieve this, we need to bound the second derivatives of the terms on the right-hand side of (13).

Lemma F.5.

Proof of Lemma F.5.

Lemma F.6.

Proof.

Our next objective is proving and identifying the constant in the equality

We will make some preparations and achieve this objective in Lemma G.5. Then we will conclude the proof of Theorem B.3.

Lemma G.1.

Proof.

Lemma G.2.

Proof.

Lemma G.3.

Proof.

Lemma G.4.

Proof.

Lemma G.5.

Proof.

We are finally ready to prove Theorem B.3.

Proof of Theorem B.3.

Our goal is to find such a trajectory that

Result I.1.

Derivation.

Result I.2.

Derivation.