CADA: Communication-Adaptive Distributed Adam

To put our work in context, we review prior contributions that we group in two categories.

We develop a new adaptive SGD algorithm for distributed learning, called Communication-Adaptive Distributed Adam (CADA). Akin to the dynamic scaling of every gradient coordinate in Adam, the key motivation of adaptive communication is that during distributed learning, not all communication rounds between the server and workers are equally important. So a natural solution is to use a condition that decides whether the communication is important or not, and then adjust the frequency of communication between a worker and the server. If some workers are not communicating, the server uses their stale information instead of the fresh ones. We will show that this adaptive communication technique can reduce the less informative communication of distributed Adam.

Analogous to the original Adam [Kingma and Ba(2014)] and its modified version AMSGrad [Reddi et al.(2018)Reddi, Kale, and Kumar], our new CADA approach also uses the exponentially weighted stochastic gradient $h^{k+1}$ as the update direction of $\theta^{k+1}$, and leverages the weighted stochastic gradient magnitude $v^{k+1}$ to inversely scale the update direction $h^{k+1}$. Different from the direct distributed implementation of Adam that incorporates the fresh (thus unbiased) stochastic gradients $\bar{\bm{\nabla}}^{k}=\frac{1}{M}\!\sum_{m\in{\cal M}}\!\!\nabla\ell(\theta^{k};\xi_{m}^{k})$, CADA exponentially combines the aggregated stale stochastic gradients $\bm{\nabla}^{k}=\frac{1}{M}\!\sum_{m\in{\cal M}}\!\!\nabla\ell(\hat{\theta}_{m}^{k};\hat{\xi}_{m}^{k})$, where $\nabla\ell(\hat{\theta}_{m}^{k};\hat{\xi}_{m}^{k})$ is either the fresh stochastic gradient $\nabla\ell(\theta^{k};\xi_{m}^{k})$, or an old copy when $\hat{\theta}_{m}^{k}\neq\theta^{k};\hat{\xi}_{m}^{k}\neq\xi_{m}^{k}$. Informally, with $\alpha_{k}>0$ denoting the stepsize at iteration $k$, CADA has the following update

where $\beta_{1},\beta_{2}>0$ are the momentum weights, $\hat{V}^{k+1}:={\text{diag}}(\hat{v}^{k+1})$ is a diagonal matrix whose diagonal vector is $\hat{v}^{k+1}:=\max\{v^{k+1},\hat{v}^{k}\}$, the constant is $\epsilon>0$, and $I$ is an identity matrix. To reduce the memory requirement of storing all the stale stochastic gradients $\{\nabla\ell(\theta^{k};\xi_{m}^{k})\}$, we can obtain $\bm{\nabla}^{k}$ by refining the previous aggregated stochastic gradients $\bm{\nabla}^{k-1}$ stored in the server via

where $\delta_{m}^{k}:=\nabla\ell(\theta^{k};\xi_{m}^{k})-\nabla\ell(\hat{\theta}_{m}^{k};\hat{\xi}_{m}^{k})$ is the stochastic gradient innovation, and ${\cal M}^{k}$ is the set of workers that upload the stochastic gradient to the server at iteration $k$. Henceforth, $\hat{\theta}_{m}^{k}=\theta^{k};\hat{\xi}_{m}^{k}=\xi_{m}^{k},\,\forall m\in{\cal M}^{k}$ and $\hat{\theta}_{m}^{k}=\hat{\theta}_{m}^{k-1};\hat{\xi}_{m}^{k}=\hat{\xi}_{m}^{k-1},\,\forall m\notin{\cal M}^{k}$. See CADA’s implementation in Figure 1.

Clearly, the selection of subset ${\cal M}^{k}$ is both critical and challenging. It is critical because it adaptively determines the number of communication rounds per iteration $|{\cal M}^{k}|$. However, it is challenging since 1) the staleness introduced in the Adam update will propagate not only through the momentum gradients but also the adaptive learning rate; 2) the importance of each communication round is dynamic, thus a fixed or nonadaptive condition is ineffective; and 3) the condition needs to be checked efficiently without extra overhead. To overcome these challenges, we develop two adaptive conditions to select ${\cal M}^{k}$ in CADA.

With details deferred to Section 2, the contributions of this paper are listed as follows.

c1) We introduce a novel communication-adaptive distributed Adam (CADA) approach that reuses stale stochastic gradients to reduce communication for distributed implementation of Adam.

c2) We develop a new Lyapunov function to establish convergence of CADA under both the nonconvex and Polyak-Łojasiewicz (PL) conditions even when the datasets are non-i.i.d. across workers. The convergence rate matches that of the original Adam.

