Original: S TOCHASTIC W EIGHT A VERAGING IN P ARALLEL : L ARGE -B ATCH T RAINING T HAT G ENERALIZES W ELL
Revised: S TOCHASTIC W EIGHT A VERAGING IN P ARALLEL : L ARGE -B ATCH T RAINING T HAT G ENERALIZES W ELL

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training.
Revised: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training.

Original: Our algorithm uses large mini-batches to compute an approximate solution quickly and then reﬁnes it by averaging the weights of multiple models computed independently and in parallel.
Revised: Our algorithm uses large mini-batches to compute an approximate solution quickly and then reﬁnes it by averaging the weights of multiple models computed independently and in parallel.

Original: The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time.
Revised: The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Stochastic gradient descent (SGD) and its variants are the de-facto methods to train deep neural networks (DNNs).
Revised: Stochastic gradient descent (SGD) and its variants are the de-facto methods to train deep neural networks (DNNs).

Original: 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.
Revised: 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.

Original: A popular strategy to accelerate DNN training is to increase the mini-batch size together with the available computational resources.
Revised: A popular strategy to accelerate DNN training is to increase the mini-batch size together with the available computational resources.

Original: Larger mini-batches produce more precise gradient estimates; these allow for higher learning rates and achieve larger reductions of the training loss per iteration.
Revised: Larger mini-batches produce more precise gradient estimates; these allow for higher learning rates and achieve larger reductions of the training loss per iteration.

Original: 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.
Revised: 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.

Original: Training with larger mini-batches requires fewer updates, thus fewer synchronization events, yielding good overall scaling behavior.
Revised: Training with larger mini-batches requires fewer updates, thus fewer synchronization events, yielding good overall scaling behavior.

Original: Even though the training loss can be reduced more efﬁciently, there is a maximum batch size after which the resulting model tends to have worse generalization performance (McCandlish et al., 2018;
Revised: Even though the training loss can be reduced more efﬁciently, there is a maximum batch size after which the resulting model tends to have worse generalization performance (McCandlish et al., 2018;

Original: Keskar et al., 2016; Hoffer et al., 2017; Golmant et al., 2018; Shallue et al., 2018).
Revised: Keskar et al., 2016; Hoffer et al., 2017; Golmant et al., 2018; Shallue et al., 2018).

Original: This phenomenon forces practitioners to use batch sizes below those that achieve the maximum throughput and limits the usefulness of large-batch training strategies.
Revised: This phenomenon forces practitioners to use batch sizes below those that achieve the maximum throughput and limits the usefulness of large-batch training strategies.

Original: 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 ﬁnal stages of a training run.
Revised: 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 ﬁnal stages of a training run.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Using these observations, we derive Stochastic Weight Averaging in Parallel (SWAP):
Revised: Using these observations, we derive Stochastic Weight Averaging in Parallel (SWAP):

Original: A simple strategy to accelerate DNN training by better utilizing available compute resources.
Revised: A simple strategy to accelerate DNN training by better utilizing available compute resources.

Original: Our algorithm is simple to implement, fast and produces good results with minor tuning.
Revised: Our algorithm is simple to implement, fast and produces good results with minor tuning.

Original: For several image classiﬁcation 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.
Revised: For several image classiﬁcation 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.

Original: We use SWAP on some of the most efﬁcient publicly available models to date, and show that it’s able to substantially reduce their training times.
Revised: We use SWAP on some of the most efﬁcient publicly available models to date, and show that it’s able to substantially reduce their training times.

Original: 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.
Revised: 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.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: The mechanism by which the training batch size affects the generalization performance is still unknown.
Revised: The mechanism by which the training batch size affects the generalization performance is still unknown.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, if ﬂatness is taken to be the curvature as measured by the second order approximation of the loss, then counterexamples exist.
Revised: However, if ﬂatness is taken to be the curvature as measured by the second order approximation of the loss, then counterexamples exist.

Original: In (Dinh et al., 2017), the authors transform a ﬂat 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.
Revised: In (Dinh et al., 2017), the authors transform a ﬂat 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.

Original: 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.
Revised: 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.

Original: For example, the accuracy begins to drop for image classiﬁcation on CIFAR10 when the batch sizes exceed 1k samples.
Revised: For example, the accuracy begins to drop for image classiﬁcation on CIFAR10 when the batch sizes exceed 1k samples.

