How to decay your learning rate

1 INTRODUCTION . Learning rate schedules play a crucial role in modern deep learning . They were originally proposed with the goal of reducing noise to ensure the convergence of SGD in convex optimization ( Bottou , 1998 ) . A variety of tuned schedules are often used , some of the most common being step-wise , linear or cosine decay . Each schedule has its own advantages and disadvantages and they all require hyperparameter tuning . Given this heterogeneity , it would be desirable to have a coherent picture of when schedules are useful and to come up with good schedules with minimal tuning . While we do not expect dependence on the initial learning rate in convex optimization , large learning rates behave quite different from small learning rates in deep learning ( Li et al. , 2020a ; Lewkowycz et al. , 2020 ) . We expect the situation to be similar for learning rate schedules : the non-convex landscape makes it desirable to reduce the learning rate as we evolve our models . The goal of this paper is to study empirically ( a ) in which situations schedules are beneficial and ( b ) when during training one should decay the learning rate . Because of the use of stochastic gradients , the loss will have a positive definite contribution due to noise . While this noise is key to explore the landscape , we want to minimize this contribution at the end of training and we will use a simple schedule as our baseline : a constant learning rate with one decay close to the end of training . Training with this smaller learning rate for a short time is expected to reduce the noise without letting the model explore the landscape too much . This is corroborated by the fact that the minimum test error often occurs almost immediately after decaying the learning rate . Part of the paper focuses on comparing the simple schedule with standard complex schedules used in the literature , studying the situations in which these complex schedules are advantageous . We find that complex schedules are considerably helpful whenever the weight norm bounces , which happens often in the usual , optimal setups . We observe that the weight norm ( sum over the squared L2-norm of the weight in each layer ) often presents a bouncing behaviour : it decreases , hits a minimum after and continues increasing . In the presence of a bouncing weight norm , we propose an automatic scheduler which performs as well as fine-tuned schedules . 1.1 OUR CONTRIBUTION . The goal of the paper is to study the benefits of learning rate schedules and when the learning rate should be decayed . We focus on the dynamics of the weight norm ( sum over layers of the square of weights ’ L2-norms ) . The paper is divided into two parts and their main conclusions are : 1 . Decaying the learning rate using complex schedules is only beneficial in the presence of weight bouncing . L2 regularization seems crucial for the weight norm to bounce and in its absence ( which is common in NLP and RL ) we don ’ t see a benefit from complex schedules . This is explored in detail in section 3 and the results are summarized in table 2 . 2 . In the presence of bouncing , one should use the weight norm to inform when to decay the learning rate . We observe that tuned step-wise schedules decay the learning rate when the weight norm converges after bouncing1 . Towards the end of training , a last decay decreases the noise . See figure 1 . We propose an Automatic , Bouncing into Equilibration Learning rate scheduler ( ABEL ) . ABEL is competitive with fine-tuned schedules and needs less tuning ( see table 1 and discussion in section 2 ) . While these two points are somewhat independent , together they provide a very general picture of when to decay the learning rate . For pedagogical purposes , we present the two points in inverse order : first reverse engineer what makes fine-tuned learning rate schedules work and then study to what extent these are helpful across tasks and domains . The origin of weight bouncing . In section 4 , we explain why the weight norm often bounces . Weight bouncing and performance . Generally speaking , weight bouncing occurs when we have non-zero L2 regularization and large enough learning rates . While L2 regularization is crucial in vision tasks , it is not found to be that beneficial in NLP or Reinforcement Learning tasks ( for example Vaswani et al . ( 2017 ) does not use L2 ) . If the weight norm does not bounce , ABEL yields the `` simple '' learning rate schedule that we expect naively from the noise reduction picture : decay the learning rate once towards the end of training . We confirm that , in the absence of bouncing , such a simple schedule is competitive with more complicated ones across a variety of tasks and architectures , see table 2 . We also see that the well known advantage of momentum compared with Adam for image classification in ImageNet ( see Agarwal et al . ( 2020 ) for example ) seems to disappear in the absence of bouncing , when we turn off L2 regularization . Weight norm bouncing thus seems empirically a necessary condition for non-trivial schedules to provide a benefit , but it is not sufficient : we observe that when the datasets are easy enough that simple schedules can get zero training error , schedules do not make a difference for the generalization performance ( models can hit zero training error and still improve the test error see Lewkowycz & Gur-Ari ( 2020 ) for example ) . 1.2 RELATED WORKS . We do not know of any explicit discussion of weight bouncing in the literature . The dynamics of deep networks with L2 regularization has drawn recent attention , see for example van Laarhoven 1We define bouncing by the monotonic decrease of the weight norm followed by a monotonic increase which occurs for a fixed learning rate as can be seen in figure 1 . ( 2017 ) ; Lewkowycz & Gur-Ari ( 2020 ) ; Li et al . ( 2020b ) ; Wan et al . ( 2020 ) ; Kunin et al . ( 2020 ) . The recent paper Wan et al . ( 2020 ) observes that the weight norm equilibration is a dynamical process ( the weight norm still changes even if the equilibrium conditions are approximately satisfied ) which happens soon after the bounce . The classic justification for schedules comes from reducing the noise in a quadratic potential ( Bottou , 1998 ) . Different schedules do not provide an advantage in convex optimization unless there is a substantial mismatch between train and test landscape ( Nakkiran , 2020 ) , however this is not the effect that we are observing in our setup : when schedules are beneficial , their training performance is substantially different , see for example figure S4 . Of course , this is not too surprising because convex optimization does not apply to deep networks . The work of Li et al . ( 2020a ) could be helpful to understand better the theory behind our phenomena , although it is not clear to us how their mechanism can generalize to multiple decays ( or other complex schedules ) . There has been lots of empirical work trying to learn schedules/optimizers see for example Maclaurin et al . ( 2015 ) ; Li & Malik ( 2016 ) ; Li et al . ( 2017 ) ; Wichrowska et al . ( 2017 ) ; Rolinek & Martius ( 2018 ) ; Qi et al . ( 2020 ) ; You et al . ( 2017 , 2020 ) . Our approach does not have an outer loop : the learning rate is decayed depending on the weight norm , which is conceptually similar to the ReduceLROnPlateau scheduler , where the learning rate is decayed when the loss plateaus which is present in most deep learning libraries . However , ReduceLROnPlateau does not perform well across our tasks . A couple of papers which thoroughly compare the performance of learning rate schedules are Shallue et al . ( 2019 ) ; Kaplan et al . ( 2020 ) . 2 AN AUTOMATIC LEARNING RATE SCHEDULE BASED ON THE WEIGHT NORM . ABEL AND ITS MOTIVATION . From the two setups in figure 1 it seems that optimal schedules tend to decay the learning rate after bouncing , when the weight norm growth slows down2 . We can use this observation to propose ABEL ( Automatic Bouncing into Equilibration Learning rate scheduler ) : a schedule which implements this behaviour automatically , see algorithm 1 . In words , we keep track of the changes in weight norm between subsequent epochs ( denoted by t ) , |wt|2 ⌘ |wt|2 |wt 1|2 . When the sign of |wt|2 flips , it necessarily means that it has gone through a local minimum : because |wt|2 < 0 initially if |wt+1|2 · |wt|2 < 0 , |wt|2 < 0 , then |wt| is a minimum : |wt+1| > |wt| < |wt 1| . After this , the weight norm grows and slows down until at some point |wt|2 is noise dominated . In this regime , |wt|2 will become negative , which we will take as our decaying condition . In order to reduce SGD noise , near the end of training we decay it one last time . In practice we do this last decay at around 85 % of the total training time and as we can see in the SM B.1 , this particular value does not really matter as long as it is a few epochs . This schedule does not require a fixed number of epochs ( as opposed to cosine decay which is strict in this regard ) : one can decide to train models for longer times by loading the checkpoint before this last noise reducing decay , or stop models early ( and do the corresponding last decay for a few epochs ) . Algorithm 1 is an implementation of the idea with the base learning rate and the decay factor as the main hyperparameters . While alternative implementations could be more explicit about the weight norm slowing down after reaching the minimum , they would likely require more hyperparameters . The definition 1 does not guarantee that the learning rate will be decayed after the weight norm reaches a minimum , in particular if the weight norm was noisy , 1 could trigger a learning rate decay without bouncing . In practice , in all the experiments that we ran ( which span variations across 4 datasets , 6 learning rates , 5 decay factors , 4 architectures ) , ABEL successfully decayed the learning rate after the weight norm reached equilibrium . We have decided to focus on the total weight norm , but one might ask what happens with the layerwise weight norm . In SM B.2 , we study the evolution of the weight norm in different layers . We focus on the 10 layers which contribute the most to the weight norm ( these layers account for 50 % of the weight norm ) . We see that most layers exhibit the same bouncing plus slowing down pattern as the total weight norm and this happens at roughly the same time scale . 2The decay rate and base learning rates are not the same between experiments , but here we focus on when the learning rate decay occurs . Algorithm 1 ABEL Scheduler # Repeat every epoch t , reached_minimum=False initially . if ( |wt|2 |wt 1|2 ) · ( |wt 1|2 |wt 2|2 ) < 0 then if reached_minimum then learning_rate = decay_factor · learning_rate reached_minimum=False else reached_minimum=True end if end if # If no fixed budget , decay for a few epochs when training is done instead , see end of sec.2 . if t = last_decay_epoch then

Summary: This work deals with the optimization of learning rate schedules for deep learning. It offers three main contributions   1. An equation relating the change in weight norm to the gradient norm, learning rate, weight decay, and the weight norm--- for weights that go into batch-norm. This equation provides intuition for the behaviour of this weight norm.   2. An algorithm and code for adjusting the learning rate by reducing it by a factor once the norm of the weights reach a maximum, then reducing it again by the same factor after about 85% of the total training time. The algorithm is benchmarked on about 14 learning tasks, within vision, RL, and NLP and shows similar results when compared to other complex learning rate schedulers algorithms. Being adaptive, their algorithm has at least one less hyperparameter compared to the competition and appears more robust.   3. Empirical tests of a conjecture that complex learning rates are only beneficial in the presence of extremums in the overall weight norms.