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 tracking Hessian sharpness across architectures and tasks. It finds that training typically reaches an Edge of Stability, where sharpness hovers near or above 2/η while loss fluctuates short-term but decreases over longer timescales.

  • Problem

    Neural-network training dynamics are poorly understood, and observed gradient-descent behavior may conflict with prevailing optimization assumptions.

  • Method

    The paper empirically trains neural networks with full-batch gradient descent and tracks Hessian sharpness across architectures and tasks.

  • Results

    Gradient descent enters an Edge of Stability where sharpness hovers at or just above 2/η, while training loss is non-monotonic short-term but decreases over long timescales.

  • Takeaways & Limitations

    The findings question whether conventional optimization assumptions describe neural-network training and motivate rigorous study of the Edge of Stability.

  • Takeaways & Limitations

    The sharpness characterization applies precisely to full-batch gradient descent, not generally to stochastic gradient descent.

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 shows that full-batch gradient descent on neural networks exhibits a broadly consistent Edge of Stability regime, challenging standard optimization assumptions. Sharpness rises toward 2/η, then hovers near or above it while training loss decreases non-monotonically over short timescales but consistently over long timescales.

  • Across architectures and tasks, sharpness continually rises while below 2/η, the instability threshold for quadratic objectives, and crossing it causes rapidly growing oscillations along the greatest-curvature direction.This progressive sharpening characterizes the approach to the Edge of Stability.
  • Full-batch gradient descent typically enters the Edge of Stability, where sharpness hovers at or just above 2/η while training loss decreases non-monotonically short-term but consistently long-term.The dynamics reflect gradient descent continually increasing sharpness while being restrained from increasing it further, allowing optimization to continue without divergence.
  • On standard architectures and CIFAR-10, avoiding the Edge of Stability requires sometimes dramatically smaller, unreasonably slow step sizes, making the regime the rule rather than the exception.At reasonable step sizes, gradient descent enters the Edge of Stability.
  • The Edge of Stability conflicts with conventional optimization assumptions based on L-smoothness, monotone descent, quadratic local models, and common step-size heuristics.The paper presents these empirical phenomena as precise targets for future theoretical study.
  • Earlier work reported related non-monotonic dynamics, instability, and sharpness behavior, but full-batch experiments had not directly observed the progressive sharpening and Edge of Stability phenomena.Prior studies primarily examined SGD or conjectured these effects without running full-batch experiments.
  • For SGD, sharpness does not flatline, but its trajectory depends strongly on step size and batch size, suggesting that Edge of Stability intuition may generalize without centering on sharpness.Existing optimization theory does not explain these dependencies.

2 BACKGROUND: STABILITY OF GRADIENT DESCENT ON QUADRATICS

Gradient descent on quadratic objectives is governed by Hessian eigenvalues: exceeding 2/η causes oscillatory divergence along the corresponding eigendirection. This stability analysis motivates applying quadratic Taylor approximations to neural-network training objectives.

  • Quadratic stability: For a quadratic objective, gradient descent diverges whenever an eigenvalue exceeds 2/η, with convex quadratics satisfying this condition if and only if.A negative eigenvalue causes divergence for any positive step size.
  • Quadratic stability: In one dimension, when a > 2/η, the error multiplier is below −1, producing oscillations around the optimum with exponentially increasing magnitude.The error evolves as (1 − ηa)^t(x0 − x∗).
  • Quadratic stability: In multiple dimensions, eigenvector coordinates evolve independently, so exceeding 2/η in any eigendirection causes divergence along that direction and overall.With eigenvalues a1 = 20 and a2 = 1, η = 0.09 converges in both directions, whereas η = 0.11 diverges along q1.
  • Momentum variants: Polyak and Nesterov momentum also diverge on quadratics when sharpness exceeds an algorithm-specific maximum stable sharpness.This discussion concerns full-batch gradient descent; analyses of SGD are discussed separately.
  • Neural-network connection: For neural-network objectives, a Hessian eigenvalue above 2/η implies divergence of gradient descent on the local quadratic Taylor approximation, with oscillatory growth along its eigenvector.The paper defines sharpness strictly as the maximum eigenvalue of the training-loss Hessian and does not claim a connection to generalization.

3 GRADIENT DESCENT ON NEURAL NETWORKS