Original: 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 signiﬁcantly improve the signal to noise ratio.
Revised: 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 signiﬁcantly improve the signal to noise ratio.

Original: In (Hoffer et al., 2017), the authors argue that, for a ﬁxed number of epochs, using a larger batch size implies fewer model updates.
Revised: In (Hoffer et al., 2017), the authors argue that, for a ﬁxed number of epochs, using a larger batch size implies fewer model updates.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Even though this large-batch strategy generates models that generalize well, it does so in more time than the smallbatch alternative.
Revised: Even though this large-batch strategy generates models that generalize well, it does so in more time than the smallbatch alternative.

Original: Irrespective of the generalization performance, the batch size also affects the optimization process.
Revised: Irrespective of the generalization performance, the batch size also affects the optimization process.

Original: 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.
Revised: 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.

Original: Methods which use adaptive batch sizes exist (Devarakonda et al., 2017; Goyal et al., 2017; Jia et al., 2018;
Revised: Methods which use adaptive batch sizes exist (Devarakonda et al., 2017; Goyal et al., 2017; Jia et al., 2018;

Original: Smith et al., 2017; You et al., 2017).
Revised: Smith et al., 2017; You et al., 2017).

Original: However, most of these methods are either designed for speciﬁc datasets or require extensive hyper-parameter tuning.
Revised: However, most of these methods are either designed for speciﬁc datasets or require extensive hyper-parameter tuning.

Original: Furthermore, they ineffectively use the computational resources by reducing the batch size during part of the training.
Revised: Furthermore, they ineffectively use the computational resources by reducing the batch size during part of the training.

Revised: 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.

Revised: Post-local SGD (Lin et al., 2018) is a variant, which reﬁnes the output of large-batch training with local-SGD.

Revised: The authors have observed that the resulting model has better generalization than the model trained with large-batches and that their scheme achieves signiﬁcant speedups.

Revised: In this manner Post-local SGD is of a very similar vein than the present work.

Original:  Stochastic weight averaging (SWA) (Izmailov et al., 2018) is a method where models are sampled from the later stages of an SGD training run.
Revised: 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.

Original: When the weights of these models are averaged, they result in a model with much better generalization properties.
Revised: When the weights of these models are averaged, they result in a model with much better generalization properties.

Original: 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).
Revised: 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).

Original: In this work, we adapt SWA to accelerate DNN training.
Revised: In this work, we adapt SWA to accelerate DNN training.

Original:  S TOCHASTIC WEIGHT AVERAGING IN PARALLEL
Revised:  S TOCHASTIC WEIGHT AVERAGING IN PARALLEL

Original: We describe SWAP as an algorithm in three phases (see Algorithm 1):
Revised: We describe SWAP as an algorithm in three phases (see Algorithm 1):

Original: In the ﬁrst phase, all workers train a single model by computing large mini-batch updates.
Revised: In the ﬁrst phase, all workers train a single model by computing large mini-batch updates.

Original: Synchronization between workers is required at each iteration and a higher learning rate is used.
Revised: Synchronization between workers is required at each iteration and a higher learning rate is used.

Original: In the second phase, each worker independently reﬁnes its copy of the model to produce a different set of weights.
Revised: In the second phase, each worker independently reﬁnes its copy of the model to produce a different set of weights.

Original: Workers use a smaller batch size, a lower learning rate, and different randomizations of the data.
Revised: Workers use a smaller batch size, a lower learning rate, and different randomizations of the data.

Original: No synchronization between workers is required in this phase.
Revised: No synchronization between workers is required in this phase.

Original: The last phase consists of averaging the weights of the resulting models and computing new batch-normalization statistics to produce the ﬁnal output.
Revised: The last phase consists of averaging the weights of the resulting models and computing new batch-normalization statistics to produce the ﬁnal output.

Original: Phase 1 is terminated before the training loss reaches zero or the training accuracy reaches 100% (for example, a few percentage points below 100% ).
Revised: Phase 1 is terminated before the training loss reaches zero or the training accuracy reaches 100% (for example, a few percentage points below 100% ).

Original: 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.
Revised: 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.

Original: However, the optimal stopping accuracy is a hyper-parameter that requires tuning.
Revised: However, the optimal stopping accuracy is a hyper-parameter that requires tuning.

Original: During phase 2, the batch size is appropriately reduced and small-batch training is performed independently and simultaneously.
Revised: During phase 2, the batch size is appropriately reduced and small-batch training is performed independently and simultaneously.

