# Adaptive Gradient Methods at the Edge of Stability

Jeremy Cohen Behrooz Ghorbani Shankar Krishnan Naman Agarwal & Sourabh
Medapati Michal Badura Daniel Suo David Cardoze & Zachary Nado George E. Dahl
Justin Gilmer

Google Research, Brain Team Work done as student researcher. Correspondence to
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Very little is known about the training dynamics of adaptive gradient methods
like Adam in deep learning. In this paper, we shed light on the behavior of
these algorithms in the full-batch and sufficiently large batch settings.
Specifically, we empirically demonstrate that during full-batch training, the
maximum eigenvalue of the _preconditioned_ Hessian typically equilibrates at a
certain numerical value — the stability threshold of a gradient descent
algorithm. For Adam with step size $\eta$ and $\beta_{1}=0.9$, this stability
threshold is $38/\eta$. Similar effects occur during minibatch training,
especially as the batch size grows. Yet, even though adaptive methods train at
the “Adaptive Edge of Stability” (AEoS), their behavior in this regime differs
in a significant way from that of non-adaptive methods at the EoS. Whereas
non-adaptive algorithms at the EoS are blocked from entering high-curvature
regions of the loss landscape, adaptive gradient methods at the AEoS keep
advancing into high-curvature regions, while adapting the preconditioner to
compensate. Our findings can serve as a foundation for the community’s future
understanding of adaptive gradient methods in deep learning.

Figure 1: Full-batch Adam trains at the Adaptive Edge of Stability (AEoS). We
train a fully-connected network on CIFAR-10 using full-batch Adam with various
hyperparameters $\eta,\beta_{1},\beta_{2}$. Observe that the maximum
eigenvalue of the preconditioned Hessian equilibrates at the numerical value
$\frac{(2+2\beta_{1})}{(1-\beta_{1})\eta}$, which is drawn as a dashed
horizontal line. However, in contrast to the non-adaptive EoS, the maximum
eigenvalue of the raw Hessian usually keeps rising at the AEoS (see Figure 8).

##  1 Introduction

Neural networks are often trained using adaptive gradient methods such as Adam
[23], rmsprop [39], or Adafactor [38]. These algorithms are variants of
preconditioned gradient descent where the preconditioner is constantly
_adapting_ to gradients from recent steps. Despite the popularity of these
algorithms, their behavior when training neural networks remains poorly
understood.

Meanwhile, a line of recent work [34, 43, 20, 21, 8, 15] has shed light on the
neural network training dynamics of _non-adaptive_ gradient descent (both with
and without momentum) in, at least, the full-batch and (sufficiently) large-
batch regimes. These works have empirically demonstrated that _dynamical
instability_ plays a key role in the training process. When training neural
networks, gradient descent is constantly attracted to regions of parameter
space with increasingly high curvature [20, 21, 8]; yet, up to a quadratic
Taylor approximation, gradient descent is an unstable dynamical system in
regions where the curvature in any direction exceeds a certain threshold — the
optimizer cannot linger for long in any such region without being expelled
[34, 43, 21]. How is this tension resolved? In the full-batch special case,
gradient descent spends the bulk of training in a regime called the Edge of
Stability (EoS) [8] in which the _sharpness_ — the maximum eigenvalue of the
training Hessian — hovers right at, or just above, the stability threshold. At
the EoS, gradient descent would still be moving into regions of higher
curvature were it not being constantly repelled from these high-curvature
regions by unstable dynamics. As we confirm below, these findings also
generalize to preconditioned gradient descent (with a static preconditioner).

However, it has not been clear whether these findings have relevance for
_adaptive_ gradient methods. Because of adaptive preconditioning, adaptive
gradient methods do not evolve as linear recurrences on the local quadratic
Taylor approximation, and thus it is not clear why their local stability would
be well-modeled by an eigenvalue condition.

In this paper, we demonstrate that the EoS phenomenon does, in fact, carry
over to the adaptive setting. Our key empirical finding is that throughout
training, the short-term stability behavior of, say, Adam is well-approximated
by that of “frozen Adam” — a version of Adam in which the preconditioner is
frozen at its current value. On the local quadratic Taylor approximation,
“frozen Adam” _does_ evolve as a linear recurrence, and is unstable whenever
the maximum eigenvalue of the preconditioned Hessian (the _preconditioned
sharpness_) exceeds the stability threshold of EMA-style heavy ball momentum,
which is $\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}$. Indeed, we observe that
during full-batch training by real Adam, the preconditioned sharpness
equilibrates at this precise numerical value (Figure 1). During minibatch
training, similar phenomena occur, paralleling the situation with non-adaptive
optimizers.

However, even though adaptive gradient methods train at the “Adaptive Edge of
Stability” (AEoS), their behavior in this regime differs in a significant way
from that of non-adaptive methods in the classical EoS regime: whereas non-
adaptive optimizers in the classical EoS regime are blocked from accessing
high-curvature regions of the loss landscape, we find that adaptive gradient
methods at the AEoS can and do enter these high-curvature regions via their
ability to adapt the preconditioner. This is especially liable to occur if the
step size is small or the preconditioner decay factor (e.g. Adam’s
$\beta_{2}$) is small. Thus, adaptive gradient methods sometimes lack the
implicit inductive bias [35] that blocks non-adaptive methods from converging
to high-curvature solutions.

##  2 Notation and Background

Optimization algorithms aim to minimize an objective function
$f:\mathbb{R}^{p}\to\mathbb{R}$ by producing a sequence of iterates
$\\{\mathbf{x}_{t}\\}\subset\mathbb{R}^{p}$. We denote an algorithm “alg” with
hyperparameters $a,b$ as $\textsc{Alg}(a,b)$. We write
$\mathbf{H}(\mathbf{x})$ for the Hessian at parameter $\mathbf{x}$, and
sometimes write $\mathbf{H}_{t}$ as an abbreviation for
$\mathbf{H}(\mathbf{x}_{t})$, the Hessian at step $t$. We denote the largest
positive eigenvalue of a matrix $\mathbf{H}$ as $\lambda_{1}(\mathbf{H})$. We
call the largest positive eigenvalue of a Hessian the _sharpness_ , and the
largest positive eigenvalue of a preconditioned Hessian the _preconditioned
sharpness_.

#### Gradient descent and momentum

The most basic optimization algorithm is $\textsc{GD}(\eta)$, defined by the
parameter update $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\,\mathbf{g}_{t+1}$,
where $\mathbf{g}_{t+1}$ is a full-batch or minibatch gradient of $f$ computed
at $\mathbf{x}_{t}$. Vanilla gradient descent can be accelerated by the use of
momentum, in which case the parameter update becomes
$\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\,\mathbf{m}_{t+1}$, where
$\mathbf{m}_{t+1}$ is a momentum vector. Momentum comes in two popular
flavors: heavy ball (HB) and Nesterov. Both of these, in turn, can be
parameterized in the “standard” fashion or the “EMA” (exponential moving
average) fashion. The momentum update for
$\textsc{StandardHB}(\eta,\beta_{1})$ is
$\mathbf{m}_{t+1}=\beta_{1}\,\mathbf{m}_{t}+\mathbf{g}_{t+1}$, whereas the
update for $\textsc{EmaHB}(\eta,\beta_{1})$ is
$\mathbf{m}_{t+1}=\beta_{1}\mathbf{m}_{t}+(1-\beta_{1})\mathbf{g}_{t+1}$.

Figure 2: We optimize the quadratic objective $f(x)=\frac{1}{2}x^{2}$ using
$\textsc{EmaHB}(\eta,0.9)$ at various $\eta$. Observe that learning rates
$\eta$ above 38 diverge.

Each of these algorithms, when run on any quadratic function (and in
particular, when run on the local quadratic Taylor approximation to a neural
training objective), evolves independently along each Hessian eigenvector as a
linear recurrence relation (Appendix A). If any Hessian eigenvalue exceeds a
certain algorithm-dependent threshold, then the linear recurrence relation for
the corresponding eigenvector is _unstable_ and diverges exponentially; this
implies that the iterate will oscillate with exponentially increasing
magnitude along this eigenvector. The stability thresholds are: $2/\eta$ for
$\textsc{GD}(\eta)$, $\frac{2+2\beta_{1}}{\eta}$ for
$\textsc{StandardHB}(\eta,\beta_{1})$, and
$\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}$ for $\textsc{EmaHB}(\eta,\beta_{1})$.

We illustrate this threshold behavior in Figure 2, where we optimize the one-
dimensional quadratic function $f(x)=\frac{1}{2}x^{2}$ using EMA-style heavy
ball momentum with $\beta_{1}=0.9$ and varying learning rates. Observe that
the iterates diverge whenever the learning rate $\eta$ exceeds
$\frac{2+2\beta_{1}}{(1-\beta_{1})}=38$.

#### Preconditioned gradient descent

