Santurkar et 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 $\theta$, and a distance $\alpha$, the effective smoothness of $f$ at parameter $\theta$ and distance $\alpha$ is defined in Santurkar et al., 2018 as
$$\sup_{\gamma \in [0, \alpha]} \frac{ \| \nabla f(\theta) - \nabla f(\theta - \gamma \nabla f(\theta)) \|_2 }{\| \gamma \nabla f(\theta) \|_2}$$

where the $\sup$ can be numerically approximated by  evaluating the given ratio at several values $\gamma$ spaced uniformly between $0$ and $\alpha$.

In Figure <ref>, 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 a distance $\alpha = \eta$.
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/\eta$ and then remains there.
Thus, at 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-BN network, confirming that BN does accelerate training.)

On a 5,000-size subset of CIFAR-10, we train a ReLU CNN both with BN (top row) and without BN (bottom row) at the same grid of step sizes.  We plot  the sharpness/smoothness (center column) as well as the effective smoothness [Santurkar et al., 2018] (right column). Observe that for both networks,  the effective smoothness hovers around zero initially, and jumps up to $2/\eta$ once gradient descent enters the Edge of Stability.

Note that this finding is actually consistent with Figure 4(c) in Santurkar et 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 $\eta = 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 Santurkar et al., 2018 in support of the argument that batch normalization improves the effective smoothness during training is their Figure 9(c).
This figure shows that a deep linear network (DLN) trained without BN has a much larger (i.e. worse) effective smoothness during training than a DLN trained with BN.
However, for this figure, the distance $\alpha$ used to compute effective smoothness was larger than the training step size $\eta$ by a factor of $30$.
The effective smoothness at distances larger than the step size does not affect  training.
We have verified that when effective smoothness is computed at a distance equal to the training step size (i.e. $\alpha = \eta$), the effective smoothness for the DLN with BN and for the DLN without BN both hover right at $2/\eta$.

Specifically, in Figure  <ref> and Figure  <ref>, we train a DLN both with and without BN (respectively), and we measure the effective smoothness at a distance $\alpha = 30 \eta$, as done in Figure 9(c) of  Santurkar et al., 2018.
We use the same experimental setup and the same step size of $\eta =$ 1e-6 as 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/\eta$ (marked by the horizontal black line), whereas when training the non-BN network, the effective smoothness is much larger.
This is consistent with Figure 9(c) in Santurkar et al., 2018.
However, in Figure  <ref> and <ref>, we measure the effective smoothness at the actual step size $\alpha = \eta$.
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/\eta$.
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 experiment possibly explains why the batch-normalized network permits training with larger step sizes.)

When training a deep linear network without BN, we measure effective smoothness at a distance $\alpha = 30 \eta$ that is 30x larger than the step size.

When training a deep linear network with BN, we measure effective smoothness at a distance $\alpha = 30 \eta$ that is 30x larger than the step size.

When training a deep linear network without BN, we measure effective smoothness at a distance $\alpha = \eta$ that is equal to the step size.

When training a deep linear network with BN, we measure effective smoothness at a distance $\alpha = \eta$ that is equal to the step size.

§ ADDITIONAL TASKS

So far, we have verified our findings on image classification.
In this appendix, we verify our findings on three additional tasks: training a Transformer on the WikiText-2 language modeling dataset (<ref>), training a one-hidden-layer network on a one-dimensional toy regression task (<ref>), and training a deep linear network (<ref>).

§.§ Transformer on WikiText-2

We consider the problem of training a Transformer on the WikText-2 word-level language modeling dataset [Merity et al., 2016].
See <ref> for full experimental details.
In Figure <ref>, we train using gradient flow (only partially, not to completion).
Observe that the sharpness continually rises.
In Figure <ref>, we train using gradient descent at a range of step sizes.
Consistent with our general findings, for each step size $\eta$, we observe that the sharpness rises to $2/\eta$ 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.

Training a Transformer using gradient flow.

Training a 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: we plot the train loss curves, with a vertical dotted line marking the iteration where the sharpness first crosses $2/\eta$.  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/\eta$.  Observe that for each step size, the sharpness rises to $2/\eta$ 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 the gradient flow trajectory, with a vertical dotted line marking the gradient descent iteration where the sharpness first crosses $2/\eta$.  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, we plot the evolution of the sharpness by “time” = $\text{iteration} \times \eta$ rather than iteration.  The black dots are the sharpness during gradient flow.

