Stochastic Weight Averaging in Parallel: Large-Batch Training That Generalizes Well
Authors: Vipul Gupta
Santiago Akle Serrano
Dennis DeCoste
Venue:  (2020)
Keywords: Large batch training
Distributed neural network training
Stochastic Weight Averaging

Abstract: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the resulting models on the computer vision datasets CIFAR10, CIFAR100, and ImageNet.

Published as a conference paper at ICLR 2020
STOCHASTIC WEIGHT AVERAGING IN PARALLEL: LARGE-BATCH TRAINING THAT GENERALIZES WELL
Vipul Gupta*
Santiago Akle Serrano* sakle@ apple.com Apple Inc.
Dennis DeCoste ddecoste ddecoste@appler @apple.com Apple Inc.
vipul gupta @ berkeley.edu Deparment of EECS, UC Berkeley
ABSTRACT
We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to ac- celerate DNN training. Our algorithm uses large mini-batches to compute an ap- proximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The resulting models generalize equally well those trained with small mini-batches but produced in sub- stantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the resulting models on the computer vision datasets CIFARIO, CIFARIOO, and ImageNet.
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 's gradient by sam- pling 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 increase the mini- 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 distributed setting, multiple nodes can compute gradient estimates simul- taneously 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 mani-batches requires fewer updates, thus fewer synchronization events, yielding good overall scaling behavior.
Even though the training loss be reduced more efficiently, there is maximum batch size after which the resulting model tends to have worse generalization performance (McCandlish et al.] 2018) Keskar al.] 2016,] Hoffer et al. 2017] Golmant et al. 2018 (Shallue al.. 2018). This phenomenon forces practitioners to use batch sizes below those that achieve the maximum throughput and limits the usefulness of latch-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 set of models sampled from the final stages of a training run. As long as the models all lie in 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 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 model trained solely with observations, we derive Stochastic Weight Averaging in Parallel (SWAP): A simple strategy to ac- celerate DNN training by better utilizing available compute resources. Our algorithm is simple to implement, fast and produces good results with minor tuning.
Using these
For several image classification tasks popular computer vision datasets (CIFARIO, CIFARIOG 100, and ImageNet), we show that SWAP achieves generalization performance comparable to models trained with small-batches but does in time similar to that of training run with 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 Apple Inc.

---

Published as a conference paper at ICLR 2020
able to substantially reduce their training times. Furthermore, are able to beat the state of the art for CIFARIO and train in 68% of the time of the winning entry of the DAWNBench competition |
2 RELATED WORK
The mechanism by which the training batch size affects the generalization performance is still un- known. 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 counterex- amples exist. In (Dinh et al. 2..7). 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 critical size without any drop in accuracy and empirically validate this claim. For example, the accuracy begins to drop for image classification 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.
CIFARÍO when the batch sizes exceed Ik samples. They
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 aare-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 more time than the small- batch 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 M is roughly equivalent to M iterations with 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. [017]] 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 al. 2019) is distributed opti- mization 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) let al. 2018) is a variant, which refines the output of large-batch training with local-SGD. The au- thors have observed that the resulting model has better generalization than the model trained with barge-batches and that their scheme achieves significant speedups. In this manner Post-local SGD of a very similar vein than the present work. However, while Post-local SGD lets the models di- verge 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 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
al.). semi-
supervised learning (Athiwaratkun et al. 2019), Bayesian inference (Maddox et al.) 2019). low- precision training (Yang al. 2019). In this work, we adapt SWA to accelerate DNN training.
'The stanford http://ttannsssaaaaaat...luhhhht wtw.://ofood...u///eccmaattt
2

---

