Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability

In plain words

Training a neural network with plain gradient descent has a hidden speed limit set by the step size. The "sharpness" is how curved the loss surface is in its steepest direction. Early in training the sharpness climbs. Once it reaches the level that the step size can handle, steps start to overshoot and bounce across that steep direction.

Training does not blow up, though. The sharpness stays pinned at roughly that limit, the loss goes up and down over short stretches, and it still trends down over the whole run. The authors call this the Edge of Stability. It is a rough-but-working regime, and it shows up in the common setups they tested.

The practical point is that the step size sets the sharpness that training settles at, and that common rules about step sizes and steadily falling loss do not hold in this regime.

TL;DR

• Full-batch gradient descent on neural networks drives the top Hessian eigenvalue (the sharpness) up to about 2/η, where η is the step size §3.1 §3.2.
• Past that point, the iterates oscillate along the sharpest direction, but the loss keeps decreasing over long timescales. The sharpness then hovers at or just above 2/η §3.2.
• The authors report this across eleven fully connected and convolutional networks on a 5,000-example CIFAR-10 subset, and on standard VGG and ResNet models on full CIFAR-10 §4 app. J app. M.
• Relevance: if you use plain gradient descent or step-size rules that assume monotone loss or L-smoothness, the paper argues those assumptions can be wrong at ordinary step sizes §5.

Problem

Convergence analyses of gradient descent usually assume either L-smoothness or monotone loss decrease. L-smoothness analyses need the sharpness to stay below 2/η along the trajectory, and monotone analyses assume the loss falls at every step §5. Step-size heuristics built on the quadratic Taylor model similarly keep the step size at or below 2 divided by the local sharpness §5.

The authors ask whether these assumptions hold for neural networks trained with full-batch gradient descent. For standard CIFAR-10 networks, a stable step size is far smaller than the step the networks train well at, which makes stable training impractical §4 app. M.

Method

Setup: full-batch gradient descent with a fixed step size η on a training loss, with the sharpness λ_t defined as the largest eigenvalue of the training loss Hessian at iteration t §1 §3. The authors also use gradient flow, approximated by fourth-order Runge-Kutta integration of the gradient flow ODE, as a reference trajectory §3.4 app. I.

1. Progressive sharpening. While λ_t < 2/η, gradient descent tends to keep increasing λ_t. Brief decreases happen, mostly early in training §3.1 fig. 3. With cross-entropy loss, λ_t usually drops near the end of training §3.1 app. C.
2. Breakeven. The first time λ_t reaches 2/η, the iterates start oscillating along the top Hessian eigenvector with growing amplitude. The loss still does not diverge §3.2 fig. 4.
3. Edge of Stability. After breakeven, λ_t stays at or just above 2/η for the rest of training. The loss is non-monotone over short windows and decreases over long ones §3.2 fig. 5. For MSE the sharpness sits a little above 2/η, and for cross-entropy the gap is larger; the authors do not know why §4.
4. Gradient flow tracking. For the architectures they tested with continuously differentiable components, gradient descent follows the gradient flow trajectory until λ reaches 2/η. Let λ_0 be the sharpness at initialization and λ_max the largest sharpness along the flow. With η < 2/λ_max, descent never enters the Edge of Stability. With η in [2/λ_max, 2/λ_0], it leaves the flow trajectory at breakeven §3.4.
5. Learning rate cut. In one run, the step size was cut from 2/200 to 2/300 at iteration 6000. The sharpness rose again right after the cut, which the authors read as 2/η acting as an active constraint fig. 6 app. O.
6. Momentum. Polyak and Nesterov momentum behave the same way, with the threshold replaced by the maximum stable sharpness in eq. 1. The sharpness rises to that value and then plateaus §2 §3.2 fig. 7.

On quadratics, gradient descent diverges along any eigenvector whose eigenvalue exceeds 2/η §2. Every second-order Taylor model of the loss has this threshold, which motivates the breakeven picture. The authors then show that the Edge of Stability itself is not a quadratic phenomenon §2 §5.

Maximum stable sharpness for momentum eq. 1:

\mathrm{MSS}_{\text{Polyak}}(\eta, \beta) = \frac{2 + 2\beta}{\eta}, \qquad \mathrm{MSS}_{\text{Nesterov}}(\eta, \beta) = \frac{1}{\eta}\cdot\frac{2 + 2\beta}{1 + 2\beta} eq. 1

Pseudocode for the training loop used to measure the Edge of Stability. The paper does not give this as an algorithm; it is a reading of the authors' code and the experiment description §3 app. I.

