Original: A L OSS C URVATURE P ERSPECTIVE ON T RAINING I NSTABILITY IN D EEP L EARNING
Revised: A L OSS C URVATURE P ERSPECTIVE ON T RAINING I NSTABILITY IN D EEP L EARNING

Original: Anonymous authors Paper under double-blind review
Revised: Firat

Original: A BSTRACT
Revised: A BSTRACT

Original: In this work, we study the evolution of the loss Hessian across many classiﬁcation tasks in order to understand the effect the curvature of the loss has on the training dynamics.
Revised: In this work, we study the evolution of the loss Hessian across many classiﬁcation tasks in order to understand the effect the curvature of the loss has on the training dynamics.

Original: 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.
Revised: 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.

Original: 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 ﬂatter regions that tolerate a higher learning rate.
Revised: 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 ﬂatter regions that tolerate a higher learning rate.

Original: 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.
Revised: 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.

Original: Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, MetaInit,
Revised: Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, MetaInit,

Original: GradInit and Fixup initialization
Revised: GradInit, and Fixup initialization.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Optimization of neural networks can easily fail.
Revised: Optimization of neural networks can easily fail.

Original: 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 conﬁguration can easily result in training that diverges.
Revised: 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 conﬁguration can easily result in training that diverges.

Original: More generally, producing a conﬁguration that strikes the right balance between stable training and rapid optimization progress on a new domain can be difﬁcult—practitioners and researchers have few reliable heuristics to guide them through the process.
Revised: More generally, producing a conﬁguration that strikes the right balance between stable training and rapid optimization progress on a new domain can be difﬁcult—practitioners and researchers have few reliable heuristics to guide them through the process.

Original: As a result, the speciﬁc hyperparameter tuning protocol has an outsized inﬂuence on the results (Choi et al., 2019; Sivaprasad et al., 2020) and successes often rely on large hyperparameter searches (Nado et al., 2021).
Revised: As a result, the speciﬁc hyperparameter tuning protocol has an outsized inﬂuence on the results (Choi et al., 2019; Sivaprasad et al., 2020) and successes often rely on large hyperparameter searches (Nado et al., 2021).

Original: 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 ﬁnding better, more scalable architectures.
Revised: 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 ﬁnding better, more scalable architectures.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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). 
Revised: 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). [Figure]

Original: A natural hypothesis is that methods which enable training at larger learning rates do so through reducing the sharpness 1 of the loss surface during training.
Revised: A natural hypothesis is that methods which enable training at larger learning rates do so through reducing the sharpness 1 of the loss surface during training.

Original: Indeed, this hypothesis has already been proposed as one of the beneﬁcial 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
Revised: Indeed, this hypothesis has already been proposed as one of the beneﬁcial effects of Batch Normalization (Ghorbani et al., 2019; Santurkar et al., 2018) and residual connections (Li et al., 2017)

Original: [Figure]
Revised: [Figure]

Original:  predict that optimization with SGD is unstable when λ 1 > 2 /η
Revised: and quadratic models of the loss surface predict that optimization with SGD is unstable when λ 1 > 2 /η

Original: (Wu et al., 2018).
Revised: (Wu et al., 2018).

Original: 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
Revised: 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

Original: [Figure]
Revised: [Figure]

Original: and do not investigate connections between sharpness, model initialization and learning rate warmup (Cohen et al., 2021; Jastrzebski et al., 2020).In this work
Revised: and do not investigate connections between sharpness, model initialization and learning rate warmup (Cohen et al., 2021; Jastrzebski et al., 2020).

Original: [Figure] we design a series of large scale experiments studying the evolution of the loss sharpness as we vary the learning rate, warmup period, initialization [Figure] and architectural choices.
Revised: 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.

Original: Our results demonstrate the central role that λ 1 plays in neural network optimization—maintaining sufﬁciently small λ 1 during optimization is a necessary condition for successful training at large learning rates.
Revised: Our results demonstrate the central role that λ 1 plays in neural network optimization—maintaining sufﬁciently small λ 1 during optimization is a necessary condition for successful training at large learning rates.

Original: Consequently, reducing λ 1 is a primary beneﬁt of proper tuning of a number of architecture and optimization hyperparameters: including model initialization, location of normalization, and warmup schedule.
Revised: Consequently, reducing λ 1 is a primary beneﬁt of proper tuning of a number of architecture and optimization hyperparameters: including model initialization, location of normalization, and warmup schedule.

Original: Speciﬁcally, we show the following:
Revised: Speciﬁcally, we show the following:

Original: • We provide large scale empirical conﬁrmation 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 (cid:46) 2 /η , where η denotes the learning rate.
Revised: • We provide large scale empirical conﬁrmation 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 (cid:46) 2 /η , where η denotes the learning rate.

Original: This corroborates the theoretical predictions of Wu et al.
Revised: This corroborates the theoretical predictions of Wu et al.

Original: and recent empirical observations of Jastrzebski et al.
Revised: and recent empirical observations of Jastrzebski et al.

Original: and
Revised: and

Original: Cohen et al.
Revised: Cohen et al.

Original: • 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.
Revised: • 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.

Original: • We show that learning rate warmup gradually reduces λ 1 during training, offering similar beneﬁts to better model initialization.
Revised: • We show that learning rate warmup gradually reduces λ 1 during training, offering similar beneﬁts to better model initialization.

Original: We connect the mechanism by which warmup operates to the dynamical stability model of Wu et al.
Revised: We connect the mechanism by which warmup operates to the dynamical stability model of Wu et al.

Original: • We show that learning rate warmup is a simple yet competitive baseline for research into better model initialization.
Revised: • We show that learning rate warmup is a simple yet competitive baseline for research into better model initialization.

Original: 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.
Revised: 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.

Original:  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.
Revised: • 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.

Original:  R ELATED W ORKUnderstanding BatchNorm
Revised:  R ELATED W ORK

Revised: Understanding BatchNorm

Original: The loss Hessian has been a central object of study for understanding optimization of neural networks.
Revised: The loss Hessian has been a central object of study for understanding optimization of neural networks.

Original: Santurkar et al.
Revised: Santurkar et al.

Original: (2018) argues that an important beneﬁt of Batch
Revised: (2018) argues that an important beneﬁt of Batch

Original: Normalization is improved smoothness of the loss surface, while Lewkowycz et al.
Revised: Normalization is improved smoothness of the loss surface, while Lewkowycz et al.

Original: (2020) notes that this is improved smoothness is only observed when higher learning rates are used in combination with Batch Normalization.
Revised: (2020) notes that this is improved smoothness is only observed when higher learning rates are used in combination with Batch Normalization.

Original: Our results are generally consistent with this current understanding of
Revised: Our results are generally consistent with this current understanding of

Original: 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).
Revised: 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).

Original: Evolution of the loss Hessian Recent research has closely studied the interaction between sharpness and learning rate.
Revised: Evolution of the loss Hessian Recent research has closely studied the interaction between sharpness and learning rate.

Original: Wu et al.
Revised: Wu et al.

Original: (2018) provides a dynamical stability model which predicts that the loss curvature at convergence must satisfy 2 λ 1 ≤ 2 /η .
Revised: (2018) provides a dynamical stability model which predicts that the loss curvature at convergence must satisfy 2 λ 1 ≤ 2 /η .

Original: Recent work has provided empirical evidence that λ 1 (cid:46) 2 /η often holds well before convergence (Cohen et al., 2021;
Revised: Recent work has provided empirical evidence that λ 1 (cid:46) 2 /η often holds well before convergence (Cohen et al., 2021;

Original: Jastrzebski et al., 2020).
Revised: Jastrzebski et al., 2020).

