Original: T HE BREAK - EVEN POINT ONTHE OPTIMIZATION TRAJECTORIES OF DEEP NEURAL NETWORKS
Revised: T HE BREAK - EVEN POINT ON OPTIMIZATION TRAJECTORIES OF DEEP NEURAL NETWORKS

Original: A BSTRACT
Revised: A BSTRACT

Original: The early phase of training of deep neural networks is critical for their ﬁnal performance.
Revised: The early phase of training of deep neural networks is critical for their ﬁnal performance.

Original: In this work, we study how the hyperparameters of stochastic gradient descent (SGD) used in the early phase of training affect the rest of the optimization trajectory.
Revised: In this work, we study how the hyperparameters of stochastic gradient descent (SGD) used in the early phase of training affect the rest of the optimization trajectory.

Original: We argue for the existence of the “break-even" point on this trajectory, beyond which the curvature of the loss surface and noise in the gradient are implicitly regularized by SGD.
Revised: We argue for the existence of the “break-even" point on this trajectory, beyond which the curvature of the loss surface and noise in the gradient are implicitly regularized by SGD.

Original: In particular, we demonstrate on multiple classiﬁcation tasks that using a large learning rate in the initial phase of training reduces the variance of gradients, and improves the conditioning of the covariance of gradients (a matrix that describes noise in the mini-batch gradient).
Revised: In particular, we demonstrate on multiple classiﬁcation tasks that using a large learning rate in the initial phase of training reduces the variance of the gradient, and improves the conditioning of the covariance of gradients.

Original: These effects are beneﬁcial from the optimization perspective and become visible after the breakeven point.
Revised: These effects are beneﬁcial from the optimization perspective and become visible after the break-even point.

Original: Complementing prior work, we also show that using a low learning rate results in bad conditioning of the loss surface even for a neural network with batch normalization layers.
Revised: Complementing prior work, we also show that using a low learning rate results in bad conditioning of the loss surface even for a neural network with batch normalization layers.

Original: In short, our work shows that key properties of the loss surface are strongly inﬂuenced by SGD in the early phase of training.
Revised: In short, our work shows that key properties of the loss surface are strongly inﬂuenced by SGD in the early phase of training.

Original: We argue that studying the impact of the identiﬁed effects on generalization is a promising future direction.
Revised: We argue that studying the impact of the identiﬁed effects on generalization is a promising future direction.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The connection between optimization and generalization of deep neural networks (DNNs) is not fully understood.
Revised: The connection between optimization and generalization of deep neural networks (DNNs) is not fully understood.

Original: For instance, using a large initial learning rate often improves generalization, which can come at the expense of the initial training loss reduction (Goodfellow et al., 2016; Li et al., 2019; Jiang et al., 2020).
Revised: For instance, using a large initial learning rate often improves generalization, which can come at the expense of the initial training loss reduction (Goodfellow et al., 2016; Li et al., 2019; Jiang et al., 2020).

Original: In contrast, using batch normalization layers typically improves both generalization and training speed of deep neural networks (Luo et al., 2019; Bjorck et al., 2018).
Revised: In contrast, using batch normalization layers typically improves both generalization and convergence speed of deep neural networks (Luo et al., 2019; Bjorck et al., 2018).

Original: These simple examples illustratefundamental limitations of our understanding of DNNs.
Revised: These simple examples illustrate limitations of our understanding of DNNs.

Original: Understanding the early phase of training has recently emerged as a promising avenue for studying the link between optimization and generalization of DNNs.
Revised: Understanding the early phase of training has recently emerged as a promising avenue for studying the link between optimization and generalization of DNNs.

Original: It has been observed that  regularization in the early phase of training is necessary to arrive at a well generalizing ﬁnal solution (Keskar et al., 2017; Sagun et al., 2017; Achille et al., 2017).
Revised: It has been observed that applying regularization in the early phase of training is necessary to arrive at a well generalizing ﬁnal solution (Keskar et al., 2017; Sagun et al., 2017; Achille et al., 2017).

Original: Another  phenomenon is that the local shape of the loss surface changes rapidly in the beginning of training (LeCun et al., 2012; Keskar et al., 2017; Achille et al., 2017; Jastrzebski et al., 2018; Fort & Ganguli, 2019).
Revised: Another observed phenomenon is that the local shape of the loss surface changes rapidly in the beginning of training (LeCun et al., 2012; Keskar et al., 2017; Achille et al., 2017; Jastrzebski et al., 2018; Fort & Ganguli, 2019).

Original: Theoretical approaches to understanding deep networks also increasingly focus on the early part of the optimization trajectory (Li et al., 2019; Arora et al., 2019).
Revised: Theoretical approaches to understanding deep networks also increasingly focus on the early part of the optimization trajectory (Li et al., 2019; Arora et al., 2019).

Original: In this work, we study the dependence of the  optimization trajectory on the early phase of training.
Revised: In this work, we study the dependence of the entire optimization trajectory on the early phase of training.

Original: We investigate noise in the mini-batch gradients using the covariance of gradients, 1 and the local 
Revised: We investigate noise in the mini-batch gradients using the covariance of gradients, 1 and the local curvature of the loss surface using the Hessian.

Revised: These two matrices capture important and ∗ Equal contribution.

Original: We deﬁne it as K =  N (cid:80) Ni=
Revised: We deﬁne it as K = 1 N (cid:80) Ni

Original: {+=
+} ( g i − g ) T ( g i − g )  g i =
Revised: {+=
+} ( g i − g ) T ( g i − g ) , g i =

Original: g
Revised: g

