Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability
Authors: Jeremy Cohen
Simran Kaur
Yuanzhi Li
J Zico Kolter
Ameet Talwalkar
Venue: ICLR 2021 Poster (2021)
Keywords: optimization
trajectory
stability
sharpness
implicit bias
implicit regularization
L-smoothness
deep learning theory
science of deep learning

Abstract: 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 maximum eigenvalue of the training loss Hessian hovers just above the value $2 / \text{(step size)}$, and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales. Since this behavior is inconsistent with several widespread presumptions in the field of optimization, our findings raise questions as to whether these presumptions are relevant to neural network training. We hope that our findings will inspire future efforts aimed at rigorously understanding optimization at the Edge of Stability.

Published as a conference paper at ICLR 2021
GRADIENT DESCENT ON NEURAL NETWORKS TYPI- CALLY OCCURS AT THE EDGE OF STABILITY
Jeremy Cohen Simran Kaur Yuanzhi Li J. Zico Kolter' and Ameet Talwalkar*
Carnegie Mellon University and: Bosch Al Determined AI
Correspondence to: jeremycohen cma.
ABSTRACT
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 maximum eigenvalue of the training loss Hessian hovers just above the value 2/(step size), and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales. Since this behavior is inconsistent with several widespread presumptions in the field of op- timization, our findings raise questions as to whether these presumptions are rel- evant to neural network training. We hope that our findings will inspire future efforts aimed at rigorously understanding optimization at the Edge of Stability.
INTRODUCTION
Neural networks are almost never trained using (full-batch) gradient descent, even though gradient descent is the conceptual basis for popular optimization algorithms such as SGD. In this paper, we train neural networks using gradient descent, and find two surprises. First, while little known about the dynamics of neural network training in general, we find that in the special case of gradient descent, there is a simple characterization that holds across a broad range of network architectures and tasks. Second, this characterization is strongly odds with prevailing beliefs in optimization.
In more detail, as we train neural networks using gradient descent with step size 7), we measure the evolution of the sharpness the maximum eigenvalue of the training loss Hessian. Empirically, the behavior of the sharpness consistent across architectures and tasks: so long as the sharpness is less than the value 2/17, it tends to continually rise (3..)). We call this phenomenon progressive sharpening. The significance of the value 2/17 is that gradient descent on quadratic objectives is unstable if the sharpness exceeds this threshold ($2). Indeed, in neural network training, if the sharpness ever crosses 2/n,. gradient descent quickly becomes destabilized - that is, the iterates start to oscillate with ever-increasing magnitude along the direction of greatest curvature. Yet once
Fully-connected net on CIFAR-10 5k subset
VGG on CIFAR-10
ResNet on CIFAR-10
train 1015
3/50
train loss
300
2/50
2/12/
320
20100
2-200
2000
300
sharpnesss
100
shoppeess
shapeees
200
50
500
1000
1500
2500
6000
2000
6,000
Iteration
Iteration
Figure 1: Gradient descent typically occurs at the Edge of Stability. On three separate architec- tures, we run gradient descent at a range of step sizes 7), and plot both the train loss (top row) and the sharpness (bottom row). For each step size 7) observe that the sharpness rises to 2/1 (marked by the horizontal dashed line of the appropriate color) and then hovers right at, or just above, this value.

---

Published as a conference paper at ICLR 2021
this happens, gradient descent does not diverge entirely stall. Instead, it enters new regime we call the Edge of Stabiliry [ ($3.)) in which (1) the sharpness hovers right at, or just above, the value 2/17; and (2) the train loss behaves non-monotonically, yet consistently decreases over long timescales. In this regime, gradient descent is constantly "trying" to increase the sharpness, but is constantly restrained from doing so. The net effect is that gradient descent continues successfully optimize the training objective, but in such a way as to avoid further increasing the sharpness
In principle, is possible run gradient descent at step sizes 1) so small that the sharpness never rises to 2/n. However, these step sizes are suboptimal from the point of view of training speed, sometimes dramatically so. In particular, for standard architectures on the standard dataset CIFAR- 10, such step sizes are so small as to be completely unreasonable - at all reasonable step sizes, gradient descent eventually enters the Edge of Stability (see $ Thus, at least for standard networks on CIFAR- - the Edge of Stability regime should be viewed as the "rule," not the "exception."
As we describe in $5 the Edge of Stability regime is inconsistent with several pieces of conventional wisdom in optimization theory: convergence analyses based on L-smoothness monotone descent, quadratic Taylor approximations as a model for local progress, and certain heuristies for step size selection. We hope that our empirical findings will both nudge the optimization community away from widespread presumptions that appear to be untrue in the case of neural network training, and also point the way forward by identifying precise empirical phenomena suitable for further study.
Certain aspects of the Edge of Stability have been observed in previous empirical studies of full- batch gradient descent (Xing al. 2018: Wu et al. 2018); our paper provides unified explanation for these observations. Furthermore, @Jsttzzbski et al. (2020) proposed a simplified model for the evolution of the sharpness during stochastic gradient descent which matches our empirical observa- tions in the special case of full-batch SGD (i.e. gradient descent). However, outside the full-batch special case, there is no evidence that their model matches experiments with any degree of quanti- tative precision, although their model does successfully predict the directional trend that large step sizes and/or small batch sizes steer SGD into regions of low sharpness. We discuss SGD at greater length in S To summarize, while the sharpness does not obey simple dynamics during SGD (as does during GD), there are indications that the "Edge of Stability" intuition might generalize somehow to SGD, just in way that does not center around the sharpness.
2 BACKGROUND: STABILITY OF GRADIENT DESCENT ON QUADRATICS
In this section, we review the stability properties of gradient descent on quadratic functions. Later, we will see that the stability of gradient descent on neural training objectives partly well-modeled by the stability of gradient descent on the quadratic Taylor approximation.
On a quadratic objective function f(x f Ax + baxx + c, gradient descent with step size n will diverge any eigenvalue of A exceeds the threshold 2/n. To see why, consider first the one- dimensional quadratic fxx .azz? + bæ + with a > 0. This function has optimum x=' = - b/a. Consider running gradient descent with step size 7) starting from zo. The update rule is xt+1
I, na(,, + 6). which means that the error It x² evolves as (Ie+1 x*) (1 na) x x²). Therefore, the error at step t is (I,
x²) = (1 - (1-na)'( (No x )"), and so the iterate at step is I (1 7 x²) +x" If a > 2/n, then (1 na) < 11, 1, the sequence {2 will oscillate around x= with ever-incrgasing magnitude, and diverge.
x²
Now consider the general d-dimensional case. Let (q;.q:) be the i-th largest cigenvalue/eigeeveccor of A. As shown in Appendix B, when the gradient descent iterates (x,} are expressed in the spe- cial coordinate system whose axes are the eigenvectors of A, each coordinate evolves separately. In particular, the coordinate for each eigenvector qi namely (qi, xr), evolves according to the dy- namies of gradient descent on a one-dimensional quadratic objective with second derivative a;
1 'Tii nomenclature was inspired by the title of et al. (020))
"In the literature, the term "sharpness" has been used reter to variety of quantities, often connected to generalization (e-g. [Keskar al. (2016)) In this paper, "sharpness" strictly means the maximum eigenvalue of the training loss Hessian. We do not claim that this quantity has any connection to generalization.
Forr convex quadraties, this is "if and only if." However, if A has negative eigenvalue, then gradient descent with any (positive) step size will diverge along the corresponding eigenvector.
2

---

Published as a conference paper at ICLR 2021
stap
step
Therefore, if a, 2/1,, then the sequence ((qx)e)) will oscillate with ever- iver-incraasing magnitude; in this case, say that the iterates {x,} diverge along the divection qi- To illustrate, Figure 2 shows a quadratic function with eigenvalues d1 20 and 42 = 1. In Figure 2ia), we run gradient descent with step size 7 = 0.09; since 0 < a2 < d1 2/m, gradient descent converges along both
8
and 92- In Figure (b)). we use step size 1) = 0.11; since O < 42 < 2/n < di, gradient descent converges along Q2 yet diverges along qi, so diverges overall.
Figure 2: Gradient descent on a quadratic with eigenvalues d1 20 and dg
Polyak momentum (Polyak 1964) and Nesterov momentum (Nesterov) 1983 Sutskever et al. 2013)
are notable variants of gradient descent which often improve the convergence speed. On quadratic functions, these two algorithms also diverge if the sharpness exceeds a certain threshold, which we call the "maximum stable sharpness," or MSS. In particular, we prove in Appendix [ that gradient descent with step size 1 and momentum parameter diverges if the sharpness exceeds:
1 PoSyuly (n, 3) (2 (2+ MSSNestero (1, B) 17 1 11++7 2
( 2 + 2,8
7
The Polyak result previously appeared in Goh (2017); the Nesterov one seems to be new. Note that
this discussion only applies to full-batch gradient descent. As we discuss in $6 several recent papers have proposed stability analyses for SGD (Wu et al., 2018 Jastrzębski et al. 2020).
Neural network training objectives are not globally quadratic. However, the second-order Taylor approximation around any point Xo in parameter space is quadratic function whose "A" matrix is this Hessian exceeds 2/n, gradient descent with step size
the Hessian at Xg. If any eigenvalue would diverge if run this quadratic function the iterates would oscillate with ever-increaste magnitude along the corresponding eigenvector. Therefore, at any point Xg in parameter space where the sharpness exceeds 2/ gradient descent with step size y would diverge if run on the quadratic Taylor approximation to the training objective around Xo.
3 GRADIENT DESCENT ON NEURAL NETWORKS
In this section, we empirically characterize the behavior of gradient descent on neural network train- ing objectives. Section 4 will show that this characterization holds broadly.
3.1 PROGRESSIVE SHARPENINGG
When training neural networks, it seems be general rule that so long the sharpness is small enough for gradient descent to be stable (< 2/n, for vanilla gradient descent), gradient descent has overwhelming tendeney to continually increase the sharpness. We call this phenomenon progressive sharpening. By "overwhelming tendency," we mean that gradient descent can occasion- ally decrease the sharpness (especially the beginning of training), but these brief decreases always seem be followed by return to continual increase. Jastrzębski et al.) (2020) previously hypothesized (in their Assumption 4) that a similar phenomenon may hold for SGD, but the evidence for, and the precise scope of, this effect are both currently far clearer for gradient descent than for SGD.
(a) Mean Squared Error (MSE) loss
(b) Cross entropy loss
222
435
2/n
2in
02.
38
10020
20002
30000
22000
33
2200 Iteration
4000
2000
4002
Iteration
Iteration
Reration
Figure 3: So long the sharpness is less than 2/7, it tends to continually increase during gradient descent. We train a network to completion (99% accuracy) using gradient descent with a very small step size. We consider both MSE loss (left) and cross-entropy loss (right).
3

---

Published as a conference paper at ICLR 2021
Progressive sharpening is illustrated in Figure 3 Here, we use (full-batch) gradient descent to train a network on a subset of 5,000 examples from CI and we monitor the evolution of the sharpness during training. The network is fully-comnected architecture with two hidden layers of width 200, and tanh activations. In Figure aa), we train using the mean squared error loss for classification (Hui & Belkin 2020), encoding the correct class with 1 and the other classes with O. We use the small step size of = 2/600, and stop when the training accuracy reaches 99%. We plot both the train loss and the sharpness, with a horizontal dashed line marking the stability threshold 2/n. Observe that the sharpness continually rises during training (except for a brief dip at the beginning). This is progressive sharpening. For this experiment, we intentionally chose step size y small enough that the sharpness remained beneath 2/1/ for the entire duration of training.
Cross-entropy. When training with cross-entropy loss, there is an exception to the rule that the sharpness tends to continually increase: with cross-entropy loss, the sharpness typically drops at the end of training. This behavior can be seen in (i), where we train the same network using the cross-entropy loss rather than MSE. This drop occurs because once most data points are classified correctly, gradient descent tries to drive the cross-entropy loss to zero by scaling up the margins, as detailed in Soudry et al. 2018). As explain in Appendix C this causes the sharpness to drop.
The effect of width. It is known that when networks parameterized in a certain way (the "NTK parameterization") are made infinitely wide, the Hessian moves vanishingly small amount during training (Jacot et al. 2018 Lee et al. 2019 Li & Liang 2018), which implies that no progressive sharpening occurs. In Appendix DJ we experiment with networks of varying width, under both NTK and standard parameterizations, We find that progressive sharpening occurs to lesser degree as net- works become increasingly wide. Nevertheless, our experiments in $ demonstrate that progressive sharpening occurs to a dramatic degree for standard architectures on the standard dataset CIFAR-1O.
We do not know why progressive sharpening occurs, or whether "sharp" solutions differ in any important way from "not sharp" solutions. These are important questions for future work. Note that Mulayoff & Michaeli (2020) studied the latter question the context of deep linear networks.
3.2 THE EDGE OF STABILITY
In the preceding section, we ran gradient descent using step sizes so small that the sharpness never reached the stability threshold 2/n. In Figure 4[a), start to train the same network at the larger step size of 1 = 0.01, and pause training once the sharpness rises to 2/1 = 200. Recall from $2 that in any region where the sharpness exceeds 2/1,, gradient descent with step size n would be unstable if run on the quadratic Taylor approximation to the training objective - the gradient descent iterates would oscillate with ever-increasing magnitude along the leading Hessian eigenvector. Empirically, we find that gradient descent on the real neural training objective behaves similarly - first. Namely, let qi be the leading Hessian eigenvector at the iteration where the sharpness reaches 2/n. In Figure [b), we resume training the network, and we monitor both the train loss and the quantity (qı, x,) for the next 215 iterations. Observe that (q) xx) oscillates with ever-increasing magnitude, similar to the divergent quadratic example in Figure 2b). At first, these oscillations are too small to affect the objective appreciably, and so the train loss continues to monotonically decrease. But eventually, these oscillations grow big enough that the train loss spikes.
train before sharpness crosses 20g
[b) train loss and (R) for first 235 steps after sharpness crosses 2/g
(c) train loss and (q) fortirst 1000 steps after sharpnets cresses 2/n
0.5
0225
0:227
0200
2.200
3.02
0.02
200
0,00
:11
[0)
0.00
1200
2002
3000
1000
150
200
300
403
coo
820
1300
sharpneet
Figure 4: Once the sharpness crosses 2/1, gradient descent becomes destabilized. We run gradi- ent descent at 17 0.01. (a) The sharpness eventually reaches 2/n. (b) Once the sharpness crosses 2/n, the iterates start to oscillate along qi with ever-increasing magnitude. (c) Somehow, GD does not diverge entirely; instead, the train loss continues decrease, albeit mon-oonotonically.
4

---

Published as a conference paper at ICLR 2021
Once gradient descent becomes destabilized in this manner, classical optimization theory gives no clues as to what will happen next. One might imagine that perhaps gradient descent might diverge entirely, or that gradient descent might stall while failing to make progress, that gradient descent jumhtjump to a flatter region and remain there. In reality, none of these outcomes occurs. In Figure 4fc), we plot both the train loss and (q1, for 1000 iterations after the sharpness first crossed 2/1-- Observe that gradient descent somehow avoids diverging entirely. Instead, after initially spiking around iteration 215, the train loss continues to decrease, albeit non-monotonically.
This numerical example is representative. In general, after the sharpness initially crosses 2/1, gradi- ent descent enters a regime we call the Edge of Stability, in which (1) the sharpness hovers right at, or just above, the value 2/ 17; and (2) the train loss behaves non mon-mooottn over short timescales, yet decreases consistently over long timescales. Indeed, in Figure 5 we run gradient descent at a range of step sizes using both MSE and cross-entropy loss. The left plane plots the train loss curves, with a vertical dotted line (of the appropriate color) marking the iteration where the sharpness first crosses 2/1z. Observe that the train loss decreases monotonically before this dotted line, but be- haves non non-monotonically afterwards. The middle plane plots the evolution of the sharpness, with a horizontal dashed line (of the appropriate color) at the value 2/1z. Observe that once the sharpness reaches 2/13. it ceases to increase further, and instead hovers right at, or just above, the value 2/n for the remainder of training. (The precise meaning of "just above" varies: in Figure [5] for MSE loss, the sharpness hovers just a minuscule amount above 2/1,. while for cross-entropy loss, the gap between the sharpness and 2/n) is small yet non -on-minisculc.)
At the Edge of Stability, gradient descent is "trying" to increase the sharpness further, but is being restrained from doing so. To demonstrate this, in Figure / we train at step size 2/200 until reaching the Edge of Stability, and then at iteration 6,000 (marked by the vertical black line), we drop the step size to 17- 2/300. Observe that after the learning rate drop. the sharpness immediately starts to crease, and only stops increasing once gradient descent is back at the Edge of Stability. Appendix O repeats this experiment on more architectures. Intuitively, gradient descent with fixed step sizes acts like a constrained optimization algorithm: the use of step size 1) imposes an implicit 2/1 constraint on the sharpness (Nar& Sastry, 2018), and at the Edge of Stability this constraint is "active."
Observe from Figure 5 that there do exist step sizes 7) (in purple) small enough that the sharpness never rises to 2/1.. We call such a step size stable. However, observe that with cross-entropy loss, takes 3700 iterations to train at the stable step size in purple, but only 1000 iterations to train the larger step size in blue. In general, we always observe that stable step sizes are suboptimal in terms
train loss
sharpness
sharpness (by time)
MSE loss
0.6
4.00
2/500
train 1955
0.4
2/300
200
0.2
2/100
0.0
5000
10000
15000
20000
0
5000
10000
15000
20000
25
75
100
125
teratio
teration
time Iteration
150
150
cross-entropy
2.0
2/150
[8 1.5
Shappeoss
1.00
Sharpress
100
2/120
2/90
10
2/60
2/30
1000
1000
3000
Figure 5: After the sharpness reaches 2/m, gradient descent enters the Edge of Stability. A network is trained with gradient descent at range of step sizes (see legend), using both MSE loss (top row) and cross-entropy (bottom row). Left: the train loss curves, with vertical dotted line at the iteration where the sharpness first crosses 2/1,. Center: the sharpness, with a horizontal dashed line at the value 2/n. Right: sharpness plotted by time (= iteration x 1)) rather than iteration.
5

---

