DESTRESS: Computation-Optimal and Communication-Efficient Decentralized Nonconvex Finite-Sum Optimization

The central contribution of this paper lies in the development of a new resource-efficient algorithm for nonconvex finite-sum optimization problems in a decentralized environment, dubbed DEcentralized STochastic REcurSive gradient methodS (DESTRESS). DESTRESS provably finds first-order stationary points of the global objective function $f({\bm{{x}}})$ with the optimal incremental first-order (IFO) oracle complexity, i.e. the complexity of evaluating sample gradients, matching state-of-the-art centralized algorithms, but at a much lower communication complexity compared to existing decentralized algorithms over a wide range of parameter regimes.

To achieve resource efficiency, DESTRESS leverages several key ideas in the algorithm design. To reduce local computation, DESTRESS harnesses the finite-sum structure of the empirical risk function by performing stochastic variance-reduced recursive gradient updates [NvDP⁺19, FLLZ18, WJZ⁺19, Li19, LR21, LBZR21, ZXG18]—an approach that is shown to be optimal in terms of IFO complexity in the centralized setting—in a randomly activated manner to further improve computational efficiency when the local sample size is limited. To reduce communication, DESTRESS employs gradient tracking [ZM10] with a few mixing rounds per iteration, which helps accelerate the convergence through better information sharing [LCCC20]; the extra mixing scheme can be implemented using Chebyshev acceleration [AS14] to further improve the communication efficiency. In a nutshell, to find an $\epsilon$-approximate first-order stationary points, i.e. ${\mathbb{\bm{E}}}\big{\|}\nabla f({\bm{{x}}}^{\mathsf{output}})\big{\|}^{2}_{2}\leq\epsilon$, where ${\bm{{x}}}^{\mathsf{output}}$ is the output of DESTRESS, and the expectation is taken with respect to the randomness of the algorithm, DESTRESS requires:

• •

  $O\big{(}m+(m/n)^{1/2}L/\epsilon\big{)}$ per-agent IFO calls,²²2The big-$O$ notation is defined in Section 1.3. which is network-independent; and

• •

  $O\Big{(}\frac{\log\big{(}(n/m)^{1/2}+2\big{)}}{(1-\alpha)^{1/2}}\cdot\big{(}(mn)^{1/2}+\frac{L}{\epsilon}\big{)}\Big{)}$ rounds of communication,

where $L$ is the smoothness parameter of the sample loss, $\alpha\in[0,1)$ is the mixing rate of the network topology, $n$ is the number of agents, and $m=N/n$ is the local sample size.

Decentralized optimization and learning have been studied extensively, with contemporary emphasis on the capabilities to scale gracefully to large-scale problems — both in terms of the size of the data and the size of the network. For the conciseness of the paper, we focus our discussions on the most relevant literature and refer interested readers to recent overviews [NRB20, XPNK20, XKK20a] for further references.

Section 2 introduces preliminary concepts and the algorithm development, Section 3 shows the theoretical performance guarantees for DESTRESS, Section 4 provides numerical evidence to support the analysis, and Section 5 concludes the paper. Proofs and experiment settings are postponed to appendices.

Throughout this paper, we use boldface letters to represent matrices and vectors. We use $\|\cdot\|_{\mathsf{op}}$ for matrix operator norm, $\otimes$ for the Kronecker product, ${\bm{{I}}}_{n}$ for the $n$-dimensional identity matrix and $\bm{1}_{n}$ for the $n$-dimensional all-one vector. For two real functions $f(\cdot)$ and $g(\cdot)$ defined on ${\mathbb{\bm{R}}}^{+}$, we say $f(x)=O\big{(}g(x)\big{)}$ or $f(x)\lesssim g(x)$ if there exists some universal constant $M>0$ such that $f(x)\leq Mg(x)$. The notation $f(x)=\Omega\big{(}g(x)\big{)}$ or $f(x)\gtrsim g(x)$ means $g(x)=O\big{(}f(x)\big{)}$.

Detailed in Algorithm 1, we propose a novel decentralized stochastic optimization algorithm, dubbed DESTRESS, for finding first-order order stationary points of nonconvex finite-sum problems. Motivated by stochastic recursive gradient methods in the centralized setting, DESTRESS has a nested loop structure:

1. 1.

   The outer loop adopts dynamic average consensus to estimate and track the global gradient $\nabla F({\bm{{x}}}^{(t)})$ at each agent in (5), where ${\bm{{x}}}^{(t)}$ is the stacked parameter estimate (cf. (4)). This helps to “reset” the stochastic gradient to a less noisy starting gradient ${\bm{{v}}}^{(t),0}={\bm{{s}}}^{(t)}$ of the inner loop. A key property of (5)—which is a direct consequence of dynamic average consensus—is that the average of ${\bm{{s}}}^{(t)}$ equals to the dynamic average of local gradients, i.e. ${\overline{\bm{s}}}^{(t)}=\tfrac{1}{n}\sum_{i\in[n]}{\bm{{s}}}_{i}^{(t)}=\tfrac{1}{n}\sum_{i\in[n]}\nabla f_{i}({\bm{{x}}}_{i}^{(t)})$.

2. 2.

   The inner loop refines the parameter estimate ${\bm{{u}}}^{(t),0}={\bm{{x}}}^{(t)}$ by performing randomly activated stochastic recursive gradient updates (6), where the stochastic recursive gradient ${\bm{{g}}}^{(t),s}$ is updated in (6b) via sampling mini-batches from activated agents’ local datasets.

To complete the last mile, inspired by [LCCC20], we allow DESTRESS to perform a few rounds of mixing or gossiping whenever communication takes place, to enable better information sharing and faster convergence. Specifically, DESTRESS performs $K_{\mathsf{out}}$ and $K_{\mathsf{in}}$ mixing steps for the outer and inner loops respectively per iteration, which is equivalent to using

as mixing matrices, and correspondingly a network with better connectivity; see (5), (6a) and (6c). Note that Algorithm 1 is written in matrix notation, where the mixing steps are described by ${\bm{{W}}}_{\mathsf{in}}\otimes{\bm{{I}}}_{n}$ or ${\bm{{W}}}_{\mathsf{out}}\otimes{\bm{{I}}}_{n}$ and applied to all agents simultaneously. The extra mixing steps can be implemented by Chebyshev acceleration [AS14] with improved communication efficiency.

Compared with existing decentralized algorithms based on stochastic variance-reduced algorithms such as D-GET [SLH20] and GT-SARAH [XKK20b], DESTRESS utilizes different gradient estimators and communication protocols: First, DESTRESS produces a sequence of reference points $x^{(t)}$—which converge to a global first-order stationary point—to “restart” the inner loops periodically using fresher information; secondly, the communication and computation in DESTRESS are paced differently due to the introduction of extra mixing, which allow a more flexible trade-off schemes between different types of resources; last but not least, the random activation of stochastic recursive gradient updates further saves local computation, especially when the local sample size is small compared to the number of agents.

We first introduce Assumption 1 and Assumption 2, which are standard assumptions imposed on the loss function. Assumption 1 implies that all local objective functions $f_{i}(\cdot)$ and the global objective function $f(\cdot)$ also have Lipschitz gradients, and Assumption 2 guarantees the absence of trivial solutions.

Due to the nonconvexity, first-order algorithms are generally guaranteed to converge to only first-order stationary points of the global loss function $f(\cdot)$, defined below in Definition 2.

Theorem 1, whose proof is deferred to Appendix B, shows that DESTRESS converges in expectation to an approximate first-order stationary point, under suitable parameter choices.

If there is only one agent, i.e. $n=1$, the mixing rate will be $\alpha=0$, we can choose $K_{\mathsf{in}}=K_{\mathsf{out}}=p=1$, and Theorem 1 reduces to [NvDP⁺19, Theorem 1], its counterpart in the centralized setting. For general decentralized settings with arbitrary mixing schedules, Theorem 1 provides a comprehensive characterization of the convergence rate, where an $\epsilon$-approximate first-order stationary point can be found in expectation in a total of

iterations; here, $T$ is the number of outer iterations and $S$ is the number of inner iterations. Clearly, a larger step size $\eta$, as allowable by (7), hints on a smaller iteration complexity, and hence a smaller IFO complexity.