Original: Cohen et al.
Revised: Cohen et al.

Original: (2021) focused on full batch training at small learning rates, and observed “progressive sharpening”, where λ 1 increases during training until λ 1 ≈ 2 . 0 /η .
Revised: (2021) focused on full batch training at small learning rates, and observed “progressive sharpening”, where λ 1 increases during training until λ 1 ≈ 2 . 0 /η .

Original: 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.
Revised: 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.

Original: (2021) argue that progressive sharpening does not occur.
Revised: (2021) argue that progressive sharpening does not occur.

Original: We note that Wu et al.
Revised: We note that Wu et al.

Original: (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.
Revised: (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.

Original: We conﬁrm this prediction holds even early in training (see Appendix Figure 16).
Revised: We conﬁrm this prediction holds even early in training (see Appendix Figure 16).

Original: Lewkowycz et al.
Revised: Lewkowycz et al.

Original: (2020) proves that for single hidden layer neural networks initialized at point with λ 1 > 2 . 

Revised: (2020) proves that for single hidden layer neural networks initialized at point with λ 1 > 2 .

Original: /η and trained with an MSE loss may enter a “catapult” regime—where the loss increases early until a ﬂatter region of the loss surface is found, with divergence occurring in cases where λ 1 greatly exceeds 2 . 0 /η .
Revised: /η and trained with an MSE loss may enter a “catapult” regime—where the loss increases early until a ﬂatter region of the loss surface is found, with divergence occurring in cases where λ 1 greatly exceeds 2 . 0 /η .

Original: In contrast to the simpliﬁed setting considered in Lewkowycz et al.
Revised: In contrast to the simpliﬁed setting considered in Lewkowycz et al.

Original: (2020), we ﬁnd that divergence may occur even though λ 1 (cid:28) 2 /η at initialization.
Revised: (2020), we ﬁnd that divergence may occur even though λ 1 (cid:28) 2 /η at initialization.

Original:  E XPERIMENTAL S ETUP
Revised:  E XPERIMENTAL S ETUP

Original: We investigate models trained on several benchmarks: CIFAR-10 (Krizhevsky, 2009) and ImageNet (Russakovsky et al., 2015) for image classiﬁcation, LM1B (Chelba et al., 2013) for Language
Revised: We investigate models trained on several benchmarks: CIFAR-10 (Krizhevsky, 2009) and ImageNet (Russakovsky et al., 2015) for image classiﬁcation, LM1B (Chelba et al., 2013) for Language

Original: Modeling, and WMT for Neural Machine Translation (NMT).
Revised: Modeling, and WMT for Neural Machine Translation (NMT).

Original: On CIFAR-10 we consider the
Revised: On CIFAR-10 we consider the

Original: WideResnet (Zagoruyko and Komodakis, 2016) and DenseNet (Huang et al., 2017) architectures, both with and without Batch Normalization.
Revised: WideResnet (Zagoruyko and Komodakis, 2016) and DenseNet (Huang et al., 2017) architectures, both with and without Batch Normalization.

Original: When training without Batch Normalization we consider several initialization strategies including the default “LeCun Normal” initialization, andrunning
Revised: We consider two variants of the DenseNet architecture. The standard variant from the open sourced code of Zhu et al. (2021) is considered in Figure 5 and

Original:  MetaInit.
Revised: Table 1. A less stable variant changes the strides in the average pooling layers to (1,1) is used for Figure 2 and is denoted as Stride-(1,1) DenseNet (see Appendix D.1 for a more detailed discussion). When training without Batch Normalization we consider several initialization strategies including the default “LeCun Normal” initialization, and running MetaInit.

Original: As a way to artiﬁcially induce worse initializations, we also consider experiments where we scale every variable produced by the default initialization by a constant factor α .
Revised: As a way to artiﬁcially induce worse initializations, we also consider experiments where we scale every variable produced by the default initialization by a constant factor α .

Original: 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.
Revised: 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.

Original: For NMT and LM1B Language Modeling, we train 6 layer Transformer models (Vaswani et al., 2017).
Revised: For NMT and LM1B Language Modeling, we trainlayer Transformer models (Vaswani et al., 2017).

Original: Inspired from Xiong et al.
Revised: Inspired from Xiong et al.

Original: (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.
Revised: (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.

Original: Each model is trained with various learning rates using cosine decay (unless mentioned explicitly).
Revised: Each model is trained with various learning rates using cosine decay (unless mentioned explicitly).

Original: For warmup experiments we use linear warmup which starts at 0 and scales linearly to a max value η before applying cosine decay.
Revised: For warmup experiments we use linear warmup which starts at 0 and scales linearly to a max value η before applying cosine decay.

Original: To measure the max eigenvalue of the loss Hessian we use the
Revised: To measure the max eigenvalue of the loss Hessian we use the

Original: Lanczos method where the number of iterations varied as needed depending on the architecture (details provided in the appendix).
Revised: Lanczos method where the number of iterations varied as needed depending on the architecture (details provided in the appendix).

Original:  E ARLY T RAINING I NSTABILITY AND THE L OSS H ESSIAN
Revised:  E ARLY T RAINING I NSTABILITY AND THE L OSS H ESSIAN

Original: In Figure 2 we plot the curvature at initialization and during training for a series of models λ
Revised: In Figure 2 we plot the curvature at initialization and during training for a series of models trained on different datasets (plots showing ﬁnal performance for all models can be found in the appendix).

Original:  , and that  are outside this region shortly before divergence.
Revised: 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 a speciﬁed 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.

Original: At initialization, some  can be successfully trained even  they  out  is  likely 
Revised: 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.

Original: 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 .
Revised: 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 .

Original: Scaling the
Revised: Scaling the

Original: WideResnet initialization up by a factor of 1 . 5 exacerbates the problem, resulting in even higher curvature at initialization and divergence when η > 10 − 2 .
Revised: WideResnet initialization up by a factor of 1 . 5 exacerbates the problem, resulting in even higher curvature at initialization and divergence when η > 10 − 2 .

Original: 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.
Revised: 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.

Original: We also observed that higher learning rates can be unlocked when the models are trained with learning rate warmup.
Revised: We also observed that higher learning rates can be unlocked when the models are trained with learning rate warmup.

Original: Warmup was particularly effective for models which exhibit large λ 1 either at initialization or early in training.
Revised: Warmup was particularly effective for models which exhibit large λ 1 either at initialization or early in training.

Original: Other models such as the post activation
Revised: Other models such as the post activation

Original: Resnet-50, and WideResnet w/ Batch Normalization did not beneﬁt from warmup at the considered learning rates (see Appendix).
Revised: Resnet-50, and WideResnet w/ Batch Normalization did not beneﬁt from warmup at the considered learning rates (see Appendix).

Original: For the  DenseNet experiments, it is noteworthy that the models with Batch Normalization actually start out with higher curvature than the non-BN variants.
Revised: 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.

Original: 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.
Revised: 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).

Original: (2019) demonstrates settings for which Batch Normalization can cause exploding gradients.

Original: We found that the Batch Normalization models were more unstable than the non-BN variants here, as some models diverged at smaller learning rates.
Revised: We found that the Batch Normalization models were more unstable than the non-BN variants here, as some models diverged at smaller learning rates.

Original: 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.
Revised: 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.

Original: This result suggests that BN still offers training stability for this model, and ﬂatter curvature mid training if trained with warmup and a higher learning rate , however no smoothness beneﬁts are observed at initialization.
Revised: This result suggests that BN still offers training stability for this model, and ﬂatter curvature mid training if trained with warmup and a higher learning rate , however no smoothness beneﬁts are observed at initialization.

Original: See Appendix D.1 for more details on this phenomenon.
Revised: See Appendix D.1 for more details on this phenomenon.

Original: 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).
Revised: 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).

Original: For the preactivation block, we also consider ﬂipping the order of the ReLU activation and batch normalization, as was considered in Brock et al.
Revised: For the preactivation block, we also consider ﬂipping the order of the ReLU activation and batch normalization, as was considered in Brock et al.

Original: We ﬁnd 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 .
Revised: We ﬁnd 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 .

Original: Notably, there are several models in our experiments which diverge despite starting out in a region where λ 1 < 2 /η .
Revised: Notably, there are several models in our experiments which diverge despite starting out in a region where λ 1 < 2 /η .

Original: This occurs for both the pre and post layernorm transformer, and the WideResnet model initialized with MetaInit.
Revised: This occurs for both the pre and post layernorm transformer, and the WideResnet model initialized with MetaInit.

Original: 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 ﬁnal observed curvature before divergence.
Revised: 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 ﬁnal observed curvature before divergence.

Original: Full training curves for these models can be found in the appendix.
Revised: Full training curves for these models can be found in the appendix.

Original: This implies that measuring λ 1 at initialization is not always sufﬁcient to predict whether or not the model will be easily trained.
Revised: This implies that measuring λ 1 at initialization is not always sufﬁcient to predict whether or not the model will be easily trained.

Original: Currently, some architectural innovations are motivated by an analysis of either gradient statistics or smoothness at initialization (Liu et al.
Revised: Currently, some architectural innovations are motivated by an analysis of either gradient statistics or smoothness at initialization (Liu et al.

Original: , 2020)—a more robust analysis would consider the evolution of these statistics under SGD.
Revised: , 2020)—a more robust analysis would consider the evolution of these statistics under SGD.

Original:  T HE I NTERACTION BETWEEN L EARNING R ATE W ARMUP , I NITIALIZATION AND C URVATURE
Revised:  T HE I NTERACTION BETWEEN L EARNING R ATE W ARMUP , I NITIALIZATION AND C URVATURE

Original: The success of learning rate warmup is inconsistent with conventional optimization wisdom, which traditionally suggests adapting the step size to the curvature (see for example the discussion aroundequation 2.4 in McCandlish et al.
Revised: The success of learning rate warmup is inconsistent with conventional optimization wisdom, which traditionally suggests adapting the step size to the curvature (see for example the discussion around

Original: However, with the understanding that λ 1 is a dynamic quantity whose evolution is tightly coupled with the learning rate schedule, the beneﬁts of a warmup period are more easily understood.

Original: We argue that the success of learning rate warmup follows naturally from two properties of training deep models:

Original: Models diverge when the learning rate is too large relative to the 2 /λ 1 bound.

Original: When the learning rate only slightly exceeds 2 /λ 1 optimization is unstable until the parameters move to a region with smaller λ 1 (Wu et al., 2018; Lewkowycz et al., 2020).

Original: Across all datasets and models, successful training occurs only when optimization enters a stable region of parameter space where λ 1 ≤ 2 /η .
Revised: Across all datasets and models, successful training occurs only when optimization enters a stable region of parameter space where λ 1 ≤ 2 /η .

Original: Models with higher initial loss curvature tend to diverge at smaller learning rates relative to models initialized with ﬂatter curvature, however larger learning rates are possible in these cases when warmup is used.
Revised: Models with higher initial loss curvature tend to diverge at smaller learning rates relative to models initialized with ﬂatter curvature, however larger learning rates are possible in these cases when warmup is used.

Original: Note, to avoid overlapping points we have applied a small deterministic shift to the x-position of points on the right hand plots.
Revised: Note, to avoid overlapping points we have applied a small deterministic shift to the x-position of points on the right hand plots.

Original: In the case of
Revised: In the case of

Original:  DenseNets w/ BN we observed some models exhibit catapult behavior (loss increases early followed by normal training), these have been marked accordingly.
Revised: Stride-(1,1) DenseNets w/ BN we observed some models exhibit catapult behavior (loss increases early followed by normal training), these have been marked accordingly.

Original: [Figure]
Revised: [Figure]

Revised: equation 2.4 in McCandlish et al.

Revised: However, with the understanding that λ 1 is a dynamic quantity whose evolution is tightly coupled with the learning rate schedule, the beneﬁts of a warmup period are more easily understood.

Revised: We argue that the success of learning rate warmup follows naturally from two properties of training deep models:

Revised: Models diverge when the learning rate is too large relative to the 2 /λ 1 bound.

Revised: When the learning rate only slightly exceeds 2 /λ 1 optimization is unstable until the parameters move to a region with smaller λ 1 (Wu et al., 2018; Lewkowycz et al., 2020).

Original: The ﬁrst criteria implies that we can’t start η off at too large of a value relative to λ 1 at initialization.
Revised: The ﬁrst criteria implies that we can’t start η off at too large of a value relative to λ 1 at initialization.

Original: The second criteria implies that gradually increasing η can gradually “push” the parameters to a region of parameter space where optimization is stable (with lower values of λ 1 ).
Revised: The second criteria implies that gradually increasing η can gradually “push” the parameters to a region of parameter space where optimization is stable (with lower values of λ 1 ).

Original: In Figure 4 there is clear evidence for this “pushing”, as during the warmup period the we see that λ 1 ≈ 2 . 

Revised: In Figure 4 there is clear evidence for this “pushing”, as during the warmup period the we see that λ 1 ≈ 2 .

Original: /η holds for a large part of the warmup phase.
Revised: /η holds for a large part of the warmup phase.

Original: Furthermore, this approximation holds even as we vary the length of the warmup period.
Revised: Furthermore, this approximation holds even as we vary the length of the warmup period.

Original: Other examples can be seen in Figure 3 (B and F), and Figure 15 in the appendix.
Revised: Other examples can be seen in Figure 3 (B and F), and Figure 15 in the appendix.

Original: Warmup is not the only method capable of reducing λ 1 during training, one can instead initialize the model in a region where λ 1 starts off small.
Revised: Warmup is not the only method capable of reducing λ 1 during training, one can instead initialize the model in a region where λ 1 starts off small.

Original: Consider for example, the points A, B and C in Figure 3.
Revised: Consider for example, the points A, B and C in Figure 3.

Original: Each point shows optimization of a non-BN WideResnet with peak learning rate of . 1 .
Revised: Each point shows optimization of a non-BN WideResnet with peak learning rate of . 1 .

Original: In (A) we see the model diverges within 3 steps without warmup using the default initialization.
Revised: In (A) we see the model diverges within 3 steps without warmup using the default initialization.

Original: In (B) we see that a linear warmup period results in λ 1 progressively decreasing until the peak step size of . 1 is reached at step 1000, with no divergence occurring.
Revised: In (B) we see that a linear warmup period results in λ 1 progressively decreasing until the peak step size of . 1 is reached at step 1000, with no divergence occurring.

Original: Finally in (C) we initialize the same model with MetaInit, at which point λ 1 is small at initialization, and the model can be trained at η = .
Revised: Finally in (C) we initialize the same model with MetaInit, at which point λ 1 is small at initialization, and the model can be trained at η = .

Original: 1 without warmup.
Revised: 1 without warmup.

Original: Similar to the aforementioned MetaInit, the success of related initialization strategies can be explained by reduced λ 1 early in training.
Revised: Similar to the aforementioned MetaInit, the success of related initialization strategies can be explained by reduced λ 1 early in training.

Original: In Figure 5 (left) we look at the evolution of λ 1 during the GradInit meta optimization process and compare this with simply training the same model using gradient clipping 3 .
Revised: In Figure 5 (left) we look at the evolution of λ 1 during the GradInit meta optimization process and compare this with simply training the same model using gradient clipping 3 .

Original: Both methods result in λ 1 decreasing dramatically, after which λ 1 hovers around 2 /η . 
Revised: Both methods result in λ 1 decreasing dramatically, after which λ 1 hovers around 2 /η . [Figure]

Revised: a factorwas enough to recover the performance the normalized model

Revised: [Table]

Revised: WMT

Original: Notably, GradInit starts regular training off at λ 1 signiﬁcantly below the 2 /η bound, however the curvature quickly increases within a few steps.
Revised: Warmup + Clip + .25x Init (ours) [Table] [Table] BLEU

Revised: [Table]

Revised: LayerNorm

Revised: [Table]

Revised: [Table]

Revised: Notably

Revised: [Table]

Revised: GradInit starts regular training off at λ 1 signiﬁcantly below the 2 /η bound, however the curvature quickly increases within a few steps.

Original: Given that initialization and warmup serve similar roles in reducing λ 1 , we expect to be able to achieve similar performance using the two methods.
Revised: Given that initialization and warmup serve similar roles in reducing λ [Table] we expect to be able to achieve similar performance using the two methods.

Original: As shown in Table 1 we can easily match key advances in this ﬁeld by applying learning rate warmup alone (see Appendix for experimental details).
Revised: As shown in Table 1 we can easily match key advances in this ﬁeld by applying learning rate warmup alone (see Appendix for experimental details).

Original: Beyond controlling λ 1 mid-training, the learning rate η controls more general conditioning measures of the loss surface. For example, in Figure 5 we observe that even the MetaInit gradient quotient — [Figure]
Revised: Beyond controlling λ 1 mid-training

Original: [Figure]
Revised: [Table]

Original: meta optimization process. Both GradInit and training with warmup are able to reduce the large curvature of model.
Revised: the learning rate η controls more general conditioning measures of the loss surface.

Original: The GradInit algorithm initializes SGD in a ﬂatter 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 MetaLoss the conditioning measure directly optimized by this initialization strategy—is controlled by η mid training.
Revised: For example, in Figure 5 we observe that even the MetaInit gradient quotient — the conditioning measure directly optimized by this initialization strategy—is controlled by η mid training.

Original: This again provides further evidence that the primary beneﬁt of this initialization method is to reduce λ 1 at initialization.
Revised: This again provides further evidence that the primary beneﬁt of this initialization method is to reduce λ 1 at initialization.

Original: As shown, any gains by optimizing the more general gradient quotient must be short lived as the initialization has no control over the long term value.
Revised: As shown, any gains by optimizing the more general gradient quotient must be short lived [Figure]

Original:  T HE E FFECTS OF C URVATURE ON B ATCH S IZE S CALING
Revised:  T HE E FFECTS OF C URVATURE ON B ATCH S IZE S CALING

Original: So far, we discussed how large loss curvature limits the range of stable learning rates for training.
Revised: So far, we discussed how large loss curvature limits the range of stable learning rates for training.

Original: In this section, we highlight how these limits on usable learning rates affect the model’s ability [Table] Transformer 6L (w/ LN) [Table]
Revised: In this section, we highlight how these limits on usable learning rates affect the model’s ability to effectively leverage larger batch sizes.

Original: to effectively leverage larger batch sizes.

Original: Previous research has studied the interplay of the loss curvature and batch size scaling from various different perspectives.
Revised: Previous research has studied the interplay of the loss curvature and batch size scaling from various different perspectives.

Original: Most notably, Shallue et al.
Revised: Most notably, Shallue et al.

Original: (2018) observe that increasing the batch size yields consistent improvements in training speed until a (problem-dependent) critical batch size is reached; increasing the batch size beyond this threshold yields diminishing improvements in training speed.
Revised: (2018) observe that increasing the batch size yields consistent improvements in training speed until a (problem-dependent) critical batch size is reached; increasing the batch size beyond this threshold yields diminishing improvements in training speed.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2019a) observe that a simple Noisy
Revised: (2019a) observe that a simple Noisy

Original: Quadratic Model (NQM) is able to capture some of the empirical behavior observed in Shallue et al.
Revised: Quadratic Model (NQM) is able to capture some of the empirical behavior observed in Shallue et al.

Original: Similarly, McCandlish et al.
Revised: Similarly, McCandlish et al.

Original: (2018) use quadratic approximations to the loss to provide a closed form expression for the critical batch size as a function of the loss Hessian and the covariance of the stochastic gradient.
Revised: (2018) use quadratic approximations to the loss to provide a closed form expression for the critical batch size as a function of the loss Hessian and the covariance of the stochastic gradient.

Original: We contribute to this literature by highlighting the role of λ 1 in the batch size scaling behavior of the model.
Revised: We contribute to this literature by highlighting the role of λ 1 in the batch size scaling behavior of the model.

Original: For this analysis, we focus on three of the WideResnet variants considered in Figure 2—the BatchNorm model (a low curvature model), the non BatchNorm model (with moderate curvature), and the non BatchNorm model with 1.5X init scaling (with high curvature).
Revised: For this analysis, we focus on three of the WideResnet variants considered in Figure 2—the BatchNorm model (a low curvature model), the non BatchNorm model (with moderate curvature), and the non BatchNorm model with 1.5X init scaling (with high curvature).

Original: We train these models while sweeping both the learning rate and the batch size.
Revised: We train these models while sweeping both the learning rate and the batch size.

Original: 4 We then measure the number of training steps required to reach 85% validation accuracy, and the optimal learning rate found for each batch size.
Revised: 4 We then measure the number of training steps required to reach 85% validation accuracy, and the optimal learning rate found for each batch size.

Original: Similar to Shallue et al.
Revised: Similar to Shallue et al.

Original: (2018), we normalize the plotted steps to 85% accuracy by the value measured at batch size 64 .
Revised: (2018), we normalize the plotted steps to 85% accuracy by the value measured at batch size 64 .

Original: The results are shown in Figure 6.
Revised: The results are shown in Figure 6.

Original: A few observations are in order: The low curvature model shows almost linear speedups in training speed as the batch size increases.
Revised: A few observations are in order: The low curvature model shows almost linear speedups in training speed as the batch size increases.

Original: In contrast, the high curvature model exhibits only minimal improvements in training speed with larger batch sizes.
Revised: In contrast, the high curvature model exhibits only minimal improvements in training speed with larger batch sizes.

Original: These scaling differences are closely mirrored by how the optimal learning rate η ∗ changes with the batch size: for the low curvature model η ∗ increases linearly with the batch size, while for the high curvature model η ∗ is ﬁxed around 3 × 10 − 3 .
Revised: These scaling differences are closely mirrored by how the optimal learning rate η ∗ changes with the batch size: for the low curvature model η ∗ increases linearly with the batch size, while for the high curvature model η ∗ is ﬁxed around 3 × 10 − 3 .

Original: Notably, for the high curvature model η ∗ is almost always the largest non-divergent value—a clear indication that the high loss curvature slows down training by preventing larger values from being used.
Revised: Notably, for the high curvature model η ∗ is almost always the largest non-divergent value—a clear indication that the high loss curvature slows down training by preventing larger values from being used.

Original: A clear picture emerges from these observations.
Revised: A clear picture emerges from these observations.

Original: 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
Revised: 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

Original: Jastrz˛ebski et al.
Revised: Jastrz˛ebski et al.

Original: ; Goyal et al.
Revised: ; Goyal et al.

Original: ; Shallue et al.
Revised: ; Shallue et al.

Original: ; McCandlish et al.
Revised: ; McCandlish et al.

Original: Our results suggest that large values of λ 1 place a sharp limit on the maximum the learning rate possible and therefore, limit the model’s ability to leverage data parallelism effectively
Revised: Our results suggest that large values of λ 1 place a sharp limit on the maximum the learning rate possible and therefore, limit the model’s ability to leverage data parallelism effectively.

Original: [Figure]
Revised: [Figure]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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 all enable higher learning rates to be used (without causing divergence) by reducing λ 1 during training.
Revised: 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 all enable higher learning rates to be used (without causing divergence) by reducing λ 1 during training.

Original: It is noteworthy that two of the most popular models we investigated (the popular post-activation variant of the Resnet-50 and the standard WideResnet 28-10) did not beneﬁt from learning rate warmup, and exhibited small values of λ 1 throughout training at the learning rates we considered.
Revised: It is noteworthy that two of the most popular models we investigated (the popular post-activation variant of the Resnet-50 and the standard WideResnet 28-10) did not beneﬁt from learning rate warmup, and exhibited small values of λ 1 throughout training at the learning rates we considered.

Original: Thus researchers and practitioners who primarily work with well-optimized architectures might never notice a beneﬁt from using warmup.
Revised: Thus researchers and practitioners who primarily work with well-optimized architectures might never notice a beneﬁt from using warmup.

Original: However, even seemingly trivial modiﬁcations to a working architecture can easily result in large values of λ 1 and thus instability early in training—a naive response to such a situation would be to dramatically reduce the learning rate or, even worse, abandon the modiﬁcation being investigated all together.
Revised: However, even seemingly trivial modiﬁcations to a working architecture can easily result in large values of λ 1 and thus instability early in training—a naive response to such a situation would be to dramatically reduce the learning rate or, even worse, abandon the modiﬁcation being investigated all together.

Original: 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.
Revised: 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.

Original: A L IMITATIONS
Revised: A L IMITATIONS

Original: Our analysis has focused primarily on models trained with SGD with momentum.
Revised: Our analysis has focused primarily on models trained with SGD with momentum.

Original: This decision was motivated to reduce additional confounds that arise when using adaptive preconditioning.
Revised: This decision was motivated to reduce additional confounds that arise when using adaptive preconditioning.

Original: Notably, it is unclear what the analogue of λ 1 ≤ 2 /η should be for a model trained with Adam.
Revised: Notably, it is unclear what the analogue of λ 1 ≤ 2 /η should be for a model trained with Adam.

Original: In the appendix, we provide evidence that loss curvature adaption to the learning rate does occur even for
Revised: In the appendix, we provide evidence that loss curvature adaption to the learning rate does occur even for

Original: Transformer models trained with Adam, and that learning rate warmup results in the similar effect of the optimization trajectory being “pushed” to ﬂatter regions.
Revised: Transformer models trained with Adam, and that learning rate warmup results in the similar effect of the optimization trajectory being “pushed” to ﬂatter regions.

Original: However we leave a deeper analysis into this for future work.
Revised: However we leave a deeper analysis into this for future work.

Original: Finally, while our experiments certainly raise questions about the efﬁcacy better model initialization has on further accelerating training, our measurements has focused primarily on the (lack of) inﬂuence initialization has on λ 1 mid training.
Revised: Finally, while our experiments certainly raise questions about the efﬁcacy better model initialization has on further accelerating training, our measurements has focused primarily on the (lack of) inﬂuence initialization has on λ 1 mid training.

Original: It is possible that better initializations could have lasting inﬂuence on the broader Hessian eigenspectrum (for example improving the ration λ k /λ 1 for smaller eigenvalues λ k ) and that our analysis is missing such an effect.
Revised: It is possible that better initializations could have lasting inﬂuence on the broader Hessian eigenspectrum (for example improving the ration λ k /λ 1 for smaller eigenvalues λ k ) and that our analysis is missing such an effect.

Original: B B RIEF REVIEW OF THE LOSS H ESSIAN , EIGENVALUES AND QUADRATIC STABILITY BOUNDS
Revised: B B RIEF REVIEW OF THE LOSS H ESSIAN , EIGENVALUES AND QUADRATIC STABILITY BOUNDS

Original: For completeness, we include a formal deﬁnition of the fundamental mathematical quantities discussed in the paper.
Revised: For completeness, we include a formal deﬁnition of the fundamental mathematical quantities discussed in the paper.

Original: We derive most of this discussion from the relevant chapters in Horn and Johnson (2012) and Boyd et al.
Revised: We derive most of this discussion from the relevant chapters in Horn and Johnson (2012) and Boyd et al.

Original: We refer the reader to these sources for a more detailed discussion.
Revised: We refer the reader to these sources for a more detailed discussion.

Original: B. T HE H ESSIAN M ATRIX
Revised: B. T HE H ESSIAN M ATRIX

Original: The second derivative or Hessian matrix of the loss function L ( · ) at a point θ ∈ R n is denoted by
Revised: The second derivative or Hessian matrix of the loss function L ( · ) at a point θ ∈ R n is denoted by

Original: H ( θ ) ∈ R n × n where ∀ 1 ≤ i, j≤ n
Revised: H ( θ ) ∈ R n × n where ∀ 1 ≤ i, j

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Moreover, by Schwarz’s theorem, if the second partial derivatives of L ( · ) are continuous at θ , the matrix H ( θ ) is symmetric.
Revised: Moreover, by Schwarz’s theorem, if the second partial derivatives of L ( · ) are continuous at θ , the matrix H ( θ ) is symmetric.

Original: This is a broad condition that holds for all the loss surfaces we examine in the main text (beyond a set of measure zero).
Revised: This is a broad condition that holds for all the loss surfaces we examine in the main text (beyond a set of measure zero).

Original: B. E IGENVALUES
Revised: B. E IGENVALUES

Original: Deﬁnition 1 Let A ∈ R n × n .
Revised: Deﬁnition 1 Let A ∈ R n × n .

Original: If a scalar λ and a nonzero vector x satisfy
Revised: If a scalar λ and a nonzero vector x satisfy

Original: [Equation]
Revised: [Equation]

Original: then λ is called an eigenvalue of A and x is called an eigenvector of A associated with λ .
Revised: then λ is called an eigenvalue of A and x is called an eigenvector of A associated with λ .

Original: If A is a symmetric real matrix (such as the Hessian matrix), A can be factored as
Revised: If A is a symmetric real matrix (such as the Hessian matrix), A can be factored as

Original: [Equation]
Revised: [Equation]

Original: where Q ∈ R n × n is an orthogonal matrix and Λ = diag ( λ 1 , . .
Revised: where Q ∈ R n × n is an orthogonal matrix and Λ = diag ( λ 1 , . .

Original: λ n ) ∈ R n × n is a real diagonal matrix.
Revised: λ n ) ∈ R n × n is a real diagonal matrix.

Original: Here, { λ i } ni =1 are all of the eigenvalues of A .
Revised: Here, { λ i } ni =1 are all of the eigenvalues of A .

Original: We order λ
Revised: We order λ

Original: i such that λ 1 ≥ λ 2 · · · ≥ λ n .
Revised: i such that λ 1 ≥ λ 2 · · · ≥ λ n .

Original: In this ordering, λ 1 corresponds to the maximum eigenvalue of A and λ n corresponds to its minimum eigenvalue.
Revised: In this ordering, λ 1 corresponds to the maximum eigenvalue of A and λ n corresponds to its minimum eigenvalue.

Original: The maximum and minimum eigenvalues of A satisfy the following important properties:
Revised: The maximum and minimum eigenvalues of A satisfy the following important properties:

Original: [Equation]
Revised: [Equation]

Original: In particular, for any x ∈ R n , we have
Revised: In particular, for any x ∈ R n , we have

Original: [Equation]
Revised: [Equation]

Original: B. S TABILITY OF G RADIENT D ESCENT FOR Q UADRATIC L OSS
Revised: B. S TABILITY OF G RADIENT D ESCENT FOR Q UADRATIC L OSS

Original: Now that we have established the basics, let’s derive the stability condition for GD applied to a quadratic loss function.
Revised: Now that we have established the basics, let’s derive the stability condition for GD applied to a quadratic loss function.

Original: Note that Wu et al.
Revised: Note that Wu et al.

Original: and Cohen et al.
Revised: and Cohen et al.

Original: (2021) provide more general bounds for the stability of SGD-type optimization algorithms.
Revised: (2021) provide more general bounds for the stability of SGD-type optimization algorithms.

Original: Here, we state & derive the stability condition for GD for the sake of completeness.
Revised: Here, we state & derive the stability condition for GD for the sake of completeness.

Original: Let
Revised: Let

Original: [Equation]
Revised: [Equation]

Original: where H is a symmetric matrix with non-negative eigenvalues.
Revised: where H is a symmetric matrix with non-negative eigenvalues.

Original: Let’s consider GD dynamics starting from a random point θ 0 .
Revised: Let’s consider GD dynamics starting from a random point θ 0 .

Original: Under GD with a ﬁxed step-size η > 0 , we have
Revised: Under GD with a ﬁxed step-size η > 0 , we have

Original: [Equation]
Revised: [Equation]

Original: Continuing this iteration to step 0 yields
Revised: Continuing this iteration to step 0 yields

Original: [Equation]
Revised: [Equation]

Original: As t → ∞ , this iteration is stable iff 5 the eigenvalues of ( I − ηH ) have absolute magnitude bounded by one, which can be stated equivalently as
Revised: As t → ∞ , this iteration is stable iff 5 the eigenvalues of ( I − ηH ) have absolute magnitude bounded by one, which can be stated equivalently as

Original: [Equation]
Revised: [Equation]

Original: which is the exact condition discussed and explored in the main text.
Revised: which is the exact condition discussed and explored in the main text.

Original: C M ISCELLANEOUS F IGURES
Revised: C M ISCELLANEOUS F IGURES

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: D P ERFORMANCE OF MODELS IN F IGURE
Revised: D P ERFORMANCE OF MODELS IN F IGURE

Original: In this this section [Figure] we plot the performance vs learning rate for all of the models shown in Figureof the main text.
Revised: In this this section, we plot the performance vs learning rate for all of the models shown in Figureof the main text.

Original: These are shown in Figures 9, 13,10, and 11.
Revised: These are shown in Figures 9, 13, and 11.

Original: For models which diverged, we plot the best test performance achieved before divergence.
Revised: For models which diverged, we plot the best test performance achieved before divergence.

Original: In all settings, high curvature affects the ﬁnal performance by limiting the use of higher learning rates.
Revised: In all settings, high curvature affects the ﬁnal performance by limiting the use of higher learning rates.

Original: We also noted several models in Figure 2 which diverged despite training starting out in a stable region of parameter space.
Revised: We also noted several models in Figure 2 which diverged despite training starting out in a stable region of parameter space.

Original: In Figure 12 we plot the evolution of the loss sharpness during training, showing that it quickly enters a region where λ 1 > 2 . 

Revised: In Figure 12 we plot the evolution of the loss sharpness during training, showing that it quickly enters a region where λ 1 > 2 .

Original: /η
Revised: /η

Original: before diverging around step 90.
Revised: before diverging around step 90.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: D. D ISCUSSION OF  D ENSE N ET E XPERIMENTS
Revised: D. D ISCUSSION OF S TRIDE -(1,1) D ENSE N ET E XPERIMENTS

Original: In this section we discuss in more detail the  DenseNet experiments shown in Figure 2 in the main text.
Revised: In this section we discuss in more detail the Stride-(1,1) DenseNet experiments shown in Figure 2 in the main text.

Revised: These experiments use a non-standard version of the DenseNet architecture where all average pooing strides are set to (1,1).

Revised: Note the experiments in Figrue 5 and Table 1 instead use the standard strides implementation from the open sourced code of Zhu et al.

Revised: The Stride-(1,1)

Original: These experiments are noteworthy because they seem todefy common intuition that BN results in ﬂatter curvature.
Revised: DenseNet architecture is noteworthy because it is a counter example to common intuition that adding Batch Normalization results in ﬂatter curvature.

Original: 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.
Revised: 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.

Original: In Figure 13 we provide a more detailed analysis to understand what is happening with BN.
Revised: In Figure 13 we provide a more detailed analysis to understand what is happening with BN.

Original: First we plot the performance of the BN vs non-BN models both with and without warmup.
Revised: First we plot the performance of the BN vs non-BN models both with and without warmup.

Original: The differences are striking.
Revised: The differences are striking.

Original: 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.
Revised: 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.

Original: However, when trained with 1000 steps of warmup the BN variants now signiﬁcantly outperform the non-BN models are all considered learning rates.
Revised: However, when trained with 1000 steps of warmup the BN variants now signiﬁcantly outperform the non-BN models are all considered learning rates.

Original: They can even be successfully trained at higher learning rates than the non-BN variants, despite the high initial curvature.
Revised: They can even be successfully trained at higher learning rates than the non-BN variants, despite the high initial curvature.

Original: 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.
Revised: 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.

Original: The BN variants all exhibit catapult behavior early—the loss increases initially until the parameters enter a region of ﬂatter curvature.
Revised: The BN variants all exhibit catapult behavior early—the loss increases initially until the parameters enter a region of ﬂatter curvature.

Original: Warmup helps the BN variants, and signiﬁcantly reduces the severity of the catapult phase while enabling faster long term training.
Revised: Warmup helps the BN variants, and signiﬁcantly reduces the severity of the catapult phase while enabling faster long term training.

Original: Additionally, when we add warmup we ﬁnd that the BN variants can now be trained at higher learning rates than the non-BN variant.
Revised: Additionally, when we add warmup we ﬁnd that the BN variants can now be trained at higher learning rates than the non-BN variant.

Original: As shown at learning rate of .22, the non-BN model diverges during the warmup phase despite lower initial curvature.
Revised: As shown at learning rate of .22, the non-BN model diverges during the warmup phase despite lower initial curvature.

Original: Based on these experiments we arrive at the following conclusions.
Revised: Based on these experiments we arrive at the following conclusions.

Original: First, adding BN to the  DenseNet architecture results in high curvature at initialization, which results in a short period of instability during training.
Revised: First, adding BN to the Stride- (1,1) DenseNet architecture results in high curvature at initialization, which results in a short period of instability during training.

Original: However, once the parameters escape this region of large curvature, the BN variant exhibits favorable training dynamics relative to the non-BN variants.
Revised: However, once the parameters escape this region of large curvature, the BN variant exhibits favorable training dynamics relative to the non-BN variants.

Original: Thus, there still seems to be beneﬁt to adding BN to DenseNet, assuming steps are taken to mitigate the initial period of high curvature.
Revised: Thus, there still seems to be beneﬁt to adding BN, assuming steps are taken to mitigate the initial period of high curvature.

Original: The fact that adding BN to a model can result in high initial curvature is not without precedent, as Yang et al.
Revised: The fact that adding BN to a model can result in high initial curvature is not without precedent, as Yang et al.

Original: (2019) observe that adding BN to deep fully connected networks can result in exploding gradients at initialization.
Revised: (2019) observe that adding BN to deep fully connected networks can result in exploding gradients at initialization.

Original: These experiments highlight one of the primary takeaways of this work: that maintaining ﬂat curvature throughout training is a necessary (not sufﬁcient) condition for stable training of neural networks.
Revised: These experiments highlight one of the primary takeaways of this work: that maintaining ﬂat curvature throughout training is a necessary (not sufﬁcient) condition for stable training of neural networks.

Original: 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).
Revised: 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).

