PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, Ameet Talwalkar

arXiv:2103.00065v3cs.LGstat.ML

TL;DR

The paper addresses limited understanding of full-batch gradient descent dynamics in neural-network training by empirically characterizing sharpness across architectures and tasks. It finds an Edge of Stability regime where sharpness hovers at or above 2/η and loss decreases non-monotonically, challenging standard optimization assumptions.

  • Problem

    The dynamics of full-batch gradient descent in neural-network training are poorly understood, despite gradient descent underlying popular optimization algorithms.

  • Method

    The paper empirically measures Hessian sharpness and training-loss evolution for gradient descent across neural-network architectures and tasks.

  • Results

    Gradient descent typically enters an Edge of Stability where sharpness hovers at or just above 2/η while training loss decreases non-monotonically over short timescales.

  • Takeaways & Limitations

    The findings challenge optimization assumptions based on L-smoothness, monotone descent, quadratic Taylor approximations, and conventional step-size selection.

  • Takeaways & Limitations

    The precise 2/η characterization applies only to full-batch gradient descent; during SGD, sharpness does not always settle at a fixed, hyperparameter-predicted value.

Abstract

from arXiv · show

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 numerical 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. Code is available at https://github.com/locuslab/edge-of-stability.

1 INTRODUCTION

The paper finds that full-batch gradient descent on neural networks exhibits a consistent Edge of Stability regime across architectures and tasks, challenging common optimization assumptions.

  • Full-batch gradient descent reveals a simple training characterization across a broad range of neural network architectures and tasks.
  • Sharpness, defined as the maximum eigenvalue of the training-loss Hessian, rises while below 2/η and then hovers at or just above that threshold.
  • In the Edge of Stability, training loss fluctuates over short timescales but continues decreasing over long timescales without divergence or stalling.
  • For standard architectures on CIFAR-10, reasonable step sizes typically enter the Edge of Stability, making it the rule rather than the exception.
  • The regime conflicts with assumptions based on L-smoothness, monotone descent, quadratic Taylor models, and conventional step-size heuristics.
  • Earlier work reported related non-monotonic optimization, instability, and approximate 2/η sharpness, including observations for SGD.

2 BACKGROUND: STABILITY OF GRADIENT DESCENT ON QUADRATICS

For quadratic objectives, gradient descent stability is governed by the Hessian eigenvalues relative to 2/η. Exceeding this threshold causes oscillatory divergence along the corresponding eigendirections.

  • The 2/η threshold is an exact stability boundary for convex quadratics, while negative eigenvalues cause divergence for every positive step size.
  • On a one-dimensional quadratic, the error is multiplied each iteration by 1 − ηa, so a > 2/η produces oscillations with exponentially increasing magnitude.
  • In multiple dimensions, each Hessian eigencoordinate evolves independently according to one-dimensional quadratic dynamics.
  • With eigenvalues a1 = 20 and a2 = 1, η = 0.09 yields convergence in both directions, whereas η = 0.11 causes divergence along q1.
  • Polyak and Nesterov momentum also have maximum stable sharpness thresholds beyond which quadratic optimization diverges.
  • A quadratic Taylor approximation of a neural-network objective has a Hessian-defined quadratic form whose eigenvalues determine gradient-descent stability.

3 GRADIENT DESCENT ON NEURAL NETWORKS

