Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithms

Introduction

Stochastic gradient descent (SGD) is the core backbone of most state-of-the-art machine learning algorithms. Due to its widespread applicability, speeding-up SGD is arguably the single most impactful and transformative problem in machine learning. Classical SGD was designed to be run on a single computing node, and its error-convergence has been extensively analyzed and improved in optimization and learning theory (Dekel et al., 2012; Ghadimi & Lan, 2013). Due to the massive training data-sets and deep neural network architectures used today, running SGD at a single node can be prohibitively slow. This calls for distributed implementations of SGD, where gradient computation and aggregation is parallelized across multiple worker nodes. Although parallelism boosts the amount of data processed per iteration, it exposes SGD to unpredictable synchronization and communication delays stemming from variability in the computing infrastructure. This work presents a unified framework called Cooperative SGD to analyze communication-efficient distributed SGD algorithms that periodically average models trained locally at different computing nodes.

Limitations of Parameter Server Framework. A commonly used method to parallelize gradient computation and process more training data per iteration is the parameter server framework (Dean et al., 2012; Li et al., 2014; Cui et al., 2014). Each of the $m$ worker nodes computes the gradients of one mini-batch of data, and a parameter server aggregates these gradients and updates the model parameters. Synchronization delays in waiting for slow workers can be alleviated via asynchronous gradient aggregation (Recht et al., 2011; Cui et al., 2014; Gupta et al., 2016; Mitliagkas et al., 2016; Dutta et al., 2018). However it is difficult to eliminate communication delays since by design, parameter server framework requires gradients and model updates to be communicated between the parameter server and workers after every iteration.

Communication-Efficient SGD. To address the limitations of the parameter server framework, recent works proposed communication-efficient SGD variants that perform more computation at worker nodes. A natural idea is to allow workers to perform $\tau$ local updates to the model instead of just computing gradients, and then periodically averaging the local models (Moritz et al., 2015; Zhang et al., 2016; Povey et al., 2014; Su & Chen, 2015; Chaudhari et al., 2017; Smith et al., 2018; Lin et al., 2018). A similar approach (averaging after several epochs) is referred to Federated Averaging (FedAvg) (McMahan et al., 2016) in recent works and it is shown to work well even for non-i.i.d. data partitions. Although extensive empirical results have validated the effectiveness of periodic averaging, rigorous theoretical understanding of how its convergence depends on the number of local updates $\tau$ is quite limited (Zhou & Cong, 2017; Yu et al., 2018; Stich, 2018).

Instead of simply averaging the local models every $\tau$ iterations, Elastic-averaging SGD (EASGD) proposed in (Zhang et al., 2015) adds a proximal term to the objective function in order to allow some slack between the models – an idea that is drawn from the Alternating Direction Method of Multipliers (ADMM) (Boyd et al., 2011; Parikh & Boyd, 2014). Although the efficiency of EASGD and its asynchronous and periodic averaging variants has been empirically validated (Zhang et al., 2015; Chaudhari et al., 2017), its convergence analysis under general convex or non-convex objectives is an open problem. The original paper (Zhang et al., 2015) only gives an analysis of vanilla EASGD for quadratic objective functions.

A different approach to reducing communication is to perform decentralized training with sparse-connected network of worker nodes. Each node only synchronizes with its neighbors, thus reducing the communication overhead significantly. Decentralized averaging has a long history in the distributed and consensus optimization community (Tsitsiklis et al., 1986; Nedic & Ozdaglar, 2009; Duchi et al., 2012; Tsianos et al., 2012; Zeng & Yin, 2016; Yuan et al., 2016; Sirb & Ye, 2018; Bijral et al., 2017). Most of these works are for gradient descent or dual averaging methods rather than stochastic gradient descent (SGD), and they do not allow workers to make local updates. Recently, decentralized averaging was successfully applied to deep learning in (Jin et al., 2016; Jiang et al., 2017; Lian et al., 2017), which also provide convergence analyses for $1$ local update per worker. It is still unclear how decentralized training compares with periodic averaging ($\tau$ updates per worker).

Main Contributions. A common thread in all the communication-efficient SGD methods described above is that they allow worker nodes to perform local model-updates and limit the synchronization/consensus between the local models. Limiting model-synchronization reduces communication overhead, but it increases model discrepancies and can give an inferior error convergence performance. Communication-efficient SGD algorithms seek to strike the best trade-off between error-convergence and communication-efficiency.

In this paper, we propose a powerful framework called Cooperative SGD that enables us to obtain an integrated analysis and comparison of communication-efficient algorithms. Existing algorithms including periodic averaging SGD, Elastic Averaging SGD, decentralized SGD are special cases of cooperative SGD, and thus can be analyzed under one single umbrella. The main contributions of this paper are:

We present the first unified convergence analysis for the cooperative SGD class (Section 4) of algorithms that subsumes periodic, elastic and decentralized averaging. The theoretical results reveal how different communication-efficient strategies influence the error-convergence performance.

In particular, we provide the first analysis of elastic-averaging SGD for non-convex objectives, and use it to determine the best elasticity parameter $\alpha$ (Section 5.1) that achieves the lowest error floor at convergence.

We obtain a new analysis and tighter error bound for periodic averaging SGD by removing the uniformly bounded gradients assumption (Section 4.3). The analysis can be applied to FedAvg with i.i.d. data partitions as well.

Based on the unified analysis, we show the first in-depth comparison between periodic/elastic-averaging with decentralized training methods and design new communication-efficient SGD variants by combining existing strategies (see Section 6).

