A Loss Curvature Perspective on Training Instabilities of Deep Learning Models

1 INTRODUCTION . Optimization of neural networks can easily fail . While recent architectural advances such as skip connections ( He et al. , 2016a ) and Batch Normalization ( Ioffe and Szegedy , 2015 ) have been applied successfully to produce architectures and hyperparameters that reliably train well , even small changes to a trainable configuration can easily result in training that diverges . More generally , producing a configuration that strikes the right balance between stable training and rapid optimization progress on a new domain can be difficult—practitioners and researchers have few reliable heuristics to guide them through the process . As a result , the specific hyperparameter tuning protocol has an outsized influence on the results ( Choi et al. , 2019 ; Sivaprasad et al. , 2020 ) and successes often rely on large hyperparameter searches ( Nado et al. , 2021 ) . Developing a principled understanding of what makes general architectures trainable would allow researchers to more reliably navigate this process and has the potential to dramatically accelerate research into finding better , more scalable architectures . The focus of the empirical investigation of this work is to better understand what limits the maximum trainable learning rate for deep learning models trained with the typical minibatch stochastic gradient descent ( SGD ) family algorithms . As part of this investigation , we examine several methods developed by the deep learning community that have enabled training at larger learning rates and improved performance . Many methods have been developed that can achieve this goal , notably normalization , learning rate warmup , gradient clipping ( Pascanu et al. , 2013 ) , and better model initializations such as Fixup ( Zhang et al. , 2019b ) , MetaInit ( Dauphin and Schoenholz , 2019 ) , and GradInit ( Zhu et al. , 2021 ) . While these methods are certainly not exactly equivalent , a key property they all have in common is that they can enable training at larger learning rates when applied to certain models ( see for example Figure 1 ) . A natural hypothesis is that methods which enable training at larger learning rates do so through reducing the sharpness1 of the loss surface during training . Indeed , this hypothesis has already been proposed as one of the beneficial effects of Batch Normalization ( Ghorbani et al. , 2019 ; Santurkar et al. , 2018 ) and residual connections ( Li et al. , 2017 ) , and quadratic models of the loss surface 1Throughout this work we will use the term sharpness to refer to the maximum eigenvalue of the loss Hessian , denoted as λ1 . See Appendix B for more details . predict that optimization with SGD is unstable when λ1 > 2/η ( Wu et al. , 2018 ) . However , recent empirical investigations into the relevance of quadratic stability bounds to neural network training have either focused on smaller models , focused on full batch training at small learning rates , and do not investigate connections between sharpness , model initialization and learning rate warmup ( Cohen et al. , 2021 ; Jastrzebski et al. , 2020 ) . In this work , we design a series of large scale experiments studying the evolution of the loss sharpness as we vary the learning rate , warmup period , initialization , and architectural choices . Our results demonstrate the central role that λ1 plays in neural network optimization—maintaining sufficiently small λ1 during optimization is a necessary condition for successful training at large learning rates . Consequently , reducing λ1 is a primary benefit of proper tuning of a number of architecture and optimization hyperparameters : including model initialization , location of normalization , and warmup schedule . Specifically , we show the following : • We provide large scale empirical confirmation that training of neural networks with SGD+momentum is stable only when the optimization trajectory primarily resides in a region of parameter space where λ1 . 2/η , where η denotes the learning rate . This corroborates the theoretical predictions of Wu et al . ( 2018 ) and recent empirical observations of Jastrzebski et al . ( 2020 ) and Cohen et al . ( 2021 ) . • We demonstrate that several successful initialization strategies for architectures without normalization operate primarily by reducing curvature early in training , enabling training at larger learning rates . • We show that learning rate warmup gradually reduces λ1 during training , offering similar benefits to better model initialization . We connect the mechanism by which warmup operates to the dynamical stability model of Wu et al . ( 2018 ) . • We show that learning rate warmup is a simple yet competitive baseline for research into better model initialization . We demonstrate that key progress in this area ( Dauphin and Schoenholz , 2019 ; Zhang et al. , 2019b ; Zhu et al. , 2021 ) can be matched by the application of learning rate warmup and/or gradient clipping alone . • Finally , we show that large loss curvature can result in poor scaling at large batch sizes and interventions designed to improve loss conditioning can drastically improve the model ’ s ability to leverage data parallelism . 2 RELATED WORK . Understanding BatchNorm The loss Hessian has been a central object of study for understanding optimization of neural networks . Santurkar et al . ( 2018 ) argues that an important benefit of Batch Normalization is improved smoothness of the loss surface , while Lewkowycz et al . ( 2020 ) notes that this is improved smoothness is only observed when higher learning rates are used in combination with Batch Normalization . Our results are generally consistent with this current understanding of Batch Normalization , however some of our experiments provide additional nuance—notably we observe several instances where models suffer from training instability ( and high loss curvature ) early in training despite using Batch Normalization ( see Section 4 ) . Evolution of the loss Hessian Recent research has closely studied the interaction between sharpness and learning rate . Wu et al . ( 2018 ) provides a dynamical stability model which predicts that the loss curvature at convergence must satisfy2 λ1 ≤ 2/η . Recent work has provided empirical evidence that λ1 . 2/η often holds well before convergence ( Cohen et al. , 2021 ; Jastrzebski et al. , 2020 ) . Cohen et al . ( 2021 ) focused on full batch training at small learning rates , and observed “ progressive sharpening ” , where λ1 increases during training until λ1 ≈ 2.0/η . We observe the progressive sharpening phenomenon also occurs for many models trained with SGD , though we do not investigate batch sizes ≤ 8 , where Cohen et al . ( 2021 ) argue that progressive sharpening does not occur . We note that Wu et al . ( 2018 ) equation 8 predicts the “ edge of stability ” is dependent on the batch size and that at small batch sizes this can be will below the 2/η bound . We confirm this prediction holds even early in training ( see Appendix Figure 16 ) . Lewkowycz et al . ( 2020 ) proves that for single hidden layer neural networks initialized at point with λ1 > 2.0/η and trained with an MSE loss may enter a “ catapult ” regime—where the loss increases early until a flatter region of the loss surface is found , with divergence occurring in cases where λ1 greatly exceeds 2.0/η . In contrast to the simplified setting considered in Lewkowycz et al . ( 2020 ) , we find that divergence may occur even though λ1 2/η at initialization . 3 EXPERIMENTAL SETUP . We investigate models trained on several benchmarks : CIFAR-10 ( Krizhevsky , 2009 ) and ImageNet ( Russakovsky et al. , 2015 ) for image classification , LM1B ( Chelba et al. , 2013 ) for Language Modeling , and WMT for Neural Machine Translation ( NMT ) . On CIFAR-10 we consider the WideResnet ( Zagoruyko and Komodakis , 2016 ) and DenseNet ( Huang et al. , 2017 ) architectures , both with and without Batch Normalization . When training without Batch Normalization we consider several initialization strategies including the default “ LeCun Normal ” initialization , and running MetaInit . As a way to artificially induce worse initializations , we also consider experiments where we scale every variable produced by the default initialization by a constant factor α . The NMT models are trained on the WMT ’ 16 EN-DE training set , tuned for hyper-parameters on the WMT ’ 16 EN-DE validation set and evaluated on the WMT ’ 14 EN-DE test set for BLEU scores . For NMT and LM1B Language Modeling , we train 6 layer Transformer models ( Vaswani et al. , 2017 ) . Inspired from Xiong et al . ( 2020 ) , we experiment with three Layer Norm settings : pre-Layer Norm , post-Layer Norm ( Liu et al. , 2020 ) and no Layer Norm for the transformer models . Each model is trained with various learning rates using cosine decay ( unless mentioned explicitly ) . For warmup experiments we use linear warmup which starts at 0 and scales linearly to a max value η before applying cosine decay . To measure the max eigenvalue of the loss Hessian we use the Lanczos method where the number of iterations varied as needed depending on the architecture ( details provided in the appendix ) . 4 EARLY TRAINING INSTABILITY AND THE LOSS HESSIAN . In Figure 2 we plot the curvature at initialization and during training for a series of models trained on different datasets ( plots showing final performance for all models can be found in the appendix ) . Each row indicates a different base model , the left column plots the curvature of the model at initialization and indicates with an ‘ X ’ whether or not the model diverges when trained without warmup . On the 2This is a simplified , potentially loose bound . See the original work for a more general bound that depends on both the loss curvature and the noise covariance matrix . right we plot the measured curvature and learning rate at a specified point during training . We observe across all datasets that successful training occurs only when optimization enters a region of parameter space where λ1 ≤ 2/η , and that divergent models are outside this region shortly before divergence . At initialization , some models can be successfully trained even when they start out in the unstable region and generally speaking , divergence is more likely for models deeper in the unstable region . For CIFAR-10 WideResnet , removing batch norm results in a model with higher curvature at initialization and results in divergent models when trained with a learning rate η > .1 . Scaling the WideResnet initialization up by a factor of 1.5 exacerbates the problem , resulting in even higher curvature at initialization and divergence when η > 10−2 . MetaInit starts the model out at a point with very small λ1 , and allows training without Batch Normalization at higher learning rates than the default initialization . We also observed that higher learning rates can be unlocked when the models are trained with learning rate warmup . Warmup was particularly effective for models which exhibit large λ1 either at initialization or early in training . Other models such as the post activation Resnet-50 , and WideResnet w/ Batch Normalization did not benefit from warmup at the considered learning rates ( see Appendix ) . For the DenseNet experiments , it is noteworthy that the models with Batch Normalization actually start out with higher curvature than the non-BN variants . This is contrary to the generally accepted narrative that Batch Normalization improves the smoothness of the loss surface ( Ghorbani et al. , 2019 ; Santurkar et al. , 2018 ) , though not without precedent as Yang et al . ( 2019 ) demonstrates settings for which Batch Normalization can cause exploding gradients . We found that the Batch Normalization models were more unstable than the non-BN variants here , as some models diverged at smaller learning rates . However , when combined with warmup the BN models were trainable at learning rates η > .1 , whereas this did not hold for the non-BN variants , which diverge both with and without warmup at these learning rates . This result suggests that BN still offers training stability for this model , and flatter curvature mid training if trained with warmup and a higher learning rate , however no smoothness benefits are observed at initialization . See Appendix D.1 for more details on this phenomenon . For Resnet-50 trained on ImageNet we compare two different residual blocks : the preactivation block ( He et al. , 2016b ) and the more commonly used post activation block ( He et al. , 2016a ) . For the preactivation block , we also consider flipping the order of the ReLU activation and batch normalization , as was considered in Brock et al . ( 2021 ) . We find that both preactivation models start out in a region of higher curvature relative to the post activation variant , and that these models diverge when η > .5 whereas the post activation variant is trainable with learning rates as large at 10 . Notably , there are several models in our experiments which diverge despite starting out in a region where λ1 < 2/η . This occurs for both the pre and post layernorm transformer , and the WideResnet model initialized with MetaInit . We found for these divergent models that the curvature rapidly increases in the initial steps of training , which is partially visible in the mid training plot where we plot the final observed curvature before divergence . Full training curves for these models can be found in the appendix . This implies that measuring λ1 at initialization is not always sufficient to predict whether or not the model will be easily trained . Currently , some architectural innovations are motivated by an analysis of either gradient statistics or smoothness at initialization ( Liu et al. , 2020 ) —a more robust analysis would consider the evolution of these statistics under SGD .

Summary: The paper aims to understand what makes general architectures trainable, specifically what limits the maximum learning rate for deep learning models trained with SGDM, from the loss curvature aspective. They empirically studied the evolution of the loss sharpness as they vary the learning rate, warmup period, initialization, and architectural choices. They show maintaining a sufficiently small \lambda_1 is a necessary condition for successful training at a large learning rate. And different methods such as initialization, learning rate warmup, and normalization all enable higher learning rates to be used by reducing \lambda_1 during training. More specifically: - Some initialization strategies for architectures without normalization actually operate by reducing curvature early in training, enabling training at larger learning rates. - learning rate warmup gradually reduces \lambda_1 during training, which is a competitive baseline in comparison with better model initialization methods.  - large loss curvature can result in poor scaling at large batch sizes and interventions designed to improve loss conditioning can improve the model’s ability to leverage data parallelism.