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

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Understanding the optimization trajectory is critical for understanding training of deep neural networks
Revised: Understanding the optimization trajectory is critical for understanding training of deep neural networks.

Original: We show how the hyperparameters of stochastic gradient descent inﬂuence the covariance of the gradients ( K  and the Hessian of the training loss ( H ) along this trajectory
Revised: We show how the hyperparameters of stochastic gradient descent inﬂuence the covariance of the gradients ( K ) and the Hessian of the training loss ( H ) along the trajectory.

Original: Based on a theoretical model , we conjecture that using a high learning rate or a small batch size leads SGD to regions of loss landscape, typically early during training, with (1) reduced spectral norm of K and (2) improved conditioning of K and H .
Revised: Based on a theoretical model, we conjecture that using a high learning rate or a small batch size leads SGD to regions of loss landscape, typically early during training, characterized by (1) reduced spectral norm of K and (2) improved conditioning of K and H .

Original: We refer to the point on the  trajectory after which these effects hold as the break-even point .
Revised: We refer to the point on the training trajectory after which these effects hold as the break-even point .

Original: We demonstrate these effects empirically for a range of deep neural networks applied tomultiple different tasks
Revised: We demonstrate these effects empirically for a range of deep neural networks applied to different tasks.

Original: Finally, we apply our analysis to networks with batch normalization (BN) layers and ﬁnd that it is necessary to use a high learning rate to improve loss surface conditioning compared to a network without BN layers.
Revised: Finally, we apply our analysis to networks with batch normalization (BN) layers and ﬁnd that it is necessary to use a higher learning rate to improve loss surface conditioning compared to a network without BN layers

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The choice of the optimization method implicitly regularizes deep neural networks (DNNs) by inﬂuencing the optimization trajectory in the loss surface (Neyshabur, 2017; Arora, 2019).
Revised: The choice of the optimization method implicitly regularizes deep neural networks (DNNs) by inﬂuencing the optimization trajectory in the loss surface (Neyshabur, 2017; Arora, 2019).

Original: In this work, we theoretically and empirically investigate how the learning rate and the batch size used at the beginning of training determine properties of the entire optimization trajectory.
Revised: In this work, we theoretically and empirically investigate how the learning rate and the batch size used at the beginning of training determine properties of the entire optimization trajectory.

Original: [Figure]
Revised: [Figure]

Original: We focus our analysis on two objects that quantify different properties of the optimization trajectory: the  Hessian of the training loss ( H), and the covariance of gradients ( K ).
Revised: We focus our analysis on two objects that quantify different properties of the optimization trajectory: the covariance of gradients ( K ) 1 , and the Hessian of the training loss ( H ).

Original: The matrix K quantiﬁes noise induced by noisy estimate of the full-batch gradient and has beenalso linked to  generalization error (Roux et al., 2008; Fort et al., 2019).
Revised: The matrix K quantiﬁes noise induced by noisy estimate of the full-batch gradient, and has been linked to the generalization error (Roux et al., 2008; Fort et al., 2019).

Original: The matrix H describes the curvature of the loss surface and is often connected to optimization speed.
Revised: The matrix H describes the curvature of the loss surface.

Original: Further, better conditioning of Hand K has been attributed as the main reason behind the efﬁcacy of batch normalization (Bjorck et al., 2018; Ghorbani et al., 2019).
Revised: Better conditioning of H has been attributed as the main reason behind the efﬁcacy of batch normalization (Bjorck et al., 2018; Ghorbani et al., 2019).

Original: Our ﬁrst and main contribution is predicting and empirically demonstrating two effectsinduced in the early phase of training  by the choice of the hyperparameters in stochastic gradient descent (SGD): (1) reduced spectral norms of K and H and (2) improved conditioning of K and H .
Revised: Our ﬁrst and main contribution is predicting, and empirically demonstrating two effects in the early phase of training inﬂuenced by the choice of the hyperparameters in stochastic gradient descent (SGD): (1) reduced spectral norms of K and H and (2) improved conditioning of K and H .

Original: These effects manifest themselves after a certain point on the optimization trajectory, to which we refer to as the break-even point .
Revised: These effects manifest themselves after a certain point on the optimization trajectory, to which we refer to as the break-even point .

Original: See Fig.
Revised: See Fig.

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

Original: We make our predictions based on a theoretical model of the initial phase of training, which incorporates recent observations on the instability and oscillations in the parameter space that characterize the learning dynamics of neural networks (Masters & Luschi, 2018;
Revised: We make our predictions based on a theoretical model of the initial phase of training, which incorporates recent observations on the instability and oscillations in the parameter space that characterize the learning dynamics of neural networks (Masters & Luschi, 2018;

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

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

Original: Delving deeper in this direction of investigation, we show that using a large learning rate is necessary to reach better-conditioned, relatively to a network without 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: Delving deeper in this direction of investigation, we show that using a large learning rate is necessary to reach better-conditioned, relatively to a network without 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: Learning dynamics and the early phase of training.
Revised: Learning dynamics and the early phase of training.

Original: Our theoretical model is motivated by recent work on the 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 the 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: ; Jastrzebski et al.
Revised: ; Jastrzebski et al.

Original: ; Lan et al.
Revised: ; Lan et al.

Original: (2019) who show that training oscillates in the parameter space.
Revised: (2019) who show that training oscillates, which can be characterized by a negative cosine between consecutive steps.

Original: Our work is closely related to Wu et al.
Revised: Our theoretical approach is closely related to Wu et al.

Original: (2018) who proposes a linear stability approach to studying how SGD selects a minimum.
Revised: While they study how SGD selects a minimum from a stability perspective, we apply their methodology to the early phase of training.

Original: In contrast to them, we study the early phase of training and  maketheoretical predictions about the whole training trajectory.
Revised: Based on their theoretical results and other recent studies, we make predictions about the whole training trajectory.

Original: In our work we argue that the initial phase of training has important implications for the rest of the trajectory.
Revised: In our work we argue that the initial phase of training determines the rest of the trajectory.

Original: This is directly related to Erhan et al.
Revised: This is directly related to Erhan et al.

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

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

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

Original: (2010) argue thatinitial 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 pre-training 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 regularizers 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: The covariance of the gradients and the Hessian.
Revised: The covariance of the gradients and the Hessian.

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

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: A similar quantity, cosine alignment between gradients computed on individual examples, was recently shown to explain some aspects of deep networks generalization (Fort et al., 2019).
Revised: A similar quantity, cosine alignment between gradients computed on different examples, was recently shown to explain some aspects of deep networks generalization (Fort et al., 2019).

Original: The second object that we study is the Hessian that quantiﬁes the loss surface shape (LeCun et al., 2012).
Revised: The second object that we study is the Hessian that quantiﬁes the loss surface shape (LeCun et al., 2012).

Original: Recent work has shown that the largest eigenvalues of H grow quickly initially, and then stabilize at a value dependent on the learning rate and the batch size (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 grow quickly in the early phase of training, and then stabilize at a value dependent on the learning rate and the batch size (Keskar et al., 2017; Sagun et al., 2017; Fort & Scherlis, 2019; Jastrzebski et al., 2018).

Original: The Hessian can be decomposed into a sum of 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 a sum of 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: While we study K , the centered version of G , K and G are typically similar due to the dominance of noise in training (Zhu et al., 2018;
Revised: While we study K , the centered version of G , K and G are typically similar due to the dominance of noise in training (Zhu et al., 2018;

Original: Thomas et al., 2019).
Revised: Thomas et al., 2019).

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

Original: Multiple prior work study the regularization effects that are attributed only to the optimization method (Neyshabur, 2017).
Revised: Multiple prior works study the regularization effects that are attributed only to the optimization method (Neyshabur, 2017).

Original: A popular research direction is to bound the generalization error based on the properties of the ﬁnal minimum such as the norm of the parameter vector or the Hessian (Bartlett et al., 2017; Keskar et al., 2017).
Revised: A popular research direction is to bound the generalization error based on the properties of the ﬁnal minimum such as the norm of the parameter vector or the Hessian (Bartlett et al., 2017; Keskar et al., 2017).

Original: Perhaps the most related work to ours is (Arora et al., 2019; Arora, 2019).
Revised: Perhaps the most related work to ours is (Arora et al., 2019; Arora, 2019).

Original: They suggest it is necessary to study the  trajectory to understand generalization of deep networks.
Revised: They suggest it is necessary to study the optimization trajectory to understand generalization of deep networks.

Original: In this vein, but in contrast to most of the previous work, we focus (1) on the implicit regularization effects that can be attributed to the GD dynamics at the beginning of training, and (2) on the covariance of gradients.
Revised: In this vein, but in contrast to most of the previous work, we focus (1) on the implicit regularization effects that can be attributed to the GD dynamics at the beginning of training, and (2) on the covariance of gradients.

Original:  T WO CONJECTURES ABOUT SGD TRAJECTORY
Revised:  T WO CONJECTURES ABOUT SGD TRAJECTORY

Original: In this section we make two conjectures about the optimization trajectory induced by SGD based on a theoretical model of the learning dynamics in the early stage of training.
Revised: In this section we make two conjectures about the optimization trajectory induced by SGD based on a theoretical model of the learning dynamics in the early stage of training.

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

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: A key object we study is the HessianH of the training loss.
Revised: A key object we study is the Hessian of the training loss ( H ).

Original: The second key object we study is the covariance of the gradients K = 1 N (cid:80) N i =

Revised: The second key object we study is the covariance of the 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  θ ( t )  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 study stability of optimization, in our case restricted to e 1 H .
Revised: we study stability of optimization, in our case restricted to e 1 H .

Original: Let us call the projection of parameters θ onto e 1 H by ψ = (cid:104) θ, e 1 H (cid:105) .
Revised: Let us call 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: Similarly to Wu et al.
Revised: Similarly to Wu et al.

Original: (2018), we say SGD is unstable along e 1 H at θ ( t ) if the norm of elements of sequence ψ ( τ + 1) =
Revised: (2018), 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 restricted from θ ( t ) only to the direction of e 1 H .

Original: Assumptions.
Revised: Assumptions.

Original: We make the following assumptions to build our model.
Revised: We make the following assumptions to construct our model.

Original: The loss surface for each example projected onto e 1 H is a quadratic function.
Revised: The loss surface for each example projected onto e 1 H is a quadratic function.

Original: This assumption was also used in Wu et al.
Revised: This assumption was also used in Wu et al.

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

Original: (2019) show  that the loss averaged over all training examples is well approximated by a quadratic function along e 1 H .
Revised: (2019) 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 , andfurthermore λ 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 Papyan (2019) who show that H can be approximated by G (non-centered K ).
Revised: This is inspired by Papyan (2019) who show that H can be approximated by G (non-centered K ).

Original: When in a region that is not stable along e 1 H , λ 1 H decreases along the training trajectory, andoverall training loss increases along the trajectory.
Revised: 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.

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  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 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 .

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

Original: Sagun et al., 2017; Jastrzebski et al., 2018; Fort & Scherlis, 2019; Fort & Ganguli, 2019) who show that in many settings λ 1 H 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 increases in the beginning of training, while at the same time the overall training loss decreases.

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 verify empirically.
Revised: Its main purpose is to make predictions about the training procedure that we verify empirically.

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  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 in real 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 real 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: 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  for the ﬁrst timeequal to 1 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 trajectory.

Original: Under the Assumption 3, λ 1 H ( t ) and λ 1 K ( t ) increase over time.
Revised: Under the 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 carried out for S .
Revised: The same reasoning can be applied to S .

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

Original: From this point on the trajectory, under Assumption 4, SGD does not enter regions where either λ 1 H ( t (cid:48) ) or λ 1 K ( t (cid:48) ) 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: From this point on the trajectory, under Assumption 4, SGD does not enter regions where either λ 1 H ( t (cid:48) ) or λ 1 K ( t (cid:48) ) 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) ) .

Original: The two conjectures about real DNNs.
Revised: The two conjectures about real DNNs.

Original: Assuming that real DNNs reach the break-even point, we make the following two conjectures, arising from our theoretical model above, about their optimization trajectory.
Revised: Assuming that real DNNs reach the break-even point, we make the following two conjectures, arising from our theoretical model above, 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 Con.
Revised: We refer to Con.

Original: 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: 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: 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 does not affect the λ iH and λ iK for i (cid:54)
Revised: It is plausible to assume that reaching the break-even point does not affect the λ iH and λ iK for i (cid:54)

Original: = 1 , because increasing their values does not impact stability along e 1 H in our theoretical model.
Revised: = 1 , because increasing their values does not impact stability along e 1 H in our theoretical model.

Original: It is also well known that  a large number of eigenvalues of H increaseinitially (Fort & Scherlis, 2019; Sagun et al., 2017).
Revised: It is also well known that initially a large number of eigenvalues of H increase (Fort & Scherlis, 2019; Sagun et al., 2017).

Original: Based on these remarks we conjecture that:
Revised: Based on these remarks 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 nonzero eigenvalues of H and K , 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 nonzero eigenvalues of H and K , respectively.

Original: Furthermore, the maximum attained values of Tr( K ) ,
Revised: Furthermore, the maximum attained values of Tr( K ) ,

Original: Tr( H ) are smaller for a larger learning rate or a smaller batch size.
Revised: 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 N − D non-zero eigenvalues, where N is the number of training points.
Revised: We consider non-zero eigenvalues in Con. 2, because K has N − D non-zero eigenvalues, where N is the number of training points.

Original: Both conjectures are only valid for learning rates and batch sizes in the range that ensures overall loss decreases.
Revised: Both conjectures are only valid for learning rates and batch sizes in the range that guarantee that training is stable after reaching the break-even point.

Revised: It is also important to stress that while our predictions are related to the curvature of the ﬁnal minimum (Wu et al., 2018; Jastrzebski et al., 2017), in contrast to prior work we are interested in the properties of the whole training trajectory.

Revised: Finally, from the optimization perspective, the effects discussed above are desirable; many papers in the optimization literature focused on either reducing variance of the mini-batch gradients (Johnson & Zhang, 2013), or on improving conditioning (Roux et al., 2008).

Revised: [Figure]

Revised: [Figure]

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 validate the two conjectures.
Revised: Next, we empirically validate 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.
Revised: Due to the space constraint we take the following approach to reporting results.

Original: In the main body of the paper, we focus on the CIFAR-10 dataset (Krizhevsky, 2009) and the IMDB dataset (Maas et al., 2011), to which we apply three architectures: a vanilla CNN (SimpleCNN) following
Revised: In the main body of the paper we focus on the CIFAR-10 dataset (Krizhevsky, 2009) and the IMDB dataset (Maas et al., 2011), to which we apply three architectures: a vanilla CNN (SimpleCNN) following

Original: Keras example (Chollet et al., 2015), ResNet-32 (He et al., 2015a), and LSTM (Hochreiter &
Revised: Keras example (Chollet et al., 2015), ResNet-32 (He et al., 2015) and LSTM (Hochreiter &

Original: Schmidhuber, 1997).
Revised: Schmidhuber, 1997).

Original: We also validate the two conjectures for DenseNet (Huang et al., 2016) on the ImageNet (Deng et al., 2009) dataset, BERT (Devlin et al., 2018) ﬁne-tuned on the MNLI[Figure] [Figure] dataset (Williams et al., 2017), and a multi-layer perceptron on the FashionMNIST dataset (Xiao et al., 2017).
Revised: We also validate the two conjectures for DenseNet (Huang et al., 2016) on the ImageNet (Deng et al., 2009) dataset, BERT (Devlin et al., 2018) ﬁne-tuned on the MNLI dataset (Williams et al., 2017), and a multi-layer perceptron on the FashionMNIST dataset (Xiao et al., 2017).

Original: These results are in the Appendix. We include all experimental details  in App. C.
Revised: These results and all experimental details are in App. C.

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 using the Lanczos algorithm on a random subset of 5% of the training set on CIFAR-10.
Revised: (2019), we estimate the top eigenvalues and eigenvectors of H using the Lanczos algorithm on a random subset of 5% of the training set on CIFAR-10.

Original: We estimate the top eigenvalues and eigenvectors of K using (in most cases) batch size of 128 and approximately 5% of the training set on CIFAR-10.
Revised: We estimate the top eigenvalues and eigenvectors of K using (in most cases) batch size of 128 and approximately 5% of the training set on CIFAR-10.

Original: Furthermore, under the assumption that noise due to sampling data is normally distributed, our approximation biases the smallest eigenvalue towards 0 .
Revised: Furthermore, under the assumption that noise due to sampling data is normally distributed, our approximation biases the smallest eigenvalue towards 0 .

Original: To counteract this, we use the 5 th smallest non-zero eigenvalue in place of λ ∗ H and λ ∗ K (see Con.
Revised: To counteract this, we use the 5 th smallest non-zero eigenvalue in place of λ ∗ H and λ ∗ K .

Original: 2 for deﬁnition).

Original: We describe the procedure in more details, as well as examine the impact of using mini-batching in computing the eigenspectrum of K , in App.
Revised: We describe the procedure in more details, as well as examine the impact of using mini-batching in computing the eigenspectrum of K , in App.

Original: B.
Revised: B.

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

Original: In this section we examine the learning dynamics in the early phase of training.
Revised: In this section 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 for η = 0 .
Revised: We analyse the evolution of λ 1 H and λ 1 K for η = 0 .

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

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

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 into 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).
Revised: This is expected: λ 1 H decays to 0 when the mean loss decays to 0 for cross entropy loss (Martens, 2016).

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: A key aspect of our model is that an increase of λ 1 K and λ 1 H translates into a decrease in stability, which we formalized as stability along e 1 H .
Revised: A key aspect of our model is that an increase of λ 1 K and λ 1 H translates into a decrease in stability, which we formalized as stability along e 1 H .