Original: ( x i y i ; θ is training curvature of the loss surface using the Hessian.

Original: These two matrices capture important and complementary aspects of optimization (Roux et al., 2008; Ghorbani et al., 2019) and generalization performance of DNNs (Jiang et al., 2020; Keskar et al., 2017; Bjorck et al., 2018; Fort et al., 2019).
Revised: ( x i , y i ; θ is complementary aspects of optimization (Roux et al., 2008; Ghorbani et al., 2019) and generalization performance of DNNs (Jiang et al., 2020; Keskar et al., 2017; Bjorck et al., 2018; Fort et al., 2019).

Original: We include a more detailed discussion in Sec. 2.
Revised: We include a more detailed discussion in Sec. 2.

Original: [Figure]
Revised: [Figure]

Original: Our ﬁrst contribution is a simpliﬁed model of the early part of the training trajectory of DNNs.
Revised: Our ﬁrst contribution is a simpliﬁed model of the early part of the training trajectory of DNNs.

Original: Based on prior empirical work (Sagun et al., 2017), we assume that the local curvature of the loss surface (the spectral norm of the Hessian) increases or decreases monotonically along the optimization trajectory.
Revised: Based on prior empirical work (Sagun et al., 2017), we assume that the local curvature of the loss surface (the spectral norm of the Hessian) increases or decreases monotonically along the optimization trajectory.

Original: Under this model, gradient descent reaches a point in the early phase of training at which it oscillates along the most curved direction of the loss surface.
Revised: Under this model, gradient descent reaches a point in the early phase of training at which it oscillates along the most curved direction of the loss surface.

Original: We call this point the break-even point and show empirical evidence of its existence in the training of actual DNNs.
Revised: We call this point the break-even point and show empirical evidence of its existence in the training of actual DNNs.

Original: Our main contribution is to state and present empirical evidence for two conjectures about the dependence of the  optimization trajectory on the early phase of training.
Revised: Our main contribution is to state and present empirical evidence for two conjectures about the dependence of the entire optimization trajectory on the early phase of training.

Original: Speciﬁcally, we conjecture that the hyperparameters of stochastic gradient descent (SGD) used before reaching the break-even point control: (1) the spectral norms of K and H , and (2) the conditioning of K and H .
Revised: Speciﬁcally, we conjecture that the hyperparameters of stochastic gradient descent (SGD) used before reaching the break-even point control: (1) the spectral norms of K and H , and (2) the conditioning of K and H .

Original: In particular, using a larger learning rate prior to reaching the break-even point reduces the spectral norm of K along the optimization trajectory (see Fig.
Revised: In particular, using a larger learning rate prior to reaching the break-even point reduces the spectral norm of K along the optimization trajectory (see Fig.

Original: 1 for an illustration of this phenomenon).
Revised: 1 for an illustration of this phenomenon).

Original: Reducing the spectral norm of K decreases the variance of the mini-batch gradient, which is linked to improved convergence speed (Johnson & Zhang, 2013).
Revised: Reducing the spectral norm of K decreases the variance of the mini-batch gradient, which has been linked to improved convergence speed (Johnson & Zhang, 2013).

Original: Finally, we apply our analysis to a network with batch normalization (BN) layers and ﬁnd that our predictions are valid in this case as well.
Revised: Finally, we apply our analysis to a network with batch normalization (BN) layers and ﬁnd that our predictions are valid in this case as well.

Original: Delving deeper in this line of investigation, we show that using a large learning rate is necessary to reach better-conditioned (relatively to a network without
Revised: Delving deeper in this line of investigation, we show that using a large learning rate is necessary to reach better-conditioned (relatively to a network without

Original: BN layers) regions of the loss surface, which was previously attributed to BN alone (Bjorck et al., 2018; Ghorbani et al., 2019; Page, 2019).
Revised: BN layers) regions of the loss surface, which was previously attributed to BN alone (Bjorck et al., 2018; Ghorbani et al., 2019; Page, 2019).

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Implicit regularization induced by the optimization method.
Revised: Implicit regularization induced by the optimization method.

Original: The choice of the optimization method implicitly affects generalization performance of deep neural networks (Neyshabur, 2017).
Revised: The choice of the optimization method implicitly affects generalization performance of deep neural networks (Neyshabur, 2017).

Original: In particular, using a large initial learning rate is known to improve generalization (Goodfellow et al., 2016; Li et al., 2019).
Revised: In particular, using a large initial learning rate is known to improve generalization (Goodfellow et al., 2016; Li et al., 2019).

Original: A classical approach to study these questions is to bound the generalization error based on measures such as the norm of the parameters at the ﬁnal minimum (Bartlett et al., 2017; Jiang et al., 2020).
Revised: A classical approach to study these questions is to bound the generalization error using measures such as the norm of the parameters at the ﬁnal minimum (Bartlett et al., 2017; Jiang et al., 2020).

Original: An emerging approach is to study the properties of the whole optimization trajectory (Arora et al., 2019).
Revised: An emerging approach is to study the properties of the whole optimization trajectory.

Original: Arora et al.
Revised: Arora et al.

Original: (2019) suggest it is necessary to study the optimization trajectory to understand optimization and generalization of deep networks.
Revised: (2019) suggest it is necessary to study the optimization trajectory to understand optimization and generalization of deep networks.

Original: Relatedly, Erhan et al. (2010); Achille et al.
Revised: In a related work, Erhan et al. ; Achille et al.

Original: (2017) show the existence of a critical period of learning.
Revised: (2017) show the existence of a critical period of learning.

Original: Erhan et al.
Revised: Erhan et al.

Original: (2010) argue that training, unless pre-training is used, is sensitive to shufﬂing of examples in the ﬁrst epochs of training.
Revised: (2010) argue that training, unless pretraining is used, is sensitive to shufﬂing of examples in the ﬁrst epochs of training.

Original: Achille et al.
Revised: Achille et al.

Original: (2017); Golatkar et al.
Revised: (2017); Golatkar et al.

Original: ; Sagun et al.
Revised: ; Sagun et al.

Original: ; Keskar et al.
Revised: ; Keskar et al.

Original: (2017) demonstrate that adding regularization in the beginning of training affects the ﬁnal generalization disproportionately more compared to doing so later.
Revised: (2017) demonstrate that adding regularization in the beginning of training affects the ﬁnal generalization disproportionately more compared to doing so later.

Original: We continue research in this direction and ask how the choice of hyperparameters in SGD in the early phase of training affects the optimization trajectory in terms of the covariance of gradients, and the Hessian.
Revised: We continue research in this direction and study how the choice of hyperparameters in SGD in the early phase of training affects the optimization trajectory in terms of the covariance of gradients, and the Hessian.

Original: The covariance of gradients and the Hessian.
Revised: The covariance of gradients and the Hessian.

Original: The Hessian quantiﬁes the local curvature of the loss surface.
Revised: The Hessian quantiﬁes the local curvature of the loss surface.

Original: Recent work has shown that the largest eigenvalues of H can grow quickly in the early phase of training (Keskar et al., 2017; Sagun et al., 2017; Fort & Scherlis, 2019; Jastrzebski et al., 2018).
Revised: Recent work has shown that the largest eigenvalues of H can grow quickly in the early phase of training (Keskar et al., 2017; Sagun et al., 2017; Fort & Scherlis, 2019; Jastrzebski et al., 2018).

Original: Keskar et al.
Revised: Keskar et al.

Original: Jastrzebski et al.
Revised: Jastrzebski et al.

Original: (2017) studied the dependence of the Hessian (at the ﬁnal minima) on the optimization hyperparameters.
Revised: (2017) studied the dependence of the Hessian (at the ﬁnal minimum) on the optimization hyperparameters.

Original: The Hessian can be decomposed into two terms, where the dominant term (at least at the end of training) is the uncentered covariance of gradients G (Sagun et al., 2017; Papyan, 2019).
Revised: The Hessian can be decomposed into two terms, where the dominant term (at least at the end of training) is the uncentered covariance of gradients G (Sagun et al., 2017; Papyan, 2019).

Original: The covariance of gradients, which we denote by K , encapsulates the geometry and the magnitude of variation in gradients across different samples.
Revised: The covariance of gradients, which we denote by K , encapsulates the geometry and the magnitude of variation in gradients across different samples.

Original: The matrix K was related to the generalization error in Roux et al.
Revised: The matrix K was related to the generalization error in Roux et al.

Original: ; Jiang et al.
Revised: ; Jiang et al.

Original: Closely related quantities, such as the cosine alignment between gradients computed on different examples, were recently shown to explain some aspects of deep networks generalization (Fort et al., 2019;
Revised: Closely related quantities, such as the cosine alignment between gradients computed on different examples, were recently shown to explain some aspects of deep networks generalization (Fort et al., 2019;

Original: Liu et al., 2020; He & Su, 2020).
Revised: Liu et al., 2020; He & Su, 2020).

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

Original: (2019) argues that in DNNs the Hessian and the covariance of gradients are close in terms of the largest eigenvalues.
Revised: (2019) argues that in DNNs the Hessian and the covariance of gradients are close in terms of the largest eigenvalues.

Original: Learning dynamics of deep neural networks.
Revised: Learning dynamics of deep neural networks.

Original: Our theoretical model is motivated by recent work on learning dynamics of neural networks (Goodfellow et al., 2014; Masters & Luschi, 2018; Wu et al., 2018;
Revised: Our theoretical model is motivated by recent work on learning dynamics of neural networks (Goodfellow et al., 2014; Masters & Luschi, 2018; Wu et al., 2018;

Original: Yao et al., 2018;
Revised: Yao et al., 2018;

Original: Xing et al., 2018; Jastrzebski et al., 2018; Lan et al., 2019).
Revised: Xing et al., 2018; Jastrzebski et al., 2018; Lan et al., 2019).

Original: We are directly inspired by Xing et al.
Revised: We are directly inspired by Xing et al.

Original: (2018) who show that for popular classiﬁcation task, the cosine of the angle between consecutive optimization steps in SGD is negative.
Revised: (2018) who show that for popular classiﬁcation benchmarks, the cosine of the angle between consecutive optimization steps in SGD is negative.

Original: Similar observations can be found in Lan et al.
Revised: Similar observations can be found in Lan et al.

Original: Our theory is closely related to Wu et al.
Revised: Our theoretical analysis is inspired by Wu et al.

Original: (2018) who study how SGD selects the ﬁnal minimum from a stability perspective.
Revised: (2018) who study how SGD selects the ﬁnal minimum from a stability perspective.

Original: We apply their methodology to the early phase of training, and make predictions about the whole training trajectory.
Revised: We apply their methodology to the early phase of training, and make predictions about the entire training trajectory.

Original:  B REAK - EVEN POINT AND  TWO CONJECTURES ABOUT SGD TRAJECTORY
Revised:  T HE BREAK - EVEN POINT AND THE TWO CONJECTURES ABOUT SGD TRAJECTORY

Original: Our overall motivation is to better understand the connection between optimization and generalization of DNNs.
Revised: Our overall motivation is to better understand the connection between optimization and generalization of DNNs.

Original: In this section we study how the covariance of gradients ( K ) and the Hessian ( H ) depend on the early phase of training.
Revised: In this section we study how the covariance of gradients ( K ) and the Hessian ( H ) depend on the early phase of training.

Original: We are inspired by recent empirical observations showing their importance for optimization and generalization of DNNs (see Sec.
Revised: We are inspired by recent empirical observations showing their importance for optimization and generalization of DNNs (see Sec.

Original: 2 for a detailed discussion).
Revised: 2 for a detailed discussion).

Original: Recent work has shown that in the early phase of training the gradient norm (Goodfellow et al., 2016; Fort & Ganguli, 2019; Liu et al., 2020) and the local curvature of the loss surface (Jastrzebski et al., 2018; Fort & Ganguli, 2019) can rapidly increase.
Revised: Recent work has shown that in the early phase of training the gradient norm (Goodfellow et al., 2016; Fort & Ganguli, 2019; Liu et al., 2020) and the local curvature of the loss surface (Jastrzebski et al., 2018; Fort & Ganguli, 2019) can rapidly increase.

Original: Informally speaking, one scenario we study here is when this initial growth is rapid enough to destabilize training.
Revised: Informally speaking, one scenario we study here is when this initial growth is rapid enough to destabilize training.

Original: Inspired by Wu et al.
Revised: Inspired by Wu et al.

Original: (2018), we formalize this intuition using concepts from dynamical stability.
Revised: (2018), we formalize this intuition using concepts from dynamical stability.

Original: Based on the developed theory, we state two conjectures about the dependence of K and H on hyperparameters of SGD, which we investigate empirically in Sec. 4.
Revised: Based on the developed analysis, we state two conjectures about the dependence of K and H on hyperparameters of SGD, which we investigate empirically in Sec. 4.

Original: Deﬁnitions.
Revised: Deﬁnitions.

Original: We begin by introducing the notation.
Revised: We begin by introducing the notation.

Original: Let us denote the loss on an example ( x , y ) by L ( x , y ; θ ) , where θ is a D -dimensional parameter vector.
Revised: Let us denote the loss on an example ( x , y ) by L ( x , y ; θ ) , where θ is a D -dimensional parameter vector.

Original: The two key objects we study are the
Revised: The two key objects we study are the

Original: Hessian of the training loss ( H ), and the covariance of gradients K = 1 N (cid:80) Ni =

Revised: Hessian of the training loss ( H ), and the covariance of gradients K = 1 N (cid:80) Ni =

Original: ( g i − g ) T ( g i − g ) ,
Revised: ( g i − g ) T ( g i − g ) ,

Original: where g i = g ( x i , y i ; θ ) is the gradient of L with respect to θ calculated on i -th example, N is the number of training examples, and g is the full-batch gradient.
Revised: where g i = g ( x i , y i ; θ ) is the gradient of L with respect to θ calculated on i -th example, N is the number of training examples, and g is the full-batch gradient.

Original: We denote the i -th normalized eigenvector and eigenvalue of a matrix A by e iA and λ iA .
Revised: We denote the i -th normalized eigenvector and eigenvalue of a matrix A by e iA and λ iA .

Original: Both H and K are computed at a given θ , but we omit this dependence in the notation.
Revised: Both H and K are computed at a given θ , but we omit this dependence in the notation.

Original: Let t index steps of optimization, and let θ ( t ) denote the parameter vector at optimization step t .
Revised: Let t index steps of optimization, and let θ ( t ) denote the parameter vector at optimization step t .

Original: Inspired by Wu et al.
Revised: Inspired by Wu et al.

Original: we introduce the following condition to quantify stability at a given θ ( t ) .
Revised: we introduce the following condition to quantify stability at a given θ ( t ) .

Original: Let us denote the projection of parameters θ onto e 1 H by ψ = (cid:104) θ, e 1 H (cid:105) .
Revised: Let us denote the projection of parameters θ onto e 1 H by ψ = (cid:104) θ, e 1 H (cid:105) .

Original: With a slight abuse of notation let g ( ψ ) =
Revised: With a slight abuse of notation let g ( ψ ) =

Original: (cid:104) g ( θ ) , e 1 H (cid:105) .
Revised: (cid:104) g ( θ ) , e 1 H (cid:105) .

Original: We say SGD is unstable along e 1 H at θ ( t ) if the norm of elements of sequence ψ ( τ + 1) =
Revised: We say SGD is unstable along e 1 H at θ ( t ) if the norm of elements of sequence ψ ( τ + 1) =

Original: ψ ( τ ) − ηg ( ψ ( τ )) diverges when τ → ∞ , where ψ (0)
Revised: ψ ( τ ) − ηg ( ψ ( τ )) diverges when τ → ∞ , where ψ (0)

Original: =
Revised: =

Original: θ ( t ) .
Revised: θ ( t ) .

Original: The sequence ψ ( τ ) represents optimization trajectory restricted from θ ( t ) only to the direction of e 1 H .
Revised: The sequence ψ ( τ ) represents optimization trajectory in which every step t (cid:48) > t is projected onto e 1 H .

Original: Assumptions.
Revised: Assumptions.

Original: Based on recent empirical studies, we make the following assumptions.
Revised: Based on recent empirical studies, we make the following assumptions.

Original: The loss surface projected onto e 1 H is a quadratic one-dimensional function of the form f ( ψ ) =
Revised: The loss surface projected onto e 1 H is a quadratic one-dimensional function of the form f ( ψ ) =

Original: (cid:80) Ni =1 ( ψ − ψ ∗ ) 2 H
Revised: (cid:80) Ni =1 ( ψ − ψ ∗ ) 2 H

Original: i .
Revised: i .

Original: The same assumption was made in Wu et al.
Revised: The same assumption was made in Wu et al.

Original: (2018), but along all directions in the weight space.
Revised: but for all directions in the weight space.

Original: Alain et al.
Revised: Alain et al.

Original: (2019) show empirically that the loss averaged over all training examples is well approximated by a quadratic function along e 1 H .
Revised: show empirically that the loss averaged over all training examples is well approximated by a quadratic function along e 1 H .

Original: The eigenvectors e 1 H and e 1 K are co-linear, i.e. e 1 H = ±
Revised: The eigenvectors e 1 H and e 1 K are co-linear, i.e. e 1 H = ±

Original: e 1 K , and λ 1 K
Revised: e 1 K , and λ 1 K

Original: = αλ 1 H for some α ∈ R .
Revised: = αλ 1 H for some α ∈ R .

Original: This is inspired by the fact that the top eigenvalues of H can be well approximated using G (non-centered K ) Papyan (2019); Sagun et al.
Revised: This is inspired by the fact that the top eigenvalues of H can be well approximated using G (non-centered K ) (Papyan, 2019; Sagun et al., 2017).

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

Original: (2019) shows direct evidence for co-linearity of the largest eigenvalues of K and H .
Revised: (2019) shows empirical evidence for co-linearity of the largest eigenvalues of K and H .

Original: When in a region that is not stable along e 1 H  , λ 1 H decreases along the training trajectory, and the distance to the minimizer along e 1 H increases.
Revised: If optimization is not stable along e 1 H at a given θ ( t ) , λ 1 H decreases in the next step, and the distance to the minimum along e 1 H increases in the next step.

Original: This is inspired by recent work showing training can escape a region with too large curvature compared to the learning rate (Zhu et al., 2018;
Revised: This is inspired by recent work showing training can escape a region with too large curvature compared to the learning rate (Zhu et al., 2018;

Original: Wu et al., 2018; Jastrzebski et al., 2018).
Revised: Wu et al., 2018; Jastrzebski et al., 2018).

Original: The spectral norm of H , λ 1 H , increases during training and the distance to the minimizer along e 1 H decreases, unless increasing λ 1 H would lead to entering a region where training is not stable along e 1 H .
Revised: The spectral norm of H , λ 1 H , increases during training and the distance to the minimum along e 1 H decreases, unless increasing λ 1 H would lead to entering a region where training is not stable along e 1 H .

Original: This is inspired by (Keskar et al., 2017; Goodfellow et al., 2016;
Revised: This is inspired by (Keskar et al., 2017; Goodfellow et al., 2016;

Original: Sagun et al., 2017; Jastrzebski et al., 2018; Fort & Scherlis, 2019; Fort & Ganguli, 2019) who show that in many settings λ 1 H or gradient norm increases in the beginning of training, while at the same time the overall training loss decreases.
Revised: Sagun et al., 2017; Jastrzebski et al., 2018; Fort & Scherlis, 2019; Fort & Ganguli, 2019) who show that in many settings λ 1 H or gradient norm increases in the beginning of training, while at the same time the overall training loss decreases.

Original: Finally, we also assume that S (cid:29) N , i.e. that the batch size is small compared to the number of training examples.
Revised: Finally, we also assume that S (cid:29) N , i.e. that the batch size is small compared to the number of training examples.

Original: These assumptions are only used to build a theoretical model for the early phase of training.
Revised: These assumptions are only used to build a theoretical model for the early phase of training.

Original: Its main purpose is to make predictions about the training procedure that we test empirically in Sec. 4.
Revised: Its main purpose is to make predictions about the training procedure that we test empirically in Sec. 4.

Original: Reaching the break-even point earlier for a larger learning rate or a smaller batch size.
Revised: Reaching the break-even point earlier for a larger learning rate or a smaller batch size.

Original: Let us restrict ourselves to the case when training is initialized at θ (0) at which SGD is stable along e 1 H (0) .
Revised: Let us restrict ourselves to the case when training is initialized at θ (0) at which SGD is stable along e 1 H (0) .

Original: 2 We aim to show that the learning rate ( η ) and the batch size ( S ) determine H and K in our model, and conjecture that the same holds empirically for actual neural networks.
Revised: 2 We aim to show that the learning rate ( η ) and the batch size ( S ) determine H and K in our model, and conjecture that the same holds empirically for realistic neural networks.

Original: Consider two optimization trajectories for η 1 and η 2 , where η 1 > η 2 , that are initialized at the same θ 0 , where optimization is stable along e 1 H ( t ) and λ 1 H ( t ) > 0 .
Revised: Consider two optimization trajectories for η 1 and η 2 , where η 1 > η 2 , that are initialized at the same θ 0 , where optimization is stable along e 1 H ( t ) and λ 1 H ( t ) > 0 .

Original: Under Assumption 1 the loss surface along e 1 H ( t ) can be expressed as f ( ψ ) =
Revised: Under Assumption 1 the loss surface along e 1 H ( t ) can be expressed as f ( ψ ) =

Original: (cid:80) Ni =1 ( ψ − ψ ∗ ) 2 H
Revised: (cid:80) Ni =1 ( ψ − ψ ∗ ) 2 H

Original: i ( t ) , where H i ( t ) ∈ R .
Revised: i ( t ) , where H i ( t ) ∈ R .

Original: It can be shown that at any iteration t the necessary and sufﬁcient condition for SGD to be stable along e 1 H ( t ) is:
Revised: It can be shown that at any iteration t the necessary and sufﬁcient condition for SGD to be stable along e 1 H ( t ) is:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: ) where N is the training set size and s ( t ) 2 =
Revised: where N is the training set size and s ( t ) 2 =

Original: Var[ H i ( t )] over the training examples.
Revised: Var[ H i ( t )] over the training examples.

Original: A proof can be found in (Wu et al., 2018).
Revised: A proof can be found in (Wu et al., 2018).

Original: We call this point on the trajectory on which the LHS of Eq.
Revised: We call this point on the trajectory on which the LHS of Eq.

Original: 1 becomes equal to 1 for the ﬁrst time the break-even point .
Revised: 1 becomes equal to 1 for the ﬁrst time the break-even point .

Original: By deﬁnition, there exists only a single break-even point on the  trajectory.
Revised: By deﬁnition, there exists only a single break-even point on the training trajectory.

Original: Under Assumption 3, λ 1 H ( t ) and λ 1 K ( t ) increase over time.
Revised: Under Assumption 3, λ 1 H ( t ) and λ 1 K ( t ) increase over time.

Original: If S = N , the break-even point is reached at λ 1 H ( t ) = 2 η .
Revised: If S = N , the break-even point is reached at λ 1 H ( t ) = 2 η .

Original: More generally, it can be shown that for η 1 , the break-even point is reached for a lower magnitude of λ 1 H ( t ) than for η 2 .
Revised: More generally, it can be shown that for η 1 , the break-even point is reached for a lower magnitude of λ 1 H ( t ) than for η 2 .

Original: The same reasoning can be repeated for S (in which case we assume N (cid:29) S ).
Revised: The same reasoning can be repeated for S (in which case we assume N (cid:29) S ).

Original: We state this formally and prove in App.
Revised: We state this formally and prove in App.

Original: B.
Revised: B.

Original: From the break-even point on the  trajectory,under Assumption 4, SGD does not enter regions where either λ 1 H or λ 1 K is larger than at the break-even point, as otherwise it would lead to increasing one of the terms in LHS of Eq. 1, and hence losing stability along e 1 H( t (cid:48) ) .
Revised: Under Assumption 4, after passing the break-even point on the training trajectory, SGD does not enter regions where either λ 1 H or λ 1 K is larger than at the break-even point, as otherwise it would lead to increasing one of the terms in LHS of Eq. 1, and hence losing stability along e 1 H .

Original: Two conjectures about real DNNs.
Revised: Two conjectures about real DNNs.

Original: Assuming that real DNNs reach the break-even point, we make the following two conjectures about their optimization trajectory.
Revised: Assuming that real DNNs reach the break-even point, we make the following two conjectures about their optimization trajectory.

Original: The most direct implication of reaching the break-even point is that λ 1 K and λ 1 H at the break-even point depend on η and S , which we formalize as:
Revised: The most direct implication of reaching the break-even point is that λ 1 K and λ 1 H at the break-even point depend on η and S , which we formalize as:

Original: Conjecture 1 (Variance reduction effect of SGD) .
Revised: Conjecture 1 (Variance reduction effect of SGD) .

Original: Along the SGD trajectory, the maximum attained values of λ 1 H and λ 1 K are smaller for a larger learning rate or a smaller batch size.
Revised: Along the SGD trajectory, the maximum attained values of λ 1 H and λ 1 K are smaller for a larger learning rate or a smaller batch size.

Original: We refer to Conjecture 1 as variance reduction effect of SGD because reducing λ 1 K can be shown to reduce the L 2 distance between the full-batch gradient, and the mini-batch gradient.
Revised: We refer to Conjecture 1 as variance reduction effect of SGD because reducing λ 1 K can be shown to reduce the L 2 distance between the full-batch gradient, and the mini-batch gradient.

Original: We expect that similar effectsshould exist for other optimization or regularization choices, which we leave  for future work.
Revised: We expect that similar effects exist for other optimization or regularization methods. We leave investigating them for future work.

Original: Next, we make another, stronger, conjecture.
Revised: Next, we make another, stronger, conjecture.

Original: It is plausible to assume that reaching the break-even point affects to a lesser degree λ iH and λ iK for i (cid:54) = 1 because increasing their values does not impact stability along e 1 H in our theoretical model.
Revised: It is plausible to assume that reaching the break-even point affects to a lesser degree λ iH and λ iK for i (cid:54) = 1 because increasing their values does not impact stability along e 1 H .

Original: Based on this we conjecture that:
Revised: Based on this we conjecture that:

Original: Conjecture 2 (Pre-conditioning effect of SGD) .
Revised: Conjecture 2 (Pre-conditioning effect of SGD) .

Original: Along the SGD trajectory, the maximum attained values of λ ∗ K λ 1 K and λ ∗ H λ 1 H are larger for a larger learning rate or a smaller batch size, where λ K ∗ and λ H ∗ are the smallest non-zero eigenvalues of K and H , respectively.
Revised: Along the SGD trajectory, the maximum attained values of λ ∗ K λ 1 K and λ ∗ H λ 1 H are larger for a larger learning rate or a smaller batch size, where λ K ∗ and λ H ∗ are the smallest non-zero eigenvalues of K and H , respectively.

Original: Furthermore, the maximum attained values of Tr( K ) and Tr( H ) are smaller for a larger learning rate or a smaller batch size.
Revised: Furthermore, the maximum attained values of Tr( K ) and Tr( H ) are smaller for a larger learning rate or a smaller batch size.

Original: We consider non-zero eigenvalues in the conjecture, because K has at most N − 1 non-zero eigenvalues, where N is the number of training points, which can be drastically smaller than D in overparametrized DNNs.
Revised: We consider non-zero eigenvalues in the conjecture, because K has at most N − 1 non-zero eigenvalues, where N is the number of training points, which can be much smaller than D in overparametrized DNNs.

Original: Both conjectures areonly valid  for learning rates and batch sizes that guarantee that training converges.
Revised: Both conjectures are valid only for learning rates and batch sizes that guarantee that training converges.

Original: Finally, from the optimization perspective the effects discussed above are desirable.
Revised: From the optimization perspective, the effects discussed above are desirable.

Original: Many papers in the optimization literature underline the importance of reducing  variance of the mini-batch gradient (Johnson & Zhang, 2013), or of conditioning of the covariance of gradients (Roux et al., 2008).
Revised: Many papers in the optimization literature underline the importance of reducing the variance of the mini-batch gradient (Johnson & Zhang, 2013) and the conditioning of the covariance of gradients (Roux et al., 2008).

Original: Thereare also  connection between these effects and generalization (Jiang et al., 2020), which we discuss towards the end of the paper.
Revised: There also exists a connection between these effects and generalization (Jiang et al., 2020), which we discuss towards the end of the paper.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section we ﬁrst analyse learning dynamics in the early phase of training.
Revised: In this section we ﬁrst analyse learning dynamics in the early phase of training.

Original: Next, we empirically investigate the two conjectures.
Revised: Next, we empirically investigate the two conjectures.

Original: In the ﬁnal part we extend our analysis to a neural network with batch normalization layers.
Revised: In the ﬁnal part we extend our analysis to a neural network with batch normalization layers.

Original: Due to the space constraint we take the following approach to reporting results. In the main body of the paper, we focus on the  CIFAR-10dataset (Krizhevsky, 2009) and the IMDB dataset (Maas et al., 2011), to which we apply three architectures: a vanilla CNN (SimpleCNN) following Keras example
Revised: We run experiments on the following datasets: CIFAR-10 (Krizhevsky, 2009), IMDB dataset (Maas et al., 2011), ImageNet (Deng et al., 2009), and MNLI (Williams et al., 2018). We apply to these datasets the following architectures: a vanilla CNN (SimpleCNN) following Keras example

Original: (Chollet et al., 2015), ResNet-32 (He et al., 2015),and LSTM (Hochreiter & Schmidhuber, 1997).
Revised: (Chollet et al., 2015), ResNet-32 (He et al., 2015), LSTM (Hochreiter & Schmidhuber, 1997),

Original: We also test the two conjectures for DenseNet (Huang et al., 2016) trained on the ImageNet (Deng et al., 2009) dataset, and BERT (Devlin et al.,2018) ﬁne-tuned on the MNLI dataset (Williams et al., 2018).
Revised: DenseNet (Huang et al., 2016), and BERT (Devlin et al., 2018).

Original: We  includein the Supplement experiments using a multi-layer perceptron trained on the FashionMNIST dataset (Xiao et al., 2017), and experiments using SGD with momentum and learning rate decay.
Revised: We also include experiments using a multi-layer perceptron trained on the FashionMNIST dataset (Xiao et al., 2017) in the Appendix.

Original: All experimental details are described in App.
Revised: All experimental details are described in App.

Original: D.
Revised: D.

Original: Following Dauphin et al.
Revised: Following Dauphin et al.

Original: Alain et al.
Revised: Alain et al.

Original: (2019), we estimate the top eigenvalues and eigenvectors of H on a small subset of the training set (e.g. 5% in the case of CIFAR-10) using the Lanczos[Figure] [Figure] algorithm (Lanczos, 1950).
Revised: (2019), we estimate the top eigenvalues and eigenvectors of H on a small subset of the training set (e.g. 5% in the case of CIFAR-10) using the Lanczos algorithm (Lanczos, 1950).

Original: Computing the full eigenspace of K is infeasible for real DNNs. Instead, we compute the covariance using mini-batch gradients.
Revised: As computing the full eigenspace of K is infeasible for real DNNs, we compute the covariance using mini-batch gradients.

Original: In App.
Revised: In App.

Original: C we argue empirically that (after normalization) this approximates well the largest eigenvalue
Revised: C we show empirically that (after normalization) this approximates well the largest eigenvalue, and we include other details on computing the eigenspaces.

Original: [Figure]
Revised: [Figure]

Original: and we include other details on computing the eigenspaces.
Revised: [Figure]

Original: 4. A CLOSER LOOK AT THE EARLY PHASE OF TRAINING
Revised: 4. A CLOSER LOOK AT THE EARLY PHASE OF TRAINING

Original: First, we examine the learning dynamics in the early phase of training.
Revised: First, we examine the learning dynamics in the early phase of training.

Original: Our goal is to verify some of the assumptions made in Sec.
Revised: Our goal is to verify some of the assumptions made in Sec.

Original: We analyse the evolution of λ 1 H and λ 1 K when using η = 0 .
Revised: We analyse the evolution of λ 1 H and λ 1 K when using η = 0 .

Original: 02 and η = 0 .
Revised: 01 and η = 0 .

Original: 002 to train SimpleCNN on the CIFAR-10 dataset.
Revised: 001 to train SimpleCNN on the CIFAR-10 dataset.

Original: We repeat this experiment 5 times using different initialization.
Revised: We repeat this experiment 5 times using different random initializations of network parameters.

Original: Visualizing the break-even point.
Revised: Visualizing the break-even point.

Original: We visualize the early part of the optimization trajectory in
Revised: We visualize the early part of the optimization trajectory in

Original: Fig.
Revised: Fig.

Original: Following Erhan et al.
Revised: Following Erhan et al.

Original: (2010), we embed the test set predictions at each step of training of SimpleCNN using UMAP (McInnes et al., 2018).
Revised: we embed the test set predictions at each step of training of SimpleCNN using UMAP (McInnes et al., 2018).

Original: Background color indicates λ 1 K (left) and the training accuracy  at the iteration with the closest embedding in Euclidean distance (right).
Revised: The background color indicates λ 1 K (left) and the training accuracy (right) at the iteration with the closest embedding in Euclidean distance.

Original: We observe that the trajectory corresponding to the lower learning rate reaches regions of the loss surface characterized by larger λ 1 K , compared to regions reached at the same training accuracy in the second trajectory.
Revised: We observe that the trajectory corresponding to the lower learning rate reaches regions of the loss surface characterized by larger λ 1 K , compared to regions reached at the same training accuracy in the second trajectory.

Original: Additionally, in Fig.
Revised: Additionally, in Fig.

Original: 3 we plot the spectrum of K (left) and H (right) at the iterations when λ K and λ H respectively reach the highest values.
Revised: 3 we plot the spectrum of K (left) and H (right) at the iterations when λ K and λ H respectively reach the highest values.

Original: We observe more outliers with the lower learning rate in the distributions of both λ K and λ H .
Revised: We observe more outliers for the lower learning rate in the distributions of both λ K and λ H .

Original: Are λ 1 K and λ 1 H correlated in the beginning of training?
Revised: Are λ 1 K and λ 1 H correlated in the beginning of training?

Original: The key assumption behind our theoretical model is that λ 1 K and λ 1 H are correlated, at least prior to reaching the break-even point.
Revised: The key assumption behind our theoretical model is that λ 1 K and λ 1 H are correlated, at least prior to reaching the break-even point.

Original: We conﬁrm this in Fig.
Revised: We conﬁrm this in Fig.

Original: The highest achieved λ 1 K and λ 1 H are larger for the smaller η .
Revised: The highest achieved λ 1 K and λ 1 H are larger for the smaller η .

Original: Additionally, we observe that after achieving the highest value of λ 1 H , further growth of λ 1 K does not translate into an increase of λ 1 H .
Revised: Additionally, we observe that after achieving the highest value of λ 1 H , further growth of λ 1 K does not translate to an increase of λ 1 H .

Original: This is expected: λ 1 H decays to 0 when the mean loss decays to 0 for cross entropy loss (Martens, 2016).[Figure] [Figure] [Figure]
Revised: This is expected as λ 1 H decays to 0 when the mean loss decays to 0 for cross entropy loss (Martens, 2016).

Original: in training speeds.

Original: By default we use

Original: S = 128 in all experiments, η = 0 .

Original: for the CIFAR-10 dataset, and η = 1 . 0 for the IMDB dataset.

Original: Does training become increasingly unstable in the early phase of training?
Revised: Does training become increasingly unstable in the early phase of training?

Original: According to Assumption 3, an increase of λ 1 K and λ 1 H translates into a decrease in stability, which we formalized as stability along e 1 H .
Revised: According to Assumption 3, an increase of λ 1 K and λ 1 H translates into a decrease in stability, which we formalized as stability along e 1 H .

Original: Computingdirectly stability along e 1 H  is computationally expensive.
Revised: Computing stability along e 1 H directly is computationally expensive.

Original: Instead, we measure a more tractable proxy.
Revised: Instead, we measure a more tractable proxy.

Original: At each iteration we measure the loss on the training set before and  after taking the step, which we denote as ∆ L (a positive value indicates a reduction of the training loss).
Revised: At each iteration we measure the loss on the training set before and [Figure] after taking the step, which we denote as ∆ L (a positive value indicates a reduction of the training loss).

Original: In Fig.
Revised: In Fig.

Original: 2 we observe that training becomes  unstable ( ∆ L starts to be increasingly oscillating) around the time λ 1 K reaches the maximum value.
Revised: 2 we observe that training becomes increasingly unstable ( ∆ L starts to take negative values) as λ 1 K reaches the maximum value.

Original: Summary.
Revised: Summary.

Original: We have shown that the early phase of training is consistent with the assumptions made in our theoretical model.
Revised: We have shown that the early phase of training is consistent with the assumptions made in our theoretical model.

Original: That is, λ 1 K and λ 1 H increase approximately proportionally to each other, which is also  correlated with a decrease of a proxy of stability.
Revised: That is, λ 1 K and λ 1 H increase approximately proportionally to each other, which is also generally correlated with a decrease of a proxy of stability.

Original: Finally, we have shown qualitatively reaching the break-even point.
Revised: Finally, we have shown qualitatively reaching the break-even point.

Original: 4. T HE VARIANCE REDUCTION AND THE PRE - CONDITIONING EFFECT OF SGD
Revised: 4. T HE VARIANCE REDUCTION AND THE PRE - CONDITIONING EFFECT OF SGD

Original: In this section we test empirically Conjecture 1 and Conjecture 2.
Revised: In this section we test empirically Conjecture 1 and Conjecture 2.

Original: For each model wepick manually  a suitable range of learning rates and batch sizes to ensure that the properties of K and H that we study have converged under a reasonable computational budget.
Revised: For each model we manually pick a suitable range of learning rates and batch sizes to ensure that the properties of K and H that we study have converged under a reasonable computational budget.

Original: We mainly focus on studying the covariance of gradients ( K ), and leave a closer investigation of the Hessian for future work.
Revised: We mainly focus on studying the covariance of gradients ( K ), and leave a closer investigation of the Hessian for future work.

Original: We use  batch size of 128 to compute K when we vary the batch size for training.
Revised: We use the batch size of 128 to compute K when we vary the batch size for training.

Original: When we vary the learning rate instead, we use the same batch size as the one used to train the model.
Revised: When we vary the learning rate instead, we use the same batch size as the one used to train the model.

Original: App.
Revised: App.

Original: C describes the remaining details on how we approximate the eigenspaces of K and H .
Revised: C describes the remaining details on how we approximate the eigenspaces of K and H .

Original: [Figure] We summarize the results for SimpleCNN, ResNet-32, LSTM, BERT, and DenseNet in Fig.
Revised: We summarize the results for SimpleCNN, ResNet-32, LSTM, BERT, and DenseNet in Fig.

Original: 4, Fig. 5, and Fig.
Revised: 4, Fig. 5, and Fig.

Original: Curves are smoothed using moving average forvisual clarity.
Revised: Curves are smoothed using moving average for clarity.

Original: Training curves and additional experiments are reported in App.
Revised: Training curves and additional experiments are reported in App.

Original: E.
Revised: E.

Original: Testing Conjecture 1.
Revised: Testing Conjecture 1.

Original: To test Conjecture 1, we examine the highest value of λ 1 K observed along the optimization trajectory.
Revised: To test Conjecture 1, we examine the highest value of λ 1 K observed along the optimization trajectory.

Original: As visible in Fig. 4, using a higher η results in λ 1 K achieving a lower  maximum during training.
Revised: As visible in Fig. 4, using a higher η results in λ 1 K achieving a lower [Figure] maximum during training.

Original: Similarly, we observe that using a higher S in SGD leads to reaching a higher maximum value of λ 1 K .
Revised: Similarly, we observe that using a higher S in SGD leads to reaching a higher maximum value of λ 1 K .

Original: For instance, for SimpleCNN (top row of Fig. 4) we observe max( λ 1 K ) = 0 .
Revised: For instance, for SimpleCNN (top row of Fig. 4) we observe max( λ 1 K ) = 0 .

Original: 68 and max( λ 1 K )
Revised: 68 and max( λ 1 K )

Original: = 3 . 30 for η = 0 .
Revised: = 3 . 30 for η = 0 .

Original: 1 and η =0 .
Revised: 1 and η =

Original: 01 , respectively.
Revised: 01 , respectively.

Original: Testing Conjecture 2.
Revised: Testing Conjecture 2.

Original: To test Conjecture 2, we compute the maximum value of λ ∗ K /λ 1 K along the optimization trajectory.
Revised: To test Conjecture 2, we compute the maximum value of λ ∗ K /λ 1 K along the optimization trajectory.

Original: It is visible in Fig.
Revised: It is visible in Fig.

Original: 4 that using a higher η results in reaching a larger maximum value of λ ∗ K /λ 1 K along the trajectory.
Revised: 4 that using a higher η results in reaching a larger maximum value of λ ∗ K /λ 1 K along the trajectory.

Original: For instance, in the case of SimpleCNN max( λ ∗ K /λ 1 K ) = 0 . 37 and max( λ ∗ K /λ 1 K ) = 0 .
Revised: For instance, in the case of SimpleCNN max( λ ∗ K /λ 1 K ) = 0 . 37 and max( λ ∗ K /λ 1 K ) = 0 .

Original: 24 for η = 0 .
Revised: 24 for η = 0 .

Original: 1 and η = 0 .
Revised: 1 and η = 0 .

Original: 01 , respectively.
Revised: 01 , respectively.

Original: A counter-intuitive effect of decreasing the batch size.
Revised: A counter-intuitive effect of decreasing the batch size.

Original: Consistently with Conjecture 2, we observe that the maximum value of Tr( K ) is smaller for the smaller batch size.
Revised: Consistently with Conjecture 2, we observe that the maximum value of Tr( K ) is smaller for the smaller batch size.

Original: In the case of SimpleCNN max(Tr( K )) = 5 . 56 and max(Tr( K ))
Revised: In the case of SimpleCNN max(Tr( K )) = 5 . 56 and max(Tr( K ))

Original: = 10 .
Revised: = 10 .

Original: 86 for S = 10 and S = 100 , respectively.
Revised: 86 for S = 10 and S = 100 , respectively.

Original: Due to space constraints we report the effect of η and S on Tr( K ) in other settings in App.
Revised: Due to space constraints we report the effect of η and S on Tr( K ) in other settings in App.

Original: E.
Revised: E.

Original: This effect is counter-intuitive because Tr( K ) is proportional to the variance of the mini-batch gradient (see also App. C).
Revised: This effect is counter-intuitive because Tr( K ) is proportional to the variance of the mini-batch gradient (see also App. C).

Original: Naturally, using a lower batch size generally increases the variance of the mini-batch gradient, and Tr( K ) .
Revised: Naturally, using a lower batch size generally increases the variance of the mini-batch gradient, and Tr( K ) .

Original: This seeming contradiction is explained by the fact that we measure Tr( K ) using a different batch size ( 128 ) than the one used to train the model.
Revised: This apparent contradiction is explained by the fact that we measure Tr( K ) using a different batch size ( 128 ) than the one used to train the model.

Original: Hence, decreasing the batch size both increases (due to approximating the gradient using fewer examples) and decreases (as predicted in Conjecture 2) the variance of the mini-batch gradient along the optimization trajectory.
Revised: Hence, decreasing the batch size both increases (due to approximating the gradient using fewer samples) and decreases (as predicted in Conjecture 2) the variance of the mini-batch gradient along the optimization trajectory.

Original: How early in training is the break-even point reached?
Revised: How early in training is the break-even point reached?

Original: We ﬁnd that λ 1 K and λ 1 H reach their highest values early in training, close to reaching 60% training accuracy on CIFAR-10, and 75% training accuracy on IMDB.
Revised: We ﬁnd that λ 1 K and λ 1 H reach their highest values early in training, close to reaching 60% training accuracy on CIFAR-10, and 75% training accuracy on IMDB.

Original: The training and validation accuracies are reported for all the experiments in App.
Revised: The training and validation accuracies are reported for all the experiments in App.

Original: E. This suggests that the break-even point is reached early in training.
Revised: E. This suggests that the break-even point is reached early in training.

Original: The Hessian.
Revised: The Hessian.

Original: In the above, we have focused on the covariance of gradients.
Revised: In the above, we have focused on the covariance of gradients.

Original: In Fig.
Revised: In Fig.

Original: 5 we report how λ 1 H depends on η and S for ResNet-32 and SimpleCNN.
Revised: 5 we report how λ 1 H depends on η and S for ResNet-32 and SimpleCNN.

Original: Consistently with prior work (Keskar et al., 2017; Jastrzebski et al., 2018), we observe that using a smaller η or using a larger S coincides with a larger peak value of λ 1 H .
Revised: Consistently with prior work (Keskar et al., 2017; Jastrzebski et al., 2018), we observe that using a smaller η or using a larger S coincides with a larger maximum value of λ 1 H .

Original: For instance, for SimpleCNN we observe max( λ 1 H ) = 26 . 27 and max( λ 1 H ) = 211 . 17 for η = 0 .
Revised: For instance, for SimpleCNN we observe max( λ 1 H ) = 26 .and max( λ 1 H ) = 211 . 17 for η = 0 .

Original: 1 and η = 0 .
Revised: 1 and η = 0 .

Original: 01 , respectively.
Revised: 01 , respectively.

Original: We leavefor future work testing predictions made in Conjecture 2 about the Hessian.
Revised: We leave testing predictions made in Conjecture 2 about the Hessian for future work.

Original: Larger scale studies.
Revised: Larger scale studies.

Original: Finally, we test the two conjectures on two larger scale settings: BERT ﬁne-tuned on the MNLI dataset, and DenseNet trained on the ImageNet dataset.
Revised: Finally, we test the two conjectures in two larger scale settings: BERT ﬁne-tuned on the MNLI dataset, and DenseNet trained on the ImageNet dataset.

Original: Due to memory constraints, we only vary the learning rate.
Revised: Due to memory constraints, we only vary the learning rate.

Original: We report results in Fig.
Revised: We report results in Fig.

Original: We observe that both[Figure] conjectures hold in these two settings.
Revised: We observe that both conjectures hold in these two settings.

Original: It is worth noting that DenseNet uses batch normalization layers.
Revised: It is worth noting that DenseNet uses batch normalization layers.

Original: In the next section we investigate closer batch-normalized networks.
Revised: In the next section we investigate closer batch-normalized networks.

Original: Summary.
Revised: Summary.

Original: In this section we have shown evidence supporting the variance reduction (Conjecture 1) and the pre-conditioning effect (Conjecture 2) of SGD on awide range of classiﬁcation tasks [Figure]
Revised: In this section we have shown evidence supporting the variance reduction (Conjecture 1) and the pre-conditioning effect (Conjecture 2) of SGD in a range of classiﬁcation tasks.

Original: We also found that the above conclusions hold for MLP trained on the Fashion MNIST dataset, SGD with momentum [Figure] and SGD with learning rate decay.
Revised: We also found that the above conclusions hold for MLP trained on the Fashion MNIST dataset, SGD with momentum, and SGD with learning rate decay.

Original: We include these results in App.
Revised: We include these results in App.

Original: E-G.
Revised: E-G.

Revised: [Figure]

Original: 4. I MPORTANCE OF LEARNING RATE FOR CONDITIONING INBATCH NORMALIZED NETWORKS 
Revised: 4. I MPORTANCE OF LEARNING RATE FOR CONDITIONING IN NETWORKS WITH BATCH NORMALIZATION LAYERS

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

Original: The loss surface of deep neural networks has been widely reported to be ill-conditioned (LeCun et al., 2012; Martens, 2016).
Revised: The loss surface of deep neural networks has been widely reported to be ill-conditioned (LeCun et al., 2012; Martens, 2016).

Original: Recently, Ghorbani et al.
Revised: Recently, Ghorbani et al.

Original: (2019); Page (2019)have argued that the key reason behind the efﬁcacy of batch normalization (Ioffe & Szegedy, 2015) is improving conditioning of the loss surface.
Revised: ; Page (2019) argued that the key reason behind the efﬁcacy of batch normalization (Ioffe & Szegedy, 2015) is improving conditioning of the loss surface.

Original: In Conjecture 2 we propose that using a high η (or a small S ) results in improving the condition number of K and H .
Revised: In Conjecture 2 we suggest that using a high η (or a small S ) results in improving the conditioning of K and H .

Original: A natural question that we investigate in this section is how the two phenomena are related.
Revised: A natural question that we investigate in this section is how the two phenomena are related.

Original: We study here the effect of learning rate, and report in App.
Revised: We study here the effect of learning rate, and report in App.

Original: H a similar study for batch size.
Revised: H an analogous study for batch size.

Original: Are the two conjectures valid in batch-normalized networks?
Revised: Are the two conjectures valid in networks with batch normalization layers?

Original: First, tofurther investigate whether our conjectures hold in batch-normalized network, we run similar experiments as in Sec.4.
Revised: First, to investigate whether our conjectures hold in networks with batch normalization layers, we run similar experiments as in Sec.

Original: on a SimpleCNN model with batch normalization layers inserted after each layer (SimpleCNN-BN), on the CIFAR-10 dataset.
Revised: 4.2 with a SimpleCNN model with batch normalization layers inserted after each layer (SimpleCNN-BN), on the CIFAR-10 dataset.

Original: We test η ∈ { 0 .
Revised: We test η ∈ { 0 .

Original: 001 , 0 . 01 , 0 . 1 , 1 . 0 } ( η = 1 . 0{+
+} leads to divergence of SimpleCNN without BN).
Revised: 001 , 0 . 01 , 0 . 1 , 1 . 0 } (using η = 1 . {+
+} leads to divergence of SimpleCNN without BN).

Original: We summarize the results in Fig.
Revised: We summarize the results in Fig.

Original: We observe that the evolution of λ ∗ K /λ 1 K and λ 1 K is consistent with both Conjecture 1 and Conjecture 2.
Revised: We observe that the evolution of λ ∗ K /λ 1 K and λ 1 K is consistent with both Conjecture 1 and Conjecture 2.

Original: A closer look at the early phase of training.
Revised: A closer look at the early phase of training.

Original: To further corroborate that our analysis applies to batch-normalized networks, we study the early phase of training of SimpleCNN-BN, complementing the results in Sec. 4.1.
Revised: To further corroborate that our analysis applies to networks with batch normalization layers, we study the early phase of training of SimpleCNN-BN, complementing the results in Sec. 4.1.

Original: We observe in Fig.
Revised: We observe in Fig.

Original: 7 (bottom) that training of SimpleCNN-BN starts in a region characterized by a relatively high λ 1 K .
Revised: 7 (bottom) that training of SimpleCNN-BN starts in a region characterized by a relatively high λ 1 K .

Original: This is consistent with prior work showing thatbatch-normalized networks  can exhibit gradient explosion in the ﬁrst iteration (Yang et al., 2019).
Revised: This is consistent with prior work showing that networks with batch normalization layers can exhibit gradient explosion in the ﬁrst iteration (Yang et al., 2019).

Original: The value of λ 1 K then decays for all but the lowest η .
Revised: The value of λ 1 K then decays for all but the lowest η .

Original: This behavior is consistent with our theoretical model.
Revised: This behavior is consistent with our theoretical model.

Original: We also track the norm of the scaling factor in  batch normalization layers, (cid:107) γ (cid:107) , in the last layer of the network in Fig.
Revised: We also track the norm of the scaling factor in the batch normalization layers, (cid:107) γ (cid:107) , in the last layer of the network in Fig.

Original: 7 (bottom).
Revised: 7 (bottom).

Original: It is visible that η = 1 . 

Revised: It is visible that η = 1 .

Original: and η = 0 .
Revised: and η = 0 .

Original: 1 initially decrease the value of (cid:107) γ (cid:107) , which we hypothesize to be one of the mechanisms by which high η steers optimization towards better conditioned regions of the loss surface in batch-normalized networks.
Revised: 1 initially decrease the value of (cid:107) γ (cid:107) , which we hypothesize to be one of the mechanisms due to which high η steers optimization towards better conditioned regions of the loss surface in batch-normalized networks.

Original: Interestingly, this seems consistent with Luo et al.
Revised: Interestingly, this seems consistent with Luo et al.

Original: (2019) who argue that using mini-batch statistics in batch normalization acts as an implicit regularizer by reducing (cid:107) γ (cid:107) .
Revised: (2019) who argue that using mini-batch statistics in batch normalization acts as an implicit regularizer by reducing (cid:107) γ (cid:107) .

Original: Using batch normalization requires using a high learning rate.
Revised: Using batch normalization requires using a high learning rate.

Original: As our conjectures hold for
Revised: As our conjectures hold for

Original: SimpleCNN-BN, a natural question is if learning can be ill-conditioned with a low learning rate even when batch normalization is used.
Revised: SimpleCNN-BN, a natural question is if the loss surface can be ill-conditioned with a low learning rate even when batch normalization is used.

Original: Ghorbani et al.
Revised: Ghorbani et al.

Original: (2019) show that without batch normalization, mini-batch gradients are largely contained in the subspace spanned by the top eigenvectors of noncentered K .
Revised: (2019) show that without batch normalization, mini-batch gradients are largely contained in the subspace spanned by the top eigenvectors of noncentered K .

Original: To answer this question we track (cid:107) g
Revised: To answer this question we track (cid:107) g

Original: (cid:107) / (cid:107) g 5 (cid:107) , where g denotes the mini-batch gradient, and g 5 denotes the mini-batch gradient projected onto the top 5 eigenvectors of K .
Revised: (cid:107) / (cid:107) g 5 (cid:107) , where g denotes the mini-batch gradient, and g 5 denotes the mini-batch gradient projected onto the top 5 eigenvectors of K .

Original: A value of (cid:107) g (cid:107) / (cid:107) g 5 (cid:107) close to 1 implies that the mini-batch gradient is mostly contained in the subspace spanned by the top 5 eigenvectors of K .
Revised: A value of (cid:107) g (cid:107) / (cid:107) g 5 (cid:107) close to 1 implies that the mini-batch gradient is mostly contained in the subspace spanned by the top 5 eigenvectors of K .

Original: We compare two settings:(1) SimpleCNN-BN optimized with η = 0 .
Revised: We compare two settings: SimpleCNN-BN optimized with η = 0 .

Original: 001 , and(2) SimpleCNN optimized with η = 0 .
Revised: 001 , and SimpleCNN optimized with η = 0 .

Original: We make three observations.
Revised: We make three observations.

Original: First, the maximum (minimum) value of (cid:107) g (cid:107) / (cid:107) g 5 (cid:107) is 1 .
Revised: First, the maximum and minimum values of (cid:107) g (cid:107) / (cid:107) g 5 (cid:107) are 1 .

Original: 90 ( 1 . 37 ) and 2 .
Revised: 90 ( 1 . 37 ) and 2 .

Original: 02 ( 1 . 09 ), respectively.
Revised: 02 ( 1 . 09 ), respectively.

Original: Second, the maximum value of λ 1 K is 12 .and 3 . 30 ,respectively.
Revised: Second, the maximum and minimum values of λ 1 K are 12 . 05 and 3 . 30 ,

Revised: respectively.

Original: Finally, λ ∗ K /λ 1 K reaches 0 .
Revised: Finally, λ ∗ K /λ 1 K reaches 0 .

Original: 343 in the ﬁrst setting, and 0 .
Revised: 343 in the ﬁrst setting, and 0 .

Original: 24 in the second setting.
Revised: 24 in the second setting.

Original: Comparing these differences to differences that are induced by using the highest η = 1 .in SimpleCNN-BN, we can conclude that using a large learning rate is necessary to observe the effect of loss smoothing which was previously attributed to batch normalization alone (Ghorbani et al., 2019; Page, 2019; Bjorck et al., 2018).
Revised: Comparing these differences to differences that are induced by using the highest η = 1 . 0 in SimpleCNN-BN, we can conclude that using a large learning rate is necessary to observe the effect of loss smoothing which was previously attributed to batch normalization alone (Ghorbani et al., 2019; Page, 2019; Bjorck et al., 2018).

Original: This might be directly related to the result that using a high learning rate is necessary to achieve good generalization when using batch normalization layers (Bjorck et al., 2018).
Revised: This might be directly related to the result that using a high learning rate is necessary to achieve good generalization when using batch normalization layers (Bjorck et al., 2018).

Original: Summary.
Revised: Summary.

Original: We have shown that the effects of the learning rate predicted in Conjecture 1 and
Revised: We have shown that the effects of the learning rate predicted in Conjecture 1 and

Original: Conjecture 2 hold for a network with batch normalization layers, and that using a high learning rate is necessary in abatch-normalized network  to improve conditioning of the loss surface, compared to conditioning of the loss surface in the same network without batch normalization layers.
Revised: Conjecture 2 hold for a network with batch normalization layers, and that using a high learning rate is necessary in a network with batch normalization layers to improve conditioning of the loss surface, compared to conditioning of the loss surface in the same network without batch normalization layers.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Based on our theoretical model, we argued for the existence of the break-even point on the optimization trajectory induced by SGD.
Revised: Based on our theoretical model, we argued for the existence of the break-even point on the optimization trajectory induced by SGD.

Original: We presented evidence that hyperparameters used in the early phase of training control the spectral norm and the conditioning of K (a matrix describing noise in the mini-batch gradients) and H (a matrix describing local curvature of the loss surface) after reaching the break-even point.
Revised: We presented evidence that hyperparameters used in the early phase of training control the spectral norm and the conditioning of K (a matrix describing noise in the mini-batch gradients) and H (a matrix describing local curvature of the loss surface) after reaching the break-even point.

Original: In particular, using a large initial learning rate steers optimization to better conditioned regions of the loss surface, which is beneﬁcial from the optimization point of view.
Revised: In particular, using a large initial learning rate steers training to better conditioned regions of the loss surface, which is beneﬁcial from the optimization point of view.

Original: A natural direction for the future is connecting our observations to recent studies on the relation of measures, such as gradient variance, to the generalization of deep networks (Li et al., 2019; Jiang et al., 2020; Fort et al., 2019).
Revised: A natural direction for the future is connecting our observations to recent studies on the relation of measures, such as gradient variance, to the generalization of deep networks (Li et al., 2019; Jiang et al., 2020; Fort et al., 2019).

Original: Our work shows that the hyperparameters of SGD control these measures after the break-even point.
Revised: Our work shows that the hyperparameters of SGD control these measures after the break-even point.

Original: Another interesting direction is to understand the connection between the existence of the break-even point and the existence of the critical learning period in training of DNNs (Achille et al., 2017).
Revised: Another interesting direction is to understand the connection between the existence of the break-even point and the existence of the critical learning period in training of DNNs (Achille et al., 2017).

Original: A CKNOWLEDGMENTS
Revised: A CKNOWLEDGMENTS

Original: KC thanksthe support by NVIDIA and eBay.
Revised: KC thanks NVIDIA and eBay for their support.

Original: SJ thanks Amos Storkey and Luke Darlow for fruitful discussions.
Revised: SJ thanks Amos Storkey and Luke Darlow for fruitful discussions.

Original: A A DDITIONAL VISUALIZATION OF THE BREAK - EVEN POINT
Revised: A A DDITIONAL VISUALIZATION OF THE BREAK - EVEN POINT

Original: We include here a similar plot as Fig. 1, but visualizing the conditioning of the covariance of gradients ( λ ∗ K /λ 1 K , left), the trace of the covariance of gradients ( Tr( K ) , middle), and the spectral norm of the Hessian ( λ 1 H , right).
Revised: We include here an analogous Figure to Fig. 1, but visualizing the conditioning of the covariance of gradients ( λ ∗ K /λ 1 K , left), the trace of the covariance of gradients ( Tr( K ) , middle), and the spectral norm of the Hessian ( λ 1 H , right).

Original: The choice of the learning rate impacts the regions of the loss surface reached by SGD in terms of these quantities.

Original: [Figure]
Revised: [Figure]

Original: B P ROOFS
Revised: B P ROOFS

Original: In this section we formally state and prove the theorem used in Sec. 3.
Revised: In this section we formally state and prove the theorem used in Sec. 3.

Original: With the deﬁnitions introduced in Sec. 3 in mind, we state the following:
Revised: With the deﬁnitions introduced in Sec. 3 in mind, we state the following:

Original: Theorem 1.
Revised: Theorem 1.

Original: Assuming that training is stable along e 1 H ( t ) at t = 0 , then λ 1 H ( t ) and λ 1 K ( t ) at which
Revised: Assuming that training is stable along e 1 H ( t ) at t = 0 , then λ 1 H ( t ) and λ 1 K ( t ) at which

Original: SGD becomes unstable along e 1 H ( t ) are smaller for a larger η or a smaller S .
Revised: SGD becomes unstable along e 1 H ( t ) are smaller for a larger η or a smaller S .

Original: Proof.
Revised: Proof.

Original: This theorem is a straightforward application of Theorem 1 from Wu et al.
Revised: This theorem is a straightforward application of Theorem 1 from Wu et al.

Original: (2018) to the early phase of training.
Revised: (2018) to the early phase of training.

Original: First, let us consider two optimization trajectories corresponding to using two different learning rates η 1 and η 2 ( η 1 > η 2 ).
Revised: First, let us consider two optimization trajectories corresponding to using two different learning rates η 1 and η 2 ( η 1 > η 2 ).

Original: For both trajectories, Theorem 1 of Wu et al.
Revised: For both trajectories, Theorem 1 of Wu et al.

Original: (2018) states that SGD is stable at an iteration t if the following inequality is satisﬁed:
Revised: (2018) states that SGD is stable at an iteration t if the following inequality is satisﬁed:

Original: [Equation]
Revised: [Equation]

Original: where s ( t ) =
Revised: where s ( t ) =

Original: Var[H i (t)] .
Revised: Var[H i (t)] .

Original: Using Assumption 1 we get Var[ g i ( ψ )] =
Revised: Using Assumption 1 we get Var[ g i ( ψ )] =

Original: Var[ ψH i ] =
Revised: Var[ ψH i ] =

Original: ψ 2 s ( t ) 2 .
Revised: ψ 2 s ( t ) 2 .

Original: Using this we can rewrite inequality (2) as:
Revised: Using this we can rewrite inequality (2) as:

Original: [Equation]
Revised: [Equation]

Original: To calculate the spectral norm of the Hessian at iteration t ∗ at which training becomes unstable we equate the two sides of inequality (3) and solve for λ 1 H ( t ∗ ) , which gives
Revised: To calculate the spectral norm of the Hessian at iteration t ∗ at which training becomes unstable we equate the two sides of inequality (3) and solve for λ 1 H ( t ∗ ) , which gives

Original: [Equation]
Revised: [Equation]

Original: where α denotes the proportionality constant from Assumption 2.
Revised: where α denotes the proportionality constant from Assumption 2.

Original: Note that if n = S the right hand side degenerates to 2 η , which completes the proof for n = S .
Revised: Note that if n = S the right hand side degenerates to 2 η , which completes the proof for n = S .

Original: To prove the general case, let denote by λ 1 H ( t ∗ 1 ) and λ 1 H
Revised: To prove the general case, let denote by λ 1 H ( t ∗ 1 ) and λ 1 H

Original: ( t ∗ 2 ) the value of λ 1 H at which training becomes unstable along e 1 H for η 1 and η 2 , respectively.
Revised: ( t ∗ 2 ) the value of λ 1 H at which training becomes unstable along e 1 H for η 1 and η 2 , respectively.

Original: Similarly, let us denote by ψ ( t ∗ 1 ) and ψ ( t ∗ 2 ) the corresponding values of ψ .
Revised: Similarly, let us denote by ψ ( t ∗ 1 ) and ψ ( t ∗ 2 ) the corresponding values of ψ .

Original: Let us assume by contradiction that λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .
Revised: Let us assume by contradiction that λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .

Original: A necessary condition for this inequality to hold is that ψ ( t ∗ 1 ) > ψ ( t ∗ 2 ) .
Revised: A necessary condition for this inequality to hold is that ψ ( t ∗ 1 ) > ψ ( t ∗ 2 ) .

Original: However, Assumption 4 implies that prior to reaching the break-even point ψ ( t ) decreases monotonically with increasing λ 1 H ( t ∗ 1 ) , which corresponds to reducing distance to the minimizer along e 1 H , which contradicts λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .
Revised: However, Assumption 4 implies that prior to reaching the break-even point ψ ( t ) decreases monotonically with increasing λ 1 H ( t ∗ 1 ) , which corresponds to reducing distance to the minimum along e 1 H , which contradicts λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .

Original: Repeating the same argument for two trajectories optimized using two different batch sizes completes the proof.
Revised: Repeating the same argument for two trajectories optimized using two different batch sizes completes the proof.

Original: It is also straightforward to extend the argument to the case when training is initialized at an unstable region along e 1 H (0) , which we formalize as follows.
Revised: It is also straightforward to extend the argument to the case when training is initialized at an unstable region along e 1 H (0) , which we formalize as follows.

Original: Theorem 2.
Revised: Theorem 2.

Original: If training is unstable along e 1 H ( t ) at t = 0 , then λ 1 H ( t ) and λ 1 K ( t ) at which SGD becomes for the ﬁrst time stable along e 1 H ( t ) are smaller for a larger η or a smaller S .
Revised: If training is unstable along e 1 H ( t ) at t = 0 , then λ 1 H ( t ) and λ 1 K ( t ) at which SGD becomes for the ﬁrst time stable along e 1 H ( t ) are smaller for a larger η or a smaller S .

Original: Proof.
Revised: Proof.

Original: We will use a similar argument as in the proof of Th.
Revised: We will use a similar argument as in the proof of Th.

Original: Let us consider two optimization trajectories corresponding to using two different learning rates η 1 and η 2 ( η 1 > η 2 ).
Revised: Let us consider two optimization trajectories corresponding to using two different learning rates η 1 and η 2 ( η 1 > η 2 ).

Original: Let λ 1 H ( t ∗ 1 ) and λ 1 H ( t ∗ 2 ) denote the spectral norm of H at the iteration at which training is for the ﬁrst time stable along e 1 H for η 1 and η 2 , respectively.
Revised: Let λ 1 H ( t ∗ 1 ) and λ 1 H ( t ∗ 2 ) denote the spectral norm of H at the iteration at which training is for the ﬁrst time stable along e 1 H for η 1 and η 2 , respectively.

Original: Following the steps in the proof of Th.
Revised: Following the steps in the proof of Th.

Original: 1 we get that
Revised: 1 we get that

Original: [Equation]
Revised: [Equation]

Original: Using the same notation as in the proof of Th.
Revised: Using the same notation as in the proof of Th.

Original: 1, let us assume by contradiction that λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .
Revised: 1, let us assume by contradiction that λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .

Original: Assumption 3 implies that ψ ( t ) increases with decreasing λ 1 H ( t ) .
Revised: Assumption 3 implies that ψ ( t ) increases with decreasing λ 1 H ( t ) .

Original: Hence, if λ 1 H ( t ∗ 2 ) is smaller, it means ψ ( t ∗ 2 ) increased to a larger value, i.e. ψ ( t ∗ 2 ) >
Revised: Hence, if λ 1 H ( t ∗ 2 ) is smaller, it means ψ ( t ∗ 2 ) increased to a larger value, i.e. ψ ( t ∗ 2 ) >

Original: ψ ( t ∗ 1 ) .
Revised: ψ ( t ∗ 1 ) .

Original: However, if ( n − S ) ( n − 1) ≈ 1 , a necessary condition for λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) inequality to hold is that ψ ( t ∗ 1 ) > ψ ( t ∗ 2 ) , which leads to a contradiction.
Revised: However, if ( n − S ) ( n − 1) ≈ 1 , a necessary condition for λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) inequality to hold is that ψ ( t ∗ 1 ) > ψ ( t ∗ 2 ) , which leads to a contradiction.

