SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum

Introduction

Distributed optimization (Chen et al. 2016; Goyal et al. 2017) is essential for training large models on large datasets (Radford et al. 2019; Liu et al. 2019; Mahajan et al. 2018b). Currently, the most widely-used approaches have workers compute small mini-batch gradients locally, in parallel, and then aggregate these using a blocking communication primitive, AllReduce, before taking an optimizer step. Communication overhead is a major issue limiting the scaling of this approach, since AllReduce must complete before every step and blocking communications are sensitive to stragglers (Dutta et al. 2018; Ferdinand et al. 2019).

Multiple complementary approaches have recently been investigated to reduce or hide communication overhead. Decentralized training (Jiang et al. 2017; Lian et al. 2017; Lian et al. 2018; Assran et al. 2019) reduces idling due to blocking and stragglers by employing approximate gradient aggregation (e.g., via gossip or distributed averaging). Approaches such as Local SGD reduce the frequency of communication by having workers perform multiple updates between each round of communication (McDonald et al. 2010; McMahan et al. 2017; Zhou & Cong 2018; Stich 2019; Yu et al. 2019b). It is also possible to combine decentralized algorithms with Local SGD (Wang & Joshi 2018; Wang et al. 2019). These approaches reduce communication overhead while injecting additional noise into the optimization process. Consequently, although they run faster than large mini-batch methods, the resulting models may not achieve the same quality in terms of training loss or generalization accuracy after the same number of iterations.

Momentum is believed to be a critical component for training deep networks, and it has been empirically demonstrated to improve both optimization and generalization (Sutskever et al. 2013). Yet, there is no consensus on how to combine momentum with communication efficient training algorithms. Momentum is typically incorporated into such approaches by having workers maintain separate buffers which are not synchronized (Lian et al. 2017; Lian et al. 2018; Assran et al. 2019; Koloskova et al. 2019a). However, recent work shows that synchronizing the momentum buffer, using periodic AllReduce or a decentralized method, leads to improvements in accuracy at the cost of doubling the communication overhead (Yu et al. 2019a). In block-wise model update filtering (BMUF), nodes perform multiple local optimization steps between communication rounds (similar to local SGD), and they also maintain a momentum buffer that is only updated after each communication round (Chen & Huo 2016). Although it is now commonly used for training speech models, there are no theoretical convergence guarantees for BMUF, and it has not been widely applied to other tasks (e.g., in computer vision or natural language processing).

Inspired by BMUF, we propose a general framework called slow momentum (SlowMo) to improve the accuracy of communication-efficient distributed training methods. SlowMo runs on top of a base algorithm, which could be local SGD or a decentralized method such as stochastic gradient push (SGP) (Nedić & Olshevsky 2016; Assran et al. 2019). Periodically, after taking some number $\tau$ of base algorithm steps, workers average their parameters using AllReduce and perform a momentum update. We demonstrate empirically that SlowMo consistently improves optimization and generalization performance across a variety of base algorithms on image classification and neural machine translation tasks—training ResNets on CIFAR-10 and ImageNet, and training a transformer on WMT’16 En-De. Ultimately, SlowMo allows us to reap the speedup and scaling performance of communication-efficient distributed methods without sacrificing as much in accuracy.

We also prove theoretical bounds showing that SlowMo converges to a stationary point of smooth non-convex functions at a rate $\mathcal{O}(1/\sqrt{mT\tau})$ after $T\tau$ total inner optimization steps and $T$ SlowMo updates with $m$ worker nodes, for a variety of base optimizers. Thus, SlowMo is order-wise no slower than stochastic gradient descent. BMUF and the recently-proposed Lookahead optimizer (Zhang et al. 2019) can be expressed through the SlowMo framework, and so our results also translate to the first theoretical convergence guarantees for both of these methods.

The Slow Momentum (SlowMo) Framework

SlowMo is a framework intended for solving stochastic optimization problems of the form

Within each outer iteration, workers first take $\tau$ steps of the base optimizer. The base optimizer could be a method which involves no communication, such as SGD (with or without momentum) or a decentralized algorithm which involves some communication, such as stochastic gradient push (SGP) (Assran et al. 2019). We denote these updates by $\bm{x}_{t,k+1}^{(i)}=\bm{x}_{t,k}^{(i)}-\gamma_{t}\bm{d}_{t,k}^{(i)}$ where $\gamma_{t}$ is the base optimizer (fast) learning rate and $\bm{d}_{t,k}^{(i)}$ is the update direction used at worker $i$. If the base optimizer is SGD then $\bm{d}_{t,k}^{(i)}=\nabla F_{i}(\bm{x}_{t,k}^{(i)};\xi_{t,k}^{(i)})$. For other base optimizers which may use additional local momentum or communication, $\bm{d}_{t,k}^{(i)}$ represents the full update applied at worker $i$ on this step. Specific examples of $\bm{d}_{t,k}^{(i)}$ for different base optimizers are presented in Table C.1 in Appendix C.