Published as a conference paper at ICLR 2021
momenium gradiens descert 1/"
Shargness
300
Fobal. P-62
shampess
cainted
0.4
200
330
Noderre $-85
Noteree 2-15
0.2
100
6500
2000
4000
5000
5030
10000
2000
4000
6000
5000
20002
Conetion
teracion
heeattio
Figure 6: Momentum. We run GD with step size 0.01 and Polyak or Nesterov momentum at various 3. For each algorithm, the horizontal dashed line marks the MSS from Equation !
Figure 7: After learning rate drop, progressive sharp- ening resumes. We start training at rj = 2/200 (orange) and then after 6000 iterations (dotted vertical black line), we cut the step size to vj = 2/300 (green). Observe that as soon the step size is cut, the sharpness starts rise.
of convergence speed. In fact, in we will see that for standard networks on CIFAR-10, stable step sizes so suboptimally small that they are completely unreasonable.
The "Edge of Stability" effect generalizes to gradient descent with momentum. In Figure 6 we train using gradient descent with step size r 0.01, and varying amounts of either Polyak Nesterov momentum. Observe that in each case, the sharpness rises until reaching the MSS given by Equation and then plateaus there. Appendix N has more momentum experiments.
In Appendix [P] we briefly examine the evolution of the next few Hessian eigenvalues during gradient descent. We find that each of these eigenvalues rises until plateauing near 2/m.
Prior work. Aspects of the Edge of Stability have been observed previously in the literature. Wu al. (2018) noted that the sharpness at the solution reached by full-batch gradient descent was not just less than 2/m,, as was expected due to stability considerations, but was mysteriously apprex- imately equal to 2/1 In retrospect, we can attribute this observation to progressive sharpening. Xing et al. (2018) observed that full- bull-batch gradient descent eventually enters a regime (the Edge of Stability) which the training loss behaves non-monotonically, and the iterates oscillate along the direction of largest curvature; however, they did not relate this regime to the sharpness. [Lekkowcz cet al. (2020) found that in neural network training, if the sharpness at initialization is larger than 2/7, then after becoming initially destabilized, gradient descent does not always diverge entirely (as the quadratic Taylor approximation would suggest), but rather sometimes "catapults" into a flatter region that is flat enough to stably accommodate the step size. It seems plausible that whichever properties of neural training objectives permit this so-called "'atapult" behavior may also be the same properties that permit successful optimization at the Edge of Stability. Indeed, optimization at the Edge of Stability can conceivably be viewed as a never-ending series of micro-catapults. As we discuss at greater length in L several papers (Jastrzebski et al., 2017 2019) have observed that large step sizes steer stochastic gradient descent into less sharp regions of the loss landscape, and Jastrzęhski et al. (2000) attributed this effect to the stability properties of SGD. Finally, our pre- cise characterization of the behavior of the sharpness during full-batch gradient descent adds to growing body of work that empirically investigates the Hcssian spectrum of neural networks (Sagun)
ct al. 2017 Ghorbani et al. 2019 Li et al. 2020a Papyan 2018 2019 2020
3.3
THE GRADIENT FLOW TRAJECTORY
In the right pane of Figure 5 we plot the evolution of the sharpness during gradient descent, with "time" iteration X 17. rather than iteration, on the x-axis. This allows us to directly compare the sharpness after, say. 100 itcrations at = 0.01 to the sharpness after 50 iterations at = 0.02; both are time Observe that when plotted by time, the sharpnesses for gradient descent different step sizes coincide until the time where each reaches 2/1,. This is because for this network, gradient descent at 7 = 0.01 and gradient descent at = 0.02 initially travel the same path (moving at a speed proportional to 17) until each reaches the point on that path where the sharpness hits 2/1,. This path is the gradient flow trajectory. The gradient flow solution at time is defined as the limit as 17 - of the gradient descent iterate at iteration t/n (if this limit exists). The empirical finding of interest is that for this particular network, gradient descent does not only track the gradient flow trajectory in the limit of infinitesimally small step sizes, but for any step size that is less than 2/ 2/sharpness.

---

Published as a conference paper at ICLR 2021
We can numerically approximate gradient flow trajectories by using the Runge-Kutta RK4 algorithm (Press et al. 1992) to numerically integrate the gradient flow ODE. Empirically, for many but not all networks studied in this paper, we find that gradient descent at any step size closely tracks the Runge-Kutta trajectory until reaching the point on that trajectory where the sharpness hits 2/1-- (This sometimes occurs even for networks with ReLU activations or max-pooling, which give rise to training objectives that are not continuously differentiable, which means that the gradient flow trajectory is not necessarily guaranteed to exist.) For such networks, the gradient flow trajectory provides a coherent framework for reasoning about which step sizes will eventually enter the Edge of Stability. Let λu be the sharpness at initialization, and let Amax be the maximum sharpness along the gradient flow trajectory. If 1) < 2/Amax, then gradient descent will stably track the gradient flow trajectory for the entire duration of training, and will never enter the Edge of Stability. On the other hand, if 1 E [2/ Xmax, 2/ then gradient descent will stably track the gradient flow trajectory only until reaching the point on that trajectory where the sharpness hits 2/ v/; shortly afterwards, gradient descent will become destabilized, depart the gradient flow trajectory, and enter the Edge of Stability.
4 FURTHER EXPERIMENTS
Section 3 focused for exposition on a single architecture and task. In this section, we show that our characterization of gradient descent holds broadly across a wide range of architectures and tasks. We detail several known caveats and qualifications in Appendix [A]
Architectures. In Appendix D we fix the task of training a 5k subset of CIFAR-10, and we systemat- ically vary the network architecture. We consider fully-connected networks, as well as convolutional networks with both max- max-poolling and average pooling. For all of these architectures, we consider tanh, ReLU, and ELU activations, and for fully-connected networks we moreover consider softplus and hardtanh- eleven networks in total. We train each network with both cross-entropy and MSE loss. In each case, we successfully reproduce Figure 5
Since batch normalization (loffe & S is known to have unusual optimization properties (Li & Arora 2019), it is natural to wonder whether our findings still hold with batch normalization. In Appendix we confirm that they do, and we reconcile this point with Santurkar et al. (2018).
Tasks. In Appendix D we verify our findings on: (1) Transformer trained on the WikiTexx-- 2 language modeling task; (2) fully-connected tanh networks with one hidden layer, trained on a synthetic one-dimensional toy regression task; and (3) deep linear networks trained on Gaussian data. In each case, the sharpness rises until hovering right at, or just above, the value 2/n-
Standard networks on CIFAR-10. In Appendix M we verify our findings on three standard ar- chitectures trained on the full CIFAR-10 dataset: a ResNet with BN, a VGG with BN, and a VGG without BN. For all three architectures, we find that progressive sharpening occurs to dramatic degree, and, relatedly, that stable step sizes are dramatically suboptimal. For example, when we train the VGG-BN to 99% accuracy using gradient flow / Runge-Kutta, Kutta, find that the sharpness rises from 6.3 at initialization to peak sharpness of 2227.6. 66. Since this is an architecture for which gradient descent closely hews the gradient flow trajectory, size for gradient descent would need to be less than 2/2227.6 = 0.000897. Training finishes at time 14.91, so gradient descent at any stable step size would require at least 14.91/0. /0.000897 16, 622 iterations. Yet empirically, this network can be trained to completion at the larger, "Edge of Stabil- ity" step size of 1) - 0.16 in just 329 iterations. Therefore, training a stable step size is suboptimal by factor of at least 16622/ /229 = 50.5. The situation is similar for the other two architectures we consider. In short, for standard architectures on the standard dataset CIFAR- 10, stable step sizes are not just suboptimally small, they are so suboptimal as to be completely unreasonable. For these networks, gradient descent at any reasonable step size eventually enters the Edge of Stability regime.
conclude that any stable step
Tracking gradient flow. Recall that for some (but not all) architectures, gradient descent closely hews to the gradient flow trajectory so long as the sharpness is less than 2/1. Among the archi- tectures considered in Appendix D we found this to be true for the architectures with continuously differentiable components, as well as some, but not all, with ReLU, hardtanh, and max max-pooling. Among the architectures in Appendix [ we found this to be true for the tanh network, but not for the deep linear network or the Transformer. Finally, we did find this to be true for the three standard architectures in Appendix M even though those architectures ReLU.
7

---

Published as a conference paper at ICLR 2021
5 DISCUSSION
We now explain why the behavior of gradient descent at the Edge of Stability contradicts several pieces of conventional wisdom in optimization.
At reasonable step sizes, gradient descent cannot be analyzed using (even local) T-smoothness
Many convergence analyses of gradient descent assume a bound on the sharpness - either globally or, at the very least, along the optimization trajectory. This condition, called L-smoothness (Nes-) terov (1998), is intended to guarantee that each gradient step will decrease the training objective by a certain amount; the weakest guarantee is that if the local sharpness is less than 2/1,, then gradient step with size is guaranteed to decrease (rather than increase) the training objective. At a bare minimum, any convergence analysis of gradient descent based on L-smoothness will require the sharpness along the optimization trajectory to be less than 2/n- Yet to the contrary, at the Edge of Stability, the sharpness hovers just above 2/n. Therefore, at any step size for which gradient descent enters the Edge of Stability (which, realistic architectures, includes any reasonable step size), gradient descent cannot be analyzed using L-smoothness, Li et al. (2020b) previously argued that convergence analyses based on L-smoothness do not apply to networks with both batch normaliza- tion and weight decay; our paper empirically extends this to neural networks without either.
T-smoothness may be inappropriate when analyzing other optimization algorithms too It is common for optimization papers seemingly motivated by deep learning to analyze algorithms un-
der the "non-convex but L-smooth" setting (Reddi et al. 2016
Agarwal et al. 2017] Zaheer et al,
2018
Zhou et al. 2018 Chen et al. 2019 Li & Orabona 2019 Ward et al. [019 You et al. 2020
Vaswani et al. 2019] Sankararaman et al. 2019 Reddi et [2021] Défossez et al. 2020 Xie et al,
2020 Liu et al., 2020 Defazio 2020). Since our experiments focus on gradient descent, it does not
necessarily follow that L-smoothness assumptions are unjustified when analyzing other optimiza- tion algorithms. However, gradient descent is arguably the simplest optimization algorithm, so we believe that the fact that (even local) L-smoothness fails even there should raise serious questions about the suitability of the L-smoothness assumption in neural network optimization more generally. In particular, the burden of proof should be on authors to empirically justify this assumption.
At reasonable step sizes, gradient descent does not monotonically decrease the training loss In neural network training, SGD does not monotonically decrease the training objective, in part due to minibatch randomness. However, it is often assumed that full-batch gradient descent would monotonically decrease the training objective, were used train neural networks. For example, Zhang et al. (2020) proposed "relaxed L-smoothness" condition that is less restrictive than standard I-smoothness, and proved a convergence guarantee for gradient descent under this condition which asserted that the training objective will decrease monotonically. Likewise, some neural network analyses such as Allen-Zhu et al.] (2019) also assert that the training objective will monotonically decrease. Yet, at the Edge of Stability, the training loss behaves non-monotonically over short timescales even as it consistently decreases over long timescales. Therefore, convergence analyses which assert monotone descent cannot possibly apply gradient descent reasonable step sizes.
The Edge of Stability is inherently non-quadratic It tempting to try to reason about the behav- ior of gradient descent on neural network training objectives by analyzing, as a proxy, the behavior of gradient descent on the local quadratic Taylor approximation (LeCun et al. 1998). However, at the Edge of Stability, the behavior of gradient descent on the real neural training objective is irrec- oncilably different from the behavior of gradient descent on the quadratic Taylor approximation: the former makes consistent (if choppy) progress, whereas the latter would diverge (and this divergence would happen quickly, we demonstrate in Appendix (E). Thus, the behavior of gradient descent at the Edge of Stability is inherently non-quadratic.
Dogma for step size selection may be unjustified An influential piece of conventional wisdom concerning step size selection has its roots in the quadratic Taylor approximation model of gradient descent. This conventional wisdom (LeCun et al.) 1993] 1998 Schaul et al. 2013) holds that the sharpness step t is de, then the current step size 7), must be set greater than 2/X, (in order to prevent divergence); and furthermore, barring additional information about the objective function, that 7/2 should optimally be set to 1/>e. Our findings complicate this conventional wisdom.
8

---

Published as a conference paper at ICLR 2021
To start, it is nearly impossible to satisfy these prescriptions with fixed step size: for any fixed (and reasonable) step size T = 17. progressive sharpening eventually drives gradient descent into regions where the sharpness is just bit greater than 2/17 which means that the step size r, is purportedly impermissible. Furthermore, in Appendix [ we try running gradient descent with the purportedly optimal ? = 1/A, rule, and find that this algorithm is soundly outperformed by the purportedly impermissible baseline of gradient descent with a fixed ft- = 1/Ao step size, where A the sharpness at initialization. The de = 1/A, rule continually anneals the step size, and in so doing ensures that the training objective will decrease at each iteration, whereas the fixed 17: = 1/Ae step size often increases the training objective. However, this non- non-monotyniccty turns out to be a worthwhile price to pay in return for the ability to take larger steps.
6 STOCHASTIC GRADIENT DESCENT
Our precise characterization of the behavior of the sharpness only applies to full-batch gradient descent. In contrast, during SGD, the sharpness does not always settle at any fixed value (Appendix G). let alone one that can be numerically predicted from the hyperparameters. Nevertheless, prior works (Jastrzebski et al., 2017; 2019 2020) have demonstrated that large step sizes do steer SGD into regions of the landscape with lower sharpness; the 2/17 rule we have identified for full-batch gradient descent is a special case of this observation. Furthermore, small batch sizes also steer SGD into regions with lower sharpness (Keskar al. 2016 Jastrzębski et al. 2017), we illustrate in Appendix G Jastrzębski et al. (2020) attributed this phenomenon to the stability properties of SGD.
Even though our findings only strictly hold for gradient descent, they may have relevance to SGD) as well. First, since gradient descent is a special case of SGD, any general characterization of the dynamics of SGD must reduce to the Edge of Stability in the full-batch special case. Second, there are indications that the Edge of Stability may have some analogue for SGD. One way to interpret our main findings is that gradient descent "acclimates" to the step size in such a way that each update sometimes increases and sometimes decrease the train loss, yet an update with a smaller step size would consistently decrease the training loss. Along similar lines, in Appendix H we demonstrate that SGD "acclimates" to the step size and batch size in such a way that each SGD update sometimes increases and sometimes decreases the training loss in expectation, yet an SGD update with smaller step size or larger batch size would consistently decrease the training loss in expectation.
In extending these findings to SGD, the question arises of how to model "stability" of SGD. This is a highly active area of research. Wu et al. (0018) proposed modeling stability in expectation, and gave a sufficient (but not necessary) ) criterion for the stability of SGD in expectation. Building on this framework, Dastrzębski al. (2020) argued that if the Hessian is aligned with the second moment matrix of per-example gradients, then SGD is stable so long as a certain expression (involving the sharpness) is below certain threshold. In the special full-batch case, their criterion reduces to the sharpness being beneath 2/1 - a constraint which we have shown is "tight" throughout training. However, in the general SGD case. there is no evidence that their stability constraint is tight through- out training. Giladi et al. (020) showed that the generalization gap in asynchronous SGD can be mostly ameliorated by setting the step size so as to ensure that stability properties in expectation remain identical to those of a well-tuned implementation of synchronous SGD. Finally, a number of papers have attempted to mathematically model the propensity of SGD to "escape from sharp
minima" (Hu et al. 2017 Zhu et al. 2019 Xie et al. 2021).
7 CONCLUSION
We have empirically demonstrated that the behavior of gradient decent on neural training objectives both surprisingly consistent across architectures and tasks, and surprisingly different from that envisioned in the conventional wisdom. Our findings raise number of questions. Why does pro- gressive sharpening occur? At the Edge of Stability, by what mechanism does gradient descent avoid diverging entirely? Since the conventional wisdom for step size selection is wrong, how should the gradient descent step size be set during deep learning? Does the "Edge of Stability" effect generalize in some way to optimization algorithms beyond gradient descent, such SGD? We hope inspire future efforts aimed at addressing these questions.

---

Published as a conference paper at ICLR 2021
8 ACKNOWLEDGEMENTS
This work was supported in part by DARPA FA875017C0141, the National Science Foundation grants IIS1705121 and IIS1838017, an Amazon Web Services Award, JP Morgan A.I. Research Faculty Award, a Carnegie Bosch Institute Research Award, Facebook Faculty Research Award, and a Block Center Grant. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of DARPA, the National Science Foundation, or any other funding agency.
REFERENCES
Naman Agarwal, Zeyuan Allen-Zhu, Brian Bullins, Elad Hazan, and Tengyu Ma. Finding approxi- mate local minima faster than gradient descent. Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pp. 1195-1199, 2017.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pp. 242-252. PMLR, 2019.
David Barrett and Benoit Dherin. Implicit gradient regularization. In International Confer- ence Learning Representations, 2021. URL http:////oppereeeeeeeeeeeeeeeeee
3q5IqUrkcF
Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Reviews, 60(2):223-311 2018.
cestou--crtisn nocedal 2018
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. In International Conference Learning Representa- tions, 2019, URL htthttt:s////pneeeeeet
wwtp:/eett/nniiiieoeeeeoeooo net forum? id=H1x -dxm21i-x==0xxxx1 x309tm
Aaron Defazio. Understanding the role of momentum in non-convex optimization: Practical insights from Iyapunov analysis, 2020.
Alexandre Défossez, Léon Bottou, Francis Bach, and Nicolas Usunier. On the convergence of adam and adagrad. arXiv preprint arXiv:2003.022300 2020.
Saber Elaydi. An introduction to difference equations. Springer Science & Business Media, 2005.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. volume 97 of Proceedings of Machine Learning Research, pp- 2232 2241, 2019.
Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry. At stability's 's edge: How adjust hyperparameters to preserve minima selection in asynchronous training of neural networks? In International Conference on Learning Representations, 2020. URL httht:::: httphhtps///peeeeeiiiiiee oopenevii.......
ket/Eorummie www.iidddkkeiiit////
Gabriel Goh. Why momentum really works. Distill, 2(4):e6, 2017.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http:
/ 7/www. //ww.deepleaznibboooooo
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. Advances in Neural Information Processing Systems, pp. 1731-1741, 2017.
Wenging Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of noncon- vex stochastic gradient descent. arXiv preprint 1rXiv:1705.0733 2017.
Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross- entropy in classification tasks, 2020.
10

---

Published as a conference paper at ICLR 2021
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings the 32nd International Conference on Machine Learning. pp. 448 456, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In Advances in neural information processing systems, pp. 8571- 1- 8580, 2018.
Arthur Jacot, Franck Gabriel, and Clement Hongler. The asymptotic spectrum of the hessian of dnn throughout training. In International Conference Learning Representations, 2020. URL
Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711004623, 2017.
Stanislaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL W
nett/Forum?1dd-8kqaa..
Stanistaw Jastrzębski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho and Krzysztof Geras The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https:// openreview. http://oooommii1110003300
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima, 2016.
LeCun, Bottou, GB Orr, and K-R Muller. Efficient backprop. Lecture notes in computer science, pp. 9-50, 1998.
Yann LeCun, Patrice Y. Simard, and Barak Pearlmutter. Automatic learning rate maximization by on-line estimation of the hessian's 's eigenvectors. In Advances in Neural Information Processing Systems 5, 1993.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in neural information processing systems, pp. 8572-8583, 2019.
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.00 02218, 2020.
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, pp. 983- 992, 2019.
Xinyan Li, Qilong Gu, Yingxue Zhou, Tiancong Chen, and Arindam Banerjec. Hessian based analysis of sgd for deep nets: Dynamics and generalization. In Proceedings of the 2020 SIAM International Conference on Data Mining, pp. 190-198. SIAM, 2020a.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradi- ent descent on structured data. Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp.
2018.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In Interna- tional Conference on Learning Representations, 2019.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconeiling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33, 2020b.
11

---

Published as a conference paper at ICLR 2021
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Toward theory of optimization for over- parameterized systems of non-linear equations: the lessons of deep learning, 2020.
James Martens. SECOND-ORDER OPTIMIZATION FOR NEURAL.NETWORSS, PhD thesis, 2016.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture
models. In International Conference Learning Representations, 2016.
Rotem Mulayoff and Tomer Michaeli. Unique properties of flat minima in deep networks, 2020.
Kamil Nar and Shankar Sastry. Step size matters in deep learning. In Advances in Neural Informa- fion Processing Systems, pp. 3436-3444, 2018.
Yurii Nesterov. Introductory lectures on convex programming volume i: Basic course. 1998.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate (1/K* 2). In Dokl. akad. nauk Sssr, volume 269, pp. 543-547, 1983.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint 17XXiv:8811.00662 2018.
Vardan Papyan, Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. arXiv preprint ar01v: 00224,, 2019.
Vardan Papyan. Traces class/cross-class structure pervade deep learning spectra, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang, Zach De Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019.
B.T. Polyak.
Some methods of speeding up the convergence iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 1964.
William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery. Numerical Recipes in C (2nd Ed.): The Art of Scientific Computing. Cambridge University Press, USA, 1992. ISBN 0521431085.
Sashank Reddi, Ahmed Hefny,, Suvrit Sra, Barnabas Poczos, and Alex Smola. Stochastic variance reduction for nonconvex optimization. In International conference on machine learning, pp. 314 323, 2016.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. httpatttsss///peee///// wwwp://eeee/frurrrreeee
ld LKFG31111300
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of
the hessian of over- over-parametrized neural networks. arXiv preprint arXiv:1706.0 140.10066.4444, 04454, 2017.
Karthik Sankararaman, Soham De, Zheng Xu, Ronny Huang, and Tom Goldstein. The impact of neural network overparameterizztion on gradient confusion and stochastic gradient descent. In International Conference on Machine Learning, 2019.
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Mądry. How does batch nor- malization help optimization? In Advances in Neural Information Processing Systems, pp. 2483- 2493, 2018
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Con- ference Machine Learning, pp. 343-351, 2013.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias gradient descent on separable data. The Journal of Machine Learning Research, 19 (1):2822-287) 2018.
12

---

Published as a conference paper at ICLR 2021
Iya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In Proceedings of the 30th International Conference on Machine Learning. 2013.
Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste- Julien. Painless stochastic gradient: Interpolation, line-search, and convergence rates. In Ad- vances Neural Information Processing Systems, pp. 3745,
Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: sharp convergence over nonconvex landscapes. In International Conference on Machine Learning, pp. 6677-6686, 2019.
Lei Wu, Chao Ma, and E Weinan. How sgd selects the global minima in over-parameterizze learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, pp- 8279 8288, 2018
Yuege Xie, Xiaoxia Wu, and Rachel Ward. Linear convergence of adaptive stochastic gradient descent. In International Conference on Artificial Intelligence and Statistics, pp. 1475-1485, 2020.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference Learning Representations, 2021. URL httpps//openreview id
iCiYCo
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXivv1802.08770, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020. URL
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive meth- ods for nonconvex optimization. In Advances in neural information processing systems, pp. 9793- 9803, 2018.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Rep- resentations, 2020. URL htttps://operrrrr
ew net. http///rrrmmmo///o wtt.rmmmmm..TT.Ax/X/pps
Dongruo Zhou, Yiqi Tang. Ziyan Yang, Yuan Cao, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. arXiv preprint arXiv:1808.05671, 2018.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochas- tic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In Proceedings of the 36th International Conference on Machine Learning. PMLR, 2019.
13

---