Original: Computing stability directly along e 1 H is computationally expensive.
Revised: Computing stability directly along e 1 H 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 reduction of the training loss).
Revised: 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 reduction of the training loss).

Original: In Fig.
Revised: In Fig.

Original: 2 we observe that training becomes unstable around the time λ 1 K reaches the maximum value.
Revised: 2 we observe that training becomes unstable when λ 1 K reaches the maximum value.

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

Original: Finally, to understand the break-even point phenomenon better, we visualize the learning dynamics leading to reaching the break-even point in our model in
Revised: Finally, to understand the break-even point phenomenon better, we visualize the learning dynamics leading to reaching the break-even point in our model 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
Revised: (2010), we embed the test set predictions at each step of training of

Original: SimpleCNN,in our case using UMAP (McInnes et al., 2018).
Revised: SimpleCNN, using UMAP (McInnes et al., 2018).

Original: Background color indicates λ 1 K  and the training accuracy at the closest iteration.
Revised: Background color indicates λ 1 K (left) and the training accuracy at the closest iteration (right).

Original: We observe that  trainingwith the  lower learning rate  η = 0 .
Revised: We observe that early in training the trajectory corresponding to a lower learning rate of η = 0 .

Original: 01 reachesearly during training (before reaching 60% training accuracy) regions of the loss surface corresponding to the same accuracy as η = 0 . 1 and at the same time larger λ 1 K .
Revised: 01 reaches regions of the loss surface corresponding to the same accuracy as η = 0 . 1 , but characterized by larger λ 1 K .

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