An alternative approach to communication-efficiency is gradient compression techniques (Wangni et al., 2017; Wen et al., 2017; Lin et al., 2017) that quantize the gradients computed by workers. Although interesting and important, this approach is beyond the scope of our paper; we focus on communication-efficiency via local updates at workers.

Preliminaries

In this section we present the update rules of existing communication-efficient SGD algorithms in terms of our notation that is used in the rest of the paper.

where $f(\cdot)$ is the loss function defined by the learning model. In the distributed setting, there are total $m$ worker machines that compute stochastic gradients in parallel. The updates can be written as

where $\eta$ is the learning rate, $\xi_{k}^{(i)}\subset\mathcal{S}$ are randomly sampled mini-batches, and $g(\mathbf{x};\xi)=\frac{1}{|\xi|}\sum_{s_{i}\in\xi}\nabla f(\mathbf{x};s_{i})$ denotes the stochastic gradient. For simplicity, we will use $g(\mathbf{x})$ instead of $g(\mathbf{x};\xi)$ in the rest of the paper.

Periodic Averaging SGD (PASGD). Local models are averaged after every $\tau$ iterations. Its update rule is

where $\mathbf{x}_{k}^{(i)}$ denotes the model parameters in the $i$-th worker and $\tau$ is defined as the communication period. The recently proposed federated learning framework (McMahan et al., 2016) also performs periodic averaging, but with non i.i.d. local datasets.

Elastic Averaging SGD (EASGD). Instead of performing a simple average of the local models, the elastic-averaging algorithm (EASGD) proposed in (Zhang et al., 2015) maintains an auxiliary variable $\mathbf{z}_{k}$. This variable is used as an anchor while updating the local models $\mathbf{x}_{k}^{(i)}$. The update rule of vanilla EASGDThe paper (Zhang et al., 2015) also presents periodic averaging and momentum variants of EASGD. However, only vanilla EASGD has been theoretically analyzed, and only for quadratic loss functions. is given by

where $\overline{\mathbf{x}}_{k}=\sum_{i=1}^{m}\mathbf{x}_{k}^{(i)}/m$. A larger value of the parameter $\alpha$ forces more consensus between the locally trained models and improves stability, but it may reduce the convergence speed – a phenomenon that is not yet well-understood.

Decentralized SGD (D-PSGD). The decentralized SGD algorithm D-PSGD (also referred as consensus-based distributed SGD), was proposed by (Jiang et al., 2017; Lian et al., 2017). Nodes perform local updates and average their models with neighboring nodes, where the network topology is captured by a mixing matrix $\mathbf{W}$. The update rule is

where $w_{ji}$ is the $(j,i)^{th}$ element of the mixing matrix $\mathbf{W}$, and it represents the contribution of node $j$ in the averaged model at node $i$.

The Cooperative SGD Framework

The Cooperative SGD algorithm is denoted by $\mathcal{A}(\tau,\mathbf{W},v)$, where $\tau$ is the number of local updates, $\mathbf{W}$ is the mixing matrix used for model averaging, and $v$ is the number of auxiliary variables. These parameters feature in the update rule as follows.

Gradients and Local Updates. In each iteration, the workers evaluate the gradient $g(\mathbf{x}_{k}^{(i)})$ for one mini-batch of data and update $\mathbf{x}_{k}^{(i)}$. The auxiliary variables are only updated by averaging a subset of the local models as described in point $3$ below. Thus, their gradients are zero, i.e., $g(\mathbf{z}_{k}^{(j)})=\mathbf{0},\forall j\in\{1,\dots,v\},\forall k$.

where the identity mixing matrix $\mathbf{I}_{(m+v)\times(m+v)}$ means that there is no inter-node communication during the $\tau$ local updates.

The update rule in terms of these matrices is

Instead of using update 10, one can use an alternative rule: $\mathbf{X}_{k+1}=\mathbf{X}_{k}\mathbf{W}_{k}-\eta\mathbf{G}_{k}$. The convergence analyses and insights in this paper can be extended to this update rule. We choose to study the update rule 10 for all existing algorithms (PASGD, EASGD, D-PSGD) since fully synchronous SGD corresponds to the special case $\mathbf{W}_{k}=\mathbf{J}$.

2 Existing Algorithms as Special Cases

We now show how existing communication-efficient algorithms are special cases of the general Cooperative SGD framework $\mathcal{A}(\tau,\mathbf{W},v)$.

Fully synchronous SGD $\Leftrightarrow\mathcal{A}(1,\mathbf{J},0)$. The local models are synchronized with all other workers after every iteration.

PASGD $\Leftrightarrow\mathcal{A}(\tau,\mathbf{J},0)$. The local models are synchronized with all other workers after every $\tau$ iterations.

EASGD $\Leftrightarrow\mathcal{A}(1,\mathbf{W}_{\alpha},1)$. In EASGD, there is one auxiliary variable. Besides, the mixing matrix is controlled by a hyper-parameter $\alpha$ as follows

One can easily validate that the updates defined in 10, 7 and 11 are equivalent to 4 and 5 when using the alternative update rule $\mathbf{X}_{k+1}=\mathbf{X}_{k}\mathbf{W}_{k}-\eta\mathbf{G}_{k}$.

D-PSGD $\Leftrightarrow\mathcal{A}(1,\mathbf{W},0)$. The mixing matrix $\mathbf{W}$ in D-PSGD is fixed as a sparse matrix. Only one local update before averaging is considered and there are no auxiliary variables. In addition to these special cases, the cooperative SGD framework allows us to design other communication-efficient SGD variants, as we describe in Section 6.

3 Communication Efficiency

The cooperative SGD framework improves the communication-efficiency of distributed SGD in three different ways, as described below. We illustrate these in Figure 1, which compares the execution timeline of cooperative SGD with fully synchronous SGD.