# Full-batch gradient descent, logging loss and sharpness  [§3] [app. I]
# Input:  params (network weights); train_set (all examples); eta (step size)
# Output: losses per step; sharpness per measurement

def train_full_batch(params, train_set, eta, steps, eig_every):
    # Start with empty logs  [code: src/gd.py#L35-L38]
    losses = []
    sharpnesses = []
    for step in range(steps):
        # Log the mean loss over all training examples  [code: src/gd.py#L41]
        losses.append(mean_loss(params, train_set))
        # Estimate the top Hessian eigenvalue every eig_every steps  [app. I]
        # on the first 5,000 examples  [code: src/gd.py#L45-L47]
        if step % eig_every == 0:
            subset = train_set[:5000]
            sharpnesses.append(top_hessian_eigenvalue(params, subset))
        # Vanilla gradient step with a fixed step size  [§2]
        # [code: src/gd.py#L63-L67]
        params = params - eta * mean_gradient(params, train_set)
    return losses, sharpnesses

Pseudocode for the sharpness estimate, which is the top eigenvalue computed from Hessian-vector products without building the Hessian app. I utilities.py.

# Top Hessian eigenvalue from Hessian-vector products  [app. I]
# Input:  params; subset (examples); count (eigenvalues wanted)
# Output: leading eigenvalue(s) of the training loss Hessian

def top_hessian_eigenvalue(params, subset, count=1):
    # Hessian-vector product, summed over batches of the subset  [app. I]
    # [code: src/utilities.py#L90-L110]
    def hessian_times(vector):
        return hessian_vector_product(params, subset, vector)
    # Lanczos iteration on the matrix-free operator
    # [code: src/utilities.py#L113-L124]
    return lanczos(hessian_times, num_params(params), count)

Pseudocode for gradient flow, used as the reference trajectory app. I flow.py.

# One fourth-order Runge-Kutta step of gradient flow  [app. I]
# Input:  params; train_set; step_size (in time units)
# Output: params advanced by step_size in gradient flow time

def rk4_flow_step(params, train_set, step_size):
    # Slope at the current point: negative full-batch gradient  [app. I]
    first_slope = -mean_gradient(params, train_set)
    # Slope at the midpoint, using the first slope  [app. I]
    second_slope = -mean_gradient(
        params + step_size / 2 * first_slope, train_set
    )
    # Slope at the midpoint, using the second slope  [app. I]
    third_slope = -mean_gradient(
        params + step_size / 2 * second_slope, train_set
    )
    # Slope at the end point, using the third slope  [app. I]
    fourth_slope = -mean_gradient(params + step_size * third_slope, train_set)
    # Weighted average of the four slopes  [app. I]
    return params + step_size / 6 * (
        first_slope + 2 * second_slope + 2 * third_slope + fourth_slope
    )

# RK step size: alpha over the most recent measured sharpness  [app. I]
# [code: src/flow.py#L105]
step_size = alpha / latest_sharpness  # alpha is 1 or 0.5 in the paper

Hyperparameters

Values stated in the paper, with the code where it confirms them. "not stated" marks gaps.

Results

Author-reported. The paper has no results tables; the numbers below come from the text and captions. Hardware: not stated (the code runs on CUDA GPUs).

Paper vs code

