GT-STORM: Taming Sample, Communication, and Memory Complexities in Decentralized Non-Convex Learning

As mentioned in Section 1, in decentralized learning, there are a set of geographically distributed computing nodes forming a network. In this paper, we represent such a networked by an undirected connected network $\mathcal{G}=(\mathcal{N},\mathcal{L})$, where $\mathcal{N}$ and $\mathcal{L}$ are the sets of nodes and edges, respectively, with $|\mathcal{N}|=m$. Each node can communicate with their neighbors via the edges in $\mathcal{L}$. The goal of decentralized learning is to use the nodes to distributively and collaboratively solve a network-wide optimization problem as follows:

where each local objective function $f_{i}(\mathbf{x})\triangleq\mathbb{E}_{\zeta\sim\mathcal{D}_{i}}f_{i}(\mathbf{x};\zeta)$ is only observable to node $i$ and not necessarily convex. Here, $\mathcal{D}_{i}$ represents the distribution of the dataset at node $i$, and $f_{i}(\mathbf{x};\zeta)$ represents a loss function that evaluates the discrepancy between the learning model’s output and the ground truth of a training sample $\zeta$. To solve Problem (1) in a decentralized fashion, a common approach is to rewrite Problem (1) in the following equivalent form:

where $\mathbf{x}\triangleq[\mathbf{x}_{1}^{\top},\cdots,\mathbf{x}_{m}^{\top}]^{\top}$ and $\mathbf{x}_{i}$ is an introduced local copy at node $i$. In Problem (2), the constraints ensure that the local copies at all nodes are equal to each other, hence the term “consensus.” Thus, Problems (1) and (2) share the same solutions. The main goal of network consensus optimization is to design an algorithm to attain an $\epsilon^{2}$-stationary point $\mathbf{x}$ defined as follows:

where $\mathbf{\bar{x}}\triangleq\frac{1}{m}\sum_{i=1}^{m}\mathbf{x}_{i}$ denotes the global average across all nodes. Different from the traditional $\epsilon^{2}$-stationary point in centralized optimization problems, the metric in Eq. (3) has two terms: the first term is the gradient magnitude for the (non-convex) global objective and the second term is the average consensus error of all local copies. To date, many decentralized algorithms have been developed to compute the $\epsilon^{2}$-stationary point (see Section 2.2). However, most of these algorithms suffer limitations in sample, communication, and memory complexities. In what follows, we formally state the definitions of sample, communication, and memory complexities used in the literature (see, e.g., (sun2019improving, )):

To make sense of these three complexity metrics into perspective, consider the standard centralized gradient descent (GD) method as an example. Note that the GD algorithm has an $O(1/T)$ convergence rate for non-convex optimization, which suggests $O(\epsilon^{-2})$ communication complexity. Also, it takes a full gradient evaluation in each iteration, i.e., $O(n)$ per-iteration sample complexity, where $n$ is the total number of samples. This implies $O(n\epsilon^{-2})$ sample complexity to converge to an $\epsilon^{2}$-stationary point. Hence, the sample complexity of GD is high if the dataset size $n$ is large.

In contrast, consider the classical stochastic gradient descent (SGD) algorithm that is widely used in machine learning. The basic idea of SGD is to lower the gradient evaluation cost by using only a mini-batch of samples in each iteration. However, due to the sample randomness in mini-batches, the convergence rate of SGD for non-convex optimization is reduced to $O(1/\sqrt{T})$ (ghadimi2013stochastic, ; bottou2018optimization, ; zhou2018new, ). Thus, to reach an $\epsilon^{2}$-stationary point $\mathbf{x}$ with $\|\nabla f(\mathbf{x})\|^{2}\leq\epsilon^{2}$, SGD has $O(\epsilon^{-4})$ sample complexity, which could be either higher or lower than the $O(n\epsilon^{-2})$ sample complexity of the GD method, depending on the relationship between $n$ and $\epsilon$. Also, for $p$-dimensional problems, both GD and SGD have memory complexity $p$, since they only need a $p$-dimensional vector to store (stochastic) gradients.

