A Loss Curvature Perspective on Training Instabilities of Deep Learning Models
Authors: Justin Gilmer
Behrooz Ghorbani
Ankush Garg
Sneha Kudugunta
Behnam Neyshabur
David Cardoze
George Edward Dahl
Zachary Nado
Orhan Firat
Venue: ICLR 2022 Poster (2022)
Keywords: Optimization
Deep Learning
Training Instability
Curvature
Loss Landscape
Hessian

Abstract: In this work, we study the evolution of the loss Hessian across many classification tasks in order to understand the effect the curvature of the loss has on the training dynamics. Whereas prior work has focused on how different learning rates affect the loss Hessian observed during training, we also analyze the effects of model initialization, architectural choices, and common training heuristics such as gradient clipping and learning rate warmup. Our results demonstrate that successful model and hyperparameter choices allow the early optimization trajectory to either avoid---or navigate out of---regions of high curvature and into flatter regions that tolerate a higher learning rate. Our results suggest a unifying perspective on how disparate mitigation strategies for training instability ultimately address the same underlying failure mode of neural network optimization, namely poor conditioning. Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, MetaInit, GradInit, and Fixup initialization.

Published as a conference paper at ICLR 2022
A LOSS CURVATURE PERSPECTIVE ON TRAINING IN- STABILITY IN DEEP LEARNING
Justin Gilmer
Behrooz Ghorbani
Ankush Garg
Sneha Kudugunta
Behnam Neyshabur
David Cardoze
George E. Dahl
Zack Nado
Orhan Firat
ABSTRACT
In this work, we study the evolution of the loss Hessian across many classification tasks in order to understand the effect the curvature of the loss has the training dynamics. Whereas prior work has focused on how different learning rates affect the loss Hessian observed during training, also analyze the effects of model initialization, architectural choices, and common training heuristics such as gradient clipping and learning rate warmup. Our results demonstrate that successful model and hyperparameter choices allow the early optimization trajectory to either avoid or navigate out of oe-regions of high curvature and into flatter regions that tolerate a higher learning rate. Our results suggest unifying perspective on how disparate mitigation strategies for training instability ultimately address the same underlying failure mode of neural network optimization, namely poor conditioning. Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, Metalnit, GradInit, and Fixup initialization.
INTRODUCTION
Optimization of neural networks easily fail. While recent architectural advances such skip connections (He et al., (2166a) and Batch Normalization (loffe and Szegedy 2015) have been applied successfully to produce architectures and hyperparameters that reliably train well, even small changes to trainable configure 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 new domain can be difficult coapuuutttpptttineeettte and researchers have few reliable heuristics to guide them through the process. As result, the specific hyperparameter tuming protocol has an outsized influence on the results (Choi et al. 2019 1219/SA (Sivaprasad et al. 2020) and successes often rely on large hyperparameter searches (Nado et al.) 2021). Developing 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 achieve this goal, notably normalization, learning rate warmup, gradient clipping (Pascanu et al.] (203)), and better model initializations such Fixup (Zhang et al. 2019b), Metalnit (Dauphin and Schoenholz) 2019). and GradInit (Zhu et al.) 2021). While these methods are certainly not exactly equivalent, key property they all have in common that they enable training at larger learning rates when applied certain models (see for example Figure
'Equal Contribution. Correspondence (gilmer, ghorbani) ) @ @gongle.com.

---

Published as a conference paper at ICLR 2022
0.7
CIFAR-10 WideResnet
10
Hessian Max Eigenvalue
Divergent Training: Max Learning Rate 0.1, No BatchNorm
10
2.0 /Lcaming Rate!
0.6
10
10
Final Training Loss
0.5
10a
10a
0.4
Globel Step
0.3
Stable Training: Same Setting as Above Plus 1930 Steps Warm-up
10a
Hossian Max Eigenvalue
0.2
2.0 (Learning Rate)
Batchloom
10a
0.1
10²
Beechkom
0.0
10
10-2
10
10 10a 10a
10a
10a
Learning Rate
Global Step
Figure 1: Left: Three different methods which when applied WideResnet 28-10 architecture (w/o Batch Normalization) enable training at larger learning rates: learning rate warmup, Metalnit, and adding normalization layers. Each point reports the final training loss after training with cosine decay for 300 epochs. The test performance of the models closely matches the training loss behavior (see Figure of the appendix). Right: The evolution of the largest eigenvalue of the Hessian throughout the training for the No- No-BatchNorm model with and without warm-up.
A natural hypothesis is that methods which enable training at larger learning rates do so through reducing the sharpness 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 al., 2018) and residual connections (Li al.) 2017). and quadratic models the loss surface predict that optimization with SGD is unstable when À > 2/11 (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 al. 2021 Jastrzebski al. 2020).
In this work, we design a series of large scale experiments studying the evolution of the loss sharpness we vary the learning rate, warmup period, initialization, and architectural choices. Our results demonstrate the central role that x plays in neural network optimization antmiizaion-mmainaaiinngg sufficiently small λ during optimization is necessary condition for successful training large learning rates. Consequently, reducing λ₁ 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 stable only when the optimization trajectory primarily resides a region of parameter space where A₁ < 2/n, where y denotes the learning rate. This corrob- orates the theoretical predictions of Wu et al. (2018) and recent empirical observations of Jastrzebski et al. (2220) 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 x during training, offering similar benefits to better model initialization. We connect the mechanism by which warmup operates to the dynamical stability model Wu et al. (2018).
We show that learning rate warmup is simple yet competitive baseline for research into better model initialization. We demonstrate that key progress in this area (Dauphin and Schoenholzz 2019 Zhang al. [019b]] Zhu et al. (021)) can be matched by the application of learning rate warmup and/or gradient clipping alone.
"Throughout this work will use the term sharpness to refer to the maximum eigenvalue of the loss Hessian, denoted Y₁ See Appendix B for more details.
2

---

