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

1 INTRODUCTION . Neural networks are almost never trained using ( full-batch ) gradient descent , even though gradient descent is the conceptual basis for popular optimization algorithms such as SGD . In this paper , we train neural networks using gradient descent , and find two surprises . First , while little is known about the dynamics of neural network training in general , we find that in the special case of gradient descent , there is a simple characterization that holds across a broad range of network architectures and tasks . Second , this characterization is strongly at odds with prevailing beliefs in optimization . In more detail , as we train neural networks using gradient descent with step size η , we measure the evolution of the sharpness — the maximum eigenvalue of the training loss Hessian . Empirically , the behavior of the sharpness is consistent across architectures and tasks : so long as the sharpness is less than the value 2/η , it tends to continually rise ( §3.1 ) . We call this phenomenon progressive sharpening . The significance of the value 2/η is that gradient descent on quadratic objectives is unstable if the sharpness exceeds this threshold ( §2 ) . Indeed , in neural network training , if the sharpness ever crosses 2/η , gradient descent quickly becomes destabilized — that is , the iterates start to oscillate with ever-increasing magnitude along the direction of greatest curvature . Yet once this happens , gradient descent does not diverge entirely or stall . Instead , it enters a new regime we call the Edge of Stability1 ( §3.2 ) , in which ( 1 ) the sharpness hovers right at , or just above , the value 2/η ; and ( 2 ) the train loss behaves non-monotonically , yet consistently decreases over long timescales . In this regime , gradient descent is constantly “ trying ” to increase the sharpness , but is constantly restrained from doing so . The net effect is that gradient descent continues to successfully optimize the training objective , but in such a way as to avoid further increasing the sharpness.2 In principle , it is possible to run gradient descent at step sizes η so small that the sharpness never rises to 2/η . However , these step sizes are suboptimal from the point of view of training speed , sometimes dramatically so . In particular , for standard architectures on the standard dataset CIFAR-10 , such step sizes are so small as to be completely unreasonable — at all reasonable step sizes , gradient descent eventually enters the Edge of Stability ( see §4 ) . Thus , at least for standard networks on CIFAR-10 , the Edge of Stability regime should be viewed as the “ rule , ” not the “ exception. ” As we describe in §5 , the Edge of Stability regime is inconsistent with several pieces of conventional wisdom in optimization theory : convergence analyses based on L-smoothness or monotone descent , quadratic Taylor approximations as a model for local progress , and certain heuristics for step size selection . We hope that our empirical findings will both nudge the optimization community away from widespread presumptions that appear to be untrue in the case of neural network training , and also point the way forward by identifying precise empirical phenomena suitable for further study . Certain aspects of the Edge of Stability have been observed in previous empirical studies of fullbatch gradient descent ( Xing et al. , 2018 ; Wu et al. , 2018 ) ; our paper provides a unified explanation for these observations . Furthermore , Jastrzębski et al . ( 2020 ) proposed a simplified model for the evolution of the sharpness during stochastic gradient descent which matches our empirical observations in the special case of full-batch SGD ( i.e . gradient descent ) . However , outside the full-batch special case , there is no evidence that their model matches experiments with any degree of quantitative precision , although their model does successfully predict the directional trend that large step sizes and/or small batch sizes steer SGD into regions of low sharpness . We discuss SGD at greater length in §6 . To summarize , while the sharpness does not obey simple dynamics during SGD ( as it does during GD ) , there are indications that the “ Edge of Stability ” intuition might generalize somehow to SGD , just in a way that does not center around the sharpness . 2 BACKGROUND : STABILITY OF GRADIENT DESCENT ON QUADRATICS . In this section , we review the stability properties of gradient descent on quadratic functions . Later , we will see that the stability of gradient descent on neural training objectives is partly well-modeled by the stability of gradient descent on the quadratic Taylor approximation . On a quadratic objective function f ( x ) = 12x TAx + bTx + c , gradient descent with step size η will diverge if3 any eigenvalue of A exceeds the threshold 2/η . To see why , consider first the onedimensional quadratic f ( x ) = 12ax 2 + bx+ c , with a > 0 . This function has optimum x∗ = −b/a . Consider running gradient descent with step size η starting from x0 . The update rule is xt+1 = xt − η ( axt + b ) , which means that the error xt − x∗ evolves as ( xt+1 − x∗ ) = ( 1− ηa ) ( xt − x∗ ) . Therefore , the error at step t is ( xt − x∗ ) = ( 1 − ηa ) t ( x0 − x∗ ) , and so the iterate at step t is xt = ( 1− ηa ) t ( x0−x∗ ) +x∗ . If a > 2/η , then ( 1− ηa ) < −1 , so the sequence { xt } will oscillate around x∗ with ever-increasing magnitude , and diverge . Now consider the general d-dimensional case . Let ( ai , qi ) be the i-th largest eigenvalue/eigenvector of A . As shown in Appendix B , when the gradient descent iterates { xt } are expressed in the special coordinate system whose axes are the eigenvectors of A , each coordinate evolves separately . In particular , the coordinate for each eigenvector qi , namely 〈qi , xt〉 , evolves according to the dynamics of gradient descent on a one-dimensional quadratic objective with second derivative ai . 1This nomenclature was inspired by the title of Giladi et al . ( 2020 ) . 2In the literature , the term “ sharpness ” has been used to refer to a variety of quantities , often connected to generalization ( e.g . Keskar et al . ( 2016 ) ) . In this paper , “ sharpness ” strictly means the maximum eigenvalue of the training loss Hessian . We do not claim that this quantity has any connection to generalization . 3For convex quadratics , this is “ if and only if. ” However , if A has a negative eigenvalue , then gradient descent with any ( positive ) step size will diverge along the corresponding eigenvector . Therefore , if ai > 2/η , then the sequence { 〈qi , xt〉 } will oscillate with ever-increasing magnitude ; in this case , we say that the iterates { xt } diverge along the direction qi . To illustrate , Figure 2 shows a quadratic function with eigenvalues a1 = 20 and a2 = 1 . In Figure 2 ( a ) , we run gradient descent with step size η = 0.09 ; since 0 < a2 < a1 < 2/η , gradient descent converges along both q1 and q2 . In Figure 2 ( b ) , we use step size η = 0.11 ; since 0 < a2 < 2/η < a1 , gradient descent converges along q2 yet diverges along q1 , so diverges overall . Polyak momentum ( Polyak , 1964 ) and Nesterov momentum ( Nesterov , 1983 ; Sutskever et al. , 2013 ) are notable variants of gradient descent which often improve the convergence speed . On quadratic functions , these two algorithms also diverge if the sharpness exceeds a certain threshold , which we call the “ maximum stable sharpness , ” or MSS . In particular , we prove in Appendix B that gradient descent with step size η and momentum parameter β diverges if the sharpness exceeds : MSSPolyak ( η , β ) = 1 η ( 2 + 2β ) , MSSNesterov ( η , β ) = 1 η ( 2 + 2β 1 + 2β ) . ( 1 ) The Polyak result previously appeared in Goh ( 2017 ) ; the Nesterov one seems to be new . Note that this discussion only applies to full-batch gradient descent . As we discuss in §6 , several recent papers have proposed stability analyses for SGD ( Wu et al. , 2018 ; Jastrzębski et al. , 2020 ) . Neural network training objectives are not globally quadratic . However , the second-order Taylor approximation around any point x0 in parameter space is a quadratic function whose “ A ” matrix is the Hessian at x0 . If any eigenvalue of this Hessian exceeds 2/η , gradient descent with step size η would diverge if run on this quadratic function — the iterates would oscillate with ever-increasing magnitude along the corresponding eigenvector . Therefore , at any point x0 in parameter space where the sharpness exceeds 2/η , gradient descent with step size η would diverge if run on the quadratic Taylor approximation to the training objective around x0 . 3 GRADIENT DESCENT ON NEURAL NETWORKS . In this section , we empirically characterize the behavior of gradient descent on neural network training objectives . Section 4 will show that this characterization holds broadly . 3.1 PROGRESSIVE SHARPENING . When training neural networks , it seems to be a general rule that so long as the sharpness is small enough for gradient descent to be stable ( < 2/η , for vanilla gradient descent ) , gradient descent has an overwhelming tendency to continually increase the sharpness . We call this phenomenon progressive sharpening . By “ overwhelming tendency , ” we mean that gradient descent can occasionally decrease the sharpness ( especially at the beginning of training ) , but these brief decreases always seem be followed by a return to continual increase . Jastrzębski et al . ( 2020 ) previously hypothesized ( in their Assumption 4 ) that a similar phenomenon may hold for SGD , but the evidence for , and the precise scope of , this effect are both currently far clearer for gradient descent than for SGD . Progressive sharpening is illustrated in Figure 3 . Here , we use ( full-batch ) gradient descent to train a network on a subset of 5,000 examples from CIFAR-10 , and we monitor the evolution of the sharpness during training . The network is a fully-connected architecture with two hidden layers of width 200 , and tanh activations . In Figure 3 ( a ) , we train using the mean squared error loss for classification ( Hui & Belkin , 2020 ) , encoding the correct class with 1 and the other classes with 0 . We use the small step size of η = 2/600 , and stop when the training accuracy reaches 99 % . We plot both the train loss and the sharpness , with a horizontal dashed line marking the stability threshold 2/η . Observe that the sharpness continually rises during training ( except for a brief dip at the beginning ) . This is progressive sharpening . For this experiment , we intentionally chose a step size η small enough that the sharpness remained beneath 2/η for the entire duration of training . Cross-entropy . When training with cross-entropy loss , there is an exception to the rule that the sharpness tends to continually increase : with cross-entropy loss , the sharpness typically drops at the end of training . This behavior can be seen in Figure 3 ( b ) , where we train the same network using the cross-entropy loss rather than MSE . This drop occurs because once most data points are classified correctly , gradient descent tries to drive the cross-entropy loss to zero by scaling up the margins , as detailed in Soudry et al . ( 2018 ) . As we explain in Appendix C , this causes the sharpness to drop . The effect of width . It is known that when networks parameterized in a certain way ( the “ NTK parameterization ” ) are made infinitely wide , the Hessian moves a vanishingly small amount during training ( Jacot et al. , 2018 ; Lee et al. , 2019 ; Li & Liang , 2018 ) , which implies that no progressive sharpening occurs . In Appendix D , we experiment with networks of varying width , under both NTK and standard parameterizations . We find that progressive sharpening occurs to a lesser degree as networks become increasingly wide . Nevertheless , our experiments in §4 demonstrate that progressive sharpening occurs to a dramatic degree for standard architectures on the standard dataset CIFAR-10 . We do not know why progressive sharpening occurs , or whether “ sharp ” solutions differ in any important way from “ not sharp ” solutions . These are important questions for future work . Note that Mulayoff & Michaeli ( 2020 ) studied the latter question in the context of deep linear networks .

Summary: This paper presents an interesting observation for GD. That is, the sharpness of the learnt model in the final phase of the training (measured by the largest eigenvalue of the training loss Hessian) hovers right at the value 2/\eta while the training loss. At the same time, the loss goes to unstable and non-monotonically decreasing. This pattern is consistent across architecture, activation functions, tasks, loss functions and BN. Comprehensive experiments are conducted to show this common observation. The paper is easy to follow.