Original: BN has clear beneﬁts of improving curvature in most cases , but it is possible to produce conﬁgurations where adding BN paradoxically results in higher initial curvature than the non-BN variant.
Revised: BN has clear beneﬁts of improving curvature in most cases , but it is possible to produce conﬁgurations where adding BN paradoxically results in higher initial curvature than the non-BN variant.

Original: In these cases training is initially unstable, but once the curvature is reduced we see beneﬁts later in training of using BN.
Revised: In these cases training is initially unstable, but once the curvature is reduced we see beneﬁts later in training of using BN.

Original: E D ETAILS ON C OMPUTING THE H ESSIAN E IGENSPECTRUM VIA L ANCZOS
Revised: E D ETAILS ON C OMPUTING THE H ESSIAN E IGENSPECTRUM VIA L ANCZOS

Original: We use Lanczos iterations to estimate the top eigenvalue of the Hessian.
Revised: We use Lanczos iterations to estimate the top eigenvalue of the Hessian.

Original: Lanczos algorithm only requires Hessian-vector products which can be efﬁciently computed via Pearlmutter’s trick Pearlmutter (1994).
Revised: Lanczos algorithm only requires Hessian-vector products which can be efﬁciently computed via Pearlmutter’s trick Pearlmutter (1994).

Original: 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.
Revised: 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.