The paper empirically characterizes full-batch gradient descent on neural networks: sharpness tends to rise below 2/η, then training typically enters the Edge of Stability, where sharpness hovers near 2/η while loss decreases non-monotonically over short timescales but consistently over long timescales.

  • 3.1 PROGRESSIVE SHARPENING: Progressive sharpening causes sharpness to continually increase while it remains below the vanilla gradient descent stability threshold 2/η.The tendency is observed across architectures and tasks, although brief decreases can occur.
  • 3.1 PROGRESSIVE SHARPENING: Cross-entropy training typically differs at the end, when sharpness drops because the loss curvature becomes small at large classifier margins.This is an exception to the usual continual increase in sharpness.
  • 3.1 PROGRESSIVE SHARPENING: 2227.6: VGG sharpness on CIFAR-10 rises from an initial value of 6.3 under gradient flow.The degree of sharpening depends on the network and task, but is dramatic for standard CIFAR-10 architectures.
  • 3.2 THE EDGE OF STABILITY: After sharpness crosses 2/η, gradient descent enters the Edge of Stability, with sharpness near 2/η and train loss decreasing non-monotonically over short timescales.The loss nevertheless decreases consistently over long timescales, rather than diverging entirely or stalling.
  • 3.2 THE EDGE OF STABILITY: Learning-rate reduction restarts progressive sharpening until sharpness reaches the new 2/η boundary, consistent with an active sharpness constraint.For η = 2/200 changed to η = 2/300, sharpness immediately increases after the drop and then stops increasing at the new Edge of Stability.
  • 3.2 THE EDGE OF STABILITY: Stable step sizes are suboptimal for convergence speed, while momentum variants plateau at their maximum stable sharpness and subsequent Hessian eigenvalues also plateau near 2/η.The paper reports this pattern for both Polyak and Nesterov momentum.

4 FURTHER EXPERIMENTS

The paper tests its gradient-descent characterization across architectures, tasks, normalization settings, and standard CIFAR-10 networks, finding repeated Edge of Stability behavior and substantial sharpening.

  • Architectures: The experiments vary eleven architectures on a 5k CIFAR-10 subset, including fully connected and convolutional networks with multiple activations and pooling choices.
  • Architectures: Batch normalization does not eliminate the reported findings, which are confirmed in additional experiments.
  • Tasks: The Edge of Stability appears across a Transformer, a synthetic regression network, and a deep linear network, with sharpness hovering at or just above 2/η.
  • Standard networks on CIFAR-10: 2227.6 peak sharpness is reached from 6.3 at initialization when a VGG-BN on CIFAR-10 is trained along a gradient-flow trajectory.
  • Standard networks on CIFAR-10: 50.5× is the reported training-iteration disadvantage of a stable step size versus η = 0.16 for that VGG-BN example.Stable-step-size training would require at least 16,622 iterations, versus 329 at the Edge of Stability step size.
  • Tracking gradient flow: Gradient descent follows gradient-flow trajectories for some architectures below the 2/η sharpness threshold, but not uniformly across all tested architectures.

5 DISCUSSION

The discussion argues that Edge of Stability dynamics conflict with smoothness-based convergence, monotone-descent assumptions, quadratic local models, and conventional step-size prescriptions.

  • L-smoothness: At reasonable step sizes, sharpness hovers just above 2/η, so convergence analyses requiring local L-smoothness cannot apply to gradient descent in this regime.
  • L-smoothness: The failure of even local L-smoothness in gradient descent raises questions about its suitability for analyzing neural network optimization more generally, though other algorithms are not directly tested.
  • Monotone descent: Short-timescale loss increases coexist with long-timescale decrease, contradicting convergence analyses that assert monotone descent for full-batch gradient descent.
  • Quadratic approximation: At the Edge of Stability, the real neural objective makes consistent but choppy progress while its local quadratic approximation would diverge.
  • Step-size selection: Progressive sharpening eventually makes any fixed reasonable step size exceed the conventional 2/λ_t prescription, while η_t = 1/λ_t is outperformed by a fixed η_t = 1/λ_0 baseline.

6 STOCHASTIC GRADIENT DESCENT

The paper limits its precise sharpness characterization to full-batch gradient descent while relating it cautiously to broader SGD observations and stability analyses.

  • Scope: The precise sharpness characterization applies only to full-batch gradient descent; SGD sharpness need not settle at a fixed, hyperparameter-predictable value.
  • SGD observations: Large step sizes and small batch sizes steer SGD toward lower-sharpness regions, with the full-batch 2/η rule presented as a special case.
  • Possible extension to SGD: The Edge of Stability may inform SGD because full-batch gradient descent is a special case of SGD, although the paper presents this as a possibility rather than a result.
  • Stability modeling: Existing SGD stability frameworks model stability in expectation or under strong alignment assumptions involving sharpness-dependent thresholds.