Full-batch gradient descent first exhibits progressive sharpening, then typically enters an Edge of Stability where sharpness remains near 2/η while loss decreases non-monotonically short-term but consistently long-term. Gradient descent initially follows a common gradient-flow trajectory until sharpness reaches this stability threshold.

  • 3.1 PROGRESSIVE SHARPENING: Before reaching 2/η, sharpness generally increases during gradient descent, although sharpening is weaker for wider, easier, or shallower networks and can reverse late under cross-entropy.For cross-entropy, increasing classifier margins reduce loss curvature and Hessian sharpness, while the maximum NTK eigenvalue continues increasing.
  • 3.2 THE EDGE OF STABILITY: After sharpness crosses 2/η, gradient descent enters the Edge of Stability: sharpness hovers near 2/η while training loss decreases non-monotonically short-term but consistently long-term.Across step sizes and both MSE and cross-entropy, sharpness stops increasing at approximately 2/η after the breakeven point.
  • 3.2 THE EDGE OF STABILITY: Reducing the step size at the Edge of Stability immediately resumes sharpening until sharpness again reaches the new 2/η constraint, indicating an active implicit constraint.With η changing from 2/200 to 2/300, sharpness increases after the drop and then stops increasing at the new threshold.
  • 3.2 THE EDGE OF STABILITY: The Edge of Stability also occurs with Polyak and Nesterov momentum, whose sharpness rises until plateauing at their maximum stable step-size thresholds.For additional Hessian eigenvalues, each rises until plateauing near 2/η.
  • 3.3 RELATED WORK: Earlier studies observed related instability and sharpness behavior, but full-batch experiments had not previously demonstrated the Edge of Stability directly; this work provides that characterization.Jastrz˛ebski et al. (2020) conjectured the full-batch phenomenon without running full-batch experiments, while Xing et al. (2018) observed non-monotone objective decreases and bouncing iterates.
  • 3.4 THE GRADIENT FLOW TRAJECTORY: For many networks, gradient descent tracks the gradient-flow trajectory until sharpness reaches 2/η; η < 2/λmax avoids the Edge of Stability, whereas larger η eventually reaches it.When η ∈ [2/λmax, 2/λ0], the trajectory is followed only until sharpness reaches 2/η, after which training destabilizes and departs from gradient flow.

4 FURTHER EXPERIMENTS

Further experiments show that the Edge of Stability characterization generalizes across architectures, normalization settings, and tasks. Standard CIFAR-10 networks further demonstrate that stable step sizes can be dramatically and impractically slower than Edge-of-Stability training.

  • Architectures: Across eleven fully connected and convolutional networks, using varied activations and cross-entropy or MSE loss, the experiments reproduce Figure 5.The architectures include max-pooling and average-pooling convolutional networks, with tanh, ReLU, ELU, softplus, and hardtanh activations.
  • Batch normalization: With batch normalization, the findings remain valid, while the experiments reconcile this result with prior observations about batch normalization’s unusual optimization properties.The batch-normalization analysis appears in Appendix K and addresses Santurkar et al. (2018) alongside Li & Arora (2019).
  • Tasks: On a Transformer, synthetic tanh regression, and deep linear networks, sharpness rises until hovering at or just above 2/η.These experiments extend the characterization beyond the single architecture and task studied in Section 3.
  • Standard networks on CIFAR-10: 2227.6 peak sharpness in VGG-BN implies stable-step training would require 16,622 iterations, versus 329 at η = 0.16, making it 50.5× slower.The three standard CIFAR-10 architectures eventually enter the Edge of Stability regime at reasonable step sizes.
  • Tracking gradient flow: Gradient descent tracks gradient flow below 2/η for continuously differentiable architectures and some nonsmooth networks, but not for deep linear networks or Transformers.It does track gradient flow for the tanh network and the three standard CIFAR-10 architectures, including their ReLU-based models.

5 DISCUSSION