Original: Papyan (2018).
Revised: Papyan (2018).

Original: For our WMT / LM1B experiments, we run the algorithm for 45 steps while for image models, we use 40 steps.
Revised: For our WMT / LM1B experiments, we run the algorithm for 45 steps while for image models, we use 40 steps.

Original: When monitoring the evolution of the top eigenvalue as a function of the number of
Revised: When monitoring the evolution of the top eigenvalue as a function of the number of

Original: Lanczos steps, in all cases except one, we observe that the algorithm converges.
Revised: Lanczos steps, in all cases except one, we observe that the algorithm converges.

Original: For the case of
Revised: For the case of

Original: Resnet with ReLU → BN ordering, due to a very small eigengap between the top eigenvalue and the bulk, the convergence is signiﬁcantly slower.
Revised: Resnet with ReLU → BN ordering, due to a very small eigengap between the top eigenvalue and the bulk, the convergence is signiﬁcantly slower.

Original: We use 200 Lanczos steps in this case to alleviate the issue.
Revised: We use 200 Lanczos steps in this case to alleviate the issue.

Original: For this model, estimating λ 1 via power iteration (as is commonly done in the deep learning literature) will incorrectly the largest negative eigenvalue, not λ 1 as desired. 
Revised: For this model, estimating λ 1 via power iteration (as is commonly done in the deep learning literature) will incorrectly the largest negative eigenvalue, not λ 1 as desired. [Figure]