Gradient descent, or its momentum variants, can be preconditioned using a
(static) preconditioner $\mathbf{P}$. In this case, the update rule becomes
$\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\,\mathbf{P}^{-1}\,\mathbf{m}_{t+1}$,
where $\mathbf{m}_{t}$ is updated according to one of the momentum rules
described above. Often, $\mathbf{P}$ is a diagonal matrix, in which case
$\mathbf{P}^{-1}$ can be interpreted as per-parameter learning rates. A
special case of particular interest in this paper is preconditioned gradient
descent with EMA-style heavy ball momentum, which we denote as
$\textsc{PrecondEmaHB}(\eta,\beta_{1},\mathbf{P})$. On a quadratic objective
with Hessian $\mathbf{H}$, preconditioned gradient descent will diverge if the
maximum eigenvalue of the _preconditioned Hessian_
$\mathbf{P}^{-1/2}\mathbf{H}\mathbf{P}^{-1/2}$ (the _preconditioned
sharpness_) exceeds the stability threshold of the non-preconditioned
algorithm. In that event, the component of the iterate that is aligned with
$(\mathbf{P}^{1/2})^{T}\mathbf{v}_{1}$, where $\mathbf{v}_{1}$ is the top
eigenvector of $\mathbf{P}^{-1/2}\mathbf{H}\mathbf{P}^{-1/2}$, will oscillate
with exponentially increasing magnitude, and diverge. Note that
$\mathbf{P}^{-1/2}\mathbf{H}\mathbf{P}^{-1/2}$ shares eigenvalues with the
similar matrix $\mathbf{P}^{-1}\mathbf{H}$, so the preconditioned sharpness
can be written as both
$\lambda_{1}(\mathbf{P}^{-1/2}\mathbf{H}\mathbf{P}^{-1/2})$ or
$\lambda_{1}(\mathbf{P}^{-1}\mathbf{H})$.

#### Adaptive gradient methods

_Adaptive_ gradient methods are variants of preconditioned gradient descent in
which the preconditioner is changing rather than fixed, i.e. the parameter
update is
$\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\mathbf{P}_{t+1}^{-1}\mathbf{m}_{t+1}$,
where $\mathbf{P}_{t+1}$ is the latest preconditioner. For example,
$\textsc{Rmsprop}(\eta,\beta_{2},\epsilon)$ [39] updates its preconditioner
according to:

$\displaystyle\boldsymbol{\nu}_{t+1}=\beta_{2}\boldsymbol{\nu}_{t}+(1-\beta_{2})\mathbf{g}_{t+1}^{\circ
2}\quad\quad\mathbf{P}_{t+1}=\text{diag}(\boldsymbol{\nu}_{t+1}^{1/2})+\epsilon\,\mathbf{I}.$
(1)

A more popular algorithm is $\textsc{Adam}(\eta,\beta_{1},\beta_{2},\epsilon)$
[23], which adds EMA-style heavy ball momentum to rmsprop. Adam has an
optional “bias correction” scheme. With no bias correction, Adam employs the
rmsprop preconditioner rule in Eq (1); with bias correction, the Adam
preconditioner is:

$\displaystyle\mathbf{P}_{t+1}$
$\displaystyle=\left(\frac{1}{1-\beta^{t+1}_{1}}\right)\;\left[\text{diag}\left(\sqrt{\frac{\boldsymbol{\nu}_{t+1}}{1-\beta_{2}^{t+1}}}\right)+\epsilon\,\mathbf{I}\right].$
(2)

In contrast to gradient descent (and preconditioned gradient descent),
adaptive gradient methods do not evolve as linear recurrence relations on
quadratic functions. Thus, it is a priori unclear whether their local
stability can be modeled using an eigenvalue condition.

##  3 Related Work

#### Training dynamics of non-adaptive gradient descent

Recent empirical studies [21, 25, 8, 24] have shed substantial light on the
training dynamics of non-adaptive gradient descent in deep learning. When
training neural networks, gradient descent tends to continually move in a
direction that increases the sharpness [20, 21, 8]. This phemonemon was dubbed
_progressive sharpening_ in [8], and remains poorly understood. In the special
case of _full-batch_ gradient descent, the sharpness usually rises past the
optimizer’s stability threshold, at which point (the “breakeven point” [21])
the optimizer becomes destabilized by exponentially growing movement along the
Hessian’s top eigenvector. On quadratic objective functions, this behavior
would lead to divergence. However, neural network training objectives are not
quadratic, and gradient descent typically does not diverge; instead, it enters
a regime called the Edge of Stability (EoS) [8] in which the sharpness hovers
just above, or oscillates around, the stability threshold. At the EoS,
gradient descent is constantly being repelled from regions of the loss
landscape with sharpness exceeding the stability threshold. Prior work has not
discussed preconditioned gradient descent, but in the next section we confirm
that the results of [8] carry over to the non-adaptive preconditioned setting.
It remains unclear how gradient descent is able to safely train at the EoS
without diverging [2, 32], though [32] has suggested that “subquadratic
growth” of the training objective may play a role. Note that these results
only apply to _full-batch_ gradient descent. In the more general case of SGD,
similar effects seem to occur [20, 21], as we review in §5.

The fact that non-adaptive gradient descent is blocked from entering sharp (as
quantified by maximum Hessian eigenvalue) regions of the loss landscape
constitutes one implicit bias [35] of non-adaptive gradient descent. It is
plausible that this implicit bias could impact generalization (e.g. see [33,
19]).

#### Understanding adaptive gradient methods

There have been a number of convergence analyses of adaptive gradient methods
[7, 48, 26, 41, 44, 11], none of which model the behavior that we observe
here. Beyond formal convergence analyses, [4, 10] proposed to model Adam as an
system of ordinary differential equations in continuous time; this approach
also cannot explain the unstable dynamics that we observe. [3] argued that
Adam should be viewed as a variant of sign gradient descent. Most relevant to
our paper, [31] conducted a qualitative study of full-batch Adam during neural
network training. They observed small bumps and large spikes in the training
loss curve; our paper explains this phenomenon as arising from dynamical
instability.

#### Implicit bias of adaptive gradient methods

Absent well-tuned regularization, adaptive gradient methods have been reported
to generalize worse than non-adaptive optimizers [42, 22, 28]. However, it has
been far from clear why, in deep learning, adaptive optimizers should even
find consistently different solutions than non-adaptive optimizers (let alone
why these different solutions should generalize worse). [42] constructed a
synthetic task where Adam provably generalizes worse than gradient descent.
However, synthetic tasks also exists where the reverse is true [1]. [40]
proved that rmsprop on homogenous neural networks is, like gradient descent
[30], implicitly biased towards maximum margin solutions. [18] argued that
adaptive methods take steps that are too large in small-Hessian-eigenvalue
directions, which are more likely to be noise than signal.

Our paper gives evidence for a different implicit bias: adaptive gradient
methods are liable to find higher-curvature solutions than non-adaptive
algorithms, since whereas non-adaptive algorithms are blocked from high-
curvature regions, adaptive optimizers can evade this restriction. [43]
briefly speculated that this might be the case, but ran no experiments with
adaptive optimizers (only L-BFGS). Meanwhile, [17] empirically observed that a
VGG trained using an adaptive optimizer had 40x the sharpness than when
trained using SGD, and [6] found that MLP-mixers and ViTs trained with Adam
converged to far sharper solutions than a ResNet trained with SGD.

[49] and [45] previously argued theoretically that SGD can more quickly escape
certain “sharp minima” than Adam, but using a different definition of
sharpness, and a mechanism tied inextricably to stochasticity. In contrast, we
demonstrate that even full-batch Adam converges to higher-curvature solutions
than full-batch gradient descent, implying that stochasticity is not crucial.

##  4 Full-batch adaptive optimizers train at the Adaptive Edge of Stability

Figure 3: “Frozen Adam” (preconditioned momentum) trains at the Edge of
Stability. We train a fully-connected network on CIFAR-10 using “frozen Adam,”
i.e. preconditioned gradient descent with EMA-style heavy ball momentum and a
fixed preconditioner $\mathbf{P}$. Consistent with [8], the preconditioned
sharpness $\lambda_{1}(\mathbf{P}^{-1}\,\mathbf{H}_{t})$ rises until
equilibrating at the stability threshold of $38/\eta$. The raw sharpness
$\lambda_{1}(\mathbf{H}_{t})$ mostly ceases to increase at the EoS.

#### Warm-up: “frozen Adam”

Before turning to adaptive preconditioning, we first confirm that prior work
[8] generalizes to the setting of preconditioned gradient descent. We train a
network using full-batch Adam (no bias correction), and extract the second
moment accumulator $\boldsymbol{\nu}_{t}$ at some midpoint step $t$ during
training. We then train the same architecture from scratch at a range of
learning rates using full-batch “frozen Adam,” i.e. Adam where the second-
moment accumulator is held fixed at that particular $\boldsymbol{\nu}_{t}$
rather than updated via Eq (1) or Eq (2). Note that “frozen Adam” is
equivalent to $\textsc{PrecondEmaHB}(\eta,\beta_{1},\mathbf{P})$ where
$\mathbf{P}=\text{diag}(\boldsymbol{\nu}_{t}^{1/2})+\epsilon\mathbf{I}$. Thus
frozen Adam is dynamically unstable on the training objective’s quadratic
Taylor approximation whenever the preconditioned sharpness
$\lambda_{1}(\mathbf{P}^{-1}\,\mathbf{H}_{t})$ exceeds the stability threshold
for EMA-style heavy ball momentum. This threshold is
$\frac{2+2\beta_{1}}{(1-\beta_{1})\eta}=\frac{38}{\eta}$ for our choice of
$\beta_{1}=0.9$. Assuming that the findings of [8] generalize from non-
preconditioned to preconditioned gradient descent, one would expect that
during full-batch training by frozen Adam, the preconditioned sharpness would
rise until plateauing at that threshold value. In Figure 3(b), we confirm that
this occurs. For each of several learning rates $\eta$, we plot the evolution
of the preconditioned sharpness $\lambda_{1}(\mathbf{P}^{-1}\mathbf{H}_{t})$
while visualizing the stability threshold of $38/\eta$ as a horizontal dashed
line. Observe that the preconditioned sharpness rises until plateauing at the
dashed line. We also plot the evolution of the “raw” sharpness
$\lambda_{1}(\mathbf{H}_{t})$ in Figure 3(c). For this network, the raw
sharpness largely ceases to increase once training enters the EoS.