Periodic Averaging. The communication delay is amortized over $\tau$ iterations and is $\tau$ times smaller than fully synchronous SGD. Moreover, periodic averaging evens out random variations in workers’ computing time, and alleviates the synchronization delay in waiting for slow workers. Observe in Figure 1 that the idle time of workers is significantly reduced.

Non-blocking Execution. Since the auxiliary variables do not compute gradients, they remain the same while worker nodes conduct local updates, that is, $\mathbf{z}_{j\tau}=\mathbf{z}_{j\tau-1}=\dots=\mathbf{z}_{(j-1)\tau+1}$ for $j\geq 1$. Thus, the worker nodes only need $\mathbf{z}_{(j-1)\tau+1}$ before the model-averaging step from $\mathbf{x}_{j\tau}$ to $\mathbf{x}_{j\tau+1}$. So, the auxiliary variables can perform and broadcast model-updates while the workers perform the next set of local updates (see Figure 1), thus reducing synchronization delay.

Group Synchronization. Lastly, instead of synchronizing with all workers, a local model just needs to exchange information with its neighbors, where the network topology is captured by the mixing matrix $\mathbf{W}$. Thus, using a sparse mixing matrix $\mathbf{W}$ reduces the overall communication delay incurred per iteration.

Unified Convergence Analysis

In this section, we present the unified convergence analysis of algorithms in cooperative SGD framework and study how the $\tau$, $\mathbf{W}$, and $v$ affect the error-convergence.

The convergence analysis is conducted under the following assumptions, which are similar to previous works on the analysis of distributed SGD (Bottou et al., 2018):

(Smoothness): $\left\|\nabla F(\mathbf{x})-\nabla F(\mathbf{y})\right\|\leq L\left\|\mathbf{x}-\mathbf{y}\right\|$;

(Lower bounded): $F(\mathbf{x})\geq F_{\text{inf}}$;

(Mixing Matrix): $\mathbf{W}\mathbf{1}_{m+v}=\mathbf{1}_{m+v},\ \mathbf{W}^{\top}=\mathbf{W}$. Besides, the magnitudes of all eigenvalues except the largest one are strictly less than $1$: $\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m+v}(\mathbf{W})|\}<\lambda_{1}(\mathbf{W})=1$.

2 Update Rule for the Averaged Model

To facilitate the convergence analysis, we firstly introduce the quantities of interests. Multiplying $\mathbf{1}_{m+v}/(m+v)$ on both sides in (10), we get

where $\mathbf{W}_{k}$ disappears due to the special property from Assumption 5: $\mathbf{W}_{k}\mathbf{1}_{m+v}=\mathbf{1}_{m+v}$. Then, define the average model and effective learning rate as

Observe that the averaged model $\mathbf{u}_{k}$ is performing perturbed stochastic gradient descent. In the sequel, we will focus on the convergence of the averaged model $\mathbf{u}_{k}$, which is common practice in distributed optimization literature (Nedic & Ozdaglar, 2009; Duchi et al., 2012; Yuan et al., 2016).

Since the objective function $F(\mathbf{x})$ is non-convex, SGD may converge to a local minimum or saddle point. Thus, the expected gradient norm is used as an indicator of convergence (Lian et al., 2015; Zeng & Yin, 2016; Bottou et al., 2018). We say the algorithm achieves an $\epsilon$-suboptimal solution if:

This condition guarantees convergence of the algorithm to a stationary point.

3 Main Results

In deep learning, it is common to keep the learning rate as a constant and decay it only when the training procedure saturates. Thus, we present the analysis for fixed learning rate case and study the error floor at convergence.

For algorithm $\mathcal{A}(\tau,\mathbf{W},v)$, suppose the total number of iterations $K$ can be divided by the communication period $\tau$. Under Assumptions 1–5 (with $\beta=0$ Constant $\beta$ in Assumption 4 only influences the constraint on the learning rate (16) and will not appear in the expression of gradient norm upper bound (17). In order to get neater results, $\beta$ is set as 0 in the main paper. In the Appendix, we provide the proof for arbitrary $\beta$.), if the learning rate satisfies

where $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m+v}(\mathbf{W})|\}$, and all local models are initialized at a same point $\mathbf{u}_{1}$, then the average-squared gradient norm after $K$ iterations is bounded as follows

where $\mathbf{u}_{k},\eta_{\text{eff}}$ are defined in (13).

All proofs are provided in the Appendix. The error floor at convergence is given by (18).

Error decomposition. It is worth noting that the upper bound (17) is decomposed into two parts. The first two terms are same as the optimization error bound in fully synchronous SGD (Bottou et al., 2018). The last term is network error, resulted from performing local updates and reducing inter-worker communication. It directly increases the error floor at convergence and is a measure of local models’ discrepancies. When all local models are fully synchronized at every iterations ($\tau=1,\zeta=0,v=0$), then the network error becomes zero.

Dependence on $\tau,\mathbf{W}$. Theorem 1 states that the error floor at convergence (18) is determined by the communication period $\tau$ and the second largest absolute eigenvalue $\zeta$ of the mixing matrix. In particular, the bound will monotonically increase along with $\tau$ and $\zeta$. The definition of $\zeta$ is common in random walks on graphs and reflects the mixing rates of different variables. When there is no communication among local workers, then $\mathbf{W}=\mathbf{I}_{m+v}$ and $\zeta=1$; When local models are fully synchronized, then $\mathbf{W}=\mathbf{J}_{m+v}$ and $\zeta=0$. Typically, a sparser matrix means a larger value of $\zeta$.