§.§ One-dimensional toy regression task

Task Our toy regression problem is to approximate  a Chebyshev polynomial using a neural network.
To generate a  toy dataset, we take 20 points spaced uniformly on the interval $[-1, 1]$, and we label them noiselessly using the Chebyshev polynomial of some degree $k$.
Note that the Chebyshev polynomial of degree $k$ is a polynomial with $k$ zeros that maps the domain $[-1, 1]$ to the range $[-1, 1]$.
Figure <ref> shows the Chebyshev datasets for degree 3, 4, and 5.

The datasets for the toy one-dimensional regression task.

Network For the network, we use a tanh architecture with one hidden layer of $h=100$ units, initialized using Xavier initialization.
We train using the MSE loss until the loss reaches 0.05.

In Figure <ref>, we fit the Chebyshev degree 3, 4, and 5  datasets using gradient flow.
Empirically, the higher the degree, the more the sharpness rises during the course of training: on the degree 3 polynomial, the sharpness rises by a factor of 1.2; on the degree 4 polynomial, by a factor of 3.2; and on the degree 5 polynomial, by a factor of 63.5.

In Figure <ref> and Figure <ref>, 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 as elsewhere in the paper.
The only difference is that for the degree 5 dataset, after the sharpness hits $2/\eta$, the training loss first undergoes a temporary period of non-monotonicity in which no progress is made, and then it decreases monotonically  until training is finished (in contrast to our other experiments where we observe that training loss behaves non-monotonically at the same time as it is consistently decreasing).

Fitting Chebyshev polynomials using gradient flow.  We use gradient flow to fit a one-hidden-layer tanh network to the Chebyshev polynomials of degrees 3, 4, and 5.  We observe that the sharpness rises more when fitting a higher-degree polynomial, likely because the dataset is more complex.

Chebyshev degree 4 (gradient descent). We fit the Chebyshev polynomial of degree 4 using gradient descent at a range of step sizes (see legend).

Chebyshev degree 5 (gradient descent). We fit the Chebyshev polynomial of degree 5 using gradient descent at a range of step sizes (see legend).

§.§ Deep linear network

The task is to map  $n$ inputs $\mathbf{x}_1, \hdots, \mathbf{x}_n \subseteq \mathbb{R}^d$ to $n$ targets $\mathbf{y}_1, \hdots, \mathbf{y}_n \subseteq \mathbb{R}^d$ using a function $f: \mathbb{R}^d \to \mathbb{R}^d$.
Error is measured using the square loss, i.e. the objective is  $\frac{1}{n} \sum_{i=1}^n \| f(\mathbf{x}_i) - \mathbf{y}_i  \|^2_2 $.
Let $\mathbf{X} \in \mathbb{R}^{n \times d}$ be the vertical stack of the inputs, and let  $\mathbf{Y} \in \mathbb{R}^{n \times d}$ be the vertical stack of the targets.
We first generate $\mathbf{X}$ as a random whitened matrix (i.e. $\frac{1}{n} \mathbf{X}^T \mathbf{X}=\mathbf{I}$).
To generate $\mathbf{X}$ as a random whitened matrix, we sample a $n \times d$ matrix of standard Gaussians, and then set $\mathbf{X}$ to be $\sqrt{n}$ times the Q factor in the QR factorization of that matrix.
We then generate $\mathbf{Y}$ via $\mathbf{Y} = \mathbf{X} \mathbf{A}^T$, where $\mathbf{A} \in \mathbb{R}^{d \times d}$ is a random matrix whose entries are sampled i.i.d from the standard normal distribution.

We use $n=50$ datapoints with a dimension of $d=50$.

The function $f: \mathbb{R}^d \to \mathbb{R}^d$ is implemented as a $L$-layer deep linear network: $f(\mathbf{x}) = \mathbf{W}_L \hdots \mathbf{W}_2 \mathbf{W}_1 \mathbf{x}$, with $\mathbf{W}_{\ell} \in \mathbb{R}^{d \times d}$.
We initialize all layers of the deep linear network using Xavier initialization: all entries of each $\mathbf{W}_{\ell}$ are drawn i.i.d from $\mathcal{N}(0, \frac{1}{d})$.

We use a network with $L=20$ layers.

