A UTO LR: A M ETHOD FOR A UTOMATIC T UNING OF L EARNING R ATE
Anonymous authors
Paper under double-blind review
A BSTRACT
One very important hyperparameter for training deep neural networks is the learning rate of the optimizer.
The choice of learning rate schedule determines the computational cost of getting close to a minima, how close you actually get to the minima, and most importantly the kind of local minima (wide/narrow) attained.
The kind of minima attained has a signiﬁcant impact on the generalization accuracy of the network.
Current systems employ hand tuned learning rate schedules, which are painstakingly tuned for each network and dataset.
Given that the state space of schedules is huge, ﬁnding a satisfactory learning rate schedule can be very time consuming.
In this paper, we present AutoLR , a method for auto-tuning the learning rate as training proceeds.
Our method works with any optimizer, and we demonstrate results on SGD, Momentum, and Adam optimizers.
We extensively evaluate AutoLR on multiple datasets, models, and across multiple optimizers.
We compare favorably against state of the art learning rate schedules for the given dataset and models, including for ImageNet on Resnet-50, Cifaron Resnet-18, and SQuAD ﬁne-tuning on BERT.
For example, AutoLR achieves an EM score of 81.2 on SQuAD v1.1 with BERT BASE compared to 80.8 reported in (Devlin et al.
by just auto-tuning the learning rate schedule.
To the best of our knowledge, this is the ﬁrst automatic learning rate tuning scheme to achieve state of the art generalization accuracy on these datasets with the given models.
I NTRODUCTION
Learning rate is one of the most important hyperparameters that impact deep neural network (DNN) training performance.
It determines how fast we get close to a minima, which in turn determines the computational cost of optimization.
It also determines how close we get to the minima, e.g. higher learning rates may get in the neighborhood of a minima much faster, but then just bounce at a height above the minima (Xing et al.
; Nar & Sastry (2018)).
The learning rate also determines the kind of minima (e.g. wide vs narrow) attained (Keskar et al.
(2016)), which has a signiﬁcant impact on the generalization accuracy.
Therefore, it is not surprising that a lot of effort has gone into automatically tuning the learning rate (Schraudolph (1999); Schaul et al.
(2013b); Rolinek & Martius (2018)).
However, till date, none of these techniques have been able to deliver state of the art test accuracy on standard benchmarks.
Instead, deep learning researchers today rely on a mixture of brute force search, augmented with simple heuristics such as using a staircase, polynomial, or exponential decay-based learning rate schedules.
This problem is further exacerbated by the fact that different optimizers such as SGD or
Adam work best on different datasets and require very different learning rate schedules.
For example, while training image datasets such as Cifar-10 and ImageNet with SGD or Momentum, a staircase schedule with high starting learning rate typically performs best.
On the other hand, with the Adam optimizer, a smaller starting learning rate is generally used.
Further, in NLP tasks such as machine translation with Adam, typically a linear warmup is followed by polynomial decay, e.g. the
Transformer network (Vaswani et al.
(2017)) uses a linear warmup followed by an inverse square root decay.
Similar schedules are used for BERT (Devlin et al.
pre-training, while BERT ﬁne-tuning uses a linear decay.
Our key idea to tackle this problem is driven by the following simple observation.
Consider training
Cifar-10/ImageNet using the typical staircase schedule.
One can observe that the initial high learning rate results in training loss stagnating after a few epochs.
However, if one were to decay the learning rate when loss stagnates, the achieved test accuracy is considerably lower as compared to letting the high learning rate continue for longer, despite no apparent improvement in training loss.
Our insight is that this initial phase of training using a high learning rate, even with no improvement in loss, is crucial for generalization and is one of the key missing pieces from prior attempts at automatic learning rate tuning.
We hypothesize that for DNNs, the number of narrow minimas far outnumber the wide minimas.
To generalize well, we want the optimizer to land in wide minimas.
An interesting intuitive observation is that a large learning rate can escape narrow minimas easily (as the optimizer can jump out of them with large steps), however once it reaches a wide minima, it is likely to get stuck in it (if the “width” of the wide minima is large compared to the step size).
The above hypothesis motivates our Explore-Exploit scheme where we force the optimizer to ﬁrst explore the landscape with a high learning rate for sometime in order to land in a wide minima.
We should give the explore phase enough time so that the probability of landing in a wide minima is high.
Once we are conﬁdent that the optimizer is stuck in the vicinity of a wide minima, we activate the Exploit phase of AutoLR .
The basic idea of the exploit phase is as follows.
We ﬁrst look at how local perturbations in the current learning rate impact the training loss.
Since we only look at local perturbations, we can model the loss as a function of these perturbations via a Taylor series expansion.
We then make a quadratic approximation, and solve for the optimal perturbation in the learning rate which will minimize the loss.
This is similar to Newton’s method, but applied only in the descent direction (section 2).
We do extensive validation that the quadratic is a good approximation here.
Although the basic idea is simple, it is complicated by the fact that deep learning relies on stochastic optimization methods, such as SGD, which causes the loss values to be noisy across mini-batches, potentially throwing off our estimates.
We discuss how we handle stochasticity in section 2.2.
We demonstrate the efﬁcacy of the explore-exploit approach by evaluating AutoLR across a wide range of models and datasets, ranging from NLP (SQuAD on BERT-base, Transformer on IWSLT) to CNNs (e.g. ImageNet on ResNet-50, Cifar-10 on ResNet18), and across multiple optimizers:
SGD, Momentum and Adam.
In all cases, AutoLR matches or beats test accuracy of state-of-theart hand-tuned learning rate schedules.
For example, on SQuAD v1.1 ﬁne-tuning with BERT BASE ,
AutoLR is able achieve an EM score of 81.2, compared to 80.8 reported in Devlin et al.
(2018) by just auto-tuning the learning rate schedule (all other parameters were unchanged).
To the best of our knowledge, AutoLR is the ﬁrst automatic learning rate scheduler to achieve state of the art results on these datasets for the given models.
We show extensive evaluation of our method in section 3.
The main contributions of our work are:
The observation that an initial Explore phase with high learning rate is crucial for good generalization in automatic learning rate schemes.
Incorporating this observation via an Explore-Exploit approach with a novel exploitation scheme for tuning learning rate using a local approximation of the optimization landscape.
The ﬁrst automatic learning rate tuning scheme that beats or achieves generalization of state of the art learning rate schedules in multiple models/datasets including ImageNet.
M ETHOD
Let L ( θ ) be the loss of the network as a function of its parameters θ .
Note that in practice, since we work with stochastic optimizers such as SGD, the loss computed in each minibatch is only an approximation of the true loss.
This distinction is important when we come to the estimation of the best learning rate, but can be ignored for the formulation below.
The loss of the network in the next time step is L ( θ − η(cid:126)d ) , where (cid:126)d is the step direction, and η is the learning rate.
If we think of L ( θ − η(cid:126)d ) as a function of the learning rate η , our task is to ﬁnd an η which minimizes this function.
Since it is hard to directly estimate L as a function of η , we instead consider what happens if we perturb η by a small amount (cid:15) , i.e we look at L ( θ − ( η + (cid:15) )
(cid:126) d ) .
Looking at this as a function of (cid:15) , and applying Taylor series expansion we get,
[Equation]
where (cid:126)g is the gradient of L w.r.t θ , and H is the Hessian.
Note that computing accurate value of the gradient (cid:126)g is expensive, as it requires going over the entire dataset (not just a minibatch), while computing the Hessian H is typically intractable for deep networks which have millions of parameters.
However, it turns out that we don’t actually need the values of (cid:126)g and H .
In fact, we can just look at ˆ L as a quadratic function of (cid:15) , i.e.
[Equation]
To compute { k 0 , k 1 , k 2 } , we simply evaluate the loss at a few values of (cid:15) and ﬁt a quadratic polynomial.
Note that we are able to get away with having to calculate the high dimensional (cid:126)g , and
H , because the loss function only looks at the effect of these ﬁrst and second order derivatives in one single direction (cid:126)d , as determined by the terms (cid:126)d T (cid:126)g and (cid:126)d T H (cid:126)d .
Once we know the values of { k 0 , k 1 , k 2 } , we can simply minimize this quadratic to get the optimal value of (cid:15) min = − k 1 2 k 2 to minimize the loss L in the next time step.
Also observe that our method only needs the search direction (cid:126)d for computing the various loss samples.
As a result, our method works for any optimizer, as long as we can access the search direction.
We have implemented our method for multiple optimizers including SGD, Momentum, and Adam.
Note that it is very important to use a perturbation (cid:15) for the Taylor expansion in equation 1, rather than just expanding on η via L ( θ − η(cid:126)d )
=
L ( θ ) − η(cid:126)d T (cid:126)g + 12 η
(cid:126)d T H (cid:126)d + O ( η 3 ) .
This is because η may not be small, causing the O ( η 3 ) error to be quite large.
See appendix A for empirical validation that a second order approximation captures the loss perturbations accurately.
[Figure]
Epsilon Thresholding.
The quadratic approximation in equation 2 yields an optimal value of (cid:15) min = − k
k 2 for the minimum loss L at the next time step.
Note that the quadratic approximation is valid only up to an O ( (cid:15) 3 ) error, while the minimization can yield (cid:15) min with large absolute values, thereby making huge errors.
That is, although our quadratic approximation is valid only for small values of | (cid:15) | , the quadratic’s minima may be beyond this approximation range, as illustrated in Figure 1.
Thus, we need to threshold (cid:15) min to a reasonable value.
Since we are interested in approximating the loss to some precision, we use a relative threshold as follows.
We approximate O ( (cid:15) 3 ) as | (cid:15) | 3 , and bound the error:
[Equation]
where r is the relative error threshold.
2. T ACKLING G ENERALIZATION
The choice of learning rate is an important factor in determining generalization quality of the trained network.
The above method for tuning the learning rate is a local method, and does not take into account the highly non-linear and non-convex global optimization landscape of deep networks.
As a result it can take globally suboptimal decisions which can impact generalization.
In this section we discuss the problem in more detail, and our solution.
Although understanding generalization of deep neural networks is an open problem, there have been interesting ﬁndings recently.
Kawaguchi (2016) found that deep neural networks have many local minimas, but all local minimas are also the global minima (also see Goodfellow et al.
Also, it is widely believed that wide minimas generalize much better than narrow minimas (Hochreiter &
Schmidhuber (1997); Keskar et al. (2016); Jastrzebski et al.
; Wang et al.
(2018)), even though they have the same training loss.
Keskar et al.
(2016) found that small batch SGD generalizes better than large batch SGD and also lands in wider minimas, suggesting that noise in SGD acts as an implicit regularizer.
Interestingly however, more recent work was been able to generalize quite well even with very large batch sizes (Goyal et al.
; McCandlish et al.
; Shallue et al.
(2018)), by scaling the learning rate linearly as a function of the batch size.
This suggests that the lack of noise in large batch SGD can be compensated with high learning rates, and thus the learning rate plays a crucial role in generalization.
Many popular learning rate schedules start the training with high learning rates, and then reduce the learning rate after every few epochs or following some hand tuned fall curve.
An interesting observation in many such examples is that even though a high learning rate stagnates after a few epochs, one still needs to run enough iterations at the higher learning rate in order to get good generalization.
For example, see Figure 2, which shows the training loss of a Resnet-18 model training on Cifar10 dataset at a ﬁxed LR of 0.1 with SGD.
The training loss stagnates after ≈epochs, however as shown in Table 1 1 , generalization continues to improve if we increase the number of epochs trained at a higher learning rate.
[Figure]
Explore Epochs
Test Accuracy 94. [Figure] the other half at 0.001.
Reported results are average over 4 runs.
To understand the above phenomena better, we run the following experiment. We train Cifar-10 on Resnet-18 using a high learning rate of 0 . 1 for only 40 epochs and then use learning rates of 10 −and 10 − 3 for 80 epochs each.
We do this training many times with different random initializations.
On an average this yields a low generalization accuracy of 94 .
58 over 20 runs.
However, we ﬁnd that in 1 of the 20 runs our test accuracy reaches 94 . 98 !
We compute the curvature of the loss surface at the end of training for the various runs and are able to validate that the high test accuracy indeed corresponds to a wider minima than the low test accuracy runs.
Speciﬁcally, we use the highest eigenvalue 2 of the Hessian at the minima as a measure of the minima width (Jastrzebski et al. Keskar et al. We ﬁnd that the high accuracy runs consistently have smaller eigenvalues compared to low accuracy runs. For example the run with highest accuracy of 94 .
98 had an eigenvalue of 0 .
02 , while the run with median accuracy of 94 .had
an eigenvalue of 0 . 05 , and the run with minimum accuracy of 94 .
32 had an eigenvalue of 0 . 10 .
Hypothesis: low density of wide minima in the loss landscape of deep neural networks.
To explain the observation that training at high learning rate for long duration consistently achieves high test accuracy while training at high learning rate for shorter duration until train loss saturates rarely achieves high test accuracy , we hypothesize that for deep neural networks, narrow minima far outnumber wide minima.
An intuitive explanation of the observed phenomena is that a large learning rate can escape narrow minima “valleys” easily (as the optimizer can jump out of them with large steps), however once it reaches a wide minima “valley”, it is likely to get stuck in it (if the “width” of the wide valley is large compared to the step size). See Jastrzebski et al. and 1 We found a bug in our previous Cifar-
Wu et al.
(2018) for a discussion on relationship between learning rate and minima shape. When the density of wide minima is low, using a high learning rate for long enough duration should still be able to ﬁnd the wide minima eventually. However, using a high learning for a short duration will only ﬁnd the wide minima with some probability based on the chance of landing in the wide minima early during training. Since we ﬁnd empirically that the probability of landing in the wide minima early in training is low (about 1/20 based on our experiments), we hypothesize that wider minima are fewer in number than narrow minima.
Please see Section 4 for another experiment in the literature that adds more evidence to this hypothesis.
Explore-Exploit: The above hypothesis motivates our Explore-Exploit scheme where we force the optimizer to ﬁrst explore the landscape with a high learning rate for sometime in order to land in the vicinity of a wide minima.
During the explore phase, we only allow a monotonic increase of LR and disallow any decrease even if suggested by the local quadratic ﬁt.
We should give the explore phase enough time so that the probability of landing in a wide minima is high.
Since the ratio of number of narrow vs wide minimas can depend on the architecture, dataset, etc., predicting the right explore phase duration is hard, and is currently a hyperparameter for AutoLR .
Once we are conﬁdent that the optimizer must be now stuck in the vicinity of a wide minima, we start the exploit phase and allow AutoLR to reduce the learning rate as dictated by the local method to make fast progress and get close to the minima of this hopefully wide minima.
Explore-Exploit improves the generalization of AutoLR signiﬁcantly.
Without an explore phase, the
AutoLR scheme will start reducing the learning rate as soon as it reaches close to a local minima, irrespective of whether it is a narrow or wide minima.
As shown in Figure 3, locally at epoch 50, it makes sense to reduce the LR as it increases the rate of descent.
Without any explore phase, AutoLR ( Eq 2), will make such a decision as well, potentially hurting generalization. [Figure] Explore Epochs [Figure]
[Figure] [Figure] [Figure] To validate the effectiveness of the explore-exploit scheme, we ﬁrst analyze a state-of-the-art learning rate schedule for Cifar-10 on Resnet-18.
This learning rate schedule trains with learning rates of 10 − 1 , 10 − 2 , 10 − 3 for 100, 50, 50 epochs each.
The ﬁrst 100 epochs at a high learning rate of 0 . 1 , can be thought of as an explore phase here, as discussed in Figure 2.
We experimented with reducing this “explore” phase, and as shown in Table 1, it is clear that explore phase plays an important role in the ﬁnal generalization accuracy.
We now evaluate explore-exploit on the same model and dataset but with learning rate determined by AutoLR in the exploit phase, and try different duration of the explore phase.
Table 2 shows the test accuracy for different number of explore epochs.
As shown, the explore phase helps with improving generalization accuracy signiﬁcantly from 93 .
77 without explore to 94 .
79 with 100 explore epochs.
Saturation Threshold:
The above discussion suggests that overall, higher learning rates are better for generalization and thus one should prefer them even though a local analysis may suggest otherwise.
To incorporate this idea into our method we employ what we call a saturation threshold .
The idea is that we want to continue with the current learning rate and not lower it, unless it has saturated in terms of loss drop per iteration.
That is, if the training loss drop rate has dropped below a threshold for the current learning rate, it suggests that the current learning rate has served its purpose, and we can move on to a lower learning rate if suggested so by AutoLR .
We can either use an absolute threshold, or a relative threshold where we use the ratio of loss-drop rate when we started using the current learning rate to the current loss-drop rate.
Finding the best absolute threshold for each model/dataset can be tricky, so we use a relative threshold, which we found easy to tune.
Also, we noticed that the loss drop rate doesn’t change drastically towards the end of the training when the loss has anyways stabilized.
This can cause a high relative saturation threshold to be too strict in the later stages of training, while it would perform well early on.
To handle this, we used a simple strategy where we choose a relative saturation threshold initially, and ﬁrst time the saturation threshold is crossed, we switch to an absolute threshold with the current loss drop date as the absolute saturation threshold.
This essentially amounts to using the relative threshold to determine a good absolute threshold value for the current model and dataset, which is then used subsequently.
Note that, although both explore phase and saturation threshold prefer higher learning rates, they serve different purposes.
While the explore phase ensures that the optimization reaches neighborhood of a wide minima before it reduces the learning rate (by forcing the optimizer to use a higher learning rate for some minimum time, possibly even after it has saturated); saturation threshold adds an overall preference for a higher learning rate until that learning rate has served its purpose and is not optimizing the loss satisfactorily.
2. H ANDLING S TOCHASTICITY
For computation of the quadratic coefﬁcients in equation 2, we need to compute the loss values ˆ
L ( (cid:15) ) at multiple values of (cid:15) .
Note that in a typical DNN setting we never compute the full loss, but only a stochastic loss based on a given minibatch.
This loss, however, can be noisy and throw off our estimate of (cid:15) min .
To handle this we follow two simple strategies.
The ﬁrst is to use a bigger minibatch (called superbatch ) for computing the loss values, and the second is to use the same superbatch for computing all the losses for a particular estimate.
The same strategy is used when computing loss drop rates for the saturation threshold check.
We use an integer multiple of minibatches to make a superbatch.
This allows us to compute the superbatch loss by simply averaging multiple single batch forward passes, and does not increase the peak GPU memory usage.
We use a conservative superbatch size of 100 for all our examples.
See section B for details on selection of superbatch size.
Finally, as a safeguard we also added a rollback policy.
In case our system makes a bad call on the learning rate change (most likely because of a bad superbatch sample), which leads to a reduction in loss drop rate, we rollback the decision and revert the state of the network to the time we made the learning rate change.
Although, rollback triggers rarely, it is helpful in preventing the optimization from going astray because of one bad decision.
E XPERIMENTS
We extensively evaluate our method on multiple networks and datasets, as well as multiple optimizers including SGD, Momentum and Adam.
We have implemented AutoLR as an optimizer in
PyTorch (Paszke et al.
(2017)), which wraps an existing optimizer, such as SGD, Adam, etc.
For our experiments, we used an out of the box policy as in Rolinek & Martius (2018), where we only change the optimizer to AutoLR , and don’t modify anything else.
We evaluate on multiple image datasets – Imagenet on Resnet-50, Cifar-10 on Resnet-18, MNIST, FashionMNIST; as well as NLP datasets – Squad v1.1 for BERT ﬁnetuning and IWSLT with transformer networks.
(Note: The source code for AutoLR will be open sourced, and a link to it has been sent to Reviewers and ACs for this submission.)
3. I MAGE N ET IMAGE CLASSIFICATION ON R ESNET -
In this experiment we trained the ImageNet dataset (Russakovsky et al.
on Resnet-50 network 3 .
We evaluated our method on SGD with momentum which performs best for this dataset.
For baseline runs, we used the standard hand-tuned learning rate schedule of 10 − 1 , 10 − 2 , 10 − 3 for 30 epochs each.
With AutoLR , we trained the network with 25 explore epochs, and used the same seed learning rate as baseline, i.e. 0.1.
Table 3 shows the training loss and test accuracies for the various runs.
As shown, we comfortably beat the test accuracy of the baseline.
Figure 4a shows the learning rate and test top-5 accuracy of AutoLR against that of the baseline.
See ﬁgure 9, for more detailed comparisons of training loss, test accuracy, and learning rate.
[Table]
[Figure]
3. C IFAR -10 IMAGE CLASSIFICATION ON R ESNET -
In this experiment we trained the Cifar-10 dataset (Krizhevsky et al.
(2009)) on Resnet-18 network (He et al.
We evaluated our method on SGD, SGD with momentum and Adam optimizers.
For baseline runs of SGD and Momentum, we used a hand-tuned learning rate schedule of 10 − 1 , 10 − 2 , 10 − 3 for 100, 50, 50 epochs, respectively; while for Adam we used a ﬁxed learning rate of 1e-3.
With AutoLR , we trained the network with 100 explore epochs, and used the same seed learning rate as baseline, i.e. 0.1 for SGD, Momentum, and 10 − 3 for Adam.
Table 4 shows the training loss and test accuracy for the various runs.
We also compare against other learning rate schedules such as one-cycle, linear decay and cosine decay for the momentum optimizer in Table 9.
Figure 4b shows the learning rate and test loss of AutoLR against that of the baseline for Momentum optimizer.
See ﬁgures 10, 11, 12 for more detailed comparisons of training loss, test accuracy, and learning rate for SGD, Momentum and Adam optimizers, respectively.
[Table]
See sections G.1 and section G.2 for results on MNIST and FashionMNIST datasets with SGD,
Momentum and Adam optimizers.
3. SQ U AD FINE - TUNING ON BERT
We now evaluate AutoLR on a few NLP tasks.
In the ﬁrst task, we ﬁne-tune the BERT BASE model (Devlin et al.
(2018)) on SQuAD v1.1 (Rajpurkar et al.
BERT ﬁne-tuning is prone to oveﬁtting because of a huge model size compared to the ﬁne-tuning dataset, and is typically run for only a few epochs.
We use the standard baseline which trains for 3 epochs with a seed learning rate of 2 e − 5 with linear decay.
With AutoLR we found that the model overﬁts much earlier, seemingly because AutoLR reduces the training loss much more quickly (see ﬁgure 14), and we thus needed to ﬁne-tune for only 2 epochs.
The AutoLR runs were trained with 2500 explore steps ( ≈ half epoch), and the same seed learning rate of 2 e − 5 as baseline.
Table 5 shows our results over 3 runs.
We achieve a best EM score of 81.2 (after 2 epochs), compared to baseline’s best of 80.7 (after 3 epochs).
We found it interesting that AutoLR was able to tune the learning rate effectively even in the small budget of 2 epochs to beat the baseline number of BERT BASE (see Table-2 of Devlin et al.
(2018) who reported an EM score of 80.8).
Figure 4c shows the learning rate and test EM score of AutoLR against that of the baseline.
See ﬁgure 14 for detailed comparisons.
We also compare against other learning rate schedules such as one-cycle, linear decay and cosine decay in Table 11.
[Table]
3. M ACHINE T RANSLATION ON T RANSFORMER N ETWORK WITH IWSLT
In the second NLP task, we train the Transformer network (Vaswani et al.
on the IWSLT
German-to-English (De-En) dataset (Cettolo et al.
with the Adam optimizer 6 .
For baseline, we used the learning rate schedule mentioned in Vaswani et al.
The baseline learning rate starts at 1 .
25 e − 7 , and is linearly increased for 4000 steps, followed by an inverse square root decay till 50 epochs.
With AutoLR , we trained the network with 8 explore epochs, and used a seed learning rate of 5e-5.
In both cases we use the model checkpoint with least loss on the validation set for computing BLEU scores on the test set.
Table 6 shows the training loss and test accuracy averaged over 3 runs.
As shown, AutoLR achieves a mean test BLEU score of 34.88, compared to 34.70 for the baseline.
Figure 4d shows the learning rate and validation perplexity of AutoLR against that of the baseline.
See ﬁgure 13 for detailed comparisons of training/validation perplexity, learning rate, etc.
We also compare against other learning rate schedules such as one-cycle, linear decay and cosine decay in Table 10.
It is interesting to note that both linear decay and cosine decay schedules outperform the baseline inverse square root decay schedule of Vaswani et al.
[Table]
Section E has more details of the one-cycle, cosine decay and linear decay baselines. Table 7 shows the test accuracy of different experiments on all learning rate schedules tried in this paper. As shown, AutoLR compares favorably against all other learning rate schedules, except slightly missing out in
Cifar-10.
For one-cycle we report results from the default policy of ﬁnding the maximum learning rate, as well as by hand tuning it (see section E for more details).
Similarly, the default cosine decay and linear decay schedules did not work well for IWSLT, so we added a warmup phase in the beginning as is done in the original paper’s baseline.
We report both these results under the default and tuned columns, respectively. An interesting observation is that both cosine and linear decay perform better than the inverse square root decay learning rate schedule suggested in the “Attention is all you need” paper. A similar observation can be made for the Bert ﬁne-tuning for SQuAD. The fact that the original authors used a suboptimal learning rate points to the general difﬁculties in learning rate tuning.
More importantly, these results show that for ﬁxed function classes of learning rate schedules, it is not just important to tune the hyperparameters of each function class, but also try out different function classes.
Since the space of function classes is huge in itself, searching for the best schedule is typically intractable.
[Table]
R ELATED W ORK
Generalization. There has been a lot of work recently on understanding the generalization characteristics of DNNs.
Kawaguchi (2016) found that DNNs have many local minima, but all local minima were also the global minima.
It has been observed by several authors that wide minima generalize better than narrow minima (Arora et al. (2018); Hochreiter & Schmidhuber (1997); Keskar et al.
; Jastrzebski et al.
; Wang et al.
but there have been other work questioning this hypothesis as well (Dinh et al.
; Golatkar et al.
Guiroy et al.
; Jastrzebski et al.
; Yoshida & Miyato (2017)).
Keskar et al.
(2016) found that small batch SGD generalizes better and lands in wider minima than large batch SGD.
However, recent work has been able to generalize quite well even with very large batch sizes (Goyal et al.
; McCandlish et al.
; Shallue et al.
(2018)), by scaling the learning rate linearly as a function of the batch size.
Jastrzebski et al.
(2019) analyze how batch size and learning rate inﬂuence the curvature of not only the SGD endpoint but also the whole trajectory.
They found that small batch or large step SGD have similar characteristics, and yield smaller and earlier peak of spectral norm as well as smaller largest eigenvalue. Dinh et al.
(2017) show analytically using model reparameterization that wide minima can be converted to sharp minima without hurting generalization.
Wang et al.
(2018) analytically show that generalization of a model is related to the Hessian and propose a new metric for the generalization capability of a model that is unaffected by model reparameterization of Dinh et al.
Yoshida & Miyato (2017) argue that regularizing the spectral norm of the weights of the neural network help them generalize better.
On the other hand, Arora et al. (2018) derive generalization bounds by showing that networks with low stable rank (high spectral norm) generalize better. Guiroy et al. (2019) looks at generalization in gradient-based meta-learning and they show experimentally that generalization and wide minima are not always correlated.
Finally, Golatkar et al.
(2019) show that regularization results in higher test accuracy speciﬁcally when it is applied during initial phase of training, similar to the importance of AutoLR ´ s explore phase during initial phase of training.
Lower density of wide minima.
Wu et al.
(2018) compares the sharpness of minima obtained by batch gradient descent (GD) with different learning rates for small neural networks on FashionMNIST and Cifar10 datasets.
They ﬁnd that GD with a given learning rate ﬁnds the theoretically sharpest feasible minima for that learning rate.
Thus, in the presence of several ﬂatter minimas, GD with lower learning rates does not ﬁnd them, leading the authors to conjecture that density of sharper minima is perhaps larger than density of wider minima.
Adaptive learning.
Many adaptive learning rate methods have been developed which adapt the learning rate on a per parameter basis, including AdaGrad (Duchi et al.
(2011)), ADADELTA (Zeiler (2012)), RMSProp (Tieleman & Hinton (2012)), and Adam (Kingma & Ba (2014)).
However, most of these methods still require specifying a global learning rate schedule which impacts the performance signiﬁcantly.
Many methods have been proposed to adaptively tune the global learning rate, such as Schaul et al.
(2013a) which makes an idealized quadratic loss function assumption near the minima, Baydin et al.
(2017) which uses gradient w.r.t the learning rate to update the learning rate in a ﬁrst order fashion, and Rolinek & Martius (2018) which again uses a ﬁrst order method along with an estimate of minimum loss achievable in a step.
While Schaul et al. (2013a) and Baydin et al.
(2017) compare to only constant learning rate baselines, Rolinek & Martius (2018) does compare to state-of-the-art learning rate schedule for Cifar-10, but is unable to match their test accuracy, although does better on train loss. Finally, approaches like one cycle learning (Smith (2018);
Xing et al.
(2018)) have been proposed that linearly increase and decrease the learning rate in a speciﬁc manner.
Based on our evaluation, one cycle scheme works reasonably well with manual tuning of the learning rate but does not always achieve state-of-the-art accuracy.
C ONCLUSIONS
We have presented AutoLR , an Explore-Exploit method for auto learning rate tuning for training deep neural networks.
Our explore phase is based on the hypothesis that narrow minimas far outnumber the wide minimas, and thus require some minimum exploration at a high learning rate to land in a wide minima region with high probability.
We do multiple validations of this hypothesis, but also plan to further study it both theoretically and empirically in future work.
Our exploit phase is based on a local quadratic approximation of loss in the search direction, as a function of perturbations in the learning rate.
We extensively validate AutoLR on both image (ImageNet, Cifar-10, MNIST,
FashionMNIST) and NLP (IWSLT, Squad) datasets on as well as multiple optimizers, and achieve or beat generalization accuracy of state of the art hand tuned learning rate schedules
A V ALIDATION OF Q UADRATIC A PPROXIMATION
We want to validate whether our quadratic approximation is a good approximation for modelling the loss as a function of the perturbation (cid:15) .
Figure 5 shows a few examples demonstrating the effectiveness of second order approximation.
As shown, the loss values at various samples overlap with the quadratic curve almost perfectly, including those not used in estimation of the quadratic (orange triangles).
Also, the estimated loss value at the quadratic minima matches the true loss value there quite well.
Figure 6 shows quadratic plots for more datasets/models.
[Figure] [Figure]
is same as ﬁgure 5.
It can be seen that the quadratic approximation works pretty well.
B S UPERBATCH SIZE SELECTION
To choose an appropriate superbatch size, we measure the standard deviation of loss as a function of superbatch size.
We compute this by evaluating the loss with 10 different randomly sampled superbatches and calculate the standard deviation.
Figure 7 shows the measurements.
We used a conservative superbatch size of 100 in all our examples, as it corresponded to low variance.
Note that higher superbatch sizes add an increased computational overhead, but since we recompute our learning rate infrequently the total overhead is not very high.
[Figure]
C C OMPUTATIONAL C OST
The primary computational cost of our method comes from computing the samples for quadratic approximation.
We use a superbatch size of 100 and pick 5 samples for approximating our quadratic, thus incurring a cost of 500 forward passes each time we need to recompute the learning rate.
Note that we recompute the learning rate only once in recompute window steps.
In most of our examples we keep the recompute window such that our learning rate recomputation occurs 2 times each epoch, except in BERT ﬁne-tuning where we only have 2-3 epochs to train, and thus want more frequent tuning of the learning rate.
In Cifar-10 on Resnet-18 for example, we used a recompute window of 1000 minibatch steps.
A backward pass is typically 2x more computationally expensive than forward pass, and 1000 minibatch steps thus cost 3000 forward pass worth of compute.
Our recomputation costs 500 forward passes of compute, and thus the computational overhead is 500 / 3000 ≈ 16 .
Note, however that in exploit phase, we do a recomputation of learning rate only when saturation threshold is crossed.
So the computational cost comes down even further and 16 .
6% overhead for
Cifar-10 is an upper bound.
Note that since the main focus of this work was to develop an automatic learning rate tuning scheme which generalizes as well as state of the art hand tuned learning rate schedules, we have not invested much effort on reducing the computational cost.
For example, a 100 superbatch size is an overkill and a size of 50 should be suitable for most examples.
Similarly, 3 samples for estimating the quadratic are enough most of the times instead of the 5 used, as the quadratic ﬁt is very close mostly.
Also, we can experiment with higher recompute window sizes to reduce the cost further.
D H YPERPARAMETERS
Table 8 shows the hyperparameters used for all our examples.
As shown, the saturation threshold and epsilon thresholds are pretty stable across all image examples, and the main parameter to tune is the explore epochs, which we typically set to around 20 − 30% of the total budget.
The reason for different scale of epsilon thresholds in NLP datasets compared to image datasets, is because they typically run on much lower learning rates (of the order of 1 e − 5 ) compared to image datasets (of the order of 1 e − 1 to 1 e − 3 ), which suggests that the optimization landscapes are more sensitive to smaller perturbations and thus need more aggressive clipping.
We have some initial ideas on computing this threshold dynamically by looking at change in loss magnitude as we increase the (cid:15) , and choosing a threshold which limits this change magnitude to some amount.
This can be computationally more expensive, as we need to search for the right epsilon, but can be done only few times in the each run.
We also noticed that the loss drop rate changes much more drastically in image datasets than NLP datasets, again most likely because of the very low learning rates used in NLP datasets.
Thus we need to use a lower saturation threshold in NLP experiments compared to image ones.
We found this hyperparameter very easy to set, by simply eyeballing at the loss drop rate changes of a trial run with ﬁxed LR.
We are also looking at ways to automate the explore epochs hyperparameter, by looking at second order information about the loss landscape, etc. to determine if we have reached a wider minima region.
See Jastrzebski et al.
for some interesting analysis on this front.
[Table]
E C OMPARISONS WITH M ORE B ASELINE L EARNING R ATE S CHEDULES
In this section we compare AutoLR against more learning rate schedules – one-cycle, linear decay and cosine decay. One-Cycle : The one-cycle learning rate schedule was proposed in Smith (2018) (also see Smith (2017)). This schedule ﬁrst chooses a maximum learning rate based on an LR Range test.
The
LR range test starts from a small learning rate and keeps increasing the learning rate until the loss starts exploding.
Smith (2018) suggests that the maximum learning rate should be chosen to be bit before the minima, in a region where the loss is still decreasing.
There is subjectivity in making this choice, although some blogs and libraries 7 suggest using a learning rate one order lower than the one at minima.
However, in our experience, we found that this choice did not perform well.
Instead, we found that values closer to the minima performed much better.
In the experiments below we tried multiple values closer to the minima where the loss was still decreasing and report the best results as one-cycle (tuned) .
Results with the default policy of 1 / 10 th the minima are reported as one-cycle (default) .
Once the maximum learning rate is chosen, the one-cycle schedule proceeds as follows.
The learning rate starts at a speciﬁed fraction 8 of the maximum learning rate and is increased linearly to the maximum learning rate for 45 percent of the training budget and then decreased linearly for the remaining 45.
For the ﬁnal 10 percent, the learning rate is reduced by a large factor (we chose a factor of 10).
For the Momentum optimizer, Smith (2018) suggests to reduce the momentum linearly when lr is increasing, and vice versa.
We used an opensource implementation 9 for our experiments.
Linear Decay : The linear decay learning rate schedule simply decays the learning rate linearly to zero starting from a seed LR.
In the case of IWSLT we found that starting from a high seed caused instability issues (most likely because of Adam optimizer.
See Liu et al.
In this case we also ran experiments with a linear warmup followed by this linear decay.
Cosine Decay : The cosine decay learning rate schedule decays the learning rate to zero following a cosine curve, starting from a seed LR.
Similar to Linear decay, in the case of IWSLT, we found that starting from a high seed caused instability issues.
In this case we also ran experiments with a linear warmup followed by this cosine decay.
[Figure]
[Figure]
[Figure]
l o ss
[Figure]
[Figure]
E. C IFAR -
Figure 8a shows the LR range test for Cifar-10 with the Resnet-18 network.
The minima occurs around learning rate of 0.09.
For the maximum learning rate, we choose 9e-3 for the default onecycle policy, but found 0.06 to perform best in our tuning.
For linear, cosine decay schedules we start with a seed learning rate of 0.1 as used in the standard baselines.
Table 9 shows the training loss and test accuracy for the various schedules.
[Table]
-10 on Resnet-18 for more learning rate schedules.
We report the mean and standard deviation over 7 runs.
For One-Cycle, we report experiments with maximum learning rate chosen via the default lr-range ﬁnder policy and via manual tuning.
E. IWSLT’14 DE-EN
Figure 8b shows the LR range test for IWSLT on the transformer networks.
The minima occurs at 9e-4.
For the maximum learning rate, we choose 9e-5 for the default one-cycle policy and found 2eto perform best in our tuning.
For linear, cosine decay schedules we ﬁrst used a seed learning rate of 5e-4 (the maximum learning rate in the schedule in Vaswani et al.
However, we found this to perform very poorly, most likely due to initial instabilities in Adam as discussed in Liu et al.
So we also ran experiments with linear warmup as in Vaswani et al.
(2017) followed by linear/cosine decays which performed much better.
Table 10 shows the training, validation perplexity and BLEU scores for the various schedules.
[Table]
[Figure]
Figure 8c show the LR range test for SQuAD ﬁne-tuning on BERT.
The minima occurs at 1e-4.
For the maximum learning rate, we choose 1e-5 for the default one-cycle policy and found 4e-5 to perform best in our tuning.
For linear, cosine decays we start with a seed learning rate of 2e-5 as used in Devlin et al.
Table 11 show the average training loss, best test EM score, and average test EM, F1 scores for the various schedules.
[Table]
F S EED S ENSITIVITY
We performed a sensitivity analysis on the importance of the choice of seed learning rate for AutoLR .
We ran AutoLR for Cifar-10 on Resnet-18 with different seeds in a 10x range.
Table 12 shows the average test accuracies over 3 runs for the different seeds.
As shown, the seed learning rate can impact the ﬁnal accuracy, but AutoLR is not highly sensitive to it, especially in the lower ranges.
[Table] G M ORE R ESULTS G. MNIST In this experiment we trained the MNIST dataset LeCun et al.
(1998) on a 4 layer network with two convolution layers and two fully connected layers 10 .
We evaluated our method on SGD, SGD with momentum and Adam optimizers.
For baseline runs, we used a ﬁxed learning rate of 0.
for SGD, Momentum, and 1e-3 for Adam; and trained for 50 epochs.
With AutoLR , we trained the network with 10 explore epochs, and used the same seed learning rate as baseline, i.e. 0.1 for
SGD, Momentum, and 1e-3 for Adam.
Table 13 shows the training loss and test accuracy for the various runs.
See ﬁgures 15, 16, 17 for more detailed comparisons of training loss, test accuracy, and learning rate for SGD, Momentum and Adam optimizers, respectively.
G. F ASHION MNIST
In this experiment we trained the Fashion MNIST dataset Xiao et al.
(2017) on the same 4 layer network as for the MNIST dataset.
We evaluated our method on SGD, SGD with momentum and
Adam optimizers.
For baseline runs, we used a ﬁxed learning rate of 0.1 for SGD, Momentum, and 1e-3 for Adam; and trained for 50 epochs.
With AutoLR , we trained the network with 10 explore epochs, and used the same seed learning rate as baseline, i.e. 0.1 for SGD, Momentum, and 1e-3 for [Table]
Adam.
Table 14 shows the training loss and test accuracy for the various runs.
See ﬁgures 18, 19,for more detailed comparisons of training loss, test accuracy [Table] and learning rate for SGD, Momentum and Adam optimizers, respectively [Table]
H D ETAILED P LOTS
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]