After the $\tau$ base optimizer steps, the workers calculate the average $\bm{x}_{t,\tau}=\bm{x}_{t,0}-\frac{\gamma_{t}}{m}\sum_{i=1}^{m}\sum_{k=0}^{\tau-1}\bm{d}_{t,k}^{(i)}$ using AllReduce (line 6), and then they perform a slow momentum update (lines 7–8),

A specific SlowMo algorithm instance is obtained by specifying the base algorithm and the hyperparameters $\alpha$, $\beta$, $\gamma$, and $\tau$. We can recover a number of existing algorithms in this framework. When the base algorithm is SGD, $\tau=1$, $\alpha=1$, and $\beta\in[0,1)$, we recover standard large mini-batch SGD with learning rate $\gamma$ and momentum $\beta$. When the base algorithm is SGD, $\tau>1$, $\alpha=1$, and $\beta=0$, we recover Local SGD (McDonald et al. 2010; Stich 2019; Yu et al. 2019b; Wang & Joshi 2018). When the base algorithm does not involve communication among workers, $\tau>1$ and $\beta>0$, we recover BMUF (Chen & Huo 2016).

We also obtain interesting novel distributed algorithms. In particular, the experiments in Section 4 demonstrate that using SlowMo with a decentralized base algorithm like SGP and reasonable values of $\tau$ consistently leads to improved optimization and generalization performance over the base method alone, without a significant increase in runtime. We also observe empirically that, for a fixed number of iterations, SlowMo combined with SGP is superior to SlowMo combined with SGD.

The above are all distributed algorithms. Perhaps surprisingly, SlowMo also encompasses a recently-introduced non-distributed method: if we have $m=1$ worker with SGD/Adam as the base algorithm, $\alpha\in(0,1]$, $\beta=0$, and $\tau>0$, we recover the Lookahead optimizer of Zhang et al. 2019, which also has a nested loop structure. Section 5 provides theoretical convergence guarantees when using the SlowMo framework to minimize smooth non-convex functions, and thus provides the first theoretical convergence guarantees in the literature for BMUF and Lookahead in this setting.

Related Work

The idea of reducing communication overhead by using AllReduce to synchronize parameters after every $\tau>0$ optimizer steps has been considered at least since the work of McDonald et al. 2010, and has been more recently referred to as Local SGD in the literature. Elastic-average SGD (Zhang et al. 2015) uses a related approach, but with a parameter server rather than AllReduce. Lin et al. 2018 apply Local SGD for distributed training of deep neural networks and propose post-local SGD, which starts by running AllReduce-SGD for some epochs before switching to Local SGD, to improve generalization at the cost of additional communication.

Decentralized methods use approximate distributed averaging over a peer-to-peer topology, rather than AllReduce. This decouples communication but also injects additional noise in the optimization process since the models at different workers are no longer precisely synchronized. Lian et al. 2017 present decentralized parallel SGD (D-PSGD), where each worker sends a copy of its model to its peers at every iteration, and show it can be faster than parameter-server and AllReduce methods for training deep neural networks. Lian et al. 2018 study an asynchronous extension, AD-PSGD. Assran et al. 2019 study stochastic gradient push (SGP), and propose its asynchronous counterpart overlap SGP (OSGP), which achieve a further speedup over D-PSGD and AD-PSGD by using less coupled communication. D-PSGD, AD-PSGD, and SGP all have similar theoretical convergence guarantees for smooth non-convex functions, showing a linear scaling relationship between the number of workers and the number of iterations to reach a neighborhood of a first-order stationary point. Although the theory for all three methods only covers the case of SGD updates without momentum, implementations use momentum locally at each worker, and workers only average their model parameters (not momentum buffers). Yu et al. 2019a prove that linear scaling holds when workers average their parameters and momentum buffers, although this doubles the communication overhead. We refer to this approach as double-averaging below.

Scaman et al. 2019 establish optimal rates of convergence for decentralized optimization methods in the deterministic, convex setting. Richards & Rebeschini 2019 provide guarantees on the generalization error of non-parametric least-squares regression trained using decentralized gradient descent, showing that there are regimes where one can achieve a linear speedup. Neither of these results apply directly to the setting considered in this paper, which focuses on smooth non-convex stochastic optimization, and extending this line of work to non-convex settings is an interesting direction.

Mahajan et al. 2018a propose an approach to distributed learning of linear classifiers (i.e., convex problems) where, in parallel, workers minimize locally formed approximate loss functions, and then the resulting minimizers are averaged to determine a descent direction. Methods which fit in the SlowMo framework, including Local SGD, BMUF (Chen & Huo 2016), and the serial Lookahead optimizer (Zhang et al. 2019), can be seen as related to this approach, where the actual loss function at each worker is used rather than an approximate one, and where the descent direction is used in a momentum update rather than a (deterministic) line search method.

