AutoLRS: Automatic Learning-Rate Schedule by Bayesian Optimization on the Fly

The learning rate ( LR ) schedule is one of the most important hyper-parameters needing careful tuning in training DNNs . However , it is also one of the least automated parts of machine learning systems and usually costs significant manual effort and computing . Though there are pre-defined LR schedules and optimizers with adaptive LR , they introduce new hyperparameters that need to be tuned separately for different tasks/datasets . In this paper , we consider the question : Can we automatically tune the LR over the course of training without human involvement ? We propose an efficient method , AutoLRS , which automatically optimizes the LR for each training stage by modeling training dynamics . AutoLRS aims to find an LR applied to every τ steps that minimizes the resulted validation loss . We solve this black-box optimization on the fly by Bayesian optimization ( BO ) . However , collecting training instances for BO requires a system to evaluate each LR queried by BO ’ s acquisition function for τ steps , which is prohibitively expensive in practice . Instead , we apply each candidate LR for only τ ′ τ steps and train an exponential model to predict the validation loss after τ steps . This mutual-training process between BO and the loss-prediction model allows us to limit the training steps invested in the BO search . We demonstrate the advantages and the generality of AutoLRS through extensive experiments of training DNNs for tasks from diverse domains using different optimizers . The LR schedules auto-generated by AutoLRS lead to a speedup of 1.22× , 1.43× , and 1.5× when training ResNet-50 , Transformer , and BERT , respectively , compared to the LR schedules in their original papers , and an average speedup of 1.31× over state-of-the-art heavily-tuned LR schedules . 1 INTRODUCTION . In the regime of deep learning , the success of training largely depends on the choice of the learning rate ( LR ) schedule , since most optimizers will have difficulty traversing a non-smooth and non-convex loss landscape with multiple local minimums and possibly saddle points ( Kawaguchi , 2016 ; Jin et al. , 2017 ; Goodfellow et al. , 2016 ; Li et al. , 2018a ) . To achieve stable and fast convergence towards a solution with good generalization performance , one has to tune the LR schedules carefully for different tasks ( Nar & Sastry , 2018 ; Jastrzębski et al. , 2017 ) . This tuning is usually non-trivial and requires many trial-and-error iterations that are computationally expensive . Moreover , the randomness of the widely-used mini-batch stochastic gradient descent ( SGD ) may introduce more uncertainty and difficulty in the tuning process . For the same reasons , it is also hard to directly formulate the search of the LR schedule as a well-posed optimization problem and address it through standard optimization . The broadly-adopted strategy is to either pick one from a family of pre-defined LR schedules or apply an optimizer that has a built-in mechanism changing the LR adaptively . However , we have a limited number of choices for pre-defined LR schedules , most of which are simple functions such as exponent or cosine and thus can not perfectly align with the non-smooth loss landscape . The latter set of adaptive optimizers , e.g. , Adam ( Kingma & Ba , 2015 ) and Adadelta ( Zeiler , 2012 ) , are extended from convex optimization and rely on strong assumptions to make the convergence properties hold . Moreover , the methods in both categories introduce new hyper-parameters that have to be tuned separately for different tasks or datasets , requiring significant human involvement . In this paper , we study the question : can we automatically tune the LR over the course of training without human involvement ? At the beginning of every τ steps ( i.e. , a “ stage ” in our method ) , we seek to identify an LR that optimizes the validation loss ( i.e. , an empirical estimate of the generalization error ) at the end of the stage . To do so , we employ Bayesian optimization ( BO ) that treats the validation loss as a black-box function of LR . BO simultaneously updates a posterior estimation of the black-box function and searches for the best LR with respect to the posterior . This approach is , however , computationally expensive since estimating the posterior needs many ( input , output ) instances of the function , and acquiring each instance costs τ steps of training . We , therefore , develop a simple yet efficient approximation : for every LR that BO decides to evaluate , we train the model by using the LR for only τ ′ τ steps and use the validation loss over the τ ′ steps to train a time-series forecasting model that provides a prediction of the validation loss after τ steps . As we will show later , an exponential model suffices to produce accurate predictions when using a small τ ′ = τ/10 . Then , AutoLRS can allow BO to explore ten different LRs in each stage and still bound the total running time to approximately twice the training cost associated with the generated schedule , i.e. , the time spent to find the stage-specific LRs is roughly equal to the time spent training the model with the identified LRs . AutoLRS does not depend on a pre-defined LR schedule , dataset , or a specified task and is compatible with almost all optimizers . Hence , it can be generally deployed across a broad range of ML tasks without much human involvement or expensive tuning over choices of LR schedules and their hyperparameters . Moreover , since it directly minimizes the validation loss , it does not only accelerate the convergence but also improves the generalization compared to just minimizing the training loss . Furthermore , AutoLRS only needs to update two extremely light-weight models , i.e. , the BO posterior and the exponential forecasting model , and it is efficient in exploring the loss landscape . Hence , it does not result in notable extra costs in either memory or computation . Note that AutoLRS searches for better LRs based on the training dynamics , which can be seen as a form of selfsupervision . The interaction between BO and the forecasting model is an example of mutual learning , where one produces training data for the other . In experiments , we apply AutoLRS to train three representative DNNs widely used in practice , i.e. , ResNet-50 ( He et al. , 2016a ) on ImageNet classification ( Russakovsky et al. , 2015 ) ; Transformer ( Vaswani et al. , 2017 ) and BERT ( Devlin et al. , 2019 ) for NLP tasks . Though they have been extensively studied and have hand-tuned LR schedules , the LR schedules computed by AutoLRS are faster than the original , hand-tuned , LR schedules by 1.22× , 1.43× , and 1.5× for training ResNet-50 , Transformer , and BERT , respectively , in terms of the training steps used to update the DNN ( i.e. , excluding the costs of the LR/hyperparameter search ) . It meanwhile achieves test-set performance better or on par with state-of-the-art results . We also carefully hand-tuned two state-of-the-art learning rate schedules , CLR ( Smith , 2017 ) and SGDR ( Loshchilov & Hutter , 2017 ) , and conducted more than ten experiments with different CLR/SGDR hyperparameters on each model . AutoLRS still has an average speedup of 1.29× and 1.34× across the three models , in terms of training steps , compared to the best CLR and SGDR LR schedules , respectively . The AutoLRS implementation is available at https : //github.com/YuchenJin/autolrs . 2 RELATED WORK . Learning rate scheduling : In contrast to traditional LR schedules with a monotone decreasing sequence of LRs and multi-step LR schedule , a recent class of LR schedules propose to apply multiple cycles of LR decay . Cyclical Learning Rate ( CLR ) changes LR from a maximal LR ( ηmax ) to a minimal LR ( ηmin ) at a pre-defined frequency and achieves faster convergence for some DNNs ( Smith , 2017 ) . The approach requires a “ LR range test ” to estimate the minimal and maximal LR . The LR range test trains the model with a linearly-increasing LR between a low LR and a high LR , and finds the LR range ( [ ηmin , ηmax ] ) over which the training loss decreases . The authors proposed three variants of CLR : triangular2 that halves the maximum LR bound after each cycle ; exp_range that exponentially reduces the maximum LR bound after each cycle ; and 1cycle containing only one triangular cycle ( Smith , 2018 ) . Similar to CLR , Stochastic Gradient Descent with Warm Restarts ( SGDR ) restarts the LR and then applies cosine annealing/decay at a pre-defined frequency ( Loshchilov & Hutter , 2017 ) . Neither CLR or SGDR is automatic , because they are quite sensitive to their hyperparameters , which require careful hand-tuning . CLR and SGDR may even cause undesirable divergence in loss during training with suboptimal hyperparameters ( see §5 ) . Learning rate adaptation with hypergradient descent : Aiming for the same goal of automatically tuning the LR , the hypergradient based technique ( Almeida et al. , 1998 ; Franceschi et al. , 2017 ; Baydin et al. , 2018 ; Donini et al. , 2020 ) optimizes the LR schedule by applying gradient descent of the objective function w.r.t . the LR during training . In addition to the initial value of the regular LR , it introduces an additional hypergradient LR whose initial value is another hyperparameter to be specified . We experimentally show that this technique is subject to overfitting , it is quite sensitive to its two hyperparameters , and it is unable to match the state-of-the-art test-set performance on the models we test ( §A.5.1 ) . We also compare its performance against AutoLRS ( §A.5.2 ) . DNN hyperparameter optimization : Automatic hyperparameter searching for DNNs has been broadly studied in recent years . When applied to learning rates , they can determine an optimized value for LR that is kept constant ( or constrained to be a pre-defined shape ) through the entire training process , as opposed to determining an LR schedule . They can be primarily categorized into Bayesian optimization based approaches ( Hutter et al. , 2011 ; Snoek et al. , 2012 ; Bergstra et al. , 2013 ) , bandit-based solutions ( Li et al. , 2017 ; 2018b ) , hybrid approaches that combine bandit-based and Bayesian optimization based approaches ( Falkner et al. , 2018 ; Zela et al. , 2018 ) , and population-based methods ( Jaderberg et al. , 2017 ; Parker-Holder et al. , 2020 ) . It might be possible to extend these techniques to determine a LR schedule with an optimized LR for each training stage , but it is not sample-efficient and time-efficient to do so since the LR schedule would correspond to hundreds or thousands of hyperparameters . Optimization methods with adaptive LR : These optimizers can adaptively adjust LR for each training step by maintaining an estimate of a better learning rate separately for each parameter in the DNN . Adagrad ( Duchi et al. , 2011 ) applies lower LRs to parameters with larger accumulated gradients and higher learning rates to the ones with smaller accumulated gradients . RMSprop ( Tieleman & Hinton , 2012 ) , AdaDelta ( Zeiler , 2012 ) , and Adam ( Kingma & Ba , 2015 ) were later proposed to address the issue in Adagrad that the model stops learning due to the continual decay of LR . These optimizers with adaptive LR are orthogonal to our automatic LR scheduler , and they still require a global learning rate schedule , which can be obtained from our AutoLRS . In particular , their default hyperparameters do not always work well and need careful tuning , e.g. , Adam ’ s default LR 0.001 performs poorly in training BERT and Transformer , and a better-tuned LR schedule can significantly reduce the training time ( §5 ) . Recent optimization methods ( Schaul et al. , 2013 ; Mahsereci & Hennig , 2015 ) proposed to remove the need for LR tuning in SGD altogether , but they are not widely used potentially due to their limited applicability and sub-optimal performance ( Baydin et al. , 2018 ) .

Summary: Training deep neural networks is typically done using gradient-based methods with either pre-defined learning-rate schedules or off-the-shelf adaptive optimizers (such as Adam). The former can not reliably align with the non-linear loss landscape, while the latter add additional hyperparameters to tune. This paper proposes an algorithm for automatically tuning a learning-rate schedule. The method works by modelling the training dynamics and adapting the learning-rate to optimize performance on the validation set. The method does end up introducing additional hyperparameters,