1) Centralized First-Order Methods with Low Complexities: Now, we review several state-of-the-art low-complexity centralized stochastic first-order methods that are related to our GT-STORM algorithm. To reduce the overall sample and communication complexities of the standard GD and SGD algorithms, a natural approach is variance reduction. Earlier works following this approach include SVRG (johnson2013accelerating, ; reddi2016stochastic, ), SAGA (defazio2014saga, ) and SCSG (lei2017non, ). These algorithms has an overall sample complexity of $O(n+n^{2/3}\epsilon^{-2})$. A more recent variance reduction method is the stochastic path-integrated differential estimator (SPIDER) (fang2018spider, ), which is based on the SARAH gradient estimator developed by Nguyen et al. (nguyen2017sarah, ). SPIDER further lowers the sample complexity to $O(n+\sqrt{n}\epsilon^{-2})$, which attains the $\Omega(\sqrt{n}\epsilon^{-2})$ theoretical lower bound for finding an $\epsilon^{2}$-stationary point for $n=O(\epsilon^{-4})$. More recently, to improve the small step-size $O(\epsilon L^{-1})$ in SPIDER, a variant called SpiderBoost was proposed in (wang2019spiderboost, ), which allows a larger constant step-size $O(L^{-1})$ while keeping the same $O(n+\sqrt{n}\epsilon^{-2})$ sample complexity. It should be noted, however, that the significantly improved sample complexity of SPIDER/SpiderBoost is due to a restrictive assumption that a universal Lipschitz smoothness constant exists for all local objectives $f(\cdot;\zeta_{i})$ $\forall i$. This means that the objectives are “similar” and there are no “outliers” in the training samples. Meanwhile, to obtain the optimal communication complexity, SpiderBoost require a (nearly) full gradient every $\sqrt{n}$ iterations and a mini-batch of stochastic gradient evaluation with batch size $\sqrt{n}$ in each iteration.

To overcome the above limitations, a hybrid stochastic gradient descent (Hybrid-SGD) method is recently proposed in (tran2019hybrid, ), where a convex combination of the SARAH estimator (nguyen2017sarah, ) and an unbiased stochastic gradient is used as the gradient estimator. The Hybrid-SGD method relaxes the universal Lipschitz constant assumption in SpiderBoost to an average Lipschitz smoothness assumption. Moreover, it only requires two samples to evaluate the gradient per iteration. As a result, Hybrid-SGD has a $O(\epsilon^{-3})$ sample complexity that is independent of dataset size. Although Hybrid-SGD is for centralized optimization, the interesting ideas therein motivate our GT-STORM approach for decentralized learning following a similar token. Interestingly, we show that in decentralized settings, our GT-STORM method can further improve the gradient evaluation to only one sample per iteration, while not degrading the communication complexity order. Lastly, we remark that all algorithms above have memory complexity at least $2p$ for $p$-dimensional problems. In contrast, GT-STORM enjoys a $p$ memory complexity.

