G RADIENT D ESCENT ON N EURAL N ETWORKS T YPICALLY O CCURS AT THE E DGE OF S TABILITY
Anonymous authors
Paper under double-blind review
A BSTRACT
We empirically demonstrate that full-batch gradient descent on neural network training objectives typically operates in a regime we call the Edge of Stability.
In this regime, the leading eigenvalue of the training loss Hessian hovers just above the value 2 / (step size), and the training loss behaves non-monotonically over short timescales, even as it consistently decreases over long timescales.
Our results expose a large gap between theory and practice in non-convex optimization: many analyses of gradient descent (those based on L -smoothness or monotonic descent) do not apply at the Edge of Stability.
Our results also shed signiﬁcant light on the dynamics of gradient descent with a ﬁxed step size, and how that algorithm differs from gradient ﬂow.
I NTRODUCTION
This paper empirically studies full-batch gradient descent with a ﬁxed step size η , applied to neural network training.
We demonstrate that, at all but the smallest learning rates, gradient descent eventually transitions into a regime we call the Edge of Stability , in which the sharpness (the largest eigenvalue of the training loss Hessian) hovers right at (or just past) the value 2 /η .
This behavior is consistent across architectures, activation functions, tasks, loss functions, and batch normalization.
An analogous phenomenon occurs with both Polyak and Nesterov momentum.
The Edge of Stability reﬂects a balance between two opposing processes: progressive sharpening and instability .
On the one hand, we demonstrate empirically that gradient descent on neural network training objectives has an overwhelming tendency to continually move in a direction of increased sharpness.
We call this phenomenon progressive sharpening .
On the other hand, gradient descent
[Figure]
with step size η is, to quadratic order, unstable in regions of the loss landscape where the sharpness exceeds the value 2 /η : in such a region, gradient descent would diverge if run on the quadratic
Taylor approximation.
The joint consequence of sharpening and instability is that during gradient descent, the sharpness increases to, and then lingers at, the value 2 /η — the Edge of Stability.
Since neural network training is poorly understood, we ﬁnd it striking that across many different architectures and tasks, for both vanilla and momentum gradient descent, there is a simple rule that can numerically predict, to respectable accuracy, the value of some statistic (the sharpness) during training.
Furthermore, our results (1) expose a signiﬁcant gap between optimization theory and practice, and (2) shed light on the dynamics of gradient descent with a non-inﬁnitesimal step size.
Optimization theory The conventional wisdom in optimization theory is that the step size η of gradient descent should be small enough that the sharpness along the optimization trajectory is less than 2 /η
(Bottou et al., 2018;
Li & Arora, 2019).
We refer to this assumption as the Bounded Curvature (BC) condition.
The BC condition ensures both (1) that gradient descent will monotonically decrease the training objective, and (2) that gradient descent is a stable algorithm.
In the case of neural network training, our paper complicates this conventional wisdom in three respects.
First, because of progressive sharpening, in order for the BC condition to hold for the entire duration of training, the step size η must be set surprisingly small.
Second, while neural networks can in principle be trained with these extremely small learning rates, we ﬁnd that they can also be successfully trained at much larger learning rates.
At these large learning rates, gradient descent eventually transitions into the Edge of Stability, a regime where gradient descent succeeds even though the
BC condition does not hold.
Third, gradient descent converges much faster at these larger learning rates; learning rates small enough for the BC condition to hold throughout training are highly suboptimal in terms of convergence speed — in a word, they are “unreasonable.”
Thus, during gradient descent training of neural networks at reasonable learning rates, the BC condition does not hold.
An extreme (but not atypical) case is the network in Figure 1: while we know that there must exist learning rates small enough for the BC condition to hold throughout training, it is not feasible on an actual computer to train this network to completion at such learning rates.
Even analyses of gradient descent that do not explicitly assume the BC condition often assert that gradient descent will decrease the training objective monotonically (Allen-Zhu et al., 2019; Zhang et al., 2020).
However, at the Edge of Stability, the training loss usually behaves non-monotonically over short timescales even as it consistently decreases over long timescales.
Thus, in deep learning, gradient descent at reasonable learning rates does not decrease the training objective monotonically.
Dynamics of gradient descent Our paper sheds signiﬁcant light on the dynamics of gradient descent on neural network training objectives, and speciﬁcally on how the trajectory of gradient descent differs from that of gradient ﬂow (gradient descent with inﬁnitesimal step sizes; see Bach (2020) for a primer).
We show that gradient ﬂow tends to continually head into sharper and sharper regions of the loss landscape (progressive sharpening).
However, gradient descent with step size η can only track the gradient ﬂow trajectory up till the point on that trajectory where the sharpness reaches 2 /η .
After this point, gradient descent departs from the gradient ﬂow trajectory and enters the Edge of
Stability.
At the Edge of Stability, gradient descent is constantly “trying” to follow gradient ﬂow into a region of increased sharpness (and indeed, this is what would happen if we were to drop the learning rate), but is being (somehow) blocked from doing so by considerations of stability.
Thus, because of progressive sharpening and instability, the optimization trajectory of gradient descent differs from that of gradient ﬂow, and the trajectory of GD with a large learning rate differs from that of GD with a small learning rate.
Note that while the trajectories of gradient descent at different stable step sizes can conceivably differ from the gradient ﬂow trajectory, and from one another, due to discretization error (Nakkiran, 2020; Barrett & Dherin, 2020), our experiments demonstrate that in practice, the predominant impact of non-inﬁnitesimal step sizes is rather due to instability.
Stochastic gradient descent We study full-batch gradient descent rather than SGD because only in the full-batch case are we (so far) able to precisely characterize the dynamics of the sharpness, and the effect of these dynamics on the optimization trajectory.
Jastrzebski et al.
(2020) recently studied similar questions for SGD, and drew much weaker conclusions (see Section 2).
However, some of our results may have analogues for SGD (Section 5).
B ACKGROUND AND RELATED WORK
Evolution of sharpness A number of recent works have empirically studied the evolution of the sharpness (and, more generally, the Hessian spectrum) during neural network training (Sagun et al., 2017; Ghorbani et al., 2019; Li et al., 2020; Papyan, 2018; 2019; 2020)
Notably, several works have observed that the hyperparameters of SGD seem to affect the sharpness along the optimization trajectory: large learning rates and small batch sizes steer SGD into less sharp regions of parameter space (Keskar et al., 2016; Jastrzebski et al., 2017; 2019); in the case of full-batch gradient descent, our paper quantiﬁes this.
Most similar to our paper is Jastrzebski et al.
(2020), who postulated a simpliﬁed model for the beginning of SGD training which assumed, among other things, that the sharpness would increase until its increase caused SGD to become destabilized.
However, while the stability properties of
GD are clear (to quadratic order, the algorithm diverges iff the sharpness exceeds 2 /η )
the stability properties of SGD are not yet well-understood; indeed, the stability of SGD only relates to the sharpness under a strong assumption not likely to be fully satisﬁed in practice.
Consequently, Jastrzebski et al.
(2020) could not (and did not) actually verify their assumption that the sharpness would increase as much as it possibly could.
Moreover, whereas Jastrzebski et al.
(2020) only conjectured the sharpness “wants to” increase at the very beginning of training, we ﬁnd that this tendency exists throughout training.
Finally, the strongest prediction in Jastrzebski et al.
(2020) is that the maximum sharpness along the SGD optimization trajectory would be lower when the learning rate is larger or the batch size is smaller.
This paper demonstrates the much stronger result that during full-batch gradient descent with step size η , the sharpness settles right at (or just past) the numerical value 2 /η .
Contrary to our results, in Ghorbani et al.
(2019) a learning rate drop caused the sharpness to decrease rather than increase.
We believe that this is because they used SGD rather than GD.
Bounded Curvature Our paper shows that gradient descent is perfectly successful even at the
Edge of Stability, where the Bounded Curvature condition does not hold.
Previously, Xing et al.
(2018) observed that full-batch gradient descent eventually enters a regime where the training loss behaves non-monotonically, but did not describe when/why this occurs, and did not discuss the implications of this observation for optimization theory.
Lewkowycz et al.
(2020) demonstrated that an initial sharpness greater than 2 /η is not necessarily fatal for gradient descent, since gradient descent can “catapult” out of that region into a much ﬂatter region, rather than diverging entirely.
However, the implication in Lewkowycz et al.
(2020) is that the actual progress would occur in regions where the sharpness remains strictly less than 2 /η .
Our experiments demonstrate otherwise.
Dynamical stability At the Edge of Stability, the stability properties of gradient descent are (somehow) constantly preventing it from following gradient ﬂow into sharper regions of the loss landscape.
Several previous works have argued that the stability properties of optimization algorithms could potentially serve as a form of implicit bias in deep learning (Nar & Sastry, 2018;
Wu et al., 2018; Giladi et al., 2020; Jastrzebski et al., 2020; Mulayoff & Michaeli, 2020).
Our contribution here is to demonstrate that due to progressive sharpening, gradient descent’s implicit stability constraint “sharpness ≤ 2 /η ” is typically “tight” throughout large portions of training.
L -smoothness
The condition we call Bounded Curvature is an instance of a more general concept in optimization called L -smoothness.
Formally, an objective function f : R p
→ R is said to be (globally)
L -smooth if the gradient map is globally L -Lipschitz, i.e. if (cid:107)∇ f ( x ) − ∇ f ( y ) (cid:107) 2 ≤
L (cid:107) x − y (cid:107) 2 for all x, y ∈ R p .
For a twice-differentiable function, this is equivalent to the condition that the Hessian’s eigenvalues are no greater than L in magnitude, i.e. (cid:107)∇ 2 f ( x ) (cid:107) 2 ≤ L for all x ∈ R p (Bubeck, 2014)
The main consequence of L -smoothness is the descent lemma : for an Lsmooth function f , a gradient descent step with learning rate η is guaranteed to decrease the training objective by at least η (cid:0) 1 − 1 2 ηL (cid:1) (cid:107)∇ f ( x ) (cid:107) 2 .
In particular, if (cid:0) 1 − 1 2 ηL (cid:1) >
⇐⇒ η < 2 /L , then a gradient step of size η is guaranteed to decrease (rather than increase) the training objective.
Using the descent lemma, one can prove convergence results for gradient descent that only require the objective function to be L -smooth, and not convex.
A famous example is the result stated in e.g.
Nesterov (1998) Eq.
1.2.12 that gradient descent with step size η = 1 /L on an L -smooth function will eventually converge to a stationary point.
The most minimal version of this kind of analysis would only assume the Bounded Curvature condition (i.e. that L -smoothness holds only near the optimization trajectory, and the function is only L -smooth enough to ensure descent).
Our paper, however, demonstrates that at reasonable learning rates, gradient descent eventually transitions into a regime where even the Bounded Curvature condition does not hold.
Since our experiments focus on gradient descent, they do not necessarily imply that L -smoothness assumptions are unjustiﬁed when analyzing other optimization algorithms.
However, gradient descent is arguably the simplest optimization algorithm, so we believe that the fact that Bounded
Curvature fails even there should raise serious questions about the suitability of the L -smoothness assumption in neural network optimization more generally.
The “non-convex but L -smooth” setting is currently ubiquitous in optimization papers that are seemingly motivated by deep learning (Reddi et al., 2016; Agarwal et al., 2017; Zaheer et al., 2018;
Zhou et al., 2018; Chen et al., 2018; Li &
Orabona, 2019; Ward et al., 2019; You et al., 2020; Vaswani et al., 2019; Sankararaman et al., 2019;
Reddi et al., 2020; D´efossez et al., 2020;
Xie et al., 2020; Liu et al., 2020).
We think that authors may want to reconsider this practice in light of our results.
S HARPENING AND INSTABILITY
In this section, we characterize the trajectory traversed by gradient descent with a ﬁxed step size η , with the aim of describing when and why gradient descent enters the Edge of Stability.
If the sharpness at initialization is greater than 2 /η , then either training will diverge entirely, or gradient descent will “locally diverge” out of the immediate neighborhood of initialization, but recover in a ﬂatter area (e.g. see Figure 18 or 19).
Lewkowycz et al.
(2020) dubbed this a “catapult.”
On the other hand, if the initial sharpness is less than 2 /η , then at ﬁrst, gradient descent will roughlytrack the gradient ﬂow trajectory.
A crucial empirical property of gradient ﬂow on neural network training objectives is that gradient ﬂow has an overwhelming tendency to continually head in a direction of increased sharpness.
We call this phenomenon progressive sharpening .
By “overwhelming tendency,” we mean that gradient ﬂow does occasionally decrease the sharpness (see e.g.
Figure 28), but these brief decreases always seem be followed by a return to continual increase.
We do not know why progressive sharpening occurs; this is an important question for future work.
In
Appendix D we take a step towards reconciling progressive sharpening with NTK theory.
Progressive sharpening presents a complication for gradient descent.
Whereas gradient ﬂow can traverse arbitrarily sharp regions of the loss landscape without difﬁculty, gradient descent is unstable (to quadratic order) in regions where the sharpness is greater than 2 /η .
By this, we mean that gradient descent would diverge if run on the quadratic Taylor approximation.
In particular, suppose that starting from any iterate, we suddenly switched from running gradient descent on the real objective to running gradient descent on the Taylor approximation.
Let ∆ t be the displacement along the leading Hessian eigenvector and let λ be the sharpness.
Then ∆ t would evolve under gradient descent as ∆ t =
c
[(1 − ηλ ) t − 1] , where c is a constant (Appendix B).
This sequence diverges iff | 1 − ηλ | > 1 , which, since λ, η > 0 , means λ > 2 /η .
Thus, if the sharpness exceeds 2 /η , then gradient descent run on the quadratic Taylor approximation would oscillate with increasing magnitude along the direction of highest curvature, leading to eventual divergence.
The learning rate η determines whether or not gradient descent remains stable for the entire duration of training.
On the one hand, if η is so small that the sharpness along the entire gradient ﬂow trajectory, from start to end, remains less than 2 /η , then gradient descent will stably track the gradient ﬂow trajectory all the way to the end.
(Put another way, this occurs if η < 2 /λ max , where λ max is the maximum sharpness along the gradient ﬂow trajectory.)
On the other hand, if η > 2 /λ max , then gradient descent will eventually reach a point where the local sharpness crosses 2 /η , causing gradient descent to become destabilized.
Once this occurs, we observe that gradient descent enters a new regime we call the Edge of Stability.
In this regime, gradient descent continues to successfully optimize the loss, but in such a way that the sharpness hovers right at (or just beyond) 2 /η instead of increasing further.
In other words, at the Edge of Stability, gradient descent is essentially executing a local search subject to an active constraint that the sharpness remain ≤ 2 /η .
We now illustrate these points with a single example; the following section will show that they hold more generally.
We consider training a network on a subset of CIFAR-10 using full-batch gradient descent.
The network is a fully-connected architecture with two hidden layers of 200 units each (656,810 total parameters), and with ELU activations (Clevert et al., 2016).
We train on a subset of
CIFAR-10 of size 5,000.
We use the square loss for classiﬁcation (Hui & Belkin, 2020), encoding the target class with 1 and the other classes with 0; our ﬁndings also hold for cross-entropy loss, but with one minor twist.
We consider training ﬁnished when the loss reaches 0.05.
We train this network at three small learning rates η ∈ { 0.001, 0.0015, 0.002 } , plotting the training loss in Figure 2(a).
These three learning rates are small enough that gradient descent roughly follows the gradient ﬂow trajectory from start to ﬁnish, traveling at a speed given by the learning rate.
To illustrate that the gradient descent at these three learning rates is following the same path, in Figure 2(b) we plot the training loss by “time” = iteration × learning rate.
This allows us to directly compare the loss after, say, 1000 iterations at training rate 1e-3 with the loss after 500 iterations of learning rate 2e-3 (both are time 1).
When plotted by time, the three loss curves coincide exactly.
Figure 2(c) plots the evolution of the sharpness.
Observe that sharpness continually rises, except for a brief downward blip at the beginning of training.
This is progressive sharpening.
Notice that the maximum sharpness along the gradient ﬂow trajectory is just a bit under 1000, which means that learning rates of η = 2 / 1000 = 0 .
002 or smaller will be stable for the entire duration of training.
Crucially, all three of our learning rates were at or below this threshold; this is the reason why gradient descent followed the gradient ﬂow trajectory from start to ﬁnish.
[Figure]
[Figure] [Figure]
We now train the same network at larger learning rates, plotting the training loss in Figure 3(a).
For each learning rate, we draw a vertical dotted line at the iteration where the sharpness ﬁrst crosses 2 /η .
Observe that shortly after the sharpness crosses 2 /η , the training loss starts to behave nonmonotonically.
In Figure 3(b), we plot the training loss curves as a function of time.
Observe that gradient descent closely tracks the gradient ﬂow trajectory up until the point where the sharpness crosses 2 /η .
After this point, gradient descent peels off from the gradient ﬂow trajectory and follows a different route.
(We include additional evidence for this claim in Appendix C.)
In Figure 3(c), we plot the evolution of the sharpness over the course of training.
For each learning rate η , we draw a horizontal line at the value 2 /η .
Observe that once the sharpness hits 2 /η , it ceases to increase further, and instead hovers right at (or just a bit beyond) 2 /η — at the Edge of Stability.
At the Edge of Stability, gradient descent is constantly “trying” to increase the sharpness, but is being somehow restrained from doing so.
If we cut the learning rate, this restraint is removed, and gradient descent is freed to follow gradient ﬂow into a region of increased sharpness.
To demonstrate this, in Figure 4, we run gradient descent until it enters the Edge of Stability, then halve the learning rate and continue training.
Observe that as soon as the learning rate is cut, the sharpness starts to increase.
It only stops increasing once gradient descent is back at the Edge of Stability.
Momentum A natural question is whether the evolution of the sharpness can be precisely characterized for optimization algorithms other than gradient descent.
Our experience suggests that the answer is no.
We believe this is because even on quadratic functions, the stability of optimization algorithms other than gradient descent is not necessarily determined by a sharpness threshold.
There is one exception: gradient descent with momentum.
Similar to vanilla gradient descent, on quadratic objectives, momentum gradient descent (with either Polyak or Nesterov momentum) acts independently along each Hessian eigenvector, and will diverge along the leading eigenvector if and only if the sharpness exceeds a certain threshold (the “maximum stable sharpness”, or MSS).
We prove in Appendix B that for GD with learning rate η and momentum parameter β , the MSS is:
[Equation]
The Polyak result previously appeared in Goh (2017); the Nesterov one seems to be new.
It is interesting to note that MSS Polyak ( η, β ) is increasing in β , whereas MSS Nesterov ( η, β ) is decreasing in β .
That is, adding Polyak momentum gives gradient descent a higher tolerance for sharpness, whereas adding Nesterov momentum gives gradient descent a lower tolerance for sharpness.
In Figure 5, we train our network to completion with varying amounts β of both Polyak and Nesterov momentum, while keeping the learning rate η held ﬁxed.
Just like in the vanilla GD case, we see that the sharpness rises to the maximum stable sharpness and then remains approximately there.
Cross-entropy loss So far, we have focused on the squared loss.
The situation is mostly the same for the commonly-used cross-entropy loss, but with one difference: for cross-entropy loss, at the end of training, we often observe that the sharpness decreases.
This behavior seems to be caused by the well-known fact that at the end of training, gradient descent continually scales up the logits so as to drive the cross-entropy loss to zero (Soudry et al., 2018).
When this occurs, the sharpness drops, because the second derivative of the cross-entropy goes to zero as its input increases.
E XPERIMENTS
We now demonstrate our main point across different architectures, activation functions, loss functions, and tasks, and on networks both with and without batch normalization.
Across all of these training setups, we will show that if the initial sharpness is less than 2 /η , then the sharpness will tend to continually rise during training (progressive sharpening) until gradient descent reaches the
Edge of Stability, a regime where the sharpness hovers right at (or just above) 2 /η and the training loss behaves non-monotonically.
Training loss curves (and accuracies) are plotted in Appendix D.
In the ﬁrst row of Figure 6, we demonstrate that our ﬁndings hold across several different architectures — a ReLU CNN, a tanh CNN, and a ReLU fully-connected network — trained using the square loss on a subset of 5k examples from CIFAR-10.
In the second row, we demonstrate that our ﬁndings hold for the more commonly-used cross-entropy loss, the only difference being that the sharpness starts to decrease towards the end of training.
(Here we cease training once the loss reaches 0.05, but if training were allowed to run longer, the sharpness would continue to drop.)
In the third row, we verify that our ﬁndings hold on the full CIFAR-10 dataset, using three small-ish architectures (where we don’t train to completion).
Figure 1 shows the same point for the larger, realistic VGG-11 architecture (where we do).
In the experiments on full CIFAR-10, we compute the sharpness on a subsampled set of 5k datapoints.
Batch normalization (Ioffe & Szegedy, 2015) is known to have unusual optimization properties (Li & Arora, 2019), so it is natural to wonder whether batch-normalized (BN) networks still operate at the Edge of Stability.
The fourth row of Figure 6 demonstrates that the sharpness in BN networks behaves exactly the same under gradient descent as it does in non-BN networks.
(As is standard practice, we use ghost batch normalization (Hoffer et al., 2017) to make full-batch gradient descent feasible with limited GPU memory.)
Appendix E reconciles this ﬁnding with Santurkar et al.
Finally, in the ﬁfth row of Figure 6, we demonstrate that our results hold on settings beyond image classiﬁcation: a Transformer trained on the WikiText-2 language modeling task; a deep linear network trained on synthetic isotropic data; and a one-hidden-layer tanh network trained on a toy one-dimensional regression problem.
C ONCLUSION
In this paper, we have empirically demonstrated that gradient ﬂow (and gradient descent with sufﬁciently small step sizes) has an overwhelming tendency to continually move in a direction of increased sharpness.
As a consequence of this progressive sharpening phenomenon, during gradient descent at reasonable learning rates, the sharpness eventually rises to 2 /η , whereupon gradient descent departs the gradient ﬂow trajectory and enters the Edge of Stability, a regime where the sharpness hovers right at, or just above 2 /η .
In this regime, gradient descent continues to successfully optimize the loss, in deﬁance of the conventional wisdom in optimization.
We think that understanding precisely how gradient descent is able to succeed at the Edge of Stability should be viewed as an important open problem in optimization theory.
While we focused in this paper on full-batch gradient descent, we believe that some of our ﬁndings may also have analogues for SGD.
First, while it is obviously normal for the training loss to behave non-monotonically during SGD, we present a surprising experiment in Appendix A which suggests that SGD on neural networks rapidly “acclimates” to the learning rate and batch size in such a way that subsequent SGD steps do not even consistently decrease the training loss in expectation , even though SGD steps with a smaller learning rate or larger batch size would consistently decrease the loss in expectation.
We view this experiment as evidence that SGD on neural networks may behave more similar to gradient descent at the Edge of Stability than to classical analyses of SGD as in Bottou et al.
Second, much like how the learning rate in gradient descent implicitly constrains the sharpness along the optimization trajectory, the experiments in Jastrzebski et al.
(2020) indicate that a similar phenomenon may be at play during SGD, though SGD seems to implicitly constrain not just the sharpness but also the variance of the stochastic gradients.
For SGD, we currently lack a precise quantitative understanding of this implicit bias.
A fascinating open problem is whether
SGD satisﬁes some “steady state” condition analogous to the sharpness ≈ 2 /η rule that we have identiﬁed for full-batch gradient descent.
A S TOCHASTIC G RADIENT D ESCENT
We have demonstrated that during full-batch gradient descent with step size η , the sharpness rises until reaching the maximum stable sharpness 2 /η , where it approximately remains for the duration of training.
Unfortunately, prior work has shown that the evolution of the sharpness during SGD is not as easily characterized.
On the one hand, just like with gradient descent, the learning rate seems to affect the sharpness along the optimization trajectory, with larger learning rates (as well as smaller batch sizes) steering SGD towards less sharp regions of the loss landscape (Jastrzebski et al., 2019; 2020).
On the other hand, in behavior different from gradient descent, during SGD the sharpness sometimes does not equilibrate at any steady-state value, much less one that we can numerically predict; for example, in Figure 5(a) of Jastrzebski et al.
(2020), one can see that the sharpness gradually drifts downwards during SGD training.
We suspect that this is partly because the sharpness does not directly affect the stability of SGD unless one makes the strong assumption that the Hessian and the covariance matrix of per-example gradients share a leading eigenvector, as in Jastrzebski et al.
That said, we do suspect that there may exist some sense in which SGD, too, operates at the Edge of Stability.
In particular, we now demonstrate that SGD seems to “acclimate” to the learning rate and batch size in such a way that an actual step sometimes increases and sometimes decreases the loss in expectation, yet a step with a larger learning rate or smaller batch size would almost always increase the loss in expectation, and a step with a smaller learning rate or a larger batch size would almost always decrease the loss in expectation.
In Figure 7, we train the ELU network from Section 3 using SGD with learning rate 0.01 and batch size 32.
We periodically compute the training loss (over the full dataset) and plot these on the left pane of Figure 7.
Observe that the training loss does not decrease monotonically, but of course this is not surprising — SGD is a random algorithm.
However, what may be more surprising is that SGD is not even decreasing the training loss in expectation .
On the right pane of Figure 7, every 500 steps during training, we use the Monte Carlo method to approximately compute the expected change in training loss that would result from an SGD step (the expectation here is over the randomness involved in selecting the minibatch).
Observe that at many points during training, an SGD step would decrease the loss (as desired) in expectation, but at other points, and SGD step would increase the loss in expectation.
In Figure 8(a), during training of the same network, we compute the expected change in training loss that would result from taking an SGD step with the same learning rate used during training, but half the batch size used during training (i.e. 16).
We observe that an SGD step with half the batch size would consistently cause an increase in the training loss in expectation.
In Figure 8(b) we repeat this experiment, but with twice the batch size used during training (i.e. 64).
Notice that an SGD step with twice the batch size would consistently cause a decrease in the training loss in expectation.
In
Figure 8(c) and (d) repeat this experiment with the learning rate; we observe that an SGD step with a larger learning rate would consistently increase the training loss in expectation, while an SGD step with a smaller learning rate would consistently decrease the training loss in expectation.
To be clear, there was nothing special in absolute terms about the hyperparameters we used.
In
Figure 9(a), we train the same network with learning rate 0.01 and batch size 16 and evaluate the expected loss change during training.
Observe that some steps increase the loss in expectation, while others decrease the loss in expectation.
Figure 9(b)-(d) repeats this experiment with: (b) learning rate 0.01 and batch size 64, (c) learning rate 0.02 and batch size 32, and (d) learning rate 0.005 and batch size 32.
We view this experiment as evidence that SGD on neural networks behaves more similar to gradient descent at the Edge of Stability, than to conventional analyses of SGD like the ones in Bottou et al.
[Figure] [Figure]
loss in expectation, while taking an SGD step with a larger learning rate or a smaller batch size would have consistently increased the loss in expectation
[Figure]
c t e d l o ss c h a n g e
[Figure]
B S TABILITY OF GRADIENT DESCENT ON QUADRATIC FUNCTIONS
This appendix describes the stability properties of gradient descent (and its momentum variants) when optimizing the quadratic objective function
[Equation]
starting from the initialization x 0 = 0 .
Our notation is intentionally chosen to resemble a secondorder Taylor approximation.
To review, vanilla gradient descent is deﬁned by the iteration:
[Equation]
Meanwhile, gradient descent with Polyak (also called “heavy ball”) momentum (Polyak, 1964;
Sutskever et al., 2013; Goodfellow et al., 2016) is deﬁned by the iteration:
[Equation]
where v t is a “velocity” vector and 0 ≤ β < 1 is the momentum coefﬁcient.
For β =
the algorithm reduces to vanilla GD.
Finally, Nesterov momentum Sutskever et al.
; Goodfellow et al.
(2016) is an adaptation of
Nesterov’s accelerated gradient (Nesterov, 1983) for deep learning deﬁned by the iteration:
[Equation]
where v t is a “velocity” vector and 0 ≤ β < 1 is the momentum coefﬁcient.
For β =
the algorithm reduces to vanilla GD.
All three of these algorithms share a special property: on quadratic functions, they act independently along each Hessian eigenvector.
That is, if we express the iterates in the Hessian eigenvector basis, then in this basis the coordinates evolve independent from one another under gradient descent.
Proposition 1.
Consider running vanilla gradient descent on the quadratic objective (2) starting from x 0 = 0 .
Let q be the leading eigenvector of H and let λ be the sharpness.
(Assume that λ > 0 .)
Then the sequence { q T x t } will diverge iff λ > 2 /η .
Note that in the main text, we referred to q T x t as ∆ t .
Proof.
The update rule for gradient descent on this quadratic function is:
[Equation]
Therefore, the quantity q T x t evolves under gradient descent as:
[Equation]
Deﬁne ˜ x t = q T x t + 1 λ q T g , and note that { q T x t } diverges
if and only if { ˜ x t } diverges.
The quantity ˜ x t evolves under gradient descent according to the simple rule:
[Equation]
The sequence ˜ x t will diverge iff | 1 − ηλ | > 1 .
Since η > 0 and λ > 0 , it is impossible for 1 − ηλ > 1 , so ˜ x t diverges if and only if 1 − ηλ < − 1 ⇐⇒ λ > 2 /η .
Remark: to recover the formula given in the main text for the evolution of q T x t , note that q T x t =
˜ x t − 1 λ q T g
=
(1 − ηλ ) t ˜ x 0 − 1 λ q T g =
((1 − ηλ ) t − 1)( 1 λ q T g ) .
We now prove analogous results for Nesterov and Polyak momentum.
Theorem 1.
Consider running Nesterov momentum on the quadratic objective (2) starting from x 0 = 0 .
Let q be the leading eigenvector of H and let λ be the sharpness.
(Assume that λ > 0 .)
Then the sequence { q T x t } will diverge iff λ > 1 η (cid:16)
+2 β 1+2 β (cid:17) .
Proof.
The update rules for Nesterov momentum on this quadratic function are:
[Equation]
Deﬁne ˜ x t = q T x t + 1 λ q T g and ˜ v t =
q T v t .
Note that q T x t diverges iff ˜ x t diverges.
The quantities ˜ x t and ˜ v t evolve under Nesterov momentum as:
[Equation]
By noting that ˜ v t = ˜ x t − ˜ x t − 1 , we can rewrite this as a recurrence in ˜ x :
[Equation]
This is a linear homogenous second-order difference equation.
By Theorem 2.37 in Elaydi (2005), since λ > 0 and η > 0 and β < 1 , this recurrence diverges if and only if λ > 1 η (cid:16) 2+2 β 1+2 β (cid:17) .
The following result previously appeared in Goh (2017).
Theorem 2.
Consider running Polyak momentum on the quadratic objective (2) starting from x 0 = 0 .
Let q be the leading eigenvector of H and let λ be the sharpness.
(Assume that λ > 0 .)
Then the sequence { q T x t } will diverge iff λ > 1 η (2 + 2 β ) .
Proof.
The update rules for Polyak momentum on this quadratic function are:
[Equation]
Deﬁne ˜ x t = q T x t + 1 λ q T g and ˜ v t =
q T v t .
Note that q T x t diverges iff ˜ x t diverges.
The quantities ˜ x t and ˜ v t evolve under Polyak momentum as:
[Equation]
By noting that ˜ v t = ˜ x t − ˜ x t − 1 , we can rewrite this as a recurrence in ˜ x :
[Equation]
This is a linear homogenous second-order difference equation.
By Theorem 2.37 in Elaydi (2005), since λ > 0 and η > 0 and β < 1 , this recurrence diverges if and only if λ > 1 η (2 + 2 β ) .
None of these three results have used the fact that ( q , λ ) are the leading eigenvector and eigenvalue; they would also apply to any eigenvector/eigenvalue pair ( q i , λ i ) .
Thus, along any direction of positive curvature (i.e. λ i > 0 ), gradient descent (vanilla, or with Polyak/Nesterov momentum) diverges iff λ
i is strictly larger than some threshold, and converges iff λ i is strictly smaller than that threshold.
If the quadratic is convex (i.e. H (cid:31) 0 ), this means that gradient descent will converge to the solution iff the sharpness is strictly less than that threshold, and will diverge iff the sharpness is strictly greater than that threshold.
If H has negative eigenvalues, then gradient descent (with a positive step size) will diverge along the directions of negative curvature.
Of course, neural network training Hessians do have negative eigenvalues (Sagun et al., 2017; Ghorbani et al., 2019), yet gradient descent does not diverge during normal training.
The reason may be that the directions of negative curvature are somewhat illusory:
Alain et al.
(2019) walked along directions of negative curvature, and found that the negative curvature did not last for very long, whereas the directions of most positive curvature lasted for longer.
In other words, the quadratic Taylor approximation was accurate over greater distances for the directions of positive curvature than for the directions of negative curvature.
C A DDITIONAL FIGURES
[Figure]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
Note that even on the three examples where the η = 0 .
001 and η = 0 .
002 networks made different predictions, the differences in the class scores were slight.
For instance, on one of these test examples, the η = 0 .
001 network predicted 4 while the η = 0 .
002 network predicted 7; the scores were:
[Table] [Table]
[Table]
D I NFINITE W IDTH L IMITS
In this appendix, we conduct an experiment which aims to reconcile progressive sharpening with the theory of inﬁnite-width neural networks.
It is well-known that, under a certain network parameterization and initialization, when training inﬁnitely wide neural networks using gradient ﬂow or gradient descent with sufﬁciently small step sizes, the parameters move almost no distance from their initial position, and therefore the Hessian does not change from initialization (Jacot et al., 2018;
Lee et al., 2019; Jacot et al., 2020).
Accordingly, one might hypothesize that progressive sharpening might attenuate as networks (with NTK parameterization) become increasingly wide.
We consider fully-connected networks under the NTK parameterization (Lee et al., 2019) with two hidden layers and tanh activations.
We consider networks with widths 100, 200, 400, and 800.
We train these four networks on the Fashion-MNIST dataset using the cross-entropy loss, with a small learning rate (aiming to mimic gradient ﬂow).
We use the same learning rate for each network (this is sensible under the NTK parameterization).
We stop each network when the loss reaches a certain threshold.
In Figure 11(b), we plot the sharpness by iteration.
One can see that sharpening happens quicker for the narrow networks than for the wider networks.
However, there is a confounding factor: as can be seen from Figure 11(a), the narrower networks train faster.
To remove this confounding factor, in
Figure 11(c), we plot sharpness as a function of train loss.
Here one can clearly see that at every train loss, the narrower (resp: wider) networks have a higher (resp: lower) sharpness.
This observation aligns with NTK theory.
[Figure]
E R ECONCLING OUR RESULTS WITH S ANTURKAR ET AL . (2018)
In this appendix, we reconcile the apparent contradiction between our results and those of the wellknown work Santurkar et al.
Aiming to explain the effectiveness of batch normalization (BN), Santurkar et al.
(2018) conjectured that BN smooths the optimization landscape, in three different senses of the word “smooth”:
BN improves the Lipschitzness of the objective function
along the optimization
trajectory.
To
demonstrate this, at regular intervals during training, Santurkar et al.
) measured the change in loss that would arise from taking a step in the negative gradient direction.
(Note that the step size used here was different from the training step size.)
BN improves the gradient predictiveness of the objective function along the optimization trajectory.
Here, gradient predictiveness is deﬁned as the (cid:96) 2 distance between the gradient at the current iterate, and the gradient after taking a very large step in the negative gradient direction.
To demonstrate this, Santurkar et al.
(2018) measured the gradient predictiveness at regular intervals during training.
BN improves the effective smoothness along the optimization trajectory, deﬁned as the Lipschitz constant of the gradient in the direction of the negative gradient.
Given an objective function f , an iterate θ , and a step size α , the effective smoothness is deﬁned as
[Equation]
To demonstrate this,
Santurkar et al.
) measured the effective smoothness at regular intervals during training
Note that Santurkar et al.
(2018) did not conjecture that BN improves the L -smoothness (a.k.a decrease the sharpness) along the optimization trajectory — only that BN improves the effective smoothness, which is the L -smoothness in the negative gradient direction.
Our main paper demonstrates that BN does not improve the L -smoothness along the optimization trajectory, at least for full-batch gradient descent.
During gradient descent with step size η , for both
BN and non-BN networks, if the sharpness is less than 2 /η , then the sharpness eventually rises to approximately 2 /η , where it remains for the duration of training.
Thus, in the long run, for both BN and non-BN networks, the smoothness settles at approximately the same value — and moreover, this is precisely the value that voids the descent lemma (and arguments based on L -smoothness).
Moreover, we now demonstrate that BN does not improve the effective smoothness either (at least during full-batch gradient descent).
In Figure 12, we train a ReLU CNN both with (top row) and without (bottom row) batch normalization, at the same grid of learning rates, using the cross-entropy loss on a subset of 5k examples from CIFAR-10.
(For the BN network we use ghost batch normalization with two ghost batches of size 2,500, since a batch with all 5k examples does not ﬁt into
GPU memory.)
We measure both the sharpness/smoothness and the effective smoothness at regular intervals during training.
(We numerically evaluate the sup in the deﬁnition of effective smoothness using a grid of 10 evenly spaced points.)
Observe that for each learning rate, the effective smoothness behaves similar to the smoothness: it rises to 2 /η
and then remains approximately there, until the end of training, when both values drop because we are using the cross-entropy loss.
Thus, BN does not improve the effective smoothness along the optimization trajectory.
Despite this, one can see that the BN network trains substantially faster than the non-BN network.
Note that this ﬁnding is actually consistent with Figure 4(c) and Figure 11(d) in Santurkar et al.
(2018), which are meant to show, respectively, that both BN and their (cid:96) p normalization schemes improve effective smoothness on a VGG network.
These two plots show that for both BN and nonBN networks trained using SGD with step size η = 0 .
1 , the effective smoothness hovers around the value 20 = 2 /η , which is exactly what we observe.
Their plots do show that the effective smoothness behaves more regularly for the batch-normalized network than for the non-BN network.
But we strongly disagree with their interpretation of these two ﬁgures as demonstrating that BN improves the effective smoothness during training.
The third and ﬁnal piece of evidence in Santurkar et al.
(2018) for the conjecture that BN improves the effective smoothness is their Figure 9(c), which shows that BN improves the effective smooth-
[Figure]
ness on a deep linear network.
For this ﬁgure, the step size α used in their computation of effective smoothness was larger than the training step size η by a factor of 3 · 10 7 .
(For the VGG network, α was larger than η by a factor of 4.)
The effective smoothness at this massive distance has no bearing on training
F E XPERIMENTAL DETAILS F. T ASKS
We used the following tasks:
• CIFAR-10 (Krizhevsky, 2009): we standardize each channel by subtracting the mean and dividing by the standard deviation.
• CIFAR-10 5k : this is a subset of CIFAR-10 of size 5,000.
In particular, we retain every 10-th train example.
• Wikitext-2 : this is a language modeling dataset.
We use bptt = 35 , i.e. we predict text in contiguous blocks of 35 words.
We use only the ﬁrst 4970 (it is divisible by 35) training examples in the dataset.
• Deep linear network task : The task is to map the matrix X ∈ R n × d = [ x 1 , . . .
, x n ] T to the matrix Y ∈ R n × d = [ y 1 , . . .
, y n ] T using a deep linear network f : R d → R d that consists of a stack of L matrices each of size d × d .
Error is measured using the squared loss: 1 n (cid:80) ni =1 (cid:107) f
( x i ) − y
i (cid:107) 22 .
We use n = d = 50 .
We generate X as a random whitened matrix (i.e. 1 n X T X = I ) and generate Y =
XA T , where A ∈ R d × d is a random matrix of standard Gaussians.
To generate X as a random whitened matrix, we sample a 50 × 50 matrix of standard Gaussians, then we let X be √ 50 times the Q factor in the QR factorization of that matrix.
• Toy regression task : This is a one-dimensional regression task using the square loss.
We sample 20 datapoints uniformly spaced between − 1 and 1 , and we label these datapoints noiselessly using the Chebyshev polynomial of degree 5.
The dataset is plotted in Figure 13.
[Figure]
A RCHITECTURES
• VGG-11 (no BN) : we took the VGG-11 sized for CIFAR-10 available at https: //github.com/chengyangfu/pytorch-vgg-cifar10 , and we removed the dropout layers.
In Figure 1, we “warm-started” optimization by running GD for 400 iterations at learning rate 0.02.
The reason we do this is that loss landscape is very ﬂat at initialization, and if we did not do this, the network trained at, say, η = 0 .
00125 would basically not budge for the ﬁrst several thousand iterations of training.
• VGG-11 (BN) : same as above (though we do not need to warm start training).
Since GPUs do not have enough memory to do a full-batch step, we use ghost batch normalization (Hoffer et al., 2017) with 50 ghost batches of size 1,000.
• ReLU CNN :
The PyTorch code for this network is: nn.
Sequential( nn.
Conv2d(3, 32, bias=
True, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.
Conv2d(32,32, bias=
True, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.
Flatten(), nn.
Linear(2048,10, bias=True) )
We use the default PyTorch initialization.
• ReLU CNN (BN)
Same as the ReLU CNN,
but we insert a nn.
BatchNorm2d layer after each ReLU.
We use 2 ghost batches of size 2,500 each.
• Tanh CNN and Tanh CNN (BN) :
same as the ReLU CNN, except we use the tanh activation instead of ReLU.
• ReLU fully-connected network The PyTorch code for this network is: nn.
Sequential( nn.Linear(3072, 100, bias=
True)
, nn.ReLU(), nn.
Linear(100, 100, bias=
True) nn.ReLU(), nn.
Linear(100, 10, bias=
True) )
We use the default PyTorch initialization.
• Tanh fully-connected network :
Same as the ReLU fully-connected network, but with tanh activations.
• ELU fully-connected network :
Same as the ReLU fully-connected network, but with ELU activations, and with hidden layer size 200 rather than 100.
(There is no nefarious reason for the difference in hidden layer size; we just ran the ReLU and tanh experiments ﬁrst, and then decided to make the network more overparameterized.)
• Deep linear network : this is a deep linear network with 20 layers, each a 50 × 50 matrix.
We initialize each matrix using Xavier initialization: each entry is drawn from N (0 , 150 ) .
• Toy 1d tanh network : The PyTorch code for this network is: nn.
Sequential( nn.
Linear(1, 100, bias=
True), nn.
Tanh(), nn.
Linear(100, 1, bias=False), )
We initialize using Xavier initialization.
• Transformer : we use the Transformer from the Pytorch language modeling tutorial: https://github.com/pytorch/examples/tree/master/word_ language_model with ninp=200, nhead=2, nhid=200, nlayers=2, dropout=0.
G T HE EFFECT OF NETWORK WIDTH
In this appendix, we study the effect of network width on progressive sharpening, by training networks of varying widths on a 5k-sized subset of CIFAR-10.
In order to study the phenomenon of progressive sharpening independent of step size, we train networks using gradient ﬂow rather than gradient descent; to train a network using gradient ﬂow, we use the Runge-Kutta RK4 algorithm to numerically integrate the gradient ﬂow ODE.
We consider fully-connected networks with two hidden layers under both NTK parameterization and standard parameterization.
NTK parameterization For NTK-parameterized networks, Figure 14 plots the evolution of the sharpness during training.
For each of widths { 32 , 64 , 128 , 256 , 512 , 1024 } , we train from ﬁve different random initializations.
We plot evolution of the sharpness for each random initialization in a different color.
Observe from Figure 14 that progressive sharpening occurs to a greater degree for narrower networks than for wide networks.
[Figure]
[Figure]
In Figure we plot statistics of the training runs depicted in Figure 14.
For a given training run, let λ 0 be the sharpness at initialization and let λ max be the maximum sharpness along the gradient ﬂow trajectory.
In Figure 15 (left), for each width, we plot the mean and standard deviation (over ﬁve different random initializations) of λ max .
Observe that as the width increases, λ max decreases in mean (the standard deviation decreases as well).
In Figure 15 (right), for each width, we plot the mean and standard deviation (over ﬁve different random initializations) of the maximum sharpness gain λ max /λ 0 .
Observe that as the width increases, the maximum sharpness gain decreases in mean;
NTK theory predicts that the maximum sharpness gain will tend toward 1 as the width → ∞ .
Standard parameterization In Figure 16 and Figure 17, we repeat the experiments in the preceding section for standard parameterized networks.
Similar to NTK parameterization, we still see that the maximum sharpness λ max is larger for narrow networks.
However, for standard parameterization, it isn’t completely clear whether or not the maximum sharpness gain λ max /λ 0 converges toas the width → ∞ ; at the very least, the rate of convergence seems to be slower than it was for NTK parameterization
[Figure]
[Figure]
[Figure]
H C OMPLETE E XPERIMENT P LOTS
In this section, we present complete plots for the experiments described in Section 4.
The purpose of these experiments is to demonstrate the following points:
When sharpness is less than 2 /η , gradient descent has an overwhelming tendency to increase the sharpness ( progressive sharpening ).
So long as the sharpness is less than 2 /η , gradient descent at different learning rates roughly follows the same path, just at different speeds.
This is because gradient descent is roughly tracking the gradient ﬂow trajectory (see the animation in Bach (2020) for a visualization of what we mean).
The way we demonstrate this is by plotting the train loss and the sharpness by time; when plotted by time, the sharpness and training loss of GD at different learning rates almost exactly coincide.
If the sharpness rises to 2 /η (see the caveat below), gradient decent subsequently enters the Edge of Stability, a regime where (1) the training loss typically behaves non-monotonically, and (2) the sharpness hovers right at, or just above, the value 2 /η , rather than increasing further.
For sufﬁciently small learning rates η , gradient descent can terminate before the sharpness reaches 2 /η .
We will demonstrate that these small learning rates are highly suboptimal in terms of convergence speed — therefore, we say that they are not “reasonable.”
On some problems
it is feasible to actually train to completion at these small learning rates.
In this case, our experiments demonstrate that larger learning rates (which enter the Edge of Stability) train faster.
On other problems, sharpening seems to occur so rapidly that training to completion at these small learning rates is not a practical experiment for us to run; however, our gradient ﬂow thought experiment from section 3 implies that these small learning rates should exist.
There is one slight caveat.
On a few optimization problems (e.g. Figure 30), we observe that the gradient descent enters the Edge of Stability regime when the sharpness is a little bit smaller than 2 /η .
This is because we are only measuring the sharpness right at the iterates themselves.
However, if the sharpness in between the iterates (which we do not measure) exceeds 2 /η , then it is possible for a gradient step to increase the training loss
[Figure]
and for gradient descent to become destabilized
[Figure]
[Figure] [Figure]
fully-connected network trained with square loss on 5k subset of CIFAR-10. .
Here we use the cross-entropy loss, so the sharpness drops near the end of training
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
full CIFAR-10.
[Figure] [Figure]
[Figure] [Figure]