Original: Repeating the same argument for batch size completes the proof.
Revised: Repeating the same argument for batch size completes the proof.

Original: C A PPROXIMATING THE EIGENSPACE OF K AND H
Revised: C A PPROXIMATING THE EIGENSPACE OF K AND H

Original: Using a small subset of the training set sufﬁces to approximate well the largest eigenvalues of the
Revised: Using a small subset of the training set sufﬁces to approximate well the largest eigenvalues of the

Original: Hessian on the CIFAR-10 dataset (Alain et al., 2019).
Revised: Hessian on the CIFAR-10 dataset (Alain et al., 2019).

Original: Following Alain et al.
Revised: Following Alain et al.

Original: we use approximately the same 5% fraction of the dataset in CIFAR-10 experiments, and  SCIPY Python package.
Revised: we use approximately the same 5% fraction of the dataset in CIFAR-10 experiments, and the SCIPY Python package.

Original: We use the same setting on the other datasets as well.
Revised: We use the same setting on the other datasets as well.

Original: Analysing the eigenspace of K is less common in deep learning.
Revised: Analysing the eigenspace of K is less common in deep learning.

Original: For the purposes of this paper, we are primarily interested in estimating λ K ∗ and λ K .
Revised: For the purposes of this paper, we are primarily interested in estimating λ K ∗ and λ K .