2) Decentralized Optimization Algorithms In the literature, many decentralized learning optimization algorithms have been proposed to solve Problem (1), e.g., first-order methods (nedic2009distributed, ; yuan2016convergence, ; shi2015extra, ; di2016next, ), prime-dual methods (sun2019distributed, ; mota2013d, ), Newton-type methods (mokhtari2016decentralized, ; eisen2017decentralized, ) (see in (nedic2018network, ; chang2020distributed, ) for comprehensive surveys). In this paper, we consider decentralized first-order methods for the non-convex network consensus optimization in (2). In the literature, the convergence rate of the well-known decentralized gradient descent (DGD) algorithm (nedic2009distributed, ) was studied in (zeng2018nonconvex, ), which showed that DGD with a constant step-size converges with an $O(1/T)$ rate to a step-size-dependent error ball around a stationary point. Later, a gradient tracking (GT) method was proposed in (di2016next, ) to find an $\epsilon^{2}$-stationary point with an $O(1/T)$ convergence rate under constant step-sizes. However, these methods require a full gradient evaluation per iteration, which yields $O(n\epsilon^{-2})$ sample complexity. To reduce the per-iteration sample complexity, stochastic gradients are adopted in the decentralized optimization, e.g., DSGD (jiang2017collaborative, ), PSGD (lian2017can, ), GNSD (lu2019gnsd, ). Due to the randomness in stochastic gradients, the convergence rate is reduced to $O(1/\sqrt{T}).$ Thus, the sample and communication complexities of these stochastic methods are $O(\epsilon^{-4})$ and $O(m^{-1}\epsilon^{-4})$, two orders of magnitude higher than their deterministic counterparts. To overcome the limitations in stochastic methods, a natural idea is to use variance reduction techniques similar to those for centralized optimization to reduce the sample and communication complexities for the non-convex network consensus optimization. So far, existing works on the decentralized stochastic variance reduction methods include DSA (mokhtari2016dsa, ), diffusion-AVRG (yuan2018variance, ) and GT-SAGA (xin2019variance, ) etc., all of which focus on convex problems. To our knowledge, the decentralized gradient estimation and tracking (D-GET) algorithm in (sun2019improving, ) is the only work for non-convex optimization. D-GET integrates the decentralized gradient tracking (lu2019gnsd, ) and the SpiderBoost gradient estimator (wang2019spiderboost, ) to obtain $O(mn+m\sqrt{n}\epsilon^{-2})$ dataset-size-dependent sample complexity and $O(\epsilon^{-2})$ communication complexity. Recall that the sample and communication complexities of GT-STORM are $O(m^{1/2}\epsilon^{-3})$ and $O(m^{-1/2}\epsilon^{-3})$, respectively. Thus, if dataset size $n=\Omega(\epsilon^{-2})$, D-GET has a higher sample complexity than GT-STORM. As an example, when $\epsilon=10^{-2}$, $n$ is on the order of $10^{4}$, which is common in modern machine learning datasets. Also, the memory complexity of D-GET is $2p$ as opposed to the $p$ memory complexity of GT-STORM. This implies a huge saving with GT-STORM if $p$ is large, e.g., $p\approx 10^{6}$ in many deep learning models.

In the literature, a standard starting point to solve Problem (2) is to reformulate the problem as (nedic2009distributed, ):

where $\mathbf{I}_{p}$ denotes the $p$-dimensional identity matrix, the operator $\otimes$ denotes the Kronecker product, and $\mathbf{W}\in\mathbb{R}^{m\times m}$ is often referred to as the consensus matrix. We let $[\mathbf{W}]_{ij}$ represent the element in the $i$-th row and the $j$-th column in $\mathbf{W}$. For Problems (4) and (2) to be equivalent, $\mathbf{W}$ should satisfy the following properties:

1. (a)

   Doubly Stochastic: $\sum_{i=1}^{m}[\mathbf{W}]_{ij}=\sum_{j=1}^{m}[\mathbf{W}]_{ij}=1$.

2. (b)

   Symmetric: $[\mathbf{W}]_{ij}=[\mathbf{W}]_{ji}$, $\forall i,j\in\mathcal{N}$.

3. (c)

   Network-Defined Sparsity Pattern: $[\mathbf{W}]_{ij}>0$ if $(i,j)\in\mathcal{L};$ otherwise $[\mathbf{W}]_{ij}=0$, $\forall i,j\in\mathcal{N}$.

The above properties imply that the eigenvalues of $\mathbf{W}$ are real and can be sorted as $-1<\lambda_{m}\leq\cdots\leq\lambda_{2}<\lambda_{1}=1$. We define the second-largest eigenvalue in magnitude of $\mathbf{W}$ as $\lambda\triangleq\max\{|\lambda_{2}|,|\lambda_{m}|\}$ for the further notation convenience. It can be seen later that $\lambda$ plays an important role in the step-size selection and the algorithm’s convergence rate.

As mentioned in Section 2.1, our GT-STORM algorithm is inspired by the GT method (di2016next, ; nedich2016geometrically, ) for reducing consensus error and the recursive variance reduction (VR) methods (fang2018spider, ; wang2019spiderboost, ) developed for centralized optimization. Specifically, in the centralized GT method, an estimator $\mathbf{y}$ is introduced to track the global gradient:

where $\mathbf{g}_{t}$ is the gradient estimation in the $t$th iteration. Meanwhile, to reduce the stochastic error, a gradient estimator $\mathbf{v}$ in VR methods is updated recursively based on a double-loop structure as follows:

where $\nabla f(\mathbf{x};\zeta)$ is the stochastic gradient dependent on parameter $\mathbf{x}$ and a data sample $\zeta,$ and $q$ is the number of the inner loop iterations. On the other hand, if $\text{mod}(t,q)=0$, $\mathbf{v}_{t}$ takes a full gradient. Note that these two estimators have a similar structure: Both are recursively updating the previous estimation based on the difference of the gradient estimations between two consecutive iterations (i.e., momentum). This motivates us to consider the following question: Could we somehow “integrate” these two methods to develop a new decentralized gradient estimator to track the global gradient and reduce the stochastic error at the same time? Unfortunately, the GT and VR estimators can not be combined straightforwardly. The major challenge lies in the structural difference in the outer loop iteration (i.e., $\text{mod}(t,q)=0$), where the VR estimator requires full gradient and does not follow the recursive updating structure.

Surprisingly, in this paper, we show that there exists an “indirect” way to achieve the salient features of both GT and VR. Our approach is to abandon the double-loop structure of VR and pursue a single-loop structure. Yet, this single-loop structure should still be able to reduce the variance and consistently track the global gradient. Specifically, we introduce a parameter $\beta_{t}\in[0,1]$ in the recursive update and integrate it with a consensus step as follows:

where $\mathbf{x}_{i,t},$ $\mathbf{v}_{i,t}$ and $\zeta_{i,t}$ are the parameter, gradient estimator, and random sample in the $t$th iteration at node $i$, respectively. Note that the estimator reduces to the classical stochastic gradient estimator when $\beta_{t}=0$. On the other hand, if we set $\beta_{t}=1$, the estimator becomes the (stochastic) gradient tracking estimator based on a single sample (implying low sample complexity). Then, the key to the success of our GT-STORM design lies in meticulously choosing parameter $\beta_{t}$ to mimic the gradient estimator technique in centralized optimization (cutkosky2019momentum, ; tran2019hybrid, ). Lastly, the local parameters can be updated by the conventional decentralized stochastic gradient descent step:

where $\eta_{t}$ is the step-size in iteration $t$. To summarize, we state our algorithm in Algorithm 1 as follows.

  Algorithm 1: Gradient-Tracking-based Stochastic Recursive Momentum Algorithm (GT-STORM).   Initialization:

1. 1.

   Choose $T>0$ and let $t=1$. Set $\mathbf{x}_{i,0}=\mathbf{x}^{0}$ at node $i$. Calculate $\mathbf{v}_{i,0}=\nabla f_{i}(\mathbf{x}_{i,0};\zeta_{i,0})$ at node $i$.

Two remarks for Algorithm 1 are in order. First, thanks to the single-loop structure, GT-STORM is easier to implement compared to the low-sample-complexity D-GET (sun2019improving, ) method, which has in a double-loop structure. Second, GT-STORM only requires $p$ memory space due to the use of only one intermediate vector $\mathbf{v}$ at each node. In contrast, the memory complexity of D-GET is $2p$ (cf. $\mathbf{y}$ and $\mathbf{v}$ in (sun2019improving, )). This 50% saving is huge particularly for deep learning models, where the number of parameters could be in the range of millions.

In this section, we will establish the complexity properties of the proposed GT-STORM algorithm. For better readability, we state the main theorem and its corollary in this section and provide the intermediate lemmas to Section 3.3. We start with the following assumptions on the global and local objectives:

In the above assumptions, (a) and (c) are standard in the stochastic non-convex optimization literature; (b) is an expected Lipschitz smoothness condition over the data distribution, which implies the conventional global Lipschitz smoothness (ghadimi2013stochastic, ) by the Jensen’s inequality. Note that (b) is weaker than the individual Lipschitz smoothness in (fang2018spider, ; wang2019spiderboost, ; sun2019improving, ): if there exists an outlier data sample, then the individual objective function might have a very large smoothness parameter while the average smoothness can still be small; (d) is equivalent to the Lipschitz continuity assumption, which is also commonly used for non-convex stochastic algorithms (zhou2018generalization, ; karimireddy2019error, ; koloskova2019decentralized, ) and is essential for analyzing the decentralized gradient descent method (yuan2016convergence, ; zeng2018nonconvex, ; jiang2017collaborative, ).¹¹1Note that under the assumption (b), as long as the parameter $\mathbf{x}$ is bounded, (d) is satisfied.