Published as a conference paper at ICLR 2021
A CAVEATS
In this appendix, we list several caveats to our generic characterization of the dynamics of full-batch gradient descent on neural network training objectives.
1. With cross-entropy loss, the sharpness often drops at the end of training As mentioned in the main text, when neural networks are trained on classification tasks using the cross-entropy loss, the sharpness frequently drops near the end of training, once the classification accuracy begins to approach 1. We explain this effect in Appendix C
2. For shallow wide networks,
simple datasets, sharpness doesn't rise that much When the network is shallow (Figure 16 17) or wide (Figure|12] 14p, or when the dataset is "easy" small (Figure T8), the sharpness may rise only a small amount over the gradient flow trajectory. For these optimization problems, "stable step sizes" (those for which gradient descent never enters the Edge of Stability) may be quite reasonable, and the range of "Edge of Stability" step sizes may be quite small.
Sharpness sometimes drops at the beginning of training We sometimes observe that the sharpness drops at the very beginning of training, as the network leaves its initialization. This was more common when training with MSE loss than with cross-entropy loss. For most networks, this drop was very slight. However, the combination of both batch normalization and MSE loss some- times caused situations where the sharpness was considerably large at initialization, and dropped precipitously soon as training began. Figure 8 illustrates one such network.
train loss
shargness
train acouracy
1.0
500
0.8
0.6
6.4
700
0.2
Lime
Uime
Une
Figure 8: With batch norm and MSE loss, the sharpness sometimes drops precipitously: We use gradient flow to train CNN with ReLU activations and batch norm using the MSE loss. Observe that the sharpness drops substantially at the beginning of training, and never recovers to its initial value. This seems to occur due to the combination of both batch normalization and MSE loss.
4. With batch normalization, need to look at sharpness between iterates As detailed in Ap- pendix K when we trained batch-normalized networks using very small step sizes 7, we sometimes observed that the sharpness the gradient descent iterates themselves plateaued below 2/1,, even though the sharpness in between the iterates plateaued just above 2/1, as expected.
5. With non-differentiable components, instability sometimes begins when the sharpness is bit less than 2/1 When training networks with ReLU hardtanh activations, some- times observed that the sharpness started to plateau (and the training loss started to behave non- monotonically) a bit before the instant when the sharpness crossed 2/T7. For example, see Figures or Figures 45 47 One potential explanation is that for such networks, the training objective the second-order Taylor approximation around given iterate may be a poor local model for the training objective at even an tiny distance away in weight space, that weight change causes some activations to switch "ReLU case" from < 0 >> 0, visa versa.
not continuously differentiable,
14

---

Published as a conference paper at ICLR 2021
B STABILITY 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
f(x) 2 Ax+b'x+
(2)
starting from the initialization X6-
To review, vanilla gradient descent is defined by the iteration:
Xt+1
V/ (/V[(xx) [xx)).
Meanwhile, gradient descent with Polyak (also called "heavy ball") momentum (Polyak Sutskever et al. 2013 Goodfellow et al. 2016) is defined by the iteration:
1964
V++1 5v, fxx)
X++1 = X₃ + V++1
where v, is a "velocity" vector and S 8 is the momentum coefficient. For B the algorithm reduces to vanilla GD.
Finally, Nesterov momentum Sutskever et al. (2013); Goodfellow al. (2016) is an adaptation of Nesterov's 's accelerated gradient (Nesterov) (983) for deep learning defined by the iteration:
Ve+1 σv, 7/ fxx + Bv:)
Xx41 X, + V
where is a "velocity" vector and 0 $ 8 is the momentum coefficient. For 8 = 0 the algorithm reduces to vanilla GD.
All three of these algorithms share special property: on quadratic functions, they act independently along each Hessian eigenvector. That is, we express the iterates the Hessian eigenvector basis, then in this basis the coordinates evolve independent from another under gradient descent.
Proposition 1. Consider running vanilla gradient descent the quadratic objective (2) starting from No. Lei (q, a) an eigenvectorfeigeevaaue pair of A. If a > 2/1,, then the sequence (q x;) will diverge.
Proof. The update rule for gradient descent on this quadratic function is:
X = X; - + b) (I- r/b.
Therefore, the quantity q" x, evolves under gradient descent as:
q" x1+1 = q" 7, Xx 1/b = (1 - na) xt 17 Tb. b.
(q Aq"A= aq)
Define ã, = -"= xt + Lq=b, and note that {q'', diverges if and only if {ū,} diverges. The quantity 7; evolves under gradient descent according to the simple rule:
$ (1 na)ãE
Since n if 2/n then (1 17a) < 1, the sequence (ø,) will diverge.
We now prove analogous results for Nesterov and Polyak momentum.
15

---

Published as a conference paper at ICLR 2021
Theorem 1. Consider running Nesterov momentum on the quadratic objective (2 any initialization, Let (q, a) be an eixenvectorleeixemaane pair of A. Ifa 1+23 sequence {q"xx) will diverge.
starting from
then the
Proof. The update rules for Nesterov momentum on this quadratic function are:
Vi+1 3(I nA)v, r/b nAx,
X++1 = X2 + Va+1-
Using the fact that V; = X2 Xt-1, we can rewrite this as a recursion in x; alone:
X1+1 x, + B( Ax) x)-1) w/b nAA,
(1 (1 BT -))x, 3(T - -1)x1111111 17b
Define 84 x²x, 2 b, and note that q² x, diverges iff 5. diverges. It can be that 3. evolves as:
ž = (1+ 8) na) - 3(1 - -1ja)++-1
This is linear homogenous second-order difference equation. By Theorem 2.37 Elaydi (2005). then this recurrence diverges.
since 0 and 3 < 1, if
2+22
1+23
The following result previously appeared Goh (0177)
Theorem 2. Consider running Polyak momentum the quadratic objective
2 starting from any + 23), then the
initialization. Let (q. a) be an eigenvectorleigemvalee pair of A. If a
sequence {q"x,) x1} will diverge.
Proof. Using the fact that recursion alone:
Xt-1, can re-write the Polyak momentum recursion
Xx+1 Xx - fxx))
-v:
= = (B(x: xx-1) f((())))
(1 B)x, 3x,-1 (xx))
For the quadratic objective 2}. this update rule amounts to:
X = (1+ 8)x - Bx (Ax; + b)
(1+ nA)x, 3x,-1 1/b.
Multiplying by qT, we obtain:
qf xi+1 = (1 3 -na)q"xx -39 x1-1 177--7/h b.
Now, define 3, q"x x, + d. Note that q" diverges iff it diverges. can be that क. evolves as:
T = (1+ 3 - - 3
This is linear homogenous second-order difference equation. By Theorem 2.37 in Elaydi 2005 since > 0 and 3 < 1, if > 1 (2 2,3) then this recurrence diverges.
16

---

Published as a conference paper at ICLR 2021
C CROSS-ENTROPY LOSS
In this appendix, we explain why the sharpness decreases at the end of training when the cross- entropy loss is used. Before considering the full multiclass case, let us first consider the simpler case of binary classification with the logistic loss.
Binary classification with logistic loss We consider a dataset (x) C Rd x {-1, 1},
where the examples are vectors in I and the labels are binary {-1,1}. We consider neural network / : IR x IP + I which maps an input x E IR and a parameter vector E Rº to prediction h(x;; 0) R. Let E R x { 1-1, 1] +
be the logistic loss function:
f(=: v) = log(1 + exp(-zy))
1
- log where p = + expl (
The second derivative of this loss function w.r.t is:
(z = p(1 p).
The full training objective is:
f()) =
(0) where f:(0) = [(f(x;));
The Hessian of this training objective the average of per-example Hessians:
(² v.)))).
For any arbitrary loss function l, we have the so-called Gauss-Newton decomposition (Martens) 2016 Bottou et al. 2018) of the per-example Hessian 4 (00);
( Vah(x;; 0) Voh(x;; 0)² + & v.) V)h(x;;0 where 2, hx;;; 0)
where Voh(x;; 0) E IR' is the gradient of the network output with respect to the weights, and refers to the derivative of f with respect to its first argument, the score.
Empirically, the first term in the decomposition (usually called the "Gauss-Newton matrix") tends to dominate the second, which implies the following "Gauss-Newton approximation" to the Hessian:
A2 f(0
(2) Voh(x; 0) Voh(x;;))
(3)
In our experience, progressive sharpening affects Voh(x;; 0) (x;; O);; That is, Vah(x;;0) 0) Voh(x;; 0) tends grow in scale continually during training. For the square loss ((2;)) = {2 - y)², the second derivative ("(z;)) = is the identity, so the f(() grows con- tinually well. training. Why is this? In Figure 9 we plot both the logistic loss l, and its second derivative l", as function of the quantity y2, which is often called the "margin."
contrast, for the logistic loss, many of the ("(2;; yi) decrease the very end of
Crucially, observe that both é and ! are decreasing in yz. Because the loss is decreasing in ye, once an example is classified correctly (i.e. y424 > 0), the training objective can be optimized fur- ther by increasing the margin y;z4. Because e" is also decreasing in yz, if the margin y/z; increases, the term E" (z(;;; yi) drops. Near the end of training, once most examples are classified correctly, gradient descent layer weight matrix. This causes the E" ("(24;)) y:) to drop. Therefore, even though progressive sharpen- ing still applies (x;(x;; 0) Voh (x; 0)² the decrease in the {" ys''s 's pulls down the leading eigenvalue of the newss-Newton matrix in equation 3
casily increase the margins of all these examples by simply scaling up the final
This effect is illustrated in Figure 10 Here, train network the binary classification task of CIFAR- 10 airplane vs. automobile, using the logistic loss. In Figure (O(e), we plot the margin y; for 10 examples in the dataset. Notice that the end of training, these margins all continually rise;
17

---

Published as a conference paper at ICLR 2021
logistic loss
logistic loss second derivative f"
24-
033
1.3
022
031
03
-2
-2
-1
y2
y2
Figure 9: Left: logistic loss function as a function of yz. Right: its second derivative l" as function of yz.
this is because gradient descent "games" the objective by increasing the margins of successfully classified examples. When the margin y;2; rises, the second derivative (z, yi) drops. This can be seen from Figure f) where plot E" (2; ys) for these same 10 examples. Now, all the while, the leading eigenvalue of the matrix I 1== Voh(x, () V) 9)* keeps rising, as can be seen from Figure 10 d). However, because the L are dropping, the leading eigenvalue of the Gauss- matrix 2i=1 (" V:) (x; 0) V)hh((x h starts to decrease at the end of training, as can be seen from the green line in Figure T Finally, since the leading eigenvalue of the Gauss-Newton matrix is an excellent approximation to the leading eigenvalue of the Hessian (i.e. the sharpness), the sharpness also drops the end of training, as can be seen from the orange line in Figure 10 c).
(c) leading eigenvalue of Hessian and Gauss-Newton matrix
(a) train loss
(b) train accuracy
1.0
0.6
0.8
Gouss
2000
4000
8000
6000
5000
step
step
step
(6) leading eigenvalue of 2 VyPDK; di Vehtu; of"
(c) margin examples
d" 10 examples
sigernluu
margin
5000
2.5
0.1
-2.5
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
step
stap
slep
Figure 10: We train a network using the logistic loss the binary classification problem of CIFAR- 10 airplane vs. automobile.
(a) The train loss. (b) The train accuracy. (c) The leading eigenvalue of the Hessian and of the Gauss-Newton matrix; observe that the latter is great approximation to the former. Gauss-Newton matrix without the t" terms; observe that this value constantly rises dip at the end of training.
(d) The leading eigenvalue of the matrix 1 > Vgh(x,; 8) Vohxx; 8)
which the
does not (e) The margin y; of 10 examples; observe that all the margins rise at (f) the value é"(2 y1) for 10 examples; observe that all of these curves decline
the end of training. at the end of training.
Multiclass classification with cross-entropy loss We consider dataset {(x,, W:) Rd x {l....,k) kj, where the examples are vectors in Rd and the labels are in kl. We consider neural network h Rdd x R + R which maps an input x Rd and parameter vector E R" to prediction h(z; 0) R Let I f1....,k) + R be the cross-entropy loss function:
exp(zy)
((2;v) log Σ/= exp(z;)
cxp(z)
= = where p Σ, exp(2;)
18

---

Published as a conference paper at ICLR 2021
The Hessian {)((zy)) E R of this loss function w.r.t the class scores z is:
v2t(z; (pa) pp".
Now, for any loss function f: R x (1....,k) R, we have the Gauss-Newton decoopostion::::::::
(i() = J V2 (zzw)) J.
V ((xsw))) (xh)(xx;;)
where z; = (x;; f) E R are the logits for example i, J. E RExP is the network output-to- weights Jacobian for example i, V2, PZz y:) E B is the Hessian of w.r. its input z₁ and (x;; 0) TRxx is the Hessian matrix of the j-th output of the network h on the i-th example.
Dropping the second term yields the Newss- approximation:
fi(0 R J V ((z v6) J.
As in the binary classification case discussed above: at the end of training, for many examples i, the Yt entry of P; will tend toward I and the other entries of p. will tend O. Once this occurs, the matrix diag(P:) P.P." will broadly decrease in scale: the diagonal entries of this matrix are of the and the off-diagonal entries are of the form -pG, which also goes to zero if p - 0, q - or - 0, q + 0 or + 1,4 - 0.
form p(1 - p). which goes to zero as p - 0 or p
This effect is illustrated in Figure Here, we train a network on CIFAR- using the cross-ent cross-entropy loss. In Figure e)..). for ten examples in the dataset (with output scores plot the margin 2, [y] maxj [.]... which is the difference between the score of the correct class y/4 and the score of the next-highest class. Observe that for all of these examples, this margin rises at the end training.
l((x)) E Rt), we
Figure I]), for those same ten examples, we plot the quantity P: P1] Observe that for all of these examples, this quantity continually decreases at the end training. Now, all the while, the leading eigenvalue of the matrix n E. J. keeps rising, as can be seen from Figure d)... d). However, because V²((z, y:) is decreasing, the leading eigenvalue of the Gauss-Newton matrix - Eー J V2 z.; yı) J; starts decrease at the end of training. as can be seen from the green line in Figure [Tc).. Finally, since the leading eigenvalue of the Gauss-Newton matrix is an excellent approximation to the leading eigenvalue of the Hessian (i.c. the sharpness), the sharpness also drops at the end of training, as can be seen from the orange line in Figure cTc).
19

---

Published as a conference paper at ICLR 2021
(c) leading algenvalue Hessian
(a) train lass
(b) train
Gauss Newton
0.8
alaunvalls
0.6
0.4
Hesslan Gaiss-Neetton
2000
4.00
6000
8000
10000
2000
4000
6000
8000
10000
2000
4000
5000 stap
8000
10000
step
step
(c) leading eigenvalue of 2
(e) margin 10 examples
1" examples
6000
2000
0.0
2000
4000
6000
8000
10000
2000
4.00 5000
8000
10000
2000
4000
6000
8200
19200
Figure 11: We train a network using the cross-entropy loss on CIFAR-10. (a) The train loss. (b) The train accuracy. (c) The leading eigenvalue of the Hessian and of the Gauss-Newton matrix; observe that the latter
great approximation to the former. (d) The leading eigenvalue of the matrix >E Jº which is the Gauss Newton matrix except the A terms; observe that this value constantly rises it does not dip at the end training. (e) The margin 23(yi] - max z of 10 examples; observe that all these margins rise the end of training. (f) the value p.[ja](1 P: (:(y:) for 10 examples; observe that all of these curves decline at the end of training.
Gauss-Newton
20

---

Published as a conference paper at ICLR 2021
D EMPIRICAL STUDY OF PROGRESSIVE SHARPENING
In this appendix, we empirically study how problem parameters such as network width, network depth, and dataset size affect the degree to which progressive sharpening occurs. To study progres- sive sharpening on its own, without the confounding factor of instability, we train neural networks using gradient flow (Runge-Kutte rather than gradient descent. Informally speaking, gradient flow does "what gradient descent would do if gradient descent didn't 't have to worry about instability."
We observe that progressive sharpening occurs greater degree: (1) for narrower networks than for wider networks (which is consistent with infinite-width NTK theory), (2) for deeper networks than for shallower networks, and (3) for larger datasets than for smaller datasets.
The effect of width When networks parameterized in a certain way (the "NTK parameterization") are made infinitely wide and trained using gradient flow, the Hessian moves vanishingly small amount during training, implying that progressive sharpening occurs (Jacot al.) 2018 Lee et al., 2019) Jacot et al. [2200] Li & Liang, 2018). Therefore, a natural hypothesis is that progressive sharpening might attenuate network width increases. We now this hypothesis.
experiment which supports
We consider fully-connected architectures with hidden layers and tanh activations, with widths (32, 64, 128, 256, 512, 1024]. We train on a size-5,000 subset of CIFAR-10 using the cross-entropy loss. We train using gradient flow (details in $I.5). We consider both NTK parameterization and standard parameterization (Lee et al. 2019).
In Figure 12 for each width, we train NTK-parameterized networks from five different random initializations, and plot the evolution of the sharpness during gradient flow. Observe that the max- imum sharpness along the gradient flow trajectory is larger narrow networks, and smaller for wide networks. (As elsewhere in this paper, note that the sharpness drops at the end of training due to the cross-entropy loss.) In Figure [3 we plot summary statisties from these training runs. Namely, define Amax the maximum sharpness over the gradient flow trajectory, and define λo as the initial sharpness. In Figure [3)a), for each width we plot the mean and standard deviation of the maximum sharpness Amax over the five different random initializations. Observe that Amax becomes smaller, on average, as the width made larger. In Figure 13),), for each width we plot the mean and standard deviation of the maximum sharpness gain Xmax Aoo over the five different random initializations. Observe that the maximum sharpness gain Amax/Ao also becomes smaller the width is made larger. NTK theory suggests that Amax/ Ån should deterministically tend as the width + 0o, and Figure 13 b) is consistent with this prediction.
In Figures and 15 we conduct similar experiments, but with standard parameterization rather than NTK parametcrization. Similar to NTK parameterization, we observe in Figure [4 that the sharpness rises more for narrow networks than for wide networks.
21

---

Published as a conference paper at ICLR 2021
width
2.0
2.0
2.0
1.5
1.5
3.5
1.0
1.0
3.0
0.5
0.5
56
105
15K
21K
255
30K
33K
5K
10K
13K
206
23K
33K
56
10K
15K
23K
25.
36K
5776
time
time
wkdth 256
width 512
Width 1024
1.5
Figure 12: NTK parameterization: evolution the sharpness. We use gradient flow to train parameereeeer networks, and we track the evolution of the sharpness during training. For each width, we train from five different random initializations (different colors). Observe that the sharpness rises more when training narrow networks than when training wide networks.
maximum sharpness along trajectory
maximum sharpness gain along trajectory
Amau/Ao (mean + sd)
1.0
40
Amar
(mean
0.5
20
32
64
128
256
512
1024
32
64
128
256 width
512
1024
width
Figure 13: NTK parameterization: summary statistics. Left: For each network width, we plot the mean and standard deviation (over the five different random initializations) of the maximum sharpness Amax along the gradient flow trajectory. Observe that Amax decreases expectation as the width increases. Right: For each network width, plot the the five different random initializations) of the maximum sharpness gain Amax/ λo along the gradient flow trajectory. Observe that Amax/ Ay decreases in expectation as the width increases. Indeed, NTK theory predicts that Amax/ do should deterministically tend to as width - oo and this plot is consistent with that prediction.
and standard deviation (over
witthhhu 32
width 64
width 128
20
110
430
400
300
330
350
200
200
220
100
100
100
1.0
25
30
25
40
Elme
Oine
Onc
Figure 14: Standard parameterization: evolution of the sharpness We gradient flow train standard-pane hunndardparrmettriieee networks, and we track the evolution of the sharpness during training. For each width, train from five different random initializations (different colors). Observe that the sharpness rises more when training narrow networks than when training wide networks.
22

---

Published as a conference paper at ICLR 2021
maximum sharpness along trajectory
max sharpness gain (mean sd)
maximum sharpness gain along trajectory
max max sharpness (mean
I
200
32
64
128
256
512
1024
32
64
128
256
512
1024
width
width
Figure 15: Standard parameterization: summary statistics. Left: For each network width, we plot the mean and standard deviation (over the five different random initializations) of the maximum sharpness Amax along the gradient flow trajectory. Observe that Amax tends to decrease in expec- tation as the width increases, though it is not clear whether this pattern still holds when moving from width 512 to width 1024 more samples are needed. Right: For each network width, plot the mean and standard deviation (over the five different random initializations) the maximum sharpness gain Amax//X Ao along the gradient flow trajectory. Observe that Amax/ Àu decreases in pectation as the width increases. It is not clear whether or not Amax Ag is deterministically tending to 1, but that does seem possible.
23

---

Published as a conference paper at ICLR 2021
Effect of depth We now explore the effect of network depth progressive sharpening. We gradient flow to train fully-connected tanh architectures of width 200 and varying depths ranging from hidden layer to hidden layers. We train a 5k subset of CIFAR-10 using both cross- entropy loss (in Figure 16) and square loss (in Figure 17). For each depth, we train from five different random initializations (different colors). Observe that progressive sharpening occurs greater degree as network depth increases.
depth
depth 2
depth 3
depth
sharpees
750
750
750
750
500
500
500
500
250
time
time
time
time
Figure 16: The effect of depth: cross-entropy. We use gradient flow train networks of var- ious depths, ranging from 1 hidden layer to 4 hidden layers, using cross-entropy loss. We train each network from five different random initializations (different colors). Observe that progressive sharpening occurs to greater degree for deeper networks.
depth
depth
depth
depth
1000
1000
1000
1000
stargners
750
500
250
250
250
250
0
25
50
75
25
50
75
50
50
100
Figure 17: The effect of depth: mean squared error. We use gradient flow to train networks of var- ious depths, ranging from hidden layer 4 hidden layers, using MSE loss. We train each network from five different random initializations (different colors). Observe that progressive sharpening occurs to greater degree for deeper networks.
Effeet of dataset size We now explore the effect of dataset size on progressive sharpening. We use gradient flow to train a network on different-sizee subsets of CIFAR-10. The network is a 2- hidden-layer, width-200 fully-connected tanh architecture, and we train using cross-entropy loss. The results are shown in Figure 18 Observe that progressive sharpening occurs to greater degree as dataset size increases.
dataset size 1k
dataset size 2k
dataset size 5k
dataset size 10k
dataset size 20k
500
500
500
500
500
400
400
Parpnes
200
200
200
200
100
100
100
100
100
0
100
time
time
time
time
time
Figure 18: subsets of CIFAR-10. Observe that progressive sharpening occurs to greater degree the dataset size increases.
The effect of dataset size. We use gradient flow to train network on varying-sized
24

---

Published as a conference paper at ICLR 2021
E SPEED OF DIVERGENCE ON QUADRATIC TAYLOR APPROXIMATION
the sharpness at some iterate is strictly greater than 2/1,, then gradient descent with step size is guaranteed diverge if on the quadratic Taylor approximation around that iterate. However, the speed of this divergence could conceivably be slow - in particular, the train loss might continue to decrease for many iterations before starts to increase. this appendix we empirically demonstrate, to the contrary. that at the Edge of Stability, gradient descent diverges quickly if, at some iterate, we start running gradient descent on the quadratic Taylor approximation around that iterate.
We consider the fully-connected tanh network from section 3 trained on 5,000-sized subsample of CIFAR-10 using both cross-entropy loss and MSE loss. some timestep to during training, we suddenly switch from running gradient descent on the real neural training objective, to running gradient descent on the quadratic Taylor approximation around the iterate at step l.- We do this for three timesteps before gradient descent has entered the Edge of Stability, and three afterwards. Figure 20 shows the results for cross-entropy loss, and Figure 21 shows the (similar) results for MSE loss. Before entering the Edge of Stability (top row), gradient descent the quadratic Taylor approximation behaves similar to gradient descent on the real neural training objective - that is, the orange line almost overlaps the blue line. Yet after entering the Edge of Stability (bottom row), gradient descent on the quadratic Taylor approximation quickly diverges, whereas gradient descent on the real neural training objective makes consistent (if choppy) progress.
In short, when gradient descent is not at the Edge of Stability, the quadratic Taylor approximation serves good model for the local progress of gradient descent. But when gradient descent is the Edge of Stability, the quadratic Taylor approximation is extremely poor model for the local progress of gradient descent. is conceivable that there exists some simple modification to the quadratic Taylor model which would fix this issue (e.g. perhaps if one ignores a certain direction, the quadratic Taylor model accurate). Nevertheless, unless/until such fix is discovered, it is unclear why quadratic Taylor approximations should yield any insight into the local behavior of gradient descent.
cross-entropy loss 2/60)
2/200)
0.6
1.5
train 1000
1.0
0.5
0.2
500
1000
1500
2000
4000
6000
B000
10000
Iteration
Feration
Figure 19: We train a neural network using cross-entropy loss (left) and MSE loss (right). In Figure 20 and 21 show what happens when, the iterations marked above by vertical dotted lines, we switch from gradient descent on the real neural training objective to gradient descent on the quadratic Taylor approximation.
25

---

Published as a conference paper at ICLR 2021
600
1.700
1.475
1.25
tain 1055
1.675
tai icos
1.450
tral: loxe
real objective quadratic
1.425
2.20
approximation
1.400
1,600
160
180
220
240
360
420
440
560
600
620
640
step 800
step 1000
step 1200
1.1
1.05
1.0
0.7
1.00
train
train
0.6
0.95
0.7
0.90
stop
stop
step
Figure 20: Cross-entropy loss (17 2/60). At six different iterations during the training of the network from Figure (marked by the vertical dotted black lines), we switch from running gradient descent the real neural training objective (for which the train loss is plotted in blue) to running gradient descent on the quadratic Taylor approximation around the current iterate (for which the train loss is plotted in orange). Top row timesteps (200, 400, 600) before gradient descent has entered the Edge of Stability; observe that the orange line (Taylor approximation) closely tracks the blue line (real objective). Bottom row are timesteps {800, 1000. 1200) during the Edge of Stability; observe that the orange line quickly diverges, whereas the blue line does not.
2000
step 3000
stop 4000
0.270
0.227
0.269
tain 1055
trin 1005
Trann lo00
0.226
real objective quadratic approximatin
0.268
0.22
0.309
0.267
0.24
0.308
0.266
1960
1980
2000
2020 2040
2960
2960
3000
3020
3040
3960
3980
4000
4020 4040
step
5000
step
6000
step
7000
0.1.9 98
0.185
0.155
train loss
0.198
train loss
0.175
train
0.150
0.192
0.170
0.145
step
step
step
Figure 21: MSE loss (1) = 2/200). At six different iterations during the training of the network from Figure 19 (marked by the vertical dotted black lines), we switch from running gradient descent on the real neural training objective (for which the train loss is plotted in blue) to running gradient descent on the quadratic Taylor approximation around the current iterate (for which the train loss is timesteps (2000, 3000, 4000) before gradient descent has entered the Edge of Stability; observe that the orange line (Taylor approximation) closely tracks the blue line (real objective). Bottom row are timesteps (5000, 6000. 7000) during the Edge of Stability; observe that the orange line quickly diverges, whereas the blue line does not.
plotted in orange). Top
26