Original: We also estimate Tr( K ) .
Revised: We also estimate Tr( K ) .

Original: It is worth noting that
Revised: It is worth noting that

Original: Tr( K ) is related to the variance in gradients as 1 D Tr( K )
Revised: Tr( K ) is related to the variance in gradients as 1 D Tr( K )

Original: = 1 D 1 N (cid:80)
Revised: = 1 D 1 N (cid:80)

Original: Ni =

Revised: Ni =

Original: ||
Revised: ||

Original: g − g i || 2 , where g is the full-batch gradient, D is the number of parameters and N is the number of datapoints.
Revised: g − g i || 2 , where g is the full-batch gradient, D is the number of parameters and N is the number of datapoints.

Original: The naive computation of the eigenspace of K is infeasible for realistically large deep networks due to the quadratic cost in the number of parameters.
Revised: The naive computation of the eigenspace of K is infeasible for realistically large deep networks due to the quadratic cost in the number of parameters.

Original: Instead, we compute K using mini-batches.
Revised: Instead, we compute K using mini-batches.

Original: To avoid storing a D × D matrix in memory, we ﬁrst sample L mini-batch gradient of size M and compute the corresponding Gram matrix K M that has entries K Mij = 1 L (cid:104) g
Revised: To avoid storing a D × D matrix in memory, we ﬁrst sample L mini-batch gradient of size M and compute the corresponding Gram matrix K M that has entries K Mij = 1 L (cid:104) g