In Figure <ref>, we train the network using gradient flow.
(Since it is unclear whether the network can be trained to zero loss, and how long this would take, we arbitrarily chose to stop training at time 100.)
In Figure <ref>, we train the network using gradient descent at a range of step sizes.
We observe mostly the same Edge of Stability behavior as elsewhere in the paper.
The only difference is that in Figure <ref>, the train loss does not really behave non-monotonically — for each step size $\eta$, there is a brief blip at some point, but otherwise, the train loss decreases monotonically.

Training a deep linear network using gradient flow.  We use gradient flow to train a 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.

Training a deep linear network using gradient descent.  We train a deep linear network using gradient descent at a range of step sizes (see legend).  (Note that in the top left pane, the blue, orange and green dotted lines are directly on top of one another.)

§ EXPERIMENTS: STANDARD ARCHITECTURES ON CIFAR-10

In this appendix, we demonstrate that our findings hold for three standard architectures on the standard dataset CIFAR-10.
The three architectures are: a VGG with batch normalization (Figures <ref>-<ref>), a VGG without batch normalization (Figures <ref>-<ref>), and a ResNet with batch normalization (Figures <ref> -<ref>).
See <ref> for full experimental details.

For each of these three architectures, we confirm our main points: (1) so long as the sharpness is less than $2/\eta$, the sharpness tends to increase; and (2) if the sharpness reaches $2/\eta$, 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/\eta$.
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/\eta$.

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 sharpening, 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 sharpness would rise “if it didn't have to worry about” instability caused by nonzero step sizes.
In Figure  <ref>, 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 to a peak  value of  2227.6.
For the other two architectures, progressive sharpening occurs to such a degree that it is not computationally 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-Kutta step size must be made small, so Runge-Kutta requires very many iterations.)
Therefore, we instead train these two networks only partially.
In Figure <ref>, 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 <ref>, 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 <ref> 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 <ref> that if $\lambda_{\max}$ is the maximum sharpness along the gradient flow trajectory, then any stable step size must be less than $2/\lambda_{\max}$.
Therefore, for these three architectures, because progressive sharpening occurs to a dramatic degree (i.e. $\lambda_{\max}$ 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/0.000897 = 16622$ iterations.
Meanwhile, we empirically observe that the network can also be trained to completion at the much larger step size of $\eta = 0.16$ in just 329 iterations.
Therefore, using a stable step size is suboptimal by a factor of at least $16622/329=50.5$.

For the other two architectures, since we are unable to train to completion using gradient flow, we are unable to obtain a 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 both: (1) a tight lower bound on the suboptimality of training partially at a 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 maximum sharpness up through this point on the gradient flow trajectory is 2461.8.
Therefore, the largest stable step size is $2/2461.8 = 0.00081$, and training to 37.1% accuracy at this step size would require $8/0.00081=9,876$ iterations.
Meanwhile, we empirically observe that the network can also be trained to 37.1% accuracy at the larger step size of $\eta = 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 a tight lower bound on the suboptimality of training to 37.1% accuracy at a stable step size.

To obtain a loose lower bound on the suboptimality of training the VGG without BN  to completion at a 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 of these lower bounds are extremely loose.)
Therefore, training this network to completion at a stable step size would require at least $8/(2/2461.8) = 9,876$ iterations (which is the same number of iterations as training to 37.1% accuracy at a stable step size).
Meanwhile, we find that the network can be trained to completion at the larger step size of $\eta = 0.16$ in just 1782 iterations.
Therefore, training to completion at a stable step size is suboptimal by a factor of at least $9,876/1782 = 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 stable step size is $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 $\eta = 2.0$ in just 99 iterations.
Therefore, when training this network to 43.2% accuracy, stable step sizes are suboptimal by a factor of at least $26,923/99=271.9$.
This is a tight lower bound on the suboptimality of training to 43.2% accuracy at a stable step size.
For the loose lower bound on the suboptimality of training to completion at a stable 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 $\eta = 2.0$ in  just 807 iterations.
Therefore, training to completion at a stable step size is suboptimal by a factor of at least $26,923/807 = 33.3$.

We train a VGG with BN to completion using gradient flow (Runge-Kutta). Observe that the sharpness rises dramatically  from 6.38 at initialization to a peak of 2227.59.

