# A PDE-based Explanation of Extreme Numerical Sensitivities and Edge of Stability in Training Neural Networks

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 59 references that are in the map.
- First posted on arXiv 4 Jun 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [Deep Learning with Limited Numerical Precision](https://mlsanity.com/paper/1502.02551/) (2015)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [PDE-Net: Learning PDEs from Data](https://mlsanity.com/paper/1710.09668/) (2017)
- [A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights](https://mlsanity.com/paper/1503.01243/) (2015)
- [Deep Neural Networks Motivated by Partial Differential Equations](https://mlsanity.com/paper/1804.04272/) (2018)
- [Adaptive Gradient Methods with Dynamic Bound of Learning Rate](https://mlsanity.com/paper/1902.09843/) (2019)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [A Variational Perspective on Accelerated Methods in Optimization](https://mlsanity.com/paper/1603.04245/) (2016)
- [Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks](https://mlsanity.com/paper/1710.11029/) (2017)
- [Towards a Mathematical Understanding of Neural Network-Based Machine Learning: what we know and what we don't](https://mlsanity.com/paper/2009.10713/) (2020)
- [Understanding the unstable convergence of gradient descent](https://mlsanity.com/paper/2204.01050/) (2022)
- [Non-convex Finite-Sum Optimization Via SCSG Methods](https://mlsanity.com/paper/1706.09156/) (2017)
- [Beyond the Quadratic Approximation: the Multiscale Structure of Neural Network Loss Landscapes](https://mlsanity.com/paper/2204.11326/) (2022)
- [Equilibrated adaptive learning rates for non-convex optimization](https://mlsanity.com/paper/1502.04390/) (2015)
- [A Lyapunov Analysis of Momentum Methods in Optimization](https://mlsanity.com/paper/1611.02635/) (2016)
- [Deep Relaxation: partial differential equations for optimizing deep neural networks](https://mlsanity.com/paper/1704.04932/) (2017)
- [Analysis of Stochastic Gradient Descent in Continuous Time](https://mlsanity.com/paper/2004.07177/) (2020)
- [Towards stability and optimality in stochastic gradient descent](https://mlsanity.com/paper/1505.02417/) (2015)
- [Channel-Directed Gradients for Optimization of Convolutional Neural Networks](https://mlsanity.com/paper/2008.10766/) (2020)

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

- [Is Flash Attention Stable?](https://mlsanity.com/paper/2405.02803/) (2024)

## Abstract (arXiv)

We discover restrained numerical instabilities in current training practices of deep networks with stochastic gradient descent (SGD), and its variants. We show numerical error (on the order of the smallest floating point bit and thus the most extreme or limiting numerical perturbations induced from floating point arithmetic in training deep nets can be amplified significantly and result in significant test accuracy variance (sensitivities), comparable to the test accuracy variance due to stochasticity in SGD. We show how this is likely traced to instabilities of the optimization dynamics that are restrained, i.e., localized over iterations and regions of the weight tensor space. We do this by presenting a theoretical framework using numerical analysis of partial differential equations (PDE), and analyzing the gradient descent PDE of convolutional neural networks (CNNs). We show that it is stable only under certain conditions on the learning rate and weight decay. We show that rather than blowing up when the conditions are violated, the instability can be restrained. We show this is a consequence of the non-linear PDE associated with the gradient descent of the CNN, whose local linearization changes when over-driving the step size of the discretization, resulting in a stabilizing effect. We link restrained instabilities to the recently discovered Edge of Stability (EoS) phenomena, in which the stable step size predicted by classical theory is exceeded while continuing to optimize the loss and still converging. Because restrained instabilities occur at the EoS, our theory provides new insights and predictions about the EoS, in particular, the role of regularization and the dependence on the network complexity.