c3) We confirm that our novel fully-adaptive CADA algorithms achieve at least $60\%$ performance gains in terms of communication upload over some popular alternatives using numerical tests on logistic regression and neural network training.

The LAG method [Chen et al.(2018)Chen, Giannakis, Sun, and Yin] modifies the distributed gradient descent update. Instead of communicating with all workers per iteration, LAG selects the subset of workers ${\cal M}^{k}$ to obtain fresh full gradients and reuses stale full gradients from others, that is

where ${\cal M}^{k}$ is adaptively decided by comparing the gradient difference $\|\nabla{\cal L}_{m}(\theta^{k})-\nabla{\cal L}_{m}(\theta^{k-\tau_{m}^{k}})\|$. Following this principle, the direct (or “naive”) stochastic version of LAG selects the subset of workers ${\cal M}^{k}$ to obtain fresh stochastic gradients $\nabla{\cal L}_{m}\big{(}\theta^{k};\xi_{m}^{k}\big{)}$, $m\in{\cal M}^{k}$. The stochastic LAG also follows the distributed SGD update, but it selects ${\cal M}^{k}$ by: if worker $m$ finds the innovation of the fresh stochastic gradient $\nabla\ell(\theta^{k};\xi_{m}^{k})$ is small such that it satisfies

where $c\geq 0$ and $d_{\rm max}$ are pre-fixed constants, then worker $m$ reuses the old gradient, $m\in{\cal M}\backslash{\cal M}^{k}$, and sets the staleness $\tau_{m}^{k+1}=\tau_{m}^{k}+1$; otherwise, worker $m$ uploads the fresh gradient, and sets $\tau_{m}^{k+1}=1$.

If the stochastic gradients were full gradients, LAG condition (5) compares the error induced by using the stale gradients and the progress of the distributed gradient descent algorithm, which has proved to be effective in skipping redundant communication [Chen et al.(2018)Chen, Giannakis, Sun, and Yin]. Nevertheless, the observation here is that the two stochastic gradients (5) are evaluated on not just two different iterates ($\theta^{k}$ and $\theta^{k-\tau_{m}^{k}}$) but also two different samples ($\xi_{m}^{k}$ and $\xi_{m}^{k-\tau_{m}^{k}}$) thus two different loss functions.

This subtle difference leads to the ineffectiveness of (5). We can see this by expanding the left-hand-side (LHS) of (5) by (see details in supplemental material)

Even if $\theta^{k}$ converges, e.g., $\theta^{k}\rightarrow\theta^{*}$, and thus the right-hand-side (RHS) of (5) $\big{\|}\theta^{k+1-d}\!-\!\theta^{k-d}\big{\|}^{2}\!\rightarrow\!0$, the LHS of (5) does not, because the variance inherited in (6a) and (6b) does not vanish yet the gradient difference at the same function (6c) diminishes. Therefore, the key insight here is that the non-diminishing variance of stochastic gradients makes the LAG rule (5) ineffective eventually. This will also be verified in our simulations when we compare CADA with stochastic LAG.

In this section, we formally develop our CADA method, and present the intuition behind its design.

The key of the CADA design is to reduce the variance of the innovation measure in the adaptive condition. We introduce two CADA variants, both of which follow the update (2), but they differ in the variance-reduced communication rules.

The first one termed CADA1 will calculate two stochastic gradient innovations with one $\tilde{\delta}_{m}^{k}:=\nabla\ell(\theta^{k};\xi_{m}^{k})-\nabla\ell(\tilde{\theta};\xi_{m}^{k})$ at the sample $\xi_{m}^{k}$, and one $\tilde{\delta}_{m}^{k-\tau_{m}^{k}}:=\nabla\ell(\theta^{k-\tau_{m}^{k}};\xi_{m}^{k-\tau_{m}^{k}})-\nabla\ell(\tilde{\theta};\xi_{m}^{k-\tau_{m}^{k}})$ at the sample $\xi_{m}^{k-\tau_{m}^{k}}$, where $\tilde{\theta}$ is a snapshot of the previous iterate $\theta$ that will be updated every $D$ iterations. As we will show in (2.2), $\tilde{\delta}_{m}^{k}-\tilde{\delta}_{m}^{k-\tau_{m}^{k}}$ can be viewed as the difference of two variance-reduced gradients calculated at $\theta^{k}$ and $\theta^{k-\tau_{m}^{k}}$. Using $\tilde{\delta}_{m}^{k}-\tilde{\delta}_{m}^{k-\tau_{m}^{k}}$ as the error induced by using stale information, CADA1 will exclude worker $m$ from ${\cal M}^{k}$ if worker $m$ finds