---

Published as a conference paper at ICLR 2021
F "OPTIMAL" STEP SIZE SELECTION
One heuristic for setting the step size of gradient descent is to set the step size at iteration to Y) 1/Ar, where A is the sharpness iteration While this heuristic is computationally impractical due to the time required to compute the sharpness at each iteration, it is often regarded as an ideal, for instance in LeCun et al. 1998) (Eq. 39). LeCun al. (1993), and Schaul al. (2013) (Eq 8). The motivation for this heuristic is: if all that is known about the training objective is that the local sharpness is X, then step size of 1/A maximizes the guaranteed decrease in the training objective that would result from taking step.
First, we demonstrate (on single numerical example) that the dynamic step size 1)t = 1/4, is outperformed by the baseline approach of gradient descent with fixed step size = 1/ where An is the sharpness at initialization. In Figure 22 we train the network from $ using both the dynamic 7 1/Ax step size heuristic as well as the baseline fixed step size of 1/ Observe that the 17 = 1/ Ao baseline outperforms the 1/4, heuristic. Intuitively, because of progressive sharpening, the n 1/A, heuristic anneals the step size, and therefore ends up taking steps that are suboptimally small. In contrast, while the т) = 1/20 baseline quickly becomes unstable, this instability is apparently worthwhile "price to pay" in return for the benefit of taking larger steps.
cross-entropy loss sharppess
loss
size
150
Boed g-1.0kk
dynamic 1.50
200
0.075
0.052
0.5
0.025
1,000
2000
3000
6010
1005
3003
4000
0.002
1000
2000
3000
steg
50ng
MSE loss
train loss
sha sprness
stop size
11a7 1005
0.0
0.4-
synamic 1AQ
faol 1 Big
shursnas
635-
0.020
0.055
0.2
0.000
3200
10030
12000
2000
10000
13200
0
5002
10200
11000
<tep
step
Figure 22: A dynamic step size m)- 1/A, underperforms a fixed step size of ne = 1/20-
Another natural idea is to dynamically set the step size iteration & 1)t 1.9/ This step size rule takes larger steps than the 1) = 1/Ae rule while still remaining stable. In Figure 23 we compare this 7)- = 1.9/X, rule to baseline approach of gradient descent with a fixed step size π) 1.9/ Ao where Ao is the sharpness at initialization. Observe that the baseline of fixed 1.9 Ào step size outperforms the dynamic 1): = 1.9/X, rule.
cross-entropy loss sharpness
step
2.2
floel q- 1.50%
0.25
oynamic 1.311
0.3
0.02
500
1002
1500
MSE loss sharppees
step
0.6
0.4
1.9L
5.02
200
0.51
0.2
10:00
5.00
Figure 23: A dynamic step size 1/2 1.9/2, underperforms fixed step size of n 19/22.
27

---

Published as a conference paper at ICLR 2021
G EVOLUTION OF SHARPNESS DURING SGD
In this appendix, we briefty illustrate how the sharpness evolves during stochastic gradient descent. In Figure 24 we train the tanh network from $ using SGD with both cross-entropy loss (top) and mean squared error (bottom). We train using range of batch sizes (different colors). We observe the following:
1. During large-batch SGD, the sharpness behaves similar to full-batch gradient descent: it rises to 2/17 (marked by the black horizontal dashed line) and then hovers just above that value.
2. Consistent with prior reports, we find that the smaller the batch size, the lower the sharpness 2020).
(Keskar et al. 2016 Jastrzębski et al., 2017 2019;
3. Notice that when training with entsspy loss batch size 8 (the blue line), the sharp- ness decreases throughout most of training. The train accuracy (not pictured) is only 66% when the sharpness starts to decrease, which suggests that the cause of this decrease unrelated to the effect described in Appendix C whereby the sharpness decreases at the end of training. Figure 5(a) of Jastrzębski et al. (2020) also depicts a network where the sharpness decreases during SGD training.
cross-entropy
focus all batch sizes
focus on smallest batch sizes
focus an largest batch sizes
BS 8
100
30
sharpness
sharpness
BS
20
65
50
128
65
256
10.
25
10000
10000
Iteration
Iteration
Iteratioo
squared
Tocus on all batch sizes
Tocus on sma smallest batch sizes
focus on largest batch sizes
BS
200
200
US B
sharpresss
sharpness
150
BS
100
40
BS
100
BS
256
50
20
20000
40000
60000
20000
40000
60000
0
2000
4000 6000 8000 10000
iteration
Reration
Iteration
Figure 24: Evolution of sharpness during SGD. We train network using SGD at various batch sizes, and plot the evolution of the sharpness. the top row, we train with cross-entropy loss and step size 0.02; in the bottom row, we train with mean squared error and step size 1 0.01. The leftmost column shows all batch sizes; the center column focuses on just the smallest; and the rightmost column focuses on just the largest. The black dashed horizontal line marks 2/1
28

---

Published as a conference paper at ICLR 2021
H SGD ACCLIMATES TO THE HYPERPARAMETERS
In this appendix, we conduct an experiment which suggests that some version of "Edge of Stability" may hold for SGD.
One way to interpret our main findings is that gradient descent "acclimates" to the step size in such way that each training update sometimes increases, and sometimes decreases, the training loss, yet an update with smaller step size would always decrease the training loss. We now demonstrate that this interpretation may generalize to SGD. In particular, we will demonstrate that SGD seems to "acclimate" to the step size and batch size in such way that an actual update sometimes increases and sometimes decreases the loss in expectation, yet a update with larger step size or smaller batch size would almost always increase the loss in expectation, and a step with a smaller step size or larger batch size would almost always decrease the loss in expectation.
In Figure 25 we train the tanh network from X with MSE loss, using SGD with step size 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 25 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 25 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 other points, and SGD step would increase the loss in expectation.
In Figure 2))a). while training that network, we compute the expected change in training loss that would result from taking an SGD step with the same step size used during training (i.e. 0.01), 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 (6)bb 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 decrease in the training loss in expectation. In Figure 26(c) and (d) repeat this experiment with the step size; we observe that an SGD step with a larger step size (0.02) would consistently increase the training loss in expectation, while an SGD step with a smaller step size (0.005) would consistently decrease the training loss in expectation.
In Figure 27 as a "control" experiment, we both train and measure the expected loss change under the following four hyperparameter settings: (step size 0.01, batch size 16), (step size 0.01, batch size 64), (step size 0.02, batch size 32), and (step size 0.005, batch size 32). In each case, we observe that, after brief period at the beginning of training, each SGD update sometime increases and sometimes decreases the training loss in expectation.
Therefore, least for this single network, we can conclude that no matter the hyperparameters, SGD quickly navigates to, and then lingers in, regions of the loss landscape in which an SGD update with those hyperparameters sometimes increases, and sometimes decreases the training loss in expectation, yet an SGD update with a smaller step size or larger batch size would consistently decrease the loss in expectation, and an SGD update with larger step size or smaller batch size would consistently increase the loss in expectation.
29

---

Published as a conference paper at ICLR 2021
training loss
expected loss change
expected loss change
0.05
training loss
0.00
0.25
-0.05
5000
10000
15000
5000
10000
15000
SGD step
SGD step
SGD does not consistently decrease the training loss in expectation. We train the tanh network from section with MSE loss, using SGD with step size 0.01 and batch size 32. Periodically during training, we compute (left) the full-batch training loss, and (right) the expected change in the full-batch training loss that would result from taking an SGD step (where the expectation is over the randomness in sampling the minibatch). Strikingly, note that after the very beginning of training, the expected loss change is sometimes negative (as desired) but oftentimes positive. See Figure 26
Figure 25:
train n- 0.01, BS-32
train 0.001, BS-32
train 0.01, ES-
train at 0.01, B5-32
evaal twice that batch size 164)
eval at half that bitch size (126(16)
eval half that step size (0.005)
eval at twice that size (0.00
8.625
0.025
0.825
0. 0.025
oxpected loua
3.000
2.002
0.002
0.002
-3.059
-3.053
5020 10000
15000
5000 20200
15020
5020 10003
15020
5000
30000 15000
500 200
560
SCD stop
560 stop
Figure 26: An SGD step with smaller learning rate or larger batch size than the ones used during training would consistently decrease the loss in expectation. At regular intervals during the training run depicted in Figure 25 (with 0.01 and batch size 32), measure the expected change in the full-batch training loss that would result from an SGD step with a different step size larger batch size would
or batch size. Observe that taking SGD step with smaller step size consistently have decreased the loss in expectation, while taking an SGD step with a larger step size smaller batch size would have consistently increased the loss in expectation.
n B5w64
eval 1..0000.1. B5w16
AND at. 0.000, B5-32 n B5
experddloss dhange
-0.02
-0.04
-0.04
15030
Figure 27: Control experiment. Above, Figure a6) as we trained a network with step size 0.01 and batch size 32, we evaluated the expected change in training loss that would result from taking an SGD step with step size 0.01 and batch size 64. Here, in Figure we train the network with step size 0.01 and batch size 64, and evaluate the expected change in training loss that would result from taking an SGD step with the same step size and batch size. We observe that SGD step using the same step size and batch size that are used during training would sometimes increase and sometimes decrease the training loss expectation. The other three panes analogous.
as a "control experiment,"
30

---

Published as a conference paper at ICLR 2021
I EXPERIMENTAL DETAILS
1.1 VARYING ARCHITECTURES ON 5K SUBSET OF CIFAR-10
Dataset. The dataset consists of the first 5,000 examples from CIFAR-IO. To preprocess the dataset, we subtracted the mean from each channel, and then divided each channel by the standard deviation (where both the mean and stddev were computed over the full CIFAR-IO dataset, not the 5k subset).
Architectures. We experimented with two architecture families: fully-connected and convolutional. For each of these two families, we experimented with several different activation functions, and for convolutional networks we experimented with both max pooling and average pooling.
The PyTorch code for c.g. the fully-connected ReLU network is as follows:
nn nn. Sew.ueiunaa
nn.Flatten (),
Ln.ein (3072, 200, bias=True), nn.ReLU()
Ln.Line (200, 200, bias=True) nn.ReLU(). nn. (200, 10, bias=True)
Networks with other activation functions would have nn nn. ReLU () replaced by nn. ELL (), nn. Tanh (), nn. Softplus (), or nn. Hardtanh ().
The PyTorch code for e.g. the convolutional ReLU network with max-pooling is as follows:
nn. Sequentiaalll
nn. nn.Convd (3, 32, bias=True, kernel_size=3, padding=1), nn.ReIU
mn.MaxPool2d12),
(n.Conv2d(32, 32, bias=True, kernel_size=3, padding=1), nn.ReLU mmMaaxPol22d
nn.Flatten()
nn.Linerr (2048, 10, bias=True)
Networks with other activation functions would have nn ReLU () replaced by nn ELU or nn nn. Tanh O) and networks with average pooling would have nn Ma.MaxP MaMMxxool22d (2) replaced by nn AngPol1d (2).
For all of these networks, we use the default PyTorch initialization. That is, both fully-connected layers and convolutional layers have the entries of their weight matrix and bias vector sampled i.i.d from Uniform(
VTan in
ma
Loss functions. For a k-class classification problem, the network outputs are z E R and the correct class is i [R:], then the squared error (MSE) loss defined as (=[1] 1)2 Ej
That is, we encode the correct class with "1" and the other classes
with "0." The cross-entropy loss is defined log Eمر
exp(zlt) exp(=(]))
I.2 STANDARD ARCHITECTURES ON CIFAR-10
To preprocess the CIFAR-10 dataset, subtracted the from each channel, and then divided each channel by the standard deviation.
Since training with full-batch gradient descent is slow, opted to experiment relatively shallow networks. The VGG networks (both with and without BN) are VGG-11's, from the
implementation here: atttas//.ttta...cco...eegggggggo
31

---

Published as a conference paper at ICLR 2021
blob/ blob/master/veg py with the dropout layers removed. The ResNet the (non-fixup) ResNet- 32 implemented here: https:::///o htttss///ittttii
/ixxup
For the two networks with batch normalization, running gradient descent with full-dataset batch normalization would not have been feasible under our GPU memory constraints. Therefore, we instead used ghost batch normalization (Hoffer al.) 2017) with 50 ghost batches of size 1,000. This means that we divided the 50,000 examples in CIFAR-10 into 50 fixed groups of size 1,000 each, and defined the overall objective function be the average of 50 fixed batch-wist objectives. To correctly compute the overall gradient of this training objective, we can just run backprop 50 times (once on each group) and average the resulting gradients.
To compute the sharpness over the full CIFAR-10 10 dataset would have been computationally expen- sive. Therefore, an approximation, we instead computed the sharpness over just the first 5,000 examples in the dataset (or, for the BN networks, over the first 5 batches out of 50).
1.3 BATCH NORMALIZATION EXPERIMENTS
We used the CNN architecture from S (described above in SL1 but with a BatchNorm2d layer inserted after each activation layer.
Since our GPUs did not have enough memory to run batch normalization with the full dataset of size 5,000, we used ghost batch normalization with five ghost batches of size 1,000 (see for details).
1.4 TRANSFORMER ON WIKITEXT-2
We used both the Transformer architecture and the preprocessing setup from the official Py- Torch (Paszke al. 2019) werd- language modeling tutorial: /https://githin com/ pytorch/ wttpecc/exxxmmeeeeseeeeee httpe//msterr/eee
wworeeerwwwrr...r word language _ocgeel___e_______ We used the settings set bytt = 35,
ninp=200, nhead=2, nhid=200, molo,allae nlayers=2, dropout=0. which means that we divided the corpus into chunks of 35 tokens, and trained the network, us- ing the negative log likelihood loss, to predict each token from the preceding tokens in the same chunk. Since computing the sharpness over the full dataset would not have been computationally practical, we computed the sharpness over subset comprising 2500 training examples.
1.5 HINGE-KUTTE
We used the "RK4" fourth-order Runge-Kutte Kutta algorithm (Press et al.) 1992) to numerically integrate the gradient flow ODE. The Runge-Kutta algorithm requires a step size. Rather than use sophisti- cated algorithm for adaptive step size control, we decided take advantage of the fact that we were already periodically computing the sharpness: at each step, we set the step size to a X, where is a tunable parameter and À is the most recent value for the sharpness. We set a = 1 or a = 0.5.
1.6 RANDOM PROJECTIONS
In order ascertain whether gradient descent at step size 1) initially followed the gradient flow trajectory (and, if so, for how long), we monitored the & distance in weight space between the gradient flow solution time and the gradient descent iterate at step t/n, One way do would be as follows: (a) when rumning gradient flow, save the weights after every At units of time, for some parameter At: (b) when running gradient descent, save the weights each (At) )-th step: (c) plot the difference between these two sequences. (Note that this approach requires At be divisible by n).)
We essentially used this approach, but with one modification: regularly saving the entire net- work weight vector would have consumed a large amount of disk space, so we instead saved low- dimensional random projections of the network weights. To be clear, let d be the number of net- work weights, and let be the number of random projections (a tunable parameter chosen such that k < d). Then first generated matrix M 2 by sampling each entry i.i.d from the standard normal distribution. During training, rather than periodically save the whole weight vector (a d-dimensional vector), premultiplied this vector by the matrix M to obtain a k-dimensional vector, and we periodically saved these vectors instead. Then we plotted the lz distance between the low-dimensional vectors from gradient flow, and the low-dimensional vectors from gradient descent.
32

