The Break-Even Point on Optimization Trajectories of Deep Neural Networks

1 INTRODUCTION . The connection between optimization and generalization of deep neural networks ( DNNs ) is not fully understood . 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 ) . 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 ) . These simple examples illustrate limitations of our understanding of DNNs . Understanding the early phase of training has recently emerged as a promising avenue for studying the link between optimization and generalization of DNNs . It has been observed that applying regularization in the early phase of training is necessary to arrive at a well generalizing final solution ( Keskar et al. , 2017 ; Sagun et al. , 2017 ; Achille et al. , 2017 ) . 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 ) . 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 ) . In this work , we study the dependence of the entire optimization trajectory on the early phase of training . 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 . These two matrices capture important and ∗Equal contribution . 1We define it as K = 1 N ∑N i=1 ( gi − g ) T ( gi − g ) , where gi = g ( xi , yi ; θ ) is the gradient of the training loss L with respect to θ on xi , N is the number of training examples , and g is the full-batch gradient . 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 ) . We include a more detailed discussion in Sec . 2 . Our first contribution is a simplified model of the early part of the training trajectory of DNNs . 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 . 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 . We call this point the break-even point and show empirical evidence of its existence in the training of actual DNNs . 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 . Specifically , 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. 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 . 1 for an illustration of this phenomenon ) . 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 ) . Finally , we apply our analysis to a network with batch normalization ( BN ) layers and find that our predictions are valid in this case as well . 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 BN layers ) regions of the loss surface , which was previously attributed to BN alone ( Bjorck et al. , 2018 ; Ghorbani et al. , 2019 ; Page , 2019 ) . 2 RELATED WORK . Implicit regularization induced by the optimization method . The choice of the optimization method implicitly affects generalization performance of deep neural networks ( Neyshabur , 2017 ) . In particular , using a large initial learning rate is known to improve generalization ( Goodfellow et al. , 2016 ; Li et al. , 2019 ) . A classical approach to study these questions is to bound the generalization error using measures such as the norm of the parameters at the final minimum ( Bartlett et al. , 2017 ; Jiang et al. , 2020 ) . An emerging approach is to study the properties of the whole optimization trajectory . Arora et al . ( 2019 ) suggest it is necessary to study the optimization trajectory to understand optimization and generalization of deep networks . In a related work , Erhan et al . ( 2010 ) ; Achille et al . ( 2017 ) show the existence of a critical period of learning . Erhan et al . ( 2010 ) argue that training , unless pretraining is used , is sensitive to shuffling of examples in the first epochs of training . Achille et al . ( 2017 ) ; Golatkar et al . ( 2019 ) ; Sagun et al . ( 2017 ) ; Keskar et al . ( 2017 ) demonstrate that adding regularization in the beginning of training affects the final generalization disproportionately more compared to doing so later . 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 . The covariance of gradients and the Hessian . The Hessian quantifies the local curvature of the loss surface . 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 ) . Keskar et al . ( 2017 ) ; Jastrzebski et al . ( 2017 ) studied the dependence of the Hessian ( at the final minimum ) on the optimization hyperparameters . 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 ) . The covariance of gradients , which we denote by K , encapsulates the geometry and the magnitude of variation in gradients across different samples . The matrix K was related to the generalization error in Roux et al . ( 2008 ) ; Jiang et al . ( 2020 ) . 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 ; Liu et al. , 2020 ; He & Su , 2020 ) . Zhang et al . ( 2019 ) argues that in DNNs the Hessian and the covariance of gradients are close in terms of the largest eigenvalues . Learning dynamics of deep neural networks . 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 ; Yao et al. , 2018 ; Xing et al. , 2018 ; Jastrzebski et al. , 2018 ; Lan et al. , 2019 ) . We are directly inspired by Xing et al . ( 2018 ) who show that for popular classification benchmarks , the cosine of the angle between consecutive optimization steps in SGD is negative . Similar observations can be found in Lan et al . ( 2019 ) . Our theoretical analysis is inspired by Wu et al . ( 2018 ) who study how SGD selects the final minimum from a stability perspective . We apply their methodology to the early phase of training , and make predictions about the entire training trajectory . 3 THE BREAK-EVEN POINT AND THE TWO CONJECTURES ABOUT SGD TRAJECTORY . Our overall motivation is to better understand the connection between optimization and generalization of DNNs . In this section we study how the covariance of gradients ( K ) and the Hessian ( H ) depend on the early phase of training . We are inspired by recent empirical observations showing their importance for optimization and generalization of DNNs ( see Sec . 2 for a detailed discussion ) . 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 . Informally speaking , one scenario we study here is when this initial growth is rapid enough to destabilize training . Inspired by Wu et al . ( 2018 ) , we formalize this intuition using concepts from dynamical stability . 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 . Definitions . We begin by introducing the notation . Let us denote the loss on an example ( x , y ) by L ( x , y ; θ ) , where θ is a D-dimensional parameter vector . The two key objects we study are the Hessian of the training loss ( H ) , and the covariance of gradients K = 1N ∑N i=1 ( gi − g ) T ( gi − g ) , where gi = g ( xi , yi ; θ ) 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 . We denote the i-th normalized eigenvector and eigenvalue of a matrix A by eiA and λ i A . Both H and K are computed at a given θ , but we omit this dependence in the notation . Let t index steps of optimization , and let θ ( t ) denote the parameter vector at optimization step t. Inspired by Wu et al . ( 2018 ) we introduce the following condition to quantify stability at a given θ ( t ) . Let us denote the projection of parameters θ onto e1H by ψ = 〈θ , e1H〉 . With a slight abuse of notation let g ( ψ ) = 〈g ( θ ) , e1H〉 . We say SGD is unstable along e1H at θ ( t ) if the norm of elements of sequence ψ ( τ + 1 ) = ψ ( τ ) − ηg ( ψ ( τ ) ) diverges when τ →∞ , where ψ ( 0 ) = θ ( t ) . The sequence ψ ( τ ) represents optimization trajectory in which every step t′ > t is projected onto e1H . Assumptions . Based on recent empirical studies , we make the following assumptions . 1 . The loss surface projected onto e1H is a quadratic one-dimensional function of the form f ( ψ ) = ∑N i=1 ( ψ − ψ∗ ) 2Hi . The same assumption was made in Wu et al . ( 2018 ) , but for all directions in the weight space . Alain et al . ( 2019 ) show empirically that the loss averaged over all training examples is well approximated by a quadratic function along e1H . 2 . The eigenvectors e1H and e 1 K are co-linear , i.e . e 1 H = ±e1K , and λ1K = αλ1H for some α ∈ R. 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 ) . Zhang et al . ( 2019 ) shows empirical evidence for co-linearity of the largest eigenvalues of K and H. 3 . If optimization is not stable along e1H at a given θ ( t ) , λ 1 H decreases in the next step , and the distance to the minimum along e1H increases in the next step . 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 ; Wu et al. , 2018 ; Jastrzebski et al. , 2018 ) . 4 . The spectral norm of H , λ1H , increases during training and the distance to the minimum along e1H decreases , unless increasing λ 1 H would lead to entering a region where training is not stable along e1H . This is inspired by ( Keskar et al. , 2017 ; Goodfellow et al. , 2016 ; Sagun et al. , 2017 ; Jastrzebski et al. , 2018 ; Fort & Scherlis , 2019 ; Fort & Ganguli , 2019 ) who show that in many settings λ1H or gradient norm increases in the beginning of training , while at the same time the overall training loss decreases . Finally , we also assume that S N , i.e . that the batch size is small compared to the number of training examples . These assumptions are only used to build a theoretical model for the early phase of training . Its main purpose is to make predictions about the training procedure that we test empirically in Sec . 4 . Reaching the break-even point earlier for a larger learning rate or a smaller batch size . Let us restrict ourselves to the case when training is initialized at θ ( 0 ) at which SGD is stable along e1H ( 0 ) . 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 . Consider two optimization trajectories for η1 and η2 , where η1 > η2 , that are initialized at the same θ0 , where optimization is stable along e1H ( t ) and λ 1 H ( t ) > 0 . Under Assumption 1 the loss surface along e1H ( t ) can be expressed as f ( ψ ) = ∑N i=1 ( ψ − ψ∗ ) 2Hi ( t ) , where Hi ( t ) ∈ R. It can be shown that at any iteration t the necessary and sufficient condition for SGD to be stable along e1H ( t ) is : ( 1− ηλ1H ( t ) ) 2 + s ( t ) 2 η2 ( N − S ) S ( N − 1 ) ≤ 1 , ( 1 ) where N is the training set size and s ( t ) 2 = Var [ Hi ( t ) ] over the training examples . A proof can be found in ( Wu et al. , 2018 ) . We call this point on the trajectory on which the LHS of Eq . 1 becomes equal to 1 for the first time the break-even point . By definition , there exists only a single break-even point on the training trajectory . Under Assumption 3 , λ1H ( t ) and λ 1 K ( t ) increase over time . If S = N , the break-even point is reached at λ1H ( t ) = 2 η . More generally , it can be shown that for η1 , the break-even point is reached 2We include a similar argument for the opposite case in App . B. for a lower magnitude of λ1H ( t ) than for η2 . The same reasoning can be repeated for S ( in which case we assume N S ) . We state this formally and prove in App . B . Under Assumption 4 , after passing the break-even point on the training trajectory , SGD does not enter regions where either λ1H 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 e1H . Two conjectures about real DNNs . Assuming that real DNNs reach the break-even point , we make the following two conjectures about their optimization trajectory . The most direct implication of reaching the break-even point is that λ1K and λ 1 H at the break-even point depend on η and S , which we formalize as : Conjecture 1 ( Variance reduction effect of SGD ) . Along the SGD trajectory , the maximum attained values of λ1H and λ 1 K are smaller for a larger learning rate or a smaller batch size . We refer to Conjecture 1 as variance reduction effect of SGD because reducing λ1K can be shown to reduce the L2 distance between the full-batch gradient , and the mini-batch gradient . We expect that similar effects exist for other optimization or regularization methods . We leave investigating them for future work . Next , we make another , stronger , conjecture . It is plausible to assume that reaching the break-even point affects to a lesser degree λiH and λ i K for i 6= 1 because increasing their values does not impact stability along e1H . Based on this we conjecture that : Conjecture 2 ( Pre-conditioning effect of SGD ) . Along the SGD trajectory , the maximum attained values of λ ∗ K λ1K and λ ∗ H λ1H 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 . Furthermore , the maximum attained values of Tr ( K ) and Tr ( H ) are smaller for a larger learning rate or a smaller batch size . 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 . Both conjectures are valid only for learning rates and batch sizes that guarantee that training converges . From the optimization perspective , the effects discussed above are desirable . 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 ) . There also exists a connection between these effects and generalization ( Jiang et al. , 2020 ) , which we discuss towards the end of the paper .

Summary: This work analyzes the optimization of deep neural networks from the point of view of the different learning trajectories obtained during different learning settings as brought about by different hyperparameters in optimization. Specifically, the authors consider how the batch size (S) and the step-size (eta) hyper-parameters modify the learning trajectory. The authors conduct their analysis for networks with and without BatchNorm.