Note that various approaches to gradient compression have been proposed to reduce the communication overhead for AllReduce and decentralized learning methods (Alistarh et al. 2007; Wen et al. 2007; Bernstein et al. 2019; Karimireddy et al. 2019; Koloskova et al. 2019b; Vogels et al. 2019). However, it is presently not clear to what extent compression may be beneficial for methods like BMUF, D-PSGD, SGP, and OSGP, which perform averaging on the model parameters rather than on gradients. Combining SlowMo with compression techniques is an interesting and important direction for future work.

Finally, although momentum methods are known to achieve accelerated rates for deterministic optimization, currently the theoretical understanding of the benefits of momentum methods (both serial and parallel) is limited (Bottou et al. 2018). Loizou & Richtárik 2017 show that accelerated convergence rates can be achieved when the objective is a quadratic finite sum (i.e., least squares problem) that satisfies an interpolation condition. Can et al. 2019 show that accelerated rates can also be achieved in the more general setting of smooth, strongly convex objectives, in a regime where the noise level is below an explicit threshold, and where the rate of convergence is measured in terms of the 1-Wasserstein metric of the distribution of trajectories produced by the momentum method. In the setting of smooth non-convex functions, Gitman et al. 2019 establish stability and asymptotic convergence results for the quasi-hyperbolic momentum method (Ma & Yarats 2019), which can be viewed as interpolating between SGD and a stochastic momentum method. Extending these results, which focus on the serial setting, to obtain decentralized momentum methods with guaranteed acceleration, is an important direction for future work.

Experimental Results

We evaluate the effectiveness of SlowMo on three datasets: image classification on CIFAR-10 and ImageNet, and neural machine translation on WMT’16-En-De. All experiments use NVIDIA DGX-1 servers as worker nodes. Each server contains $8$ NVIDIA V100 GPUs and the servers are internetworked via commodity 10 Gbps Ethernet.

On CIFAR-10 (Krizhevsky et al. 2009), we train a ResNet-18 (He et al. 2016) using $32$ V100 GPUs, located on $32$ different worker nodes. The total mini-batch size is $4096$, and we train for 200 epochs. The learning rate ($\gamma_{t}$) linearly increases during the first $5$ epochs, following the warm-up strategy in Goyal et al. 2017, and then decays by a factor of $10$ at epochs $100,150$, and $175$. The (fast) learning rate was tuned separately for each base optimizer. All experiments were run $5$ times with different random seeds, and the mean metrics are reported.

On ImageNet (Krizhevsky et al. 2012), we train a ResNet-50 (He et al. 2016) using $32$ worker nodes (i.e., $256$ GPUs). The total mini-batch size is $8192$, and we train for 90 epochs. The learning rate schedule is identical to (Goyal et al. 2017), i.e., linear warm-up in the first $5$ epochs and decay by a factor of $10$ at epochs $30$, $60$ and $80$.

On WMT’16-En-De, we train a transformer model (Vaswani et al. 2017) using $8$ worker nodes (i.e., $64$ GPUs). The model is trained with $200$k token batches, and we train for $25$ epochs. We follow the experimental setting of Ott et al. 2018.

For each task, we consider several baselines: (i) Local SGD /Local Adam, where worker nodes independently run single-node SGD (with Nesterov momentum) or Adam and periodically average model parameters; (ii) stochastic gradient push (SGP), the state-of-the-art synchronous decentralized training method; and (iii) Overlap-SGP (OSGP), an asynchronous version of SGP. For each baseline, we examine its performance with and without SlowMo. Recall that Local SGD and Local Adam with SlowMo are equivalent to BMUF. Local SGD and Local Adam do not involve communication during the inner loop (base optimizer) updates, while SGP and OSGP involve gossiping with one peer at every step. In addition, we also evaluate the performance of AR-SGD/AR-Adam, the traditional AllReduce implementation of parallel SGD/Adam. Details of all baseline methods are provided in Appendices A and C.

In general, the hyperparameters of SlowMo (slow learning rate $\alpha$, slow momentum $\beta$, and number of inner loop steps $\tau$) need to be tuned for each base optimizer and task. The results in Table 1 all use $\alpha=1$, which we found to be consistently the best. For Local SGD (with or without SlowMo), we set $\tau=12$, and for all other baseline methods we use $\tau=48$. Using $\tau>12$ for Local SGD resulted in significantly worse loss/accuracy on ImageNet and WMT’16 En-De.

Note also that all of our baselines (or base optimizers) leverage a local momentum scheme, following previous works (Assran et al. 2019; Koloskova et al. 2019a). When using these methods with SlowMo, there are different ways to handle the base algorithm local momentum buffers at the beginning of each outer loop (line 2 in Algorithm 1): zeroing, averaging among workers, or maintaining the current local value. Section B.4 provides an empirical comparison. For the experiments reported here, when using SGD with Nesterov momentum as the base algorithm (CIFAR-10 and Imagenet) we zero the base algorithm buffer, and when using Adam as the base algorithm (WMT’16 En-De) we maintain the current value of the Adam buffers. We also tried to apply SlowMo on top of AR-SGD base optimizer, but we did not observe any improvement in that setting.