Besides, since the network error bound is linear to $\tau$ but proportional to $(1+\zeta^{2})/(1-\zeta^{2})$, as shown in Figure 2, it is more sensitive to the changes in communication period. In Figure 3, we evaluate various hyper-parameter settings for training VGGNet (Simonyan & Zisserman, 2014) for classification of the CIFAR10 dataset (Krizhevsky, 2009). As suggested by 18, the empirical results show that a higher network error (larger $\tau$ or larger $\zeta$) leads to a higher error floor at convergence.

Dependence on $v$. Note that the effective learning rate (13) is determined by the number of auxiliary variables. Using more auxiliary variables results in smaller effective learning rate, since they update only through model averaging. Consequently, it may slow down the optimization progress (increase the first term in (17)) while enable smaller error floor at convergence (reduce the second term in (17)).

Finite horizon result. If $K$ is decided preemptively, then with a proper learning rate, we obtain the following bound. A similar approach also appears in (Ghadimi & Lan, 2013; Lian et al., 2017; Yu et al., 2018; Bernstein et al., 2018).

For algorithm $\mathcal{A}(\tau,\mathbf{W},v)$, under Assumption 1–5, if the learning rate is $\eta=\frac{m+v}{Lm}\sqrt{\frac{m}{K}}$, the average-squared gradient norm after $K$ iterations is bounded by

if the total iterations $K$ is sufficiently large: $K\geq 10m[(1+\frac{v}{m})\frac{\tau}{1-\zeta}]^{2}$. Furthermore, if $K\geq(m+v)^{2}m[(1+\frac{v}{m})\frac{\tau}{1-\zeta}]^{2}$, then the average-squared gradient norm will be bounded by $2[L(F(\mathbf{x}_{1})-F_{\text{inf}})+\sigma^{2}]/\sqrt{mK}$.

By directly setting $\mathbf{W}=\mathbf{J}$ (i.e., $\zeta=0$) and $v=0$ in Corollary 1, one can obtain the result for PASGD. Comparing to previous results on non-convex objectives (Yu et al., 2018), we remove the uniformly bounded gradients assumption. To obtain an error bound in the form $C/\sqrt{mK}$ for some constant $C$, our result shows $\tau$ can be large up to $\sqrt{K/m^{3}}$ instead of $(K/m^{3})^{1/4}$ (Yu et al., 2018).

Novel Analyses of Existing Algorithms

Using the unified analysis of cooperative SGD presented in Theorem 1, one can directly derive novel analyses of EASGD, PASGD and D-PSGD. The general framework also provides new insights such as the best choice of parameter $\alpha$ in EASGD (see Lemma 1).

Recall that EASGD uses hyper-parameter $\alpha$ to control the eigenvalues of mixing matrix. For $\mathbf{W}_{\alpha}$ defined in (11), the second largest eigenvalue magnitude is

In order to let $\mathbf{W}_{\alpha}$ satisfy the conditions in Assumption 5, it is required that $\zeta<1$, namely $0\leq\alpha<2/(m+1)$. This condition suggests that $\alpha$ can be selected in a broader range than the original paper (Zhang et al., 2015) suggested ($0\leq\alpha<1/m$).

Intuitively, a larger $\alpha$ forces more consensus between the locally trained models and improves stability. However, from equation 20, we observe that there exists an optimal $\alpha$ that minimizes the value of $\zeta$.

If $\alpha=2/(m+2)$, then the second largest absolute eigenvalue of $\mathbf{W}_{\alpha}$, given in 20, achieves the minimal value $m/(m+2)$.

Accordingly, by choosing the best $\alpha$, the error floor at convergence can also be minimized. To be specific, we have the following theorem.

When $\alpha$ is set to $2/(m+2)$ as suggested by Lemma 1, the error of EASGD can be bounded as follows:

where $\mathbf{u}_{k}$ and $\eta_{\text{eff}}$ are defined in (13).

To the best of our knowledge, this theorem is the first convergence result for EASGD with general objectives and also the first theoretical justification for the best choice of $\alpha$. By setting $\eta_{\text{eff}}=\frac{1}{L}\sqrt{\frac{m}{K}}$, one can also obtain a finite horizon result as Corollary 1.

𝑚20.22/(m+2)=0.2, which performs better than the empirical choice $\alpha=0.9/m=0.1125$ suggested in (Zhang et al., 2015). The best choice of $\alpha$ yields the lowest training loss and the least discrepancies between workers and auxiliary variable. Empirical validation. As shown in Figure 4, the best choice $\alpha=2/(m+2)=0.2$ yields fastest convergence and least discrepancies between workers and the auxiliary variable. When $\alpha$ is greater than $2/(m+1)\approx 0.2222$, we observe the algorithm cannot converge. Furthermore, in Figure 4(c), we show the benefit of non-blocking execution. By overlapping the broadcast of auxiliary variable and workers computation, it directly reduces about $67\%$ training time.

2 PASGD 𝒜​(τ,𝐉,0)𝒜𝜏𝐉0\mathcal{A}(\tau,\mathbf{J},0) Vs. D-PSGD 𝒜​(1,𝐖,0)𝒜1𝐖0\mathcal{A}(1,\mathbf{W},0)

The general framework enables easy comparisons between different communication reduction strategies. Here, we compare periodic communication and group synchronization strategies. Note that when PASGD $\mathcal{A}(\tau,\mathbf{J},0)$ and D-PSGD $\mathcal{A}(1,\mathbf{W},0)$ have the same error floor at convergence, we have

