The large learning rate phase of deep learning

The choice of initial learning rate can have a profound effect on the performance of deep networks . We present empirical evidence that networks exhibit sharply distinct behaviors at small and large learning rates . In the small learning rate phase , training can be understood using the existing theory of infinitely wide neural networks . At large learning rates , we find that networks exhibit qualitatively distinct phenomena that can not be explained by existing theory : The loss grows during the early part of training , and optimization eventually converges to a flatter minimum . Furthermore , we find that the optimal performance is often found in the large learning rate phase . To better understand this behavior we analyze the dynamics of a two-layer linear network and prove that it exhibits these different phases . We find good agreement between our analysis and the training dynamics observed in realistic deep learning settings . 1 INTRODUCTION . Deep learning has shown remarkable success across a variety of tasks . At the same time , our theoretical understanding of deep learning methods remains limited . In particular , the interplay between training dynamics , properties of the learned network , and generalization remains a largely open problem . In tackling this problem , much progress has been made by studying deep neural networks whose hidden layers are wide . In the limit of infinite width , connections between stochastic gradient descent ( SGD ) dynamics of neural networks , compositional kernels , and linear models have been made . These connections hold when the learning rate is sufficiently small . However , a theory of the dynamics of deep networks that operate outside this regime remains largely open . In this work , we present evidence that SGD dynamics change significantly when the learning rate is above a critical value , ηcrit , determined by the local curvature of the loss landscape at initialization . These dynamics are stable above the critical learning rate , up to a maximum learning rate ηmax . Training at these large learning rates results in different signatures than observed for learning rates η < ηcrit : the loss initially increases and peaks before decreasing again , and the local curvature drops significantly early in training . We typically find that the best performance is obtained when training above the critical learning rate . Empirically , we find these two learning rate regimes are robust , holding across a variety of architectural and data settings . Figure 1 highlights our key findings . We now describe the main contributions of this work . 1.1 TRAINING WITH A LARGE LEARNING RATE LEADS TO A CATAPULT EFFECT . Consider a deep network defined by the network function f ( θ , x ) , where θ are the model parameters and x the input . We define the curvature λt at training step t to be the max eigenvalue of the Fisher Information Matrix , Ft : = Ex [ ∇θf ( θt , x ) ∇θf ( θt , x ) T ] Amari et al . ( 2000 ) ; Karakida et al . ( 2018 ) . Equivalently , λt is the max eigenvalue of the Neural Tangent Kernel Jacot et al . ( 2018 ) . Figure 2 shows the results of training several deep networks with mean squared error ( MSE ) loss using SGD with a range of learning rates . The loss and curvature are measured at every step during the early part of training . We notice the following effects , which occur when the learning rate is above the critical value ηcrit = 2/λ0 , where λ0 is the curvature at initialization.1 1 . During the first few steps of training , the loss grows significantly compared to its initial value before it begins decreasing . We call this the catapult effect . 2 . Over the same time frame , the curvature decreases until it is below 2/η . We can build intuition for these effects using loss landscape considerations . Consider a linear model where the curvature of the loss landscape is given by λ0 . Here , curvature means the largest eigenvalue of the linear model kernel . The model can be trained using gradient descent as long as the learning rate η obeys η < 2/λ0 . When η > 2/λ0 , the loss diverges and optimization fails . Next , consider a deep network . If we train the model with learning rate η > ηcrit , we may again expect the loss to grow initially , assuming the curvature is approximately constant in the neighborhood of the initial point in parameter space . This is the effect observed in Figure 2 . However , unlike the linear case , optimization may still succeed if gradient descent is able to navigate to an area of the landscape that has lower curvature λ , such that η < 2/λ . This is indeed what we observe in practice . In Figure 1 we show that optimal performance typically occurs when a network is trained in the large learning rate regime . As discussed further in Section 2 , this is true even when the compute budget for smaller learning rates is increased to account for the smaller step size . This is consistent with previous observations in the literature , which showed a correlation between performance and the flatness of the minimum ( Keskar et al. , 2016 ) . 1.2 AT LARGE WIDTH , A SHARP DISTINCTION BETWEEN LEARNING RATES REGIMES . The large width limit of deep networks has been shown to lead to simplified training dynamics that are amenable to theoretical study , as in the case of the Neural Tangent Kernel ( Jacot et al. , 2018 ) . In this work we show that the distinction between small and large learning rates becomes sharply defined at large width . This can be seen in Figures 2c , 2f , which show the curvature of sufficiently wide networks after the initial part of training , as a function of learning rate . When η < ηcrit the curvature is approximately independent of the learning rate , while for η > ηcrit the curvature is lower than 2/η . 1The critical learning rate depends on the scale of initialization through λ0 . In Section 3 we analyze the gradient descent dynamics of 2-layer linear networks , and find that they exhibit similar behavior that can be understood in the large width limit . When training with small learning rates , existing theory can describe the training dynamics of such networks . And , we present new theoretical results that explain the behavior at large learning rates . To summarize our findings , in the large width limit we identify two separate learning rate regimes , or phases , with the following characteristics . Lazy phase : η < 2/λ0 . For sufficiently small learning rate , the curvature λt at training step t remains constant throughout training , and the model becomes equivalent to a linear model ( Jacot et al. , 2018 ; Lee et al. , 2019 ) . The model converges to a nearby point in parameter space , and this behavior is sometimes called lazy training ( Du et al. , 2019 ; Zou et al. , 2018 ; Allen-Zhu et al. , 2019 ; Li & Liang , 2018 ; Chizat et al. , 2019 ) . Catapult phase : ηcrit < η < ηmax . At large learning rates the loss grows to be of order the width n over a number of training steps that is of order log ( n ) . During the same period , the curvature decreases until it is below 2/η . Beyond this point the loss decreases and training converges , ultimately reaching a flat minimum ( relative to that reached in the lazy phase ) . The gradient descent dynamics in this phase are visualized in Figure 1 and in Figure S1 . The maximum learning rate ηmax ( beyond which training no longer converges ) depends on the setup . In our theoretical model the maximum learning rate is ηmax = 4/λ0 . For ReLU networks , we find empirically that ηmax ≈ 12/λ0 . 1.3 LIMITATIONS . Our empirical results and our theoretical analysis focus on the case of MSE loss and on training with vanilla SGD , and do not extend to the case of cross-entropy loss , or to training with other optimizers such as momentum . Additionally , our theoretical analysis involves a 2-layer linear network , and does not apply to the case of networks with non-linearities such as ReLU . 1.4 RELATED WORKS . Our work builds on several existing results , which we now briefly review . The existing theory of infinite width networks is insufficient to describe large learning rates . A recent body of work has investigated the gradient descent dynamics of deep networks in the limit of infinite width ( Daniely , 2017 ; Jacot et al. , 2018 ; Lee et al. , 2019 ; Du et al. , 2019 ; Zou et al. , 2018 ; Allen-Zhu et al. , 2019 ; Li & Liang , 2018 ; Chizat et al. , 2019 ; Mei et al. , 2018 ; Rotskoff & Vanden-Eijnden , 2018 ; Sirignano & Spiliopoulos , 2018 ; Woodworth et al. , 2019 ; Naveh et al . ; Xiao et al. , 2019 ) . Of particular relevance is the work by Jacot et al . ( 2018 ) showing that gradient flow in the space of functions is governed by a dynamical quantity called the Neural Tangent Kernel ( NTK ) which is fixed at its initial value in this limit . Lee et al . ( 2019 ) showed this result is equivalent to training the linearization of a model around its initialization in parameter space . Finally , moving away from the strict limit of infinite width by working perturbatively , Dyer & Gur-Ari ( 2020 ) ; Huang & Yau ( 2019 ) introduced an approach to computing the finite-width corrections to network evolution . Despite this progress , in many practical deep learning settings , the neural network has finite width and evolves nontrivially , with a large change in its associated Neural Tangent Kernel . Depending on the architecture and hyperparameters , such networks may give superior performance . Prior work has compared the performance of finite-width , SGD-trained deep networks with the infinite-width kernels derived from the networks ( Lee et al. , 2018 ; Novak et al. , 2019 ; Arora et al. , 2019 ) . Performance gaps are observed in some cases , notably in convolutional networks , implying that existing infinite-width theory is insufficient to explain the performance of deep networks in such settings where the network evolves nontrivially . Large learning rate SGD improves generalization . SGD training with large initial learning rates often leads to improved performance over training with small initial learning rates ( see ( Li et al. , 2019 ; Leclerc & Madry , 2020 ; Xie et al. , 2020 ; Frankle et al. , 2020 ; Jastrzebski et al. , 2020 ) for recent discussions ) . It has been suggested that one of the mechanisms underlying the benefit of large learning rates is that noise from SGD leads to flat minima , and that flat minima generalize better than sharp minima ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2016 ; Smith & Le , 2018 ; Jiang et al. , 2020 ; Park et al. , 2019 ) ( though see Dinh et al . ( 2017 ) for discussion of some caveats ) . According to this suggestion , training with a large learning rate ( or with a small batch size ) can improve performance because it leads to more stochasticity during training ( Smith & Le , 2018 ; Mandt et al. , 2017 ; Smith et al. , 2017 ; Smith et al. , 2018 ) . We develop a connection between large learning rate and flatness of minima in models trained via SGD . Unlike the relationship explored in most previous work though , this connection is not driven by SGD noise , but arises solely as a result of training with a large initial learning rate , and holds even for full batch gradient descent .

Summary: The paper is a detailed account of how large of a learning rate a given mode can take when it is trained by constant step size SGD. Many papers investigated the effect of learning rate (and batch size and width etc…) to the final accuracy. In general, the findings indicate increased accuracy up to a certain threshold (where the algorithm doesn’t converge any more). Such findings are abundant in the literature. However, the present paper investigates the path that SGD takes to get to find those ‘good’ performing points in the weight space. Two phases emerge: monotonic decay in loss vs catapult regime and the latter appears to perform the best.