Optimization and Generalization Performance. Table 1 shows the best training loss and the validation accuracy/BLEU score for each baseline, with and without SlowMo. Using SlowMo consistently improves both the optimization and generalization performance across all training tasks and baseline algorithms. Figure 2 presents validation error/loss per epoch to give a sense of convergence speed. Observe that SGP with SlowMo substantially improves convergence, compared to SGP alone. We observe a similar phenomenon when comparing the training curves; see Appendix B.

Communication Cost. Table 2(b) shows the average training time per iteration on ImageNet and WMT’16. For SGP/OSGP, since the additional communication cost due to averaging in line 6 of Algorithm 1 is amortized over $\tau=48$ iterations, SlowMo maintains nearly the same speed as the corresponding base algorithm. For methods like Local SGD or Local Adam, which already compute an exact average every $\tau$ iterations, using SlowMo (i.e., using $\beta>0$) does not increase the amount of communication. In other words, using SlowMo on top of the base algorithm improves training/validation accuracy at a negligible additional communication cost.

Effects of $\tau$. The most important hyper-parameter in SlowMo is the number of base optimizer steps $\tau$ before each SlowMo update, since it influences both the accuracy and the training time. Figure 3 presents the validation accuracy and average iteration time of SGP-SlowMo for different values of $\tau$ on ImageNet and WMT’16. It can be observed that the validation performance does not monotonically increase or decrease with $\tau$. Instead, there is a best value. On both ImageNet and WMT’16, we find $\tau=48$ to be a good tradeoff between speed and accuracy. Moreover, SlowMo is pretty robust to the choice of $\tau$; even if $\tau=96$ for ImageNet and $\tau=192$ for WMT’16, SGP with SlowMo achieves better validation accuracy/loss than SGP alone.

We further investigate the effect of other hyperparameters (the slow learning rate $\alpha$, slow momentum $\beta$) as well as the different strategies for handling base algorithm buffers in Appendix B.

Comparison with Double-Averaging Momentum. As mentioned in Section 3, Yu et al. 2019a propose an alternative momentum scheme, double-averaging, to improve the convergence of Local SGD and D-PSGD. We empirically compare it with SlowMo in terms of the validation accuracy and average training time per iteration on ImageNet. When the base algorithm is SGP, double averaging achieves $75.54\%$ validation accuracy and takes $402$ ms per iteration on average, while SlowMo-SGP ($\tau=48$) reaches $75.73\%$ validation accuracy while taking $302$ ms per iteration on average. Similarly, when the baseline algorithm is Local SGD with $\tau=12$, double-averaging reaches $72.04\%$ and takes 405 ms per iteration, while SlowMo reaches $73.24\%$ and takes only $282$ ms per iteration.

Theoretical Results

In order to generalize the analysis to various base algorithms, we define $\bm{d}_{t,k}=\frac{1}{m}\sum_{i=1}^{m}\bm{d}_{t,k}^{(i)}$ as the average descent direction across the $m$ workers and make the following assumption.

As mentioned in Section 2, the analytic form of $\bm{d}_{t,k}$ depends on the choice of base algorithm. Therefore, the value of $V$ also changes. For instance, when the base algorithm is Local-SGD, then $\bm{d}_{t,k}=\frac{1}{m}\sum_{i=1}^{m}\nabla F_{i}(\bm{x}_{t,k}^{(i)};\xi_{t,k}^{(i)})$. It follows that

The above value ($V=\sigma^{2}/m$) can also be applied to other base algorithms, such as D-PSGD, SGP, and OSGP. More details are provided in Appendix C.

Our main convergence result is stated next. Proofs of all results in this section appear in Appendix D.

Suppose all worker nodes start from the same initial point $\bm{x}_{0,0}$, and the initial slow momentum is $\bm{u}_{0}=\bm{0}$. If we set $\alpha$, $\beta$, $\gamma_{t}=\gamma$, $\tau$ and $T$ so that $\frac{\alpha\gamma}{1-\beta}=\sqrt{\frac{m}{\tau T}}$ and the total iterations $\tau T$ satisfies $\tau T\geq mL^{2}\left(1+\sqrt{3}\max\left\{\frac{3\tau(1-\beta-\alpha)}{\alpha},\frac{4\tau\beta}{1-\beta},1\right\}\right)$, then under 1, 2 and 3, we have that:

where $f_{\text{inf}}=\inf_{\bm{x}}f(\bm{x})$.

Consistent with AR-SGD. Recall that AR-SGD is equivalent to taking $\tau=1$, $\alpha=1$, and $\beta=0$ and using SGD with learning rate $\gamma$ as the base optimizer. In this case, all terms on the RHS but the first one vanish, $V=\sigma^{2}/m$, and (5) is identical to the well-known rate of $\mathcal{O}(1/\sqrt{mT\tau})$ for SGD.