Original: It is well-known that Lanczos algorithm can suffer from numerical instabilities caused by ﬁniteprecision arithmetic.
Revised: It is well-known that Lanczos algorithm can suffer from numerical instabilities caused by ﬁniteprecision arithmetic.

Original: To alleviate these issues, Lanczos vectors are stored in ﬂoat64 accuracy and we perform reorthogonalized at each step of the algorithm. [Figure]
Revised: To alleviate these issues, Lanczos vectors are stored in ﬂoat64 accuracy and we perform reorthogonalized at each step of the algorithm

Original: F T RAINING D ETAILS FOR T ABLEF. N EURAL M ACHINE T RANSLATION
Revised: F T RAINING D ETAILS FOR T ABLEF. N EURAL M ACHINE T RANSLATION

Original: Neural Machine Translation experiments are based on the Transformer models (Vaswani et al., 2017).
Revised: Neural Machine Translation experiments are based on the Transformer models (Vaswani et al., 2017).

Original: We use separate embeddings on encoder and decoder, and a common word piece vocabulary of size 32000.
Revised: We use separate embeddings on encoder and decoder, and a common word piece vocabulary of size 32000.

Original: For depth, we use 6 layers on both encoder and decoder.
Revised: For depth, we use 6 layers on both encoder and decoder.