Original: i − g, g j − g (cid:105) , where g is the full-batch gradient, which we estimate based on the L mini-batches.
Revised: i − g, g j − g (cid:105) , where g is the full-batch gradient, which we estimate based on the L mini-batches.

Original: To compute the eigenspace of K M we use SVD routine from the NumPy package.
Revised: To compute the eigenspace of K M we use SVD routine from the NumPy package.

Original: We take the ( L − 1) th smallest eigenvalue as the smallest non-zero eigenvalue of K M (covariance matrix computed using L observations has by deﬁnition L − 1 non-zero eigenvalues).
Revised: We take the ( L − 1) th smallest eigenvalue as the smallest non-zero eigenvalue of K M (covariance matrix computed using L observations has by deﬁnition L − 1 non-zero eigenvalues).

Original: Papyan (2019); Fort & Ganguli (2019) show that the top eigenvalues of H emerge due to clustering of gradients of the logits.
Revised: Papyan (2019); Fort & Ganguli (2019) show that the top eigenvalues of H emerge due to clustering of gradients of the logits.

Original: See also Fort & Ganguli (2019).
Revised: See also Fort & Ganguli (2019).

Original: Based on the proximity of the largest eigenvalues of H and K , it suggests that the top eigenvalue of K might be similar to that of K M .
Revised: Based on the proximity of the largest eigenvalues of H and K , this observation suggests that the top eigenvalue of K might be similar to that of K M .