#### Adaptive preconditioning: Adam

Adam is a more complex dynamical system than “frozen Adam,” because its
preconditioner evolves in respose to recent gradients; consequently, its
dynamics on the local quadratic Taylor approximation does not reduce to a
linear recurrence. Throughout physics and engineering, a time-tested way to
understand complex systems is to approximate them by simpler ones that still
capture some behavior of interest. Consider Adam at any particular step
$t_{0}$, and let $\mathbf{P}_{t_{0}}$ be Adam’s preconditioner at that step.
Over short time horizons $t>t_{0}$, one can approximate Adam by
$\textsc{PrecondEmaHB}(\eta,\beta_{1},\mathbf{P}_{t_{0}})$, i.e. “frozen Adam”
where the preconditioner is frozen at $\mathbf{P}_{t_{0}}$. On the local
quadratic Taylor approximation, this simpler algorithm is unstable if the
preconditioned sharpness
$\lambda_{1}(\mathbf{P}_{t_{0}}^{-1}\mathbf{H}_{t_{0}})$ exceeds the stability
threshold of heavy-ball EMA momentum, which is
$\frac{2+2\beta_{1}}{(1-\beta_{1})\eta}$. Of course, whether this
approximation is useful is an empirical question. And empirically, we now show
that Adam has similar stability properties as its frozen counterpart.

In Figure 1, we train a fully-connected network on the full CIFAR-10 dataset
using full-batch Adam (no bias correction). In Figure 1(a), we vary the
learning rate $\eta$ while fixing $\beta_{2}=0.999$ and $\beta_{1}=0.9$.
Observe that for every learning rate $\eta$, the preconditioned sharpness
$\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}_{t})$ equilibrates right around the
stability threshold of frozen Adam. In Figure 1(b), we vary $\beta_{1}$ as we
fix $\eta=0.0001$ and $\beta_{2}=0.999$. Observe that the same phenomenon
occurs. Finally, in Figure 1(c), we vary $\beta_{2}\in\\{0.9,0.99,0.999\\}$ as
we fix $\beta_{1}=0.9$ and $\eta\in\\{$5e-5, 4e-4$\\}$. The same phenomenon
occurs, even when $\beta_{2}$ takes the relatively small value of 0.9, which
is when one might expect Adam’s behavior to differ the most from frozen
Adam’s, because the preconditioner adapts faster. Thus, we see that Adam
becomes unstable whenever its frozen counterpart would become unstable.
Nevertheless, in §6, we will see that Adam’s behavior at the AEoS differs in
other respects from that of frozen Adam at the EoS.

#### Other architectures

In Figure 4, we verify that the full-batch phenomenon generalizes to other
computer vision architectures. In particular, we consider (a) a CNN on
CIFAR-10; (b) an un-normalized Wide ResNet (WRN) [47] on CIFAR-10; and (c) a
(batch-normalized) WRN on CIFAR-100. Furthermore, Figure 7 in the subsequent
section verifies that our _minibatch_ findings apply to transformers on WMT
machine translation.

Figure 4: The phenomenon generalizes to other architectures. We train three
vision architectures using full-batch Adam with $\beta_{1}=0.9$ and
$\beta_{2}=0.999$ at a range of learning rates (colors). In each case, the
preconditioned sharpness equilibrates at the threshold $38/\eta$. Each network
was trained until either reaching a milestone train loss value, or until
reaching a step limit.

#### Other full-batch adaptive gradient methods

In Figure 5, we train the same fully-connected network on CIFAR-10 using eight
adaptive gradient algorithms in the full-batch regime: (a) Adam with bias
correction, which essentially decreases the preconditioner on a schedule in
the early stage of training; (b) AdamW [29], which employs decoupled weight
decay; (c) Adafactor [38], which maintains a factored layer-wise
preconditioner; (d) Amsgrad [37], which ensures that the preconditioner is
entrywise non-decreasing over time; (e) Padam [5], which employs an exponent
smaller than $\frac{1}{2}$ (we use $0.25$) in the definition of the
preconditioner Eq (1); (f) Nadam [12], which uses Nesterov rather than heavy
ball momentum; (g) rmsprop [39], which lacks momentum; and finally (h) Adagrad
[13], which preconditions using the sum of past squared gradients, rather than
an exponential moving average. In all cases, we observe that the
preconditioned sharpness $\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}_{t})$
equilibrates around the appropriate stability threshold (written in
parentheses).

Figure 5: Other adaptive gradient algorithms train at the AEoS. We train a FC
network on CIFAR-10 using eight adaptive optimizers in full-batch mode. We
train each algorithm at five learning rates (colors). In each case, the
preconditioned sharpness equilibrates at, or just above, the stability
threshold (written in parentheses). Note that the qualitative behavior of the
preconditioned sharpness depends on the presence and type of momentum; see
Appendix C.

#### A corner case

When running full-batch Adam at extremely small learning rates, we sometimes
observe that the preconditioned sharpness equilibrates _short of_ the
stability threshold. We elaborate on this corner case, and offer an
explanation, in Appendix D.

##  5 The minibatch setting

We now move beyond the full-batch setting to the more general setting of
minibatch training. In the case of gradient descent and minibatch SGD, it is
clear from prior work [20, 21] that during minibatch training, the sharpness
is subject to similar effects as during full-batch training. For one, provided
that training is successful, the sharpness never ventures more than a bit
above the stability threshold of the corresponding full-batch algorithm [21].
This can be explained by the fact that SGD is unstable in expectation whenever
GD is unstable [14]. Yet additional factors seem to also be at play in the
case of SGD. Specifically, it has been observed that at small batch sizes
(where there is more gradient noise), the mid-training sharpness is smaller.
One hypothesized explanation [43, 21] is that in the presence of gradient
noise, SGD becomes unstable at lower sharpnesses.

Figure 6: During minibatch Adam, the preconditioned sharpness behaves
analogous to the sharpness during minibatch SGD. We train a Resnet-50 on
ImageNet using minibatch Adam. The preconditioned sharpness is (1) below the
full-batch stability threshold (pictured as a horizontal line), (2) smaller
when the batch size is smaller, and (3) smaller when the learning rate is
larger.

We now demonstrate that the behavior of the preconditioned sharpness during
minibatch Adam parallels that of the sharpness during minibatch SGD. Namely,
we observe that during minibatch Adam, the preconditioned sharpness (1) never
rises more than a bit beyond the stability threshold of the full-batch
algorithm, provided that training succeeds; (2) tends to be smaller, mid-
training, when the learning rate is large; and (3) tends to be smaller, mid-
training, when the batch size is small.

In Figure 6, we train a Resnet-50 on ImageNet at three batch sizes. Observe
that for the largest batch size considered, the preconditioned sharpness
behaves similar to the full-batch setting (almost rising to the full-batch
stability threshold); yet, as the batch sizes decreases, the trajectory of the
preconditioned sharpness shifts downwards.

In Figure 7, we train both a pre-LayerNorm (Pre-LN) and a post-LayerNorm
(Post-LN) Transformer on the En->De WMT translation task. Following standard
practice [36], we train using a linear warmup schedule lasting 40k steps. We
train at both a smaller batch size of 1024 and a larger batch size of 10240.
The dotted line depicts the stability threshold of $38/\eta$; this threshold
is constantly decreasing as the learning rate is warmed up. Observe that the
preconditioned sharpness initially rises until reaching either the stability
threshold (larger batch size), or a value just short of the stability
threshold (smaller batch size). Then, as learning rate warmup continues, the
preconditioned sharpness decreases to track the stability threshold. In
effect, the preconditioned sharpness is being “pushed down” by Adam’s learning
rate warmup. Note that [15] previously showed that a warmup schedule for
momentum SGD likewise “pushed down” the (unpreconditioned) sharpness in
Transformer training. The Post-LN Transformer is known to be less stable than
the Pre-LN variant [46, 27]; this instability can be seen in the spike in
$\lambda_{1}(\mathbf{P}^{-1}\mathbf{H})$ (and training loss) towards the end
of the warmup period.

Figure 7: Learning rate warmup gradually reduces the preconditioned sharpness
during optimization. Left: The evolution of the preconditioned sharpness for a
6L-6L encoder-decoder Pre-LN transformer trained on the WMT En=>De task at
$\eta=.001$ with a linear warmup period of 40000 steps. Right: Same as left,
but with a PostLN Transformer. In both cases the preconditioned curvature
closely tracks the $38/\eta$ bound during warmup, however there is a
noticeable gap at the smaller batch size. The PostLN Transformer training
fails late in the warmmup period.

##  6 A closer look at the AEoS

In this section, we take a closer look at Adam’s behavior at the AEoS. We will
see that this behavior sometimes differs substantially from that of non-
adaptive optimizers. In particular, whereas non-adaptive optimizers at the EoS
are blocked from entering high-curvature regions of parameter space, adaptive
gradient methods at the AEoS can and do enter high-curvature regions via their
ability to adapt the preconditioner.