Equation (22) provides a threshold for $\zeta$. As long as $\zeta\leq\zeta_{\tau}$, D-PSGD $\mathcal{A}(1,\mathbf{W},0)$ would perform better than PASGD $\mathcal{A}(\tau,\mathbf{J},0)$ in terms of the worst-case final error at convergence. Along with the increase of $\tau$, the value of threshold $\zeta_{\tau}$ rapidly converges to 1. Therefore, when $\tau$ becomes large, D-PSGD has a lower error floor in a very broad range of $\zeta$.

As for communication efficiency, the benefit of group synchronization relies on the number of workers. It at most reduces the communication overhead by $m$ times, since at least one connection should be preserved for each worker. As the mixing matrix affects the communication delay implicitly, it is not trivial to design a good mixing matrix that not only has small eigenvalues but also enables efficient implementation. On the contrary, periodic averaging has higher flexibility without such limitations. If we set $\tau\geq m$, then PASGD always has shorter training time than D-PSGD.

Designing New Communication-Efficient SGD Algorithms

As shown in Section 5, the Cooperative SGD framework enables us to analyze and compare existing communication-efficient SGD algorithms such as PASGD, EASGD and D-PSGD. The Cooperative SGD framework can also be used to design new algorithms that combine the communication-efficiency strategies adopted by these algorithms.

From Section 5.2 we see that D-PSGD has superior convergence performance, while PASGD can easily control the communication delay and provide higher throughput. We propose using a combination of these called decentralized periodic averaging SGD $\mathcal{A}(\tau,\mathbf{W},0)$ with carefully chosen $\tau$ and $\mathbf{W}$. For a small number of well-connected workers, larger $\tau$ is more preferable. For a large number of workers, using a sparse mixing matrix $\mathbf{W}$ and small $\tau$ gives better convergence. For a fixed topology worker network where $\mathbf{W}$ is prescribed, increasing the communication period can be an effective way to speedup the decentralized training. In Figure 5, we implemented the algorithm with $7$ worker nodes and evaluated it on CIFAR10 dataset. The observation is decentralized periodic averaging with $\tau=15,\zeta=0.75$ achieves significant speedup over the pure D-PSGD algorithm as well as similar throughput as pure PASGD with a larger communication period $\tau=50$.

2 Generalized Elastic Averaging

In generalized elastic averaging $\mathcal{A}(1,\mathbf{W}^{\prime},1)$, we modify decentralized SGD with mixing matrix $\mathbf{W}$ by adding an auxiliary variable (with elasticity parameter $\alpha$) stored at a new node that is connected to all $m$ worker nodes. Recall that a sparse mixing matrix $\mathbf{W}$ can reduce communication delay, but it may have large $\zeta$ that leads to inferior convergence. Introducing the auxiliary variable results in the mixing matrix $\mathbf{W}^{\prime}$ shown in (23) below. The second largest eigenvalue of this matrix is $(1-\alpha)$ lower than $\zeta$ as shown by Lemma 2.

Suppose there is a $m$-dimension symmetric matrix $\mathbf{W}$ such that $\mathbf{W}\mathbf{1}=\mathbf{1}$, and its eigen-values satisfy $-1\leq\lambda_{m}(\mathbf{W})\leq\cdots\leq\lambda_{1}(\mathbf{W})\leq 1$. Let $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m}(\mathbf{W})|\}$. Then, for matrix $\mathbf{W}^{\prime}$ which is defined as:

Setting $\alpha=\frac{1+\zeta}{m+1+\zeta}$ yields the minimum $\zeta^{\prime}=\frac{m\zeta}{m+1+\zeta}$.

The proof is given in the Appendix. Lemma 2 implies that by setting $\alpha=\frac{1+\zeta}{m+1+\zeta}$, the new algorithm $\mathcal{A}(1,\mathbf{W}^{\prime},1)$ gives a lower error bound at convergence as compared to D-PSGD $\mathcal{A}(1,\mathbf{W},0)$ as $\zeta^{\prime}<\zeta$. Furthermore, since the updates and broadcast of the auxiliary variable can overlap with the local computation at workers (as explained in Section 3.3), we do not expect an increase in the training time. Thus, adding an auxiliary variable is a highly effective method to increase the consensus between loosely connected workers.

3 Hierarchical Averaging

Based on the analysis of Cooperative SGD, we believe that a hierarchical averaging framework will aptly capture the benefits of all the communication-efficiency strategies discussed in this paper. In particular, consider that workers are divided into groups that cannot directly communicate with each other, as shown in Figure 6(a). Local models in each group will be averaged via an auxiliary node. Inter-auxiliary node communication can occur concurrently with local updates at workers, as illustrated in Figure 6(b). Our unified convergence analysis can be applied to this hierarchical averaging model and ongoing research includes finding the node structure that gives the best convergence.

Concluding Remarks

We propose a communication-efficient SGD framework called Cooperative SGD that combines the periodic, decentralized, and elastic model-averaging strategies to reduce inter-node communication via local updates at worker nodes. By analyzing cooperative SGD for general non-convex objectives, we provide strong convergence guarantees for existing communication-efficient SGD variants, and to the best of our knowledge, the first general analysis of elastic-averaging SGD. Furthermore, the cooperative SGD framework greatly enlarges the design space of communication-efficient SGD algorithms. We present some promising new ideas such as decentralized periodic averaging, generalized elastic-averaging and hierarchical averaging that can strike a good trade-off between convergence speed and communication efficiency. However, further exploration of the communication-efficient SGD design space and analyses of new variants is ripe for future investigation.

Acknowledgments

The authors thank Anit Kumar Sahu for his suggestions and feedback. This work was partially supported by the CMU Dean’s fellowship and an IBM Faculty Award. The experiments were conducted on the ORCA cluster provided by the Parallel Data Lab at CMU, and on Amazon AWS (supported by an AWS credit grant).