As an example, we provide the convergence analysis for the extreme case of Local SGD, where there is no communication between nodes during each inner iteration. Intuitively, using other base algorithms should only make this term smaller since they involve more communication than Local SGD.

Under the same conditions as 1, if the inner algorithm is Local SGD and there exists a positive finite constant $\zeta$ such that $\frac{1}{m}\sum_{i=1}^{m}\left\|\nabla f(\mathbf{x})-\nabla f_{i}(\mathbf{x})\right\|^{2}\leq\zeta^{2}$, then

Linear speedup. Corollary 1 shows that when the total number of steps $T\tau$ is sufficiently large: $T\geq m^{3}\tau^{3}$, the convergence rate will be dominated by the first term $\mathcal{O}(1/\sqrt{mT\tau})$. That is, in order to achieve an $\epsilon$ error, the algorithm requires $m$ times less steps when using $m$ times more worker nodes. This also recovers the same rate as AR-SGD.

Extension to single-node case. As mentioned in Section 2, when there is only one node and the slow momentum factor is $\beta=0$, the SlowMo-SGD is the Lookahead optimizer. One can directly apply 1 to this special case and get the following corollary.

Under the same conditions as 1, if the inner optimizer is AR-SGD and $\beta=0$, then one can obtain the following upper bound:

Faster SlowMo: Removing the Periodic AllReduce

SlowMo helps improve both the optimization and generalization of communication-efficient algorithms. When the base optimizer is SGP or OSGP, SlowMo also comes at the expense of higher communication cost, since it requires performing an exact average every $\tau$ iterations. In Local SGD/Local Adam, an exact average is also required every $\tau$ iterations, hence in comparison, using SlowMo does not increase the amount of communication required. Although the communication cost can be amortized, here we go one step further and propose a SGP-SlowMo variant, named SGP-SlowMo-noaverage, where we remove the exact average when we perform the SlowMo update, i.e we skip line 6 in Algorithm 1. We empirically evaluate this variant on the ImageNet and WMT’16 datasets, using $\alpha=1$, $\beta=0.6$ and $\tau=48$.

Surprisingly, we observe that SGP-SlowMo-noaverage achieves similar performances on Imagenet ($75.78\%$, compared to $75.73\%$ for SGP-SlowMo) and only slightly degrades the validation NLL on WMT’16 ($2.11$, compared to $2.10$), while preserving the iteration time of the base algorithm ($298$ ms per iteration on ImageNet and $1227$ ms per iteration on WMT’16) since this variant does not require additional communication. These results suggest that the slow momentum updates, and not the momentum buffer synchronization, contribute the most to the performance gain of SlowMo. We leave further investigation of SlowMo-SGP-noaverage for future work.

Concluding Remarks

In this paper, we propose a general momentum framework, SlowMo, for communication-efficient distributed optimization algorithms. SlowMo can be built on the top of SGD, as well as decentralized methods, such as SGP and (asynchronous) OSGP. On three different deep learning tasks, we empirically show that SlowMo consistently improves the optimization and generalization performance of the corresponding baseline algorithm while maintaining a similar level of communication efficiency. Moreover, we establish a convergence guarantee for SlowMo, showing that it converges to a stationary point of smooth and non-convex objective functions. Since BMUF (Chen & Huo 2016) can be expressed through SlowMo framework (by setting the base optimizer to be Local SGD or Local Adam), to the best of our knowledge, we provide the first convergence guarantee for BMUF in the literature.

References

Appendix A Experiment Details

All methods are implemented in PyTorch 1.0 (Paszke et al. 2017), and our experiments use CUDA 9.2, CUDNN 7.3, and NCCL 2.2.13. The ImageNet experiments build on the example from https://github.com/pytorch/examples/imagenet. The WMT’16 En-De experiments build on https://github.com/pytorch/fairseq. For SGP and OSGP we use the implementations available at https://github.com/facebookresearch/stochastic_gradient_push.

A.2 CIFAR-10

For the CIFAR-10 experiments, we train a ResNet-18, the implementation of which is available at https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py. In all base algorithms, we use a Nesterov momentum parameter of $0.9$ and set the weight decay factor as $10^{-4}$. For each base algorithm, we tune the (fast) learning rate from $\{0.01,0.025,0.05,0.1,0.15\}$ and linearly scale it with the number of workers (i.e., $32$). We found that, with a total batch size $4096$, the best learning rate for AR-SGD is $0.01$, for OSGP/SGP is $0.05$, and for Local SGD is $0.025$.

When applying SlowMo to these base algorithms, we fix $\alpha=1$ and $\tau=12$ and tune the value of $\beta$ from $\{0.4,0.5,0.6,0.7,0.8\}$. It turns out that for OSGP, SGP, and Local SGD, the best values of $\beta$ are all equal to $0.7$. More discussion on the effects of $\alpha$ and $\beta$ can be found in Section B.3.

