Gradient descent stability

Overview

Plain gradient descent has a speed limit set by its step size η. The sharpness λ_t is the top eigenvalue of the training-loss Hessian at iteration t. Once λ_t reaches 2/η, the iterates start to oscillate along the sharpest direction, but they do not diverge. Sharpness then stays at or just above 2/η, and the loss rises and falls over short windows while still decreasing over long ones 2103.00065 §3.2. The authors call this the Edge of Stability.

The common analyses this challenges assume either that sharpness stays below 2/η along the trajectory, or that the loss decreases at every step 2103.00065 §5. The paper's question is whether those assumptions hold for neural networks trained with full-batch gradient descent.

Variants

The first four variants come from one paper. Its results split the behaviour by step size and optimizer, so they are regimes of one phenomenon, not competing methods. The next three come from three other papers. Each measures something different: gradient agreement under Adam, gradient spikes in ViT training, and loss spikes from float32 rounding. None of them is a sharpness result, so none confirms or refutes the 2/η picture by itself.

Stable regime: η below 2/λ_max

Along the gradient flow trajectory, the largest sharpness is λ_max. With η < 2/λ_max, descent follows the flow and never enters the Edge of Stability 2103.00065 §3.4. For the networks tested, the step that keeps training stable is far smaller than the step they train well at 2103.00065 §4 2103.00065 app. M. Card: arXiv:2103.00065 paper card.

Edge of Stability: η between 2/λ_max and 2/λ_0

With η in this range, descent leaves the flow trajectory at breakeven, where sharpness reaches 2/η, and then sits at that level 2103.00065 §3.4 2103.00065 §3.2. The paper reports this for fully connected, convolutional and batch-normalized networks 2103.00065 §4. For MSE the sharpness settles slightly above 2/η. For cross-entropy the gap is larger, and the authors do not know why 2103.00065 §4.

Momentum: a higher threshold

With Polyak or Nesterov momentum β, the threshold moves from 2/η to a maximum stable sharpness. The sharpness rises to that value and then plateaus 2103.00065 §3.2 2103.00065 eq. 1. The formulas are in the implementation sketch.

Changing the step size during training

A learning-rate cut in one run made sharpness rise again right after the cut. The authors read 2/η as an active constraint on the sharpness 2103.00065 fig. 6 2103.00065 app. O. A toy comparison in the paper found that a fixed rule η = 1/λ_0 beat an "optimal" annealing rule based on the current sharpness 2103.00065 §5 2103.00065 app. F. This is one toy setting, so treat it as a single data point.

Stochastic gradient descent

The paper's 2/η picture is for full-batch gradient descent. For SGD, sharpness does not settle at a fixed value. The authors argue for an analogue in expectation but do not show the 2/η threshold for SGD 2103.00065 §6 2103.00065 app. G 2103.00065 app. H.

Gradient signal-to-noise under full-batch Adam

This variant does not measure sharpness. A third paper trains physics-informed neural networks (PINNs) with full-batch Adam and tracks a batch-wise gradient signal-to-noise ratio (SNR): the norm of the mean batch gradient divided by the norm of the batch standard deviation 2403.18494 §2.1 [arxiv:2403.18494 eq. 7]. The reported training has three phases: fitting, diffusion, and a late "diffusion equilibrium", where the SNR is flat and the test error falls fastest 2403.18494 §1.4 2403.18494 §4 2403.18494 fig. 2. The authors tie that equilibrium to residuals that are homogeneous across collocation points 2403.18494 §2.4. Their reweighting method, residual-based attention, is covered under adaptive loss weighting. Its reported Allen-Cahn speed-up comes from one run 2403.18494 fig. 4. The paper tests five problems, and its code covers three of them. Card: Learning in PINNs: Phase transition, diffusion equilibrium, and generalization.

Gradient spikes in self-supervised ViT training

This variant does not measure sharpness. A second paper trains Vision Transformers with MoCo v3, a contrastive self-supervised method, and reports training instability. Sudden gradient spikes produce dips in the training curves, and the authors say the dips can be small enough that a sparse check misses them. A dense kNN monitor was needed to see them 2104.02057 §4.1 2104.02057 app. A. Runs can also lose about 1 to 3 points of accuracy without diverging 2104.02057 §1. Card: Empirical Study of Training Self-Supervised Vision Transformers.