References

Appendix A Convergence of PASGD and D-PSGD

By directly setting $\mathbf{W}=\mathbf{J}$ (i.e., $\zeta=0$) and $v=0$ in Theorem 1, one can obtain the convergence guarantee for PASGD. Comparing to previous results on non-convex objectives (Zhou & Cong, 2017; Yu et al., 2018), our result removes the uniformly bounded gradients assumption and provides a tighter upper bound.

For $\mathcal{A}(\tau,\mathbf{J},0)$, under the same assumptions as Theorem 1, if the learning rate satisfies $\eta L+\eta^{2}L^{2}\tau(\tau-1)\leq 1$, then we have

The notable insight provided by Corollary 2 is there exists a trade-off between the error-convergence and communication-efficiency. While a larger communication period leads to higher error at convergence, it directly reduces the communication delay by $\tau$ times and enables higher throughput. The primary advantage of PASGD is that one can easily change the communication period and find the best one that has the fastest convergence rate with respect to wall-clock time. The best value of $\tau$ should depend on the network bandwidth/latency and vary in different environments.

Empirical validation. In Figure 7, we show this trade-off in PASGD with different learning rate choices. One can see that even though PASGD with $\tau=100$ finishes the training first, it has the highest loss after the same number of iterations. Comparing Figure 7 (a) and (b), observe that the small learning rate reduces the gap between different communication periods. This phenomenon has already been discussed in Theorem 1: small learning rate can alleviate the relative effect of the network error term. Besides, for completeness, we present the test accuracy of PASGD in Figure 7 (c). The interesting observation is that PASGD with large communication period has better generalization performance than fully synchronous SGD.

Similarly, setting $\tau=1$ and $v=0$ in Theorem 1, we get the convergence guarantee for D-PSGD, which is consistent to (Lian et al., 2017).

For $\mathcal{A}(1,\mathbf{W},0)$, under the same assumptions as Theorem 1, if the learning rate satisfies

where $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m}(\mathbf{W})|\}$, then we have

We implemented a ring-connected D-PSGD with $4$ workers. As shown in Figure 7, the number of workers is so few that the communication reduction effect is quite limited (saving about $10\%$ training time while PASGD ($\tau=10$) reduces about $67\%$ time).

Appendix B Proof Preliminaries

Similar to $\mathbf{X}_{k}$ and $\mathbf{G}_{k}$, we stack all full batch gradients in a $d\times(m+v)$ dimension matrix:

Accordingly, the Frobenius norm of full batch gradients is $\left\|\nabla F(\mathbf{X}_{k})\right\|_{\tiny{\mbox{F}}}^{2}=\sum_{i=1}^{m}\left\|\nabla F(\mathbf{x}_{k}^{(i)})\right\|^{2}$. In order to facilitate reading, the definitions of matrix Frobenius norm and operator norm are also provided here.

The Frobenius norm defined for $\mathbf{A}\in M_{n}$ by

The operator norm defined for $\mathbf{A}\in M_{n}$ by

All notations used in the proof are listed below.

Appendix C A Supporting Lemma for Theorem 1

Before providing the proof of Theorem 1, we prefer to first present an important lemma that describes the basic intuition for the convergence of cooperative SGD: the discrepancies of local models have a negative impact on the convergence. The proof of Theorem 1 will be built upon this lemma.

For algorithm $\mathcal{A}(\tau,\mathbf{W},v)$, under Assumption 1–5, if the learning rate satisfies $\eta_{\text{eff}}L(1+\beta/m)\leq 1$ and all local model parameters are initialized at the same point $\mathbf{x}_{1}$, then the average-squared gradient after $K$ iterations is bounded as follows

where $\mathbf{u}_{k},\eta_{\text{eff}}$ are defined in (13) and both $\mathbf{I}$ and $\mathbf{J}$ are $(m+v)\times(m+v)$ matrices.

Under Assumption 3 and 4, we have the following variance bound for the averaged stochastic gradient:

According to the definition of $\mathcal{G}_{k},\mathcal{H}_{k}$ 29, we have

where equation (38) is due to $\{\xi_{k}^{(i)}\}$ are independent random variables. Now, directly applying Assumption 3 and 4 to (38), one can observe that all cross terms are zero. Then, we have

Under Assumption 3, the expected inner product between stochastic gradient and full batch gradient can be expanded as

where equation (44) comes from $2\mathbf{a}^{\top}\mathbf{b}=\left\|\mathbf{a}\right\|^{2}+\left\|\mathbf{b}\right\|^{2}-\left\|\mathbf{a}-\mathbf{b}\right\|^{2}$. ∎

Under Assumption 3 and 4, the squared norm of stochastic gradient can be bounded as

Since $\mathbf{E}_{k}[\mathcal{G}_{k}]=\mathcal{H}_{k}$, then we have

where (48) follows Lemma 4 and (49) comes from the convexity of vector norm and Jensen’s inequality:

C.1.2 Proof of Lemma 3

According to Lipschitz continuous gradient assumption, we have

After minor rearranging and according to the definition of Frobenius norm, it is easy to show

Taking the total expectation and averaging over all iterates, we have

If the effective learning rate satisfies $\eta_{\text{eff}}L(\beta/m+1)\leq 1$, then

Recalling the definition $\mathbf{u}_{k}=\mathbf{X}_{k}\mathbf{1}_{m+v}/(m+v)$ and adding a positive term to the RHS, one can get

where $\mathbf{I},\mathbf{J}$ are $(m+v)\times(m+v)$ matrices. Plugging the inequality (60) into (57), we complete the proof.

Appendix D Proof of Theorem 1: Convergence of Cooperative SGD