Published as a conference paper at ICLR 2022
Finally, 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 central object of study for understanding optimization of neural networks. [Santurar et al. (2018) argues that an important benefit of Batch Normalization is improved smoothness of the loss surface, while Lewkowye 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 quance-notabl 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. WW et al. (2018) provides a dynamical stability model which predicts that the loss curvature at convergence must satisfy À S 2/n.. Recent work has provided empirical evidence
that X1 x 2/n 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/n. 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 small batch sizes this can be will below the 2/1 bound. We confirm this prediction holds even early in training (see Appendix Figure [6). Lewkowycz et al. (2020) proves that for single hidden layer neural networks initialized at point with A > 2.0/17 17 and trained with an MSE loss may enter a "'atapult" regime- regime-where -where the loss increases early until a flatter region of the loss surface found, with divergence occurring in cases where A greatly exceeds 2.0/n. In contrast to the simplified setting considered in Lewkowycz et al. (2020), we find that divergence may occur even though A < 2/1 at initialization.
EXPERIMENTAL SETUP
We investigate models trained on several benchmarks: CIFAR-10 (Krizhessky
2009 and ImageNet
(Russakovsky et al. 2015) for image classification, LMIB (Chelba et al. 2013) for Language Modeling, and WMT for Neural Machine Translation (NMT). On CIFAR- WideResnet (Zagoruyko and Komodakis 2016) and DenseNet (Huang et al. 2017) architectures, both with and without Batch Normalization. We consider two variants of the DenseNet architecture. The standard variant from the open sourced code of Zhu al. (2021) is considered in Figure and Jnnd Table([] A less stable variant changes the strides in the average pooling layers to (1,1) is used for Figure
10 we consider the
lnd and denoted as Stride-(1,1) DenseNet (see Appendix [.I for a more detailed discussion). When training without Batch Normalization we consider several initialization strategies including the default "LeCun Normal" initialization, and running Metalnit. As way to artificially induce worse initializations, we also consider experiments where we scale every variable produced by the default initialization by a constant factor a. 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 LMIB Language Modeling, we train layer Transformer models (Vaswani et al. 2017). Inspired from Xiong et al. (2020), we experiment with three Layer Norm settings: pre- 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 use linear warmup which starts 0 and scales linearly to max value v) before applying cosine decay. To measure the max eigenvalue of the loss Hessian we use the
This is simplified, potentially loose bound. See the original work for more general bound that depends on both the loss curvature and the noise covariance matrix.
3

---

Published as a conference paper at ICLR 2022
Lanezos 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 right we plot the measured curvature and learning rate at specified point during training. We observe across all datasets that successful training occurs only when optimization enters a region of parameter space where \ S 2/n,, 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 learning rate 17 > .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 7) > 10 Metalnit starts the model out at a point with very small Å₁ and allows training without Batch Normalization at higher learning rates than the default initialization. We also observed that higher learningg rates can be unlocked when the models are trained with learning rate warmup. Warmup was particularly effective for models which exhibit large a 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 Stride-(1,1) 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 let al., 2019; (annurkar et al. 2018). We found that the Batch Normalization models were more unstable than the non-BN variants here, as some models diverged smaller learning rates. However, when combined with warmup the BN models were trainable at learning rates 1) > .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 higher learning rate, however no smoothness benefits are observed at initialization. See Appendix D.I 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., (006a)). For the preactivation block, we also consider flipping the order of the ReLU activation and batch normalization, as was considered in Brock et al. (02I). 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 1) > .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 x 2/11. This occurs for both the pre and post layeroom transformer, and the WideResnet model initialized with Metalnit. 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 plot the final observed curvature before divergence. Full training curves for these models can be found in the appendix. This implies that measuring A 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 initialization (Liu et al. 2020) a more robust analysis would consider the evolution of these statistics under SGD,
5
THE INTERACTION BETWEEN LEARNING RATE WARMUP, INITIALIZATION AND CURVATURE
The success of learning rate warmup is inconsistent with conventional optimization wisdom, which traditionally suggests adapting the step size the curvature (see for example the discussion around
4

---

Published as a conference paper at ICLR 2022
Resnet-50 Imagenet
102
Prx Achvatin dredl 2/t.arin
10a
Eigeruslue at maalization
11008 Eigenvatue
10*
20"
102
201-
1o1
10%
10a
10
10
10
Step 11000 Learning
Stride (1,1) DenseNet Cifar- -0
10%
10a
123 uams 31131
Eigenvalue Initaalration
204
LG IS
10
X
102-
13 Lave Co ORtEcC
2000
101
10-1
10
10
10
100
10
10
10"
Step 2000 Leamnng Rate
Transformer LM1B
Mox Elgenvelus atiiiaiizaaion
Uneigad gagggiiioo varigg
Stap 60000 Max Eigervatue
sata
202
10²
101
101
109
10 10a
10-2
10-3
10-2
10-1 Step 60000 Learning
10"
Peak Learning
Figure 2: Measurements of the Hessian max eigenvalue different models initialization (left) and during training (right). Divergent models are indicated with an X. On the left, we plot A at initialization along with the peak learning used during training. On the right, we plot and À at specified step in training. For divergent models on the right plot, we record the last learning rate and max eigenvalue that occur before divergence is detected (defined as observing a NaN in the loss). Across all datasets and models, successful training occurs only when optimization enters stable region of parameter space where A S 2/1j. Models with higher initial loss curvature tend to diverge smaller learning rates relative models initialized with flatter curvature, however larger learning rates are possible in these cases when warmup is used. Note, to avoid overlapping points we have applied small deterministic shift the x-position of points the right hand plots. In the case of Stride-(1.) DenseNets w/ BN we observed some models exhibit catapult behavior (loss increases early followed by normal training), these have been marked accordingly.
5

---

Published as a conference paper at ICLR 2022
Wideelenmel Cilan-10
107
30"
10"
201
10
Stro 2000 HoN 110
101
10
X
20'
10t
20
D
Stap 5500
A
B
Hesstan Eigenxalus
10"
10
162
10%
D
Gibbel Step
Clobal Siez
Figure 3: Top row: Measured À at initialization (left) and mid training (right) for variants of the WideResnet 28- 10 model on C0f Both BatchNorm and Metalnit reduce è initialization. Bottom Row: Full time evolution of λ and for select runs (labeled with letters in the top row). A: The non-BN variant diverges when trained at r = .1. B: Same peak as (A), but when warmup is used Å gradually decreases without diverging. C: Metalnit allows training without warmup v by starting x small. D: An example of model diverging despite starting out in the "stable" region. E: An example of progressive sharpening when training with small 17. F: Learning rate warmup recovers from an even poorer initialization.
equation 2.4 in McCandlish et al. (2018). However, with the understanding that λ is dynamic quantity whose evolution is tightly coupled with the learning rate schedule, the benefits of warmup period are more easily understood. We argue that the success of learning rate warmup follows naturally from two properties of training deep models:
1. Models diverge when the learning rate is too large relative to the 2/ bound.
2. When the learning rate only slightly exceeds 2/ A optimization is unstable until the parame- ters move to region with smaller A (Wu et al. 2018 Lewkowycz et al., 2020).
The first criteria implies that can't start off at too large of a value relative X at initialization. The second criteria implies that gradually increasing 1) can gradually "push" the parameters to region of parameter space where optimization is stable (with lower values of A1). In Figure there is clear evidence for this "pushing". as during the warmup period the we see that \ su 2.0/n holds for large part of the warmup phase. Furthermore, this approximation holds even as we vary the length of the warmup period. Other examples can be seen in Figure and F), and Figure [5] in the appendix.
Warmup is not the only method capable of reducing A during training, one can instead initialize the model in region where A starts off small. Consider for example, the points A, B and C in Figure 3 Each point shows optimization of non-BN WideResnet with peak learning rate of .1. (A) we see the model diverges within 3 steps without warmup using the default initialization. In (B) we see that a linear warmup period results in a₁ progressively decreasing until the peak step size of is reached step 1000, with divergence occurring. Finally in (C) initialize the same model with Metalnit, at which point A is small at initialization, and the model can be trained at .1 without warmup.
Similar to the aforementioned Metalnit, the success of related initialization strategies can be explained by reduced A₁ early in training. In Figure 5 (left) (left) we look at the evolution of X during the GradInit meta optimization process and compare this with simply training the same model using gradient clipping Both methods result in Xj decreasing dramatically, after which x hovers around 2/n,
Similar to warmup, gradient clipping reduces the step size in regions of large curvature.

---

Published as a conference paper at ICLR 2022
6-Layer Trersformer, WMT
ImaggeNet
Warmup
6:00 Steps
1000 Warmup Steps
40000 Steps
sco0 Warmup Steps
Max Nession big. Anw
40000 Staps
107
10000 War Warmup Steps
2.0
100
10-
10a
10a
109 10"
10"1
10a
10a
10
10
102
10a
10a
10a
10'
10*
Training Steps
Training Steps
Figure 4: Learning rate warmup "pushes" the optimization trajectory towards regions of flat- curvature. Solid lines correspond to the maximum eigenvalue the Hessian throughout the training. Dashed lines correspond to 2/1. During the warmup period, A fluctuates close to the 2/n bound, The WMT models diverge for 0, 400 and 4000 warmup steps. The pre-activation Resnet-50 model diverges quickly without warmup.
Table 1: Learning rate warmup can match the performance of recent advances in initialization research. The non-warmup results were all taken from their respective original works, Metalnit from Dauphin and Schoenholz (2019), GradInit from et al. (202T) and Fixup from (Zhang et al. (2019b). We found all cases that warmup and gradient clipping could leveraged to recover from bad initializations, offering similar benefits to these better initialization strategies. The DenseNet-100 Kaiming-Cui Clip 1.0 result is taken directly from Zhu et al. (02I), we found using the same codebase that changing the gradient clipping 6.0 was enough to closely match the GradInit results. For the 6-layer Transformer w/o LayerNorm we found the default initialization to exhibit extreme curvature, with initial gradient for any training to occur at this point, however scaling the default initialization down by a factor was enough recover the performance the normalized model.
the order of 1020 We found that warmup insufficient in this
Model
Dataset
Method
Ace
10n
CIFAR-10 10
Kaiming-Cuin 1.0
93.97
DenseNet-100
CIFAR-10
GradInit
94.85
DenseNet-100
CIFAR-10
Kaiming-Clin 6.0 (ours)
94.65
WideResnet 28-10 (w/o BN)
CIFAR-10
Warmup 1000
97.2
WideResnet 28-10 (w/o BN)
CIFAR-10
Metalnit
97.1
Resnet-50 (w/o BN)
ImageNet
Fixup
76.0
Resnet-50 (w/o BN)
ImageNet
Metalnit
76.0
Resnet-50 (w/o BN)
ImageNet
GradInit
76.2
Resnet-50 (w/o BN)
ImageNet
Warmup 1000 (ours)
76.2
Transformer 6L (w/o LN)
WMT
Warmup Clip + .25x Init (ours)
27.10 (BLEU)
Transformer 6L (w/ LN)
WMT
LayerNorm
27.01 (BLEU)
Notably, GradInit starts regular training off at A significantly below the 2/n bound, however the curvature quickly increases within a few steps. Given that initialization and warmup serve similar roles in reducing Å₁ we expect to be able achieve similar performance using the two methods. As shown in Table I we can easily match key advances in this field by applying learning rate warmup alone (see Appendix for experimental details).
Beyond controlling A mid-training, the learning rate 7 controls more general conditioning measures of the loss surface. For example, in Figure we observe that even the Metalnit gradient quotient- the conditioning measure directly optimized by this initialization strategy is controlled by mid training. This again provides further evidence that the primary benefit of this initialization method is to reduce X at initialization. As shown, any gains by optimizing the more general gradient quotient must be short lived as the initialization has control over the long term value.
7

---

Published as a conference paper at ICLR 2022
Training with Gradient Clioping
Gradinit Ncta Optimization
10
Training after Gradint
(Learning
Hessin Max Elgervelue
107
10²
Metaloss
100
10
Warmup
Training Warmup Step
10a
Training after Mezainit Step 0
Trainnng after Metaintt int Step 2000
100
25
50
75
100
125
150
175
200
10
10-2 Pesk Learning Rate
10
10
Step
Figure 5: The mid-training conditioning is determined by the learning rate, not on the initial- ization method used. (a): We plot the Hessian max eigenvalue during training for two models, DenseNet model trained with learning rate warmup and gradient clipping, and the same model initial- ized with GradInit. We also plot the Hessian max eigenvalue during the GradInit meta optimization process. Both GradInit and training with warmup are able to reduce the large curvature of the model. The GradInit algorithm initializes SGD flatter region, however this effect only last for a few steps of SGD, at which point the curvature of both the GradInit model and the model trained with warmup are nearly identical. (b): We plot the Metalnit MetaLoss during training for two groups of models, ones initialized with Metalnit and ones trained with learning rate warmup. Although the two groups start training with dramatically different values of the MetaLoss, after 3000 steps of SGD the MetaLoss of the two groups of models is almost completely determined by the learning rate used in training, not on the model initialization.
THE EFFECTS OF CURVATURE ON BATCH SIZE SCALING
So far, we discussed how large loss curvature limits the range of stable learning rates for training. In this section, highlight how these limits usable learning rates affect the model's ability to effectively leverage larger batch sizes. Previous research has studied the interplay of the loss curvature and batch size scaling from various different perspectives. Most notably, [Shallue al. (2018) observe that increasing the batch size yields consistent improvements in training speed until a (problem-dependeen) critical batch size reached; increasing the batch size beyond this threshold yields diminishing improvements in training speed. ZZhang et al. (2019a) observe that a simple Noisy Quadratic Model (NQM) able to capture some of the empirical behavior observed in Shallue et al. (2018). Similarly, McCandlish et al. (2018) use quadratic approximations to the loss to provide closed form expression for the critical batch size as a function of the loss Hessian and the covariance of the stochastic gradient. We contribute to this literature by highlighting the role of X in the batch size scaling behavior of the model.
For this analysis. we focus three of the WideResnet variants considered in Figure the Batch- Norm model low curvature model), the non BatchNorm model (with moderate curvature), and the non BatchNorm model with 1.5X init scaling (with high curvature). We train these models I7 while sweeping both the learning rate and the batch size We then measure the number of training steps required reach 85% validation accuracy, and the optimal learning rate found for each batch size. Similar to Shalllu et al. (018)) we normalize the plotted steps to 85% accuracy by the value measured at batch size 64.
The results are shown in Figure 6 A few observations are in order: The low curvature model shows almost linear speedups in training speed the batch size increases. In contrast, the high curvature model exhibits only minimal improvements in training speed with larger batch sizes. These scaling differences are closely mirrored by how the optimal learning rate 7)" changes with the batch size: for the low curvature model ry increases linearly with the batch size, while for the high curvature model 17" fixed around 10 Notably, for the high curvature model 77 is almost always the
We sweep for the optimal learning rate on log-scale grid between 10 and 1. For batch size, sweep over powers 2 from 16 to 4096.
8