The stabiliser the paper tests is to freeze the patch projection at its random initialisation, using a stop-gradient 2104.02057 §4.2. In the ablations, the frozen projection scored higher than the learned one at each of the three base learning rates tried: 70.8 vs 70.4 at 0.5e-4, 72.8 vs 72.2 at 1.0e-4, and 73.4 vs 71.7 at 1.5e-4 (linear probing top-1, %) 2104.02057 fig. 5. The paper says the trick reduces instability but does not remove it, and that training can still become unstable if the learning rate is too large 2104.02057 §4.2. Results come from ViT-B/16 on ImageNet-1k with 100-epoch ablations, so they are not a general recipe.

Loss spikes from float32 rounding

This variant does not measure sharpness either. A fourth paper trains with unregularised cross-entropy and Adam, and attributes periodic loss spikes (the Slingshot pattern) to float32 rounding in the loss, not to the optimiser 2605.06152 §1 §3.1. The mechanism has steps. Once the margin between the correct logit and the largest wrong logit exceeds about 16, the log-sum-exp denominator absorbs the other classes, and the correct-class gradient rounds to exactly zero 2605.06152 §3.1 eq. 1 eq. 2. The wrong-class gradients stay non-zero and break the zero-sum property, so the classifier mean and feature mean drift and grow together. The authors call this numerical feature inflation 2605.06152 §3.2 [thm. 2]. When the margins fall back under the threshold, the correct-class gradient returns. Adam's second moment is small after the long quiet period, so the step is large. The size jump at the spike is about 50 times the previous epoch's update 2605.06152 §3.3, and §4.1 gives more than 40 times 2605.06152 §4.1.

The authors' mitigations come from small setups. A float64 loss, a zero-sum projection of the logit gradient, or BatchNorm directly before the classifier removed the spikes in their runs 2605.06152 §4.2 §4.1 fig. 3c. A LayerNorm before the classifier did not prevent them fig. 3b. In their setting, Adam ε of 1e-5 removed the spikes and 1e-6 reduced them 2605.06152 §4.2 fig. 3. Label smoothing removed the precision-driven spikes but added spikes of the edge-of-stability type 2605.06152 §4.2 app. A.2. ResNet18 on CIFAR-10 showed no spike tab. 1. The tested settings are modular division, CIFAR-10 and one nanoGPT run, and the paper reports no accuracy numbers for the mitigations. Card: Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes.

When to use which

Monitoring sharpness costs Hessian-vector products on a subset at each measurement. The paper's full-data flow runs were slow enough that some were stopped early 2103.00065 app. M.

Implementation sketch

Full-batch training loop with sharpness logging, then the regime check and the momentum threshold. The paper estimates sharpness on the first 5,000 training examples 2103.00065 app. I. It gives no tolerance for "hovers", so choose and report your own criterion.

# Full-batch gradient descent with sharpness logging
# [arxiv:2103.00065 §3] [arxiv:2103.00065 app. I]
# Input:  params, train_set, eta (fixed step size), steps, eig_every
# Output: loss log and sharpness log

def train_with_sharpness(params, train_set, eta, steps, eig_every):
    # Start with empty logs
    losses = []
    sharpnesses = []
    for step in range(steps):
        # Loss over the full training set [arxiv:2103.00065 §3]
        losses.append(mean_loss(params, train_set))
        # Sharpness: top Hessian eigenvalue on a subset, via
        # Hessian-vector products and Lanczos [arxiv:2103.00065 app. I]
        if step % eig_every == 0:
            sharpnesses.append(top_hessian_eig(params, train_set[:5000]))
        # Plain gradient step with fixed eta [arxiv:2103.00065 §2]
        params = params - eta * mean_gradient(params, train_set)
    return losses, sharpnesses

# Classify the regime from the latest sharpness [arxiv:2103.00065 §3.2]
# Input:  sharp (latest measurement), eta (step size)
# Output: a label for the regime

def regime(sharp, eta):
    # Below the threshold, sharpness is still climbing [arxiv:2103.00065 §3.1]
    if sharp < 2 / eta:
        return "progressive sharpening"
    # At or just above the threshold: Edge of Stability [arxiv:2103.00065 §3.2]
    return "edge of stability"

# Momentum threshold in place of 2/eta [arxiv:2103.00065 eq. 1]
# Input:  eta (step size), beta (momentum), nesterov (flag)
# Output: maximum stable sharpness for this optimizer

def mss_threshold(eta, beta, nesterov=False):
    # Polyak momentum [arxiv:2103.00065 eq. 1]
    if not nesterov:
        return (2 + 2 * beta) / eta
    # Nesterov momentum [arxiv:2103.00065 eq. 1]
    return (1 / eta) * (2 + 2 * beta) / (1 + 2 * beta)