Figure 8: The sharpness continues to rise even as Adam trains at the AEoS. We
train a fully-connected network on CIFAR-10 using full-batch Adam with
$\beta_{1}=0.9$, $\beta_{2}=0.99$, and four learning rates. Observe that while
the _preconditioned_ sharpness flatlines, the _raw_ sharpness continues to
rise.

In Figure 8, we train a fully-connected network on CIFAR-10 using full-batch
Adam at four learning rates. Observe that even though the preconditioned
sharpness $\lambda_{1}(\mathbf{P}^{-1}_{t}\mathbf{H}_{t})$ flatlines, the raw
sharpness $\lambda_{1}(\mathbf{H}_{t})$ continues to rise while the network is
training at the AEoS. This behavior stands in stark contrast to the behavior
of frozen Adam at the EoS in Figure 3; there, the raw sharpness ceased to
appreciably increase after training had entered the EoS.

Why might this be the case? A crucial distinction between adaptive and non-
adaptive optimizers is that the latter are necessarily repelled from high-
curvature regions of parameter space. For example,
$\textsc{EmaHB}(\eta,\beta_{1})$ is repelled from the set
$\\{\mathbf{x}:\lambda_{1}(\mathbf{H}(\mathbf{x}))>\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}\\}$,
and $\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$ is repelled from the set
$\\{\mathbf{x}:\lambda_{1}(\mathbf{P}^{-1}\mathbf{H}(\mathbf{x}))>\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}\\}$.
By contrast, while Adam’s trajectory
$\\{(\mathbf{P}_{t},\mathbf{x}_{t})\\}_{t}$ must approximately respect the
constraint
$\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}(\mathbf{x}_{t}))\leq\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}$,
Adam is allowed to change the preconditioner, and hence can train stably in
regions of arbitrarily high curvature.

To gain more insight into the behavior of Adam at the AEoS, in Figure 9 we
“zoom in” on 200 steps of Adam training, for both a small learning rate (5e-5,
the blue line in Figure 8) and a large learning rate (4e-4, the red line in
Figure 8). In the top row, which corresponds to the small learning rate, we
observe the following cycle. Whenever the preconditioned sharpness is below
the stability threshold, both the sharpness and preconditioned sharpness rise.
Eventually, the preconditioned sharpness crosses the stability threshold — an
event which we mark with a vertical dotted red line. This causes Adam to
oscillate explosively along a certain direction, as discussed in §2 (see
Appendix E for more evidence). The large gradients associated with this
explosive growth increase the second-moment accumulator $\boldsymbol{\nu}_{t}$
and hence Adam’s preconditioner $\mathbf{P}_{t}$. This increase in
$\mathbf{P}_{t}$ in turn causes the preconditioned sharpness
$\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}_{t})$ to decrease until it drops
below the stability threshold, at which point the large gradients go away. The
cycle then repeats itself. By this mechanism, the preconditioned sharpness
equilibrates at the stability threshold, but the raw sharpness continues to
grow. In other words, instability generates large gradients, which in turn
increase the preconditioner, allowing Adam to continue advancing into regions
of higher sharpness.

Figure 9: Zoom in on 200 steps of training. For two learning rates from Figure
9 (top row = 5e-5, bottom row = 2e-4), we zoom in on 200 steps of training,
plotting five important quantities. The red vertical lines mark the moments at
which the preconditioned sharpness crosses the stability threshold.

The bottom row of Figure 9 depicts 200 steps of Adam at a larger learning
rate. At this larger learning rate, Adam behaves somewhat differently. As
above, we see that whenever the preconditioned sharpness is below the
threshold, both the sharpness and preconditioned sharpness rise. And, as
above, we see that after the preconditioned sharpness crosses the threshold
(red vertical lines), the gradient norm begins to increase explosively.
However at the larger learning rate, this instability is “resolved”
differently than at the small learning rate. In particular, for this large
learning rate, we observe that the instability is resolved when both the
sharpness and the preconditioned sharpness drop, indicating that the optimizer
has escaped to a region where the curvature in many different directions (not
just the maximum eigenvector of the preconditioned Hessian) is lower. This
behavior is similar to the “catapult” effect studied in [25] for non-adaptive
optimizers.

#### The role of Adam’s hyperparameters

The preceding discussion suggests that Adam’s behavior at the AEoS — and in
particular, the extent to which the stability-based constraint on
$\lambda_{1}(\mathbf{P}^{-1}_{t}\mathbf{H}_{t})$ is maintained by adapting
$\mathbf{P}$ vs. by constraining $\mathbf{H}$ — may depend on Adam’s
hyperparameters. To test this hypothesis, in Figure 10(a), we use full-batch
Adam at a range of learning rates to train a Wide Resnet on CIFAR-100 until
reaching a training loss value of 0.1. We plot the sharpness
$\lambda_{1}(\mathbf{H}(\mathbf{x}^{*}))$ at the solution $\mathbf{x}^{*}$, as
a function of the learning rate. Observe that full-batch Adam at small
learning rates finds sharper solutions. In Figure 10(b), we confirm that this
trend extends to minibatch Adam at the large batch size of 4096. Meanwhile, in
Figure 10(c), we sweep full-batch Adam at a range of values of $\beta_{2}$. We
observe that when $\beta_{2}$ is small, Adam tends to find sharper solutions.
Figure 10(d) demonstrates that this trend extends to minibatch Adam at batch
size 4096. Incidentally, we show in Appendix F that the models trained with
higher $\eta$ or higher $\beta_{2}$ generalize better.

Figure 10: Adam finds sharper solutions when $\eta$ or $\beta_{2}$ is small.
We train a WRN on CIFAR-100 until reaching train loss 0.1, and we record the
sharpness $\lambda_{1}(\mathbf{H}(\mathbf{x}^{*}))$ at the solution
$\mathbf{x}^{*}$. In (a, b) we sweep the learning rate $\eta$, while in (c, d)
we sweep $\beta_{2}$. In (a, c) we train in full-batch mode, while in (b, d)
we use the large batch size 4096.

One speculative hypothesis is that Adam’s ability to resolve instabilities via
preconditioner adaptation is contingent on how fast the preconditioner adapts
(smaller $\beta_{2}$ = faster adaptation) relative to the speed at which the
optimizer is moving (larger $\eta$ = faster movement).

##  7 Adam finds sharper solutions than momentum

Adam and other adaptive gradient methods are able to train stably in
arbitrarily sharp regions of parameter space, whereas momentum gradient
descent is repelled from regions of parameter space where the sharpness
exceeds the algorithm’s stability threshold. We now demonstrate that as a
consequence, Adam sometimes converges to far sharper solutions than momentum
GD. In Figure 11, we train a WRN on CIFAR-100 (left two panes) and a fully-
connected network on CIFAR-10 (right two panes) using both Adam and momentum
GD. It is impossible to make a blanket claim that Adam finds sharper solutions
than momentum GD, since momentum at a small learning rate finds sharper
solutions than Adam at a large one; however, we observe that _for a given
training speed_ , Adam converges to a sharper solution.

Figure 11: At comparable training speeds, Adam finds sharper solutions than
momentum gradient descent. We train networks using both Adam (blue) and
momentum (orange) at a range of learning rates. For each algorithm and
learning rate, we plot the maximum Hessian eigenvalue at the solution as a
function of the number of steps to reach the solution. For both algorithms,
smaller learning rates (= longer training speeds) find sharper solutions; but
in general, Adam finds a much sharper solution than momentum GD for the same
training time.

##  8 Conclusion

In a recent year, Adam [23] was the most highly cited scientific paper across
all fields of science [9]. Yet, very little is known concretely about the
training dynamics of Adam or other adaptive gradient methods in deep learning.
Since the training dynamics affect both the optimization and generalization
performance of neural networks in mysterious ways, any basic knowledge about
the training dynamics could pay dividends for future research. This paper has
demonstrated that in the special setting of full-batch training, there is an
equilibrium rule that is typically maintained throughout training. We have
provided evidence that a similar equilibrium may exist in the more practical
minibatch setting, and we hope that an analogous equilibrium rule will
eventually be found.