Original: To investigate this we run the following experiment on the CIFAR-10 dataset using SimpleCNN.
Revised: To investigate this, we run the following experiment on the CIFAR-10 dataset using SimpleCNN.

Original: We estimate λ 1 K using M = 1 and M = 128 , in both cases using the whole training set.
Revised: We estimate λ 1 K using M = 1 and M = 128 , in both cases using the whole training set.

Original: We subsample the dataset to 10% to speed up the computation.
Revised: We subsample the dataset to 10% to speed up the computation.

Original: Fig.
Revised: Fig.

Original: 9 shows a strong correlation between λ 1 K computed with M = 1 and with M = 128 , for three different learning rates.
Revised: 9 shows a strong correlation between λ 1 K computed with M = 1 and with M = 128 , for three different learning rates.

Original: [Figure]
Revised: [Figure]

Original: In all experiments we use L = 25 , and for the IMDB dataset we increase the number of mini-batch gradients to L = 200 to reduce noise (same conclusion hold for L = 25 ).
Revised: In all experiments we use L = 25 , and for the IMDB dataset we increase the number of mini-batch gradients to L = 200 to reduce noise (same conclusion hold for L = 25 ).

Original: For instance, on the
Revised: For instance, on the

Original: CIFAR-10 dataset this amounts to using approximately 5% of the training set.
Revised: CIFAR-10 dataset this amounts to using approximately 5% of the training set.