Original: Here, each worker (or a subset of them) performs training using all the data, but sampling in different random order.
Revised: Here, each worker (or a subset of them) performs training using all the data, but sampling in different random order.

Original: Thus, after the end of the training process, each worker (or subset) will have produced a different model.
Revised: Thus, after the end of the training process, each worker (or subset) will have produced a different model.

Original: Figure 1 plots the accuracies and learning-rate schedules for a run of SWAP.
Revised: Figure 1 plots the accuracies and learning-rate schedules for a run of SWAP.

Original: During the large-batch phase (phase 1), all workers share a common model and have the same generalization performance.
Revised: During the large-batch phase (phase 1), all workers share a common model and have the same generalization performance.

Original: 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.
Revised: 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.

Original: We also plot the test-accuracy of the averaged model that would result were we to stop phase 2 at that point.
Revised: We also plot the test-accuracy of the averaged model that would result were we to stop phase 2 at that point.

Original: Note that the averaged model performs consistently better than each individual model .
Revised: Note that the averaged model performs consistently better than each individual model .

Original: [Figure]
Revised: [Figure]

Original:  L OSS L ANDSCAPE V ISUALIZATION AROUND SWAP ITERATES
Revised:  L OSS L ANDSCAPE V ISUALIZATION AROUND SWAP ITERATES

Original: 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.
Revised: 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.

Original: 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 .
Revised: 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 .

Original: We plot the loss value generated by model θ = θ 1 + αu + βv at the location ( α , β ).
Revised: We plot the loss value generated by model θ = θ 1 + αu + βv at the location ( α , β ).

Original: To plot a loss value, we ﬁrst 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.
Revised: To plot a loss value, we ﬁrst 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.

Original: In Figure 2, we plot the training and testing error for the CIFAR10 dataset.
Revised: In Figure 2, we plot the training and testing error for the CIFAR10 dataset.

Original: Here ‘LB’ marks the output of phase one, ‘SGD’ the output of a single worker after phase two, and ‘SWAP’ the ﬁnal model.
Revised: Here ‘LB’ marks the output of phase one, ‘SGD’ the output of a single worker after phase two, and ‘SWAP’ the ﬁnal model.

Original: Color codes correspond to error measures at the points interpolated on the plane.
Revised: Color codes correspond to error measures at the points interpolated on the plane.

Original: 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’) and the output of one of the workers of phase 2 (‘SGD’) lie in the outer edges of the basin.
Revised: In Figure 2a, we observe that the level-sets of the training error (restricted to this plane) form an almost convex decay do η do B B W g w gradient end

Revised: 17 end do do

Revised: θ wt

Revised: ; /* first order method

Revised: update at local worker */

Revised: end 26 end /

Revised: * We get W

Revised: different models at the end of phase

Revised: 28 Compute

Revised: Result:

Revised: [Table]

Revised: basin and that both the output of phase 1 (‘LB’) 2 and the output of one of the workers of phase(‘SGD’) lie in the outer edges of the basin.

Original: Importantly, during phase 2 the model traversed to a different side of the basin (and not to the center).
Revised: Importantly, during phase 2 the model traversed to a different side of the basin (and not to the center).

Original: Also, the ﬁnal model (‘SWAP’) is closer to the center of the basin.
Revised: Also, the ﬁnal model (‘SWAP’) is closer to the center of the basin.

Original: 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.
Revised: 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.

Original: But, since the ‘SWAP’ point is closer to the center of the basin, it is less affected by the change in topology.
Revised: But, since the ‘SWAP’ point is closer to the center of the basin, it is less affected by the change in topology.

Original: In Figure 3, we neglect the ‘LB’ point and plot the plane spanned by three workers ‘SGD1’, ‘SGD2’, ‘SGD3’.
Revised: In Figure 3, we neglect the ‘LB’ point and plot the plane spanned by three workers ‘SGD1’, ‘SGD2’, ‘SGD3’.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: For reference, we have also plotted the best model that can be generated by this region of the plane.do do
Revised: For reference, we have also plotted the best model that can be generated by this region of the plane.

Original: LB
Revised: LB

Original: SWAP
Revised: SWAP

Original: SGD 0.
Revised: SGD 0.

Original: > 2.
Revised: > 2.

Original: (a) Train Error (%) 0.
Revised: [Figure]

Original: LB
Revised: LB

Original: SWAP
Revised: SWAP