Original: 3 we plot the spectrum of K (left) and H at the iteration when λ K and λ H , respectively, reach the highest value.
Revised: 3 we plot the spectrum of K (left) and H at the iteration when λ K and λ H , respectively, reach the highest value.

Original: We can observe that for the lower learning rate the distribution has more outliers.
Revised: We can observe that for the lower learning rate the distribution has more outliers.

Original: Summary.
Revised: Summary.

Original: We have shown that the dynamics of the early phase of training is consistent with the assumptions made in our  model.
Revised: We have shown that the dynamics of 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 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. V ARIANCE REDUCTION AND PRE - CONDITIONING EFFECT OF SGD
Revised: 4. V ARIANCE REDUCTION AND PRE - CONDITIONING EFFECT OF SGD

Original: In this section we validate empirically Con.
Revised: In this section we validate empirically Con.

Original: 1 and Con.
Revised: 1 and Con.

Original: 2 in three settings.
Revised: 2 in three settings.

Original: For each model we pick manually a suitable range of learning rates and batch sizes to ensure that the properties of K and H that weexamine have converged in a reasonable computational budget and training is stable; we use 200 epochs on CIFAR-10 and 50 epochs on IMDB.
Revised: For each model we pick manually a suitable range of learning rates and batch sizes to ensure that the properties of K and H that we converged in a reasonable computational budget and training is stable. We use 200 epochs on CIFAR-10 and 50 epochs on IMDB.