The Edge of Stability challenges conventional optimization assumptions: L-smoothness and monotone descent fail at reasonable step sizes, while quadratic approximations and standard step-size prescriptions misrepresent neural-network training.

  • At reasonable step sizes, gradient descent cannot be analyzed using (even local) L-smoothness: At reasonable step sizes, sharpness hovers above 2/η, so gradient descent cannot be analyzed using even local L-smoothness; this extends prior limitations to networks without batch normalization or weight decay.Li et al. (2020b) previously identified this issue for networks using both batch normalization and weight decay.
  • L-smoothness may be inappropriate when analyzing other optimization algorithms too: The failure of even local L-smoothness for gradient descent raises broader questions about its suitability for neural-network optimization and places the burden on authors to justify it empirically.The experiments do not establish that L-smoothness is unjustified for every other optimization algorithm.
  • At reasonable step sizes, gradient descent does not monotonically decrease the training loss: At reasonable step sizes, full-batch gradient descent decreases training loss over long timescales but behaves non-monotonically over short timescales, invalidating monotone-descent analyses.This contrasts with assumptions in Zhang et al. (2020) and Allen-Zhu et al. (2019).
  • The Edge of Stability is inherently non-quadratic: At the Edge of Stability, real neural-network objectives make consistent but choppy progress, whereas their local quadratic approximations diverge quickly, making the regime inherently non-quadratic.The divergence of the quadratic approximation is demonstrated in Appendix E.
  • Dogma for step size selection may be unjustified: Progressive sharpening eventually makes every fixed reasonable step size exceed the prescribed 2/λ_t limit, challenging standard step-size selection dogma.The conventional prescriptions are associated with LeCun et al. (1993; 1998) and Schaul et al. (2013).
  • Dogma for step size selection may be unjustified: The fixed η_t = 1/λ_0 baseline outperformed the continually annealed η_t = 1/λ_t rule, suggesting that tolerating non-monotonicity enables larger and more effective steps.The comparison is reported in Appendix F; the annealed rule decreases the objective every iteration, whereas the fixed rule often increases it.

6 STOCHASTIC GRADIENT DESCENT

The paper’s precise sharpness characterization applies only to full-batch gradient descent, but evidence suggests an analogous acclimation phenomenon may extend to SGD. Large step sizes and small batch sizes steer SGD toward lower-sharpness regions, while existing stability criteria remain insufficiently validated as tight throughout training.

  • Unlike full-batch gradient descent, SGD sharpness need not settle at a fixed hyperparameter-predictable value, although larger steps and smaller batches steer SGD toward lower sharpness.The full-batch 2/η rule is presented as a special case of this broader observation.
  • SGD acclimates to step size and batch size so updates sometimes increase and sometimes decrease training loss in expectation, whereas smaller steps or larger batches consistently decrease it in expectation.This suggests an analogue of the Edge of Stability for SGD, though the findings are not a strict characterization.
  • Existing SGD stability models include expectation-based criteria and sharpness thresholds, but the proposed constraint has no evidence of remaining tight throughout general SGD training.Under strong alignment assumptions, Jastrz˛ebski et al. (2020) reduce the criterion to sharpness below 2/η in the full-batch case, where the paper finds the constraint tight.

7 CONCLUSION … 5. With non-differentiable components, instability sometimes begins when the sharpness is a bit less than 2/η

The paper finds that full-batch gradient descent dynamics are consistent across architectures and tasks yet differ from conventional optimization expectations, while identifying caveats that qualify this characterization. These findings motivate questions about progressive sharpening, divergence avoidance, step-size selection, and generalization beyond gradient descent.

  • 7 CONCLUSION: The authors report surprisingly consistent gradient-descent behavior across architectures and tasks, while emphasizing open questions about its mechanism, step-size selection, and applicability beyond gradient descent.They specifically ask why progressive sharpening occurs, how divergence is avoided at the Edge of Stability, and whether the effect extends to SGD.
  • A CAVEATS: The appendix qualifies the paper’s generic characterization by collecting caveats about full-batch gradient-descent dynamics on neural-network training objectives.The subsequent subsections specify when sharpness falls, rises only slightly, must be measured between iterates, or crosses the instability threshold earlier than expected.
  • 1. With cross-entropy loss, the sharpness often drops at the end of training: With cross-entropy classification, sharpness frequently drops near the end of training as classification accuracy approaches 1.The paper explains this end-of-training effect in Appendix C.
  • 2. For shallow or wide networks, or on simple datasets, sharpness doesn’t rise that much: For shallow or wide networks and easy or small datasets, sharpness may rise only slightly, making stable step sizes reasonable and the Edge-of-Stability range small.The cited examples include Figures 12–14 and 16–18.
  • 3. Sharpness sometimes drops at the beginning of training: Sharpness sometimes drops immediately after initialization, especially with MSE loss combined with batch normalization, where the decline can be precipitous.For most networks the initial drop is slight; Figure 8 illustrates a case where sharpness never returns to its initial value.
  • 4. With batch normalization, need to look at sharpness between iterates: With batch normalization and very small η, sharpness at gradient-descent iterates can plateau below 2/η while between-iterate sharpness plateaus just above 2/η.Thus, diagnosing the Edge of Stability may require examining sharpness between iterates rather than only at the iterates.
  • 5. With non-differentiable components, instability sometimes begins when the sharpness is a bit less than 2/η: With ReLU or hardtanh activations, non-monotonic loss behavior can begin when sharpness is slightly below 2/η because the objective is not continuously differentiable.Activation switches can make the second-order Taylor approximation inaccurate even for tiny weight changes.

