Budgeted Training: Rethinking Deep Neural Network Training Under Resource Constraints

In most practical settings and theoretical analyses , one assumes that a model can be trained until convergence . However , the growing complexity of machine learning datasets and models may violate such assumptions . Indeed , current approaches for hyper-parameter tuning and neural architecture search tend to be limited by practical resource constraints . Therefore , we introduce a formal setting for studying training under the non-asymptotic , resource-constrained regime , i.e. , budgeted training . We analyze the following problem : “ given a dataset , algorithm , and fixed resource budget , what is the best achievable performance ? ” We focus on the number of optimization iterations as the representative resource . Under such a setting , we show that it is critical to adjust the learning rate schedule according to the given budget . Among budget-aware learning schedules , we find simple linear decay to be both robust and high-performing . We support our claim through extensive experiments with state-of-the-art models on ImageNet ( image classification ) , Kinetics ( video classification ) , MS COCO ( object detection and instance segmentation ) , and Cityscapes ( semantic segmentation ) . We also analyze our results and find that the key to a good schedule is budgeted convergence , a phenomenon whereby the gradient vanishes at the end of each allowed budget . We also revisit existing approaches for fast convergence and show that budgetaware learning schedules readily outperform such approaches under ( the practical but under-explored ) budgeted training setting . 1 INTRODUCTION . Deep neural networks have made an undeniable impact in advancing the state-of-the-art for many machine learning tasks . Improvements have been particularly transformative in computer vision ( Huang et al. , 2017b ; He et al. , 2017 ) . Much of these performance improvements were enabled by an ever-increasing amount of labeled visual data ( Russakovsky et al. , 2015 ; Kuznetsova et al. , 2018 ) and innovations in training architectures ( Krizhevsky et al. , 2012 ; He et al. , 2016 ) . However , as training datasets continue to grow in size , we argue that an additional limiting factor is that of resource constraints for training . Conservative prognostications of dataset sizes – particularly for practical endeavors such as self-driving cars ( Bojarski et al. , 2016 ) , assistive medical robots ( Taylor et al. , 2008 ) , and medical analysis ( Fatima & Pasha , 2017 ) – suggest one will train on datasets orders of magnitude larger than those that are publicly available today . Such planning efforts will become more and more crucial , because in the limit , it might not even be practical to visit every training example before running out of resources ( Bottou , 1998 ; Rai et al. , 2009 ) . We note that resource-constrained training already is implicitly widespread , as the vast majority of practitioners have access to limited compute . This is particularly true for those pursuing research directions that require a massive number of training runs , such as hyper-parameter tuning ( Li et al. , 2017 ) and neural architecture search ( Zoph & Le , 2017 ; Cao et al. , 2019 ; Liu et al. , 2019 ) . †Work done while at Argo AI . Instead of asking “ what is the best performance one can achieve given this data and algorithm ? ” , which has been the primary focus in the field so far , we decorate this question with budgeted training constraints as follows : “ what is the best performance one can achieve given this data and algorithm within the allowed budget ? ” . Here , the allowed budget refers to a limitation on the total time , compute , or cost spent on training . More specifically , we focus on limiting the number of iterations . This allows us to abstract out the specific constraint without loss of generality since any one of the aforementioned constraints could be converted to a finite iteration limit . We make the underlying assumption that the network architecture is constant throughout training , though it may be interesting to entertain changes in architecture during training ( Rusu et al. , 2016 ; Wang et al. , 2017 ) . Much of the theoretical analysis of optimization algorithms focuses on asymptotic convergence and optimality ( Robbins & Monro , 1951 ; Nemirovski et al. , 2009 ; Bottou et al. , 2018 ) , which implicitly makes use of an infinite compute budget . That said , there exists a wide body of work ( Zinkevich , 2003 ; Kingma & Ba , 2015 ; Reddi et al. , 2018 ; Luo et al. , 2019 ) that provide performance bounds which depend on the iteration number , which apply even in the non-asymptotic regime . Our work differs in its exploration of maximizing performance for a fixed number of iterations . Importantly , the globally optimal solution may not even be achievable in our budgeted setting . Given a limited budget , one obvious strategy might be data subsampling ( Bachem et al. , 2017 ; Sener & Savarese , 2018 ) . However , we discover that a much more effective , simpler , and under-explored strategy is adopting budget-aware learning rate schedules — if we know that we are limited to a single epoch , one should tune the learning schedule accordingly . Such budget-aware schedules have been proposed in previous work ( Feyzmahdavian et al. , 2016 ; Lian et al. , 2017 ) , but often for a fixed learning rate that depends on dataset statistics . In this paper , we specifically point out linearly decaying the learning rate to 0 at the end of the budget , may be more robust than more complicated strategies suggested in prior work . Though we are motivated by budget-aware training , we find that a linear schedule is quite competitive for general learning settings as well . We verify our findings with state-of-the-art models on ImageNet ( image classification ) , Kinetics ( video classification ) , MS COCO ( object detection and instance segmentation ) , and Cityscapes ( semantic segmentation ) . We conduct several diagnostic experiments that analyze learning rate decays under the budgeted setting . We first observe a statistical correlation between the learning rate and the full gradient magnitude ( over the entire dataset ) . Decreasing the learning rate empirically results in a decrease in the full gradient magnitude . Eventually , as the former goes to zero , the latter vanishes as well , suggesting that the optimization has reached a critical point , if not a local minimum1 . We call this phenomenon budgeted convergence and we find it generalizes across budgets . On one hand , it implies that one should decay the learning rate to zero at the end of the training , even given a small budget . On the other hand , it implies one should not aggressively decay the learning rate early in the optimization ( such as the case with an exponential schedule ) since this may slow down later progress . Finally , we show that linear budget-aware schedules outperform recently-proposed fast-converging methods that make use of adaptive learning rates and restarts . Our main contributions are as follows : 1Whether such a solution is exactly a local minimum or not is debatable ( see Sec 2 ) . • We introduce a formal setting for budgeted training based on training iterations and provide an alternative perspective for existing learning rate schedules . • We discover that budget-aware schedules are handy solutions to budgeted training . Specif- ically , our proposed linear schedule is more simple , robust , and effective than prior approaches , for both budgeted and general training . • We provide an empirical justification of the effectiveness of learning rate decay based on the correlation between the learning rate and the full gradient norm . 2 RELATED WORK . Learning rates . Stochastic gradient descent dates back to Robbins & Monro ( 1951 ) . The core is its update step : wt = wt−1 − αtgt , where t ( from 1 to T ) is the iteration , w are the parameters to be learned , g is the gradient estimator for the objective function2 F , and αt is the learning rate , also known as step size . Given base learning rate α0 , we can define the ratio βt = αt/α0 . Then the set of { βt } Tt=1 is called the learning rate schedule , which specifies how the learning rate should vary over the course of training . Our definition differs slighter from prior art as it separates the base learning rate and learning rate schedule . Learning rates are well studied for ( strongly ) convex cost surfaces and we include a brief review in Appendix H. Learning rate schedule for deep learning . In deep learning , there is no consensus on the exact role of the learning rate . Most theoretical analysis makes the assumption of a small and constant learning rate ( Du et al. , 2018a ; b ; Hardt et al. , 2016 ) . For variable rates , one hypothesis is that large rates help move the optimization over large energy barriers while small rates help converge to a local minimum ( Loshchilov & Hutter , 2017 ; Huang et al. , 2017a ; Kleinberg et al. , 2018 ) . Such hypothesis is questioned by recent analysis on mode connectivity , which has revealed that there does exist a descent path between solutions that were previously thought to be isolated local minima ( Garipov et al. , 2018 ; Dräxler et al. , 2018 ; Gotmare et al. , 2019 ) . Despite a lack of theoretical explanation , the community has adopted a variety of heuristic schedules for practical purposes , two of which are particularly common : • step decay : drop the learning rate by a multiplicative factor γ after every d epochs . The default for γ is 0.1 , but d varies significantly across tasks . • exponential : βt = γt . There is no default parameter for γ and it requires manual tuning . State-of-the-art codebases for standard vision benchmarks tend to employ step decay ( Xie & Tu , 2015 ; Huang et al. , 2017b ; He et al. , 2017 ; Carreira & Zisserman , 2017 ; Wang et al. , 2018 ; Yin et al. , 2019 ; Ma et al. , 2019 ) , whereas exponential decay has been successfully used to train Inception networks ( Szegedy et al. , 2015 ; 2016 ; 2017 ) . In spite of their prevalence , these heuristics have not been well studied . Recent work proposes several new schedules ( Loshchilov & Hutter , 2017 ; Smith , 2017 ; Hsueh et al. , 2019 ) , but much of this past work limits their evaluation to CIFAR and ImageNet . For example , SGDR ( Loshchilov & Hutter , 2017 ) advocates for learning-rate restarts based on the results on CIFAR , however , we find the unexplained form of cosine decay in SGDR is more effective than the restart technique . Notably , Mishkin et al . ( 2017 ) demonstrate the effectiveness of linear rate decay with CaffeNet on downsized ImageNet . In our work , we rigorously evaluate on 5 standard vision benchmarks with state-of-the-art networks and under various budgets . Gotmare et al . ( 2019 ) also analyze learning rate restarts and in addition , the warm-up technique , but do not analyze the specific form of learning rate decay . Adaptive learning rates . Adaptive learning rate methods ( Tieleman & Hinton , 2012 ; Kingma & Ba , 2015 ; Reddi et al. , 2018 ; Luo et al. , 2019 ) adjust the learning rate according to the local statistics of the cost surface . Despite having better theoretical bounds under certain conditions , they do not generalize as well as momentum SGD for benchmark tasks that are much larger than CIFAR ( Wilson et al. , 2017 ) . We offer new insights by evaluating them under the budgeted setting . We show fast descent can be trivially achieved through budget-aware schedules and aggressive early descent is not desirable for achieving good performance in the end . 2Note that g can be based on a single example , a mini-batch , the full training set , or the true data distribution . In most practical settings , momentum SGD is used , but we omit the momentum here for simplicity .

Summary: This work presents a simple technique for tuning the learning rate for Neural Network training when under a "budget" -- the budget here is specified as a fixed number of epochs that is expected to be a small fraction of the total number of epochs required to achieve maximum accuracy. The main contribution of this paper is in showing that a simpler linear decay schedule that goes to zero at the end of the proposed budget achieves good performance. The paper proposes a framework called budget-aware schedule which represents any learning rate schedule where the ratio of learning rate at time `'t' base learning rate is only a function of the ratio of 't' to total budget 'T'. In this family of schedules, the paper shows that a simple linear decay works best for all budgets. In the appendix, the authors compare their proposed schedule with adaptive techniques and show that under a given budget, it outperforms latets adaptive techniques like adabound, amsgrad, etc.