Original: For width, we experiment with two models, namely Transformer-Base and Transformer-Wide.
Revised: For width, we experiment with two models, namely Transformer-Base and Transformer-Wide.

Original: For Transformer-Base, we use word embeddings with 512 dimensions, 8 heads and 2048 feed-forward dimension.
Revised: For Transformer-Base, we use word embeddings with 512 dimensions, 8 heads and 2048 feed-forward dimension.

Original: For TransformerWide we use word embeddings with 1024 dimensions, 16 heads and 4096 feed-forward dimension.
Revised: For TransformerWide we use word embeddings with 1024 dimensions, 16 heads and 4096 feed-forward dimension.

Original: The experiments reported in Figures 3 and 5 use Transformer-Base.
Revised: The experiments reported in Figures 3 and 5 use Transformer-Base.

Original: The experiments reported in
Revised: The experiments reported in

Original: Table 1 use Transformer-Wide models trained with Adam (Kingma and Ba, 2014).
Revised: Table 1 use Transformer-Wide models trained with Adam (Kingma and Ba, 2014).

Original: 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 1.
Revised: 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 1.

Original: All the models are trained for 60 epochs at batch size of 1024 for Transformer-Base models, and batch size of 512 for Transformer-Big models.
Revised: All the models are trained for 60 epochs at batch size of 1024 for Transformer-Base models, and batch size of 512 for Transformer-Big models.