In (7), we use the change of parameter $\theta^{k}$ averaged over the past $d_{\rm max}$ consecutive iterations to measure the progress of algorithm. Intuitively, if (7) is satisfied, the error induced by using stale information will not large affect the learning algorithm. In this case, worker $m$ does not upload, and the staleness of information from worker $m$ increases by $\tau_{m}^{k+1}=\tau_{m}^{k}+1$; otherwise, worker $m$ belongs to ${\cal M}^{k}$, uploads the stochastic gradient innovation $\delta_{m}^{k}$, and resets $\tau_{m}^{k+1}=1$.

The rationale of CADA1. In contrast to the non-vanishing variance in LAG rule (see (6)), the CADA1 rule (7) reduces its inherent variance. To see this, we can decompose the LHS of (7) as the difference of two variance reduced stochastic gradients at iteration $k$ and $k-\tau_{m}^{k}$. Using the stochastic gradient in SVRG as an example [Johnson and Zhang(2013)], the innovation can be written as

Define the minimizer of (1) as $\theta^{\star}$. With derivations given in the supplementary document, the expectation of the LHS of (7) can be upper-bounded by

If $\theta^{k}$ converges, e.g., $\theta^{k},\theta^{k-\tau_{m}^{k}},\tilde{\theta}\rightarrow\theta^{*}$, the RHS of (9) diminishes, and thus the LHS of (7) diminishes. This is in contrast to the LAG rule (6) lower-bounded by a non-vanishing value. Notice that while enjoying the benefit of variance reduction, our communication rule does not need to repeatedly calculate the full gradient $\nabla{\cal L}_{m}(\tilde{\theta})$, which is only used for illustration purpose.

In addition to (7), the second rule is termed CADA2. The key difference relative to CADA1 is that CADA2 uses $\nabla\ell(\theta^{k};\xi_{m}^{k})-\nabla\ell(\theta^{k-\tau_{m}^{k}}_{m};\xi_{m}^{k})$ to estimate the error of using stale information. CADA2 will reuse the stale stochastic gradient $\nabla\ell(\theta^{k-\tau_{m}^{k}}_{m};\xi_{m}^{k-\tau_{m}^{k}})$ or exclude worker $m$ from ${\cal M}^{k}$ if worker $m$ finds

If (10) is satisfied, then worker $m$ does not upload, and the staleness increases by $\tau_{m}^{k+1}=\tau_{m}^{k}+1$; otherwise, worker $m$ uploads the stochastic gradient innovation $\delta_{m}^{k}$, and resets the staleness as $\tau_{m}^{k+1}=1$. Notice that different from the naive LAG (5), the CADA condition (10) is evaluated at two different iterates but on the same sample $\xi_{m}^{k}$.

The rationale of CADA2. Similar to CADA1, the CADA2 rule (10) also reduces its inherent variance, since the LHS of (10) can be written as the difference between a variance reduced stochastic gradient and a deterministic gradient, that is

With derivations deferred to the supplementary document, similar to (9) we also have that $\mathbb{E}[\|\nabla\ell(\theta^{k};\xi_{m}^{k})-\nabla\ell(\theta^{k-\tau_{m}^{k}};\xi_{m}^{k})\|^{2}]\rightarrow 0$ as the iterate $\theta^{k}\rightarrow\theta^{\star}$.

For either (7) or (10), worker $m$ can check it locally with small memory cost by recursively updating the RHS of (7) or (10). In addition, worker $m$ will update the stochastic gradient if the staleness satisfies $\tau_{m}^{k}\geq D$. We summarize CADA in Algorithm 1.

Computational and memory cost of CADA. In CADA, checking (7) and (10) will double the computational cost (gradient evaluation) per iteration. Aware of this fact, we have compared the number of iterations and gradient evaluations in simulations (see Figures 2-5), which will demonstrate that CADA requires fewer iterations and also fewer gradient queries to achieve a target accuracy. Thus the extra computation is small. In addition, the extra memory for large $d_{\rm max}$ is low. To compute the RHS of (7) or (10), each worker only stores the norm of model changes ($d_{\rm max}$ scalars).

The convergence results of CADA critically builds on the subsequent Lyapunov analysis. We will start with analyzing the expected descent in terms of $\mathcal{L}(\theta^{k})$ by applying one step CADA update.