Original: SGD 5.
Revised: SGD 5.

Original: > 8.
Revised: > 8.

Original: [Figure]
Revised: [Figure]

Revised: SGD

Revised: SGD

Revised: SGD

Revised: SWAP 0.

Revised: > 2.

Revised: [Figure]

Revised: SGD

Revised: SGD

Revised: SGD

Revised: SWAP

Original:  4. S AMPLING FROM INDEPENDENT RUNS OF SGD OR SAMPLING FROMONE
Revised: BEST 5. > 8. [Figure] 4. S AMPLING FROM INDEPENDENT RUNS OF SGD OR SAMPLING FROM

Original:  In (Mandt et al., 2017), the authors argue that in the later stages of SGD the weight iterates behave similar to an Ornstein Uhlenbeck process.
Revised: ONE In (Mandt et al., 2017), the authors argue that in the later stages of SGD the weight iterates behave similar to an Ornstein Uhlenbeck process.

Original: So, by maintaining a constant learning rate the SGD iterates should reach a stationary distribution that is similar to a high-dimensional Gaussian.
Revised: So, by maintaining a constant learning rate the SGD iterates should reach a stationary distribution that is similar to a high-dimensional Gaussian.

Original: This distribution is centered at the local minimum, has a covariance that grows proportionally with the learning rate, inversely proportional to the batch size and has a shape that depends on both the
Revised: This distribution is centered at the local minimum, has a covariance that grows proportionally with the learning rate, inversely proportional to the batch size and has a shape that depends on both the

Original: Hessian of the mean loss and covariance of the gradient.
Revised: Hessian of the mean loss and covariance of the gradient.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure] [Figure] [Figure] The authors of (Izmailov et al., 2018) argue that by virtue of being a high dimensional Gaussian all the mass of the distribution is concentrated near the ‘shell’ of the ellipsoid, and therefore, it is unlikely for SGD to access the interior.
Revised: The authors of (Izmailov et al., 2018) argue that by virtue of being a high dimensional Gaussian all the mass of the distribution is concentrated near the ‘shell’ of the ellipsoid, and therefore, it is unlikely for SGD to access the interior.

Original: 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 be closer to the center.
Revised: 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 be closer to the center.

Original: Without any further assumptions, we can justify sampling from different SGD runs (as done in phase 2 during SWAP).
Revised: Without any further assumptions, we can justify sampling from different SGD runs (as done in phase 2 during SWAP).

Original: 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 to the same stationary distribution, and each run can generate independent samples from it.
Revised: 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 to the same stationary distribution, and each run can generate independent samples from it.

Original: 4. O RTHOGONALITY OF THE GRADIENT AND THE DIRECTION TO THE CENTER OF BASIN
Revised: 4. O RTHOGONALITY OF THE GRADIENT AND THE DIRECTION TO THE CENTER OF BASIN

Original: To win some intuition on the advantage that SWA and SWAP have over SGD, we measure the cosine similarity between the gradient descent direction, − g i , and the direction towards the output of SWAP, ∆ θ = θ swap − θ i .
Revised: To win some intuition on the advantage that SWA and SWAP have over SGD, we measure the cosine similarity between the gradient descent direction, − g i , and the direction towards the output of SWAP, ∆ θ = θ swap − θ i .

Original: In Figure 4, we see that the cosine similarity, (cid:104) ∆ θ, − g i (cid:105) (cid:107) g
Revised: In Figure 4, we see that the cosine similarity, (cid:104) ∆ [Table] [Table]

Original: i (cid:107)(cid:107) ∆ θ (cid:107) , decreases as the training enters its later stages.
Revised: , decreases as the training enters its later stages.

Original: 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.
Revised: 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.

Original: However, averaging samples from different sides of the basin can (and does) make faster progress towards the center. 
Revised: However, averaging samples from different sides of the basin can (and does) make faster progress towards the center. [Figure]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section we evaluate the performance of SWAP for image classiﬁcation tasks on the CIFAR10,
Revised: In this section we evaluate the performance of SWAP for image classiﬁcation tasks on the CIFAR10,

Original: CIFAR100, and ImageNet datasets.
Revised: CIFAR100, and ImageNet datasets.

Original: 5. CIFAR10 AND CIFAR 
Revised: 5. CIFAR10 AND CIFAR