## References

  * [1] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang.  Disentangling adaptive gradient methods from learning rates.  arXiv preprint arXiv:2002.11803, 2020.
  * [2] Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra.  Understanding the unstable convergence of gradient descent.  arXiv preprint arXiv:2204.01050, 2022.
  * [3] Lukas Balles and Philipp Hennig.  Dissecting adam: The sign, magnitude and variance of stochastic gradients.  In International Conference on Machine Learning, pages 404–413. PMLR, 2018.
  * [4] Anas Barakat and Pascal Bianchi.  Convergence and dynamical behavior of the adam algorithm for non-convex stochastic optimization.  arXiv preprint arXiv:1810.02263, 2018.
  * [5] Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu.  Closing the generalization gap of adaptive gradient methods in training deep neural networks.  arXiv preprint arXiv:1806.06763, 2018.
  * [6] Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong.  When vision transformers outperform resnets without pre-training or strong data augmentations.  arXiv preprint arXiv:2106.01548, 2021.
  * [7] Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong.  On the convergence of a class of adam-type algorithms for non-convex optimization.  arXiv preprint arXiv:1808.02941, 2018.
  * [8] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet S. Talwalkar.  Gradient descent on neural networks typically occurs at the edge of stability.  ArXiv, abs/2103.00065, 2021.
  * [9] Bec Crew.  Google scholar reveals its most influential papers for 2020, Jul 2020\.
  * [10] André Belotto da Silva and Maxime Gazeau.  A general system of differential equations to model first-order adaptive algorithms.  J. Mach. Learn. Res., 21:129–1, 2020.
  * [11] Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier.  A simple convergence proof of adam and adagrad.  arXiv preprint arXiv:2003.02395, 2020.
  * [12] Timothy Dozat.  Incorporating nesterov momentum into adam.  ICLR 2016, workshop track, 2016.
  * [13] John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  Journal of Machine Learning Research, 12(61):2121–2159, 2011.
  * [14] Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry.  At stability’s edge: How to adjust hyperparameters to preserve minima selection in asynchronous training of neural networks?  arXiv preprint arXiv:1909.12340, 2019.
  * [15] Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Dahl, Zachary Nado, and Orhan Firat.  A loss curvature perspective on training instability in deep learning.  arXiv preprint arXiv:2110.04369, 2021.
  * [16] Gabriel Goh.  Why momentum really works.  Distill, 2017.
  * [17] Diego Granziol.  Flatness is a false friend.  arXiv preprint arXiv:2006.09091, 2020.
  * [18] Diego Granziol, Xingchen Wan, Samuel Albanie, and Stephen Roberts.  Explaining the adaptive generalisation gap.  arXiv preprint arXiv:2011.08181, 2020.
  * [19] Stanislaw Jastrzebski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras.  Catastrophic fisher explosion: Early phase fisher matrix impacts generalization.  In International Conference on Machine Learning, pages 4772–4784. PMLR, 2021.
  * [20] Stanisław Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey.  On the relation between the sharpest directions of dnn loss and the sgd step length.  arXiv preprint arXiv:1807.05031, 2018.
  * [21] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras.  The break-even point on optimization trajectories of deep neural networks.  arXiv preprint arXiv:2002.09572, 2020.
  * [22] Nitish Shirish Keskar and Richard Socher.  Improving generalization performance by switching from adam to sgd.  arXiv preprint arXiv:1712.07628, 2017.
  * [23] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [24] Sungyoon Lee, Jinseong Park, and Jaewook Lee.  Implicit jacobian regularization weighted with impurity of probability output, 2022.
  * [25] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari.  The large learning rate phase of deep learning: the catapult mechanism.  arXiv preprint arXiv:2003.02218, 2020.
  * [26] Xiaoyu Li and Francesco Orabona.  On the convergence of stochastic gradient descent with adaptive stepsizes.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 983–992. PMLR, 2019.
  * [27] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han.  Understanding the difficulty of training transformers.  arXiv preprint arXiv:2004.08249, 2020.
  * [28] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  arXiv preprint arXiv:1608.03983, 2016.
  * [29] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [30] Kaifeng Lyu and Jian Li.  Gradient descent maximizes the margin of homogeneous neural networks.  arXiv preprint arXiv:1906.05890, 2019.
  * [31] Chao Ma, Lei Wu, et al.  A qualitative study of the dynamic behavior for adaptive gradient algorithms.  arXiv preprint arXiv:2009.06125, 2020.
  * [32] Chao Ma, Lei Wu, and Lexing Ying.  The multiscale structure of neural network loss functions: The effect on optimization and origin.  arXiv preprint arXiv:2204.11326, 2022.
  * [33] Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry.  The implicit bias of minima stability: A view from function space.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
  * [34] Kamil Nar and Shankar Sastry.  Step size matters in deep learning.  Advances in Neural Information Processing Systems, 31, 2018.
  * [35] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro.  In search of the real inductive bias: On the role of implicit regularization in deep learning.  arXiv preprint arXiv:1412.6614, 2014.
  * [36] Martin Popel and Ondřej Bojar.  Training tips for the transformer model.  arXiv preprint arXiv:1804.00247, 2018.
  * [37] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.  On the convergence of adam and beyond.  In International Conference on Learning Representations, 2018.
  * [38] Noam Shazeer and Mitchell Stern.  Adafactor: Adaptive learning rates with sublinear memory cost.  In International Conference on Machine Learning, pages 4596–4604. PMLR, 2018.
  * [39] Tijmen Tieleman and Geoffrey Hinton.  Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude.  Coursera: Neural Networks for Machine Learning, 2012.
  * [40] Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu.  The implicit bias for adaptive optimization algorithms on homogeneous neural networks.  In International Conference on Machine Learning, pages 10849–10858. PMLR, 2021.
  * [41] Rachel Ward, Xiaoxia Wu, and Leon Bottou.  Adagrad stepsizes: Sharp convergence over nonconvex landscapes.  In International Conference on Machine Learning, pages 6677–6686. PMLR, 2019.
  * [42] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht.  The marginal value of adaptive gradient methods in machine learning.  Advances in neural information processing systems, 30, 2017.
  * [43] Lei Wu, Chao Ma, and Weinan E.  How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [44] Yuege Xie, Xiaoxia Wu, and Rachel Ward.  Linear convergence of adaptive stochastic gradient descent.  In International Conference on Artificial Intelligence and Statistics, pages 1475–1485. PMLR, 2020.
  * [45] Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama.  Adai: Separating the effects of adaptive learning rate and momentum inertia.  arXiv preprint arXiv:2006.15815, 2020.
  * [46] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In International Conference on Machine Learning, pages 10524–10533. PMLR, 2020.
  * [47] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  arXiv preprint arXiv:1605.07146, 2016.
  * [48] Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu.  On the convergence of adaptive gradient methods for nonconvex optimization.  arXiv preprint arXiv:1808.05671, 2018.
  * [49] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al.  Towards theoretically understanding why sgd generalizes better than adam in deep learning.  Advances in Neural Information Processing Systems, 33:21285–21296, 2020.

##  Appendix A Mathematical background

###  A.1 Stability of gradient descent algorithms on quadratics

We now review the stability properties of gradient descent algorithms on
quadratic objective functions. We first refer the reader to Appendix A of [8];
that appendix derives the stability thresholds for $\textsc{GD}(\eta)$,
$\textsc{StandardHB}(\eta,\beta_{1})$, and
$\textsc{StandardNesterov}(\eta,\beta_{1})$. In this appendix, we derive the
stability thresholds for $\textsc{EmaHB}(\eta,\beta_{1})$, and
$\textsc{EmaNesterov}(\eta,\beta_{1})$. These derivations are minor variants
of the corresponding derivations in [8], since the EMA-style momentum
algorithms yield the same recurrence relation as the standard-style momentum
algorithms, except with a learning rate scaled by $(1-\beta_{1})$.

Algorithm | Update rule | Stability threshold | Reference for threshold
---|---|---|---
$\textsc{GD}(\eta)$ | $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\mathbf{g}_{t+1}$ | $2/\eta$ |  common knowledge, reproduced in [8] for completeness
$\textsc{StandardHB}(\eta,\beta_{1})$ |  $\mathbf{m}_{t+1}=\beta_{1}\mathbf{m}_{t}+\mathbf{g}_{t+1}$ $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\mathbf{m}_{t+1}$ | $\frac{2+2\beta_{1}}{\eta}$ |  [16], reproduced in [8] for completeness
$\textsc{EmaHB}(\eta,\beta_{1})$ |  $\mathbf{m}_{t+1}=\beta_{1}\mathbf{m}_{t}+(1-\beta_{1})\mathbf{g}_{t+1}$ $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\mathbf{m}_{t+1}$ | $\frac{2+2\beta_{1}}{\eta(1-\beta_{1})}$ | this paper
$\textsc{StandardNesterov}(\eta,\beta_{1})$ |  $\mathbf{m}_{t+1}=\beta_{1}\mathbf{m}_{t}+\mathbf{g}_{t+1}$ $\hat{\mathbf{m}}_{t+1}=\beta_{1}\mathbf{m}_{t+1}+\mathbf{g}_{t+1}$ $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\hat{\mathbf{m}}_{t+1}$ | $\frac{2+2\beta_{1}}{\eta(1+2\beta_{1})}$ | [8]
$\textsc{EmaNesterov}(\eta,\beta_{1})$ |  $\mathbf{m}_{t+1}=\beta_{1}\mathbf{m}_{t}+(1-\beta_{1})\mathbf{g}_{t+1}$ $\hat{\mathbf{m}}_{t+1}=\beta_{1}\mathbf{m}_{t+1}+(1-\beta_{1})\mathbf{g}_{t+1}$ $\mathbf{x}_{t+1}=\mathbf{x}_{t}-\eta\hat{\mathbf{m}}_{t+1}$ | $\frac{2+2\beta_{1}}{\eta(1-\beta_{1})(1+2\beta_{1})}$ | this paper
Table 1: Gradient descent algorithms and their stability thresholds.

######  Lemma 1.

Consider running $\textsc{EmaNesterov}(\eta,\beta_{1})$ on a quadratic
objective
$f(\mathbf{x})=\frac{1}{2}\mathbf{x}^{T}\mathbf{A}\mathbf{x}+\mathbf{b}^{T}\mathbf{x}+c$
starting from any initialization. Let $(\mathbf{q},a)$ be an
eigenvector/eigenvalue pair of $\mathbf{A}$. If
$a>\frac{1}{\eta}\left(\frac{2+2\beta_{1}}{(1-\beta_{1})(1+2\beta_{1})}\right)$,
then the sequence $\\{\mathbf{q}^{T}\mathbf{x}\\}$ will diverge.