There are two conditions in (7). On one end, $K_{\mathsf{in}}$ needs to be large enough (i.e., perform more rounds of extra mixing) to counter the effect when $p$ is small (i.e., we compute less stochastic gradients every iteration), or when $\alpha$ is close to $1$ (i.e., the network is poorly connected). On the other end, the step size $\eta$ needs to be small enough to account for the requirement of the step size in the centralized setting, as well as the effect of imperfect communication due to decentralization. For well-connected networks where $\alpha\ll 1$, the terms introduced by the decentralized setting will diminish—indicating the iteration complexity is close to that of the centralized setting. For poorly-connected networks, carefully designing the mixing matrix and other parameters can ensure a desirable trade-off between convergence speed and communication cost. The following corollary provides specific parameter choices for DESTRESS to achieve the optimal per-agent IFO complexity. The proof is deferred to Appendix C.

As elaborated in Section 1.1, DESTRESS achieves a network-independent IFO complexity that matches the optimal complexity in the centralized setting. In addition, when the accuracy $\epsilon\lesssim L/(mn)^{1/2}$, DESTRESS reaches a communication complexity of $O\big{(}\frac{1}{(1-\alpha)^{1/2}}\cdot\frac{L}{\epsilon}\big{)}$, which is independent of the sample size.

It is worthwhile to further highlight the role of the random activation probability $p$ in achieving the optimal IFO by allowing “fractional” batch size. Note that the batch size is set as $b=\left\lceil\sqrt{m/n}\right\rceil$, where $m$ is the local sample size, and $n$ is the number of agents.

1. 1.

   When the local sample size is large, i.e. $m\geq n$, we can approximate $b\approx\sqrt{m/n}$ and $p\approx 1$. In fact, Corollary 1 continues to hold with $p=1$ in this regime.

2. 2.

   However, when the number of agents is large, i.e. $n>m$, the batch size $b=1$ and $p=\sqrt{m/n}<1$, which mitigates the potential computation waste by only selecting a subset of agents to perform local computation, compared to the case when we naively set $p=1$.

Therefore, by introducing random activation, we can view $pb=\sqrt{m/n}$ as the effective batch size at each agent, which allows fractional values and leads to the optimal IFO complexity in all scenarios.

To begin with, we employ logistic regression with nonconvex regularization to solve a binary classification problem using the Gisette dataset.⁴⁴4The dataset can be accessed at https://archive.ics.uci.edu/ml/datasets/Gisette. We split the Gisette dataset to $n=20$ agents, where each agent receives $m=300$ training samples of dimension $d=5000$. The sample loss function is given as

where $\{\bm{f},l\}$ represents a training tuple, $\bm{f}\in{\mathbb{\bm{R}}}^{d}$ is the feature vector and $l\in\{0,1\}$ is the label, and $\lambda$ is the regularization parameter. For this experiment, we set $\lambda=0.01$.

Figure 1 shows the loss and testing accuracy for all algorithms. DESTRESS significantly outperforms other algorithms both in terms of communication and computation. It is worth noting that, DSGD converges very fast at the beginning of training, but cannot sustain the progress due to the diminishing schedule of step sizes. On the contrary, the variance-reduced algorithms can converge with a constant step size, and hence converge better overall. Moreover, due to the refined gradient estimation and information mixing designs, DESTRESS can bear a larger step size than GT-SARAH, which leads to the fastest convergence and best overall performance. In addition, a larger number of extra mixing steps leads to a better performance when the graph topology becomes less connected.

Next, we compare the performance of DESTRESS with comparisons to DSGD and GT-SARAH for training a one-hidden-layer neural network with $64$ hidden neurons and sigmoid activations for classifying the MNIST dataset [Den12]. We evenly split $60,000$ training samples to $20$ agents at random. Figure 2 plots the training loss and testing accuracy against the number of communication rounds and gradient evaluations for all algorithms. Again, DESTRESS significantly outperforms other algorithms in terms of computation and communication costs due to the larger step size and extra mixing, which validates our theoretical analysis.