Original: We use dropout of 0.1, label smoothing of 0.1 and no weight decay for all these models.
Revised: We use dropout of 0.1, label smoothing of 0.1 and no weight decay for all these models.

Original: F. D ENSE N ETS
Revised: F. D ENSE N ETS

Original: In Table 1 the ResNet-50 (w/o BN) architecture was trained for 100 epochs at batch size 512, with l regularization of 5e-5, dropout of .3.
Revised: In Table 1 the ResNet-50 (w/o BN) architecture was trained for 100 epochs at batch size 512, with l regularization of 5e-5, dropout of .3.

Original: It was trained with SGD with nesterov momentum of .9 and learning rate of .2.
Revised: It was trained with SGD with nesterov momentum of .9 and learning rate of .2.

Original: We applied gradient clipping at global l2 norm of 5 and used linear learning rate warmup with warmup period of 1000 steps.
Revised: We applied gradient clipping at global l2 norm of 5 and used linear learning rate warmup with warmup period of 1000 steps.

Original: For Table 1, the DenseNet-100 model was trained using the Gradinit codebase 6 by modifying the supplied DenseNet script to apply gradient clipping of norm 6 and to use the default initialization instead of GradInit.
Revised: For Table 1, the DenseNet-100 model was trained using the Gradinit codebase 6 by modifying the supplied DenseNet script to apply gradient clipping of norm 6 and to use the default initialization instead of GradInit.

