Stochastic Weight Averaging in Parallel: Large-Batch Training That Generalizes Well

1 INTRODUCTION . Stochastic gradient descent ( SGD ) and its variants are the de-facto methods to train deep neural networks ( DNNs ) . Each iteration of SGD computes an estimate of the objective ’ s gradient by sampling a mini-batch of the available training data and computing the gradient of the loss restricted to the sampled data . A popular strategy to accelerate DNN training is to increase the mini-batch size together with the available computational resources . Larger mini-batches produce more precise gradient estimates ; these allow for higher learning rates and achieve larger reductions of the training loss per iteration . In a distributed setting , multiple nodes can compute gradient estimates simultaneously on disjoint subsets of the mini-batch and produce a consensus estimate by averaging all estimates , with one synchronization event per iteration . Training with larger mini-batches requires fewer updates , thus fewer synchronization events , yielding good overall scaling behavior . Even though the training loss can be reduced more efficiently , there is a maximum batch size after which the resulting model tends to have worse generalization performance ( McCandlish et al. , 2018 ; Keskar et al. , 2016 ; Hoffer et al. , 2017 ; Golmant et al. , 2018 ; Shallue et al. , 2018 ) . This phenomenon forces practitioners to use batch sizes below those that achieve the maximum throughput and limits the usefulness of large-batch training strategies . Stochastic Weight Averaging ( SWA ) ( Izmailov et al. , 2018 ) is a method that produces models with good generalization performance by averaging the weights of a set of models sampled from the final stages of a training run . As long as the models all lie in a region where the population loss is mostly convex , the average model can behave well , and in practice , it does . We have observed that if instead of sampling multiple models from a sequence generated by SGD , we generate multiple independent SGD sequences and average models from each , the resulting model achieves similar generalization performance . Furthermore , if all the independent sequences use small-batches , but start from a model trained with large-batches , the resulting model achieves generalization performance comparable with a model trained solely with small-batches . Using these observations , we derive Stochastic Weight Averaging in Parallel ( SWAP ) : A simple strategy to accelerate DNN training by better utilizing available compute resources . Our algorithm is simple to implement , fast and produces good results with minor tuning . For several image classification tasks on popular computer vision datasets ( CIFAR10 , CIFAR100 , and ImageNet ) , we show that SWAP achieves generalization performance comparable to models trained with small-batches but does so in time similar to that of a training run with large-batches . We use SWAP on some of the most efficient publicly available models to date , and show that it ’ s ∗Equal contribution †Work done during an internship at Apple Inc. able to substantially reduce their training times . Furthermore , we are able to beat the state of the art for CIFAR10 and train in 68 % of the time of the winning entry of the DAWNBench competition.1 2 RELATED WORK . The mechanism by which the training batch size affects the generalization performance is still unknown . A popular explanation is that because of the reduced noise , a model trained using larger mini-batches is more likely to get stuck in a sharper global minima . In ( Keskar et al. , 2016 ) , the authors argue that sharp minima are sensitive to variations in the data because slight shifts in the location of the minimizer will result in large increases in average loss value . However , if flatness is taken to be the curvature as measured by the second order approximation of the loss , then counterexamples exist . In ( Dinh et al. , 2017 ) , the authors transform a flat minimizer into a sharp one without changing the behavior of the model , and in ( Li et al. , 2018 ) , the authors show the reverse behavior when weight-decay is not used . In ( McCandlish et al. , 2018 ) , the authors predict that the batch size can be increased up to a critical size without any drop in accuracy and empirically validate this claim . For example , the accuracy begins to drop for image classification on CIFAR10 when the batch sizes exceed 1k samples . They postulate that when the batch size is large , the mini-batch gradient is close to the full gradient , and further increasing the batch size will not significantly improve the signal to noise ratio . In ( Hoffer et al. , 2017 ) , the authors argue that , for a fixed number of epochs , using a larger batch size implies fewer model updates . They argue that changing the number of updates impacts the distance the weights travel away from their initialization and that this distance determines the generalization performance . They show that by training with large-batches for longer times ( thus increasing the number of updates ) , the generalization performance of the model is recovered . Even though this large-batch strategy generates models that generalize well , it does so in more time than the smallbatch alternative . Irrespective of the generalization performance , the batch size also affects the optimization process . In ( Ma et al. , 2017 ) , the authors show that for convex functions in the over-parameterized setting , there is a critical batch size below which an iteration with a batch size of M is roughly equivalent to M iterations with a batch size of one , and batch-sizes larger than M do not improve the rate of convergence . Methods which use adaptive batch sizes exist ( Devarakonda et al. , 2017 ; Goyal et al. , 2017 ; Jia et al. , 2018 ; Smith et al. , 2017 ; You et al. , 2017 ) . However , most of these methods are either designed for specific datasets or require extensive hyper-parameter tuning . Furthermore , they ineffectively use the computational resources by reducing the batch size during part of the training . Local SGD ( Zhang et al. , 2016 ; Stich , 2018 ; Li et al. , 2019 ; Yu et al. , 2019 ) is a distributed optimization algorithm that trades off gradient precision with communication costs by allowing workers to independently update their models for a few steps before synchronizing . Post-local SGD ( Lin et al. , 2018 ) is a variant , which refines the output of large-batch training with local-SGD . The authors have observed that the resulting model has better generalization than the model trained with large-batches and that their scheme achieves significant speedups . In this manner Post-local SGD is of a very similar vein than the present work . However , while Post-local SGD lets the models diverge for T iterations where T is in the order of tens , SWAP averges the models once after multiple epochs . For example , in our Imagenet exeperiments ( see Sec . 5 ) we average our models after tens of thousands of updates , while Post-local SGD does after at most 32 . Because of this difference , we believe that the mechanisms that power the success of SWAP and Post-local SGD must be different and point to different phenomena in DNN optimization . Stochastic weight averaging ( SWA ) ( Izmailov et al. , 2018 ) is a method where models are sampled from the later stages of an SGD training run . When the weights of these models are averaged , they result in a model with much better generalization properties . This strategy is very effective and has been adopted in multiple domains : deep reinforcement learning ( Nikishin et al . ) , semisupervised learning ( Athiwaratkun et al. , 2019 ) , Bayesian inference ( Maddox et al. , 2019 ) , lowprecision training ( Yang et al. , 2019 ) . In this work , we adapt SWA to accelerate DNN training . 1The https : //dawn.cs.stanford.edu/benchmark/ 3 STOCHASTIC WEIGHT AVERAGING IN PARALLEL . We describe SWAP as an algorithm in three phases ( see Algorithm 1 ) : In the first phase , all workers train a single model by computing large mini-batch updates . Synchronization between workers is required at each iteration and a higher learning rate is used . In the second phase , each worker independently refines its copy of the model to produce a different set of weights . Workers use a smaller batch size , a lower learning rate , and different randomizations of the data . No synchronization between workers is required in this phase . The last phase consists of averaging the weights of the resulting models and computing new batch-normalization statistics to produce the final output . Phase 1 is terminated before the training loss reaches zero or the training accuracy reaches 100 % ( for example , a few percentage points below 100 % ) . We believe that stopping early precludes the optimization from getting stuck at a location where the gradients are too small and allows the following stage to improve the generalization performance . However , the optimal stopping accuracy is a hyper-parameter that requires tuning . During phase 2 , the batch size is appropriately reduced and small-batch training is performed independently and simultaneously . Here , each worker ( or a subset of them ) performs training using all the data , but sampling in different random order . Thus , after the end of the training process , each worker ( or subset ) will have produced a different model . Figure 1 plots the accuracies and learning-rate schedules for a run of SWAP . During the large-batch phase ( phase 1 ) , all workers share a common model and have the same generalization performance . During the small-batch phase ( phase 2 ) the learning rates for all the workers are the same but their testing accuracies differ as the stochasticity causes the models to diverge from each other . We also plot the test-accuracy of the averaged model that would result were we to stop phase 2 at that point . Note that the averaged model performs consistently better than each individual model . 4 LOSS LANDSCAPE VISUALIZATION AROUND SWAP ITERATES . To visualize the mechanism behind SWAP , we plot the error achieved by our test network on a plane that contains the outputs of the three different phases of the algorithm . Inspired by ( Garipov et al. , 2018 ) and ( Izmailov et al. , 2018 ) , we pick orthogonal vectors u , v that span the plane which contains θ1 , θ2 , θ3 . We plot the loss value generated by model θ = θ1+αu+βv at the location ( α , β ) . To plot a loss value , we first generate a weight vector θ , compute the batch-norm statistics for that model ( through one pass over the training data ) , and then evaluate the test and train accuracies . In Figure 2 , we plot the training and testing error for the CIFAR10 dataset . Here ‘ LB ’ marks the output of phase one , ‘ SGD ’ the output of a single worker after phase two , and ‘ SWAP ’ the final Algorithm 1 : Stochastic Weight Averaging in Parallel ( SWAP ) 1 Number of workers W ; Weight initialization θ0 ; t = 0 2 Training accuracy , τ , at which to exit phase one 3 Learning rate schedules LR1 and LR2 for phase one and two , respectively 4 Mini-batch sizes B1 and B2 for phase one and two , respectively 5 Gradient of loss function for sample i at weight θ : gi 6 SGDUpdate ( · ) : A function that updates the weights using SGD with momentum and weight decay 7 Phase 1 : 8 while Training accuracy ≤ τ do 9 η ← LR1 ( t ) 10 for w in [ 0 , ... , W − 1 ] In parallel do 11 Bw ← random sub-sample of training data with size B1W 12 gw ← W|B1| ∑ i∈Bw g i worker gradient 13 end 14 gt ← 1W ∑ gw synchronization of worker gradients 15 θt+1 ← θt + SGDUpdate ( ηt , gt , gt−1 , · · · ) ; / * first order method update * / 16 t = t+ 1 ; T = t 17 end 18 Phase 2 : 19 for t in [ T , T +Q ] do 20 η ← LR2 ( t− T ) 21 for w in [ 0 , ... , W − 1 ] In parallel do 22 Bw ← random sub-sample of training data with size B2 23 gw ← 1|B2| ∑ i∈Bw g i worker gradient 24 θwt+1 ← θwt + SGDUpdate ( ηt , gwt , gwt−1 , · · · ) ; / * first order method update at local worker * / 25 end 26 end / * We get W different models at the end of phase 2 * / 27 Phase 3 : θ̂ ` ← 1W ∑ θiT+Q produce averaged model 28 Compute batch-norm statistics for θ̂ ` to produce θ ` Result : Final model θ ` model . Color codes correspond to error measures at the points interpolated on the plane . In Figure 2a , we observe that the level-sets of the training error ( restricted to this plane ) form an almost convex basin and that both the output of phase 1 ( ‘ LB ’ ) 2 and the output of one of the workers of phase 2 ( ‘ SGD ’ ) lie in the outer edges of the basin . Importantly , during phase 2 the model traversed to a different side of the basin ( and not to the center ) . Also , the final model ( ‘ SWAP ’ ) is closer to the center of the basin . When we visualize these three points on the test loss landscape ( Figure 2b ) , we observe that the variations in the topology of the basin cause the ‘ LB ’ and ‘ SGD ’ points to fall in regions of higher error . But , since the ‘ SWAP ’ point is closer to the center of the basin , it is less affected by the change in topology . In Figure 3 , we neglect the ‘ LB ’ point and plot the plane spanned by three workers ‘ SGD1 ’ , ‘ SGD2 ’ , ‘ SGD3 ’ . In Figure 3a , we can observe that these points lie at different sides of the training error basin while ‘ SWAP ’ is closer to the center . In Figure 3b , we observe that the change in topology causes the worker points to lie in regions of higher testing errors than ‘ SWAP ’ , which is again close to the center of both basins . For reference , we have also plotted the best model that can be generated by this region of the plane . 2Recall that the weights ‘ LB ’ are obtained by stopping the large-batch training early in phase 1 . Hence , the training error for ‘ LB ’ is worse than ‘ SGD ’ and ‘ SWAP ’ .

Summary: This paper proposes a parallel version of the stochastic weight averaging method. It utilizes to phases to train the DNN. The first phase consists of distributed large-batch training, where the learning rate is scaled linearly with respect to the scale of the batch size. The second phase consists of using small batches with SWA to obtain the final model. Experiments verify that this method is able to achieve similar generalization performance as small-batch methods in less training time. A comparison against small-batch SWA, large-batch SWA, etc.