Published as a conference paper at ICLR 2020
3 STOCHASTIC WEIGHT AVERAGING IN PARALLEL
We describe SWAP as an algorithm in three phases (see Algorithm D: In the first phase, all work- ers train 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 different set of weights. Workers use smaller batch size, a lower learning rate, and different randomizations of the data. No synchroniza- tion 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 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 location where the gradients are too small and allows the fol- lowing 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 inde- pendently and simultaneously. Here, each worker (or 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.
plots the accuracies and learning-rate schedules for run of SWAP. During the large- batch phase (phase 1), all workers share 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 stop phase at that point. Note that the averaged model performs consistently better than each individual model.
Figure I
Worker
Worker
Leaning
Esocis
Large batch Mhsee
Smalt batch
Testing sccurscy
Worber
Worker
Average
Fpocte
Figure 1: Learning schedules and CIFARIO test accuracies for workers participating in SWAP. The large- batch phase with synchronized models is followed by the small-batch phase with diverging independent models. The test accuracy of the averaged weight model computed by averaging the independent models com- puting the test loss for the resulting model.
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 0,,02, 0g. We plot the loss value generated by model 0 = 01 ouu+ 3e at the location (a, B3. To plot a loss value, we first generate a weight vector 8, 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 CIFARIO dataset. Here 'LB' marks the output of phase one, 'SGD' the output of single worker after phase two, and 'SWAP" the final
3

---

Published as a conference paper at ICLR 2020
Algorithm 1: Stochastic Weight Averaging in Parallel (SWAP)
Number of workers W; Weight initialization lo: /
Training accuracy, 7, which exit phase
Learning rate schedules LR and LRR for phase one and two, respectively
Mini-batch sizes B1 and B3 for phase one and two, respectively
Gradient of loss function for sample at weight 0: g'
SGDUpdate(-) A function that updates the weights using SGD with momentum and weight decay
Phase 1:
while Training accuracy S r do
LR (t)
10
for w [0, W- 1] In parallel do
B + random sub-sample training data with size Bi W
11
g" W Eign= g worker gradient
12
B,
end
14
9t.
W Σ 9 synchronization of worker gradients
Oc+1 O, + SGDUpdate(ns It. It- /* first order methed update x/
16
t=t =t
end
18
Phase 2:
for I T, T Q) do
19
20
+ LR T)
21
for in [0, ..., W In parallel do
22
8" + random sub-sample of training data with size B2
9th B. E g6 worker gradient
23
24
0+11 c A. + SGDUpdateenee . 9"-1"""));
/* first order method
update at local worker #/
end
end
/ We get different models at the end of phase
Phase 3: 02 - , t produce averaged model
28 Compute batch-norm statistics for 0, to produce O, Result: Final model 0ę
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 ("LB') P and the output of of the workers of phase ('SGD') lie in the outer edges of the basin. Importantly, during phase 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 the loss landscape (Figure 2b observe that the variations in the topology of the basin cause the 'LB' and 'SGD' points to Tall in regions of higher error. But, since the 'SWAP' point is closer to the center of the basin, it 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 3aj we can observe that these points lie different sides of the training error basin while 'SWAP" closer to the center. In Figure 33b we observe that the change 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.
Recall that the weights 'IBB' are obtained by stopping the large-batch training early in phase 1. Hence, the training error for 'IB'' is worse than 'SGD' and 'SWAP.
4

---

Published as a conference paper at ICLR 2020
25
X62
23
12
26
47
54
12
1.1
11
501
(a) Train Error (%)
(h) Test Error (%)
Figure 2: CIFARIO train and error restricted a 2D plane spanned by the output of phase ('LB'), one the outputs phase 2 ("SGD') and the averaged model ('SWAP").
9002
15
12
72
15
68
SAMP
SNMP
50
12
10
13
25
(a) Train Error (%)
(b) Test Error (%)
Figure 3: CIFARIO train and test error restricted to 2D plane spanned by the output of three workers after phase 2 ('SGDI', 'SGD2', 'SGD3') and location of the average model ('SWAP'). The minimum test achievable for models restricted this region of the plane (marked BEST).
4.1 SAMPLING FROM INDEPENDENT RUNS OF SGD OR SAMPLING FROM ONE
In (Mandt al. 2017). the authors argue that in the later stages of SGD the weight iterates behave similar to an Ornstein Uhlenbeck process. So, by maintaining constant learning rate the SGD iterates should reach a stationary distribution that is similar distribution is centered at the local minimum, has a covariance that grows proportionally with the learning rate, inversely proportional the batch size and has shape that depends on both the Hessian of the mean loss and covariance of the gradient.
high-dimensional Gaussian. This
The authors (Izmailov al., 2018) argue that by virtue of being a high dimensional Gaussian all the mass of the distribution concentrated near the 'shell' of the ellipsoid, and therefore, it is unlikely for SGD to access the interior. They further argue that by sampling weights from an SGD run (leaving enough time steps between them) will choose weights that are spread out on the surface of this ellipsoid and their average will closer to the center.
can justify sampling from different SGD runs (as done in phase 2 during SWAP). As long as all runs start in the same basin of attraction, and provided the model from Mandt et al. 2017) holds, all runs will converge the same stationary distribution, and each run can generate independent samples from
Without any further assumptions,
4.2 ORTHIOGONALITY OF THE GRADIENT AND THE DIRECTION TO THE CENTER OF BASIN
To win intuition the advantage that SWA and SWAP have over SGD, measure the cosine similarity between the gradient descent direction, -t. and the direction towards the output -3) decreases as of SWAP, A Oswap 0, In Figure 4 we see that the cosine similarity, [9 the training enters its later stages. We believe that towards the end of training. the angle between the gradient direction and the directions toward the center of the basin is large, therefore the process moves mostly orthogonally to the basin, and progress slows. However, averaging samples from different sides of the basin can (and does) make faster progress towards the center.
5

---

Published as a conference paper at ICLR 2020
(65
[511
0.4-
Fpoch
Figure Cosine similarity between direction of gradient descent and A6
EXPERIMENTS
In this section we evaluate the performance of SWAP for image classification tasks on the CIFARIO, CIFARI00, and ImageNet datasets.
5.1 CIFARIO AND CIFARI00
For the experiments in this subsection, we found the best hyper-paramettrs using grid searches (see Appendix A for details), We train using mini-batch SGD with Nesterov momentum (set to 0.9) and weight decay of 5x 10-4. We augment the data using cutout (DeVries & Taylor 2017) and use fast- to-train custom ResNet from a submission 3 the DAWNBench leaderboard (Coleman al.). All experiments were run one machine with 8 NVIDIA Tesla V100 GPUs and use Horovod (Sergeev) Del Balso 2018) to distribute the computation. All statistics were collected over 10 different runs.
For these experiments, we used the following settings-SWAP phase one: 4096 samples per batch using GPUs (512 samples per GPU). Phase one terminated when the training accuracy reaches 98% (on average 108 epochs). SWAP phase two: 8 workers with one GPU each and 512 samples per batch for 30 epochs. The experiment that uses only bageebbtcches had 4096 samples per batch across 8 GPUs and is for 150 epochs. The experiments that use only small-batche had 512 samples per batch on 2 GPUs and is trained for 100 epochs.
CIFARIO:
Table compares the best test accuracies and corresponding training times for models trained with smalh-batch only, with large-batch only, and with SWAP. We report the average accuracy of the workers before averaging and the accuracy of the final model.
CIFARIO
Test Accuracy (%)
Training Time (sec)
SGD (stchl-bacch)
95.24 + 0.09
254.12 = 0.62
SGD (large-batch)
94.77 0.23
132.62 1.09
SWAP (before averaging)
94.70 + 0.20
167.57 3.25
SWAP (after averaging)
95.23 + 0.08
169.20 = 3.25
Table Training Statisties for CIFAR10
For these experiments, we use the following Settings-SWAL phase one: 2048 samples per batch using GPUs (256 samples per GPU). Phase one exits when the training accuracy reaches 90% (on average 112 epochs). SWAP phase two: 8 workers with one GPU each and 128 samples per batch, training for for 10 epochs. The experiments that use only large- batch training were run for 150 epochs with batches of 2048 8 GPUs The experiments that only small-batch were trained for 150 epochs using batches of 128 on GPU.
CIFARI00:
CIFAR100
Test Accuracy (%)
Training Time (sec)
SGD (stall-batch)
77.01 0.25
573.76 2.25
SGD (large-batch)
75.84 + 0.35
116.13 = 1.35
SWAP (before averaging)
75.74 0.15
123.11 = 1.85
SWAP (after averaging)
78.18 + 0.21
125.34 = 1.85
Table 2: Training Statistics for CIFAR100
Table 2 compares the best test accuracies and corresponding training times for models trained with only small- (for 150 epochs), with only large-batches (for 150 epochs), and with SWAP.

---

Published as a conference paper at ICLR 2020
For SWAP, report test accuracies obtained using the last SGD iterate before averaging, and test accuracy of the final model obtained after averaging. We observe significant improvement in test accuracies after averaging the models.
For both CIFAR 10 and COFARIOO, training with small-batches achieves higher testing accuracy than training with large-batches but takes much longer to train. SWAP, however, terminates in time comparable to the large-batch run but achieves accuracies on par (or better) than small batch training.
Achieving state of the art training speeds for CIFAR10: At the time of writing the front-runner of the DAWNBench competition takes 37 seconds with 4 Tesla V100 GPUs to train CIFARIO 10 to 94% test accuracy. Using SWAP with 8 Tesla V100 GPUs, phase batch size of 2048 samples and 28 epochs, and a phase two batch size of 256 samples for one epoch is able to reach the same accuracy in 27 seconds.
5.2 EXPERIMENTS ON IMAGENET
We use SWAP to accelerate a publicly available fast-to-train ImageNet model with published learn- ing rate and batch size schedules | The default settings for this code modify the learning-rates and batch sizes throughout the optimization (see 3.. Our small- batch experiments train ImageNet for 28 epochs using the published schedules with no modification and are run on 8 Tesla V100 GPUs. Our large-batch experiments modify the schedules by doubling the batch size and doubling the learning rates (see Figure (5) and are run on 16 Tesla V100 GPUs. For SWAP phase 1, we use the large-batch settings for 22 epochs, and for SWAP phase 2, we run two independent workers each with 8 GPUs using the settings for small-batche for epochs.
We observe that doubling the batch size reduces the Topl and Top5 test accuracies with respect to the small-batch run. SWAP, however, recovers the generalization performance at substantially reduced training times. Our results compiled in Table 3 (the statistics were collected over runs). We believe it's worthy of mention that these accelerations were achieved with no tuning other than increasing the learning rates proportionally to the increase in batch size and reverting to the original schedule when transitioning between phases.
Original
8000
Orignal
2x Larger nini- butch
2x Larger mini batch SWAP
Learnng nat wheeldee
7000
6000
Barchsite
2000
5600
1000
30
15
20
25
Fpoch
Epoch
(a) Learning rate schedule
(b) Batch sizes across epochs for ImageNet
Figure 5: Learning rate and mini-batch schedules used for ImageNet. The original schedule for 8 GPUs was double the batch size,
taken from existing DAWNBench submission. For larger batch experiment, double the number GPUs and double the learning rate of the original schedule, For SWAP, we switch from the modified schedule to the original schedule as we move from phase to phase 2.
ImageNet
Top1 Accuracy (%)
Top5 Accuracy (%)
Training Time (min)
SGD (small-batch)
76.14 + 0.07
93.30 + 0.07
235.29 + 0.33
SGD (large-batch)
75.86 + 0.03
92.98 0.06
127.20 0.78
SWAP (before averaging)
75.96 + 0.02
93.15 + 0.02
149.12 + 0.55
SWAP (after averaging)
76.19 + 0.03
93.32 + 0.02
156.55 0.56
Table 3: Training Statistics for ImageNet
7

---

Published as a conference paper at ICLR 2020
12
12-
Warier
Wwww
10-
1.0
0.6-
Laannng
0.-
0.2
0.2
02-
0.0-
0.8
50
102
125
175
Epodis
Lagg
Smead banch
(a) Large-batch SWA
(b) Large-batch training followed
(c) Small-batch SWA
by SWA with small-batches
Figure 6: Illustration of SWA with different batch sizes
5.3 EMPIRICAL COMPARISON OF SWA AND SWAP
We now compare SWAP with SWA: the sequential weight averaging algorithm from Izmailov al. (018). For the experiments in this section, we use the CIFAR100 dataset. We sample the same number of models for both SWA and SWAP and maintain the same number of epochs per sample. For SWA, we sample each model with 10 epochs in-between and average them to get the final model. For SWAP, we run 8 independent workers for 10 epochs each and use their average as the final model.
Large-batch SWA: We explore if SWA can recover the test accuracy of small-batch training on a large-batch training run. We use the same (large) batch size throughout. We follow an initial training cycle with cyclic learning rates (with cycles of 10 epochs) to sample 8 models (one from the end of each cycle). See Figure 6a for illustration of the learning rate schedule.
As expected we observe that the large-batch training run achieves lower training accuracy, but sur- prisingly SWA was unable to improve it (see Table 4 row 1).
Large-batch followed by small-batch SWA: We evaluate the effect of executing SWA using small- We interrupt the large-batch phase the same accuracy
batches after large-batch training interrupt phase of our CIFAR (0 experiment (Table 2}. In this case, the small-batch phase uses single worker and samples the models sequentially. SWA able reach the test accuracy of small-batch run but requires more than three times longer than SWAP to compute the model (see Table row 2). An illustration of the learning rate schedule provided in Figure 6b
We start the SWA cyclic learning rate schedule from the best model found by solely small-batch training (table [ row 1). Since the cycle length and cycle count are fixed, the only free parameter is the peak learning rate. We select this using a grid-search. Once the SWA schedule specified, we re-use the peak learning rate settings in SWAP. We start phase two from the model that was generated as the output of phase for the experiment on section 5.1 reported on table 2 rows and 4. With these settings, small-batch SWA achieves better accuracy than SWAP (by around 0.9%) at 6.8x more training time.
Small- SWA and SWAP:
Next, we wish to explore the speed-up that SWAP achieves over SWA if the precision of SWA is set target. To that end, relax the constraints SWAP. By increasing the phase two schedule from one 10 epoch cycle to two 20 epoch cycles and sampling two models from each worker (16 models) the resulting model achieved a test accuracy of 79. 11%1% in 241 seconds or 3.5x less time.
Test accuracy
Test accuracy after averaging (%)
Training Time (sec)
CIFARI00
before averaging (%)
Large-batch SWA Large-batch followed by small-batch SWA
76.06 + 0.25
76.00 0.31
376.4 2.25
76.26 + 0.35
78.12 + 0.14
398.0 + 1.35
Small- batch SWA
76.80 + 0.15
79.09 + 0.19
848.6 + 5.61
SWAP (10 small-batch epochs)
75.74 0.15
78.18 0.21
125.30 + 1.85
SWAP (40 small-batch epochs)
76.19
79. + 0.12
241.54 + 1.62
Table 4: Comparison: SWA versus SWAP
Availlabl http://ggthabcccmccsscommmmmmnnn
old
8

---

Published as a conference paper at ICLR 2020
6 CONCLUSIONS AND FUTURE WORK
We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm that uses a variant of Stochastic Weight Averaging (SWA) to improve the generalization performance of a model trained with large mini-batches. Our algorithm uses large mini-batches compute an approximate solution quickly and then refines it by averaging the weights of multiple models trained using small-batches The final model obtained after averaging has good generalization performance and is trained shorter time. We believe that this variant and this application of SWA are novel.
We observed that using large-batches in the initial stages of training does not preelude the models from achieving good generalization performance. That is, by refining the output of large-batch run, with models sampled sequentially as in SWA or in parallel as in SWAP, the resulting model is able perform well as the models trained using small-batches only. We confirm this in the image classification datasets CIFAR1O, CIFAR and ImageNet.
Through visualizations,
complement the existing evidence that averaged weights are closer to the center of a training loss basin than the models produced by stochastic gradient descent. It's interesting note that the basin into which the large mini-batch run is converging the same basin where the refined models are found. So, it is possible that regions with bad and good generalization performance are connected through regions of low training loss and, more so, that both belong to an almost convex basin. Our method requires the choice of (at least) one more hyperparameter: the transition point between the large-batc's batch and small-batch. For our experiments, we chose this by using grid search. A principled method to choose the transition point will be the focus future work.
seems to be
In future work we intend to explore the behavior of SWAP when used with other optimization schemes, such as Layer-wise Adaptive Rate Scaling (LARS) You al. 2017), mixed-precision training Jia et al. (2018), post-local SGD (Lin al. 2018) NovoGrad (Ginsburg al. 2019)) The design of SWAP allows us to substitute any of these for the large-batch stage, for example, we can use local SGD to accelerate the first stage of SWAP by reducing the communication overhead.

---

Published as a conference paper at ICLR 2020
REFERENCES
Ben Athiwaratkun, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson There are many consistent explanations of unlabeled data: Why you should average. In International Confer- ence Learning Representations, 2019. URL http:::/oppereviieeeee EkgKBhA5Y7
mettforummic
Cody Coleman, Deepak Narayanan, Daniel Kang, Tian Zhao, Jian Zhang, Luigi Nardi, Peter Bailís, Kunle Olukotun, Chris Ré, and Matei Zaharia. Dawnbench: An end-to-end deep learning bench- mark and competition.
Aditya Devarakonda, Maxim Naumov, and Michael Garland. Adabatch: Adaptive batch sizes for training deep neural networks. CoRR, abs/1712.02029, 2017. URL http:://axiv.orq/
abs/ 11t..1112.co 02029
Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint (oXiv:1708.045555 2017.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1019-1028. JMLR. org. 2017.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. In Advances in Neural Information Processing Systems, pp. 8789-8798, 2018
Boris Ginsburg, Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, and Jonathan M Cohen. Stochastic gradient methods with layer- wise adaptive moments for training of deep networks. arXiv preprint arXiv: 1286, 2019.
Noah Golmant, Nikita Vemuri, Zhewei Yao, Vladimir Feinberg, Amir Gholami, Kai Rothauge, Michael W Mahoney, and Joseph Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint /811 /8111244 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 706.02677, 2017.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. In NIPS, 2017.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407 2018.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang. Liwei Yu, al. Highly scalable deep learning training system with mixxed-recisioon Training imagenet in four minutes. arXiv preprint arXiv: 1205, 2018.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv::609; 04836, 2016.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. In Advances in Neural Information Processing Systems, pp. 6389-6399, 2018.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. arXiv preprint
078873 2019.
Tao Lin, Sebastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-bul batches, use local sgd. arXiv preprint 188iv 072 17, 2018.
10

---

Published as a conference paper at ICLR 2020
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of sgd in modern over-parametrized learning. arXiv preprint arXiv:1712.06559, 2017.
Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. arXiv preprint arXiv: (.02v:9902.0270000 2019.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research,
2017.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. An empirical model of large-batch training. arXiv preprint arXiv: //200118/..066/2, 2018.
Evgenii Nikishin, Pavel Izmailov, Ben Athiwaratkun, Dmitrii Podoprikhin, Timur Garipov, Pavel Shvechikov, Dmitry Vetrov, and Andrew Gordon Wilson. Improving stability in deep reinforce- ment learning with weight averaging.
Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning tensor- flow. arXiv preprint arXiv: (alivvv800.0579999 2018.
Christopher J Shallue, Jaehoon Lee, Joe Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv:1811.03660 2018.
Samuel L Smith, Picter-Jan Kindermans, Chris Ying, and Quoc V Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv:1711.0044 1.00899 2017.
Sebastian U Stich. Local sgd
converges fast and communicates little.
arXiv preprint
arXiv:1805.09767, 2018
Guandao Yang. Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. SWALP: Stochastic weight averaging in low precision training. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7015-7024, Long Beach, California, USA, 09-15 Jun 2019, PMLR.
Yang You, Igor Gitman, and Boris Ginsburg. Scaling sgd batch size to 32k for imagenet training. 2017.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Jian Zhang, Christopher De Sa, Ioannis Mitliagkas, and Christopher Ré. Parallel sgd: When does 2016
averaging help? ArXiv,

---

Published as a conference paper at ICLR 2020
A HYPERPARAMETERS FOR CIFARIO AND CIFARI00 EXPERIMENTS
We provide the parameters used the experiments of Section 5.1 These were obtained by doing independent grid searches for each experiment. For all CIFAR experiments, the momentum and weight decay constants were kept at 0.9 and 5 x 10 10-4 respectively. Tables 6 list the remaining hyperparameters. When a stopping accuracy of 100% listed, we mean that the maximum number of epochs were used.
SGD
SGD (large-batch)
SWAP (Phase 1)
SWAP
CIFARIO
(small-batch)
(Phase
Batch-size
512
4096
4096
512
Learning-rate Peak
0.3
1.2
1.2
0.12
Maximum Epochs
100
150
150
30
Warm-up Epochs
30
30
30
0
GPUs used per model
2
8
8
Stopping Accuracy (%)
100
100
98
100
Table Hyperparameters obtained using tuning for CIFAR10
SGD (small-batch)
SGD
SWAP (Phase 1)
SWAP
CIFAR100
(large-batch)
(Phase 2)
Batch-size
128
2048
2048
128
Learning-rate Peak
0.2
1.2
1.2
0.05
Total Epochs
150
150
150
30
Warm-up Epochs
60
45
45
0
GPUs used per model
1
8
8
Stopping Accuracy (%)
100
100
90
100
Table 6: Hyperparameters obtained using tuning for CIFAR100
12


Reviews:
[{"review_id": "Ske2EIQ0KH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In the paper, the authors propose a novel three-stage training strategy for deep learning models, training using large batch, training using small-batch locally and then aggregate models. Experimental results show that the proposed method converges faster than compared methods.  I have the following concerns:\n\n1) In figure 1, It looks like that the local models in worker 1, 2, 3 can reach the same testing accuracy without average. What is the meaning of computing average in this case?\n\n2) In the paper, authors mentioned that “Note that there exist training schemes in the literature that train on even larger batch sizes such as 32k (You et al., 2017; Jia et al., 2018), but these methods require a lot of hyperparameter tuning specific to the dataset.” As far as I know, they just need to tune warmup steps and peak learning rate, which is also required in the paper. In the proposed method, it is required to tune the switch point between phase 1 and phase 2.\n\n3) In the experiment, it also uses warmup for small batch size, is it necessary?\n\n4) Does the large batch training use layer-wise learning rate scaling? From my point of view, it is better to use it in the large-batch training. \n\n5) A guideline about how to select the switch point between phase 1 and phase 2 should be given if it takes time to tune it. \n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "S1eVmCKvKB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nStrengths:\n\nThe proposed algorithm is a natural extension to SWA, and appears to give good generalization performance. It is able to utilize large-batch training effectively, which is perhaps surprising given the amount of tuning necessary in Shallue, et al. (2018) in order to achieve good performance. The experiments are well-detailed, and some interesting visualizations, graphs, and empirical analyses are provided.\n\nWeaknesses:\n\nI think that this paper is fairly complete; the only question is whether or not it contains enough novelty, as it is a natural extension of SWA to the parallelized setting. No theoretical analysis of the algorithm is given.\n\nSome questions and comments:\n\n- How sensitive is the algorithm to the choice of the transition point? How was the transition point tuned?\n- How large of a batch size can one use before this approach breaks down or is no longer efficient?\n- In Figure 2, LB is shown to obtain worse training error than SGD. What is the reason for this? This seems contrary to theory (assuming one has already converged to a neighborhood of the solution).\n- The authors comment in the conclusion that \"it is possible that regions with bad and good generalization performance are connected through regions of low training loss\". Can one check if this is related to the invariances described in Dinh, et al. (2017)? \n- What is the relationship between SWAP and methods on local SGD?\n\nThis paper is missing some classical optimization references on increasing batch sizes, which has been well-studied within that literature:\n\n[1] Byrd, Richard H., et al. \"Sample size selection in optimization methods for machine learning.\" Mathematical programming 134.1 (2012): 127-155.\n[2] Bollapragada, Raghu, Richard Byrd, and Jorge Nocedal. \"Adaptive sampling strategies for stochastic optimization.\" SIAM Journal on Optimization 28.4 (2018): 3312-3343.\n[3] Bollapragada, Raghu, et al. \"A progressive batching L-BFGS method for machine learning.\" arXiv preprint arXiv:1802.05374 (2018). \n[4] Friedlander, Michael P., and Mark Schmidt. \"Hybrid deterministic-stochastic methods for data fitting.\" SIAM Journal on Scientific Computing 34.3 (2012): A1380-A1405.\n\nAlthough developing some theory for this algorithm would be beneficial, this paper performs a comprehensive set of experiments and is well-written. For these reasons, I'm inclined to accept the paper.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "HJxGKR8pYH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a 2-stage SGD variant that improves the generalization. The experiments show good performance.\nHowever, there are some weakness in this paper:\n\n1. (Minor issue) The Update() function in Algorithm 1 seems to be something very general. However, it seems that Update() is simply SGD or SGD with (Nesterov) momentum, according to Section 5.1. Furthermore, the authors never explicitly explain what exact Update() function is used, which is very unfriendly to the readers.\n\n2. The major issue is that the proposed algorithm and the contribution (improvement of generalization) is not novel. Phase 2 of Algorithm 2 is called local SGD, proposed in [1,2]. Local SGD also has variants with Polyak momentum and Nesterov momentum [3]. Furthermore, [4] has already proposed an algorithm, post-local SGD, which is basically the same as Algorithm 1 in this paper (run fully synchronous SGD first and then local SGD). Note that [4] also shows that post-local SGD converges to flatter minima, and results in better generalization. Please correct me if I'm wrong, and explain the difference between Algorithm 1 and (post-)local SGD in details.\n\n\n----------\nReference\n\n[1] Stich, Sebastian U.. “Local SGD Converges Fast and Communicates Little.” ArXiv abs/1805.09767 (2018).\n[2] Yu, Hao et al. “Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning.” AAAI (2018).\n[3] Yu, Hao et al. “On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization.” ICML (2019).\n[4] Lin, Tao et al. “Don't Use Large Mini-Batches, Use Local SGD.” ArXiv abs/1808.07217 (2018).\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors proposed a simple and effective approach to parallel training based on stochastic weight averaging. Moreover, the authors have carefully addressed the reviewer comments in the discussion period, particularly the relation to local SGD, to the satisfaction of reviewers. Local SGD mimics sequential SGD with noise induced by lack of synchronization, whereas SWAP averages multiple samples from a stationary distribution, and synchronizes at the end. Please clarify these points and carefully account for reviewer comments in the final version. Overall, the proposed approach will make an excellent addition to the program, both elegant and practically useful.
Decision: Accept (Poster)