Original: For the experiments in this subsection, we found the best hyper-parameters using grid searches (see
Revised: For the experiments in this subsection, we found the best hyper-parameters using grid searches (see

Original: Appendix A for details).
Revised: Appendix A for details).

Original: We train using mini-batch SGD with Nesterov momentum (set to 0.9) and weight decay of 5 × 10 − 4 .
Revised: We train using mini-batch SGD with Nesterov momentum (set to 0.9) and weight decay of 5 × 10 − 4 .

Original: We augment the data using cutout (DeVries & Taylor, 2017) and use a fastto-train custom ResNet 9 from a submission 2 to the DAWNBench leaderboard (Coleman et al.).
Revised: We augment the data using cutout (DeVries & Taylor, 2017) and use a fastto-train custom ResNet 9 from a submission 3 to the DAWNBench leaderboard (Coleman et al.).

Original: All experiments were run on one machine with 8 NVIDIA Tesla V100 GPUs and use Horovod (Sergeev & Del Balso, 2018) to distribute the computation.
Revised: All experiments were run on one machine with 8 NVIDIA Tesla V100 GPUs and use Horovod (Sergeev & Del Balso, 2018) to distribute the computation.

Original: All statistics were collected over 10 different runs.
Revised: All statistics were collected over 10 different runs.

Original: CIFAR10 : For these experiments, we used the following settings—SWAP phase one: 4096 samples per batch using 8 GPUs (512 samples per GPU).
Revised: CIFAR10 : For these experiments, we used the following settings—SWAP phase one: 4096 samples per batch using 8 GPUs (512 samples per GPU).

Original: Phase one is terminated when the training accuracy reaches 98% (on average 108 epochs).
Revised: Phase one is terminated when the training accuracy reaches 98% (on average 108 epochs).

Original: SWAP phase two: 8 workers with one GPU each andsamples per batch for 30 epochs.
Revised: SWAP phase two: 8 workers with one GPU each andsamples per batch for 30 epochs.

Original: The experiment that uses only large-batches had 4096 samples per batch across 8 GPUs and is run for 150 epochs.
Revised: The experiment that uses only large-batches had 4096 samples per batch across 8 GPUs and is run for 150 epochs.

Original: The experiments that use only small-batches had 512 samples per batch on 2 GPUs and is trained for 100 epochs.
Revised: The experiments that use only small-batches had 512 samples per batch on 2 GPUs and is trained for 100 epochs.

Original: Table 1 compares the best test accuracies and corresponding training times for models trained with small-batch only, with large-batch only, and with SWAP.
Revised: Table 1 compares the best test accuracies and corresponding training times for models trained with small-batch only, with large-batch only, and with SWAP.

Original: We report the average accuracy of the workers before averaging and the accuracy of the ﬁnal model.
Revised: We report the average accuracy of the workers before averaging and the accuracy of the ﬁnal model.

Original: [Table]
Revised: [Table]

Original: CIFAR100 : For these experiments, we use the following settings—SWAP phase one: 2048 samples per batch using 8 GPUs (256 samples per GPU).
Revised: CIFAR100 : For these experiments, we use the following settings—SWAP phase one: 2048 samples per batch using 8 GPUs (256 samples per GPU).

Original: Phase one exits when the training accuracy reaches 90% (on average 112 epochs).
Revised: Phase one exits when the training accuracy reaches 90% (on average 112 epochs).

Original: SWAP phase two: 8 workers with one GPU each and 128 samples per batch, training for for 10 epochs.
Revised: SWAP phase two: 8 workers with one GPU each and 128 samples per batch, training for for 10 epochs.

Original: The experiments that use only large-batch training were run for 150 epochs with batches of 2048 on 8 GPUs The experiments that use only small-batch were trained for 150 epochs using batches of 128 on 1 GPU.
Revised: The experiments that use only large-batch training were run for 150 epochs with batches of 2048 on 8 GPUs The experiments that use only small-batch were trained for 150 epochs using batches of 128 on 1 GPU.

Original: [Table]
Revised: [Table]

Original: Table 2 compares the best test accuracies and corresponding training times for models trained with only small-batches (for 150 epochs), with only large-batches (for 150 epochs), and with SWAP.
Revised: Table 2 compares the best test accuracies and corresponding training times for models trained with only small-batches (for 150 epochs), with only large-batches (for 150 epochs), and with SWAP.

Original: For SWAP, we report test accuracies obtained using the last SGD iterate before averaging, and test accuracy of the ﬁnal model obtained after averaging.
Revised: For SWAP, we report test accuracies obtained using the last SGD iterate before averaging, and test accuracy of the ﬁnal model obtained after averaging.