A.3 ImageNet

For the ImageNet experiments, we use the same learning-rate, schedule, momentum, and weight decay as those suggested in Goyal et al. 2017 for SGD. In particular, we use ResNet50 (He et al. 2016) and train it for 90 epochs with a reference learning-rate of 0.1 with respect to a 256 sample batch, and scale this linearly with the batch-size. We decay the learning-rate by a factor of 10 at epochs 30, 60, 80. We use a Nesterov momentum parameter of 0.9. We use weight decay $10^{-4}$.

When using SlowMo, we set the slow learning rate to $\alpha=1$ and explore different numbers of inner steps, $\tau\in\{12,48\}$ and different slow momentum value $\beta\in\{0.5,0.6,0.7\}$ when the base optimizer is SGP/OSGP and $\beta=0.7$ when the base optimizer is LocalSGD. We also explore a larger set of $\tau$ values $\tau\in\{12,24,48,96\}$ in the ablation experiments.

A.4 WMT16 En-De

For the WMT16 En-De experiments, we follow the experimental protocol described in Ott et al. 2018. All experiments are based on the big transformer model (Vaswani et al. 2017) with 6 blocks in the encoder and decoder networks. For these experiments, the base optimizer used is Adam (Kingma & Ba 2015) with beta1 $=0.9$, beta2 $=0.98$, and $\epsilon=10^{-8}$ and trained for 25 epochs. We use the same learning rate schedule as Ott et al. 2018, i.e., the learning rate increases linearly for $4,000$ steps to $10^{-3}$, after which it is decayed proportionally to the inverse square root of the number of steps. We use label smoothing with weight 0.1 for the uniform prior distribution.

For SlowMo, we explore $\{0.5,1.0\}$ as the slow learning rate $\alpha$. We observe that $\alpha=1$ gives better performance and therefore report results for $\alpha=1$ unless stated otherwise. We explore different numbers of inner steps, $\tau\in\{12,48\}$ and different slow momentum value $\beta\in\{0.1,0.3,0.5,0.6,0.7\}$. We also explore a larger set of $\tau$ values, i.e. $\tau\in\{12,24,48,96,192\}$, in the ablation experiments.

Appendix B Additional Empirical Results

We show the Validation NLL on WMT’16 En-De in Table B.1, corresponding to the experiments in Table 1. We observe that SlowMo improves the validation NLL (along with BLEU score) of SGP and Local Adam.

B.2 Additional Training Curves

We present the training loss-versus-epochs curves in Figure B.1, corresponding to the validation curves in Figure 2. It can be observed that SlowMo substantially improves the convergence speed of SGP.

B.3 Effect of Slow Learning Rate α\alpha and Slow Momentum Factor β\beta

In this section we evaluate the impact of slow learning rate $\alpha$ and slow momentum $\beta$ hyperameters.

In Figure 2(a), we perform a parameter sweep over $\alpha$ and $\beta$ on CIFAR-10 dataset, using OSGP as the base algorithm of SlowMo. One can observe that when the value of $\beta$ is fixed, $\alpha=1$ always gives the highest validation accuracy; when the value of $\alpha$ is fixed, there is a best value of $\beta$ ranging from $0.4$ to $0.8$.

We further validate this claim on the WMT’16-En-De dataset. Figure 2(b) shows that $\alpha=1$ gives lower validation loss than $\alpha=0.5$ for fixed $\beta$ when using SGP or Local Adam as the base algorithms. When running SlowMo-Adam with $\beta>0.5$ and $\alpha=1.0$, or with $\beta>0.7$ and $\alpha=0.5$, the validation loss was substantially worse and so is not plotted here. Motivated by the above observations, we stick to fix $\alpha=1$ and fine-tune $\beta$ for SlowMo methods on all training tasks.

B.4 Base Optimizer Momentum Buffer Strategies

As described in Section 2, the base optimizer may have some associated buffers. SGD with momentum uses a momentum buffer, and Adam tracks estimates of the first and second moments of the gradient. The slow momentum buffer is updated every $\tau$ steps according to Eq.Equation 2. There are several strategies that can be used to update the base optimizer buffers at the outer iteration level (line 2 in Algorithm 1). Here, we explore three strategies: 1) reset the base optimizer buffers to zero ; 2) maintain the base optimizer buffers to their current values; 3) average the base optimizer buffers across workers, which requires additional communications. We evaluate the impact of these strategies on ImageNet and WMT’16 in Table B.2 and Table B.3.

On ImageNet, we observe that the different buffer strategies achieve similar training and validation performance. However, the averaging strategy comes at the cost of higher communication overhead (an additional call to AllReduce for each buffer averaged). Based on these results, we choose the reset strategy as the default in our experiments.