---

Published as a conference paper at ICLR 2022
largest non-divergent value clear indication that the high loss curvature slows down training by preventing larger values from being used.
A clear picture emerges from these observations. Previous research suggests that in order to effectively leverage larger batch sizes, one has to increase the learning rate in tandem with the batch size Jastrzebski et al. (2017); Goyal al. (2017)| [Shallue et al. (2018); McCandlish (2088). Our results suggest that large values of X place a sharp limit on the maximum the learning rate possible and therefore, limit the model's ability to leverage data parallelism effectively.
1.500
B.8002
DM
2.700
0.2414
Srrtt Radt
2.385
0.0125
(13tpp
0.195
0.0220
0.000
0.0063
Batch Size
32
Stapi to 55% AcCurNY
35400
Stepe to 85 AcTurNC
40393
11682
12965
2514
3249
22342
1203
35134
2343
356
10:4
30253
Figure 6: Large curvature models scale poorly with batch size. The four plots explore how the WideResnet models from Figure 2 scale with batch size. We look at three models, a low curvature model (WideResnet with BatchNorm) medium curvature model (WideResnet without BatchNorm), and a high curvature model (WideResnet w/o BatchNorm and 1.5 init scaling). The low curvature model exhibits almost perfect (linear) scaling the batch size increases, with the optimal learning rate increasing almost linearly with the batch size. The high curvature model shows almost no speedups at larger batch sizes, with the optimal learning rate fixed at the largest value with does not diverge. Top left: Steps required for each model to reach 85% accuracy, normalized by the steps required at batch size 64. Top Right: Optimal learning rate for each batch size. Bottom Left: Steps to 85% accuracy for each learning rate, broken down by the batch size for the BatchNorm model. Bottom Right: Steps 85% accuracy for the non BatchNorm model with 1.5X init scaling.
CONCLUSION
Through extensive empirical experiments measuring the evolution of the loss sharpness during training, we have demonstrated how different methods such as initialization, learning rate warmup, and normalization enable higher learning rates to be used (without causing divergence) by reducing X during training. It is noteworthy that two of the most popular models investigated (the popular nost-activation variant of the Resnet-50 and the standard WideResnet 28-10) 10) did not benefit from learning rate warmup, and exhibited small values of X throughout training the learning rates considered. Thus researchers and practitioners who primarily work with well-optimized architectures might never notice benefit from using warmup. However, even seemingly trivial modifications to working architecture can casily result in large values of x and thus instability early in training a naive response such a situation would be dramatically reduce the learning rate or, even worse, abandon the modification being investigated all together. We hope the perspective presented in this work can help future researchers better navigate such situations, either through investigating different initializations, applying warmup and gradient clipping, or changing the location of normalization layers in the model.

---

Published as a conference paper at ICLR 2022
REFERENCES
S. Boyd, S. P. Boyd, and I.. Vandenberghe. Convex optimization. Cambridge university press, 2004.
A. Brock, S. De, and S. Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets. arXiv preprint arXiv:2 arXiv:2101.08692, 08692, 2021.
C. Chelba, T. Mikolov, M. Schuster, Q. Ge, T. Brants, and P. Koehn. One billion word benchmark for measuring progress in statistical language modeling. INTERSPEECH, 2013.
D. Choi, Shallue, Z. Nado, J. Lee, C. J. Maddison, and G. E. Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint 19:.v////000050444 2019.
J. M. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. arXiv preprint arXiv: 203.: 000600 2021.
Y. N. Dauphin and S. Schoenholz. Metainit: Initializing learning by learning to initialize. Advances in Neural Information Processing Systems, 32:12645-122 12657, 2019.
B. Ghorbani. S. Krishnan, and Y. Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2241. PMLR, 2019.
P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 7706002677 2017.
K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 770-778, 2016a.
K. He, X. Zhang, S. Ren, and Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645. Springer, 2016b.
R. A. Horn and C. R. Johnson. Matrix analysis. Cambridge university press, 2012.
G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708. 2017.
S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. PMLR, 2015.
S. Jastrzebski, 7. Kenton, D. Arpit, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv: 2017.
S. Jastrzebski, M. Szymczak, S. Fort, D. Arpit, J. Tabor, K. Cho, and K. Geras. The break-even point optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09572 2020.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint 14X::: 2014.
A. Krizhevsky. Learning multiple layers features from tiny images. Technical report, University of Toronto, 2009.
A. Lewkowycz, Y. Bahri, E. Dyer, J. Sohl-Dickstein, and G. Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/6 2020.
H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein. Visualizing the loss landscape of neural nets. arXiv preprini arXiv: 1712:07122.09111 2017.
L. Liu, X. Liu, J. Gao, W. Chen, and J. Han. Understanding the difficulty of training transformers. arXiv preprint 2=Xiv:2004.08249, 2020.
McCandlish, J. Kaplan, D. Amodei, and O. D. Team. An empirical model of large-batch training. arXiv preprint arXiv: (a/i.0///200660 2018.
10

---

Published as a conference paper at ICLR 2022
Z. Nado, J. Gilmer, C. J. Shallue, R. Anil, and G. E. Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. CoRR, abs/2102.06356, 2021. URL htttss://arxis org/ httpg//b//////////////////////// 2102. 06356
V. Papyan. The full spectrum of deepnet hessians scale: Dynamics with sgd training and sample size. arXiv preprini arXiv: /81 1. 1.07062, 2018.
R. Pascanu, T. Mikolov, and Y. Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318. PMLR, 2013.
B. A. Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160, 1994.
O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 1 (15(3)221-2555 2015.
S. Santurkar, D. Tsipras, A. Ilyas, and A. Madry. How does batch normalization help optimization? arXiv preprint arXiv: 1805. 1604, 2018.
C. J. Shallue, J. Lee, J. Antognini, J. Sohl-Dickstein, R. Frostig, and G. E. Dahl. Measuring the effects of data parallelism neural network training. arXiv preprint arXi::1811.03600 2018.
P. T. Sivaprasad, F. Mai, T. Vogels, M. Jaggi, and F. Fleuret. Optimizer benchmarking needs to account for hyperparameter tuning. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9036-9045. PMLR, 13-18 Jul 2020. URL http:://proceeiingg mlr ppress/ul1S/iivazzssa222 0hmmmlla...mmmmmmmmmmmmmm
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polo- sukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017. URL
htw.:roocceeiinn............ neurips cc/ cc/paper//011/fi1ee Srtee2225777ae110000333ccaaaaaa2
L. Wu, C. Ma, and W. E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 8289-8298, 2018
R. Xiong, Y. Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 10524-10533. PMLR, 2020.
G. Yang, J. Pennington, V. Rao, Sohl-Dickstein, and S. S. Schoenholz. mean field theory of batch normalization. arXiv preprint (00.:v:902.00/222.0 2019.
S. Zagoruyko and N. Komodakis. Wide residual networks. arXiv preprint 1xXvv:1605.07446 2016.
G. Zhang, L. Li, Z. Nado, J. Martens, S. Sachdeva, G. Dahl, C. Shallue, and R. B. Grosse. Which algorithmic choices matter which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32:81 12:8196-8207, 2019a.
H. Zhang, M. Cisse, N. Dauphin, and D. Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:
2017.
H. Zhang, Y. N. Dauphin, and T. Ma. Fixup initialization: Residual learning without normalization. arXiv preprint 1rXiv: 000021, 2019b.
C. Zhu, R. Ni, Z. Xu, Kong, W. R. Huang, and Goldstein. Gradinit: Learning to initialize neural networks for stable and efficient training. arXiv preprint
2021.
11

---

Published as a conference paper at ICLR 2022
A LIMITATIONS
Our analysis has focused primarily on models trained with SGD with momentum. This decision was motivated reduce additional confounds that arise when using adaptive preconditioning. Notably, it is unclear what the analogue of X I 2/1 should be for model trained with Adam. In the appendix, we provide evidence that loss curvature adaption to the learning rate does occur even for Transformer models trained with Adam, and that learning rate warmup results in the similar effect of the optimization trajectory being "pushed" to flatter regions. However we leave a deeper analysis into this for future work.
Finally, while our experiments certainly raise questions about the efficacy better model initialization has on further accelerating training. measurements has focused primarily the (lack of) influence initialization has on À1 mid training. It is possible that better initializations could have lasting influence the broader Hessian eigenspectrum (for example improving the ration As/Ar A for smaller eigenvalues ) and that our analysis is missing such an effect.
B BRIEF REVIEW OF THE LOSS HESSIAN, EIGENVALUES AND QUADRATIC STABILITY BOUNDS
For completeness, we include a formal definition of the fundamental mathematical quantities dis- cussed in the paper. We derive most of this discussion from the relevant chapters Horn and Johnson (2012) and Boyd et al. 2004). We refer the reader to these sources for a more detailed discussion.
B. THE HESSIAN MATRIX
The second derivative or Hessian matrix of the loss function L(-) at a point 0 E R" is denoted by H (0) _=x= where V1 <i Sn
LL)
H());,, = 00,08,
(1)
Moreover, by Schwarz's theorem, if the second partial derivatives of L(-) are continuous at 0, the matrix H (0) symmetric. This a broad condition that holds for the loss surfaces we examine in the main text (beyond set of measure zero).
B.2 EIGENVALUES
Definition Let A p If scalar À and a nonzero vector satisfy
Ar= Ar, X C, C" 0
(2)
then X is called an eigenvalue of A and is called an eigenvector of associated with X.
symmetric real matrix (such as the Hessian matrix), can be factored as
A = QAQ""
(3)
where Q E Ruxn is an orthogonal matrix and A matrix. Here, {A,)%1 are all of the eigenvalues of A. We order ); such that A 2 Az... 2 An. In this ordering, Å corresponds to the maximum eigenvalue of A and Xn corresponds to its minimum eigenvalue.
E Rax is real diagonal
The maximum and minimum eigenvalues of A satisfy the following important properties:
xJ AAr
TT Axx
sup I
An inf 270
x x
x²xx
In particular, for any R", we have
An |100||
< xx
Ax
X 1|
12