Original: In experiments that vary the batch size we use M = 128 .
Revised: In experiments that vary the batch size we use M = 128 .

Original: Otherwise, we use the value M as the batch size used to train the model.
Revised: Otherwise, we use the value M as the batch size used to train the model.

Original: D E XPERIMENTAL DETAILS FOR S EC . 4. 
Revised: D E XPERIMENTAL DETAILS FOR S EC . 4.

Original: In this section we describe all the details for experiments in Sec. 4.2.
Revised: In this section we describe all the details for experiments in Sec. 4.2.

Original: ResNet-32 on CIFAR.
Revised: ResNet-32 on CIFAR-10.

Original: ResNet-32 (He et al., 2015) is trained for 200 epochs with a batch size equal to 128 on the CIFAR-10 dataset.
Revised: ResNet-32 (He et al., 2015) is trained for 200 epochs with a batch size equal to 128 on the CIFAR-10 dataset.

Original: Standard data augmentation and preprocessing is applied.
Revised: Standard data augmentation and preprocessing is applied.

Original: Following He et al.
Revised: Following He et al.

Original: , we regularize the model using weight decay 0 . 0001 .
Revised: , we regularize the model using weight decay 0 . 0001 .

Original: We apply it to all convolutional kernels.
Revised: We apply weight decay to all convolutional kernels.

Original: When varying the batch size, we use learning rate of 0 . 05 .
Revised: When varying the batch size, we use learning rate of 0 . 05 .

Original: When varying the learning rate, we use batch size of 128 .
Revised: When varying the learning rate, we use batch size of 128 .

Original: SimpleCNN on CIFAR.
Revised: SimpleCNN on CIFAR-10.

Original: The network used is a simple convolutional network with 4 convolutional layers from Keras examples repository (Chollet et al., 2015).
Revised: SimpleCNN is a simple convolutional network with four convolutional layers based on Keras examples repository (Chollet et al., 2015).

Original: The architecture is as follows.
Revised: The architecture is as follows.

Original: After the ﬁrst two convolutional layers with 32 ﬁlters, andafter the last two convolutional layers withﬁlters, we include  max pooling layer with a window size of 2.
Revised: The ﬁrst two convolutional layers have 32 ﬁlters, and the last two convolutional layers haveﬁlters. After each pair of convolutional layers, we include a max pooling layer with a window size of 2.

