Wide-minima Density Hypothesis and the Explore-Exploit Learning Rate Schedule

1 INTRODUCTION . One of the fascinating properties of deep neural networks ( DNNs ) is their ability to generalize well , i.e. , deliver high accuracy on the unseen test dataset . It is well-known that the learning rate ( LR ) schedules play an important role in the generalization performance ( Keskar et al. , 2016 ; Wu et al. , 2018 ; Goyal et al. , 2017 ) . In this paper , we study the question , what are the key properties of a learning rate schedule that help DNNs generalize well during training ? We start with a series of experiments training Resnet18 on Cifar-10 over 200 epochs . We vary the number of epochs trained at a high LR of 0.1 , called the explore epochs , from 0 to 100 and divide up the remaining epochs equally for training with LRs of 0.01 and 0.001 . Note that the training loss typically stagnates around 50 epochs with 0.1 LR . Despite that , we find that as the number of explore epochs increase to 100 , the average test accuracy also increases . We also find that the minima found in higher test accuracy runs are wider than the minima from lower test accuracy runs , corroborating past work on wide-minima and generalization ( Keskar et al. , 2016 ; Hochreiter & Schmidhuber , 1997 ; Jastrzebski et al. , 2017 ; Wang et al. , 2018 ) . Moreover , what was particularly surprising was that , even when using fewer explore epochs , a few runs out of many trials still resulted in high test accuracies ! Thus , we not only find that an initial exploration phase with a high learning rate is essential to the good generalization of DNNs , but that this exploration phase needs to be run for sufficient time , even if the training loss stagnates much earlier . Further , we find that , even when the exploration phase is not given sufficient time , a few runs still see high test accuracy values . To explain these observations , we hypothesize that , in the DNN loss landscape , the density of narrow minima is significantly higher than that of wide minima . A large learning rate can escape narrow minima 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 ) . With fewer explore epochs , a large learning rate might still get lucky occasionally in finding a wide minima but invariably finds only a narrower minima due to their higher density . As the explore duration increase , the probability of eventually landing in a wide minima also increase . Thus , a minimum duration of explore is necessary to land in a wide minimum with high probability . Heuristic-based LR decay schemes such as cosine decay ( Loshchilov & Hutter , 2016 ) implicitly maintain a higher LR for longer than schemes like linear decay . Thus , the hypothesis also explains cosine decay ’ s better generalization compared to linear decay . Moreover , the hypothesis enables a principled learning rate schedule design that explicitly accounts for the requisite explore duration . Motivated by the hypothesis , we design a novel Explore-Exploit learning rate schedule , where the initial explore phase optimizes at a high learning rate in order to arrive in the vicinity of a wide minimum . This is followed by an exploit phase which descends to the bottom of this wide minimum . We give explore phase enough time so that the probability of landing in a wide minima is high . For the exploit phase , we experimented with multiple schemes , and found a simple , parameterless , linear decay to zero to be effective . Thus , our proposed learning rate schedule optimizes at a constant high learning rate for some minimum time , followed by a linear decay to zero . We call this learning rate schedule the Knee schedule . We extensively evaluate the Knee schedule across a wide range of models and datasets , ranging from NLP ( BERT pre-training , Transformer on WMT ’ 14 ( EN-DE ) and IWSLT ’ 14 ( DE-EN ) ) to CNNs ( ImageNet on ResNet-50 , Cifar-10 on ResNet18 ) , and spanning multiple optimizers : SGD Momentum , Adam , RAdam , and LAMB . In all cases , Knee schedule improves the test accuracy of state-of-the-art hand-tuned learning rate schedules , when trained using the original training budget . The explore duration is a hyper-parameter in Knee schedule but even if we set the explore duration to a fixed 50 % fraction of total training budget , we find that it still outperforms prior schemes . We also experimented with reducing the training budget , and found that Knee schedule can achieve the same accuracy as the baseline under significantly reduced training budgets . For the BERTLARGE pretraining , WMT ’ 14 ( EN-DE ) and ImageNet experiments , we are able to train in 33 % , 57 % and 44 % less training budget , respectively , for the same test accuracy . This corresponds to significant savings in GPU compute , e.g . savings of over 1000 V100 GPU-hours for BERTLARGE pretraining . The main contributions of our work are : 1 . A hypothesis of lower density of wide minima in the DNN loss landscape , backed by extensive experiments , that explains why a high learning rate needs to be maintained for sufficient duration to achieve good generalization . 2 . The hypothesis also explains the good performance of heuristic-based schemes such as cosine decay , and promotes a principled design of learning rate decay schemes . 3 . Motivated by the hypothesis , we design an Explore-Exploit learning rate schedule called Knee schedule that outperforms prior heuristic-based learning rate schedules , including achieving state-of-the-art results in IWSLT ’ 14 ( DE-EN ) and WMT ’ 14 ( DE-EN ) datasets . 2 WIDE-MINIMA DENSITY HYPOTHESIS . Many popular learning rate ( LR ) schedules , such as the step decay schedules for image datasets , start the training with high LR , and then reduce the LR periodically . For example , consider the case of Cifar-10 on Resnet-18 , trained using a typical step LR schedule of 0.1 , 0.01 , 0.001 for 100 , 50 , 50 epochs each . In many such schedules , even though training loss stagnates after several epochs of high LR , one still needs to continue training at high LR in order to get good generalization . For example , Figure 1 shows the training loss for Cifar-10 on Resnet-18 , trained with a fixed LR of 0.1 ( orange curve ) , compared to a model trained via a step schedule with LR reduced at epoch 50 ( blue curve ) . As can be seen from the figure , the training loss stagnates after ≈ 50 epochs for the orange curve , and locally it makes sense to reduce the learning rate to decrease the loss . However , as shown in Table 1 , generalization is directly correlated with duration of training at high LR , with the highest test accuracy achieved when the high LR is used for 100 epochs , well past the point where training loss stagnates . To understand the above phenomena , we perform another experiment . We train Cifar-10 on Resnet18 for 200 epochs , using a high LR of 0.1 for only 30 epochs and then use LR of 0.01 and 0.001 for 85 epochs each . We repeat this training 50 times with different random weight initializations . On an average , as expected , this training yields a low test accuracy of 94.81 . However , in 1 of the 50 runs , we find that the test accuracy reaches 95.24 , even higher than the average accuracy of 95.1 obtained while training at high LR for 100 epochs ! Hypothesis . To explain the above observations , i.e. , using a high learning rate for short duration results in low average test accuracy with rare occurrences of high test accuracy , while using the same high learning rate for long duration achieves high average test accuracy , we introduce a new hypothesis . We hypothesize that , in the DNN loss landscape , the density of narrow minima is significantly higher than that of wide minima . An intuitive explanation of why high LRs are necessary to locate wide minima then follows : a large LR 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 ) . For example , see Wu et al . ( 2018 ) for a result showing that large LRs are unstable at narrow minima and thus don ’ t converge to them . Thus the optimizer , when running at a high LR , jumps from one narrow minimum region to another , until it lands in a wide minimum region where it then gets stuck . Now , the probability of an optimization step landing in a wide minima is a direct function of the proportion of wide minima compared to that of narrow minima . Thus , if our hypothesis is true , i.e. , wide minima are much fewer than narrow minima , this probability is very low , and the optimizer needs to take a lot of steps to have a high probability of eventually landing in a wide minimum . This explains the observation in Table 1 , where the average accuracy continues to improve as we increase the number of high LR training steps . The hypothesis also explains why very few ( just 1 ) of the 50 runs trained at 0.1 LR for 30- epochs also managed to attain high accuracy – they just got lucky probabilistically and landed in a wide minimum even with a shorter duration . To validate this hypothesis further , we run experiments similar to the one in Table 1 . Specifically , we train Cifar-10 on Resnet-18 model for 200 epochs using a standard step schedule with LR of 0.1 , 0.01 , 0.001 . We vary the number of epochs trained using the high LR of 0.1 , called the explore epochs , from 30 to 100 epochs , and divide up the rest of the training equally between 0.01 and 0.001 . For each experimental setting , we conduct 50 random trials and plot the distributions of final test accuracy and the minima sharpness as defined by the metric in Keskar et al . ( 2016 ) . If our hypothesis is true , then the more you explore , the higher the probability of landing ( and getting stuck ) in a wide minima region , which should cause the distribution to tighten and move towards wider minima ( lower sharpness ) , as the number of explore steps increase . This is exactly what is observed in Figure 2 . Also since wide minima correlate with higher test accuracy , we should see the test accuracy distribution move towards higher accuracy and sharpen , as the number of explore steps increase . This is confirmed as well in Figure 3 . Finally , to verify whether explore at high LR is essential , we train Cifar-10 for 10,000 epochs at a fixed lower LR of 0.001 . The training converged but the final test accuracy was only 93.9 . Thus , even training 50x longer at low LR is not sufficient , adding more evidence to the hypothesis . Multi-scale . Given the importance of explore at high LR , a natural question that may arise is whether explore is necessary at smaller LR as well . To answer this , we train the same network for a total of 200 epochs with an initial high LR of 0.1 for 100 epochs , but now we vary the number of epochs trained with the LR of 0.01 ( we call this finer-scale explore ) , and train with LR of 0.001 for the remaining epochs . As can be seen from Table 2 , although the final training loss remains similar , we find that finer-scale explore also plays a role similar to the initial explore in determining the final test accuracy . This indicates that our hypothesis about density of wide/narrow regions indeed holds at multiple scales .

Summary: This work studies the problem of how to define learning rate schedules when training deep models so that the models better generalize.  To this end, the paper proposes and evaluates a learning rate schedule that consists of two stages (knee schedule).  A first stage of exploration adoptes a high learning rate. This initial stage is followed by a second stage where the learning rate decreases in a linear way. Extensive experimental results, both in text and image data, show that the proposed scheme allows one to train faster or to obtain better results with a fixed computational budget. The proposed learning schedule leads to SOTA results on IWSLT’14 (DE-EN) and WMT’14 (DE-EN) datasets.