###### Proof.

First, we re-write $\textsc{EmaNesterov}(\eta,\beta_{1})$ as a recursion in
$\mathbf{x}_{t}$ alone:

$\mathbf{x}_{t+1}=(1+\beta_{1})\mathbf{x}_{t}-\beta_{1}\mathbf{x}_{t-1}-\eta(1-\beta_{1})(1+\beta_{1})\nabla
f(\mathbf{x}_{t})+\eta(1-\beta_{1})\nabla f(\mathbf{x}_{t-1})$

For the quadratic objective, we have $\nabla
f(\mathbf{x}_{t})=\mathbf{A}\mathbf{x}_{t}+\mathbf{b}$, so this update
becomes:

$\mathbf{x}_{t+1}=(1+\beta_{1})(\mathbf{I}-\eta\mathbf{A})\mathbf{x}_{t}-\beta_{1}(\mathbf{I}-\eta(1-\beta_{1})\mathbf{A})\mathbf{x}_{t-1}-\eta(1-\beta_{1})\mathbf{b}.$

This is exactly the recurrence in Theorem 1 of [8], except with
$(1-\beta_{1})\eta$ in place of $\eta$. Thus, by the same logic as in Theorem
1 of [8], $\mathbf{q}^{T}\mathbf{x}_{t}$ diverges if
$a>\frac{1}{\eta}\left(\frac{2+2\beta_{1}}{(1-\beta_{1})(1+2\beta_{1})}\right)$.
∎

######  Lemma 2.

Consider running $\textsc{EmaHB}(\eta,\beta_{1})$ on a quadratic objective
$f(\mathbf{x})=\frac{1}{2}\mathbf{x}^{T}\mathbf{A}\mathbf{x}+\mathbf{b}^{T}\mathbf{x}+c$
starting from any initialization. Let $(\mathbf{q},a)$ be an
eigenvector/eigenvalue pair of $\mathbf{A}$. If
$a>\frac{1}{\eta}\left(\frac{2+2\beta_{1}}{1-\beta_{1}}\right)$, then the
sequence $\\{\mathbf{q}^{T}\mathbf{x}\\}$ will diverge.

###### Proof.

First, we re-write $\textsc{EmaHB}(\eta,\beta_{1})$ as a recursion in
$\mathbf{x}_{t}$ alone:

$\mathbf{x}_{t+1}=(1+\beta_{1})\mathbf{x}_{t}-\beta_{1}\mathbf{x}_{t-1}-\eta(1-\beta_{1})\nabla
f(\mathbf{x}_{t}).$

For the quadratic objective, we have $\nabla
f(\mathbf{x}_{t})=\mathbf{A}\mathbf{x}_{t}+\mathbf{b}$, so this update
becomes:

$\mathbf{x}_{t+1}=\left[(1+\beta_{1})\mathbf{I}-\eta(1-\beta_{1})\mathbf{A}\right]\mathbf{x}_{t}-\beta_{1}\mathbf{x}_{t-1}-\eta(1-\beta_{1})\mathbf{b}.$

This is exactly the recurrence in Theorem 2 of [8], except with
$(1-\beta_{1})\eta$ in place of $\eta$. Thus, by the same logic as in Theorem
2 of [8], $\mathbf{q}^{T}\mathbf{x}_{t}$ diverges if
$a>\frac{1}{\eta}\left(\frac{2+2\beta_{1}}{1-\beta_{1}}\right)$. ∎

###  A.2 Stability of preconditioned gradient descent

As discussed in §2, gradient descent can be preconditioned using a
preconditioner $\mathbf{P}$. For example, the preconditioned version of
$\textsc{EmaHB}(\eta,\beta_{1})$, which we denote
$\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$ is:

$\displaystyle\mathbf{m}_{t+1}$
$\displaystyle=\beta_{1}\mathbf{m}_{t}+(1-\beta_{1})\mathbf{g}_{t+1},$
$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=\mathbf{x}_{t}-\eta\,\mathbf{P}^{-1}\,\mathbf{m}_{t+1}.$

On quadratic objective functions, the stability behavior of preconditioned
gradient descent algorithms differs slightly from that of their non-
preconditioned counterparts. Whereas non-preconditioned gradient descent
algorithms become unstable when the maximum eigenvalue of the Hessian
$\mathbf{H}$ exceeds a certain stability threshold, preconditioned gradient
descent algorithms become unstable when the maximum eigenvalue of the
_preconditioned Hessian_ $\mathbf{P}^{-1/2}\mathbf{H}\mathbf{P}^{-1/2}$
exceeds the threshold of their non-precondintioned counterpart.

The simplest way to see this is to note that preconditioned gradient descent
algorithms are isomorphic to their non-preconditioned counterparts up to a
certain reparameterization of the training objective. In particular,
preconditioned gradient descent on an objective $f$ using preconditioner
$\mathbf{P}$ is isomorphic to non-preconditioned gradient descent on the
reparameterized objective
$\hat{f}(\mathbf{x})=f(\mathbf{P}^{-1/2}\mathbf{x})$, in the sense that the
iterates of the two algorithms are related by a linear map. Below, we make
this rigorous for $\textsc{EmaHB}(\eta,\beta_{1})$ and
$\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$, but note that one can carry
out the same argument for the other momentum algorithms and their
preconditioned counterparts:

######  Proposition 1.

Let $\\{(\mathbf{x}_{t},\mathbf{m}_{t})\\}$ denote the iterates of
$\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$ on the objective function
$f(\mathbf{x})$, and let
$\\{(\tilde{\mathbf{x}}_{t},\tilde{\mathbf{m}}_{t})\\}$ denote the iterates of
$\textsc{EmaHB}(\eta,\beta_{1})$ on the reparameterized objective function
$\tilde{f}(\mathbf{x})=f(\mathbf{P}^{-1/2}\mathbf{x})$ starting from the
initialization
$(\tilde{\mathbf{x}}_{0},\tilde{\mathbf{m}_{0}})=(\mathbf{P}^{1/2}\mathbf{x}_{0},\mathbf{P}^{-1/2}\mathbf{m}_{0})$.
Then we have the equivalence
$\tilde{\mathbf{x}}_{t}=\mathbf{P}^{1/2}\mathbf{x}_{t}$ and
$\tilde{\mathbf{m}}_{t}=\mathbf{P}^{-1/2}\mathbf{m}_{t}$ for all steps $t$.

###### Proof.

The equivalence is true at initialization, by definition. Thus, it remains to
show that if the equivalence is true at step $t$, it is true at step $t+1$.

The updates for $\\{(\mathbf{x}_{t+1},\mathbf{m}_{t+1})\\}$ are:

$\displaystyle\mathbf{m}_{t+1}$
$\displaystyle=\beta_{1}\mathbf{m}_{t}+(1-\beta_{1})\nabla f(\mathbf{x}_{t}),$
$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=\mathbf{x}_{t}-\eta\mathbf{P}^{-1}\mathbf{m}_{t+1}.$

Meanwhile, the updates for
$\\{(\tilde{\mathbf{x}}_{t+1},\tilde{\mathbf{m}}_{t+1})\\}$ are:

$\displaystyle\tilde{\mathbf{m}}_{t+1}$
$\displaystyle=\beta_{1}\tilde{\mathbf{m}}_{t}+(1-\beta_{1})\nabla\tilde{f}(\tilde{\mathbf{x}}_{t})$
$\displaystyle=\beta_{1}\tilde{\mathbf{m}}_{t}+(1-\beta_{1})\mathbf{P}^{-1/2}\nabla
f(\mathbf{P}^{-1/2}\tilde{\mathbf{x}}_{t}),$
$\displaystyle\tilde{\mathbf{x}}_{t+1}$
$\displaystyle=\tilde{\mathbf{x}}_{t}-\eta\tilde{\mathbf{m}}_{t+1}.$

To verify the equivalence
$\tilde{\mathbf{m}}_{t+1}=\mathbf{P}^{-1/2}\mathbf{m}_{t+1}$, observe that

$\displaystyle\mathbf{P}^{-1/2}\,\mathbf{m}_{t+1}$
$\displaystyle=\beta_{1}\,\mathbf{P}^{-1/2}\mathbf{m}_{t}+(1-\beta_{1})\mathbf{P}^{-1/2}\nabla
f(\mathbf{x}_{t})$
$\displaystyle=\beta_{1}\tilde{\mathbf{m}_{t}}+(1-\beta_{1})\mathbf{P}^{-1/2}\nabla
f(\mathbf{P}^{-1/2}\tilde{\mathbf{x}}_{t})$
$\displaystyle=\tilde{\mathbf{m}}_{t+1}.$

where the two replacements in the middle step are due to the inductive
hypothesis holding for step $t$.

To verify the equivalence
$\tilde{\mathbf{x}}_{t+1}=\mathbf{P}^{1/2}\mathbf{x}_{t+1}$, observe that

$\displaystyle\mathbf{P}^{1/2}\mathbf{x}_{t+1}$
$\displaystyle=\mathbf{P}^{1/2}\mathbf{x}_{t}-\eta\mathbf{P}^{-1/2}\mathbf{m}_{t+1}$
$\displaystyle=\tilde{\mathbf{x}}_{t}-\eta\tilde{\mathbf{m}_{t}}$
$\displaystyle=\tilde{\mathbf{x}}_{t+1}.$