We train a VGG with BN to completion using gradient descent at different step sizes (see legend in the top right pane). Top left: we plot the train loss, with a vertical dotted line marking the iteration where the sharpness first crosses $2/\eta$.   Top center: we plot the evolution of the sharpness, with a horizontal dashed line (of the appropriate color) marking the value $2/\eta$.  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/\eta$, and the gradient flow solution at time $t$, with a vertical dotted line (of the appropriate color) marking the time when the sharpness crosses $2/\eta$. Observe that the distance between gradient descent and gradient flow is almost zero before this instant, but starts to rise shortly afterwards. Bottom center: we plot the sharpness by time (= iteration $\times \eta$) rather than iteration. The black dots are the sharpness of the gradient flow trajectory, which shoots up immediately.  Bottom right: we plot the test accuracy.

We train a 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 at 37.1% accuracy. We train this network only partway because training this network to completion would be too computationally expensive: Runge-Kutta runs very slowly when the sharpness is high (because it 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 is a long way to go).

We train a VGG without BN to completion using gradient descent  at a range of step sizes (see legend in the top right pane).  Refer to the Figure <ref> caption for more information.

We train a ResNet to 43.2% accuracy using gradient flow (Runge-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: Runge-Kutta runs very slowly when the sharpness is high (because it is forced to take small steps) and for this network the sharpness is extremely high when the train accuracy is only 42% (which means that there is a long way to go).

We train a ResNet to completion using gradient descent  at a range of step sizes (see legend in the top right pane).  Refer to the Figure <ref> caption for more information.

§ EXPERIMENTS: MOMENTUM

This appendix contains systematic experiments for gradient descent with Polyak momentum and Nesterov momentum.
Our aim is to demonstrate that the sharpness rises until reaching the maximum stable sharpness (MSS) given by (<ref>), and then either plateaus just above that value, or oscillates around that value.

We experiment on a 5k-sized subset of CIFAR-10, using four architectures: a tanh fully-connected network (section <ref>), a ReLU fully-connected network (section <ref>), a tanh convolutional network (section <ref>), and a ReLU convolutional network (section <ref>).
For each of these four architectures, we experiment with both the square loss (for classification) and cross-entropy loss.
For each architecture and loss function, we experiment with both Polyak momentum at $\beta = 0.9$, and  gradient descent with Nesterov momentum at $\beta = 0.9$.
We run gradient descent at a range of several step sizes which were chosen by hand so that the MSS's are approximately spaced evenly.
Note that for Polyak momentum with step size $\eta$ and momentum parameter $\beta = 0.9$, the MSS is $\frac{2 + 2 \beta}{\eta} = \frac{3.8}{\eta}$.
For Netsterov momentum with step size $\eta$ and momentum parameter $\beta = 0.9$, the MSS is $\frac{2 + 2 \beta}{\eta (1 + 2 \beta)} \approx \frac{1.35714}{\eta}$.
We run gradient descent until reaching 99% accuracy.

We find that the sharpness rises until reaching the maximum stable sharpness (MSS) given by (<ref>), and then either plateaus just above that value, or oscillates around that value.
Sometimes these oscillations are rapid (e.g. Figure <ref>), sometimes they are a bit slower (e.g. Figure <ref>), and sometimes they are slow (e.g. Figure <ref>).

§.§ 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 can 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.

§.§.§ Square loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§.§ Cross-entropy loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§ 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 can 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.

§.§.§ Square loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§.§ Cross-entropy loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§ 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 can 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.

§.§.§ Square loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§.§ Cross-entropy loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§ 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 can 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.

§.§.§ Square loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§.§.§ Cross-entropy loss

Gradient descent with Polyak momentum, $\beta = 0.9$.

Gradient descent with Nesterov momentum, $\beta = 0.9$.

§ 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 as the step size is cut, and only stops increasing once gradient descent is back at the Edge of Stability (or training is finished).
As a consequence of this experiment, one can interpret the Edge of Stability as a regime in which gradient descent is constantly “trying” to increase the sharpness beyond $2 / \eta$, but is constantly being blocked from doing so.
Our experiments focus on image classification on a 5k-sized subset of CIFAR-10.
We study two architectures (a fully-connected tanh network and a convolutional ReLU network) and two loss functions (squared loss and cross-entropy loss).

§.§ Fully-connected tanh network: square loss

Top row: we train a network using gradient descent at a range of step sizes $\eta$ (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 is cut.  This iteration was chosen by hand, but was not cherry-picked. We use different colors (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 / \eta$ 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 / \eta$ for the new step size $\eta$ (or training finishes).

§.§ Fully-connected tanh network: cross-entropy loss

Refer to the caption for Figure <ref>.

§.§ Convolutional ReLU network: square loss

Refer to the caption for Figure <ref>.

§.§ Convolutional ReLU network: cross-entropy loss

Refer to the caption for Figure <ref>.

§ OTHER EIGENVALUES

Throughout most of this paper, we have studied the evolution of the maximum Hessian eigenvalue during gradient descent.
In this appendix, we examine the evolution of the top six eigenvalues.
While training the network from <ref>, 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 <ref>.
Observe that each of the top six eigenvalues rises and then plateaus.
The precise details differ between MSE loss  and cross-entropy loss.
For MSE loss, each eigenvalue rises past $2/\eta$, and then plateaus just above that value.
In contrast, for cross-entropy loss, some of the lesser eigenvalues plateau below $2/\eta$.

Evolution of top six eigenvalues.  For both cross-entropy loss (top) and MSE loss (bottom), we run gradient descent at four different step sizes (columns), and plot the evolution of the top six eigenvalues (different colors; see legend).

§ RELATION TO JASTRZĘBSKI ET AL., 2020

We now elaborate on the relationship between our work and Jastrzębski et al., 2020, which conjectured (but didn't observe) the Edge of Stability phenomenon for full-batch gradient descent.

Jastrzębski et al., 2020 aimed to explain the observation [Jastrzębski et al., 2019] that the  SGD learning rate and batch size hyperparameters strongly influence the local geometry along the optimization trajectory, with large learning rates and small batch sizes guiding SGD into regions of the loss landscape characterized by lower sharpness and lower variance in the stochastic gradients.
However, whereas we focus in this paper on the simple yet unrealistic setting of full-batch gradient descent,   Jastrzębski et al., 2020 ambitiously attempted to tackle the realistic yet complex setting of SGD.
Towards that end, in their Assumptions 1-4,  Jastrzębski et al., 2020 conjectured a simplified model for the evolution of certain quantities during SGD training of neural networks.

For the special case of full-batch gradient descent, their Assumption #4 amounts to an assertion that gradient descent will continually increase the sharpness so long as the sharpness remains beneath $2/\eta$, and their Assumption #3 amounts to an assertion that if gradient descent finds itself in a region with sharpness above $2/\eta$, it will subsequently decrease the sharpness.
Taken together, these two assumptions constitute a conjecture that gradient descent trains at the Edge of Stability.

However, Jastrzębski et al., 2020 did not train any neural networks using full-batch gradient descent, and therefore did not observe Edge of Stability (the sharpness hovering at or just above $2/\eta$).
Furthermore, neither Jastrzębski et al., 2020 nor the papers cited therein conducted any controlled experiments to verify their Assumptions 1-4.
Indeed, Jastrzębski et al., 2020 did not claim that Assumptions 1-4 are literally true — and Assumption 1, which asserts that all example-wise loss functions share the same minimum along the sharpest direction, is clearly not literally true.
(Assumption 1 is only required for SGD.)

The main focus of Jastrzębski et al., 2020 was stochastic gradient descent.
For SGD, Jastrzębski et al., 2020 essentially conjectured that SGD trains at the “edge of stability,” where the stability of SGD is determined by their Equation 1.
In the special case of full-batch training, their Equation 1 reduces to “sharpness < $2/\eta$”, a constraint which we have shown is tight throughout training.
However, in the more general case of SGD, their Equation 1 does not appear to be anywhere close to tight throughout training.
We attribute this to deficiencies in their stability model of SGD.
One issue is that their Assumptions 1 and 2 may not be satisfied in practice.
Another issue is that stability in expectation may not describe the typical behavior of a random variable (see <https://openreview.net/forum?id=jh-rTtvkGeM   noteId=KHKfRvNx0A>).

In summary, whereas we have restricted our study to the simple yet unrealistic algorithm of full-batch gradient descent — where we are able to make very precise statements about the behavior of the sharpness —  Jastrzębski et al., 2020 ambitiously studied the more realistic yet complex algorithm of SGD.
In the special case of full-batch training, Jastrzębski et al., 2020  conjectured the Edge of Stability behavior that we observe.
Outside the special case of full-batch training, the “simplified model” in Jastrzębski et al., 2020 does not seem to be numerically accurate.