7 CONCLUSION

The paper finds that full-batch gradient descent has consistent but unconventional dynamics, while identifying several scope boundaries and open questions for this characterization.

  • Gradient descent behavior is surprisingly consistent across architectures and tasks yet differs from conventional optimization assumptions.
  • The Edge of Stability characterization raises open questions about progressive sharpening, divergence avoidance, step-size selection, and generalization beyond gradient descent.
  • Cross-entropy training can produce a late sharpness drop as classification accuracy approaches 1.
  • For shallow or wide networks, easy or small datasets, progressive sharpening may be weak, making stable step sizes reasonable and Edge-of-Stability step sizes narrow.
  • Batch normalization combined with MSE loss can cause a large initial sharpness followed by a precipitous drop.
  • With very small step sizes, sharpness at gradient descent iterates may plateau below 2/η while between-iterate sharpness remains just above 2/η.
  • For ReLU or hardtanh networks, sharpening and non-monotone loss behavior may begin before sharpness crosses 2/η.

B STABILITY OF GRADIENT DESCENT ON QUADRATIC FUNCTIONS

This appendix analyzes gradient descent and momentum methods on quadratic objectives, showing that sufficiently large curvature along an eigenvector causes divergence.

  • Vanilla gradient descent is defined by an iterative update using the gradient and step size η.
  • Polyak and Nesterov momentum introduce a velocity vector and coefficient β, reducing to vanilla gradient descent when β = 0.
  • On quadratic functions, vanilla gradient descent, Polyak momentum, and Nesterov momentum evolve independently along Hessian eigenvectors.
  • Vanilla gradient descent diverges along an eigenvector when its eigenvalue a exceeds 2/η.
  • Nesterov momentum diverges on a quadratic objective when an eigenvalue exceeds its stated stability threshold.
  • Polyak momentum diverges when an eigenvalue a exceeds 1/η(2 + 2β).

C CROSS-ENTROPY LOSS

The appendix explains late sharpness decreases under logistic and cross-entropy losses through increasing classification margins, decreasing loss curvature, and the Gauss-Newton approximation.

  • The Gauss-Newton term typically dominates the Hessian decomposition, motivating a Gauss-Newton approximation to the Hessian.
  • Binary classification with logistic loss: For logistic loss, both the loss and its second derivative decrease as the margin yz increases.
  • Binary classification with logistic loss: After examples are classified correctly, increasing their margins lowers ℓ′′, reducing curvature near the end of training.
  • Binary classification with logistic loss: In logistic-loss experiments, the Jacobian-based matrix keeps rising while the curvature-weighted Gauss-Newton matrix and sharpness decrease late in training.
  • Multiclass classification with cross-entropy loss: For cross-entropy, increasing class margins drives p_i[y_i](1 − p_i[y_i]) downward near the end of training.
  • Multiclass classification with cross-entropy loss: In cross-entropy experiments, rising Jacobian structure combines with decreasing loss curvature, causing the Gauss-Newton eigenvalue and sharpness to fall late in training.

D EMPIRICAL STUDY OF PROGRESSIVE SHARPENING

Gradient-flow experiments show that progressive sharpening is stronger in narrower, deeper, and larger-data settings, while quadratic Taylor models fail at the Edge of Stability and some step-size heuristics underperform.

  • Gradient flow isolates progressive sharpening from instability by approximating what gradient descent would do without instability constraints.
  • Progressive sharpening occurs more strongly for narrower, deeper, and larger-dataset training problems.
  • The effect of width: Across NTK and standard parameterizations, narrow networks exhibit larger sharpness increases than wide networks.
  • The effect of width: The maximum sharpness tends to decrease in expectation as network width increases, although the widest standard-parameterization comparison remains uncertain.
  • Effect of depth: Deeper networks show greater progressive sharpening under both cross-entropy and mean squared error losses.
  • When sharpness exceeds 2/η, gradient descent on the quadratic Taylor approximation quickly diverges at the Edge of Stability, unlike training on the real neural objective.
  • The quadratic Taylor approximation is a poor model of local progress at the Edge of Stability, despite working well away from that regime.
  • The fixed step size ηt = 1/λ0 outperforms the dynamic ηt = 1/λt heuristic in a single numerical example.

