Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability
Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel pith:566GR33Srecord.jsonopen to challenge →
read the original 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 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.
This paper has not been read by Pith yet.
Forward citations
Cited by 24 Pith papers
-
The Spectral Lifecycle of Transformer Training: Transient Compression Waves, Persistent Spectral Gradients, and the Q/K--V Asymmetry
Transformer weight spectra exhibit transient compression waves that propagate layer-wise, persistent non-monotonic depth gradients in power-law exponents, and Q/K-V asymmetry, with the spectral exponent alpha predicti...
-
Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective
Muon moves faster along signal river directions early but converges slower or oscillates near optima than GD due to orthogonal updates removing scale information, supporting two-stage optimization.
-
Algebraic Dead Directions in LayerNorm Transformers: A Forward-Pass-Only Diagnostic at LLM Scale
The normalized inverse-scale direction of LayerNorm's affine parameters is an exact algebraic kernel of the post-final-norm centred activation covariance for any input distribution in LayerNorm transformers.
-
Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network
Derives an effective free energy model and mean-field kinetic equation for gradient descent dynamics at the edge of stability in two-layer networks, interpreted as a Wasserstein-2 gradient flow.
-
Why Muon Outperforms Adam: A Curvature Perspective
Muon outperforms Adam by reducing curvature penalty via lower Normalized Directional Sharpness, as shown via Taylor approximation on LLM training and proven on stylized quadratic problems with heterogeneous curvature.
-
Convergence of difference inclusions via a diameter criterion
A diameter criterion tied to a potential function certifies convergence of difference inclusions, enabling discrete proofs for first-order optimization methods with diminishing steps.
-
Muon is Not That Special: Random or Inverted Spectra Work Just as Well
Muon succeeds by guaranteeing local step-size optimality rather than by tracking any ideal global geometry, as random-spectrum and quasi-norm variants match its performance on language models.
-
Spectral Dynamics in Deep Networks: Feature Learning, Outlier Escape, and Learning Rate Transfer
A two-level DMFT predicts width-consistent outlier escape and hyperparameter transfer under μP in deep networks, with bulk restructuring dominating for tasks with many outputs.
-
Layer-wise Geometric Approximation Rates for Deep Networks
A shared mixed-activation network of width 2dN+d+2 yields layer-wise L^p approximation rates bounded by the modulus of continuity at geometric scale N^{-ℓ}, reducing to (2d+1)N^{-ℓ} for 1-Lipschitz targets.
-
Momentum Further Constrains Sharpness at the Edge of Stochastic Stability
Momentum SGD exhibits two distinct EoSS regimes for batch sharpness, stabilizing at 2(1-β)/η for small batches and 2(1+β)/η for large batches, aligning with linear stability thresholds.
-
Flat Channels to Infinity in Neural Loss Landscapes
Neural loss landscapes contain flat channels to infinity along which gradient flow leads pairs of neurons to implement gated linear units.
-
Radial Suppression Accelerates Algorithmic Generalization: A Geometric Analysis of Delayed Generalization
A norm penalty constraining activations to a sqrt(d)-radius hypersphere accelerates grokking by up to 6x on modular arithmetic via radial suppression in activation dynamics.
-
Evaluation-Strategy Gap in Fault Diagnosis of Deep Learning Programs
Using a corpus of 5542 fault-injected traces from 38 DL programs, the study finds a 0.19 balanced accuracy gap in fault diagnosis between within-program and cross-program evaluation caused by program-specific feature ...
-
Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods
Lower bounds establish that heavy-ball momentum extends the compute-efficient batch-size window by sqrt(kappa) over SGD in linear regression, with accelerated SGD showing spectrum-dependent CE-serial runtime tradeoffs.
-
Preserving Plasticity in Continual Learning via Dynamical Isometry
Dynamical isometry (Jacobian singular values near 1) preserves plasticity in continual learning; an isometry-promoting regularizer and decoupled AdamO optimizer match or beat prior methods on supervised and RL benchmarks.
-
Factor Augmented High-Dimensional SGD
Proposes Factor-Augmented SGD that runs on streaming high-dimensional data and supplies the first convergence analysis explicitly accounting for latent-factor estimation error.
-
Does Weight Decay Enhance Training Stability?
Weight decay slows progressive sharpening at the edge of stability, inducing damped oscillations in CNNs and a phase transition to sub-2/η sharpness in MLPs driven by parameter-sharpness gradient alignment, yielding m...
-
Spectral Dynamics in Deep Networks: Feature Learning, Outlier Escape, and Learning Rate Transfer
A two-level DMFT tracks bulk and outlier spectral dynamics in wide networks, predicting width-consistent outlier growth and hyperparameter transfer under muP scaling for deep linear nets while noting bulk restructurin...
-
Layer-wise Geometric Approximation Rates for Deep Networks
A fixed-width, shared-activation deep network architecture is constructed so that every intermediate readout approximates the target function at a geometrically decaying rate proportional to depth.
-
Generalization at the Edge of Stability
Training at the edge of stability causes neural network optimizers to converge on fractal attractors whose effective dimension, measured via a new sharpness dimension from the Hessian spectrum, bounds generalization e...
-
Contribution of task-irrelevant stimuli to drift of neural representations
Task-irrelevant stimuli create long-term representational drift in task-relevant features, with drift rate increasing with variance and dimension of the irrelevant subspace, across Hebbian and gradient-based learning.
-
The Hidden Power of Scaling Factor in LoRA Optimization
Alpha in LoRA outperforms learning-rate scaling, follows a square-root law with rank, and enables a minimalist LoRA-alpha method that improves performance across tasks.
-
A Physics-Inspired Optimizer: Velocity Regularized Adam
VRAdam hybridizes Adam's per-parameter adaptation with a physics-inspired velocity regularizer to stabilize training at the edge of stability, delivering better empirical performance than AdamW and O(ln(N)/sqrt(N)) co...
-
There Will Be a Scientific Theory of Deep Learning
A mechanics of the learning process is emerging in deep learning theory, characterized by dynamics, coarse statistics, and falsifiable predictions across idealized settings, limits, laws, hyperparameters, and universa...
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.