B STABILITY OF GRADIENT DESCENT ON QUADRATIC FUNCTIONS

On quadratic objectives, vanilla, Polyak, and Nesterov momentum evolve independently along Hessian eigenvectors. Eigen-directions become divergent when their curvature exceeds the method-specific stability threshold.

  • Shared quadratic structure: All three methods act independently along each Hessian eigenvector, so quadratic dynamics reduce to separate scalar recurrences in the eigenvector basis.The methods considered are vanilla gradient descent, Polyak momentum, and Nesterov momentum.
  • Vanilla gradient descent: For vanilla gradient descent, any Hessian eigenvalue a > 2/η causes the corresponding coordinate sequence {qT xt} to diverge.This follows because the scalar update factor satisfies (1 − ηa) < −1.
  • Nesterov momentum: For Nesterov momentum, the corresponding eigenvector coordinate sequence {qT xt} diverges when its eigenvalue satisfies the theorem’s stated instability condition.The proof reduces the dynamics to a homogeneous second-order difference equation whose recurrence diverges under that condition.
  • Polyak momentum: For Polyak momentum, an eigenvalue a > 1/η(2 + 2β) causes the corresponding coordinate sequence {qT xt} to diverge.The result applies from any initialization with 0 ≤ β < 1 and was previously reported in Tugay & Tanik (1989); Goh (2017).

C CROSS-ENTROPY LOSS

With logistic or cross-entropy loss, sharpness can decrease at the end of training because loss-curvature factors shrink as correctly classified examples’ margins increase, even while network Jacobian factors continue growing. The Gauss–Newton approximation explains this behavior in binary and multiclass experiments.

  • Mechanism: Decreasing loss curvature offsets growing Jacobian factors, causing the Gauss–Newton leading eigenvalue and Hessian sharpness to decline late in training.The Gauss–Newton term empirically dominates the Hessian, while progressive sharpening increases the Jacobian outer-product factors.
  • Binary classification with logistic loss: In Figure 10, binary CIFAR-10 margins rise and logistic-loss second derivatives fall late in training, while the Jacobian-factor eigenvalue keeps rising and sharpness decreases.The Gauss–Newton leading eigenvalue closely approximates the Hessian leading eigenvalue in this experiment.
  • Multiclass classification with cross-entropy loss: For multiclass cross-entropy, probabilities approach one for the correct class and zero for others, broadly shrinking the class-score Hessian matrix.Its diagonal terms have the form p(1 − p), and its off-diagonal terms also vanish as class probabilities become increasingly decisive.
  • Multiclass classification with cross-entropy loss: In Figure 11, CIFAR-10 margins rise and p_i[y_i](1 − p_i[y_i]) declines late in training, while the Jacobian-factor eigenvalue rises and Hessian sharpness falls.The leading eigenvalue of the Gauss–Newton matrix remains an excellent approximation to the Hessian’s leading eigenvalue.

D EMPIRICAL STUDY OF PROGRESSIVE SHARPENING … G EVOLUTION OF SHARPNESS DURING SGD