---

Published as a conference paper at ICLR 2022
B.3 STABILITY OF GRADIENT DESCENT FOR QUADRATIC LOSS
Now that we have established the basics, let's derive the stability condition for GD applied to quadratic loss function. Note that Wu al. (2018) and Cohen al. (2021) provide general bounds for the stability of SGD-type optimization algorithms. Here, we state & derive the stability condition for GD for the sake of completeness.
Let
((0) 6" H0,
2
where is symmetric matrix with non-negative eigenvalues. Let's consider GD dynamics starting from a random point Oo. Under GD with a fixed step-size 1) 0, we have
0 = 0, L(,,) = 0g nH0, = (1 -1
Continuing this iteration to step yields
= - -
(5)
As t oc, this iteration is stable if the eigenvalues of (I nH) have absolute magnitude bounded by one, which can be stated equivalently as
2 A1
which is the exact condition discussed and explored in the main text.
"As Oo randomly chosen, assume has non-zero overlap with all eigenvectors.
13

---

Published as a conference paper at ICLR 2022
C MISCELLANEOUS FIGURES
0.7
No BatchNorm
Warmup
0.6
Metalnit
0.5
Test Los
0.3
0.4
0.2
0.3
BatchNorm
Warmup
0.1
Metalnit
BetchNorm
10
10-2
Max Learning Rate
Max Learning Rate
Figure 7: The test-time behavior of the different WideResnet models closely mirrors their training loss dynamics (Figure left).
L-yye Post Transformer, Trained with Momentum on WMT
2/r (dashed) !H!
Evolution Training Loss
No warm-up
No warm-up
4000
Anar(M)
80
60
10
40
20
Ad
10a
101
10'
10
10
10
107
10*
Training Steps
Training Steps
Figure The curvature at initialization can be unreliable for predicting training stability. Left: Maximum eigenvalues of the Hessian throughout training. Dashed line correspond to 2/1,. Right: Evolution the training loss for different models. The model with no warm-up (red) suffers from training instabilities even though it is below the stability threshold initialization.
D PERFORMANCE OF MODELS FIGURE 2
In this this section, plot the performance vs learning rate for all of the models shown in Figure of the main text. These are shown in Figures 9 [3 10 and IT For models which diverged, we plot the best test performance achieved before divergence. In all settings, high curvature affects the final performance by limiting the use of higher learning rates.
We also noted several models in Figure which diverged despite training starting out in stable region of parameter space. In Figure 12 we plot the evolution the loss sharpness during training, showing that it quickly enters region where \ 2.0/n before diverging around step 90.
14

---

Published as a conference paper at ICLR 2022
dicehannd Clacip Verietaa Mithout Warmas
Wedellazra Cite Vananta
Beirklere
Figure 9: Performance of the WideResnet 28-10 models trained Cifar10 shown in Figure of the main text. (left) Performance of models trained without learning rate warmup. (right) Performance with learning rate warmup. When warmup is applied, all models reach comparable performance (though additional regularization is needed to match the generalization of the Batch Normalized models).
Fra Activaaion mek 25
Fia Acth BAL--33
Stargee
Port
Port Loantigg
Figure 10: Performance the Resnet50 models trained on Imagenet shown in Figure 2 of the main text. (left) Performance of models trained without learning rate warmup. (right) Performance with learning rate warmup.
Oitive Shomes
Traneberee dborea
Past Nar--aa Tianaterre Learning
Loarning
Pert
Figure 11: Performance of the Transformer models trained on LMIB shown in Figure 2 of the main text. (left) Performance of models trained without learning rate warmup. (right) Performance with learning rate warmup.
15

---

Published as a conference paper at ICLR 2022
WideResnet Cifar10 with Metainit Initialization
Hessian Max Eigenvalue
2.0 (Learning Rate)
10a
10a
10
20
30
40
50
70
80
Step
Figure 12: Training curve of the WideResnet 28-10 which diverged despite starting in region with A < 2.0/7. The parameters quickly leave the stable region and end up in region of higher curvature before diverging.
16

---

Published as a conference paper at ICLR 2022
D.1 DISCUSSION OF STRIDE-(L,) (1,1) DENSENET EXPERIMENTS
In this section we discuss in more detail the Stride-(1,1) DenseNet experiments shown in Figure 2 in the main text. These experiments use a non-standard version of the DenseNet architecture where all average pooing strides are set (1.1). Note the experiments in Figrue [ and Table I Tinstead use the standard strides implementation from the open sourced code of Zhu et al. (202)) The Stride-(1.1) DenseNet architecture noteworthy because it is counter example to common intuition that adding Batch Normalization results in flatter curvature. As shown in Figure 2 (left), the BN variants all have high curvature at initialization, however the right hand side plot shows that the mid training curvature becomes comparable to the non-BN variants. In Figure 13 we provide a more detailed analysis to understand what is happening with BN.
First we plot the performance of the BN vs non-BN models both with and without warmup. The differences are striking. Without warmup, we see the BN performance is highly stochastic, some trials outperform the non-BN variants, while some trials underperform the non-BN variants. However, when trained with 1000 steps of warmup the BN variants now significantly outperform the non-BN models are all considered learning rates. They can even be successfully trained at higher learning rates than the non-BN variants, despite the high initial curvature. To provide further detail, we show the training curves of select individual runs, both the evolution of the training loss and the evolution of curvature. The BN variants all exhibit catapult behavior early- the loss increases initially until the parameters enter region of flatter curvature. Warmup helps the BN variants, and significantly reduces the severity of the catapult phase while enabling faster long term training. Additionally, when we add warmup we find that the BN variants can now be trained at higher learning rates than the non-BN variant. As shown at learning rate of .2, the non-BN BN model diverges during the warmup phase despite lower initial curvature.
Based on these experiments arrive at the following conclusions. First, adding BN to the Stride- (1,1) DenseNet architecture results in high curvature at initialization, which results in short period of instability during training. However, once the parameters escape this region of large curvature, the BN variant exhibits favorable training dynamics relative to the non-BN variants. Thus, there still seems to be benefit adding BN, assuming steps are taken to mitigate the initial period of high curvature. The fact that adding BN to a model can result in high initial curvature is not without precedent, as Yang al. (2099) observe that adding BN to deep fully connected networks can result in exploding gradients at initialization.
These experiments highlight one of the primary takeaways of this work: that maintaining flat curvature throughout training is a necessary (not sufficient) condition for stable training of neural networks. Thus it is not the presence of BN that is necessary for stable training, instead BN is generally a useful tool for reducing curvature (and thus stabilizing training). BN has clear benefits of improving curvature in most cases, but it possible to produce configurations where adding BN paradoxically results in higher initial curvature than the non-BN variant. In these cases training is initially unstable, but once the curvature is reduced we see benefits later in training of using BN.
E DETAILS ON COMPUTING THE HESSIAN EIGENSPECTRUM VIA LANCZOS
We use Lanczos iterations to estimate the top eigenvalue of the Hessian. Lanczos algorithm only requires Hessian-vector products which can be efficiently computed via Pearlmutter's trick |Pear mutter (194)). Previous research has demonstrated that this approach provides a robust and scalable framework to examine the eigenvalues of the Hessian for large neural networks Ghorbani et al. (2019
Papyan 2018).
For our WMT / LMIB experiments, we run the algorithm for 45 steps while for image models, we use 40 steps. When monitoring the evolution of the top eigenvalue as a function of the number of Lanezos steps, in all cases except one, we observe that the algorithm converges. For the case of Resnet with ReI.U-BN ordering, due to a very small eigengap between the top eigenvalue and the bulk, the convergence is significantly slower. We use 200 Lanczos steps in this case alleviate the issue. For this model, estimating λ via power iteration (as is commonly done in the deep learning literature) will incorrectly the largest negative cigenvalue, not \ as desired.
17

---

Published as a conference paper at ICLR 2022
107
Denselet Veriants Without Warmup
107
Densellet Variams With Warmuup
ayeer
ayere
Layers
Layers BN
Layers
Layers
Layers
Layers
Leaninng
dentoon
Pest: leaming
Peak Leamning Rate: 0.01 No Warmup
Layars mt
Layars SN
Layera
Leeaming Raatl
Henee Max geroooe
Layars
Peak Learning Rate:
2011 syers NJ RN
100
201 syeers 90h AN
Isee Nn RONK Elge Mali
102
Trainrg
3014
201 Layars wd
Zleamingg hanel
201 Layers
10'
10
Learning
Warmup Staos:
BN
(Lemingg Rate)
Heeesee Mox tiqenattee
201 Layers wib
204
102
130
12
107 20 10 10
Figure 13: Performance of the Stride-(1,1) DenseNet models trained Cifar10 shown in Figure of the main text. Top left: Performance models trained without learning rate warmup. The performance of BN has high stochasticity due to the large initial curvature, depending on how well the model recovers from the catapult phase. Top Right: Performance of models trained with 1000 steps of warmup. The BN variants now consistently outperform the non-BNN variants, and can even be trained at higher learning rates despite the larger initial curvature. Second Row: Training curves and evolution of > for the BN and non BN model learning rate .01 without warmup. The BN variant exhibits catapult behavior early but recovers once the curvature drops later in training. Third Row: Same as second row, but now with 1000 steps of warmup. The catapult behavior the BN model is reduced, and training is faster after the catapult phase. Bottom Row: Comparing BN with non BN at a higher learning rate .22 with 1000 steps of warmup. The BN model still has catapult behavior, but after recovering trains successfully once the curvature is reduced. The non-BN model diverges during the warmup phase just before step 1000.
It is well- well-known that Lanczos algorithm suffer from numerical instabilities caused by finite- precision arithmetic. To alleviate these issues, Lanczos vectors are stored float64 accuracy and we perform reorthogonalized each step of the algorithm.
18

---

Published as a conference paper at ICLR 2022
F TRAINING DETAILS FOR TABLE
F.1 NEURAL MACHINE TRANSLATION
Neural Machine Translation experiments are based on the Transformer models (Vaswani et al. 2017). We separate embeddings encoder and decoder, and common word piece vocabulary of size 32000. For depth, we use layers on both encoder and decoder. For width, we experiment with two models, namely Transformer-Bass and Transformer-Wide For Transformer-Base, we use word embeddings with 512 dimensions, 8 heads and 2048 feed-forward dimension. For Transformer- Wide we use word embeddings with 1024 dimensions, 16 heads and 4096 feed feed-forward dimension. The experiments reported in Figures 3 and 5 use Transformer-Base The experiments reported in Table use Transformer- Wide models trained with Adam (Kingma and Ba 2014). We sweep over warm-up, learning rate, gradient clipping and init_ scaling and optimize for validation loss to evaluate performance on test set BLEU reported in Table []] All the models are trained for 60 epochs batch size of 1024 for Transformer-Bas Base models, and batch size of 512 for Transformer-Bis models. We use dropout of 0. label smoothing of 0.1 and no weight decay for all these models.
F.2 DENSENETS
In Table T the ResNet-50 (w/ BN) architecture was trained for 100 epochs at batch size 512, with 12 regularization of 5e-5, dropout of .3. It was trained with SGD with nesterov momentum of .9 and learning rate of 2. We applied gradient clipping at global 12 norm of 5 and used linear learning rate warmup with warmup period of 1000 steps.
For Table 1, the DenseNet- 100 model was trained using the Gradinit codebase by modifying the supplied DenseNet script apply gradient clipping of norm 6 and use the default initialization instead of GradInit.
G TRAINING DETAILS FOR FIGURE 2
The WideResnet-28- 10 models trained with batch size of 1024 for 300 epochs. We applied the MixUp augmentation et al.) 2017). For learning rate warmup we used 1000 steps of linear warmup until the peak learning rate is achieved, at which point the learning rate is decayed according to the cosine schedule.
The Stride-(1.1) DenseNet models were trained with batch size of 512 using the SGD optimizer with momentum of 0.9, weight decay of Se-4, L2 regularization of le-4 and warmup of 1000 steps (for the models where warmup is used) followed by cosine decay. The models were trained for 200 epochs. For the DenseNet architecture we used growth rate of 32 and reduction of 0.5.
The Resnet-50 models were trained with batch size of 2048 using the SGD optimizer with nesterov momentum of .9. The learning rate schedule was the same as in the WideResnet case, with linear warmup of 1000 steps followed by cosine decay. We applied label smoothing of and used the standard flip plus crop for data augmentation.
The Transformer models on LMIB were trained at batch size 1024 using SGD with nesterov momentum of .9. We use embedding dimension of 512, 6 layers with 8 heads and MLP hidden dimension of 1024. The attention dropout rate was .. The learning rate schedule followed the same recipe in the Resnet cases.
H CURVATURE ADAPTATION WITH THE ADAM OPTIMIZER
The discussion the main text focused primarily on models trained with SGD and momentum. In this appendix, we briefly examine if similar conclusions hold for optimizers such Adam that use preconditioning.
is unclear a priori whether or not curvature adapation to the learning rate should occur for optimizers which apply preconditioning. However, given that Adam is diagonal preconditioner applied to non-diagonal Hessian, there may be some similar effects observed.
19