Original: In ResNet-32 we remove batch normalization.
Revised: In ResNet-32 we remove batch normalization.

Original: [Figure]

Original: We summarize the resultsin Fig. 4, Fig. 5 for SimpleCNN, ResNet-32, and LSTM.
Revised: We summarize the results for SimpleCNN, ResNet-32, and LSTM in Fig. 4 and Fig.

Original: Curves are smoothed with a moving average.
Revised: Curves are smoothed with a moving average.

Original: The training curves, as well as experiments for other architectures and datasets (including a DenseNet on ImageNet and BERT on MNLI) can be found in App.
Revised: The training curves, as well as experiments for other architectures and datasets (including a DenseNet on ImageNet and BERT on MNLI) can be found in App.

Original: D.
Revised: D.

Original: Null hypothesis.
Revised: Null hypothesis.

Original: A natural assumption is that the choice of η or S does not inﬂuence K and H along the optimization trajectory.
Revised: A natural assumption is that the choice of η or S does not inﬂuence K and H along the optimization trajectory.

Original: In particular, it is not self-evident that using a high η , or a small
Revised: In particular, it is not self-evident that using a high η , or a small

Original: S , would steer optimization towards better conditioned regions of the loss surface.
Revised: S , would steer optimization towards better conditioned regions of the loss surface.

