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: Learning rate schedule plays an important role in DL, which has a large influence over the final performance. Though there have been lots of schedules, achieving SOTA performance still requires careful hand-tuned schedule that may be case by case. Compared with previous learning rate schedules, authors first conjectured that the number of wide minima is significantly lower than the number of sharp minima, and then  proposed to use a large learning rate at the initialization phase for sufficient exploration to achieve a wide minima, which may achieve better generalization performance. Extensive experiments validate the proposed learning rate schedule.