# Training on the Edge of Stability Is Caused by Layerwise Jacobian Alignment

arXiv: https://arxiv.org/abs/2406.00127 · HTML: https://mlsanity.com/paper/2406.00127/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 41 references that are in the map.
- First posted on arXiv 31 May 2024.

## Builds on (most cited first, 24 of 24)

- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Deep Learning Scaling is Predictable, Empirically](https://mlsanity.com/paper/1712.00409/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond](https://mlsanity.com/paper/1611.07476/) (2016)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)
- [Understanding the unstable convergence of gradient descent](https://mlsanity.com/paper/2204.01050/) (2022)
- [Measurements of Three-Level Hierarchical Structure in the Outliers in the Spectrum of Deepnet Hessians](https://mlsanity.com/paper/1901.08244/) (2019)
- [Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability](https://mlsanity.com/paper/2207.12678/) (2022)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)
- [Beyond the Quadratic Approximation: the Multiscale Structure of Neural Network Loss Landscapes](https://mlsanity.com/paper/2204.11326/) (2022)
- [Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability](https://mlsanity.com/paper/2305.11788/) (2023)
- [Large Stepsize Gradient Descent for Logistic Loss: Non-Monotonicity of the Loss Improves Optimization Efficiency](https://mlsanity.com/paper/2402.15926/) (2024)
- [Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory](https://mlsanity.com/paper/2307.04204/) (2023)
- [Sharpness-Aware Minimization and the Edge of Stability](https://mlsanity.com/paper/2309.12488/) (2023)
- [Outliers with Opposing Signals Have an Outsized Effect on Neural Network Optimization](https://mlsanity.com/paper/2311.04163/) (2023)

## Built on by (most cited first, 1 of 1)

- [PreMoE: Proactive Inference for Efficient Mixture-of-Experts](https://mlsanity.com/paper/2505.17639/) (2025)

## Abstract (arXiv)

During neural network training, the sharpness of the Hessian matrix of the training loss rises until training is on the edge of stability. As a result, even nonstochastic gradient descent does not accurately model the underlying dynamical system defined by the gradient flow of the training loss. We use an exponential Euler solver to train the network without entering the edge of stability, so that we accurately approximate the true gradient descent dynamics. We demonstrate experimentally that the increase in the sharpness of the Hessian matrix is caused by the layerwise Jacobian matrices of the network becoming aligned, so that a small change in the network preactivations near the inputs of the network can cause a large change in the outputs of the network. We further demonstrate that the degree of alignment scales with the size of the dataset by a power law with a coefficient of determination between 0.74 and 0.98.