---

Published as a conference paper at ICLR 2022
Consider simple quadratic loss
L( = 0² НО, >
where optimization is performed via preconditioned gradient descent with a fixed diagonal precondi- tioning matrix D:
0: = 0:-1 7 11-1-1n-ーv-((-11
0:-1 7D (HOr-1)
(T (1-1) nn 0)
(I-nD-11)) Oo
As such, this simple model would suggest that the max eigenvalue of the following matrix may be related to training instability of models trained with Adam
mma D H) Amaz 1/2 1/3).
(6)
While (6) does not take into the account the effects of adaptive preconditioning or momentum, find some empirical evidence that this approximation provides understanding into the stability of the optimization.
Figure 14 below examines the evolution of A 1DD-1/2HD-11 112) for three Transformer models trained with Adam and different warm-up lengths. Here, D diagonal matrix with Di, = /Corrected Adam grad squared EMA + Adam e. We observe that, similar to the models trained with momentum, the maximum (preconditioned) Hessian eigenvalue adapts the warm-up schedule (green and red markers). We notice that -perhaps due to the effect of momentum or adaptive preconditioning the threshold 2/n does not seem to be aligning with the data well. Instead, empirically corrected threshold 40/1) seems to fit the data better. We observe that instabilities model training coincide exactly with A (D2
1/2) crossing the empirically corrected threshold.
These observations suggest that some of the insights discussed in the main text seem to carry to the case of adaptive optimizers. We leave further exploration of this more complex setting to future work.
6-Layer Pro- LN Transtormer, Trained with Adam on WMT
Maximum Eigenvalue Preconditioned Hessian
Evolution of the Training Loss
No Warm-up
No Warm-up
10:0
4000
Steps
Steps
Steps
9
10a
Training Loss
8
10c
5-
10a
4-
10a
10a
10 10a
10a
10
Training Steps
Training Steps
Figure 14: Instabilities in model training are reflected in the loss curvature even for models trained with Adam. Left: Maximum eigenvalues of the preconditioned Hessian throughout the training. Dashed line correspond to 40/1, Right: Evolution of the training loss for different models. Training becomes unstable exactly when the eigenvalues cross the threshold.
20