Original: Conjecture 1.
Revised: Conjecture 1.

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

Original: As visible in Fig.
Revised: As visible in Fig.

Original: 4 using a higher η results in λ 1 K achieving a lower maximum.
Revised: 4 using a higher η results in λ 1 K achieving a lower maximum.

Original: For instance max( λ 1 K ) = 0 . 87 and max( λ 1 K ) = 3 .
Revised: For instance max( λ 1 K ) = 0 . 87 and max( λ 1 K ) = 3 .

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

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

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

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

Original: Recall that we compute λ 1 K using a constant batch size of 128 .
Revised: Recall that we compute λ 1 K using a constant batch size of 128 .

Original: While we say that low S leads to variance reduction (lower maximum λ 1 K ), this is not contradictory to the fact that increasing S generally decreases variance of mini-batch gradients.
Revised: While we say that low S leads to variance reduction (lower maximum λ 1 K ), this is not contradictory to the fact that increasing S generally decreases variance of mini-batch gradients.

Original: Conjecture 2.
Revised: Conjecture 2.

Original: To test Conjecture 2 we compute the maximum value of λ ∗ K /λ 1 K along the optimization trajectory; as discussed before we use the 5 th smallest eigenvalue in place of λ ∗ K .
Revised: To test Conjecture 2 we compute the maximum value of λ ∗ K /λ 1 K along the optimization trajectory; as discussed before we use the 5 th smallest eigenvalue in place of λ ∗ K .

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

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

Original: For instance, in the case of SimpleCNN (top row of Fig. 4) max( λ ∗ K /λ 1 K ) = 0 .
Revised: For instance, in the case of SimpleCNN (top row of Fig. 4) max( λ ∗ K /λ 1 K ) = 0 .

Original: 45 and max( λ ∗ K /λ 1 K ) = 0 .
Revised: 45 and max( λ ∗ K /λ 1 K ) = 0 .

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

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

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

Original: Similarly, we can conclude from these plots that using a higher S leads to lower max( λ ∗ K /λ 1 K ) . 
Revised: Similarly, we can conclude from these plots that using a higher S leads to lower max( λ ∗ K /λ 1 K ) . [Figure] better conditioning).

Revised: These effects occur early in training.

Revised: The vertical lines marks the ﬁrst epoch at which training accuracy is above 60% for CIFAR-10, and above for 75% for IMDB.

Revised: We use base S = 128 in all experiments.

Original:  Due to space constraints we move Figures showing the effect of η and S on Tr( K ) to App.3.
Revised: We use base η = 0 . for CIFAR-10 and η = 0 . 1 for IMDB. Due to space constraints we move Figures showing the effect of η and S on Tr( K ) to App.

Original:  We observe that the maximum of Tr( K ) depends on η and S in the same way as λ 1 K .
Revised: D. We observe that the maximum of Tr( K ) depends on η and S in the same way as λ 1 K .

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

Original: How λ 1 H and λ 1 K depend on η and S at the end of training was already studied by Jastrzebski et al.
Revised: How λ 1 H and λ 1 K depend on η and S at the end of training was already studied by Jastrzebski et al.

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

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

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

Original: See also the vertical lines in Fig.
Revised: See also the vertical lines in Fig.

Original: Other experiments.
Revised: Other experiments.

Original: We report how λ 1 H depends on η and S for ResNet-32 and SimpleCNN in
Revised: We report how λ 1 H depends on η and S for ResNet-32 and SimpleCNN in

Original: Fig.
Revised: Fig.

Original: We observe that the conclusions carry over to λ 1 H , which is consistent with experiments in Jastrzebski et al.
Revised: We observe that the conclusions also apply to λ 1 H , which is consistent with experiments in Jastrzebski et al.

Original: We found the effect  on λ ∗ H /λ 1 H of η and S to be weaker.
Revised: We found the effect of η and S on λ ∗ H /λ 1 H of η and S to be weaker.

Original: This might be because, in contrast to K , we approximate the spectrum of H using only the top ﬁve eigenvalues (see App.
Revised: This might be because, in contrast to K , we approximate the spectrum of H using only the top ﬁve eigenvalues (see App.

Original: B for details).
Revised: B for details).

Original: Summary
Revised: Summary

Original: In this section we have demonstrated the variance reduction (Conjecture 1) and the preconditioning effect (Conjecture 2) of SGD.
Revised: In this section we have demonstrated the variance reduction (Conjecture 1) and the preconditioning effect (Conjecture 2) of SGD.

Original: Furthermore, we have shown these effects occur early in training.
Revised: Furthermore, we have shown these effects occur early in training.

Original: We also found that conclusions carry over to other settings including BERT on MNLI and
Revised: We also found that conclusions also apply to other settings including BERT on MNLI and

Original: DenseNet on ImageNet (see App. D).
Revised: DenseNet on ImageNet (see App. D).

Original: [Figure]

Original: Trajectories corresponding to higher learning rates ( η ) or smaller batch size ( S ) are characterized by a lower maximum λ 1 H .

Original: Vertical lines mark the ﬁrst epoch at which training accuracy is above 60% .

