Don't Use Large Mini-batches, Use Local SGD

1 INTRODUCTION . Fast and efficient training of large scale deep-learning models relies on distributed hardware and on distributed optimization algorithms . For efficient use of system resources , these algorithms crucially must ( i ) enable parallelization while being communication efficient , and ( ii ) exhibit good generalization behaviour , i.e . good performance on unseen data ( test-set ) . Most machine learning applications currently depend on stochastic gradient descent ( SGD ) ( Robbins & Monro , 1985 ) and in particular its mini-batch variant ( Bottou , 2010 ; Dekel et al. , 2012 ) . However , this algorithm faces generalization difficulties in the regime of very large batch sizes as we will review now . Mini-batch SGD . For a sum-structured optimization problem of the form minwPRd 1N °N i 1 fipwq where w P Rd denotes the parameters of the model and fi : Rd Ñ R the loss function of the i-th training example , the mini-batch SGD update for K ¥ 1 workers is given as wpt 1q : wptq γptq 1 K °K k 1 1 B ° iPIk ptq ∇fi wptq , ( 1 ) where γptq ¡ 0 denotes the learning rate and Ikptq rN s the subset ( mini-batch ) of training datapoints selected by worker k ( typically selected uniformly at random from the locally available datapoints on worker k ) . For convenience , we will assume the same batch size B per worker . Local SGD . Motivated to better balance the available system resources ( computation vs. communication ) , local SGD ( a.k.a . local-update SGD , parallel SGD , or federated averaging ) has recently attracted increased research interest ( Mcdonald et al. , 2009 ; Zinkevich et al. , 2010 ; McDonald et al. , 2010 ; Zhang et al. , 2014 ; 2016 ; McMahan et al. , 2017 ) . In local SGD , each worker k P rKs evolves a local model by performing H sequential SGD updates with mini-batch size Bloc , before communication ( synchronization by averaging ) among the workers . Formally , wkptq h 1 : w k ptq h γptq 1 Bloc ° iPIk ptq h ∇fi wkptq h , wkpt 1q : 1 K °K k 1w k ptq H , ( 2 ) where wkptq h denotes the local model on machine k after t global synchronization rounds and subsequent h P rHs local steps ( Ikptq h is defined analogously ) . Mini-batch SGD is a special case if H 1 and Bloc B . Furthermore , the communication patterns of both algorithms are identical if B HBloc . However , as the updates ( eq . ( 2 ) ) are different from the mini-batch updates ( eq . ( 1 ) ) for any H ¡ 1 , the generalization behavior ( test error ) of both algorithms is expected to be different . Large batch SGD . Recent schemes for scaling training to a large number of workers rely on standard mini-batch SGD ( 1 ) with very large overall batch sizes ( Shallue et al. , 2018 ; You et al. , 2018 ; Goyal et al. , 2017 ) , i.e . increasing the global batch size linearly with the number of workers K. However , the batch size interacts differently with the overall system efficiency on one hand , and with generalization performance on the other hand . While a larger batch size in general increases throughput , it may negatively affect the final accuracy on both the train- and test-set . Two main scenarios of particular interest can be decoupled as follows : Scenario 1 . The communication restricted setting , where the synchronization time is much higher than the gradient computation time . In this case the batch size of best efficiency is typically large , and is achieved by using partial computation ( gradient accumulation ) while waiting for communication . We in particular study the interesting case when the mini-batch sizes of both algorithms satisfy the relation B HBloc , as in this case both algorithms ( local SGD and minibatch SGD ) evaluate the same number of stochastic gradients between synchronization steps . Scenario 2 . The regime of poor generalization of large-batch SGD , that is the use of very large overall batches ( often a significant fraction of the training set size ) , which is known to cause drastically decreased generalization performance ( Chen & Huo , 2016 ; Keskar et al. , 2017 ; Hoffer et al. , 2017 ; Shallue et al. , 2018 ; Golmant et al. , 2018 ) . If sticking to standard mini-batch SGD and maintaining the level of parallelisation , the batch size B would have to be reduced below the device ( locally optimal ) capacity in order to alleviate this generalization issue , which however impacts training time1 . Main Results . Key aspects of the empirical performance of local SGD compared to mini-batch baselines are illustrated in Figure 1 . In scenario 1 ) , comparing local SGD with H 4 ( A4 ) with mini-batch SGD of same effective batch size B 4Bloc ( A3 ) reveals a stark difference , both in terms of train and test error ( local SGD achieves lower training loss and higher test accuracy ) . This motivates the use of local SGD as an alternative to large-batch training—a hypothesis that we confirm in our experiments . Further , in scenario 2 ) , mini-batch SGD with smaller batch size B Bloc ( A2 ) is observed to suffer from poor generalization , although the training curve matches the single- 1Note that in terms of efficiency on current GPUs , the computation time on device for small batch sizes is not constant but scales non-linearly with B , as shown in Table 7 in Appendix A. machine baseline ( A1 ) . The generalization gap can thus not be explained as an optimization issue alone . Our proposed post-local SGD ( A5 ) ( defined by starting local SGD from the model obtained by large-batch SGD ( A2 ) at epoch 150 ) closes this generalization gap with the single-machine baseline ( A1 ) and is also more communication efficient than the mini-batch competitors . In direct comparison , post-local SGD is more communication-efficient than mini-batch SGD ( while less than local SGD ) . It achieves better generalization performance than both these algorithms . Contributions . Our main contributions can thus be summarized as follows : • Trade-offs in Local SGD : We provide the first comprehensive empirically study of the trade-offs in local SGD for deep learning—when varying the number of workers K , number of local steps H and mini-batch sizes—for both scenarios 1 ) on communication efficiency and 2 ) on generalization . • Post-local SGD : We propose post-local SGD , a simple but very efficient training scheme to address the current generalization issue of large-batch training . It allows us to scale the training to much higher number of parallel devices . Large batches trained by post-local SGD enjoy improved communication efficiency , while at the same time strongly outperforming most competing small and large batch baselines in terms of accuracy . Our empirical experiments on standard benchmarks show that post-local SGD can reach flatter minima than large-batch SGD on those problems . 2 RELATED WORK . The generalization gap in large-batch training . State-of-the-art distributed deep learning frameworks ( Abadi et al. , 2016 ; Paszke et al. , 2017 ; Seide & Agarwal , 2016 ) resort to synchronized large-batch SGD training , allowing scaling by adding more computational units and performing data-parallel synchronous SGD with mini-batches divided between devices . Training with large batch size ( e.g . batch size > 103 on ImageNet ) typically degrades the performance both in terms of training and test error ( often denoted as generalization gap ) ( Chen & Huo , 2016 ; Li , 2017 ; Li et al. , 2014 ; Keskar et al. , 2017 ; Shallue et al. , 2018 ; McCandlish et al. , 2018 ; Golmant et al. , 2018 ; Masters & Luschi , 2018 ) . Goyal et al . ( 2017 ) argue that the test error degrades because of optimization issues and propose to use a “ learning rate warm-up ” phase with linear scaling of the step-size . You et al . ( 2017a ) propose Layer-wise Adaptive Rate Scaling ( LARS ) to scale to larger mini-batch size , but the generalization gap does not vanish . Hoffer et al . ( 2017 ) argue that the generalization gap can be closed when increasing the number of iterations along with the batch size . However , this diminishes the efficiency gains of parallel training . Keskar et al . ( 2017 ) empirically show that larger batch sizes correlate with sharper minima ( Hochreiter & Schmidhuber , 1997 ) found by SGD and that flat minima are preferred for better generalization . This interpretation—despite being debated in Dinh et al . ( 2017 ) —was further developed in Hoffer et al . ( 2017 ) ; Yao et al . ( 2018 ) ; Izmailov et al . ( 2018 ) . Neelakantan et al . ( 2015 ) propose to add isotropic white noise to the gradients to avoid over-fitting and better optimization . Zhu et al . ( 2019 ) ; Xing et al . ( 2018 ) further analyze “ structured ” anisotropic noise and highlight the importance of anisotropic noise ( over isotropic noise ) for improving generalization . The importance of the scale of the noise for non-convex optimization has also been studied in Smith & Le ( 2018 ) ; Chaudhari & Soatto ( 2018 ) . Wen et al . ( 2019 ) propose to inject noise ( sampled from the expensive empirical Fisher matrix ) to large-batch SGD . However , to our best knowledge , none of the prior work ( except our post-local SGD ) can provide a computation efficient way to inject noise to achieve as good generalization performance as small-batch SGD , for both of CIFAR and ImageNet experiments . Local SGD and convergence theory . While mini-batch SGD is very well studied ( Zinkevich et al. , 2010 ; Dekel et al. , 2012 ; Takáč et al. , 2013 ) , the theoretical foundations of local SGD variants are still developing . Jain et al . ( 2018 ) study one-shot averaging on quadratic functions and Bijral et al . ( 2016 ) study local SGD in the setting of a general graph of workers . A main research question is whether local SGD provides a linear speedup with respect to the number of workers K , similar to mini-batch SGD . Recent work partially confirms this , under the assumption that H is not too large compared to the total iterations T . Stich ( 2019 ) and Patel & Dieuleveut ( 2019 ) show convergence at rate O pKTHBlocq 1 on strongly convex and smooth objective functions when H OpT 1 { 2q . For smooth non-convex objective functions , Zhou & Cong ( 2018 ) show a rate of O pKTBlocq 1 { 2 ( for the decrement of the stochastic gradient ) , Yu et al . ( 2019 ) give an improved result O pHKTBlocq 1 { 2 when H OpT 1 { 4q . For a discussion of more recent theoretical results we refer to ( Kairouz et al. , 2019 ) . Alistarh et al . ( 2018 ) study convergence under adversarial delays . Zhang et al . ( 2016 ) empirically study the effect of the averaging frequency on the quality of the solution for some problem cases and observe that more frequent averaging at the beginning of the optimization can help . Similarly , Bijral et al . ( 2016 ) argue to average more frequently at the beginning .

Summary: This paper proposes a variant of local SGD, post-local SGD, for distributed training of deep neural networks. It targets to mitigate the generalization gap caused by large batch training. The idea is straightforward and easy to understand-- start the training with standard mini-batch SGD and later switch to local SGD. The rationale behind this scheme is that switching to local SGD helps the training converge to flatter minima compared to using large-batch SGD, which correlates with sharper minima, and that helps close the generation gap. Switching to local SGD at the second phase also helps improve communication efficiency by reducing the amortized communication volume. The authors perform empirical studies using ResNet and DenseNet to conclude that post-local SGD outperforms large-batch SGD in terms of generalization performance while also with improved communication efficiency.