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: This paper analyzes the effect of choosing a large step-size on the generalization of deep networks. They suggest that starting with a large learning rate, the loss initially increases before converging to a flatter minimum with improved generalization (catapult effect). When the learning rate is above a certain threshold, the authors observed this catapult effect along with a decrease in the curvature of the landscape. These observations were empirically demonstrated through various experiments and analytically studied for a two-layer linear network.