Across these empirical appendices, progressive sharpening is stronger in narrower, deeper, and larger-data settings, while Edge-of-Stability dynamics undermine quadratic local models and favor fixed over sharpness-tracking step sizes. During SGD, sharpness typically approaches and hovers just above 2/η for large batches, but can decrease for very small batches.

  • D EMPIRICAL STUDY OF PROGRESSIVE SHARPENING: Progressive sharpening increases with depth and dataset size, but decreases with network width under both NTK and standard parameterizations.Under NTK parameterization, λmax and λmax/λ0 decrease as width increases, consistent with λmax/λ0 tending to 1 at infinite width; standard parameterization shows the same width trend, with uncertainty at width 512–1024.
  • E SPEED OF DIVERGENCE ON QUADRATIC TAYLOR APPROXIMATION: At the Edge of Stability, gradient descent on a quadratic Taylor approximation diverges quickly, whereas the real neural objective continues making choppy progress.Before entering the Edge of Stability, the quadratic approximation tracks the real objective; afterward, it becomes an extremely poor model of local progress.
  • F “OPTIMAL” STEP SIZE SELECTION: The fixed step size ηt = 1/λ0 outperforms the dynamic ηt = 1/λt heuristic because progressive sharpening makes sharpness-tracking steps increasingly small.The same result holds for the larger-step comparison: fixed ηt = 1.9/λ0 outperforms dynamic ηt = 1.9/λt.
  • G EVOLUTION OF SHARPNESS DURING SGD: During large-batch SGD, sharpness rises to 2/η and then hovers just above it, paralleling full-batch gradient descent.The experiments use cross-entropy with η = 0.02 and mean squared error with η = 0.01 across multiple batch sizes.
  • G EVOLUTION OF SHARPNESS DURING SGD: Smaller SGD batch sizes produce lower sharpness, consistent with prior reports [Keskar et al., 2016; Jastrz˛ebski et al., 2017; 2019; 2020].This comparison is made across the batch-size conditions in the SGD sharpness experiments.
  • G EVOLUTION OF SHARPNESS DURING SGD: With cross-entropy and batch size 8, sharpness decreases through most training, beginning when train accuracy is only 66%.This decrease is distinguished from the end-of-training decline attributed elsewhere to cross-entropy behavior, and a similar SGD pattern was reported by Jastrz˛ebski et al. (2020).

H SGD ACCLIMATES TO THE HYPERPARAMETERS … I.6 RANDOM PROJECTIONS

The appendix finds that SGD acclimates to its step size and batch size, while documenting experimental setups spanning architectures, normalization, language modeling, Runge–Kutta integration, and random-projection comparisons. Across these settings, the methods use computational approximations such as ghost batches, subset sharpness, adaptive numerical integration, and low-dimensional weight projections.

  • H SGD ACCLIMATES TO THE HYPERPARAMETERS: SGD updates under step size 0.01 and batch size 32 sometimes decrease and sometimes increase full-dataset training loss in expectation.This behavior is observed by Monte Carlo estimation during training, rather than being attributed only to minibatch randomness.
  • H SGD ACCLIMATES TO THE HYPERPARAMETERS: Changing either hyperparameter reverses the expected-loss behavior: smaller steps or larger batches consistently decrease loss, whereas larger steps or smaller batches consistently increase it.The tested alternatives were step sizes 0.005 and 0.02 and batch sizes 16 and 64, compared with the training settings 0.01 and 32.
  • H SGD ACCLIMATES TO THE HYPERPARAMETERS: Across four matched training-and-measurement controls, SGD updates after an initial period still sometimes increase and sometimes decrease training loss in expectation.The control settings were (0.01,16), (0.01,64), (0.02,32), and (0.005,32).
  • I.1 VARYING ARCHITECTURES ON 5K SUBSET OF CIFAR-10: Experiments covered fully connected and convolutional networks on 5,000 CIFAR-10 examples, varying activation functions and pooling choices.The architectures included ReLU, ELU, Tanh, Softplus, and Hardtanh variants, with max or average pooling for convolutional networks.
  • I.2 STANDARD ARCHITECTURES ON CIFAR-10: Standard CIFAR-10 experiments used shallow VGG-11 and ResNet-32 networks, with ghost batch normalization for batch-normalized models and sharpness estimated on data subsets.The ghost-batch objective averaged 50 fixed groups of 1,000 examples; sharpness used the first 5,000 examples or five batches for normalized networks.
  • I.3 BATCH NORMALIZATION EXPERIMENTS: Additional setups inserted ghost batch normalization into CNNs, trained a specified Transformer on WikiText-2, and estimated sharpness on computationally manageable subsets.The Transformer used the PyTorch word-level language-modeling setup with negative log likelihood and 2,500 training examples for sharpness.
  • I.5 RUNGE-KUTTA: Gradient flow was integrated with RK4 using a sharpness-adaptive step size α/λ, with α set to 1 or 0.5.Here λ denotes the most recent sharpness value.
  • I.6 RANDOM PROJECTIONS: Gradient-flow and gradient-descent trajectories were compared through ℓ2 distances between low-dimensional random projections of their weight vectors.A Gaussian matrix projected the d-dimensional weights into k dimensions, with k much smaller than d, reducing storage requirements.