---

Published as a conference paper at ICLR 2022
Warmup Steps
10a
3000 Warmup Steps
10000 Warmup Steps
2.0/ (Learning Rate)
Hessian Max Eigenvalue
10
10
10a
109
10
10a
10a
Step
Figure 15: Further evidence of learning rate warmup "pushing" the optimization trajectory towards regions with reduced Å1. Note the rate at which a₁ changes closely matches depends on the length of the warmup. Model shown is the non-BN WideResnet (standard init) trained at batch size 2048 with peak learning rate of .1
Batch Size 32
10a
Batch Size 256
Batch Size 2048
2.0 (Learning Rate)
Hessian Max genevalue
10a
10²
10a
10a
10a
10a
10a
10a
10a
10a
Step
Figure 16: Wu et al. (0018) predicts that the stability bound will become tighter than 2/1 at smaller batch sizes. This plot confirms holds not only at convergence (as implied by their Theorem 1), but early in training as well. For small batch size (32), the mid training curvature hovers significantly below the 2.0/1) approximation. As the batch size increases to 2048, the mid training curvature is larger, approaching the 2.0/ n approximation. All curves show the WRN without batch norm, trained with the same learning schedule using 1000 steps of warmup.
21

---

Published as a conference paper at ICLR 2022
I COMPUTE RESOURCES USED
Nearly all experiments utilized the Google Cloud Platform with v2 cloud TPUs except for the following: The Figure 2 Resnet-50 and SIride-(1.1) Densenet experiments utilized the v3 cloud TPU, while the GradInit code was run on a cloud machine with a single V100 GPU. The Figure experiments were done parallel using up to 50 v2 TPUs concurrently over the period of few days. Additionally, all the Machine Translation models were trained on v3 cloud TPUs.
22


