A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes

1 INTRODUCTION . In recent years , hardware systems employing GPUs and TPUs have enabled neural network training programs to process dramatically more data in parallel than ever before . The most popular way to exploit these systems is to increase the batch size in the optimization algorithm ( i.e . the number of training examples processed per training step ) . On many workloads , modern systems can scale to larger batch sizes without significantly increasing the time per step ( Jouppi et al. , 2017 ; Wang et al. , 2019 ) , thus proportionally increasing the number of training examples processed per second . If researchers can use this increased throughput to reduce the time required to train each neural network , then they should achieve better results by training larger models , using larger datasets , and by exploring new ideas more rapidly . As the capacity for data parallelism continues to increase , practitioners can take their existing , well-tuned training configurations and re-train with larger batch sizes , hoping to achieve the same performance in less training time ( e.g . Ying et al. , 2018 ) . On an idealized data-parallel system with negligible overhead from increasing the batch size , they might hope to achieve perfect scaling , a proportional reduction in training time as the batch size increases . However , achieving perfect scaling is not always straightforward . Changing the batch size changes the training dynamics , requiring the training hyperparameters ( e.g . learning rate ) to be carefully re-tuned in order to maintain the same level of validation performance.1 In addition , smaller batch sizes provide implicit regularization from gradient noise ( Smith et al. , 2021 ) that may need to be replaced by other forms of regularization when the batch size is increased . Finally , even with perfect tuning , increasing the batch size eventually produces diminishing returns . After a critical batch size , the number of training steps can not be decreased in proportion to the batch size – the number of epochs must increase to match the validation performance of the smaller batch size . See Shallue et al . 2019 for a survey of the effects of data parallelism on neural network training . Once these effects are taken into account , there is no strong evidence that increasing the batch size degrades the maximum achievable performance on any workload . At the same time , the ever-increasing capacity for data parallelism presents opportunities for new regularization techniques that can replace the gradient 1 Although there are heuristics for adjusting the learning rate as the batch size changes , these heuristics inevitably break down sufficiently far from the initial batch size and it is also not clear how to apply them to other training hyperparameters ( e.g . momentum ) , as detailed in Shallue et al . ( 2019 ) . noise of smaller batch sizes and new optimization algorithms that can extend perfect scaling to larger batch sizes by using more sophisticated gradient information ( Zhang et al. , 2019 ) . You et al . ( 2017 ) proposed the LARS optimization algorithm in the hope of speeding up neural network training by exploiting larger batch sizes . LARS is a variant of stochastic gradient descent ( SGD ) with momentum ( Polyak , 1964 ) that applies layer-wise normalization before applying each gradient update . Although it is difficult to draw strong conclusions from the results presented in the LARS paper , 2 the MLPerf3 Training benchmark4 adopted LARS as one of two allowed algorithms in the closed division for ResNet-50 on ImageNet and it became the de facto standard algorithm for that benchmark task . With MLPerf entrants competing to find the fastest-training hyperparameters for LARS , the first place submissions in the two most recent MLPerf Training competitions used LARS to achieve record training speeds with batch sizes of 32,678 and 65,536 , respectively . No publications or competitive submissions to MLPerf have attempted to match these results with a standard optimizer ( e.g . Momentum or Adam ) . However , MLPerf entrants do not have a strong incentive ( nor are necessarily permitted by the rules ) to explore other algorithms because MLPerf Training is a systems benchmark that requires algorithmic equivalence between submissions to make fair comparisons . Moreover , since the main justification for LARS is its excellent performance on ResNet-50 at large batch sizes , more work is needed to quantify any benefit of LARS over standard algorithms at any batch size . You et al . ( 2019 ) later proposed the LAMB optimizer to speed up pre-training for BERT ( Devlin et al. , 2018 ) using larger batch sizes after concluding that LARS was not effective across workloads . LAMB is a variant of Adam ( Kingma & Ba , 2014 ) that adds a similar layer-wise normalization step to LARS . You et al . ( 2019 ) used LAMB for BERT pre-training with batch sizes up to 65,536 and claimed that Adam can not match the performance of LAMB beyond batch size 16,384 . In this paper , we demonstrate that standard optimizers , without any layer-wise normalization techniques , can match or improve upon the large batch size results used to justify LARS and LAMB . In Section 2 , we show that Nesterov momentum ( Nesterov , 1983 ) matches the performance of LARS on the ResNet-50 benchmark with batch size 32,768 . We are the first to match this result with a standard optimizer . In Section 3 , contradicting the claims in You et al . ( 2019 ) , we show that Adam obtains better BERT pre-training results than LAMB at the largest batch sizes , resulting in better downstream performance metrics after fine-tuning . In addition , we establish a new state-of-the-art for BERT pretraining speed , reaching an F1 score of 90.46 in 7,818 steps using Adam at batch size 65,536 ( we report training speed in steps because our focus is algorithmic efficiency , but since we compare LARS and LAMB to simpler optimizers , fewer training steps corresponds to faster wall-time in an optimized implementation – our BERT result with Adam also improves upon the wall-time record of LAMB reported in You et al . 2019 ) . Taken together , our results establish stronger training speed baselines for these tasks and batch sizes , which we hope will assist future work aiming to accelerate training using larger batch sizes . In addition to the contributions mentioned above , we demonstrate several key effects that are often overlooked by studies aiming to establish the superiority of new optimization algorithms . We show that future work must carefully disentangle regularization and optimization effects when comparing a new optimizer to baselines . We also report several under-documented details used to generate the best LARS and LAMB results , a reminder that future comparisons should document any novel tricks and include them in baselines . Finally , our results add to existing evidence in the literature on the difficulty of performing independently rigorous hyperparameter tuning for optimizers and baselines . In particular , we show that the optimal shape of the learning rate schedule is optimizer-dependent ( in addition to the scale ) , and that differences in the schedule can dominate optimizer comparisons at smaller step budgets and become less important at larger step budgets . 1.1 RELATED WORK . Shallue et al . ( 2019 ) and Zhang et al . ( 2019 ) explored the effects of data parallelism on neural network training for different optimizers , finding no evidence that larger batch sizes degrade performance and 2 The modified AlexNet on ImageNet benchmark did not have well-established accuracy targets from prior work and LARS used a more general learning rate schedule than the momentum baseline . For ResNet-50 on ImageNet , LARS achieved sub-par accuracy numbers and was not compared to any other optimizer at the same batch size , leaving open the possibility that a generic optimizer would scale just as well as LARS . 3 MLPerf is a trademark of MLCommons.org . 4 https : //mlperf.org/training-overview demonstrating that different optimizers can achieve perfect scaling up to different critical batch sizes . Geiping et al . ( 2021 ) found that full batch gradient descent can achieve competitive performance on CIFAR-10 . You et al . ( 2017 ; 2019 ) developed the LARS and LAMB optimizers in the hope of speeding up training by achieving perfect scaling beyond standard optimizers . Many other recent papers have proposed new optimization algorithms for generic batch sizes or larger batch sizes ( see Schmidt et al. , 2020 ) . Choi et al . ( 2019 ) , Sivaprasad et al . ( 2020 ) , and Schmidt et al . ( 2020 ) demonstrated the difficulties with fairly comparing optimizers , showing that the hyperparameter tuning protocol is a key determinant of optimizer rankings . The MLPerf Training benchmark ( Mattson et al. , 2019 ) provides a competitive ranking of neural network training systems , but does not shed much light on the relative performance of optimizers because entrants are limited in the algorithms they can use and the hyperparameters they can tune . 2 MATCHING LARS ON IMAGENET . The MLPerf training benchmark for ResNet-50 v1.5 on ImageNet ( Mattson et al. , 2019 ) aims to reach 75.9 % validation accuracy in the shortest possible wall-clock time . In the closed division of the competition , entrants must choose between two optimizers , SGD with momentum or LARS , and are only allowed to tune a specified subset of the optimization hyperparameters , with the remaining hyperparameter values set by the competition rules.5 The winning entries in the two most recent competitions used LARS with batch size 32,768 for 72 training epochs6 and LARS with batch size 65,536 for 88 training epochs,7 respectively . Kumar et al . ( 2019 ) later improved the training time for batch size 32,768 by reaching the target accuracy in 64 epochs . These are currently the fastest published results on the ResNet-50 benchmark . However , it has been unclear whether LARS was necessary to achieve these training speeds since no recent published results or competitive MLPerf submissions have used another optimizer . In this section , we describe how we matched the 64 epoch , 32,768 batch size result of LARS using standard Nesterov momentum.8 A fair benchmark of training algorithms or hardware systems must account for stochasticity in individual training runs . In the MLPerf competition , the benchmark metric is the mean wall-clock time of 5 trials after the fastest and slowest trials are excluded . Only 4 out of the 5 trials need to reach the target accuracy and there is no explicit limit on the number of times an entrant can try a different set of 5 trials . Since our goal is to compare algorithms , rather than systems , we aim to match the LARS result in terms of training steps instead ( but since Nesterov momentum is computationally simpler than LARS , this would also correspond to faster wall-clock time on an optimized system ) . Specifically , we measure the median validation accuracy over 50 training runs with a fixed budget of 2,512 training steps9 at a batch size of 32,768 . When we ran the published LARS training pipeline,10 LARS achieved a median accuracy of 75.97 % and reached the target in 35 out of 50 trials . We consider the LARS result to be matched by another optimizer if the median over 50 trials exceeds the target of 75.9 % . 2.1 NESTEROV MOMENTUM AT BATCH SIZE 32K . This section describes how we used the standard Nesterov momentum optimizer to train the ResNet50 v1.5 on ImageNet to 75.9 % validation accuracy in 2,512 update steps at a batch size of 32,768 , matching the best published LARS result at this batch size . Although we implemented our own training program , the only logical changes we made to the published LARS pipeline were to the optimizer and the optimization hyperparameters . Our model implementation and data pre-processing pipeline were identical to those required under the MLPerf closed division rules ( see Appendix B ) . We present two Nesterov momentum hyperparameter configurations that achieve comparable performance to LARS . Configuration A achieved a median accuracy of 75.97 % ( the same as LARS ) and reached the target accuracy in 34 out of 50 trials . Configuration B is a modified version of Configuration A designed to make as few changes as possible to the LARS hyperparameters ; it 5 https : //git.io/JtknD 6 https : //mlperf.org/training-results-0-6 7 https : //mlperf.org/training-results-0-7 8 The 88 epoch , 65,536 batch size result is faster in terms of wall-clock time but requires more training epochs , indicating that it is beyond LARS ’ s perfect scaling regime . Although LARS obtains diminishing returns when increasing the batch size from 32,768 to 65,536 , future work could investigate whether Nesterov momentum drops off more or less rapidly than LARS . 9 Corresponding to 64 training epochs in Kumar et al . ( 2019 ) . 10 https : //git.io/JtsLQ achieved a median accuracy of 75.92 % and reached the target in 29 out of 50 trials . See Appendix E.1 for the complete hyperparameter configurations . To achieve these results , we tuned the hyperparameters of the training pipeline from scratch using Nesterov momentum . We ran a series of experiments , each of which searched over a hand-designed hyperparameter search space using quasi-random search ( Bousquet et al. , 2017 ) . Between each experiment , we modified the previous search space and/or tweaked the training program to include optimization tricks and non-default hyperparameter values we discovered in the state-of-the-art LARS pipeline . The full sequence of experiments we ran , including the number of trials , hyperparameters tuned , and search space ranges , are provided in Appendix E.4 . Once we had matched the LARS result with Configuration A , we tried setting each hyperparameter to its value in the LARS pipeline in order to find the minimal set of changes that still achieved the target result , producing Configuration B . The remainder of this section describes the hyperparameters we tuned and the techniques we applied on the journey to these results .

Summary: The authors detail the significant effort required to reproduce results from the original LARS and LAMB papers.  Having reproduced results for LARS (for ResNet-50 training on ImageNet) and LAMB (for BERT pretraining), the authors then show that standard/older optimizers (Nesterov momentum and Adam) can match or even exceed LARS and LAMB at large batch sizes.  The authors provide in-depth details on both the difficulties required for them to reproduce the original results (i.e., fixing discrepancies between details in the publications and available source-code from online implementations) and their hyperparameter optimization for Nesterov Momentum and Adam on the datasets evaluated.  Also, insights and lessons are provided, both on the importance of hyperparameter tuning for comparison of DL solvers and on practitioners attempting to demonstrate performance gains for a new method (are not merely due to better tuning or regularization tricks).