Assume the rows of matirx $\mathbf{A}$ are denoted by $\mathbf{a}_{1}^{\top},\dots,\mathbf{a}_{d}^{\top}$ and $\mathcal{I}=\{i\in[1,d]:\left\|\mathbf{a}_{i}\right\|\neq 0\}$. Then, we have

where the last inequality follows the definition of matrix operator norm. ∎

Suppose there is a $m\times m$ matrix $\mathbf{W}$ that satisfies Assumption 5. Then

where $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m}(\mathbf{W})|\}$.

Since $\mathbf{W}$ is a real symmetric matrix, then it can be decomposed as $\mathbf{W}=\mathbf{Q}\mathbf{\Lambda}\mathbf{Q}^{\top}$, where $\mathbf{Q}$ is an orthogonal matrix and $\mathbf{\Lambda}=\operatorname{diag}\{\lambda_{1}(\mathbf{W}),\lambda_{2}(\mathbf{W}),\dots,\lambda_{m}(\mathbf{W})\}$. In particular, since the largest eigenvalue of $\mathbf{W}$ is $1$ and $\mathbf{W}\mathbf{1}=\mathbf{1}$, the corresponding eigenvector (i.e., the first column of $\mathbf{Q}$) is $\frac{\mathbf{1}}{\sqrt{m}}$. Similarly, matrix $\mathbf{J}$ can be decomposed as $\mathbf{Q}\mathbf{\Lambda}_{0}\mathbf{Q}^{\top}$ where $\mathbf{\Lambda}_{0}=\operatorname{diag}\{1,0,\dots,0\}$. Then, we have

According to the definition of matrix operator norm,

Since $\mathbf{W}^{2j}-\mathbf{J}=\mathbf{Q}\left(\mathbf{\Lambda}^{2j}-\mathbf{\Lambda}_{0}\right)\mathbf{Q}^{\top}$, the maximal eigenvalue will be $\max\{0,\lambda_{2}(\mathbf{W})^{2j},\dots,\lambda_{m}(\mathbf{W})^{2j}\}=\zeta^{2j}$. As a consequence, we have $\left\|\mathbf{W}^{j}-\mathbf{J}\right\|_{\tiny{\mbox{op}}}=\sqrt{\lambda_{\text{max}}(\mathbf{W}^{2j}-\mathbf{J})}=\zeta^{j}$.

D.2 Proof of Theorem 1

Recall the intermediate result (56) in the proof of Lemma 3:

Our goal is to provide an upper bound for the network error term $\frac{L^{2}}{Km}\sum_{k=1}^{K}\left\|\mathbf{X}_{k}(\mathbf{I}-\mathbf{J})\right\|_{\tiny{\mbox{F}}}^{2}$. First of all, let us derive a specific expression for $\mathbf{X}_{k}(\mathbf{I}-\mathbf{J})$.

According to the update rule (10) in Section 3, one can observe that

where (75) follows the special property of doubly stochastic matrix: $\mathbf{W}_{k-1}\mathbf{J}=\mathbf{J}\mathbf{W}_{k-1}=\mathbf{J}$ and hence $(\mathbf{I}-\mathbf{J})\mathbf{W}_{k-1}=\mathbf{W}_{k-1}(\mathbf{I}-\mathbf{J})$. Then, expanding the expression of $\mathbf{X}_{k-1}$, we have

Repeating the same procedure for $\mathbf{X}_{k-2},\mathbf{X}_{k-3},\dots,\mathbf{X}_{2}$, finally we get

where $\mathbf{\Phi}_{s,k-1}=\prod_{l=s}^{k-1}\mathbf{W}_{l}$. Since all optimization variables are initialized at the same point $\mathbf{X}_{1}(\mathbf{I}-\mathbf{J})=0$, the squared norm of the network error term can be directly written as

Then, let us take a closer look at the expression of $\mathbf{\Phi}_{s,k-1}$. Without loss of generality, assume $k=j\tau+i$, where $j$ denotes the index of communication rounds and $i$ denotes the index of local updates. As a result, matrix $\mathbf{\Phi}_{s,k-1}$ can be expressed as follows:

For the ease of writing, define accumulated stochastic gradient within one local update period as $\mathbf{Y}_{r}=\sum_{s=r\tau+1}^{(r+1)\tau}\mathbf{G}_{s}$ for $0\leq r<j$ and $\mathbf{Y}_{j}=\sum_{s=j\tau+1}^{j\tau+i-1}\mathbf{G}_{s}$. Similarly, define accumulated full batch gradient $\mathbf{Q}_{r}=\sum_{s=r\tau+1}^{(r+1)\tau}\nabla F(\mathbf{X}_{s})$ for $0\leq r<j$ and $\mathbf{Q}_{j}=\sum_{s=j\tau+1}^{j\tau+i-1}\nabla F(\mathbf{X}_{s})$. Accordingly, we have

Note that the network error term can be decomposed into two parts:

where 88 follows $\left\|a+b\right\|^{2}\leq 2\left\|a\right\|^{2}+2\left\|b\right\|^{2}$. Next, we are going to separately provide bounds for $T_{1}$ and $T_{2}$. Recall that we are interested in the average of all iterates $\frac{L^{2}}{Km}\sum_{k=1}^{K}\left\|\mathbf{X}_{k}(\mathbf{I}-\mathbf{J})\right\|_{\tiny{\mbox{F}}}^{2}$. Accordingly, we will also derive the bounds for $\frac{L^{2}}{Km}\sum_{k=1}^{K}T_{1}$ and $\frac{L^{2}}{Km}\sum_{k=1}^{K}T_{2}$.