---

Published as a conference paper at ICLR 2021
J EXPERIMENTS, VARY ARCHITECTURES
In this appendix, we fix the task as that of fitting a 5,000-sized subset of CIFAR-10, and we verify our main findings across broad range of architectures.
Procedure We consider connecco networks and convolutional networks, the latter with both max-pooling and average pooling. For all of these, we consider tanh, ReLJU, and ELU activations, and for cully-connected networks we moreover consider softplus and hardtanh activations. We train each network with both cross-entropy and MSE loss. See $1.1 for full experimental details.
In each case, we first use the Runge-Kutta method to numerically integrate the gradient flow ODE (see $1.5 for details). For architectures that give rise to continuously differentiable training objec- tives, the gradient flow ODE is guaranteed to have unique solution (which we call the gradient flow trajectory), and Runge-Kutte Kutta will return a numerical approximation to this solution. On the other hand, for architectures with ReLU, hardtanh, or max-pooling, the training objective is not continu- ously differentiable, so the gradient flow ODE does not necessarily have a unique solution, and there no guarantees priori what Runge- Kutta will return (more on this below under the "findings" heading), Still, in both cases, since our implementation of Runge- Kutta automatically adjusts the step size based on the local sharpness in order to remain stable, the Runge- trajectory can be roughly viewed as "what gradient descent would do if instability was not an issue."
We then run gradient descent a range of step sizes. These step sizes rj were chosen by hand so that the quantity 2/1 would be spaced uniformly between Ao (the sharpness at initialization) and Amax (the maximum sharpness along the Runge- Kutte-Kuttaaa trajectory).
Results During Runge-Kutta, we observe that the sharpness tends to continually increase dur- ing training (progressive sharpening). with the exception that when cross-entropy loss is used, the sharpness decreases at the very end of training, as explained in Appendix
During gradient descent with step size 1, we observe that once the sharpness reaches 2/1,, it ceases to increase much further, and instead hovers right at, or just above, the value 2/n. For reasons unknown, it tends to be true that for MSE loss, the sharpness hovers just a tiny bit above the value 2/n, while for cross-entropy loss the gap between the sharpness and the value 2/1 is bit larger.
For each step size, we monitor the distance between the gradient descent trajectory and the Runge- Kutta trajectory that is, we monitor the distance between the Runge-Kutta iterate time / and the gradient descent iterate at step t/1 (see $1.6 for details). Empirically, for architectures that give rise to continuously differentiable training objectives, we observe that this distance is nearly zero before the sharpness hits 2/1,, and it starts to climb immediately afterwards. This means that gradient descent closely tracks the gradient flow trajectory so long as the sharpness remains less than 2/1.. Note that this finding was not a foregone conclusion: gradient descent is guaranteed to track the gradient flow trajectory the limit of infinitesimal step sizes (since gradient descent is the forward Euler discretization of the gradient flow ODE), but for non-infinitesimal step sizes, there is discretization error, which is studied in Barrett & Dherin (202I). Our empirical finding is essentially that this discretization error is small compared to the difference between trajectories caused by instability.
On the other hand, for architectures with non-differentiable components such ReLU max- pooling, we sometimes observe that gradient descent tracks the Runge-Kutta trajectory so long as the sharpness remains less than 2/1,, but we also sometimes observe that the gradient descent trajectories differ from one another (and from Runge-Kutta) from the beginning of training. In the former case, can infer that the gradient flow trajectory apparently does exist, and is returned by Runge-Kutta in the latter case, we can infer that either (a) the gradient flow trajectory does not exist, or (b) that does exist (and is returned by Kunge-Kutta), but the step sizes we used for gradient descent were too large to track it.
33

---

Published as a conference paper at ICLR 2021
J.1. FULLY-CONNECTED TANH NETWORK
J.1.1
SQUARE LOSS
train loss
sharpness
train accuracy
500
400
Stsrpnes
tran
0.4
300
0.6
0.4
20
40
6.0
80
20
40 60
80
60
80
Figure 28: Gradient flow. We train the network to 99% accuracy with gradient flow, by using the Runge-Kutta method to discretize the gradient flow ODE (details in $I.5). We plot the train loss (left), sharpness (center), and train accuracy (right) over time. Observe that the sharpness tends to continually increase (except for slight decrease at initialization).
(Iberation)
sharpness (iteration)
accuracy (Iteration)
1.0
0.8
005
sharphes
300
0.4
denn
200
0.4
0-
2200
100
0.2
2/400
5000
10000
15000
5000
10000
15000
5000
10000
15000
Iteration
Steation
iberation
distance from gradient flow
sharpness (time)
test acturacy (iteration)
40
0.8
distarce
30
charprezz
300
0.6
200
0.2
100
0.0
10000
time
time
iteration
Figure 29: Gradient descent. We train the network to 99% accuracy using gradient descent range of step sizes 17. Top left: we plot the train loss curves, with a vertical line marking the iteration where the sharpness first crosses 2/1j. Observe that the train loss monotonically decreases before this point, but behaves non-monotonically afterwards. Top middle: we plot the sharpness (measured at regular intervals during training). For each step size, the horizontal dashed line of the appropriate color marks the maximum stable sharpness 2/n. Observe that the sharpness tends to increase during training until reaching the value 2/n, and then hovers just bit above that value. Bottom left: we track the & distance between (random projections of) the gradient flow iterate at time t and the gradient descent iterate iteration t/n (details in $1.6). For each step size 1), the vertical dotted line marks the time when the sharpness first crosses 2/n. Observe that the distance is essentially zero until this time, and starts to grow afterwards. From this, we can conclude that gradient descent closely tracks the gradient flow trajectory (moving at speed 17) until the point that trajectory where the sharpness reaches 2/n. Bottom middle: to further visualize the previous point, we plot the evolution of sharpness during gradient descent, but with time (= iteration step size) the x-axis rather than iteration. We plot the gradient flow sharpness in black. Observe that the gradient descent sharpness matches the gradient flow sharpness until reaching the value 2/n-
34

---

Published as a conference paper at ICLR 2021
J.1.2
CROSS-ENTROPY LOSS
train loss
Sharpness
train accuracy
150
2.0
0.8
train l035
sharpics
100
0.6
0.4
20
30
40
10
20
30
40
50
10
20
30
40
50
time
Figure 30: Gradient flow, Refer to the Figure 28 caption for more information. Additionally, in this figure the sharpness drops at the end of training due to the cross-entropy loss.
train (Iteration)
sharpness (Iteration)
train accuracy (tteration)
200
0.8
troin 1035
1.5
Sharpnez
0-
1.0
30
0= 2/0
0- 2/80
0.2
2/110
1000
15:00
2000
2500
1000
1500
2000
2500
2000
2500
distance from gradient flow
sharpness (time)
test accuracy (iteration)
150
30-
sharpnecs
100
0.0
20
0.0
10
20
40
60
500
1000
25000
2000
2500
time
time
Berrtion
Figure 31: Gradient descent. Refer the Figure 29 caption for more information.
35

---

Published as a conference paper at ICLR 2021
J.2
FULLY-CONNECTED ELU NETWORK
J.2.
SQUARE LOSS
train loss
sharpness
train accuracy
0.6
600
Itsiprezs
0.8
400
0.6
ton
200
20
60
20
60
20
60
Figure 32: Gradient flow. Refer to the Figure 28 caption for more information.
train loss (iteration)
sharpness (iteration)
train accuracy (Iteration)
1.0
500
train 1055
480
300
0.6
n= 2/20
0.4
0-
0.2
2/400
0.0
5000
10000
15000
20000
5000
10000
15000
20000
5000
10000
15000
20000
Iteration
Iteration
Iteratioo
gradient
Sharppess (timc)
accuracy
1.0
0.8
Sharproc
600
0.4
20-
200
0.2
100
125
5000
10000
15000
20000
Figure 33: Gradient descent. Refer to the Figure 29 caption for more information.
36

---

Published as a conference paper at ICLR 2021
J.2.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
2.0
0.8
shargneze
100
0.6
train
0.5
10
20
40
10
20
30
40
10
20
30
40
time
Figure 34: Gradient flow, Refer to the Figure 28 caption for more information. Additionally, in this figure the sharpness drops at the end of training due to the cross-entropy loss.
train loss (Iteratioo)
sharpness (Iteration)
train accuracy (Iteration)
200
1.0
0.8
1035
1.5
Sharpeis
troin
3.0
100
0= 2/60
0-
2/00
0.2
2/120
1500
2000
2500
500
1000
1500
2000
2500
Itoration
distance from gea flow
sharpness (time)
test accuracy (Iteration)
8.0
0.8
share 02222222
150
60
0.0
0.4
0.0
20
30
40
60
500
1000
1500
2000
2500
time
time
Eeration
Figure 35: Gradient descent. Refer the Figure 29 caption for more information.
37

---

Published as a conference paper at ICLR 2021
J.3
FULLY-CONNECTED SOFTPLUS NETWORK
J.3.1
SQUARE LOSS
train loss
sharpness
train accuracy
1000
0.8
0.8
0.0
600
0.4
0.2
50
100
150
100
150
100
150
Figure 36: Gradient flow. Refer to the Figure 28 caption for more information.
train loss (iteration)
sharpness [iteratonn]
train accuracy (Iteration)
1,00
600
1.0
trai lo0s
0.75
400
0.6
0.50
0.4
n- 2/00
2/400
200
0.00
0.0
10000
20000
30000
40000
50008
0
10000 20000
300ce 40000 Iteration
50000
0
10000
20000
10000
40000 50000
Iteration
Iteration
gradient
sharpness (timc)
accuracy (iteration)
60
1000
1.0
0.8
800
40
scouracy
6000
0.4
29
0.2
50
200
150
50
150
200
10000
20000
30000
40000 50000
Figure 37: Gradient descent. Refer to the Figure 29 caption for more information.
38

---

Published as a conference paper at ICLR 2021
J.3.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
2,0
200
0.8
shargnez
150
train
100
0.4
0.5
0.2
20
40
20
40
60
20
40
60
Figure 38: Gradient flow. Refer to the Figure 28 caption for more information. Additionally, in this figure the sharpness drops at the end of training due to the cross-entropy loss.
train (Iteration)
sharpness (Iteration)
train accuracy (iteration)
0.8
2.5
2/20
1.0
100
2/.00
0.5
0.2
1) 2/280
2000
4000
5000
6000
2000
4000
1000
3000
4000
5000 6000
destance from gradient flow
sharpness (time)
test aecuracyy (Iteration)
200
disancee
100
15.
0.6
50
0.0
20
30
60
20
40
60 time
80
100
120
1000
2000
3000
4000
5000
6000
Lime
Iteation
Figure 39: Gradient descent. Refer the Figure 29 caption for more information.
39

---

Published as a conference paper at ICLR 2021
J.4 FULLY-CONNECTED RELU NETWORK
Note that since the ReLU activation function is not continuously differentiable, the training objective continuously differentiable, and a unique gradient flow trajectory is not guaranteed exist.
J.4.1
SQUARE LOSS
train
sharpness
accuracy
1.0
150
sharpnics
0.3
1.0
0.6
0.2
20
40
40
60
Figure 40: Runge-Kutta. We train the network to 99% accuracy using the Runge- (Since unique gradient flow trajectory is not guaranted to exist, we hesitate to call this "gradient flow"; Kange-Kutta, should essentially be viewed as gradient descent with a very small step size.) Observe that the sharpness tends to continually increase.
Runge-Kutta Kutta algorithm.
train less (iteration)
sharpness (iteration)
train accuracy (Iteration)
0.6
150
0.8
logs
tran
0.4
100
0.6
2 2/60
2/90
0.2
50
2/120
0.0
0.0
1000
ieration
Reeatio
Reration
distance from gradient Rlw
sharpness (time)
test accuracy (iteration)
1.0
150
0.4
30
40
60
20
40
60
100
120
1000
2000
3000
4000
5000
time
iteration
Figure 41: Gradient descent. All panes except bottom left: refer to the Figure 29 caption for more information. Bottom left: We track the &g distance between (random projections of) the Runge-Kutte iterate at time and the gradient descent iterate at iteration l/1). For each step size 7). the vertical dotted loss marks the time when the sharpness first crosses 2/n, In contrast Figure 29 here the distance between gradient descent and gradient flow starts to noticeably grow from the beginning of training. From this we conclude that for this architecture, gradient descent does not track the Runge- Kutta trajectory at first. Furthermore, for this architecture, observe that the training loss starts behaving non-monotonically (and the sharpness begins to plateau) before the sharpness hits 2/m n. As discussed in Appendix A ("Caveats""), we believe that this is also due to the fact that ReLU is non-differentiable.
40

---

Published as a conference paper at ICLR 2021
J.4.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
100
2.0
shargneze
75
0.8
train
0.4
0.5
25
0.2
0-
10
20
30
40
10
20
30
40
10
30
40
Figure 42: Runge-Kutta. Refer to the Figure 40 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
(Iteration)
sharpness [tteratoon]
train acouracy (Iteration
2.0
0.8
1.5
trin
1.0
0
2/40
0.5
0.2
2/80
500
500
1000
1500
distance from gradient flow
sharpness (time)
test accuracy (iteration)
disancce
0.0
50
50
0.0
20
30
40
20
40
500
1000
1500
time
time
ieeratio
Figure 43: Gradient descent. Refer the Figure 41 caption for more information.
41

---

Published as a conference paper at ICLR 2021
J.5
FULLY-CONNECTED HARD TANH NETWORK
Note that since the hardtanh function is not continuously differentiable, the training objective is not continuously differentiable, and unique gradient flow trajectory is guaranteed exist.
J.5.1
SQUARE LOSS
train
sharpness
accuracy
1.0
800
train 1055
0.4
sharpnece
600
0.6
400
0.4
200
20
60
20
80
Figure 44: Runge-Kutta, Refer to the Figure 40 caption for more information.
train loss (iteration)
sharpness (iteration]
train (iteration)
800
1.0
1055
shappezs
tran
0.4
0- 2/400
200
2- 2/800
10000
20000
30000
10000
20000
30000
10000
20000
30000
iteration
steration
Iteration
distancc from gradient flow
sharpness (time)
test accuracy (Iteration)
1.0
desance
40
Sharpness
800
0.8
0.6
20
200
0.2
0.0
20
10000
20000
30000
Ieration
Figure 45: Gradient descent. Refer to the Figure 41 caption for more information.
42

---

Published as a conference paper at ICLR 2021
J.5.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
200
2.0
train loss
shargncze
150
0.8
100
0.6
10
20
30
40
10
20
40
10
20
30
40
Figure 46: Runge-Kutta. Refer to the Figure 40 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
sharpness (Iteration)
train accuracy (Iteration)
1.0
0.8
1035
1.5
Sharpness
100
n=2/4
troin
1.0
2-200
2/1.0
0.2
2/230
500
1500
2500
2000
500
1000
1500
2000
1000
1500
2000
2500
3000
Iteration
distance from gradient flow
sharpness (time)
test accuracy (iteration)
150
0.8
30-
sharewee
0.0
100
20
0.4
0.0
20
30
40
20
30
40
50
500
1000
2500
2000
2500
3000
time
tmme
Iteraiin
Figure 47: Gradient descent. Refer the Figure 41 caption for more information.
43

---

Published as a conference paper at ICLR 2021
J.6 CONVOLUTIONAL TANH NETWORK WITH MAX POOLING
Note that since max-pooling is not continuously differentiable, the training objective is not continu- ously differentiable, and so unique gradient flow trajectory is not guaranteed to exist.
J.6.1
SQUARE LOSS
train
sharpness
accuracy
0.8
train 1035
shargnics
100
0.6
0.3
75
0.2
2000
100
200
100
Figure 48: Runge-Kutta. We train the network to 99% accuracy using the Runge- (Since unique gradient flow trajectory is not guaranted to exist, we hesitate to call this "gradient flow.") Observe that the sharpness tends to continually increase.
Kunge-Kutta Kutta algorithm.
train loss (Iteration)
sharpness (iteration)
accuracy (Iteration
0.6
120
1.0
8 2/40
tain 103s
0.4
Sharpiess
1.00
2/00
0.4
2/120
5000
1
1000
15000
5000
10000
15000
5000
10000
15000
Iteration
iteration
teration
distance from gradient flow
sharpness (time)
test accuracy (iteration)
20
125
0.8
15
sharpeooss
1:00
0.6
30
0.2
0.0
time
time
Iteraion
Figure 49: All panes except bottom left: refer to the Figure 29
caption for more information.
Bottom left: We track the & distance between (random projections of) the Runge- Runge-Kutta Kutta iterate at time & and the gradient descent iterate iteration t/n. For each step size 17, the vertical dotted loss marks the time when the sharpness first crosses 2/1j. Observe that the distance is essentially zero until this time, and starts grow afterwards. From this, we can conclude that even though the training objective is not differentiable, gradient descent does closely track the Runge- trajectory (moving at a speed 17) until the point on that trajectory where the sharpness reaches 2/1.
44

---

Published as a conference paper at ICLR 2021
J.6.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
2.0
train l05s
shargoedss
60
0.8
0.2
25
50
75
100
125
25
50
75
100
125
25
50
75
100
125
time
Figure 50: Runge-Kutta. Refer to the Figure 48 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
(teeattonn)
sharpness (Iteration)
train accuracy (Iteration)
60
0.8
troin 1935
1.5
40
3.0
n- 2710
20
0.2
n-
2770
1000
3000
4000
2000
3000
1000
3000
4000
5000
distance from gradient flow
sharpness (time)
test accuracy (iteration)
60
30-
0.6
80
20
20
0.0
25
50
75 time
100
125
50
75 time
125
1000
2000
3000
4000
5000
Iteratio
Figure 51: Gradient descent. Refer the Figure 49 caption for more information.
45

---

Published as a conference paper at ICLR 2021
J.7 CONVOLUTIONAL ELU NETWORK WITH MAX POOLING
Note that since max-pooling is not continuously differentiable, the training objective is not continu- ously differentiable, and so unique gradient flow trajectory is not guaranteed to exist.
J.7.1
SQUARE LOSS
sharpness
accuracy
1000
0.8
train 0055
750
0.6
500
230
100
150
100
50
Figure 52: Runge-Kutta, Refer to the Figure 48
caption for more information.
train (iteration)
sharpness (eteration)
train (iteration)
0.6
800
1.0
6- 2/20
tran 1085
shappezs
0.6
0-
6 2/600
600
0.4
n-
0.2
0.0
20000 40000 60000 8000010000
20000 40000 60000 80000 80000 100000 0
0
20000 40000
50000 800ce 100000
Steraiio
iteration
heeation
distance from gradient flow
sharpness (time)
test accuracy (Iteration)
1.0
1000
20
0.8
destnce
15-
0.6
0.4
250
0.2
0.0
150
150
200
20000 40000
80000 100000
Reration
Figure 53: Gradient descent. Refer to the Figure 49
caption for more information.
46

---

Published as a conference paper at ICLR 2021
J.7.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
2.0
1000
0.8
0.6
500
0.5
0.2
10
20
30
40
10
20
40
50
10
20
30
40
50
time
Figure 54: Runge-Kutta. Refer to the Figure 48 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
(Iteration)
sharpness (Iteration)
accuracy (Iteration)
1.0
1000
0.8
troin 1035
1.5
1.0
21400
0 2700
0.2
2/1000
5000
30000
15000
20000
25000
5000
10000
15000
20000
25000
5000
10000
15000
20000 25000
discance from gradient flow
sharpness ess (tme)
test accuracy (iteration)
0.8
1000
30-
0.6
20
500
0.4
0.0
40
50
10
20
30
5000
10000 15000 20000
25000
tmme
Eime
Reration
Figure 55: Gradient descent. Refer the Figure 49 caption for more information.
47

---

Published as a conference paper at ICLR 2021
J.8
CONVOLUTIONAL RELU NETWORK WITH MAX POOLING
Note that since ReLU and max-pooling are not continuously differentiable, the training objective is continuously differentiable, and unique gradient flow trajectory guaranteed exist.
J.8.1
SQUARE LOSS
sharpness
accuracy
0.5
0.8
1035
0.4
0.0
train
0.3
0.2
150
100
50
Figure 56: Runge-Kutta. Refer to the Figure 48
caption for more information.
less (iteration)
sharpness (iteration)
train (iteration)
1.0
0 22200
tein lo0s
0.4
shappezz
27700
400
0.4
200
20000 40000 60000 80000
100000
20000 40000
60000
8c000 100000
20000
40000
60000
80000 100000
steation
iteration
iteratioo
distance from gradient flow
sharpness (Ume)
test accuracy (tteratio)
1.0
40
0.8
desance
30-
sharpnoss
1000
0.6
500
10
0.2
0.0
1.00
200
250
300
20000
40000
60000
80000
100000
iteration
Figure 57: Gradient descent. Refer to the Figure 49
caption for more information.
48

---

Published as a conference paper at ICLR 2021
J.8.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
2000
2.0
1500
0.8
1000
0.4
0.5
0.2
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
Figure 58: Runge-Kutta. Refer to the Figure 48 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
sharpness (Iteration)
train accuracy (Iteration)
1.0
1000
0.8
1935
train
1.0
500
0- 2/00
0-
0.2
2700
10000
15000
20000
5000
15000
20000
15000
20000
Iteration
distance from gradient flow
sharpness (time)
test accuracy (Iteration)
2000
15:0
0.8
30-
0.6
20-
0.4
500
0.0
100
30
50
20
5000
10000
15000
20000
time
time
Iteraion
Figure 59: Gradient descent. Refer the Figure 49 caption for more information.
49

---

Published as a conference paper at ICLR 2021
J.9
CONVOLUTIONAL TANH NETWORK WITH AVERAGE POOLING
J.9.1
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
1.0
1035
sharpniss
30
0.8
20
0.4
0.5
0.2
50
100
150
200
250
300
50
100
150 time
200
250
300
100
150
200
250
300
time
time
Figure 60: Gradient flow. Refer to the Figure 28 caption for more information. Additionally, in this figure the sharpness drops at the end of training due to the cross-entrope loss.
train loss (iteration)
sharpness (tteration)
train (iteration)
1.0
40
2.0
1035
sharonuas
0.6
138
20
0.4
0= 210
10
Q- 2/30
1000
2000
1000
4000
1000
2000
3000
4000
0
10000
2000
3000
4000
lberation
Iteratio
iteration
distance from gradient
sharpness (time)
test accuracy (iteration)
40
0.8
distance
sharpnoss
30
20
0.4
10
0.2
0.0
250
300
2000
2000
3000
4000
Iteration
Figure 61: Gradient descent. Refer to the Figure caption for more information.
50

---

Published as a conference paper at ICLR 2021
J.10
CONVOLUTIONAL ELU NETWORK WITH AVERAGE POOLING
J.10.1
SQUARE LOSS
train loss
sharpmess
train accuracy
1.0
1000
0.8
0.6
190
0.3
500
0.2
100
200
300
400
500
1.00
200
300
400
500
1.00
200
500
400
500
Figure 62: Gradient flow. Refer to the Figure 28 caption for more information.
train loss (iteration)
sharpness (iteration)
train accuracy (iteration)
1.0
0=20
train 1655
accura.r
0.6
0- 2200
0.4
100
0.0
50000
100000 150000
200000
0
50000
100000
150000
200000
0
50000
100000
150000 200000
iteration
Iteration
Iteration
gradient flow
sharpness (timc)
accuracy (iteration)
50
1.0
40
1000
0.8
shargness
750
accuracy
0.4
300
0.2
100
500
1000
1500
50000
100000
150000
200000
Figure 63: Gradient descent. Refer to the Figure 29 caption for more information.
51

---

Published as a conference paper at ICLR 2021
J.10.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
1500
0.8
0.6
500
0.5
20
40
60
80
100
20
40
80
100
20
40
60
80
100
Figure 64: Gradient flow, Refer to the Figure 28 caption for more information. Additionally, in this figure the sharpness drops at the end of training due to the cross-entropy loss.
train (Iteration)
sharpness (fteration)
train accuracy (iteratioo)
1000
1.0
750
0.8
troin 1035
1.5
3.0
2250
2/500
0.2
2/750
40000
10000
20000
30000
40000
10000
20000
30000
40000
Iteration
di distance from gradient flow
sharpness (t me)
test accuracy (Iteration)
1500
0.8
30-
1000
0.6
20
0.4
0.0
40
60
100
20
40
60
80
10000
20000 iteratin
30000
40000
time
time
Figure 65: Gradient descent. Refer the Figure 29 caption for more information.
52

---

Published as a conference paper at ICLR 2021
J.11 CONVOLUTIONAL RELU NETWORK WITH AVERAGE POOLING
Note that since ReLU is not continuously differentiable, the training objective not continuously differentiable, and unique gradient flow trajectory is guaranteed exist.
J.11.1
SQUARE LOSS
sharpness
accuracy
0.5
2000
0.8
1500
0.0
train
0.3
0.2
100
4.00
200
400
200
400
time
Figure 66: Runge-Kutta. Refer to the Figure 48
caption for more information.
train loss (iteration)
sharpness (iteration]
train (iteration)
0.6
800
1.0
6- 2/50
0-
trin los3
0.4
0.6
6 2200
0.4
2/600
0.2
50000 100000150000 200000 250000
50000 100000100020000200000
5,000 100001500020000300000
Ieration
iterat
iteration
distance from gradient flow
sharpness (time)
test accuracy (Iteration)
1.0
2000
20
0.8
Jestance
1.5
sharpness
1500
0.6
1.00
500
0.2
1.00
200
800
1000
0.0
50000 1000001 50000000002226000000000000000000 50000
Figure 67: Gradient descent. Refer to the Figure 49 caption for more information.
53

---

Published as a conference paper at ICLR 2021
J.11.2
CROSS-ENTROPY LOSS
train loss
sharpness
train accuracy
0.8
0.6
train
0.4
0.5
0.2
20
40
100 60
80
20
40
60
80
100
20
60
80
100
time
Figure 68: Runge-Kutta. Refer to the Figure 48 caption for more information. Additionally, this figure the sharpness drops at the end of training due to the cross-entropy loss.
(iteration)
sharpness (Iteration)
accuracy (iteration)
1.0
1000
0.8
train 1035
1.5
n-
1.0
21400
0.2
21800
10000
20000
30000
40000
50000
10000
20000
30000
40000
10000
20000
30000
40000 50000
distance from gradient flow
sharpness (time)
test accuracy (Iteration)
2000
0.8
15
0.6
30
2000
0.4
0.0
60 time
80
100
75
100
125
10000
20000
30000
40000
50000
time
iteration
Figure 69: Gradient descent. Refer the Figure 49 caption for more information.
54

---

Published as a conference paper at ICLR 2021
K BATCH NORMALIZATION EXPERIMENTS
In this appendix, we demonstrate that our findings hold for networks that are trained with batch normalization (BN) (loffe & Szegedy 2015). We experiment on a size-5,000 subset of CIFAR- 10, and we consider convolutional networks with three different activation functions: ELU (Figure tanh (Figure 72 73) and ReLU (Figure 7 See $1.3 for experimental details.
70 71
Empirically, our findings hold for batch-normalized networks. The one catch is that when training batch-normalized networks at very small step sizes, it is apparently inadequate to measure the sharp- ness directly at the iterates themselves, as we do elsewhere in the paper. Namely, observe that in Figure 7I])). when we run gradient descent at the red step size, the sharpness (measured directly at the iterates) plateaus bit beneath the value 2/m. At first, this might sound puzzling: after all, if the sharpness is less than 2/1 then gradient descent should be stable. The explanation is that the sharp- ness in between successive iterates does in fact cross 2/n. In Figure ()). track the maximum sharpness on the path "in between" successive iterates. (To estimate the maximum sharpness be- tween pair of successive iterates, we compute the sharpness at grid of eight points spaced evenly between them, and then take the maximum of these values.) Observe that this quantity does rise to 2/n and hover there. We do not know why measuring the sharpness between iterates is necessary for aatch-normalized networks, whereas for non-BN networks it suffices to measure the sharpness only at the iterates themselves.
In K. we reconcile these findings with Santurkar et al. 2018).
train 1oss
sharpness
train accuracy
1.50
1500
0.75
1000
0.25
0.00
Une
1876
Figure 70: We train a ELU CNN (+ BN) using gradient flow.
sharpnesx
2000
2000
3:00
500
15020
Wim
(d) train acournry literation)
(e) sharpness terate literationn
(1) distance from gradient flow
3.0
20-
0.5
3000
ataree 20
300 677
Iaratoo 19020
Iaration
We train an ELU CNN (+ BN) using gradient decent at range of step sizes. (a) we plot the train loss, with a vertical dotted line marking the iteration where the sharpness on the path first crosses 2/n. The inset shows that the red curve is indeed behaving non-monotonically. (b) we track the sharpness "between iterates." " This means that instead of computing the sharpness right at the iterates themselves (as we do elsewhere in the paper, and in pane (e) here), we compute the maximum sharpness on the line between between successive iterates. Observe that this quantity rises 2/1 (marked by the horizontal dashed line) and then hovers right at, or just above that value. (c) we plot the same quantity by "time" (= iteration x step size) rather than iteration. (e) we plot the sharpness at the iterates themselves. Note that for the red step size, this quantity plateaus at a value that is beneath 2/1. (f) distance from the gradient flow trajectory.
Figure 71:
55