On WMT’16, we find that the reset buffer strategy underperforms both the maintain and average approaches. When using Adam as base optimizer, reseting the second moment to zeros hurts the optimization performance. This is not surprising since it is recognized that warming up the Adam buffer is important. Averaging buffers achieves the best results but comes at a significantly higher communication cost. We therefore select the maintain strategy as the default one when using Adam.

B.5 Standard Deviations on CIFAR-10

Since experiments for CIFAR-10 were ran for $5$ times with different random seeds, here we report the standard deviations on the validation accuracy in Table B.4, as a complementary to Table 1.

Appendix C Baseline Algorithms

In this section, we give a detailed description of each baseline algorithms used throughout the paper, provide theoretical justification on how to incorporate the update rules of D-PSGD, SGP and OSGP into the analysis of SlowMo, and also derive the analytic form of $V$ used in 3 for each method. A summary of the base optimizer update directions is given in Table C.1

Algorithms 2 and 3 present the pseudo-code of SGP and OSGP (Assran et al. 2019). To be consistent with the experimental results of the original paper, we also use local Nesterov momentum for each worker node. The communication topology among worker nodes is a time-varying directed exponential graph Assran et al. 2019. That is, if all nodes are ordered sequentially, then, according to their rank ($0,1,\dots,m-1$), each node periodically communicates with peers that are $2^{0},2^{1},\dots,2^{\lfloor\log_{2}(m-1)\rfloor}$ hops away. We let each node only send and receive a single message (i.e., communicate with 1 peer) at each iteration.

Note that although the implementation of SGP is with Nesterov momentum, the theoretical analysis in Assran et al. 2019 only considers the vanilla case where there is no momentum. Accordingly, the update rule can be written in a matrix form as

If we multiply a vector $\bm{1}/m$ on both sides of the update rule (9), we have

where $\bm{x}_{k}=\bm{X}_{k}\bm{1}/m$ denotes the average model across all worker nodes. Recall that in SlowMo, we rewrite the updates of the base algorithm at the $k$th steps of the $t$th outer iteration as

Similarly, for OSGP, one can repeat the above procedure again. But the definition of $\bm{X}_{k},\bm{Z}_{k}$ and $\nabla\bm{F}(\bm{Z}_{k})$ will change, in order to account for the delayed messages. In this case, we still have the update rule Eq. (10). But $\bm{x}_{k}$ is no longer the averaged model across all nodes. It also involves delayed model parameters. We refer the interested reader to Assran et al. 2019 for futher details.

C.2 D-PSGD

In the case of decentralized parallel SGD (D-PSGD), proposed in Lian et al. 2017, the update rule is quite similar to SGP. However, the communication topology among worker nodes is an undirected graph. Hence, the mixing matrix $\bm{P}_{k}$ is doubly-stochastic. Each node will exchange the model parameters with its neighbors. The update rule can be written as

As a result, using the same technique as (12)-(14), we have $V=\sigma^{2}/m$ for D-PSGD.

C.3 Local SGD

We further present the pseudo-code of Local SGD in Algorithm 4, and the pseudo-code of double-averaging momentum scheme in Algorithm 5.

Appendix D Proofs

To begin, recall that $\bm{d}_{t,k}=\frac{1}{m}\sum_{i=1}^{m}\bm{d}_{t,k}^{(i)}$, and that the local updates are

for $k=0,\dots,\tau-1$, followed by an averaging step to obtain $\bm{x}_{t,\tau}=\frac{1}{m}\sum_{i=1}^{m}\bm{x}_{t,\tau}^{(i)}$. Therefore, we can write the update rule of the base optimizer as

Combining this with Equations 2 and 3, we have

Let $\bm{y}_{t,0}=\bm{x}_{t,0}+\frac{\beta}{1-\beta}(\bm{x}_{t,0}-\bm{x}_{t-1,0}),\forall t$. Then by rearranging terms we get

Now, let us further extend the auxiliary sequence to all values of $k\neq 0$ as follows:

It is easy to show that $\bm{y}_{t,\tau}=\bm{y}_{t+1,0}$. In the sequel, we will analyze the convergence of sequence $\{\bm{y}_{t,k}\}$ instead of $\{\bm{x}_{t,k}\}$.

D.2 Preliminaries

In the table below, we list all notations used in this paper.

Throughout the theoretical analysis, we will repeatedly use the following facts:

Fact 1: $\left<{\mathbf{a}},\,{\mathbf{b}}\right>=\frac{1}{2}\left\|\mathbf{a}\right\|^{2}+\frac{1}{2}\left\|\mathbf{b}\right\|^{2}-\frac{1}{2}\left\|\mathbf{a}-\mathbf{b}\right\|^{2}$;

Fact 2: According to Young’s Inequality, for any $a>0$, we have

Fact 3: $\left\|\mathbf{a}+\mathbf{b}\right\|^{2}\leq 2\left\|\mathbf{a}\right\|^{2}+2\left\|\mathbf{b}\right\|^{2}$;