J EXPERIMENTS: VARY ARCHITECTURES … J.11.2 CROSS-ENTROPY LOSS

Across fully connected and convolutional architectures, gradient descent generally tracks the Runge–Kutta trajectory until sharpness reaches 2/η, then departs as sharpness hovers near or above that threshold. Non-differentiable components weaken this pattern, while cross-entropy produces a late-training sharpness decrease in the Runge–Kutta trajectories.

  • J EXPERIMENTS: VARY ARCHITECTURES: The experiments span fully connected and convolutional networks with multiple activations, pooling choices, and both cross-entropy and MSE losses on 5,000 CIFAR-10 examples.The appendix varies tanh, ReLU, ELU, softplus, and hardtanh, including max-pooling and average-pooling convolutional networks.
  • J EXPERIMENTS: VARY ARCHITECTURES: Runge–Kutta sharpness generally increases throughout training, except for a late decrease when cross-entropy loss is used.Runge–Kutta uses adaptive step sizes based on local sharpness and approximates gradient flow for continuously differentiable objectives.
  • J EXPERIMENTS: VARY ARCHITECTURES: Gradient descent tracks gradient flow until sharpness reaches 2/η, after which the trajectories diverge and sharpness hovers near or above 2/η.For continuously differentiable objectives, the trajectory distance is nearly zero below 2/η and begins increasing immediately afterward; MSE typically produces a tiny overshoot, whereas cross-entropy produces a larger gap.
  • J.1–J.3 FULLY-CONNECTED SMOOTH NETWORKS: The cited figures report gradient-flow and gradient-descent results, including a cross-entropy gradient-flow result where sharpness drops at the end of training.
  • J.4 FULLY-CONNECTED RELU NETWORK: For fully connected ReLU networks, gradient descent can diverge from Runge–Kutta from the beginning and become non-monotone before sharpness reaches 2/η.Because ReLU is not continuously differentiable, a unique gradient-flow trajectory is not guaranteed; the reported square-loss case shows early trajectory separation and an early sharpness plateau.
  • J.5 FULLY-CONNECTED HARD TANH NETWORK: For hardtanh networks, the same non-differentiability caveat applies, so a unique gradient-flow trajectory is not guaranteed, while the square- and cross-entropy panels extend the corresponding comparisons.The hardtanh training objective is not continuously differentiable, and its Runge–Kutta and gradient-descent results are presented separately for both losses.
  • J.6–J.8 CONVOLUTIONAL MAX-POOLING NETWORKS: Convolutional networks with max-pooling show both standard threshold tracking and architecture-dependent deviations, while non-differentiability prevents guaranteeing a unique gradient-flow trajectory.The tanh max-pooling square-loss case tracks Runge–Kutta until 2/η, whereas ReLU can fail to track from the outset; ELU and cross-entropy variants provide the corresponding extensions.
  • J.9–J.11 CONVOLUTIONAL AVERAGE-POOLING NETWORKS: For ReLU, because it is not continuously differentiable, the training objective is not continuously differentiable, so a unique gradient-flow trajectory is not guaranteed; cross-entropy sharpness drops at training’s end.

K BATCH NORMALIZATION EXPERIMENTS … L.1 TRANSFORMER ON WIKITEXT-2

Across batch-normalized networks and additional tasks, gradient descent generally reaches the Edge of Stability, where sharpness or effective smoothness hovers at or just above 2/η. Batch normalization does not reduce these quantities along full-batch trajectories, while the Transformer experiment also shows non-monotonic loss after this threshold and no initial tracking of gradient flow.

  • K BATCH NORMALIZATION EXPERIMENTS: For BN networks at very small step sizes, sharpness measured between successive iterates reaches and hovers at 2/η even when sharpness measured at the iterates plateaus below 2/η.The between-iterate quantity is estimated from eight evenly spaced points between successive iterates.
  • K.1 RELATION TO SANTURKAR ET AL. (2018): For both BN and non-BN networks, sharpness hovers at or just above 2/η during full-batch gradient descent, so BN does not decrease sharpness along the optimization trajectory.This behavior holds for convolutional networks with ELU, tanh, and ReLU activations on a 5,000-example CIFAR-10 subset.
  • K.1 RELATION TO SANTURKAR ET AL. (2018): Once gradient descent enters the Edge of Stability, effective smoothness jumps from around zero to 2/η for both BN and non-BN ReLU CNNs, although BN trains faster.The comparison uses the same step-size grid on a 5,000-example CIFAR-10 subset.
  • K.1 RELATION TO SANTURKAR ET AL. (2018): With effective smoothness measured at the actual step size α = η, deep linear networks with and without BN both hover at 2/η, providing no evidence that BN improves smoothness or effective smoothness.At α = 30η, the BN and non-BN networks differ, but that distance exceeds the training step size and does not affect training.
  • L ADDITIONAL TASKS: The appendix extends the findings beyond image classification to Transformer language modeling on WikiText-2, one-hidden-layer toy regression, and deep linear-network training.These tasks are identified as sections L.1, L.2, and L.3, respectively.
  • L.1 TRANSFORMER ON WIKITEXT-2: In the WikiText-2 Transformer, sharpness rises to 2/η and then hovers at or just above it for each step size, while loss becomes non-monotonic after crossing the threshold.Gradient flow sharpness continually rises, and gradient descent does not closely track the gradient-flow trajectory at the beginning of training.