Original: [Figure]

Original: [Figure]

Original: 4. I MPORTANCE OF η FOR CONDITIONING IN BATCH NORMALIZED NETWORKS
Revised: 4. I MPORTANCE OF η FOR CONDITIONING IN BATCH NORMALIZED NETWORKS

Original: The loss surface of deep networks has been widely reported to be ill-conditioned, which is the key motivation behind using second order  methods in deep learning (LeCun et al., 2012;
Revised: The loss surface of deep networks has been widely reported to be ill-conditioned, which is the key motivation behind using second order optimization methods in deep learning (LeCun et al., 2012;

Original: Martens & Grosse, 2015).
Revised: Martens & Grosse, 2015).

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

Original: Our Conjecture 2 is that using a high η or a low S results as well in improved conditioning.
Revised: Our Conjecture 2 is that using a high η or a low S results as well in improved [Figure] [Figure] [Figure]

Revised: || ||

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

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

Original: First, to investigate whether our conjectures hold in batch normalized network, we run similar experiments as in Sec.
Revised: First, to investigate whether our conjectures hold in batch-normalized network, we run similar experiments as in Sec.

Original: 4.2 on a SimpleCNN model with batch normalization layers inserted after each convolutional layer (SimpleCNNBN), using the CIFAR-10 dataset.
Revised: 4.2 on a SimpleCNN model with batch normalization layers inserted after each convolutional layer (SimpleCNNBN), using 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 3 .
Revised: 001 , 0 . 01 , 0 . 1 , 1 . 0 } ( η = 1 . 0 leads to divergence of SimpleCNN without BN).

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

Original: The evolution of λ ∗ K /λ 1 K and λ 1 K show that both Conjecture 1 and Conjecture 2 hold in this setting.
Revised: The evolution of λ ∗ K /λ 1 K and λ 1 K show that both Conjecture 1 and Conjecture 2 hold in this setting.

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
Revised: To further corroborate that our analysis applies to

Original: BN networks, we study the early phase of training of a network with batch normalization layers, complementing the results in Sec. 4.1.
Revised: BN networks, we study the early phase of training of a network with batch normalization layers, complementing the results in Sec. 4.1.

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

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

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

Original: λ 1 K then decays for all but the lowest η .
Revised: λ 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 BN, (cid:107) γ (cid:107) , in the last layer of the network in Fig.
Revised: We also track the norm of the scaling factor in BN, (cid:107) γ (cid:107) , in the last layer of the network in Fig.

Original: 6 (bottom).
Revised: 6 (bottom).

Original: It is visible that η = 1  .[-
-]
Revised: It is visible that η = 1 .and η = 0 .[-
-]

Original: and η = 0 .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 BN networks.
Revised: 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 BN networks.

Original: BN requires using a high learning rate.
Revised: BN requires using a high learning rate.

Original: As our conjectures hold for BN network, a natural question is if learning can be ill-conditioned with a low learning rate even when BN is used.
Revised: As our conjectures hold for BN network, a natural question is if learning can be ill-conditioned with a low learning rate even when BN is used.

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

Original: (2019) show that without BN, mini-batch gradients are largely contained in the subspace spanned by the top eigenvectors of noncentered K .
Revised: (2019) show that without BN, 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: (1) SimpleCNN-BN optimized with η = 0 .

Original: 001 , and (2) SimpleCNN optimized with η = 0 .
Revised: 001 , and (2) 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 (minimum) value of (cid:107) g (cid:107) / (cid:107) g 5 (cid:107) is 1 .

Original: 90 ( 1 . 37 ) and 1 . 

Revised: 90 ( 1 . 37 ) and 1 .

Original: ( 1 . 12 ), respectively.
Revised: ( 1 . 12 ), respectively.

Original: Second, the maximum value of λ 1 K is 10 .and 16 , respectively.
Revised: Second, the maximum value of λ 1 K is 10 .and 16 , respectively.

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

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

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

Original: Comparing these differences to differences that are induced by using the highest η = 1 . 0 in
Revised: Comparing these differences to differences that are induced by using the highest η = 1 . 0 in

Original: 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 BN alone (Ghorbani et al., 2019; Page, 2019;
Revised: 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 BN alone (Ghorbani et al., 2019; Page, 2019;

Original: Bjorck et al., 2018).
Revised: Bjorck et al., 2018).

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

Original: Summary.
Revised: Summary.

Original: We have shown that the effects of changing the learning rate described in Con.
Revised: We have shown that the effects of changing the learning rate described in Con.

Original: 1 and
Revised: 1 and

Original: Con.
Revised: Con.

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

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Based on a theoretical model, we conjectured and empirically argued for the existence of the breakeven point on the optimization trajectory induced by SGD.
Revised: Based on a theoretical model we conjectured and empirically argued for the existence of the breakeven point on the optimization trajectory induced by SGD.

Original: Next, we demonstrated that using a high learning rate or a small batch size in SGD has two effects on K and H along the trajectory that we referred to as (1) variance reduction and (2) pre-conditioning.
Revised: Next, we demonstrated that using a high learning rate or a small batch size in SGD has two effects on K and H along the trajectory that we referred to as (1) variance reduction and (2) pre-conditioning.

Original: There are many potential implications of the existence of the break-even point.
Revised: There are many potential implications of the existence of the break-even point.

Original: We investigated one in particular, and demonstrated that using a high learning rate is necessary to achieve the loss smoothing effects previously attributed to batch normalization alone.
Revised: We investigated one in particular and demonstrated that using a high learning rate is necessary to achieve the loss smoothing effects previously attributed to batch normalization alone.