Fact 4: Suppose $\{a_{i}\}_{i=1}^{k}$ is a set of non-negative scalars and $s=\sum_{i=1}^{k}a_{i}$. Then according to Jensen’s Inequality, we have

D.3 General Treatment

Since each local objective $f_{i}$ is $L$-smooth, the function $f=\frac{1}{m}\sum_{i=1}^{m}f_{i}$ is also $L$-smooth. It follows that

where Equation 27 comes from Fact 4 Equation 24 , and $a>0$ is constant. For simplicity, we directly set $a=0.5$. Eqn. Equation 28 uses Fact 2 $\left<{\mathbf{a}},\,{\mathbf{b}}\right>=\frac{1}{2}\left\|\mathbf{a}\right\|^{2}+\frac{1}{2}\left\|\mathbf{b}\right\|^{2}-\frac{1}{2}\left\|\mathbf{a}-\mathbf{b}\right\|^{2}$. Furthermore, according to the definition of $\bm{y}_{t,k}$, it can be shown that

Substituting Equation 32 into Equation 29, it follows that

Moreover, for the second term in Equation 26, we have

Then, plugging Equations 33 and 34 into Equation 26,

where $\gamma_{\text{eff}}=\alpha\gamma/(1-\beta)$. Taking the total expectation,

Summing from $k=0$ to $k=\tau-1$, we have

where $N_{1}(t,k)$, $N_{2}(t,k)$, and $N_{3}(t)$ are as defined in (36). Summing from $t=0$ to $t=T-1$ and dividing both side by total iterations $K=\tau T$,

Now, we are going to further expand the expressions of the last two terms in Equation 39.

Using the fact $\left\|a+b\right\|^{2}\leq 2\left\|a\right\|^{2}+2\left\|b\right\|^{2}$, we have

where the last inequality comes from Fact 3. Then, taking the total expectation and summing over the $t$-th outer iteration,

where Equation 43 uses the following fact:

As a result, we end up with the following

where $K=\tau T$ denotes the total steps.

D.3.2 Bounding N3​(t)N_{3}(t)

From the update rule Equations 18, 2 and 3, we have

For the first term $T_{1}$, taking the total expectation, we get

where Equations 57 and 58 are derived using the same routine as Equations 46, 47 and 48. Similarly, for the second term $T_{2}$ in Equation 55, according to Fact 4,

Substituting Equations 59 and 62 back into Equation 55 and summing over the $t$-th outer iteration, we have

D.3.3 Final Results

Plugging Equations 51 and 68 back into Equation 39, one can obtain

where $C_{1}=1/2-\gamma_{\text{eff}}L-4\gamma_{\text{eff}}^{2}L^{2}\tau(\tau-1)(1-\beta-\alpha)^{2}/\alpha^{2}-8\gamma_{\text{eff}}^{2}L^{2}\tau^{2}\beta^{2}/(1-\beta)^{2}$. When the constants satisfy

Furthermore, since $\bm{y}_{0,0}=\bm{x}_{0,0}-\beta\bm{x}_{-1,0}/(1-\beta)=\bm{x}_{0,0}$ and $f(\bm{y}_{T,0})\geq f_{\text{inf}}$, the above upper bound can be simplified as

If we set $\gamma_{\text{eff}}=\sqrt{\frac{m}{K}}$, then

When $\gamma_{\text{eff}}=\frac{\sqrt{m}}{\sqrt{K}}$, the constraint can be rewritten as

Therefore, when $K\geq mL^{2}\left(1+\sqrt{3}\max\left\{\frac{3\tau(1-\beta-\alpha)}{\alpha},\frac{4\tau\beta}{1-\beta},1\right\}\right)$, the condition Equation 79 must be satisfied.

D.4 Special Case 1: Blockwise Model Update Filtering (BMUF)

In this case, the inner optimizer is Local-SGD. That is,

Since all worker nodes are averaged after every $\tau$ iterations, we have $\bm{x}_{t,0}^{(i)}=\bm{x}_{t,0},\forall i$. Besides, it is easy to validate that $V=\sigma^{2}/m$.

According to previous literature on the convergence of Local-SGD (Wang & Joshi 2018; Yu et al. 2019a), we can directly get the following upper bound.

When $\gamma L\tau\leq\frac{1}{6}$, we have $1/(1-12\gamma^{2}L^{2}\tau^{2})\geq 3/2$. It follows that

Substituting Equation 88 into Equation 73 and setting $\frac{\alpha}{1-\beta}\gamma L=\sqrt{m/K}$, we have

D.5 Special Case 2: Lookahead

It can be observed that when $\alpha=1$ or $\tau=1$, the above upper bound reduces to the case of vanilla mini-batch SGD. If we set $\alpha\gamma L=\sqrt{1/K}$, then we have

If the total iterations $K$ is sufficiently large, then the first term will dominate the convergence rate.