Reviews:
[{"review_id": "txj2fhH_xNj", "reviewer": "Reviewer_P6at", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Comments :\n- Results are well presented and yield a rather coherent picture. In particular, I like the intuition that warmup brings to a region of smaller curvature, which fits in well with the catapult mechanism. I also appreciate the batch size scaling experiments, which go against the idea that the only quantity that matters is the ratio of learning rate / batch size. \n- The observation that lambda_1 follows 2/lr in presence of warmup could be better connected to Cohen et al. with a sentence such as “Cohen et al. showed that lambda_1 increases up to 2/lr then plateaus at the edge of stability, in the constant lr setup : we show that this holds even in presence of scheduling”\n- The figures are extremely long to process for the various PDF readers I tried. This is likely due to an excessive number of datapoints in the logscale figures (happened to me in the past) : consider subsampling the data at large values.\n\nQuestions :\n- Fig. 3 : Could one try to understand why some of the models diverge even though eta<2/lambda_1 ?\n- Fig. 3 : Multiplying the init variance by 1.5 seems to increase the sharpness by a factor between 5 and 100. Why such a variability ? Can one predict this factor from the number of layers in the model ? \n- If increasing init variance increases sharpness, couldn’t one simply reduce the sharpness by reducing init variance, hence enabling larger learning rates ? This could be connected to the question of lazy vs feature learning regimes", "summary_of_the_paper": "This paper presents a comprehensive set of experiments showing how the maximal learning rate to train modern architectures depends on various algorithmic choices.", "summary_of_the_review": "Although the results presented are not particularly groundbreaking, the paper is rather well written and easy to follow, and gives some potentially useful insights on an important topic : learning rate scheduling. Therefore, I think this paper could be of interest to the community.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "njiO8TRXQeK", "reviewer": "Reviewer_oiwU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper presents the necessity of low curvature at the early stage of training for stable neural network training with large learning rate, some notable observations includes “the mid-training conditioning is determined by the learning rate, not on the initialization method used”; “Learning rate warmup can match the performance of recent advances in initialization research”. It seems to suggest that the weight initialization does not matter that much and warmup can be mitigated.\n\nHowever, the paper is kind of hard to read as there are too many empirical observations presented but not quite well organized. For example, the most measurement \\lamda_1 is is not clearly introduced and there is only a small footprint saying lambda_1 refers to maximum eigenvalue of the loss Hessian. The reason why 2/\\eta is chosen is also not clear.\n\nThe curvature and sharpness: the authors use the term sharpness to denote the max eigenvalues of loss Hessian to denote the curvature or the the loss surface, however, there are other sharpness/smoothness measurements such as [1, 2], which were used to denote the generalization ability of the loss surface at the end of training. It is easy to get confused with those terms for different conditions. Are there any connections? If so, why not study them as well? It would be great to know how those metrics are correlated with the curvature. If not, it would be better to make clear the difference and not use the terms interchangeably.\n\n\nThe authors have shown that many techniques actually made low curvature. However, curvature itself can not reliably determine the trainability. As noted by the authors,  some models can be successfully trained even when they start out in the unstable region and measuring \\lamda_1 at initialization is not always sufficient to predict whether or not the model will be easily trained. \n\nOn the other hand, the DenseNet experiments says that the models with Batch Normalization actually start out with higher curvature than the non-BN variants, which suggests that curvature may not be able to explain the effects of BN quite well and no smoothness benefits are observed at initialization. I would hope there are more investigations here as BN is still a critical module and the authors observation contradicts previous belief.\n\n\n[1] Keskar et al, On large-batch training for deep learning: Generalization gap and sharp minima. ICLR 2017\n[1] Jiang et al, Fantastic Generalization Measures and Where to Find Them, ICLR 2020", "summary_of_the_paper": "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:\n- Some initialization strategies for architectures without normalization actually operate by reducing curvature early in training, enabling training at larger learning rates.\n- learning rate warmup gradually reduces \\lambda_1 during training, which is a competitive baseline in comparison with better model initialization methods. \n- 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.", "summary_of_the_review": "The authors present investigations regarding the connections between curvature and many initialization methods and techniques. However, no clear conclusions are made and some investigations are not deep enough. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LxZXnbsRUEw", "reviewer": "Reviewer_RGH4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n* The paper includes quite a lot of experimental results and definitely provides valuable empirical observations (see them in the brief summary above). It gives an insightful and convincing understanding of the learning rate warmup strategy. It also points out several failure cases where good curvature does not guarantee convergence in training, which is equally valuable. \n\nHere are the questions for each section:\n\n* Section 4.\nI see slightly different max eigenvalue at initialization for same architecture with different learning rate -- is it fully due to the randomness in the initialization? Or do I miss something here?\n\n* Section 5. \nExplain MetaLoss briefly. \n\n* Section 6. \n  1. Given Figure 3, why would the authors not include MetaInit as another baseline to test batch size scaling?\n  2. In Figure 6, what learning rate is used for the Top left plot? \n I'd guess it's the optimal (peak) learning rate per batch size, but it's not fully clear. If not, please clarify and better help me understand how the learning rate is chosen; If yes, it looks like BN with batch size 128, 256, 512 has a log-linear improvement from the top left plot, but not very log-linear if looking at the bottom left plot: improvement from 128 to 256 is much larger than the improvement from 256 to 512. Similar results for NoBN 1.5x init. So I may be missing something here. ", "summary_of_the_paper": "The paper explores training instabilities of deep learning models by monitoring the max eigenvalue of the Hessian in i) different architectures; ii) different optimization tricks; iii) different stages of training. \n\nThe paper is a mix of several empirical observations: \n\n(Section 4) One of the main observations is that the necessary (not sufficient) requirement for a \"successful\" training is to have *relatively* low max eigenvalue through the training process, where *relative* is roughly determined by the inverse of the learning rate. \n\n(Section 5) Second main observation is specific to the learning rate warmup strategy: the warmup strategy pushes max eigenvalue to the boundary below $2/\\eta$, and with a better understanding of this optimization strategy, the authors show comparable performance in multiple datasets compared with other optimization tricks. \n\n(Section 6) Lastly, the authors explore batch size scaling with different scales of max eigenvalue, focusing on comparing BN (with small curvature only for that specific task), NoBN (larger curvature), NoBN 1.5xinit (highest curvature). The experiment shows for small curvature setting, both batch size and learning rate can scale quite well. ", "summary_of_the_review": "Overall, this is a good empirical paper with lots of experimental results and good insights. The learning rate warmup part is the most insightful among all. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "A-f5ca5al7u", "reviewer": "Reviewer_vdgP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper seems well written and manages I think to condense a significant amount of empirical work in the short limit we have. I think the analysis that has been done across multiple architectures and models provides a good basis for the authors claims in the text. \n\nStrengths:\n* The paper contains a very solid set of large scale and in-depth evaluation, which is something well needed for this kind of empirical work.\n* The different analysis done well presents the vast amount of results and the key takeaways from the experiments\n\nWeaknesses: \n * Since we do observe models that diverge despite beginning in stable regions (Fig. 3 D) it seems that there the initial warnup has a significant effect of where it takes the model beyond its initialization, which is not too well discussed.\n * In the Dense Net experiment in Table 1, there is no reported results for including warmup\n\n\n", "summary_of_the_paper": "The paper performs extensive empirical investigation into factors that play important role in making neural network trainable. One of the key quantity that the authors find to give significant insights is the largest eigenvalue of the Hessian matrix. They show that models that train successfully tend to have learning rates close to the well known critical bound for GD - 2 / \\lambda. Further investigation to various architectural choices - such as normalizations and initialization techniques indicate that the value of \\lambda behave somewhat consistently after the initial training period for well training models. In addition, they demonstrate that learning rate wramup can provide a significant improvement by pushing parameters in regions with lower values of \\lambda. ", "summary_of_the_review": "Strong empirical paper which well presents many experiments and observations of the interplay of the maximum eigenvalue of the Hessian and learning dynamics.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a thorough study of the evolution of Hessian depending on a wide variety of aspects such as initialization, architectural choices, and common training heuristics. The paper makes a number of interesting observations. Some of them are not really new but overall, the experimental evaluation of the paper makes it a valuable resource for the community.

The reviewers are overall quite positive. One reviewer notes that more investigation of the behavior of batch-normalization is required. I encourage the author to address this concern in the final manuscript. There is a lot of recent work on batch-normalization that might be worth discussing, e.g.:
Training BatchNorm and Only BatchNorm: On the Expressive Power of Random Features in CNNs
Jonathan Frankle, David J. Schwab, Ari S. Morcos

Batch Normalization Provably Avoids Rank Collapse for Randomly Initialised Deep Networks
Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, Aurelien Lucchi

A Quantitative Analysis of the Effect of Batch Normalization on Gradient Descent
Yongqiang Cai, Qianxiao Li, Zuowei Shen
Decision: Accept (Poster)