Original: We observe signiﬁcant improvement in test accuracies after averaging the models.
Revised: We observe signiﬁcant improvement in test accuracies after averaging the models.

Original: For both CIFAR 10 and CIFAR100, training with small-batches achieves higher testing accuracy than training with large-batches but takes much longer to train.
Revised: For both CIFAR 10 and CIFAR100, training with small-batches achieves higher testing accuracy than training with large-batches but takes much longer to train.

Original: SWAP, however, terminates in time comparable to the large-batch run but achieves accuracies on par (or better) than small batch training.
Revised: SWAP, however, terminates in time comparable to the large-batch run but achieves accuracies on par (or better) than small batch training.

Original: Achieving state of the art training speeds for CIFAR10 :
Revised: Achieving state of the art training speeds for CIFAR10 :

Original: At the time of writing the front-runner of the DAWNBench competition takes 37 seconds with 4 Tesla V100 GPUs to train CIFAR10 to 94% test accuracy.
Revised: At the time of writing the front-runner of the DAWNBench competition takes 37 seconds with 4 Tesla V100 GPUs to train CIFAR10 to 94% test accuracy.

Original: Using SWAP with 8 Tesla V100 GPUs, a phase one 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.
Revised: Using SWAP with 8 Tesla V100 GPUs, a phase one 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.

Original: 5. E XPERIMENTS ON I MAGE N ET
Revised: 5. E XPERIMENTS ON I MAGE N ET

Original: We use SWAP to accelerate a publicly available fast-to-train ImageNet model with published learning rate and batch size schedules 3 .
Revised: We use SWAP to accelerate a publicly available fast-to-train ImageNet model with published learning rate and batch size schedules 4 .

Original: The default settings for this code modify the learning-rates and batch sizes throughout the optimization (see Figure 5).
Revised: The default settings for this code modify the learning-rates and batch sizes throughout the optimization (see Figure 5).

Original: Our small-batch experiments train ImageNet for 28 epochs using the published schedules with no modiﬁcation and are run on 8 Tesla V
Revised: Our small-batch experiments train ImageNet for 28 epochs using the published schedules with no modiﬁcation and are run on 8 Tesla V

Original: GPUs.
Revised: GPUs.

Original: 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.
Revised: 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.

Original: 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-batches for 6 epochs.
Revised: 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-batches for 6 epochs.

Original: We observe that doubling the batch size reduces the Top1 and Top5 test accuracies with respect to the small-batch run.
Revised: We observe that doubling the batch size reduces the Top1 and Top5 test accuracies with respect to the small-batch run.

Original: SWAP, however, recovers the generalization performance at substantially reduced training times.
Revised: SWAP, however, recovers the generalization performance at substantially reduced training times.

Original: Our results are compiled in Table 3 (the statistics were collected overruns).
Revised: Our results are compiled in Table 3 (the statistics were collected overruns).

Original: 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.
Revised: 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: 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 speciﬁc to the dataset.

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: [Table]
Revised: [Table]

Original: after averaging)

Original: [Table]
Revised: [Figure]

Original: 5. E MPIRICAL COMPARISON OF SWA AND SWAP
Revised: 5. E MPIRICAL COMPARISON OF SWA AND SWAP

Original: We now compare SWAP with SWA: the sequential weight averaging algorithm from Izmailov et al.
Revised: We now compare SWAP with SWA: the sequential weight averaging algorithm from Izmailov et al.

Original: For the experiments in this section, we use the CIFAR100 dataset.
Revised: For the experiments in this section, we use the CIFAR100 dataset.

Original: We sample the same number of models for both SWA and SWAP and maintain the same number of epochs per sample.
Revised: We sample the same number of models for both SWA and SWAP and maintain the same number of epochs per sample.

Original: For SWA, we sample each model with 10 epochs in-between and average them to get the ﬁnal model.
Revised: For SWA, we sample each model with 10 epochs in-between and average them to get the ﬁnal model.

Original: For SWAP, we run 8 independent workers for 10 epochs each and use their average as the ﬁnal model.
Revised: For SWAP, we run 8 independent workers for 10 epochs each and use their average as the ﬁnal model.

Original: Large-batch SWA : We explore if SWA can recover the test accuracy of small-batch training on a large-batch training run.
Revised: Large-batch SWA : We explore if SWA can recover the test accuracy of small-batch training on a large-batch training run.