Original: Additionally, the break-even occurs typically early during training, which might be related to the recently discovered phenomenon of the critical learning period in training of deep networks (Achille et al. Golatkar et al., 2019).
Revised: Additionally, the break-even occurs typically early during training, which might be related to the recently discovered phenomenon of the critical learning period in training of deep networks (Achille et al., 2017; Golatkar et al.

Original: We plan to investigate this connection in the future
Revised: We plan to investigate this connection in the future

Original: A P ROOFS
Revised: A P ROOFS

Original: In this section we will formally state and prove the theorem used informally in Sec. 3.
Revised: In this section we will formally state and prove the theorem used informally 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 straightforwarrd implication of Theorem 1 from Wu et al.
Revised: This theorem is a straightforward application of Theorem 1 from Wu et al.

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]

Original: [Equation]
Revised: [Equation]

Original: At the iteration t ∗ at which training becomes unstable along e 1 H ( t ∗ ) , λ 1 H ( t ∗ ) can be calculated from inequality (3) as
Revised: At the iteration t ∗ at which training becomes unstable along e 1 H ( t ∗ ) , λ 1 H ( t ∗ ) can be calculated from inequality (3) as

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 in the case n = S .
Revised: Note that if n = S the right hand side degenerates to 2 η , which completes the proof in the case 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 value of ψ at which training becomes unstable.
Revised: Similarly, let us denote by ψ ( t ∗ 1 ) and ψ ( t ∗ 2 ) the value of ψ at which training becomes unstable.

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, ψ ( t ) by Assumption 4 decreases with increasing λ 1 H , which corresponds to reducing the distance to the minimizer along e 1 H .
Revised: However, ψ ( t ) by Assumption 4 decreases with increasing λ 1 H , which corresponds to reducing the distance to the minimizer along e 1 H .

Original: Hence, that would imply λ 1 H ( t ∗ 1 ) <
Revised: Hence, that would imply λ 1 H ( t ∗ 1 ) <

Original: λ 1 H ( t ∗ 2 ) , which contradicts λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .
Revised: λ 1 H ( t ∗ 2 ) , which contradicts λ 1 H ( t ∗ 1 ) > λ 1 H ( t ∗ 2 ) .

Original: Repeating the same argument for batch size completes the proof.
Revised: Repeating the same argument for batch size 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 ) has reached a smaller value, it means ψ ( t ∗ 2 ) has increased to a larger value, i.e. ψ ( t ∗ 2 ) > ψ
Revised: Hence, if λ 1 H ( t ∗ 2 ) has reached a smaller value, it means ψ ( t ∗ 2 ) has increased to a larger value, i.e. ψ ( t ∗ 2 ) > ψ

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

Original: However, 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, 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: B A PPROXIMATING THE EIGENSPACE OF K AND H
Revised: B A PPROXIMATING THE EIGENSPACE OF K AND H

Original: A small subset of the training set approximates well the largest eigenvalues of the true Hessian on the CIFAR-10 dataset (Alain et al., 2019).
Revised: A small subset of the training set approximates well the largest eigenvalues of the true Hessian on the CIFAR-10 dataset (Alain et al., 2019).

Original: This might be due to the hierarchical structure of the
Revised: This might be due to the hierarchical structure of the

Original: Hessian (Papyan, 2019; Fort & Ganguli, 2019); the largest eigenvalues are connected to the class structure in the data.
Revised: Hessian (Papyan, 2019; Fort & Ganguli, 2019); the largest eigenvalues are connected to the class structure in the data.

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 SCIPY Python package.

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

Original: We sample L mini-batch gradient of size M .
Revised: We sample L mini-batch gradient of size M .

Original: Then following Papyan (2019) we compute the corresponding Gram matrix ˆ K that has entries ˆK ij = (cid:104) g
Revised: Then following Papyan (2019) we compute the corresponding Gram matrix ˆ K that has entries ˆK ij = (cid:104) g

Original: i − g, g j − g (cid:105) , where g is the full batch gradient.
Revised: i − g, g j − g (cid:105) , where g is the full batch gradient.

Original: Note that ˆK ij is only L × L dimensional.
Revised: Note that ˆK ij is only L × L dimensional.

Original: It can be shown that in expectation of mini-batches ˆ K has the same eigenspectrum as
Revised: It can be shown that in expectation of mini-batches ˆ K has the same eigenspectrum as

Original: K .
Revised: K .

Original: In all experiments we ﬁx L = 25 .
Revised: In all experiments we ﬁx L = 25 .

Original: When comparing different batch sizes we use M = 128 , otherwise we use the same batch size as the one used to train the model.
Revised: When comparing different batch sizes we use M = 128 , otherwise we use the same batch size as the one used to train the model.

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

Original: To compute Tr( K ) we compute the trace of ˆK .
Revised: To compute Tr( K ) we compute the trace of ˆK .

Original: Due to large computational cost of estimating top eigenvalues using the Lanczos algorithm, we approximate in most experiments the Tr( H ) using only the top 5 eigenvalues of H .
Revised: Due to large computational cost of estimating top eigenvalues using the Lanczos algorithm, we approximate in most experiments the Tr( H ) using only the top 5 eigenvalues of H .

Original: A natural question is whether using M = 128 approximates well the underlying K .
Revised: A natural question is whether using M = 128 approximates well the underlying K .