where the two replacements in the middle step are due to the inductive
hypothesis holding for step $t$.

∎

Due to Proposition 1, we know that
$\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$ on a quadratic function

$f(\mathbf{x})=\frac{1}{2}\mathbf{x}^{T}\mathbf{A}\mathbf{x}+\mathbf{b}^{T}\mathbf{x}+c$

has the same stability behavior as $\textsc{EmaHB}(\eta,\beta_{1})$ on the
reparameterized quadratic

$\tilde{f}(\mathbf{x})=\frac{1}{2}\mathbf{x}^{T}(\mathbf{P}^{-1/2}\mathbf{A}\mathbf{P}^{-1/2})\mathbf{x}+\mathbf{b}^{T}\mathbf{P}^{-1/2}\mathbf{x}+c.$

Thus, $\textsc{PreconEmaHB}(\eta,\beta_{1},\mathbf{P})$ diverges on $f$
whenever the maximum eigenvalue of
$(\mathbf{P}^{-1/2}\mathbf{A}\mathbf{P}^{-1/2})$ exceeds the stability
threshold of $\textsc{EmaHB}(\eta,\beta_{1})$. When this occurs, the component
of $\mathbf{P}^{1/2}\mathbf{x}_{t}$ that is aligned with $\mathbf{v}_{1}$, the
top eigenvector of the preconditioned Hessian, will oscillate with
exponentially increasing magnitude. Equivalently, we can say that the
component of $\mathbf{x}_{t}$ that is aligned with
${(\mathbf{P}^{1/2})}^{T}\mathbf{v}_{1}$ will oscillate with exponentially
increasing magnitude.

##  Appendix B Adaptive preconditioning

In Figure 12, we optimize the objective function $f(x)=\frac{1}{2}x^{2}$ using
rmsprop with learning rates that would be divergent if the preconditioner
$\nu_{t}$ were held fixed at its initial value of $\nu_{0}=1.0$. Yet since the
preconditioner is adaptive rather than fixed, rmsprop only oscillates
explosively for a finite number of steps before the large gradients associated
with this explosive movement cause $\nu_{t}$ to increase, which decreases
rmsprop’s effective learning rate and halts divergence.

Figure 12: Adaptive preconditioning often prevents rmsprop from diverging. We
optimize the objective function $f(x)=\frac{1}{2}x^{2}$ using rmsprop. In the
left column, we hold $\eta$ fixed at $\eta=2.2$ and vary $\beta_{2}$; in the
right column, we hold $\beta_{2}$ fixed at $\beta_{2}=0.999$ and vary $\eta$.
Observe that rmsprop only oscillates explosively for a finite number of steps
before the large gradients associated with this explosive movement cause
$\nu_{t}$ to increase, which decreases rmsprop’s effective learning rate and
halts divergence.

##  Appendix C Qualitative aspects of the AEoS

Looking closely at Figure 5, one can discern that the preconditioned sharpness
behaves in qualitatively different ways depending on the type of momentum. For
rmsprop and Adagrad, which do not employ any momentum, the preconditioned
sharpness hovers above the stability threshold. For Nadam, too, which employs
Nesterov-style momentum, the preconditioned sharpness also hovers above the
stability threshold. In contrast, for the other algorithms, which employ
heavy-ball style momentum, the preconditioned sharpness does not hover above
the stability threshold — instead, the optimizer is rapidly thrown into
flatter regions whenever the preconditioned sharpness crosses the threshold;
as a result, the preconditioned sharpness oscillates around the stability
threshold instead of hovering above it.

We now (1) confirm that these behaviors also occur in the simpler setting of
non-adaptive gradient descent, and (2) suggest a speculative explanation.

#### Non-adaptive gradient descent

The pattern described above also holds for non-adaptive gradient descent. For
example, Figures 29, 95, and 96 in [8] depict the evolution of the sharpness
as a network is trained using vanilla gradient descent, heavy ball momentum,
and Nesterov momentum, respectively. In both Figure 29 (vanilla) and Figure 96
(Nesterov), the sharpness hovers above the stability threshold, whereas in
Figure 95 (heavy ball), the sharpness oscillates around the threshold.

We now reproduce this observation using the same network from Figure 5. In
Figure 13, we train the fully-connected network on CIFAR-10 using vanilla
gradient descent $\textsc{GD}(\eta)$, heavy ball gradient descent
$\textsc{StandardHB}(\eta,0.9)$, and Nesterov gradient descent
$\textsc{StandardNesterov}(\eta,0.9)$, each at a range of learning rates.
Observe that for vanilla and Nesterov gradient descent, the sharpness hovers
above the stability threshold (except at the smallest learning rate
considered), whereas for heavy ball gradient descent, the sharpness oscillates
around the stability threshold.

Figure 13: Similar to Figure 5, the behavior of the sharpness during (non-
adaptive) gradient descent depends on the type of momentum. We train a fully-
connected network on CIFAR-10 using three non-adaptive algorithms: vanilla
gradient descent ($\eta$ in logspace between 0.02 and 0.2), heavy ball
momentum ($\beta_{1}=0.9$, and $\eta$ in logspace between 0.002 and 0.02), and
Nesterov momentum ($\beta_{1}=0.9$, and $\eta$ in logspace between 0.002 and
0.02). Observe that for vanilla GD and Nesterov, the sharpness hovers above
the stability threshold, whereas for heavy ball, it oscillates around the
stability threshold. This pattern parallels that in Figure 5 for Adam.

#### A speculative explanation

We now give a speculative explanation for the above observations. One
difference between heavy ball gradient descent, on the one hand, and Nesterov
and vanilla gradient descent, on the other, is that heavy ball momentum can be
much less robust to increases in the sharpness beyond the optimizer’s
stability threshold. By way of background, recall from [8] that on quadratic
functions, all three of these algorithms act independently along each Hessian
eigenvector; in particular, the component of the iterate that is aligned with
each Hessian eigenvector evolves according to a linear recurrence relation
whose coefficients are determined by the corresponding Hessian eigenvalue
$\lambda$ [8]. For example, the recurrence relation for $\textsc{GD}(0.2)$ is
$x_{t+1}=(1-0.2\,\lambda)x_{t}$, the recurrence relation for
$\textsc{StandardHB}(0.02,0.9)$ is
$x_{t+1}=(1.9-0.02\lambda)x_{t}-0.9x_{t-1}$, and the recurrence relation for
$\textsc{StandardNesterov}(0.02,0.9)$ is
$x_{t+1}=1.9(1-0.02\lambda)x_{t}-0.9(1-0.02\lambda)x_{t-1}$. Asymptotically
(i.e. over many optimizer steps), the growth rate of a linear recurrence
relation is determined by the magnitude of the largest root of the
characteristic polynomial. (Note that for a recurrence of the form
$x_{t+1}=ax_{t}+bx_{t-1}$, the characteristic polynomial is $x^{2}-ax-b$.) The
stability threshold $\lambda^{*}$ of an optimizer is the value of $\lambda$
for which the magnitude of the largest-magnitude root is exactly 1. The
optimizer’s robustness to barely-unstable eigenvalues ($\lambda$ just above
$\lambda^{*}$) can be quantified by the asymptotic growth rate at at those
eigenvalues.

For example, consider the three optimization algorithms $\textsc{GD}(0.2)$,
$\textsc{StandardHB}(0.02,0.9)$, and $\textsc{StandardNesterov}(0.02,0.9)$.
These algorithms all share the same effective learning rate of 0.2, and they
correspond to the purple lines in Figure 13. In Figure 14, we visualize the
asymptotic growth factor of these algorithms as a function of the input
eigenvalue $\lambda$. The stability threshold $\lambda^{*}$ is marked with a
vertical dotted black line (of course, the asymptotic growth factor at
$\lambda^{*}$ is always 1), and we plot the asymptotic growth factor for
eigenvalues up to 1.1 times the stability threshold.

Crucially, observe that the asymptotic growth rate for
$\textsc{StandardHB}(0.02,0.9)$ at an eigenvalue 1.1 times its stability
threshold is much higher than for $\textsc{GD}(0.02)$ and
$\textsc{StandardNesterov}(0.02,0.9)$. This means that whenever StandardHB
finds itself in a region of parameter space where the sharpness is 1.1 times
its stability threshold, it diverges much faster than $\textsc{GD}(0.02)$ or
$\textsc{StandardNesterov}(0.02,0.9)$ do in an analogous situation. One can
easily imagine that this may be related to our observation that the latter two
algorithms can train at the AEoS with the sharpness hovering a bit above the
stability threshold, whereas heavy ball momentum is flung into a flatter
region whenever the sharpness crosses the stability threshold even by a small
amount.

Figure 14: Vanilla gradient descent and Nesterov momentum are both more robust
than heavy ball momentum to eigenvalues shortly above the stability threshold.
We plot the asymptotic growth factors of three algorithms: vanilla GD, heavy
ball GD, and Nesterov GD. That is, for each possible Hessian eigenvalue, we
plot the multiplicative constant which the corresponding component of the
iterate will be asymptotically multiplied by, over many steps of gradient
descent. The vertical dotted line marks each algorithm’s stability threshold,
and we plot the asymptotic growth factor (AGF) for eigenvalues up to 1.1 times
that stability threshold. Observe that for heavy ball momentum, the AGF at 1.1
times the stability threshold is much higher than for Nesterov momentum or
vanilla gradient descent.