Checked the official repository at commit 0d8ceeb3ad98e28930762205934f0a941faeb479 (the authors' repository; the abstract links it). The repository has no LICENSE file, so code_license is recorded as none stated.

• The paper's cross-entropy formula has the wrong denominator: it sums over j ≠ i and omits the correct class app. I, whereas the code uses the standard softmax cross-entropy utilities.py.
• The paper says the CIFAR-10 mean and standard deviation are per channel, computed over the full dataset app. I. The code computes a mean and standard deviation for each of the 3,072 input positions, using only the 50,000 training images cifar.py cifar.py.
• The sharpness is described as the maximum eigenvalue §1. The code calls SciPy's eigsh with its default which='LM' (largest magnitude) utilities.py. See Pitfalls.
• The Transformer experiment is not in the repository: the model and the WikiText loader are commented out archs.py data.py.
• The paper's text says VGG-11 app. I, and the code uses VGG-11 configuration A vgg.py. The caption of Figure 7.3 says VGG-19 fig. 7.3.
• Fully connected, convolutional and loss definitions match the paper: 200-unit hidden layers, the ½-squared-error MSE, and the default initialization archs.py utilities.py.
• The Runge-Kutta step (α / sharpness, with α = 1 or 0.5 in the paper and default α = 1.0 in the code) matches flow.py.

Pitfalls

• Eigenvalue selection. The paper's sharpness is the largest eigenvalue. SciPy's default which='LM' returns the eigenvalues of largest magnitude. If the Hessian has a large negative eigenvalue, that value can be reported as the sharpness. Use an algebraically largest selection when you reproduce this utilities.py.
• Subset estimate. The training loss uses the full training set, but the sharpness is estimated on the first 5,000 examples. For the Transformer, the subset is 2,500 examples app. I gd.py gd.py. The paper calls this an approximation for the standard networks. Its size relative to the full set is not studied.
• Measurement interval not stated. The paper does not say how often the sharpness is measured. The README's examples use 50 or 100 steps README.md.
• Batch norm. The physical batch size sets the ghost-batch size for batch-normalized networks, and the default is 1,000 gd.py. Changing it changes the batch statistics and therefore the trained model. The paper also reports that for batch-norm networks the sharpness at the iterates can plateau below 2/η while the sharpness between iterates sits above it app. A, caveat 4.
• Early instability for ReLU and hardtanh. With non-differentiable activations, the sharpness can start to plateau slightly below 2/η app. A, caveat 5. Do not treat 2/η as an exact trigger for these networks.
• Normalization. The per-feature statistics in the code differ from the per-channel statistics in the paper, and the code computes them on the training split only cifar.py. Reproduced loss values will differ slightly.
• Gradient flow stalls without sharpness updates. The gradient flow step size is α divided by the most recent measured sharpness. The sharpness starts at infinity and is updated only when the eigenvalue estimate runs flow.py flow.py flow.py. If the eigenvalue interval is disabled, the step size is zero, and the Runge-Kutta time loop never terminates flow.py. Keep measurements on for the flow run.
• Stale sharpness in the flow step. Between measurements, the Runge-Kutta step size is based on the last measured sharpness. The sharpness can rise in that window, so the step can be too large. The code adds a maximum step size as a stopgap flow.py README.md.
• Inconsistent gradient flow numbers for VGG with BN. Section 4 says training finishes at time 14.91 §4. The appendix says gradient flow terminates at time 15.66 app. M. The initial sharpness appears as 6.3 §4, 6.38 app. M, and 6.4 in the figure caption fig. 7.3. Use the values that match your own run.
• Stable step rounding. 2/2227.6 is about 0.0008978. The paper rounds it to 0.000897 and gets 16,622 iterations. The unrounded value gives about 16,610 iterations. Compute from 2/λ_max, not from the rounded number.
• The 50.5× figure is a lower bound. It divides flow time by the largest stable step, which assumes gradient descent at that step follows the gradient flow trajectory. The paper observes this for these networks §4. The 329-iteration run at η = 0.16 is a different trajectory, covering about 52.6 time units (0.16 × 329).
• Edge of Stability has no numeric test. The paper says the sharpness "hovers" at or just above 2/η but gives no tolerance. Choose and report your own criterion when you detect the regime.
• Step-size grids are not listed. The exact step sizes in the sweeps, the number of steps per run and the number of seeds are not stated for most figures app. J.
• Transformer not reproducible from the repo. The model code is commented out archs.py. The paper's Transformer results need an independent implementation.
• Figure numbers. Figure numbers in this card follow the order of the LaTeX source. The experiment figures use a custom counter (7.1 to 7.3). The compiled PDF was not checked against these numbers.
• Scope of the claim. The reported pattern is for full-batch gradient descent. For SGD, the paper shows that the sharpness does not settle at a fixed value, and it argues for an analogue in expectation without showing the 2/η threshold §6 app. G app. H.
• Cost. Full-batch sharpness tracking needs Hessian-vector products over the subset at each measurement. The paper's full-data flow runs were slow enough that some runs had to stop early app. M.

Lineage

• Builds on: Jastrzebski et al. 2020 (the "breakeven" conjecture for SGD, which the paper states is the basis for progressive sharpening and the 2/η point in the full-batch case) §3.1 §3.3; Wu et al. 2018 (dynamical instability as implicit regularization; sharpness near 2/η at convergence) §3.3; Xing et al. 2018 (non-monotone loss with bouncing between valley walls) §3.3; Lewkowycz et al. 2020 (catapult behavior when initial sharpness exceeds 2/η) §3.3.
• Compared against: convergence analyses based on L-smoothness and on monotone decrease (for example Zhang et al. 2019 and Allen-Zhu et al. 2018), and quadratic-model step-size rules (LeCun et al.) §5.
• Superseded by: not known from content already in this wiki.