Original: G T RAINING D ETAILS FOR F IGURE
Revised: G T RAINING D ETAILS FOR F IGURE

Original: The WideResnet-28-10 models were trained with batch size of 1024 for 300 epochs.
Revised: The WideResnet-28-10 models were trained with batch size of 1024 for 300 epochs.

Original: We applied the
Revised: We applied the

Original: MixUp augmentation(Zhang et al., 2017).
Revised: MixUp augmentation(Zhang et al., 2017).

Original: 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.
Revised: 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.

Original: The  DenseNet models were trained with batch size of 512 using the SGD optimizer with momentum of 0.9, weight decay of 5e-4, L2 regularization of 1e-4 and warmup of 1000 steps (for the models where warmup is used) followed by cosine decay.
Revised: 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 5e-4, L2 regularization of 1e-4 and warmup of 1000 steps (for the models where warmup is used) followed by cosine decay.

Original: The models were trained for 200 epochs.
Revised: The models were trained for 200 epochs.

Original: For the DenseNet architecture we used growth_rate of 32 and reduction of 0.5.
Revised: For the DenseNet architecture we used growth_rate of 32 and reduction of 0.5.

Original: The Resnet-50 models were trained with batch size of 2048 using the SGD optimizer with nesterov momentum of .9.
Revised: The Resnet-50 models were trained with batch size of 2048 using the SGD optimizer with nesterov momentum of .9.

Original: The learning rate schedule was the same as in the WideResnet case, with linear warmup of 1000 steps followed by cosine decay.
Revised: The learning rate schedule was the same as in the WideResnet case, with linear warmup of 1000 steps followed by cosine decay.

Original: We applied label smoothing of .1 and used the standard ﬂip plus crop for data augmentation.
Revised: We applied label smoothing of .1 and used the standard ﬂip plus crop for data augmentation.

Original: The Transformer models on LM1B were trained at batch size 1024 using SGD with nesterov momentum of .9.
Revised: The Transformer models on LM1B were trained at batch size 1024 using SGD with nesterov momentum of .9.

Original: We use embedding dimension of 512, 6 layers with 8 heads and MLP hidden dimension of 1024.
Revised: We use embedding dimension of 512, 6 layers with 8 heads and MLP hidden dimension of 1024.

Original: The attention dropout rate was .1.
Revised: The attention dropout rate was .1.

Original: The learning rate schedule followed the same recipe as in the Resnet cases.
Revised: The learning rate schedule followed the same recipe as in the Resnet cases.

Original: H C URVATURE A DAPTATION WITH THE A DAM O PTIMIZER
Revised: H C URVATURE A DAPTATION WITH THE A DAM O PTIMIZER

Original: The discussion in the main text focused primarily on models trained with SGD and momentum.
Revised: The discussion in the main text focused primarily on models trained with SGD and momentum.

Original: In this appendix, we brieﬂy examine if similar conclusions hold for optimizers such as Adam that use preconditioning.
Revised: In this appendix, we brieﬂy examine if similar conclusions hold for optimizers such as Adam that use preconditioning.

Original: It is unclear a priori whether or not curvature adapation to the learning rate should occur for optimizers which apply preconditioning.
Revised: It is unclear a priori whether or not curvature adapation to the learning rate should occur for optimizers which apply preconditioning.

Original: However, given that Adam is a diagonal preconditioner applied to a non-diagonal Hessian, there may be some similar effects observed.
Revised: However, given that Adam is a diagonal preconditioner applied to a non-diagonal Hessian, there may be some similar effects observed.

Original: Consider a simple quadratic loss
Revised: Consider a simple quadratic loss

Original: [Equation]
Revised: [Equation]

Original: where optimization is performed via preconditioned gradient descent with a ﬁxed diagonal preconditioning matrix D :
Revised: where optimization is performed via preconditioned gradient descent with a ﬁxed diagonal preconditioning matrix D :

Original: [Equation]
Revised: [Equation]

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: While (6) does not take into the account the effects of adaptive preconditioning or momentum, we ﬁnd some empirical evidence that this approximation provides understanding into the stability of the optimization.
Revised: While (6) does not take into the account the effects of adaptive preconditioning or momentum, we ﬁnd some empirical evidence that this approximation provides understanding into the stability of the optimization.

Original: Figure 14 below examines the evolution of λ max ( D − 1 / 2 HD − 1 / 2 ) for three Transformer models trained with Adam and different warm-up lengths.
Revised: Figure 14 below examines the evolution of λ max ( D − 1 / 2 HD − 1 / 2 ) for three Transformer models trained with Adam and different warm-up lengths.

Original: Here, D is a diagonal matrix with D i,i = √
Revised: Here, D is a diagonal matrix with D i,i = √

Original: Corrected Adam grad squared EMA + Adam (cid:15) .
Revised: Corrected Adam grad squared EMA + Adam (cid:15) .

Original: We observe that, similar to the models trained with momentum, the maximum (preconditioned) Hessian eigenvalue adapts to the warm-up schedule (green and red markers).
Revised: We observe that, similar to the models trained with momentum, the maximum (preconditioned) Hessian eigenvalue adapts to the warm-up schedule (green and red markers).

Original: We notice that –perhaps due to the effect of momentum or adaptive preconditioning – the threshold 2 /η does not seem to be aligning with the data well.
Revised: We notice that –perhaps due to the effect of momentum or adaptive preconditioning – the threshold 2 /η does not seem to be aligning with the data well.

Original: Instead, an empirically corrected threshold 40 /η seems to ﬁt the data better.
Revised: Instead, an empirically corrected threshold 40 /η seems to ﬁt the data better.

Original: We observe that instabilities in model training coincide exactly with λ max ( D − 1 / 2 HD − 1 / 2 ) crossing the empirically corrected threshold.
Revised: We observe that instabilities in model training coincide exactly with λ max ( D − 1 / 2 HD − 1 / 2 ) crossing the empirically corrected threshold.

Original: These observations suggest that some of the insights discussed in the main text seem to carry over to the case of adaptive optimizers.
Revised: These observations suggest that some of the insights discussed in the main text seem to carry over to the case of adaptive optimizers.

Original: We leave further exploration of this more complex setting to future work.
Revised: We leave further exploration of this more complex setting to future work.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: I C OMPUTE R ESOURCES U SED
Revised: I C OMPUTE R ESOURCES U SED

Original: Nearly all experiments utilized the Google Cloud Platform with v2 cloud TPUs except for the following: The Figure 2 Resnet-50 and  Densenet experiments utilized the v3 cloudTPU, while the
Revised: Nearly all experiments utilized the Google Cloud Platform with v2 cloud TPUs except for the following: The Figure 2 Resnet-50 and Stride-(1,1) Densenet experiments utilized the v3 cloud

Original:  GradInit code was run on a cloud machine with a single V100 GPU.
Revised: TPU, while the GradInit code was run on a cloud machine with a single V100 GPU.

Original: The Figure 2 experiments were done in parallel using up to 50 v2 TPUs concurrently over the period of a few days.
Revised: The Figureexperiments were done in parallel using up to 50 v2 TPUs concurrently over the period of a few days.

Original: Additionally, all the Machine Translation models were trained on v3 cloud TPUs.
Revised: Additionally, all the Machine Translation models were trained on v3 cloud TPUs.