Original: To investigate this we compare λ 1 K evaluated using either M = 128 or M = 1 .
Revised: To investigate this we compare λ 1 K evaluated using either M = 128 or M = 1 .

Original: We adapt L to keep the same number of N = LM of used training examples in the computation.
Revised: We adapt L to keep the same number of N = LM of used training examples in the computation.

Original: We test this on SimpleCNN model with η = 0 .
Revised: We test this on SimpleCNN model with η = 0 .

Original: We show that the two are strongly correlated in Fig.
Revised: We show that the two are strongly correlated in Fig.

Original: [Figure]
Revised: [Figure]

Original: C E XPERIMENTAL DETAILS
Revised: C E XPERIMENTAL DETAILS

Original: In this section we describe all the details for experiments in the main text and in the Appendix.
Revised: In this section we describe all the details for experiments in the main text and in the Appendix.

Original: ResNet-32 on CIFAR The model trained is the ResNet-32 (He et al., 2015b).
Revised: ResNet-32 on CIFAR The model trained is the ResNet-32 (He et al., 2015).

Original: The network is trained for 200 epochs with a batch size equal to 128.
Revised: The network is trained for 200 epochs with a batch size equal to 128.

Original: The dataset used is CIFAR-10.
Revised: The dataset used is CIFAR-10.

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

Original: The default learning rate is 0 . 05 , and the default batch size is 128 .
Revised: The default learning rate is 0 . 05 , and the default batch size is 128 .

Original: Weight decay 0 . 0001 is applied to all convolutional layers.
Revised: Weight decay 0 . 0001 is applied to all convolutional layers.

Original: SimpleCNN on CIFAR
Revised: SimpleCNN on CIFAR

Original: The network used is a simple convolutional network based on example from Keras repository (Chollet et al., 2015).
Revised: The network used is a simple convolutional network based on example from Keras repository (Chollet et al., 2015).

Original: First, the data is passed through two convolutional layers with 32 ﬁlters, both using same padding, ReLU activations, and the kernel size of 3x3.
Revised: First, the data is passed through two convolutional layers with 32 ﬁlters, both using same padding, ReLU activations, and the kernel size of 3x3.

Original: Next, the second pair of layers is used with 64 ﬁlters (the rest of parameters are the same).
Revised: Next, the second pair of layers is used with 64 ﬁlters (the rest of parameters are the same).

Original: Each pair ends with a max-pooling layer with a window of size 2.
Revised: Each pair ends with a max-pooling layer with a window of size 2.

Original: Before the classiﬁcation layer, a densely connected layer with 128 units is used.
Revised: Before the classiﬁcation layer, a densely connected layer with 128 units is used.

Original: The data follows the same scheme as in the ResNetexperiment.
Revised: The data follows the same scheme as in the ResNetexperiment.

Original: The default learning rate is 0 . 05 , and the default batch size is 128 .
Revised: The default learning rate is 0 . 05 , and the default batch size is 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 BERT-base (Devlin et al., 2018), pretrained on multilingual data 4 .

Original: The model is trained on MultiNLI dataset (Williams et al., 2018) with the maximum sentence length equal to 40.
Revised: The model is trained on 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 with a batch size of 32.

Original: MLP on FMNIST This experiment is using a Multi Layer Perceptron network with two hidden layers of size 300 and 100, both with ReLU activations.
Revised: MLP on FMNIST This experiment is using a Multi Layer Perceptron network with two hidden layers of size 300 and 100, both with ReLU activations.

Original: The data is normalized to the [0 , 1] range and no augmentation is being used.
Revised: The data is normalized to the [0 , 1] range and no augmentation is being used.

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: 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 length of the sequence equal to 80.

Original: The model is trained for 100 epochs with a batch size of 128.
Revised: The model is trained for 100 epochs with a batch size of 128.

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 (Huang et al., 2016).

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 is being used.
Revised: The images are centered, but no augmentation is being used.

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

Original: Neither dropout nor weight decay is used for training.
Revised: Neither dropout nor weight decay is used for training.

Original: D A DDITIONAL EXPERIMENTS FOR S EC . 4. 
Revised: D A DDITIONAL EXPERIMENTS FOR S EC . 4.

Original: In this section we repeat experiments from Sec.
Revised: In this section we repeat experiments from Sec.

Original: 4.2 in various other settings, as well as include additional data from settings already included in the main text.
Revised: 4.2 in various other settings, as well as include additional data from settings already included in the main text.

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

Original: In Fig.
Revised: In Fig.

Original: 8 and Fig.
Revised: 8 and Fig.

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

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

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

Original: In Fig.
Revised: In Fig.

Original: 8 and Fig.
Revised: 8 and Fig.

Original: 9 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all runs on the SimpleCNN model and the CIFAR-10 dataset .
Revised: 9 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all runs on the SimpleCNN model and 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: 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 runs on the LSTM model and the IMDB dataset.
Revised: 13 we report accuracy on the training set and the validation set, λ 1 H , and Tr( K ) for all runs on the LSTM model and the IMDB dataset.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

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

Original: In Fig.
Revised: In Fig.

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

Original: 15 we report results for the BERT model on the MNLI dataset.
Revised: 15 we report results for the BERT model on the MNLI dataset.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

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

Original: In Fig.
Revised: In Fig.

Original: 16 and Fig.
Revised: 16 and Fig.

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

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

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

Original: In Fig.
Revised: In Fig.

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

Original: 19 we report results for the DenseNet-121 model on the ImageNet dataset
Revised: 19 we report results for the DenseNet-121 model on the ImageNet dataset

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