Original: After each layer we include  ReLU nonlinearity.
Revised: After each layer we include a ReLU nonlinearity.

Original: Before the classiﬁcation layer, a densely connected layer with 128 units is used.
Revised: The output of the ﬁnal convolutional layer is processed by a densely connected layer with 128 units and ReLU nonlinearity.

Original: When varying the batch size, we use a learning rate of 0 . 05 .
Revised: When varying the batch size, we use a learning rate of 0 . 05 .

Original: When varying the learning rate, we use a batch size of 128 .
Revised: When varying the learning rate, we use a batch size of 128 .

Original: BERT on MNLI.
Revised: BERT on MNLI.

Original: The model used in this experiment is  BERT-base (Devlin et al., 2018), pretrained on multilingual data 4 .
Revised: The model used in this experiment is the BERT-base from Devlin et al. (2018), pretrained on multilingual data 4 .

Original: The model is trained on the MultiNLI dataset (Williams et al., 2018) with the maximum sentence length equal to 40.
Revised: The model is trained on the MultiNLI dataset (Williams et al., 2018) with the maximum sentence length equal to 40.

Original: The network is trained for 20 epochs with a batch size of 32.
Revised: The network is trained for 20 epochs using a batch size of 32.

Original: Experiments are repeated with three different seeds that control initialization and data shufﬂing.
Revised: Experiments are repeated with three different seeds that control initialization and data shufﬂing.

Original: MLP on FMNIST.
Revised: MLP on FashionMNIST.

Original: This experiment is using a Multi Layer Perceptron network with two hidden layers of size 300 and 100, both with ReLU activations.
Revised: This experiment is using a multi-layer perceptron with two hidden layers of size 300 and 100, both with ReLU activations.

Original: The data is normalized to the [0 , 1] range.
Revised: The data is normalized to the [0 , 1] range.

Original: The network is trained with a batch size of 64 for 200 epochs.
Revised: The network is trained with a batch size of 64 for 200 epochs.

Original: [Table]

Original: LSTM on IMDB.
Revised: LSTM on IMDB.

Original: The network used in this experiment consists of an embedding layer followed by an LSTM with 100 hidden units.
Revised: The network used in this experiment consists of an embedding layer followed by an LSTM with 100 hidden units.

Original: We use vocabulary size of 20000 words and the maximum  length of the sequence equal to 80.
Revised: We use vocabulary size of 20000 words and the maximum bert multi_cased_ L-12_H- length of the sequence equal to 80.

Original: The model is trained for 100 epochs.
Revised: The model is trained for 100 epochs.

Original: When varying the learning rate, we usea constant batch size of 128.
Revised: When varying the learning rate, we use batch size of 128.

Original: When varying the batch size, we usea learning rate of 1.0.
Revised: When varying the batch size, we use learning rate of 1.0.

Original: Experiments are repeated with two different seeds that control initialization and data shufﬂing.
Revised: Experiments are repeated with two different seeds that control initialization and data shufﬂing.

Original: DenseNet on ImageNet.
Revised: DenseNet on ImageNet.

Original: The network used is the DenseNet-121 (Huang et al., 2016).
Revised: The network used is the DenseNet-121 from Huang et al.

Original: The dataset used is the ILSVRC 2012 (Russakovsky et al., 2015).
Revised: The dataset used is the ILSVRC 2012 (Russakovsky et al., 2015).

Original: The images are centered, but no augmentation isbeing used.
Revised: The images are centered and normalized. No data augmentation is used.

Original: Due to large computational cost, the network is trained only for 10 epochs with a batch size of 32.
Revised: Due to large computational cost, the network is trained only for 10 epochs using a batch size of 32.

Original: E A DDITIONAL EXPERIMENTS FOR S EC . 4.2.
Revised: E A DDITIONAL EXPERIMENTS FOR S EC . 4.2.

Original: In this section we include additional data for experiments in Sec.
Revised: In this section we include additional data for experiments in Sec.

Original: 4.2, as well as include experiments using MLP trained on the Fashion MNIST dataset.
Revised: 4.2, as well as include experiments using MLP trained on the Fashion MNIST dataset.

Original: SimpleCNN on CIFAR-10.
Revised: SimpleCNN on CIFAR-10.

Original: In Fig.
Revised: In Fig.

Original: 12 and Fig.
Revised: 12 and Fig.

Original: 13 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all experiments with SimpleCNN model on the CIFAR-10 dataset
Revised: 13 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all experiments with SimpleCNN model on the CIFAR-10 dataset

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: ResNet-32 on CIFAR-10.
Revised: ResNet-32 on CIFAR-10.

Original: In Fig.
Revised: In Fig.

Original: 12 and Fig.
Revised: 12 and Fig.

Original: 13 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all experiments with ResNet-32 on the CIFAR-10 dataset .
Revised: 13 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all experiments with ResNet-32 on the CIFAR-10 dataset.

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

Original: LSTM on IMDB.
Revised: LSTM on IMDB.

Original: In Fig.
Revised: In Fig.

Original: 14 and Fig.
Revised: 14 and Fig.

Original: 15 we report accuracy on the training set and the validation set, λ 1 H [Figure]
Revised: 15 we report accuracy on the training set and the validation set, λ 1 H , and Tr(

Original: and Tr( K ) for all experiments with LSTM on the IMDB dataset.
Revised: ) for all experiments with LSTM on the IMDB dataset.

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

Original: BERT on MNLI.
Revised: BERT on MNLI.

Original: In Fig.
Revised: In Fig.

Original: 16 we report accuracy on the training set and the validationset and
Revised: 16 we report accuracy on the training set and the validation

Original: Tr( K ) for BERT model on theMNLI dataset.
Revised: Tr BERT model on the

Original: [Figure]
Revised: [Figure]

Original: DenseNet on ImageNet.
Revised: DenseNet on ImageNet.

Original: In Fig.
Revised: In Fig.

Original: 17 we report accuracy on the training set and the validation set and Tr( K ) for DenseNet on the ImageNet dataset.
Revised: 17 we report accuracy on the training set and the validation set and Tr( K ) for DenseNet on the ImageNet dataset.

Original: [Figure]
Revised: [Figure]

Original: MLP on FMNIST.
Revised: MLP on FashionMNIST.

Original: In Fig.
Revised: In Fig.

Original: 18 and Fig.
Revised: 18 and Fig.

Original: 19 we report results for MLP model on the FMNIST dataset.
Revised: 19 we report results for MLP model on the FashionMNIST dataset.

Original: We observe that all conclucions carry over to this setting
Revised: We observe that all conclusions carry over to this setting

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

Original: F A DDITIONAL EXPERIMENTS FOR SGD WITH 
Revised: F A DDITIONAL EXPERIMENTS FOR SGD WITH MOMENTUM

Original: MOMENTUM

Original: Here, we run the same experiment as in Sec.
Revised: Here, we run the same experiment as in Sec.

Original: 4.2 using SimpleCNN on the CIFAR-10 dataset.
Revised: 4.2 using SimpleCNN on the CIFAR-10 dataset.

Original: Instead of varying the learning rate, we test different values of the momentum β parameter in the range of 0 . 1 , 0 .
Revised: Instead of varying the learning rate, we test different values of the momentum β parameter in the range of 0 . 1 , 0 .

Original: 5 and 0 .
Revised: 5 and 0 .

Original: We can observe that Conjecture 1 and Conjecture 2 generalize to momentum in the sense that using a higher momentum has an analogous effect to using a higher learning rate, or using a smaller batch size in SGD.
Revised: We can observe that Conjecture 1 and Conjecture 2 generalize to momentum in the sense that using a higher momentum has an analogous effect to using a higher learning rate, or using a smaller batch size in SGD.

Original: We report the results in Fig.
Revised: We report the results in Fig.

Original: 20 and Fig.
Revised: 20 and Fig.

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

Original: Next, to study whether our Conjectures are also valid for SGD with momentum, but held constant, we run the same experiment as in Sec.
Revised: Next, to study whether our Conjectures are also valid for SGD with momentum, but held constant, we run the same experiment as in Sec.

Original: 4.2 using SimpleCNN on the CIFAR-10 dataset.
Revised: 4.2 using SimpleCNN on the CIFAR-10 dataset.

Original: For all runs, we set momentum to 0 .
Revised: For all runs, we set momentum to 0 .

Original: Learning rate 0 .
Revised: Learning rate 0 .

Original: 1 diverged training, so we include only 0 . 01 and 0 . 001 .
Revised: 1 diverged training, so we include only 0 . 01 and 0 . 001 .

Original: We can observe that both Conjectures generalize to this setting.
Revised: We can observe that both Conjectures generalize to this setting.

Original: We report the results in Fig.
Revised: We report the results in Fig.

Original: 22 and
Revised: 22 and

Original: Fig. 23.
Revised: Fig. 23.

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

Original: G A DDITIONAL EXPERIMENTS FOR SGD WITH LEARNING RATE DECAY
Revised: G A DDITIONAL EXPERIMENTS FOR SGD WITH LEARNING RATE DECAY

Original: To understand the effect of learning rate decay, we run the same experiment as in Sec.
Revised: To understand the effect of learning rate decay, we run the same experiment as in Sec.

Original: 4.2 using
Revised: 4.2 using

Original: SimpleCNN on the CIFAR-10 dataset.
Revised: SimpleCNN on the CIFAR-10 dataset.

Original: Additionally, we divide the learning rate by the factor ofafter 100th epoch.
Revised: Additionally, we divide the learning rate by the factor ofafter 100th epoch.

Original: We can observe that Conjecture 1 and Conjecture 2 generalize to scenario with learning rate schedule in the sense that changing learning rate doesn’t change the relative ordering of the maximum λ 1 K and λ ∗ K /λ 1 K .
Revised: We can observe that Conjecture 1 and Conjecture 2 generalize to scenario with learning rate schedule in the sense that changing learning rate doesn’t change the relative ordering of the maximum λ 1 K and λ ∗ K /λ 1 K .

Original: We report the results in Fig.
Revised: We report the results in Fig.

Original: 24 and Fig.
Revised: 24 and Fig.

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

Original: H A DDITIONAL EXPERIMENTS FOR S IMPLE CNN-BN
Revised: H A DDITIONAL EXPERIMENTS FOR S IMPLE CNN-BN

Original: To further explore the connection between our conjectures and the effects of batch-normalization layers on the conditioning of the loss surface
Revised: To further explore the connection between our conjectures and the effects of batch-normalization layers on the conditioning of the loss surface

Original: [Figure]
Revised: [Figure]

Original: we repeat here experiments from Sec. 4.
Revised: we repeat here experiments from Sec. 4.

Original: [Figure]
Revised: [Figure]

Original: but varying the batch size.
Revised: but varying the batch size.

Original: Fig.
Revised: Fig.

Original: 26 summarizes the results.
Revised: 26 summarizes the results.

Original: On the whole
Revised: On the whole

Original: [Figure]
Revised: [Figure]

Original: conclusions from Sec.
Revised: conclusions from Sec.

Original: 4.3 carry over to this setting in the sense that decreasing the batch size has a similar effect on the studied metrics as increasing the learning rate.
Revised: 4.3 carry over to this setting in the sense that decreasing the batch size has a similar effect on the studied metrics as increasing the learning rate.

Original: One exception is the experiment using the smallest batch size of 10 .
Revised: One exception is the experiment using the smallest batch size of 10 .

Original: In this case, the maximum values of (cid:107) g (cid:107) (cid:107) g 5 (cid:107) and λ ∗ K /λ 1 K are smaller than in the experiments using larger batch sizes
Revised: In this case, the maximum values of (cid:107) g (cid:107) (cid:107) g 5 (cid:107) and λ ∗ K /λ 1 K are smaller than in the experiments using larger batch sizes

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