L.2 ONE-DIMENSIONAL TOY REGRESSION TASK · L.3 DEEP LINEAR NETWORK

The toy regression experiments use tanh networks to fit noiseless Chebyshev polynomials, while the deep linear experiments test the same training dynamics on a 20-layer network. Both settings largely reproduce Edge of Stability behavior, with task-dependent deviations in loss monotonicity.

  • L.2 ONE-DIMENSIONAL TOY REGRESSION TASK: The toy task fits 20 noiseless, uniformly spaced points on [−1, 1] labeled by Chebyshev polynomials, using a 100-unit one-hidden-layer tanh network trained to loss 0.05.The datasets use polynomial degrees 3, 4, and 5, and the network uses Xavier initialization.
  • L.2 ONE-DIMENSIONAL TOY REGRESSION TASK: Sharpness rises more for higher-degree Chebyshev polynomials, increasing 1.2×, 3.2×, and 63.5× for degrees 3, 4, and 5, respectively.These results come from gradient-flow fits of the one-hidden-layer tanh network.
  • L.2 ONE-DIMENSIONAL TOY REGRESSION TASK: For Chebyshev degree 5, after sharpness reaches 2/η, loss briefly becomes non-monotonic without progress, then decreases monotonically until training ends.Degree-4 and degree-5 gradient-descent fits otherwise show mostly the same Edge of Stability behavior reported elsewhere.
  • L.3 DEEP LINEAR NETWORK: The deep linear task maps whitened 50-dimensional inputs to targets generated as Y = XA^T, with error measured by square loss.The dataset contains n = 50 datapoints and uses a random Gaussian matrix A.
  • L.3 DEEP LINEAR NETWORK: The network is a 20-layer deep linear composition f(x) = W_L...W_2W_1x, with all layers initialized using Xavier initialization.Each layer is a d × d matrix, and the experiment uses L = 20 layers.
  • L.3 DEEP LINEAR NETWORK: Deep linear-network training also shows mostly the same Edge of Stability behavior, but loss has only a brief blip before otherwise decreasing monotonically for each step size.The experiment compares gradient descent across a range of step sizes; gradient-flow training stops at time 100.

M EXPERIMENTS: STANDARD ARCHITECTURES ON CIFAR-10

Experiments on VGG and ResNet architectures with and without batch normalization on CIFAR-10 reproduce the Edge of Stability findings, including progressive sharpening and severe inefficiency of stable step sizes. Gradient descent tracks gradient flow until sharpness reaches 2/η, after which loss becomes locally non-monotonic while decreasing over longer timescales.

  • Main findings: Across all three CIFAR-10 architectures, sharpness increases below 2/η and then hovers at or just above 2/η in the Edge of Stability.At this transition, training loss becomes non-monotonic over short timescales but continues decreasing over long timescales.
  • Main findings: Gradient descent closely follows Runge-Kutta until sharpness reaches 2/η, despite the architectures using the non-smooth ReLU activation.This agreement holds across VGG with batch normalization, VGG without batch normalization, and ResNet with batch normalization.
  • Stable step sizes: Stable step sizes are extraordinarily small because the maximum sharpness becomes extraordinarily large, making stable-step training require many iterations.The experiments therefore compare stable-step iteration requirements with substantially larger-step training for each architecture.
  • Progressive sharpening: Sharpness rises from 6.38 to 2227.6 for VGG with BN, from 0.64 to 2461.78 for VGG without BN, and from 1.07 to 760.6 for ResNet.Runge-Kutta reaches completion only for VGG with BN; the other architectures are stopped early because high sharpness makes integration computationally infeasible.
  • Stable step sizes: VGG with BN requires at least 16,622 stable-step iterations versus 329 iterations at η = 0.16, making stable training suboptimal by at least 50.5×.The stable-step bound uses maximum sharpness 2227.6 and largest stable step size 0.000897.
  • Stable step sizes: VGG without BN is at least 27.8× suboptimal to 37.1% accuracy and at least 5.54× suboptimal to completion when using a stable step size.The comparisons use 9,876 stable-step iterations versus 355 iterations to 37.1% accuracy and 1,782 iterations to completion at η = 0.16.
  • Stable step sizes: ResNet is at least 271.9× suboptimal to 43.2% accuracy and at least 33.3× suboptimal to completion when using a stable step size.The comparisons use 26,923 stable-step iterations versus 99 iterations to 43.2% accuracy and 807 iterations to completion at η = 2.0.