Original: We use the same (large) batch size throughout.
Revised: We use the same (large) batch size throughout.

Original: 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).
Revised: 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).

Original: See Figure 6a for an illustration of the learning rate schedule.
Revised: See Figure 6a for an illustration of the learning rate schedule.

Original: As expected we observe that the large-batch training run achieves lower training accuracy, but surprisingly SWA was unable to improve it (see Table 4, row 1).
Revised: As expected we observe that the large-batch training run achieves lower training accuracy, but surprisingly SWA was unable to improve it (see Table 4, row 1).

Original: Large-batch followed by small-batch SWA : We evaluate the effect of executing SWA using smallbatches after a large-batch training run.
Revised: Large-batch followed by small-batch SWA : We evaluate the effect of executing SWA using smallbatches after a large-batch training run.

Original: We interrupt the large-batch phase at the same accuracy we interrupt phase 1 of our CIFAR100 experiment (Table 2).
Revised: We interrupt the large-batch phase at the same accuracy we interrupt phase 1 of our CIFAR100 experiment (Table 2).

Original: In this case, the small-batch phase uses a single worker and samples the models sequentially.
Revised: In this case, the small-batch phase uses a single worker and samples the models sequentially.

Original: SWA is able to reach the test accuracy of a small-batch run but requires more than three times longer than SWAP to compute the model (see
Revised: SWA is able to reach the test accuracy of a small-batch run but requires more than three times longer than SWAP to compute the model (see

Original: Table 4, row 2).
Revised: Table 4, row 2).

Original: An illustration of the learning rate schedule is provided in Figure 6b.
Revised: An illustration of the learning rate schedule is provided in Figure 6b.

Original: Small-batch SWA and SWAP
Revised: Small-batch SWA and SWAP

Original: : We start the SWA cyclic learning rate schedule from the best model found by solely small-batch training (table 2, row 1).
Revised: : We start the SWA cyclic learning rate schedule from the best model found by solely small-batch training (table 2, row 1).

Original: Since the cycle length and cycle count are ﬁxed, the only free parameter is the peak learning rate.
Revised: Since the cycle length and cycle count are ﬁxed, the only free parameter is the peak learning rate.

Original: We select this using a grid-search.
Revised: We select this using a grid-search.

Original: Once the SWA schedule is speciﬁed, we re-use the peak learning rate settings in SWAP.
Revised: Once the SWA schedule is speciﬁed, we re-use the peak learning rate settings in SWAP.

Original: We start phase[Figure] two from the model that was generated as the output of phase 1 for the experiment on section 5.
Revised: We start phase two from the model that was generated as the output of phase 1 for the experiment on section 5.

Original: reported on table 2 rows 3 and 4.
Revised: reported on table 2 rows 3 and 4.

Original: With these settings, small-batch SWA achieves better accuracy than SWAP (by around ∼ 0 .
Revised: With these settings, small-batch SWA achieves better accuracy than SWAP (by around ∼ 0 .

Original: 9% ) at 6.8x more training time.
Revised: 9% ) at 6.8x more training time.

Original: Next, we wish to explore the speed-up that SWAP achieves over SWA if the precision of SWA is set as a target.
Revised: Next, we wish to explore the speed-up that SWAP achieves over SWA if the precision of SWA is set as a target.

Original: To that end, we relax the constraints on SWAP.
Revised: To that end, we relax the constraints on SWAP.

Original: By increasing the phase two schedule from one 10 epoch cycle to two 20 epoch cycles and sampling two models from each worker ( models) the resulting model achieved a test accuracy of 79 .
Revised: By increasing the phase two schedule from one 10 epoch cycle to two 20 epoch cycles and sampling two models from each worker ( models) the resulting model achieved a test accuracy of 79 .

Original: 11% in 241 seconds or 3 .
Revised: 11% in 241 seconds or 3 .

Original: 5 x less time.
Revised: 5 x less time.

Original: CIFAR
Revised: [Table]

Original: Test accuracy before averaging (%)
Revised: before averaging [Table]

Original: Test accuracy after

Original: [Table]
Revised: [Table]

Original: Large-batch SWA
Revised: Large-batch SWA

Original: [Table]
Revised: [Table]

Original: Large-batch followed by small-batch SWA
Revised: Large-batch followed by small-batch SWA