Lemma 3.1 contains four terms in the RHS of (3.1): the first two terms quantify the correlations between the gradient direction $\nabla\mathcal{L}(\theta^{k})$ and the stale stochastic gradient $\bm{\nabla}^{k}$ as well as the state momentum stochastic gradient $h^{k}$; the third term captures the drift of two consecutive iterates; and, the last term estimates the maximum drift of the adaptive stepsizes over $D+1$ iterations.

From Lemma 3.1, analyzing the progress of $\mathcal{L}(\theta^{k})$ under CADA is challenging especially when the effects of staleness and the momentum couple with each other. Because the the state momentum gradient $h^{k}$ is recursively updated by $\bm{\nabla}^{k}$, we will first need the following lemma to characterize the regularity of the stale aggregated stochastic gradients $\bm{\nabla}^{k}$, which lays the theoretical foundation for incorporating the properly controlled staleness into the Adam’s momentum update.

Lemma 3.2 justifies the relevance of the stale yet properly selected stochastic gradients. Intuitively, the first term in the RHS of (3.2) resembles the descent of using SGD with the unbiased stochastic gradient, and the second and third terms will diminish if the stepsizes are diminishing since $\mathbb{E}\left[\|\theta^{k}-\theta^{k-1}\|^{2}\right]={\cal O}(\alpha_{k}^{2})$. This is achieved by our designed communication rules.

In view of Lemmas 3.1 and 3.2, we introduce the following Lyapunov function:

where $\theta^{\star}$ is the solution of (1), $\{b_{k}\}_{k=1}^{K}$ and $\{\rho_{d}\}_{d=1}^{D}$ are constants that will be specified in the proof.

The design of Lyapunov function in (3.1) is motivated by the progress of $\mathcal{L}(\theta^{k})$ in Lemmas 3.1-3.2, and also coupled with our communication rules (7) and (10) that contain the parameter difference term. We find this new Lyapunov function can lead to a much simple proof of Adam and AMSGrad, which is of independent interest. The following lemma captures the progress of the Lyapunov function.

The first term in the RHS of (15) is strictly negative, and the second term is positive but potentially small since it is ${\cal O}(\alpha_{k}^{2})$ with $\alpha_{k}\rightarrow 0$. This implies that the function ${\cal V}^{k}$ will eventually converge if we choose the stepsizes appropriately. Lemma 3.3 is a generalization of SGD’s descent lemma. If we set $\beta_{1}=\beta_{2}=0$ in (2) and $b_{k}=0,\rho_{d}=0,\,\forall d,k$ in (3.1), then Lemma 3.3 reduces to that of SGD in terms of $\mathcal{L}(\theta^{k})$; see e.g., [Bottou et al.(2018)Bottou, Curtis, and Nocedal, Lemma 4.4].

Building upon our Lyapunov analysis, we first present the convergence in nonconvex case.

From Theorem 3.4, the convergence rate of CADA in terms of the average gradient norms is ${\cal O}(1/\sqrt{K})$, which matches that of the plain-vanilla Adam [Reddi et al.(2018)Reddi, Kale, and Kumar, Chen et al.(2019)Chen, Liu, Sun, and Hong]. Unfortunately, due to the complicated nature of Adam-type analysis, the bound in (16) does not achieve the linear speed-up as analyzed for asynchronous nonadaptive SGD such as [Lian et al.(2016)Lian, Zhang, Hsieh, Huang, and Liu]. However, our analysis is tailored for adaptive SGD and does not make any assumption on the asynchrony, e.g., the set of uploading workers are independent from the past or even independent and identically distributed.

Next we present the convergence results under a slightly stronger assumption on the loss ${\cal L}(\theta)$.

The PL condition is weaker than the strongly convexity, which does not even require convexity [Karimi et al.(2016)Karimi, Nutini, and Schmidt]. And it is satisfied by a wider range of problems such as least squares for underdetermined linear systems, logistic regression, and also certain types of neural networks.

We next establish the convergence of CADA under this condition.

Theorem 3.5 implies that under the PL-condition of the loss function, the CADA algorithm can achieve the global convergence in terms of the loss function, with a fast rate ${\cal O}(1/K)$. Compared with the previous analysis for LAG [Chen et al.(2018)Chen, Giannakis, Sun, and Yin], as we highlighted in Section 3.1, the analysis for CADA is more involved, since it needs to deal with not only the outdated gradients but also the stochastic momentum gradients and the adaptive matrix learning rates. We tackle this issue by i) considering a new set of communication rules (7) and (10) with reduced variance; and, ii) incorporating the effect of momentum gradients and the drift of adaptive learning rates in the new Lyapunov function (3.1).