# Batch-wise gradient SNR over the weights [arxiv:2403.18494 eq. 7]
# Input:  weights, batches (list of index sets at the same parameters)
# Output: scalar SNR

def batch_snr(weights, batches):
    # Gradient of each batch, with no optimizer step [arxiv:2403.18494 §2.1]
    grads = [batch_gradient(weights, b) for b in batches]
    # Per-element mean and population std (ddof = 0) over batches
    # [arxiv:2403.18494 §2.1] [arxiv:2403.18494 app. B]
    mu = mean(grads, axis=0)
    sigma = std(grads, axis=0)
    # Ratio of global norms over all weight elements [arxiv:2403.18494 eq. 7]
    return norm(mu) / norm(sigma)

# Float32 absorption check for one example [arxiv:2605.06152 §3.1] [eq. 1]
# Input:  logits (K,) in float32; r = index of the correct class; p = 24
# Output: True if the correct-class gradient is expected to round to zero

def correct_grad_absorbed(logits, r, p=24):
    # Largest logit among the wrong classes
    wrong_max = max(logits[k] for k in range(len(logits)) if k != r)
    # Margin between the correct logit and the largest wrong logit
    margin = logits[r] - wrong_max
    # Threshold (p - 1) * ln 2, about 16 for float32 [arxiv:2605.06152 §3.1]
    # Depends on the kernel and dtype; check the loss path [tab. 4]
    return margin > (p - 1) * log(2)

# Patch projection frozen at its random start, via a stop-gradient
# [arxiv:2104.02057 §4.2]
# Input:  images (N, 3, 224, 224); w_patch (random init, never trained)
# Output: patch tokens (N, 196, d)

def patch_embed(images, w_patch):
    # Project 16 x 16 patches; the weights receive no gradient
    # [arxiv:2104.02057 §4.2]
    return project(images, stop_gradient(w_patch))

Contradictions

No direct disagreement on a shared quantity is known. The four papers measure different things: the first tracks the Hessian eigenvalue of full-batch gradient descent, the second tracks gradient spikes and kNN accuracy dips in MoCo v3 ViT training, the third tracks batch-wise gradient SNR under full-batch Adam on PINNs, and the fourth tracks float32 absorption and feature drift in cross-entropy training. Do not read one as confirming or refuting another.

The float32 paper and the sharpness paper explain loss behaviour differently. The float32 paper says its trigger is rounding and that the drift is a precondition for the spike, not a side effect of curvature 2605.06152 §1 app. A.2. The sharpness paper describes the oscillation at 2/η that does not diverge 2103.00065 §3.2. The two use different optimizers, losses and scales. The float32 paper's label-smoothing runs show edge-of-stability-type spikes 2605.06152 §4.2 app. A.2, so the two mechanisms can appear in one setting. Neither side is resolved here.

Within the float32 paper, the growth at the spike is about 50 times in §3.3 and more than 40 times in §4.1 2605.06152 §3.3 §4.1. The card records both, and the figure is not a constant to reuse.

Within the PINN paper, the text gives a smoothing window of 100 past iterations 2403.18494 §3.1, while its notebook uses a window of 60 stored rows. The stored rows are not every iteration, so the two windows are not the same length. SNR curves reproduced from the code may differ from the paper's; the paper card lists the details.

Inside the first paper, two numbers for the same VGG-with-BN gradient-flow run differ between sections: training ends at time 14.91 in the main text, but at time 15.66 in the appendix 2103.00065 §4 2103.00065 app. M. The initial sharpness is also given as 6.3 in the text and 6.38 in the appendix 2103.00065 §4 2103.00065 app. M. The paper's code may also differ from its text on normalization and eigenvalue selection; the paper card lists the details.

Papers

• arXiv:2103.00065, Cohen et al.: defines the 2/η threshold and the Edge of Stability for full-batch gradient descent, with gradient-flow tracking and momentum thresholds.
• arXiv:2104.02057, Chen et al.: studies MoCo v3 training of ViTs, reports gradient-spike instability hidden by final accuracy, and tests a frozen random patch projection as a stabiliser.
• arXiv:2403.18494, Anagnostopoulos et al.: tracks batch-wise gradient SNR in full-batch Adam on PINNs, names a "diffusion equilibrium" phase, and proposes residual-based attention.
• arXiv:2605.06152, Liu et al.: attributes Slingshot loss spikes in unregularised cross-entropy training to float32 absorption, names numerical feature inflation, and tests float64, zero-sum and BatchNorm mitigations.