---

Published as a conference paper at ICLR 2021
train loss
shargness
train accuracy
0.50
1000
500
2
20
15
20
13
15
20
0.00
10
15
Figure 72: We train a tanh CNN (+ BN) using gradient flow.
train (Iteration)
iiteration
2000
2200
33:00
10th
1000
2200
3005
19383
13000
20000
5000
10020
23000
20000
Ewratoo
6ime
accursry
siarppesss
17.0
Sitrree
20
22200
2/506
7/1540
2000
500
20
77000
0.0
20000
Breiiin
8mm
Figure 73: We train tanh CNN (+ BN) using gradient decent.
train
sharpmess
train accuracy
1000
300
0.50
0.25
2
0
Figure 74: We train ReLU CNN (+ BN) using gradient flow.
(2) train loso (iteation)
shar press between Iterates (toration) ono
aharpneds between iterates (amme
500
20 steps
200
1030
2000
3000
1000
2000
3000
Iaration
(d) train lnteration)
herate Interationn
(1) distance gradient
15-
000
2010
2/100
400
0- 24210
o- 24000
200
5.0
1030
2000
3000
1000
2000
3200
Excation
Umno
Figure 75: We train a ReLU CNN (+ BN) using gradient descent.
K.1 RELATION TO SANTURKAR
2018
We have demonstrated that the sharpness hovers right at (or just above) the value 2/n when both BN and non-BN networks are trained using gradient descent at reasonable step sizes. Therefore, at least the case of ball-batch gradient descent, cannot be said that batch normalization decreases the sharpness (i.e. improves the local [-smoothness) along the optimization trajectory.
56

---

Published as a conference paper at ICLR 2021
Santurkar al. (2018) argued that batch normalization improves the effective smoothness along the optimization trajectory, where effective smoothness is defined as the Lipschitz constant of the gradient in the update direction (i.e. the negative gradient direction, for full-batch GD). That is, given an objective function f. an iterate 0, and a distance o, the effective smoothness of f at parameter and distance is defined in Santurkar et al. (2018)
[V
V f())ll
sup rE[0,a]
|V)lll))lz
where the sup can be numerically approximated by evaluating the given ratio at several values y spaced uniformly between 0 and
In Figure 76 we train two ReLU CNNs - one with BN, one without - at the same set of step sizes, and we monitor both the sharpness (i.e. the L-smoothness) and the effective smoothness. When computing the effective smoothness, we use distance a = 1/. Observe that for both the BN and the non-BN network, the effective smoothness initially hovers around zero, but once gradient descent enters the Edge of Stability, the effective smoothness jumps to the value 2/1) and then remains there. Thus, least for full-batch gradient descent on this particular architecture, batch normalization does not improve the effective smoothness along the optimization trajectory. (Despite this, note that for each step size, the BN network trains faster than the non-BX network, confirming that BN does accelerate training.)
CNN with Normalization
train loss
sharpness (smoothness)
effective smoothness
22200
1000
1000
2-
2m 22600
6
750
750
20000
500
500
250
250
0-
500
1000
1500
2000
2500
3000
10000
1500
2500
300n
1000
1500
2500
ReLU CNN without Batch Normalization sharpness (smoothness)
train loss
effective smoothness
1000
1000
2- 2/400
1.5
27600
750
750
1.0
O-
500
500
0.5
250
250
0
5000
10000
15000
20000
5000
10000
15000
20000
5000
10000
15000
20000
Figure 76: On a 5,000-size subset of CIFAR-10, we train ReLU CNN both with BN (top row) and without BN (bottom row) at the same grid of step sizes. We plot the sharpnes//mmothnnss (center column) as well as the effective smoothness (Santurkar al. (018) (right column). Observe that for both networks, the effective smoothness hovers around zero initially, and jumps up to 2/n once gradient descent enters the Edge of Stability.
Note that this finding is actually consistent with Figure 4(c) Santurkar al. (2018), which is meant to show that BN improves effective smoothness when training a VGG network using SGD. Their Figure 4(c) shows that during SGD with step size 1 0.1, the effective smoothness hovers around the value 20 for both the BN and the non-BN network. Since 20 2/ (0.1)) this is fully consistent with our findings (though they use SGD rather than full-batch GD). Figure 4(c) does show that the effective smoothness behaves more regularly for the BN network than for the non-BN network. But we disagree with their interpretation of this figure as demonstrating that BN improves the effective smoothness during training.
The other piece of evidence in /Santurar et al. (0118) in support of the argument that batch normal- ization improves the effective smoothness during training is their Figure 9c)) This figure shows that a deep linear network (DLN) trained without BN has a much larger (i.e. worse) effective smooth- ness during training than DLN trained with BN. However, for this figure, the distance a used to compute effective smoothness was larger than the training step size v by a factor of 30. The effec- tive smoothness at distances larger than the step size does not affect training. We have verified that
57

---

Published as a conference paper at ICLR 2021
when effective smoothness is computed distance equal to the training step size (i.e. 1)) the effective smoothness for the DLN with BN and for the DLN without BN both hover right at 2/7).
Specifically, in Figure 77 and Figure 78 train a DLN both with and without BN (respectively), and we measure the effective smoothness at a distance a = 301, as done in Figure 9() Santurkar et al. (2018). We use the same experimental setup and the same step size of 17 = 1e-6 they do, and we repeat the experiment across four random seeds. Observe that when training the BN network, the effective smoothness hovers right at 2/11 (marked by the horizontal black line), whereas when training the non-BN network, the effective smoothness is much larger. This is consistent with Figure (c) in Santurkar al. (2018). However, in Figure 79 and 80 we measure the effective smoothness at the actual step size a 17. When effective smoothness Is computed in this way, we observe that for both the network with BN and the network without BN, the effective smoothness hovers right at 2/n. Therefore, we conclude that there is no evidence that the use of batch normalization improves either the smoothness or the effective smoothness along the optimization trajectory. (That said, this network permits training with larger step
experiment possibly explains why the sizes.)
smnodiniess
1011
2/g
20
2/n
101
2/g
2/
effective
20a
10*
10a
101
10a
100
16a
5000 lberatioo
10000
5000
10000
5000 Iberation
10000
5000 iteration
10000
Reration
Figure 77: When training a deep linear network without BN, we measure effective smoothness at a distance a = 301y that is 30x larger than the step size.
seed 0
seed
seod 2
seed
effectie
shooddness
10
2/n
1011
25m
101
2/g
10
2in
10a
10a
168
5000 Iteration
10000
5000
10000
5000 Iberation
10000
5000 Iteration
10000
iteration
Figure 78: When training a deep linear network with BN, we measure effective smoothness at a distance 30n that is 30x larger than the step size.
Seed
seed 1
seed
Seed
offective
1013
2/g
1033
---
2/n
10t3
---
2/n
1033
---
2/n
10P
109
10a.
10a
107
10a
10'
10a
5000
10000
10000
5000
10000
5000 Iteration
10000
Iteration
iteratioo
Iberation
Figure 79: When training a deep linear network without BN, we measure effective smoothness at distance a = 7 that is equal to the step size.
58

---

Published as a conference paper at ICLR 2021
seed 0
seed
seed
seed
effecivve
smicodiiess
1011
1011
21g
10
10th
10a
16a
5000 lteratiion
10000
5000
10000
5000 Iberation
10000
5000 Iteration
10000
Itertioo
Figure 80: When training a deep linear network with BN, we measure effective smoothness at distance 77 that is equal to the step size.
59

---

Published as a conference paper at ICLR 2021
L ADDITIONAL TASKS
So far, we have verified our findings on image classification and language modeling. In this ap- pendix, we verify our findings on three additional tasks: training Transformer on the WikiText-2 language modeling dataset (LI), training one-hidden-layer network on one-dimensional toy regression task (L.2) and training deep linear network (L.3).
L. TRANSFORMER ON WIKITEXT-2
We consider the problem of training Transformer on the WikText-2 word-level language modeling dataset (Merity et al., 2016). See $1.4 for full experimental details. Figure 81 we train using gradient flow (only partially, not to completion). Observe that the sharpness continually rises. In Figure 82 we train using gradient descent at a range of step sizes. Consistent with our general findings, for each step size 7/, we observe that the sharpness rises to 2/n and then hovers right at, or just above, that value. However, for this Transformer, we do not observe that gradient descent closely tracks the gradient flow trajectory at the beginning of training.
train loss
sharpness
10
20
50
100 time
150
200
50
100 time
150
200
Figure 81: Training Transformer using gradient flow.
train loss (iteration)
sharpness (iteration)
0-
n- 2/20
train
2140
2500
5000
7500
10000
12500
15000
2300
5000
7500
10000
12500 15000
iteration
Deration
distance from gradient flow
sharpness (time)
60
astanoc
1.-
pharpiccss
40
50
100
150
200
50
100
150
200
time
time
Figure 82: Training Transformer using gradient descent. We train a Transformer for 15,000 iterations on the WikiText-2 language modeling dataset. (For this problem, training to completion would not be computationally practical.) Top left: plot the train loss curves, with a vertical dotted line marking the iteration where the sharpness first crosses 2/1,. The train loss decreases monotonically before this line, but behaves non-monotonically afterwards. Top right: we plot the evolution of the sharpness, with a horizontal dashed line marking the value 2/1,. Observe that for each step size, the sharpness rises to 2/1 and then hovers right at, or just above, that value. Bottom left: for the initial phase of training. we plot the distance between the gradient descent trajectory and gradient flow trajectory, with a vertical dotted line marking the gradient descent iteration where the sharpness first crosses 2/n- Observe that the distance begins to rise from the start of training, indicating that for this architecture, gradient descent does not track the gradient flow trajectory initially. Bottom right: for the initial phase of training, plot the evolution of the sharpness by "time" = iteration x y rather than iteration. The black dots are the sharpness during gradient flow.
60

---

Published as a conference paper at ICLR 2021
L.2 ONE- -I---EIEEIINN TOY REGRESSION TASK
Task Our toy regression problem is to approximate Chebyshev polynomial using a neural net- work. To generate toy dataset, take 20 points spaced uniformly the interval [-1, 1]. and label them noiselessly using the Chebyshev polynomial of some degree k. Note that the Chebyshev polynomial of degree & is a polynomial with zeros that maps the domain [-1, 1] the range [-1, 1]. Figure 83 shows the Chebyshev datasets for degree 3, 4, and 5.
Chebyshev degree
Chebyshev degree
Chebyshev degree
0.5
0.5
0.5
0.0
0.0
0.0
0.5
-0.5
-0.5
-1.0
-1.0
-1.0
-1.0
-0.5
0.5
1.0
-10
-0.5
0.0
0.5
-1.0
-0.5
0.5
1.0
Figure 83: The datasets for the toy one-dimensional regression task.
Network For the network, initialized using Xavier initialization. We train using the MSE loss until the loss reaches 0.05.
tanh architecture with hidden layer of h 100 units,
Results In Figure 84] we fit the Chebyshev degree 3, 4, and 5 datasets using gradient flow. Em- pirically, the higher the degree, the more the sharpness rises during the course of training: on the degree polynomial, the sharpness rises by factor 1.2; on degree polynomial, by factor 3.2; and on the degree 5 polynomial, by factor of 63.5.
In Figure 85 and Figure 86 we fit the degree 4 and 5 datasets using gradient descent at a range of step sizes. We observe mostly the same Edge of Stability behavior elsewhere in the paper. The only difference is that for the degree dataset, after the sharpness hits 2/1, the training loss first un- dergoes temporary period of con-monotonicity in which progress is made, and then decreases monotonically until training is finished (in contrast to our other experiments where we observe that training loss behaves non-monotonically the same time it consistently decreasing).
61

---

Published as a conference paper at ICLR 2021
Chebyshev degree
sharpness
0.2
0.1
0.0
10
20
30
40
50
0
10
20
30
40
50
time
time
Chebyshev degree 4
train lo5s
sharpness
sharpness
0.1
0.0
time
time
Chebyshev degree
train loss
sharpner pness
sharppeess
0.1
2000
4000
6000
20000
4000
6000
Figure 84: Fitting Chebyshev polynomials using gradient flow. We gradient flow to fit mme-hidden-layer tanh network to the Chebyshev polynomials degrees 3, 4, and 5. We observe that the sharpness rises more when fitting higher-degree polynomial, likely because the dataset is more complex.
train loss (iteration)
sharpness (iteration)
10
0.3
train loss
sharpness
2/6
0.2
J) 2/8
2/10
0.1
0.0
500
1000
1500
2000
500
1000
1500
2000
Iteration
Iteration
distance from gradient flow
sharpness (time)
distance
60
sharpness
40
20
0
200
400
200
600
time
time
Figure 85: Chebyshev degree (gradient descent). We fit the Chebyshev polynomial of degree using gradient descent at a range of step sizes (see legend).
62