Original: [Table]
Revised: [Table]

Original: Small-batch SWA
Revised: Small-batch SWA

Original: [Table]
Revised: [Table]

Original:  (10 small-batch epochs)
Revised: SWAP (10 small-batch epochs)

Original: [Table]
Revised: [Table]

Revised: SWAP (40 small-batch epochs)

Original: [Table]
Revised: [Table]

Original:  C ONCLUSIONS AND F UTURE W ORK
Revised:  C ONCLUSIONS AND F UTURE W ORK

Original: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm that uses a variant of
Revised: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm that uses a variant of

Original: Stochastic Weight Averaging (SWA) to improve the generalization performance of a model trained with large mini-batches.
Revised: Stochastic Weight Averaging (SWA) to improve the generalization performance of a model trained with large mini-batches.

Original: Our algorithm uses large mini-batches to compute an approximate solution quickly and then reﬁnes it by averaging the weights of multiple models trained using small-batches.
Revised: Our algorithm uses large mini-batches to compute an approximate solution quickly and then reﬁnes it by averaging the weights of multiple models trained using small-batches.

Original: The ﬁnal model obtained after averaging has good generalization performance and is trained in a shorter time.
Revised: The ﬁnal model obtained after averaging has good generalization performance and is trained in a shorter time.

Original: We believe that this variant and this application of SWA are novel.
Revised: We believe that this variant and this application of SWA are novel.

Original: We observed that using large-batches in the initial stages of training does not preclude the models from achieving good generalization performance.
Revised: We observed that using large-batches in the initial stages of training does not preclude the models from achieving good generalization performance.

Original: That is, by reﬁning the output of a large-batch run, with models sampled sequentially as in SWA or in parallel as in SWAP, the resulting model is able to perform as well as the models trained using small-batches only.
Revised: That is, by reﬁning the output of a large-batch run, with models sampled sequentially as in SWA or in parallel as in SWAP, the resulting model is able to perform as well as the models trained using small-batches only.

Original: We conﬁrm this in the image classiﬁcation datasets CIFAR10, CIFAR100, and ImageNet.
Revised: We conﬁrm this in the image classiﬁcation datasets CIFAR10, CIFAR100, and ImageNet.

Original: Through visualizations, we 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.
Revised: Through visualizations, we 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.

Original: It’s interesting to note that the basin into which the large mini-batch run is converging to seems to be the same basin where the reﬁned models are found.
Revised: It’s interesting to note that the basin into which the large mini-batch run is converging to seems to be the same basin where the reﬁned models are found.

Original: 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.
Revised: 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.

Original: Our method requires the choice of (at least) one more hyperparameter: the transition point between the large-batch and small-batch.
Revised: Our method requires the choice of (at least) one more hyperparameter: the transition point between the large-batch and small-batch.

Original: For our experiments, we chose this by using a grid search.
Revised: For our experiments, we chose this by using a grid search.

Original: A principled method to choose the transition point will be the focus of future work.
Revised: A principled method to choose the transition point will be the focus of future work.

Revised: 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)

Revised: (You et al., 2017), mixed-precision training Jia et al.

Revised: (2018), post-local SGD (Lin et al., 2018) or NovoGrad (Ginsburg et al., 2019).

Revised: 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 ﬁrst stage of SWAP by reducing the communication overhead.

Original: A H YPERPARAMETERS FOR CIFAR10 AND CIFAR100 E XPERIMENTS
Revised: A H YPERPARAMETERS FOR CIFAR10 AND CIFAR100 E XPERIMENTS

Original: We provide the parameters used in the experiments of Section 5.1.
Revised: We provide the parameters used in the experiments of Section 5.1.

Original: These were obtained by doing independent grid searches for each experiment.
Revised: These were obtained by doing independent grid searches for each experiment.

Original: For all CIFAR experiments, the momentum and weight decay constants were kept at 0.9 and 5 × 10 − 4 respectively.
Revised: For all CIFAR experiments, the momentum and weight decay constants were kept at 0.9 and 5 × 10 − 4 respectively.

Original: Tables 5 and 6 list the remaining hyperparameters.
Revised: Tables 5 and 6 list the remaining hyperparameters.

Original: When a stopping accuracy of 100% is listed, we mean that the maximum number of epochs were used.
Revised: When a stopping accuracy of 100% is listed, we mean that the maximum number of epochs were used.

Original: [Table] [Table]
Revised: [Table] [Table]