N EXPERIMENTS: MOMENTUM … Q RELATION TO JASTRZ ˛EBSKI ET AL. (2020)

Across momentum, learning-rate-drop, eigenvalue, and related-work analyses, the appendix finds that sharpness is driven toward the relevant stability threshold, while clarifying the scope and limitations of the comparison with Jastrzębski et al. (2020).

  • N EXPERIMENTS: MOMENTUM; N.1–N.4 ARCHITECTURES; SQUARE-LOSS AND CROSS-ENTROPY SUBSECTIONS: Sharpness rises to the maximum stable sharpness, then plateaus just above it or oscillates around it under Polyak or Nesterov momentum.The experiments use β = 0.9 across four architectures, two losses, and multiple step sizes on a 5k-sized CIFAR-10 subset.
  • N.1–N.4 ARCHITECTURES; N.1.1–N.4.2 LOSS SUBSECTIONS: The momentum experiments span tanh and ReLU fully-connected and convolutional networks with both square and cross-entropy losses, reaching 99% accuracy.Polyak momentum uses β = 0.9, while Nesterov momentum has MSS approximately 1.35714 for β = 0.9 under the reported normalization.
  • O EXPERIMENTS: LEARNING RATE DROP; O.1–O.4 ARCHITECTURES AND LOSSES: After a learning-rate drop, sharpness immediately increases until reaching the new maximum stable sharpness 2/η, unless training finishes first.This supports interpreting the Edge of Stability as a regime where gradient descent continually tries to increase sharpness beyond 2/η but is blocked from doing so.
  • P OTHER EIGENVALUES: The top six Hessian eigenvalues generally rise and plateau; for MSE all exceed 2/η, whereas some lower eigenvalues remain below 2/η for cross-entropy.The comparison uses four step sizes for each of cross-entropy and MSE loss.
  • Q RELATION TO JASTRZ ˛EBSKI ET AL. (2020): Jastrz˛ebski et al. (2020) conjectured Edge of Stability behavior for full-batch training through assumptions predicting sharpness increases below 2/η and decreases above it.Their work primarily targeted the more realistic but more complex setting of SGD, whereas this paper studies full-batch gradient descent.
  • Q RELATION TO JASTRZ ˛EBSKI ET AL. (2020): Jastrz˛ebski et al. (2020) did not observe Edge of Stability in full-batch neural-network experiments or conduct controlled tests of Assumptions 1–4.The paper notes that those assumptions were not claimed to be literally true, and that Assumption 1 is clearly false outside its role in the SGD model.
  • Q RELATION TO JASTRZ ˛EBSKI ET AL. (2020): For full-batch training, Jastrz˛ebski et al. (2020)’s stability condition reduces to sharpness < 2/η, which this paper finds is tight throughout training.For general SGD, however, the condition is not close to tight, which the authors attribute to deficiencies in the stochastic stability model, including potentially invalid assumptions and reliance on stability in expectation.
  • Q RELATION TO JASTRZ ˛EBSKI ET AL. (2020): Overall, the paper observes the conjectured Edge of Stability behavior in the full-batch special case but finds Jastrz˛ebski et al. (2020)’s simplified model numerically inaccurate beyond that case.The distinction reflects the paper’s precise analysis of an unrealistic full-batch algorithm versus the prior work’s broader but more complex SGD setting.
Loading 2103.00065v3…