##  Appendix D Corner case

While we generally observe that during full-batch Adam, the preconditioned
sharpness equilibrates at the stability threshold of frozen Adam, we have also
consistently observed one notable corner case where this does not occur.
Namely, when running Adam at _extremely small_ learning rates, we often find
that the preconditioned sharpness either flatlines short of the stability
threshold, or increases at a very slow rate while remaining below the
threshold. These behaviors are evident in Figure 15, in which we train a CNN
on CIFAR-10 using full-batch Adam at a range of very small learning rates.
(This task is the same one as in Figure 4, left panel, just at smaller
learning rates.) Observe that at these small learning rates, the
preconditioned sharpness fails to rise to the stability threshold.

Figure 15: When running Adam at _extremely small_ learning rates, the
sharpness can fail to rise all the way to the stability threshold. We train a
CNN on CIFAR-10 using full-batch Adam at a range of extremely small learning
rates. Observe that the preconditioned sharpness fails to rise all the way to
the stability threshold.

Why does the preconditioned sharpness not rise to the stability threshold? One
clue is that this behavior seems to be related to adaptivity. In Figure 16, we
train using full-batch Adam at the green learning rate from Figure 15, and at
step 6000 of training, we suddenly freeze the preconditioner (i.e. we stop
updating $\boldsymbol{\nu}_{t}$). The train loss and preconditioned sharpness
of frozen Adam are plotted in pink. Observe that upon freezing the
preconditioner, the preconditioned sharpness rapidly rises all the way to the
stability threshold. Thus, we can infer that it was due to Adam’s adaptivity
that the preconditioned sharpness was located below the stability threshold in
the first place.

Figure 16: If we freeze the preconditioner, the preconditioned sharpness
immediately rises to the stability threshold. In green, we train a CNN on
CIFAR-10 using full-batch Adam at step size $\eta=2.236$e-6 (the green
learning rate from Figure 15). In pink, we freeze the preconditioner at step
6,000. Observe that once the preconditioner is frozen, the preconditioned
sharpness rises immediately to the stability threshold, suggesting that
adaptive preconditioning was preventing it from rising to the stability
threshold in the first place.

We hypothesize that this behavior is related to the following property of
momentum algorithms: momentum algorithms can generate large gradients even
when they are stable. For example, in Figure 17, we optimize the quadratic
objective $f(x)=\frac{1}{2}x^{2}$ using EMA heavy ball momentum at
$\beta_{1}=0.9$ and three different stable learning rates (columns). We set
the initial iterate $x_{0}$ to 1, and the initial momentum vector $m_{0}$ to
0.2. We run $\textsc{EmaHB}(\eta,0.9)$ for thirty steps, plotting the iterate
(top row), the momentum vector (middle row), and the gradient squared norm
(bottom row). Observe that the gradient squared norm increases dramatically
from its initial value.

We therefore hypothesize that the preconditioned sharpness
$\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}_{t})$ fails to rise all the way to
the stability threshold because the large gradients generated by the momentum
algorithm (even while it is stable) are causing the preconditioner
$\mathbf{P}_{t}$ to constantly increase. This, in turn “cancels out” the
effect of progressive sharpening, and causes the preconditioned sharpness
$\lambda_{1}(\mathbf{P}_{t}^{-1}\mathbf{H}_{t})$ to remain short of the
stability threshold.

Figure 17: Momentum can generate large gradients even when stable. We optimize
the function $f(x)=\frac{1}{2}x^{2}$ using 30 steps of
$\textsc{EmaHB}(\eta,0.9)$ at various stable learning rates $\eta$. The top
row plots the iterate, the second row the momentum vector (really a scalar),
and the third row the squared gradient norm. Observe that the gradient norm
increases from its initial value, sometimes dramatically. This does not happen
with vanilla gradient descent.

##  Appendix E Further plots of zoom in

Figure 18: Explosive growth in the (preconditioned) gradient is limited to the
top eigenvectors of the preconditioned Hessian. We train a fully-connected
network on CIFAR-10 using full-batch Adam. In the top row, we plot the train
loss. In the middle row, we plot the evolution of the top four eigenvalues of
the preconditioned Hessian $\mathbf{P}_{t}^{-1}\mathbf{H}_{t}$. In the bottom
row, we plot (1) in blue, the squared norm of the “semi-preconditioned
gradient” $\mathbf{P}^{1/2}\mathbf{g}_{t}$, which is the quantity that is
expected to grow explosively when the algorithm is unstable; in orange, the
squared norm of the portion of $\mathbf{P}^{1/2}\mathbf{g}_{t}$ that is
aligned with the top four eigenvectors of $\mathbf{P}_{t}^{-1}\mathbf{H}_{t}$;
and in green, the squared norm of the remainder of
$\mathbf{P}^{1/2}\mathbf{g}_{t}$. Observe that almost all of the spikes in the
gradient norm can be attributed to the top eigenvectors of the preconditioned
Hessian.

##  Appendix F Test accuracy of Wide ResNet

In §6, we saw that Adam tended to find solutions with higher sharpness
(maximum Hessian eigenvalue) whenever the learning rate $\eta$ was smaller or
the decay factor $\beta_{2}$ was smaller. We now demonstrate that the test
error tends to be higher in these same situations. In Figure 19, we plot the
test error of the Wide ResNet from Figure 10 as a function of $\eta$ and
$\beta_{2}$. (For each training run, we select the test error at the moment
the train loss first drops below the threshold of 0.05.) As in Figure 10, we
experiment with both the full-batch and large batch (batch size 4096)
settings. Observe that across both these settings, the test error tends to be
higher whenever $\eta$ is small (subplots (a) or (b)) or whenever $\beta_{2}$
is small (subplots (c) or (d)). We emphasize that we do not intend to claim
that there is necessarily a causal connection between the test error and the
sharpness — we only mean to point out that the Adam’s implicit bias towards
low-curvature regions when $\eta$ is high or $\beta_{2}$ is high coincides
with superior generalization performance.

Figure 19: Adam tends to generalize better when $\eta$ or $\beta_{2}$ is
small. We train a Wide ResNet on CIFAR-100 using full-batch (panels a, c) or
large-batch (panels b, d) Adam at various $\eta$ (panels a, b) or $\beta_{2}$
(panels b, d). We plot the test error, measured when the training loss first
drops below the threshold 0.05.

##  Appendix G Experimental details

###  G.1 Architectures

#### Fully-connected network

This network has 5 hidden layers of width 200 each. The activation function is
tanh, and the weights are initialized with variance 1 / fan_in. The fully-
connected layers have bias parameters, which are initialized to zero.

#### CNN

This network repeats (Conv => ReLU => maxpool) three times, then flattens the
activations and does FC => ReLU => FC => ReLU => FC. The first conv layer has
64 filters, a kernel size of 5, “valid” padding, a 3x3 window size, and stride
2. The second conv layer has 96 filters, a kernel size of 3, “valid” padding,
a 3x3 window size, and stride 2. The third conv layer has 128 filters, a
kernel size of 3, “same” padding, a 3x3 window size, and stride 2. The first
fully-connected layer has width 512, and the second has width 256. All biases
are initialized to zero, and all weights are initialized with variance 1 /
fan_in.

#### Wide ResNet (normalized)

This Wide ResNet has four blocks per group, and a channel multiplier of 10.

#### Wide ResNet (un-normalized)

When we remove batch normalization, we insert a learnable scalar multiplier
(initialized to 1) after the residual branch of each block.

###  G.2 Figures

#### Figure 2

We extracted the fixed preconditioner from step 1100 of a network trained
using full-batch Adam at learning rate 2e-5.

#### Figure 4

In the left pane, we trained for 2000 steps or until reaching a threshold loss
value of 0.05. In the middle pane, we trained for 5000 steps. In the right
pane, we trained until reaching a threshold loss value of 0.05.

#### Figure 5

Adam with bias correction: $\beta_{1}=0.9$, $\beta_{2}=0.999$,
$\epsilon=$1e-7, learning rates were:

np.logspace(np.log10(3e-5), np.log10(1e-3), 5)

AdamW: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=$1e-7, learning rates
were:

np.logspace(np.log10(3e-5), np.log10(1e-3), 5)

Adafactor: The decay rate was 0.8, the momentum was 0.9, epsilon was 1e-7, and
the learning rates were

np.logspace(-5, -3, 5)

Amsgrad: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=$1e-7, bias correction
was off, and learning rates were:

np.logspace(-5, -3, 5)

Padam: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=$1e-7, bias correction
was off, the Padam exponent was 0.25, and learning rates were:

np.logspace(-3, -1, 5)

Nadam: $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=$1e-7, bias correction
was off, and learning rates were:

np.logspace(-5, -3, 5)

rmsprop: decay was 0.995, epsilon was 1e-7, and learning rates were:

np.logspace(-4, -3, 5)

Adagrad: Learning rates were np.logspace(-2, -1, 5).

#### Figure 10

For the learning rate sweeps we fixed $\beta_{2}=0.999$, and for the
$\beta_{2}$ sweeps we fixed $\eta=$ 1e-4.

#### Figure 11

For Adam we used $\beta_{1}=0.9$ and $\beta_{2}=0.999$; for momentum we used
$\beta_{1}=0.9$.