H SGD ACCLIMATES TO THE HYPERPARAMETERS

SGD appears to acclimate to its step size and batch size, entering regions where updates sometimes increase and sometimes decrease expected training loss. Smaller steps or larger batches would instead usually decrease the loss in expectation.

  • SGD ACCLIMATES TO THE HYPERPARAMETERS: SGD updates sometimes increase and sometimes decrease the full-batch training loss in expectation during training.This was observed for a tanh network trained with MSE loss, step size 0.01, and batch size 32.
  • SGD ACCLIMATES TO THE HYPERPARAMETERS: A batch size of 16 with step size 0.01 consistently increases training loss in expectation, whereas batch size 64 would decrease it.
  • SGD ACCLIMATES TO THE HYPERPARAMETERS: After a brief initial period, each tested hyperparameter setting produces SGD updates that sometimes increase and sometimes decrease expected training loss.The settings tested were (0.01, 16), (0.01, 64), (0.02, 32), and (0.005, 32) for step size and batch size.
  • SGD ACCLIMATES TO THE HYPERPARAMETERS: For this network, SGD appears to settle into regions where smaller steps or larger batches consistently decrease expected loss, while larger steps or smaller batches increase it.
  • Experimental setup: The experiments use fully connected and convolutional architectures with multiple activations and pooling choices, including CIFAR-10 preprocessing and ghost batch normalization.

J EXPERIMENTS: VARY ARCHITECTURES

Across varied architectures, gradient flow exhibits progressive sharpening, while gradient descent sharpness rises toward 2/η and then hovers near that value. Gradient descent tracks gradient flow before reaching this threshold, but this correspondence is less reliable for nondifferentiable architectures.

  • Gradient flow: Across the tested architectures and tasks, gradient-flow sharpness generally increases during training, except for an end-of-training decrease with cross-entropy loss.
  • Gradient descent: Once sharpness reaches 2/η, gradient descent sharpness ceases increasing substantially and hovers at or just above 2/η.For MSE, the gap is typically tiny; for cross-entropy, it is somewhat larger for reasons the authors do not know.
  • Trajectory comparison: Gradient descent nearly tracks the Runge-Kutta trajectory while sharpness remains below 2/η, then diverges from it immediately afterward.The comparison uses the Runge-Kutta iterate at time t and the gradient-descent iterate at step t/η.
  • Nondifferentiable architectures: For ReLU, hardtanh, or max-pooling architectures, gradient-descent trajectories sometimes differ from Runge-Kutta from the beginning.This may reflect nonexistence of a unique gradient-flow trajectory or step sizes too large to track it.

K BATCH NORMALIZATION EXPERIMENTS

The Edge of Stability findings extend to batch-normalized networks and additional tasks. Batch normalization does not reduce sharpness along full-batch gradient-descent trajectories, and effective smoothness reaches 2/η after entering the regime.

  • Batch-normalized networks: The findings hold for ELU, tanh, and ReLU convolutional networks trained with batch normalization on a 5,000-example CIFAR-10 subset.
  • Sharpness measurement: For batch-normalized networks at very small step sizes, measuring sharpness between iterates can be more informative than measuring it directly at iterates.
  • Sharpness and batch normalization: With or without batch normalization, full-batch gradient descent sharpness hovers at or just above 2/η at reasonable step sizes.
  • Effective smoothness: For both BN and non-BN ReLU CNNs, effective smoothness initially stays near zero, then jumps to 2/η upon entering the Edge of Stability.
  • Deep linear networks: At distance α = η, effective smoothness for deep linear networks with and without BN both hovers at 2/η.
  • Additional tasks: On WikiText-2, Transformer sharpness rises to 2/η and then hovers near that value under gradient descent.

L.2 ONE-DIMENSIONAL TOY REGRESSION TASK