---

Published as a conference paper at ICLR 2021
train loss (iteration)
sharpness (iteration)
200
0.3
sharpness
n=
2/100
train
0.2
100
n 2/150
2/200
0
200000 400000 600000 800000
0
200000 400000 600000 800000 iteration
iteration
distance from gradient flow
sharpness (time)
150
sharpness
200
distance
100
50
100
0
2000
4000
6000
0
2000
4000
6000
8000
time
time
Figure 86: Chebyshev degree (gradient descent). We fit the Chebyshev polynomial of degree using gradient descent at range of step sizes (see legend).
L.3 DEEP LINEAR NETWORK
Task The task is to map inputs X....., Xn. C Rd to targets y1, . ,n C I using function f: Rd - Rd. Error measured using the square loss, i.e. the objective is Let X E RRd be the vertical stack of the inputs, and let Y E Raxd be the vertical stack of the targets. We first generate X random whitened matrix (i.e. _ = D). To generate X random whitened matrix, sample n matrix of standard Gaussians, and then times the Q factor in the QR factorization of that matrix. We then generate Y via Y = XAT, where A dd is random matrix whose entries are sampled i.i.d from the standard normal distribution.
(lf(x)) yillă.
to be y
We use = 50 datapoints with a dimension of d = 50.
Network The function f Rd + Rd is implemented as a L-layer deep linear network: f(x) W.... W2Wix with W, Raxd We initialize all layers of the deep linear network using Xavier initialization: all entries each W. are drawn i.i.d from N(0, }).
We network with L 20 layers.
Results In Figure 87 we train the network using gradient flow. (Since is unclear whether the network be trained zero loss, and how long this would take, arbitrarily chose stop training at time 100.) In Figure 88 we train the network using gradient descent at a range of step sizes. We observe mostly the same Edge of Stability behavior as elsewhere the paper. The only difference is that in Figure 38 the train loss does not really behave non-monotonically - for each step size 7) there is brief blip at some point, but otherwise, the train loss decreases monotonically.
63

---

Published as a conference paper at ICLR 2021
train loss
sharpness
1000
sharpners
2000
1000
10
time
time
Figure 87: Training a deep linear network using gradient flow. We use gradient flow to train deep linear network. Since it is unclear whether this network can be trained to zero loss (or how long that would take), we arbitrarily chose to stop training at time 100.
train loss (iteration)
sharpmess (iteration)
100
2/1000
2/1500
train
50
2/2000
1000
2/2500
0
5000 1000015000 20000
0
5000
10000 15000 2000025000
iteration
iteration
distance from gradient
sharpness (time)
3000
distance
sharpnoes
2000
1000
time
time
Figure 88: Training a deep linear network using gradient descent. We train a deep linear network using gradient descent a range of step sizes (see legend). (Note that the top left pane, the blue, orange and green dotted lines are directly on top of one another.)
64

---

Published as a conference paper at ICLR 2021
M EXPERIMENTS: STANDARD ARCHITECTURES ON CIFAR-10
In this appendix, we demonstrate that our findings hold for three standard architectures on the stan- dard dataset CIFAR-10. 10. The three architectures are: a VGG with batch normalization (Figures VGG without batch normalization (Figures (99gures/99922) and ResNet with batch normalization
89 90
(Figures 93 94.. See $1. for full experimental details.
confirm our main points: (1) so long the sharpness is less than 2/17, the sharpness tends to increase; and (2) if the sharpness reaches 2/1,, gradient descent enters a regime (the Edge of Stability) in which (a) the training loss behaves non-monotonically, yet consistently decreases over long timescales, and (b) the sharpness hovers right at, or just above, the value 2/1z. Moreover, we observe that even though these architectures use the ReLU activation function (which is not continuously differentiable), gradient descent closely tracks the Runge- Kutta- trajectory until reaching the point on that trajectory where the sharpness hits 2/1-
For each of these three architectures,
Furthermore, we observe that for these three standard architectures, the following additional points hold: (1) progressive sharpening occurs to a dramatic degree, and (2) stable step sizes are so small as to be completely unreasonable.
Progressive sharpening occurs to a dramatic degree To assess the degree of progressive sharp- ening, we train these networks using Runge-Kutta gradient flow, which can be viewed as gradient descent with an infinitesimally small step size. (In practice, the Runge- Kutta algorithm does have a step size parameter, and throughout training, we periodically adjust this step size in order to ensure that the algorithm remains stable.) Intuitively, training with gradient flow tell us how far the sharp- ness would rise "if it didn't itt't have to worry about" instability caused by nonzero step sizes. In Figure 89] we train the VGG with BN to completion (99% training accuracy) using Runge- Kutta/ / gradient flow, and find that the sharpness rises from its initial value of 6.38 a peak value of 2227.6. For the other two architectures, progressive sharpening occurs to such degree that it is not computation- ally feasible for to train using Runge- Kutta- / gradient flow all the way to completion. (The reason is that in regions where the sharpness is high, the Runge-Kutte step size must be made small, so Runge-Kutte Kutta requires very many iterations.) Therefore, we instead train these two networks only partially. In Figure 91 for the VGG without BN, we find that the sharpness rises from its initial value of 0.64 to the value 2461.78 at 37.1 % accuracy, when we stop training. In Figure 93] for the ResNet, we find that the sharpness rises from its initial value of 1.07 to the value 760.6 at 43.2% accuracy, when we stop training. Thus, even though we observed in Appendix D that progressive sharpening attenuates as the width of fully-connected networks is made larger, it appears that either: (1) this does not happen for modern families of architectures such as ResNet and VGG, or (2) this does happen for modern families of architectures, but practical network widths lie on the narrow end of the scale.
Stable step sizes are so small as to be unreasonable Recall from 3.3 that if Amax is the maximum sharpness along the gradient flow trajectory, then any stable step size must be less than 2/Amax- Therefore, for these three architectures, because progressive sharpening occurs to dramatic degree (i.e. Amax is extraordinarily large), any stable step size must be extraordinarily small, which means that training will require many iterations. Yet, at the same time, we find that these three networks can be successfully trained in far fewer iterations by using larger step step size. This means that training at a stable step size is extremely suboptimal. We now elaborate on this point:
VGG with BN. For this network, gradient flow terminates at time 15.66, and the maximum sharpness along the gradient flow trajectory is 2227.6. Therefore, the largest stable step size is 2/2227.59 = 0.000897, and training to completion at this step size would take 14.91 00.000000097 = 16622 iterations. Mcanwhile, we empirically observe that the network can also be trained to com- pletion at the much larger step size of n = 0.16 in just 329 itcrations. Therefore, using a stable step size is suboptimal by a factor of at least 06622/320 50.5.
For the other two architectures, since we are unable to train to completion using gradient flow, we unable to obtain tight lower bound for the number of iterations required to run gradient descent to completion at a stable step size. Therefore, by extension, we are unable to compute a tight lower bound for the suboptimality factor of stable step sizes. As a substitute, we will instead compute
65

---

Published as a conference paper at ICLR 2021
both: (1) tight lower bound the suboptimality of training partially at stable step size, and (2) a very loose lower bound on the suboptimality of training to completion at a stable step size.
VGG without BN. For this network. gradient flow reaches 37.1% accuracy at time 8, and the maxi- mum sharpness up through this point on the gradient flow trajectory is 2461.8. Therefore, the largest 0.00081, and training to 37.1% % accuracy at this step size would quire 8/ $/0.00081 = 9, 876 iterations. Meanwhile, we empirically observe that the network can also be trained 37.1% % accuracy at the larger step size of r 0.16 in just 355 iterations. Therefore, when training this network to 37% accuracy, stable step sizes are suboptimal by a factor of at least 9876/355 27.8. This is tight lower bound on the suboptimality of training to 37. 1% accuracy at a stable step size.
stable step size is 2/2461.8
To obtain a loose lower bound on the suboptimality of training the VGG without BN to completion at stable step size, we note that (a) since the maximum sharpness up through time 8 is 2461.8, the maximum sharpness along the entire gradient flow trajectory must be at least 2461.8; and (b) since by time 8 gradient flow has only attained 37.1% training accuracy, the time to reach 99% accuracy (i.e. completion) must be at least 8. (Note that both these lower bounds are extremely loose.) Therefore, training this network to completion at stable step size would require at least 8/22/246188 = 9, 876 iterations (which is the same number iterations as training to 37.1%% % accuracy at stable step size). Meanwhile, the larger step size of r) = 0.16 in just 1782 iterations. Therefore, training to completion at a stable step size is suboptimal by a factor of least 9, 876/177777
find that the network can be trained to completion at
5.54.
ResNet. For this network, gradient flow reaches 43.2% accuracy at time 70, and the maximum sharpness up through this point on the gradient flow trajectory is 760.6. Therefore, the largest sta- ble step size is 2/ 2/760.6 = 0.0026, and training to 43.2% accuracy at this step size would require 70/0.0026 26, 923 iterations. Meanwhile, we empirically observe that the network can also be trained to 43.2% accuracy at the larger step size of = 2.0 in just 99 iterations. Therefore, when training this network 43.2% accuracy, stable step sizes are suboptimal by factor of at least 26, 923/99- = 271.9. This is tight lower bound on the suboptimality of training to 43.2% accuracy a stable step size. For the loose lower bound on the suboptimality of training to completion at sta- ble step size, note that (by similar reasoning as the VGG-without-BN above), training to completion at a stable step size must require at least 26, 923 iterations. Meanwhile, we find that the network can be trained to completion at the larger step size of 77 2.0 in just 807 iterations. Therefore, training to completion at a stable step size is suboptimal by a factor at least 26, 923
023/807~ 8807 33.3.
66

---

Published as a conference paper at ICLR 2021
sharpness
accuracy
1.00
2.0
train
1500
0.75
1.0
1000
0.50
0.25
15.0
Figure 89: We train a VGG with BN to completion using gradient flow KCutae-Kutta). Observe that the sharpness rises dramatically from 6.38 initialization peak of 2227.59.
sharpness (iteration)
accuracy
1.0
2.0
200
sharpreess
0.6-
9- 2/25
200
2- 2/100
0.5-
0.2-
2/200
500
1000
1500
250
750
1000
1250
1500
1000
1500
Itcration
Itcration
from gradient
sharpness (time)
accuracy (Iteration)
1.0
docance
10
Sharpnees
150
0.6
200
0.4
0.0
15
1000
iteratio
Figure 90: We train a VGG with BN to completion using gradient descent at different step sizes (see legend in the top right pane). Top left: plot the train loss, with vertical dotted line marking the iteration where the sharpness first crosses 2/1). Top center: we plot the evolution of the sharpness, with horizontal dashed line (of the appropriate color) marking the value 2/1j. Top right: we plot the train accuracy. Bottom left: for the initial phase of training, we monitor the distance between the gradient descent iterate at iteration t/17, and the gradient flow solution time f, with vertical dotted line (of the appropriate color) marking the time when the sharpness crosses 2/n. Observe that the distance between gradient descent and gradient flow almost zero before this instant, but starts to rise shortly afterwards. Bottom center: plot the sharpness by time (= iteration xT) rather than iteration. The black dots are the sharpness the gradient flow trajectory. which shoots up immediately. Bottom right: we plot the test accuracy.
67

---

Published as a conference paper at ICLR 2021
train loss
sharpness
accuracy
2000
0.35
sharpeaas
accumacy
1000
0.25
0.20
Umz
time
3mc
Figure 91: We train VGG without BN to 37.1% accuracy using gradient flow (Runge-Kutta), Observe that the sharpness rises dramatically from 0.64 at initialization to 2461.78 78 at 37.1% % accu- racy. We train this network only partway because training this network to completion would be too computationally expensive: Runge- runs very slowly when the sharpness high (hecause is forced to take small steps) and for this network the sharpness is extremely high when the train accuracy is only 37% (which means that there a long way to go).
train loss lIteration)
sharpness (iteration)
train accuracy (iteration)
3.0
300
Sharpnesse
200
0.6
7-
2,25
1.0
2.50
0.2
2/200
2000
4000
8000
RD00
2000
4000
6DCD
8000
2000
4000
6000
RD00
Beraiion
Reration
leration
distance from gradient flow
sharpness (time)
test accuracy (iteration)
10.0
1.0
300
Sharpnoss
200
0.6
2.5-
0.2
2
50
100
150
2000
4000
6000
8000
tmm
time
karation
Figure 92: We train a VGG without BN to completion using gradient descent at range of step sizes (see legend in the top right pane). Refer to the Figure 900 caption for more information.
train
sharpness
accuracy
0.4
train tast
Sharpness
5.0
sccuresy
0.3
8:0
1.8
200
1.5
0.1
time
Figure 93: We train ResNet to 43.2% accuracy using gradient flow (Cunge-Kutta)) Observe that the sharpness rises dramatically from 1.07 at initialization to 760.63 at 43.2% accuracy. We train this network only partway because training this network to completion would be too computationally expensive: Kunge-Kutta small steps) and for this network the sharpness is extremely high when the train accuracy only 42% (which means that there is long way to go).
very slowly when the sharpness is high (because it is forced to take
68

---

Published as a conference paper at ICLR 2021
train loss (iteration)
sharpness (iteration)
train accuracy (iteration)
100
1.0
2.0
2/10
0.2
2140
20000
400n
6000
2000
6000
6000
2000
4000
6000
teration
Reration
Itenitin
gradient flow
sharpness (time)
test acouracy (Iteration)
100
1.0
dscorcc
Sharpnes
0.6
0.2
10
20
30
40
100
200
300
400
500
2000
4300
6000
time
time
Ieratio
Figure 94: We train ResNet to completion using gradient descent at a range of step sizes (see legend in the top right pane). Refer to the Figure 90) caption for more information.
69

---