where (90) follows Lemma 7, (91) comes from Lemma 9. Recall that $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m+v}(\mathbf{W})|\}$. Then for any $0\leq r<j$,

Now we show that the cross terms are zero. For any $s<l$, according to Assumption 4, one can obtain

where (100) is according to Assumption 4. Using the same technique, one can obtain that

Substituting 101 and 102 back into (92), we have

where (104) follows the summation formula of power series:

Next, summing over all iterates in the $j$-th local update period (from $i=1$ to $i=\tau$):

Then, summing over all periods from $j=0$ to $j=K/\tau-1$, where $K$ is the total iterations:

Expanding the summation in (109), we have

For the second term in (88), since $\left\|\mathbf{A}\right\|_{\tiny{\mbox{F}}}^{2}=\operatorname{Tr}(\mathbf{A}^{\top}\mathbf{A})$, we have

According to Lemma 8, the trace can be bounded as:

where (116) follows Lemma 7 and (117) is because of $2ab\leq a^{2}+b^{2}$. Then, it follows that

where (119) uses the fact that indices $n$ and $l$ are symmetric and (122) is according to the summation formula of power series:

where (127) follows the convexity of Frobenius norm and Jensen’s inequality. Next, summing over all iterates in the $j$-th period, we can get

D.2.4 Final result.

According to 88, 113 and 134, the network error can be bounded as

Substituting the expression of network error back to inequality (56), we obtain

where $\eta_{\text{eff}}=m\eta/(m+v)$ and $\zeta=\max\{|\lambda_{2}(\mathbf{W})|,|\lambda_{m+v}(\mathbf{W})|\}$. Setting $\beta=0$, the condition on learning rate (138) can be further simplified as follows:

Appendix E Proof of Corollary 1 (Finite Horizon Result)

Directly substituting $\eta=\frac{m+v}{Lm}\sqrt{\frac{m}{K}}$ into (139), we have

Note that the learning rate should satisfy the condition in (143). That is, the total iterations should satisfy:

When $K$ is sufficiently large, the first term can be arbitrarily small. In particular, when $K>4m$, the first term will be smaller than $1/2$. Then, it is enough to show the second term is smaller than $1/2$ as well.

Here, we complete the proof of the first part. Furthermore, when the communication period and total iterations satisfy

then the last term in (144) is smaller than the second term. As a result, we have

In order to get a lower bound on $K$ from (148), it is enough to show

Once $m+v\geq\sqrt{10}\approx 3.1$, (153) is more strict than (147).

Appendix F Proof of Lemma 1 and Theorem 2: Best Choice of α𝛼\alpha in EASGD

Recall that in EASGD, $\zeta=\max\{|1-\alpha|,|1-(m+1)\alpha|\}$. It is straightforward to show that

When $\alpha=\frac{2}{m+2}$, one can get the minimal value of $\zeta$, which equals to $1-\alpha=(m+1)\alpha-1=\frac{m}{m+2}$. Then, substituting $\zeta=\frac{m}{m+2},\tau=1,v=0$ into Theorem 1, we complete the proof of Theorem 2.

Appendix G Proof of Lemma 2: Generalized Elastic Averaging

Lemma 2 is built upon a known result about the eigenvalues of block matrices.

Let $\mathbf{A}$ be a symmetric $m\times m$ matrix with eigenvalues $\lambda_{1},\lambda_{2},\dots,\lambda_{m}$, let $\mathbf{u},\left\|\mathbf{u}\right\|=1$, be a unit eigenvector corresponding to $\lambda_{1}$; let $\mathbf{B}$ be a symmetric $n\times n$ matrix with eigenvalues $\beta_{1},\beta_{2},\dots,\beta_{n}$, let $\mathbf{v},\left\|\mathbf{v}\right\|=1$, be a unit eigenvector corresponding to $\beta_{1}$. Then for any $\rho$, the matrix

has eigenvalues $\lambda_{2},\dots,\lambda_{m},\beta_{2},\beta_{n},\gamma_{1},\gamma_{2}$, where $\gamma_{1},\gamma_{2}$ are eigenvalues of the matrix:

In our case, recall the definition of $\mathbf{W}^{\prime}$:

In order to apply Lemma 10, let us set $\mathbf{A}=(1-\alpha)\mathbf{W}$. Accordingly, the eigenvalues of $\mathbf{A}$ are $1-\alpha,(1-\alpha)\lambda_{2},\dots,(1-\alpha)\lambda_{m}$. The eigenvector corresponding to $1-\alpha$ is $\frac{\mathbf{1}}{\sqrt{m}}$. Moreover, set $B=1-m\alpha$. Then, it has only one eigenvalue $1-m\alpha$ and the corresponding eigenvector is scalar $1$. Substituting $\mathbf{A},B$ into $\mathbf{W}^{\prime}$, we have

According to Lemma 10, the eigenvalues of $\mathbf{W}^{\prime}$ are $(1-\alpha)\lambda_{2},\dots,(1-\alpha)\lambda_{m},\gamma_{1},\gamma_{2}$, where $\gamma_{1},\gamma_{2}$ are eigenvalues of the matrix:

The above equation yields $\gamma_{1}=1,\gamma_{2}=1-(m+1)\alpha$.

Finally, we have $\zeta^{\prime}=\max\{|(1-\alpha)\lambda_{2}|,|(1-\alpha)\lambda_{m}|,|1-(m+1)\alpha|\}=\max\{(1-\alpha)\zeta,|1-(m+1)\alpha|\}$. As a consequence, when $(1-\alpha)\zeta=(m+1)\alpha-1$, i.e., $\alpha=\frac{1+\zeta}{m+1+\zeta}$, the value of $\zeta^{\prime}$ is minimized.