The toy regression experiments fit Chebyshev polynomials with a one-hidden-layer tanh network and observe progressive sharpening that grows with polynomial degree. Gradient descent mostly shows Edge of Stability behavior, with task-specific loss dynamics.

  • Task: The toy task uses 20 uniformly spaced points on [−1, 1], labeled noiselessly by Chebyshev polynomials of degrees 3, 4, and 5.
  • Network: A one-hidden-layer tanh network with 100 hidden units is trained using MSE loss until the loss reaches 0.05.
  • Gradient flow results: The sharpness rises by factors of 1.2, 3.2, and 63.5 for degree-3, degree-4, and degree-5 polynomials, respectively.
  • Gradient descent results: Gradient descent on the degree-4 and degree-5 tasks mostly exhibits Edge of Stability behavior across step sizes.
  • Gradient descent results: For degree 5, the loss briefly becomes non-monotonic without progress after sharpness reaches 2/η, then decreases monotonically until training ends.
  • Deep linear network: A 20-layer deep linear network also mostly exhibits Edge of Stability behavior, although its loss is nearly monotonic apart from a brief blip.

M EXPERIMENTS: STANDARD ARCHITECTURES ON CIFAR-10

Experiments across three standard architectures on CIFAR-10 confirm progressive sharpening and Edge of Stability behavior. Stable step sizes are often dramatically slower than larger step sizes used in practice.

  • Experimental scope: Three architectures—VGG with BN, VGG without BN, and ResNet—reproduce the paper’s main sharpening and Edge of Stability findings.The experiments use standard architectures and CIFAR-10.
  • Edge of Stability: Sharpness rises while below 2/η, then hovers at or just above 2/η as training loss becomes locally non-monotonic but decreases over longer timescales.This pattern is observed across all three architectures.
  • Progressive sharpening: Progressive sharpening is dramatic, making any stable step size extraordinarily small and requiring many iterations.The stable-step-size bound is determined by the maximum sharpness along the gradient-flow trajectory.
  • Stable step-size comparison: 0.000897 was the largest stable step size for VGG with BN, while η = 0.16 reached completion in 329 iterations.The stable-step-size trajectory required 16,622 iterations, making stable training substantially slower.
  • Stable step-size comparison: 27.8× was the lower-bound suboptimality factor for training VGG without BN to 37.1% accuracy at a stable step size.The comparison is 9,876 iterations at the stable step size versus 355 iterations at η = 0.16.
  • Stable step-size comparison: 271.9× was the lower-bound suboptimality factor for training ResNet to 43.2% accuracy at a stable step size.The comparison is 26,923 iterations at the stable step size versus 99 iterations at η = 2.0.

N EXPERIMENTS: MOMENTUM

Momentum experiments examine Polyak and Nesterov methods across four architectures and two losses. Sharpness rises to the maximum stable sharpness, then plateaus just above it or oscillates around it.

  • Momentum dynamics: Sharpness rises until reaching the maximum stable sharpness, then either plateaus just above it or oscillates around it.Oscillations can be rapid, slower, or slow depending on the experiment.
  • Experimental scope: The experiments cover four architectures, two loss functions, and both Polyak and Nesterov momentum with β = 0.9.Training uses a 5k-sized CIFAR-10 subset and runs until 99% accuracy.
  • Momentum dynamics: For Nesterov momentum, the maximum stable sharpness is approximately 1.35714/η when β = 0.9.The appendix gives this expression as the MSS for Nesterov momentum.
  • Step-size intervention: Cutting the step size after reaching the Edge of Stability causes sharpness to increase until the system returns to that regime or training finishes.This experiment supports interpreting the regime as continual pressure to increase sharpness constrained by stability.
  • Other eigenvalues: The top six Hessian eigenvalues generally rise and plateau, although cross-entropy can leave some lesser eigenvalues below 2/η.The appendix compares cross-entropy and MSE across four step sizes.
  • Relation to prior work: Earlier work conjectured Edge of Stability behavior but did not observe it in full-batch neural-network experiments or verify its assumptions through controlled tests.The present work directly studies full-batch gradient descent, whereas Jastrzębski et al. focused mainly on SGD.
Loading 2103.00065v3…