Published as a conference paper at ICLR 2021
N EXPERIMENTS: MOMENTUM
This appendix contains systematic experiments for gradient descent with Polyak momentum and Nesterov momentum. Our aim is demonstrate that the sharpness rises until reaching the maxi- mum stable sharpness (MSS) given by equation and then either plateaus just above that value, or oscillates around that value.
We experiment on 5k-sized subset of CIFAR-IC 10, using four architectures: a tanh fully-connected network (section N.I), I. ReLU fully-connected network (section N.2)). a tanh convolutional network (section N. and a ReLU convolutional network (section N.4). For each of these four architectures, we experiment with both the square loss (for classification) and cross-entropy loss. For each archi- tecture and loss function, we experiment with both Polyak momentum at / 0.9, and gradient de- scent with Nesterov momentum at 0.9. We run gradient descent at range of several step sizes which were chosen by hand so that the MSS's s are approximately spaced evenly. Note that for Polyak momentum with step size 17 and momentum parameter 0.9, the MSS is
2+
3.8 For Nets-
terov momentum with step size y and momentum parameter 3 0.9, the MSS is 2+29 " We gradient descent until reaching 99% accuracy.
m1+223 R 1.35714
We find that the sharpness rises until reaching the maximum stable sharpness (MSS) given by equa- tion and then either plateaus just above that value, or oscillates around that value. Sometimes these oscillations are rapid (e.g. Figure 96), sometimes they are a bit slower (e.g. Figure [03), and sometimes they slow (e.g. Figure T09p.
70

---

Published as a conference paper at ICLR 2021
N.1 FULLY-CONNECTED TANH NETWORK
In the leftmost plot, the vertical dotted line marks the iteration where the sharpness first crosses the MSS. (Note that unlike vanilla gradient descent, momentum gradient descent sometimes cause the train loss to increase even when the algorithm is stable (Goh) 2017).) In the middle plot, the horizontal dashed line marks the MSS.
N.1.1 SQUARE LOSS
(iteration)
sharpness (iteration]
accuracy (iteration)
000
1.0
0.8
train 1005
500
7 0.04
0.4
0.02
7 0.013333
100
0.2
0.0
0.0
iteration
Feer Peraaton
Figure 95: Gradient descent with Polyak momentum, 3 = 0.9.
train (Iteration)
sharpness (Iteration)
accuracy (Iteration)
1.0
train l055
Porpnez
300
0.6
200
0.0065
100
0.2
0.00325
500
1000
2500
2000
2500
500
1000
1500
2000
2500
1000
1500
2000
2500
Figure 96: Gradient descent with Nesterov momentum, 8 = 0.9.
N.1.2 CROSS-ENTROPY LOSS
train loss (iteration)
sharpness (/teration)
train accuracy (iteration)
100
1an 1955
0.6
1.0
7
0,0606
0.2
3 0.04
50
1000
150
50
100
150
50
100
150
iteration
iteration
Iteration
Figure 97: Gradient descent with Polyak momentum, 3 0.9.
train lass (lteration)
sharpness (Iteration)
train accuracy (Iteration)
1.00
0.8
train 1025
3.5
sharpnice
3.0
0.4
0.0216
0.0
100
200
300
400
100
200
300
400
100
200
300
400
Iteration
Iteration
Iteration
Figure 98: Gradient descent with Nesterov momentum, 8 0.9.
71

---

Published as a conference paper at ICLR 2021
N.2 FULLY-CONNECTED RELU NETWORK
In the leftmost plot, the vertical dotted line marks the iteration where the sharpness first crosses the MSS. (Note that unlike vanilla gradient descent, momentum gradient descent sometimes cause the train loss to increase even when the algorithm is stable (Goh) 2017).
In the middle plot, the horizontal dashed line marks the MSS.
N.2.1
SQUARE LOSS
sharpness (iteration)
accuracy
1.0
100
198
sharprics
0.6
train
0.2
50
044
6 0,0444
0.0666
0.2
0.0333
300
150
150
250
300
iterstion
Figure 99: Gradient descent with Polyak momentum, 3 = 0.9.
train loss (Iteration)
sharpness (iteration)
accurary (Iteration)
125
1.0
trah lo00
0.4
0.6
0.0433
0.0216
0.2
0.0108
600
200
400
200
400
600
Iteration
iteration
Ioeraioo
Figure 100: Gradient descent with Nesterov momentum, / = 0.9.
N.2.2 CROSS-ENTROPE LOSS
train loss (iteration)
sharpness [/teration]
train accuracy (iteration)
100
0.8
brah 1038
Sharpness
0.6
3.0
0,0666
0.2
7 0.05
50
1.00
150
200
50
100
150
200
30
100
330
200
Iteration
iteration
Iteratio
Figure 101: Gradient descent with Polyak momentum, 0.9.
train los (lteratioo)
sharpness (Iteration)
train accuracy (Iteration)
100
0.8
mm
unn 1025
3.5
0.065
3.0
0.4
25
0.0216
0.0
0.0
100
200
300
50
100
150
200
250
300
100
200
300
Reration
Zeration
Iteraion
Figure 102: Gradient descent with Nesterov momentum, 9 0.9.
72

---

Published as a conference paper at ICLR 2021
N.3 CONVOLUTIONAL TANH NETWORK
In the leftmost plot, the vertical dotted line marks the iteration where the sharpness first crosses the MSS. (Note that unlike vanilla gradient descent, momentum gradient descent sometimes cause the train loss to increase even when the algorithm is stable (Goh) 2017).
In the middle plot, the horizontal dashed line marks the MSS.
N.3.1 SQUARE LOSS
(iveration)
sharpness (iteration)
accuracy (iteration)
0.6
125
1.0
1995
sharpece
0.6
train
0.4
0.0571
n-
0.2
0.0307
400
800
1000
200
200
800
Figure 103: Gradient descent with Polyak momentum, / = 0.9.
train loss (Iteration)
sharpness (iteration)
accurary (Iteration)
1.0
trah lo5s
Sharpicc
0.6
0.0325
0.2
0.2
0.013
20
500
1000
1500
2000
500
10000
2000
500
1000
1500
2000
Iteration
Figure 104: Gradient descent with Nesterov momentum, / = 0.9.
N.3.2 CROSS-ENTROPY LOSS
train loss (iteration)
sharpness (iteration]
train accuracy (Iterati on)
0.8
brah 1238
Shargnass
0.6
3.0
7
0.0066
0.2
7 0.05
50
1.00
150
200
250
300
50
100
150
200
250
300
0
50
100
130
200
230
300
Itoration
Iteration
Iteration
Figure 105: Gradient descent with Polyak momentum, 0.9.
train loss (lteration)
sharpness (Iteration)
train accuracy (ieeratonn)
0.0
unn 1025
1.5
Sharpness
40
3.0
0.4
0
20
0.2
0.032
0.0
0.0
200
400
600
400
600
200
400
600
iteratin
Seration
Iteration
Figure 106: Gradient descent with Nesterov momentum, 8 0.9.
73

---

Published as a conference paper at ICLR 2021
N.4 CONVOLUTIONAL RELU NETWORK
In the leftmost plot, the vertical dotted line marks the iteration where the sharpness first crosses the MSS. (Note that unlike vanilla gradient descent, momentum gradient descent sometimes cause the train loss to increase even when the algorithm is stable (Goh) 2017).) In the middle plot, the horizontal dashed line marks the MSS.
N.4.1
SQUARE LOSS
sharpness (iteration]
accuracy (iteration)
1000
1.0
0.8
0.4
0.0
500
0.4
0.0133
250
0.00666
0.2
0.004
0.0
0.0
2000
4000
2000
6000
6000
iberatinn
Reration
Figure 107: Gradient descent with Polyak momentum, B = 0.9.
(teeaation)
sharpness (Iteration)
accuracy (Iteration)
1.0
1000
train lo55
0.4
0.6
n=
500
00216
0.2
0,001083
5000
20000
15000
20000
25000
5000
15000 20000
25:00
5000
10000 15000
20000 25000
Figure 108: Gradient descent with Nesterov momentum, / = 0.9.
N.4.2 CROSS-ENTROPY LOSS
train loss (iteration)
sharpness (iteration)
train accuracy (iteration)
1.0
1500
195 1055
0.6
0.04
1.0
0
500
0.2
0 0.00266
1000
2000 Iteration
3000
4000
1000
2000
3000
4000
1000
2000
3000
4000
iteration
iterat
Figure 109: Gradient descent with Polyak momentum, 0.9.
train loss (Iteration)
sharpness (|teration)
train accuracy (Iteration)
1500
0.8
train 1025
3.5
1000
0.013
3.0
0.4
500
0.0013
1 0.000866
0.0
2000
4000
6000
1000
2000
3000
4000
5000
6000
2000
4000
5000
Iteration
Iteratioo
Iteration
Figure 110: Gradient descent with Nesterov momentum, 8 0.9.
74

---

Published as a conference paper at ICLR 2021
O EXPERIMENTS: LEARNING RATE DROP
In this appendix, we run gradient descent until reaching the Edge of Stability, and then we cut the step size. We will see that the sharpness starts increasing as soon the step size is cut, and only stops increasing once gradient descent is back at the Edge of Stability (or training is finished), As consequence of this experiment, one can interpret the Edge Stability as a regime in which gradient descent is constantly "trying" to increase the sharpness beyond 2/1,, but is constantly being blocked from doing so. Our experiments focus on image classification on 5k-sized subset of CIFAR- 10. We study two architectures (a fully-connected tanh network and convolutional ReL. network) and two loss functions (squared loss and cross-entropy loss).
75

---

Published as a conference paper at ICLR 2021
0.1
FULLY-CONNECTED TANH NETWORK: SQUARE LOSS
Gradient descent with fixed step size
train loss (Iteration)
sharpness [iteration]
train accuracy (Iteration)
1.0
0.6
400
0.8
025
0.4
sharphezz
2/100
200
0.4
0= 2/20
2/300
100
0.2
6 2/400
5000
10000
15000
5000
10000
15000
5000
10000
15000
Iteration
Iteration
Ieeraion
Drop 7 from to iteration 2000
400
0.8
tra los5
0.4
Fargness
300
0.0
200
100
0.2
4000
2000
4000
2000
4000 Iteration
6000
Drop 1? from 2/200 to 2/300 at Iteration 6000
560
1.0
train 1005
300
0.6
200
0.4
0.2
0
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
teration
Beration
Iteration
Drop from 2/300
2/400 iteration 10000
0.5
500
1.0
460
0.8
train 095
300
200
100
0.2
2500
5000
7500
10000
12500
2500
5,000
75:00 10000
12500
2500 5000
7500 10000 12500
iteration
iteraioo
Figure 111: Top row: we train a network using gradient descent at a range of step sizes 17 (see the legend in the right pane). Bottom three rows: Once gradient descent has reached the Edge of Stability, we cut the step size. The black vertical dotted line marks the iteration where the step size cut. This iteration (consistent with the legend) to plot the train loss, sharpness, and train accuracy before and after the learning rate drop. In the middle sharpness plot, the two horizontal lines mark the maximum stable sharpness 2/n for both the old and new step size. Takeaway: once we decrease the step size, the sharpness immediately starts to increase until it reaches the maximum stable sharpness 2/n for the new step size 17 (or training finishes).
chosen by hand, but not cherry-picked. We different colors
76

---

Published as a conference paper at ICLR 2021
O.2
FULLY-CONNECTED TANH NETWORK: CROSS-ENTROPY LOSS
Gradient step size ?
train loss (iteration)
sharpness (iteration)
train accuracy (iteration)
1.0
2.0
1035
sharpnes
0.6
0- 2/10
1.0
0.4
2/50
7 2/80
0.0
0.0
1000
1500
2000
2500
1000
1500
2000
2500
1500
2000
2500
iteration
Reration
Beration
Drop from 2/20 2/50 at iteration 300
150
1.0
2.0
trann 1035
3.5
sharpness
100
1.0
50
0.4
0.5
0.0
200
4.00
600
800
1000
200
400
660
800
0
200
400
800
800
1000
Iteration
teration
Iberaaioo
Drop from
0.8
tran 1055
sharpness
0.2
1000
1500
1000
1500
500
1000
1500
iteration
Drap from 2/80 to 2/110 at iteration on 1200
150
1.0
2.0
train 1003
3.5
sharpeeus
1.0
0.6
2.0
0.4
50
0.5
0
0.0
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Itaration
Reratien
Ieeation
Drop 7 from 2/110 to 2/140 iteration 2000
150
1.0
0.8
train 1055
Sharpnics
100
3.0
0.4
0.5
0.2
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
500
1000
1500
2000
2500
3000
Iteration
Iteration
keration
Figure 112: Refer to the caption for Figure
|1I
77

---

Published as a conference paper at ICLR 2021
O.3
CONVOLUTIONAL RELU NETWORK: SQUARE LOSS
Gradient descent with fixed step size
train loss (Iteration)
sharpness (Iteration)
train accuracy (iteration)
1.0
0.6
0 2/200
600
0.8
sharpnezz
s8n
0.4
400
0.4
4 2700
0.2
0.2
20000 40000
60000
80000
100000
20000
40000
60000
80000
100000
20000
40000
60000 Iteration
80000 100000
leration
Iteration
Drop
10000
800
0.8
tra 1035
0.4
accurccy
0.0
0.3
0.2
200
0.2
10000
20000
30000
40000
50000
10000
20000
30000
40000
50000
10000
20000
30000
40000
50000
Drop from 2/450 to 2/700 at iteration 30000
1000
1.0
600
0.6
tran
400
0.4
0.0
20000
40000 lberation
60000
80000
20000
40000 Seration
60000
80000
20000
40000
60000
80000
Reration
Drop from 2/700
2/900 iteration 60000
1000
1.0
0.5
800
0.8
train 1025
0.4
0.3
0.4
0.2
200
0.2
20000
40000 60000 Iteration
nocoo 100000
20000 40000 80000 nоооn 100000
20000
40000
60000
80000
100000
iteration
iteration
Figure 113: Refer to the caption for Figure 111
78

---

Published as a conference paper at ICLR 2021
0.4
CONVOLUTIONAL RELU NETWORK:
LOSS
step size ?
train loss (iteration)
sharpness (iteration)
train accuracy (iteration)
2.0
1000
1035
sharpnizz
0.6
1.0
500
0.4
2/100
7- 2/400
0.0
10000
150000
20000
5000
10000
20000
0.0
5000
10000
15000
20000
Ieeratin
Reration
iteration
Drop from 2/100 to 2/400 at iter ation 1000
1500
1.0
2.0
train 1035
3.5
1000
0.6
500
0.4
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
0
2000
4000
6000
8000
18000
iteration
Iteration
Iteration
1)
Iteration
7500
0.8
tran 1055
0.6
3.0
0.5
0.2
2500
5000
10000
12500 15000
2500
7500 10000 12500 15000
2500
7500
20000
15000
Drop n from 2/700 to 01000 at iteration 7500
1500
1.0
2.0
train 1003
3.5
Sharpness
1000
0.6
500
0.4
0.0
5000
10000
15000
20000
5000
10000
25000
20000
5000
10000
15000
20000
Iteration
Reeaten
teration
Figure 114: Refer to the caption for Figure 111
79

---

Published as a conference paper at ICLR 2021
P OTHER EIGENVALUES
Throughout most of this paper, we have studied the evolution the maximum Hessian eigenvalue during gradient descent. In this appendix, we examine the evolution of the top six eigenvalues. While training the network from f we monitor the evolution of the top six Hessian eigenvalues. For each of cross-entropy loss and MSE loss, we train at four different step sizes. The results are shown in Figure [13| Observe that each of the top six eigenvalues rises and then plateaus. The precise details ditfer between MSE loss and cross-entropy loss. For MSE loss, each eigenvalue rises past 2/1,, and then plateaus just above that value. In contrast, for cross-entropy loss, some of the lesser eigenvalues plateau below 2/1)-
1.0
step size 2210
250
step kizw
sieo sire
150
step 2/120
cross-entropy
1:2
325
125
125
1.00
300
100
200
400
600
500
5:00
1003
1500
302
1000
1500
2000
500
1000
1500
2000
2006 3302
mean squared error
5.50
2teo size 27100
340
stop sibe 2/200
step sico 2/300
200p
size 24400
4.0
320
300
60
-200
40
100
1:00
sic
1990 1560 2000 Barsdan
2520
3000
5/0
1006
15:00
2000
10:31
2000
1000
1000
2000
5992
0.02
lar
Figure 115: Evolution of top six eigenvalues. For both cross-entropy loss (top) and MSE loss (bottom), we run gradient descent four different step sizes (columns), and plot the evolution of top six eigenvalues (different colors; see legend).
80


Reviews:
[{"review_id": "t_Jh7YExp5s", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThis paper presents an interesting observation for GD. That is, the sharpness of the learnt model in the final phase of the training (measured by the largest eigenvalue of the training loss Hessian) hovers right at the value 2/\\eta while the training loss. At the same time, the loss goes to unstable and non-monotonically decreasing. This pattern is consistent across architecture, activation functions, tasks, loss functions and BN. Comprehensive experiments are conducted to show this common observation. The paper is easy to follow.\n\nBesides the empirical results in the main body, authors give insightful discussions in Intro and related work section. Specifically, authors propose a novel guess, that GD eventually transitions to “Edge of stability”, where GD can finally succeed with non-small enough step size. Although I am not sure how GD can do this, the concept of “Edge of stability” is still attractive. \n\nI have two concerns for this work. \n1) Authors did not investigate why sharpness finally hover over 2/\\eta. Is it a trivial consequence followed by some relationship between the update rule of GD and the definition of sharpness, without any condition? Even if yes, we may further think about how to leverage it along the existing discussions in this paper. Hope to have authors' feedbacks on this later. \n2) Given people use SGD to train neural networks, discussions about the insight from the observation of GD to SGD will enhance the impact of this paper. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "khKWeJYqVFV", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work identifies a new empirical phenomenon in the training dynamics of deep nets: when trained with full-batch GD, the curvature of the train loss increases up to a critical value of 2/(step size), at which point it plateaus for the remainder of training.\nThis phenomenon is demonstrated robustly for networks trained with MSE loss, across various architectures and datasets, and a slightly weaker version of this holds for cross-entropy loss as well.\nThis work contributes to our understanding of deep network dynamics -- it is a precise and apparently robust phenomenon that was surprisingly not noticed before (perhaps because of the requirement of GD vs SGD). \nIn terms of impact: This work will be instructive for DL optimization theory, since it points out that certain assumptions which are usually made in theoretical works (e.g. step size << curvature) are far from true in practice -- moreover, it guides theory towards more realistic assumptions.\nIt may also have later impact in practice, by leading to a better understanding of the interaction between optimization algorithm, step size, and architecture. Thus I recommend acceptance.\n\nWeaknesses and desired clarifications:\n- It should be mentioned more prominently that these results are primarily for networks trained with MSE loss -- and that a similar but weaker phenomena holds for cross-entropy loss. \n- Why is the main example in Section 3 given for a non-standard network for CIFAR-10? A 2-layer MLP with ELU activation. Why not a standard network with standard activation? (VGG-11 or ResNet-18, etc).\n- The distinction between SGD and GD seems crucial for this phenomenon, so more discussion would be good. In particular, as noted in the related works, some papers using SGD claim an opposite effect. This is especially important to clarify since SGD is most often used in practice. If time allows, experiments with increasing batch size could shed light on the importance of GD vs SGD.\n- The Related Works is currently written as an account of what previous works do *not* do, as opposed to what they do. It would help contextualize this work to relate it to prior works which are consistent (or inconsistent) with this phenomena -- especially works studying the Hessian of deep nets. Some of the mechanisms proposed in prior works (eg Lewkowycz et al 2020 and works on deep linear networks) may also be helpful to understand the phenomena in this work.\n\n\n\nComments which do not affect the score:\n\n- I wonder if you have measured the 2nd eigenvalue during training as well? In particular, after the 1st eigenvalue has saturated at 2/eta, does the 2nd eigenvalue also \"progressively sharpen\" up to 2/eta ? (And so on for later eigenvals).\n- I am glad to see the experiments on deep linear networks, it suggests that it may be possible to theoretically understand this phenomenon in such simple settings. This would be a nice topic for future work.", "extra_scores": {}, "extra_text": {}}, {"review_id": "kaAKJ6gMfb", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis submission numerically shows that during exploring the neural network landscape,  GD flow keeps increasing the sharpness.  As a result, GD with a fixed learning rate will exhibit two phases during the dynamics.  Denote by $\\eta$ the fixed learning rate.  In the first phase, GD follows closely to the GD flow, and it finally converges to a region where the sharpness is roughly $2/\\eta$.  Then, it transits into the second phase during which the sharpness hovers right at or above $2/\\eta$. In the second phase, GD cannot increase the sharpness anymore due to the dynamical stability constraint. Thus, the authors name it the Edge of Stability phase.  What is interesting is that in the edge of stability phase, the loss is still decreasing steadily although not monotonically. \n\nPros:\n\nI enjoy reading this submission. It is clearly written and the numerical evaluation is also sufficient.  To my best of knowledge, the observation that the edge of stability happens during the whole late phase of GD dynamics is new.  It reveals a very complicated dynamical behavior of GD for training neural networks, which has not been systematically investigated before.   Thus, I think this submission made a very important and original contribution to the understanding of GD dynamics  in deep learning. \n\nCons:\n\nThe relationship with the previous study on the dynamical stability of (S)GD is not sufficient discussed. In my opinion, just saying ''previous works have argued that the stability properties of optimization algorithms could potentially serve as a form of implicit bias in deep learning'' is obviously not precise and enough. A large number of numerical results in [1,2] already showed that the edge of stability happens for the convergent solutions, which implies that the edge of stability must happen at least in the very late phase of GD dynamics.   The new finds of this submission are that the edge of stability actually holds for a large portion of GD dynamics, which is very unexpected.   The authors should explicitly mention that the edge of stability was already observed in these previous works. Giving the right credit to the right references does not harm the contribution of this submission.  Especially, the jargon \"Edge of stability\" was first used in [1], and the authors even did not mention it. \n\n\n[1] Giladi, Niv, et al. \"At Stability's Edge: How to Adjust Hyperparameters to Preserve Minima Selection in Asynchronous Training of Neural Networks?.\" arXiv preprint arXiv:1909.12340 (2019).\n\n[2] Wu, Lei, Chao Ma, and E. Weinan. \"How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective.\" Advances in Neural Information Processing Systems. 2018.", "extra_scores": {}, "extra_text": {}}, {"review_id": "OpSGDsEc9dT", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of results: This empirical paper finds that deep neural network \"sharpness\" (as measured by the top eigenvalue of the Hessian) tends to saturate at or hover just above the value 2/\\eta, where \\eta is the step size in gradient descent (GD), during the course of optimization. This is accompanied by non-monotonicity in the loss. (Results also hold for gradient descent with momentum.) This phenomena occurs for full-batch GD, a variety of tasks, and two different loss functions (square loss and cross-entropy, although in the latter case late-time dynamics of the sharpness is different, with the sharpness decreasing). The result also holds across a variety of architectures (VGG-11, with and without batch norm, convolutional networks and fully-connected networks with different nonlinearities, a deep linear network, a Transformer model -- although I comment on the architecture dependence below.) The authors refer to this phenomenon of sharpness hovering at or above the 2/\\eta bound as optimization on the \"edge of stability.\" The authors posit that this observation goes against our current understanding of optimization in deep learning: (i) theoretical work may use assumptions such as monotonic descent or L-smoothness, which is in violation of the \"edge of stability\" regime where the loss behavior is not monotonic; (ii) relatedly, it corresponds to a regime of instability in a quadratic Taylor approximation, so this is not a good assumption either. The empirical observations do not appear to straightway carry over to SGD (although there is some discussion of similarities in the appendix) -- the authors leave this is as an open problem.\n\nQuality and clarity: The work is of good quality. The experiments that are presented clearly demonstrate the described phenomenon, attempt to sample from a variety of problems (ranging from a simple 1D regression task to training VGG-11 on CIFAR-10 in Fig.1), and are well-explained. \n\nOriginality: The observed phenomena (of sharpness progressively increasing and then hovering at or near 2/\\eta through much of GD dynamics) appears new to the best of my knowledge. (However, prior papers have pointed out that GD optimization can stably proceed for learning rates above 2/\\eta and are accompanied with nonmonotonic loss: one example is Lewkowycz, et al. (2020).)\n\nSignificance: I think the observations are somewhat important, although perhaps not as much as the authors seem to stress in writing (e.g. in the abstract, \"Our results ... shed significant light on the dynamics of gradient descent with a fixed step size ...\" -- this paper focuses only on empirical observations rather than explanation of the mechanism behind the stability). The results are for full-batch GD which somewhat limits the applicability of the results to practice where SGD is more common. Nonetheless, I do think understanding how optimization is stabilized in this regime is an interesting problem.\n\nOther comments: I think that a key point not fully understood (at least, as expressed in the writing) or investigated in the experiments is the role of network width (i.e. I'm skeptical that the observations will be \"consistent across architectures\" as expressed in the opening paragraph in the Introduction if wider networks are also investigated empirically). This is a main reason for not giving the paper a higher score. In Appendix D, the authors try to reconcile the results with existing theory on infinite-width limits. Fig. 11(c) shows that across networks of varying width but trained at the same learning rate, wider networks end up with smaller values of sharpness at the end of training (here, say that all experiments are stopped at the same value of the training loss). Hence if a narrower network saturates the 2/\\eta bound at late times in gradient descent, the wider networks will fall below this value. (This is, of course, consistent with no evolution in the curvature in the infinite width limit.) That is to say, whether or not the \"edge of stability\" regime is reached depends quite strongly on how wide the networks are. \n\nWhile some realistic networks are explored (e.g. VGG-11, *narrow* Transformers), which is a positive point, all of the remaining networks used in the experiments (as far as I can gather) are on the narrow side (e.g. layer widths of order 100 or 200). I think a shortcoming of the paper is that the strong width (i.e. architecture) dependence of the phenomenon is not fully appreciated or discussed by the authors (e.g. for instance, by discussing in the main text that it only sets in for narrow networks) or investigated empirically. (Alternatively, noting for readers that all the networks chosen are rather narrow, out of transparency.)  I note that the authors do mention in Appendix D: \"...one might hypothesize that progressive sharpening might attenuate as networks (with NTK parameterization) become increasingly wide.\" However, I don't think that NTK parameterization is necessary for this to be true. In short, I think the width dependence of the phenomenon is an important factor that affects the significance and applicability of the observations and could have been treated with greater transparency (with additional experiments and additions to the main text and abstract).\n\nA comment on relation to prior work: the authors write that Lewkowycz, et al. (2020) imply that \"actual progress would occur in regions where the sharpness remains strictly less than 2/\\eta. Our experiments demonstrate otherwise.\" I don't believe this is a conclusion of that paper (progress happens when sharpness is above 2/\\eta). (Note also that the paper tends to study wide networks.) Could the authors elaborate on what they mean here?", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper demonstrates that Gradient Descents generally operates in a regime where the spectral norm of the Hessian is as large as possible given the learning rate. 

The paper presents a very thorough empirical demonstration of the central claim, which was appreciated by the reviewers.

A central issue to me in accepting the work was its novelty. Prior work has shown very closely related effects for SGD. The reviewers appreciated in discussions the novelty of the precise claim about the spectral norm hovering at around $\frac{2}{\eta}$. R4 and R2 also raised the issue that the related work discussion is not sufficient. Please make sure that you discuss very carefully related work in the paper, including a more detailed discussion in the Introduction.

The two key issues raised by R3, who voted for rejection, were that (1) the work studies Gradient Descent (rather than SGD), and (2) lack of theory. I agree with these concerns. Perhaps the Authors should address (1) by citing more carefully prior work that shows that a similar phenomenon does seem to happen in training with SGD. As for (2), I agree here with R1,R2 and R4 that empirical evaluation is a key strength of the paper. 

Based on the above, it is my pleasure to recommend the acceptance of the paper. Thank you for submitting your work to ICLR, and please make sure you address all remarks of the reviewers in the camera-ready version.
Decision: Accept (Poster)