For convenience, in the subsequent analysis, we define $\tilde{\mathbf{W}}=\mathbf{W}\otimes\mathbf{I}_{m},$ $\mathbf{g}_{i,t}=\nabla f_{i}(\mathbf{x}_{i,t}),$ $\mathbf{u}_{i,t}=\nabla f_{i}(\mathbf{x}_{i,t};\zeta_{i,t})$, $\mathbf{w}_{i,t}=\nabla f_{i}(\mathbf{x}_{i,t};\zeta_{i,t})-\nabla f_{i}(\mathbf{x}_{i,t-1};\zeta_{i,t})$ and $\mathbf{a}_{t}=[\mathbf{a}_{1,t}^{\top},\cdots,\mathbf{a}_{m,t}^{\top}]^{\top}$ and $\bar{\mathbf{a}}_{t}=\frac{1}{m}\sum_{i=1}^{m}\mathbf{a}_{i,t},$ for $\mathbf{a}\in\{\mathbf{x},\mathbf{u},\mathbf{w},\mathbf{v},\mathbf{g}\}$. Then, the algorithm can be compactly rewritten in the following matrix-vector form:

Furthermore, since $\mathbf{1}^{\top}\mathbf{W}=\mathbf{1}^{\top},$ we have $\mathbf{\bar{x}}_{t}=\mathbf{\bar{x}}_{t-1}-\eta_{t-1}\bar{\mathbf{v}}_{t-1},$ $\bar{\mathbf{v}}_{t}=\beta_{t}{\bar{\mathbf{v}}}_{t-1}+\beta_{t}\bar{\mathbf{w}}_{t}+(1-\beta_{t})\bar{\mathbf{u}}_{t}.$ We first state the convergence result for Algorithm 1 as follows:

In Theorem 1, $c_{0}$ and $c_{1}$ are two constants depending on the network topology, which in turn will affect the step-size and convergence: with a sparse network, i.e., $\lambda$ is close to but not exactly one (recall that $\lambda=\max\{|\lambda_{2}|,|\lambda_{m}|\}$). In order for $1-(1+c_{1})\lambda^{2}-\frac{1}{c_{0}}>0$ to hold, $c_{0}$ needs to be large and $c_{1}$ needs be close to zero, which leads to small $k_{1},$ $k_{2}$ and $k_{3}.$ Note that the step-size $\eta_{t}$ is of the order $O(t^{-1/3}),$ which is larger than the $O(t^{-1/2})$ order for the classical decentralized SGD algorithms. With this larger step-size, the convergence rate is $O(t^{-2/3})$ and faster than the rate $O(t^{-1/2})$ for the decentralized SGD algorithms. Based on Theorem 1, we have the sample and communication complexity results for Algorithm 1:

Due to space limitation, we provide a proof sketch for Theorem 1 here and relegate the details to the appendices. First, we bound the error of gradient estimator $\mathbb{E}[\|\mathbf{v}_{t}-\mathbf{g}_{t}\|^{2}]$ as follows:

It can be seen that the upper bound depends on the error in the previous step with a factor $\beta_{t}^{2}$. This will be helpful when we construct a potential function. Then, according to the algorithm updates (9)–(10), we show the following descent inequality:

We remark that the right-hand-side (RHS) of the above inequality contains the consensus error of local parameters $\sum_{t=0}^{T}\mathbb{E}[\|\mathbf{x}_{t}-\mathbf{1}\otimes\mathbf{\bar{x}}_{t}\|^{2}]$, which makes the analysis more difficult than that of the centralized optimization. Next, we prove the contraction of iterations in the following lemma, which is useful in analyzing the decentralized gradient tracking algorithms.

Finally, we define a potential function in (4), based on which we prove the convergence bound:

Finally, by properly selecting the parameters, constants $C_{1},$ $C_{2}$ and $C_{3}$ can be made non-negative, which leads to Theorem 1.