A Unified Momentum-based Paradigm of Decentralized SGD for Non-Convex Models and Heterogeneous Data

TL;DR

This paper introduces a unified momentum-based framework, UMP, for decentralized SGD addressing non-convexity and data heterogeneity, with proven convergence and significant empirical accuracy improvements.

Contribution

The paper proposes UMP, a general framework with two algorithms, D-SUM and GT-DSUM, unifying momentum-based decentralized SGD methods with theoretical convergence guarantees.

Findings

Convergence guarantees for non-convex objectives.

Up to 57.6% accuracy improvement in experiments.

Unified framework covers various momentum-based variants.

Abstract

Emerging distributed applications recently boosted the development of decentralized machine learning, especially in IoT and edge computing fields. In real-world scenarios, the common problems of non-convexity and data heterogeneity result in inefficiency, performance degradation, and development stagnation. The bulk of studies concentrates on one of the issues mentioned above without having a more general framework that has been proven optimal. To this end, we propose a unified paradigm called UMP, which comprises two algorithms, D-SUM and GT-DSUM, based on the momentum technique with decentralized stochastic gradient descent(SGD). The former provides a convergence guarantee for general non-convex objectives. At the same time, the latter is extended by introducing gradient tracking, which estimates the global optimization direction to mitigate data heterogeneity(i.e., distribution drift). We can cover most momentum-based variants based on the classical heavy ball or Nesterov's acceleration with different parameters in UMP. In theory, we rigorously provide the convergence analysis of these two approaches for non-convex objectives and conduct extensive experiments, demonstrating a significant improvement in model accuracy by up to 57.6% compared to other methods in practice.

Tables

Table 1. Table 1: The testing accuracy with different algorithms on various training benchmarks and different degrees of non-IID.

Datasets	Algorithms	Testing Accuracy ($\%$)
		non-IID $=0.1$	non-IID $=1$	non-IID $=10$
MNIST LeCun et al. (1998)	Local SGD w/ momentum	$95.66\pm 0.21$	$97.99\pm 0.03$	$98.39\pm 0.03$
	QG-DSGDm	$96.02\pm 0.19$	$97.46\pm 1.36$	$98.21\pm 0.04$
	SlowMo	$97.32\pm 0.02$	$97.93\pm 0.07$	$98.34\pm 0.06$
	D-SUM (ours)	$\text{97.89}\pm 0.21$	$\text{98.77}\pm 0.04$	$\text{98.94}\pm 0.01$
	GT-DSUM (ours)	$97.51\pm 0.61$	$98.70\pm 0.01$	$98.82\pm 0.03$
EMNIST Cohen et al. (2017)	Local SGD w/ momentum	$45.90\pm 1.21$	$36.77\pm 0.13$	$38.29\pm 0.03$
	QG-DSGDm	$46.03\pm 0.6$	$46.02\pm 0.12$	$36.72\pm 0.02$
	SlowMo	$45.52\pm 0.03$	$37.11\pm 0.01$	$37.50\pm 0.0$
	D-SUM (ours)	$49.68\pm 0.43$	$49.75\pm 0.05$	$42.50\pm 0.01$
	GT-DSUM (ours)	$\text{50.49}\pm 0.82$	$\text{50.25}\pm 0.07$	$\text{51.87}\pm 0.02$
CIFAR10 Krizhevsky et al. (2009)	Local SGD w/ momentum	$22.94\pm 1.11$	$42.93\pm 0.85$	$52.82\pm 0.01$
	QG-DSGDm	$26.34\pm 1.42$	$49.12\pm 038$	$54.03\pm 0.24$
	SlowMo	$31.06\pm 1.27$	$50.46\pm 0.04$	$55.50\pm 0.10$
	DSUM (ours)	$31.16\pm 1.27$	$54.34\pm 0.11$	$57.59\pm 1.05$
	GT-DSUM (ours)	$\text{36.16}\pm 0.74$	$\text{56.95}\pm 1.56$	$\text{59.34}\pm 1.55$
AG NEWS Zhang et al. (2015)	Local SGD w/ momentum	$75.51\pm 0.44$	$77.98\pm 0.39$	$80.66\pm 0.02$
	QG-DSGDm	$78.82\pm 0.31$	$79.33\pm 0.38$	$82.24\pm 0.02$
	SlowMo	$82.57\pm 0.03$	$83.17\pm 0.01$	$83.79\pm 0.01$
	DSUM (ours)	$84.13\pm 0.55$	$85.46\pm 0.31$	$87.52\pm 0.04$
	GT-DSUM (ours)	$\text{84.29}\pm 0.37$	$\text{87.59}\pm 0.18$	$\text{89.07}\pm 0.04$

Table 2. Table 2: The impact of α 𝛼 \alpha for D-SUM and GT-DSUM on the test accuracy with non-IID = 1 absent 1 =1 . “ ⋆ ⋆ \star ” indicates non-convergence.

Datasets	Methods	The test accuracy ($\%$) evaluated on different $\alpha$ under the non-IID $=1$ case
		$\alpha =0$	$\alpha =0.5$	$\alpha =1$	$\alpha =2$	$\alpha =3$	$\alpha =4$	$\alpha =5$	$\alpha =8$	$\alpha =10$	$\alpha =15$
MNIST	D-SUM	$98.05$	$98.40$	$98.57$	$98.76$	$98.85$	$98.76$	$98.85$	99.09	$98.87$	$92.86$
	GT-DSUM	$98.18$	$98.50$	$98.70$	$98.70$	98.80	$\star$	$\star$	$\star$	$\star$	$\star$
EMNIST	D-SUM	$37.1$	$43.58$	$35.58$	$49.75$	55.32	$49.80$	$43.00$	$48.47$	$53.04$	$\star$
	GT-DSUM	$33.72$	$46.06$	$47.10$	50.25	$39.84$	$\star$	$\star$	$\star$	$\star$	$\star$
CIFAR10	D-SUM	$47.10$	$50.85$	$51.68$	$50.32$	54.83	$51.10$	$51.53$	$50.68$	$\star$	$\star$
	GT-DSUM	$45.98$	$49.80$	$50.55$	$54.77$	57.58	$54.43$	$\star$	$\star$	$\star$	$\star$
AG NEWS	D-SUM	$79.16$	$81.78$	$83.72$	$85.46$	$86.38$	$86.36$	$88.20$	$87.07$	88.82	$88.56$
	GT-DSUM	$78.90$	$83.12$	$85.34$	87.59	$86.89$	$77.67$	$\star$	$\star$	$\star$	$\star$

Table 3. Table 3: Datasets and Models

Dataset	Task	Training samples	Testing samples	Classes	Model
MNIST LeCun et al. [1998]	Handwritten character recognition (CV)	$60,000$	$10,000$	$10$	LeNet described in Table 4
EMNIST Cohen et al. [2017]	Handwritten character recognition (CV)	$731,668$	$82,587$	$62$	CNN described in Table 5
CIFAR10 Krizhevsky et al. [2009]	Image classification (CV)	$50,000$	$10,000$	$10$	LeNet described in Table 4
AG NEWS Zhang et al. [2015]	Text classification (NLP)	$120,000$	$7,600$	$4$	RNN described in Table 6

Table 4. Table 4: LeNet model on MNIST and CIFAR10.

Layer	Output Shape	Hyperparameters	Activation
Conv2d	$(28,28,6)$	kernel size $=5$	ReLU
MaxPool2d	$(14,14,6)$	pool size $=2$
Con2d	$(10,10,16)$	kernel size $=5$	ReLU
MaxPool2d	$(5,5,16)$	pool size $=2$
Flatten	$400$
Dense	$120$
Dense	$84$
Dropout	$84$	$p=0.5$
Dense	$10$

Table 5. Table 5: CNN model on EMNIST.

Layer	Output Shape	Hyperparameters	Activation
Conv2d	$(26,26,32)$	kernel size $=3$, strides $=1$
Conv2d	$(24,24,64)$	kernel size $=3$, strides $=1$	ReLU
MaxPool2d	$(12,12,64)$	pool size $=2$
Dropout	$(12,12,64)$	$p=0.25$
Flatten	$9,216$
Dense	$128$
Dropout	$128$	$p=0.5$
Dense	$62$		softmax

Table 6. Table 6: RNN model on AG NEWS.

Layer	Hyperparameters
EmbeddingBag	embeddings $=95,812$, dimension $=64$
Dense	in_features $=64$, out_features $=4$
Dropout	$p=0.5$

Equations

$$f^{\ast }=\underset{\mathbf{x}\in {\mathbb{R}}^d}{\min }\left[f(\mathbf{x}):=\frac{1}{n}\sum _{i=1}^n\left[f_i(\mathbf{x})={\mathbb{E}}_{{\xi }_i\sim {\mathcal{D}}_i}{F}_i(\mathbf{x},{\xi }_i)\right]\right].$$

$$\begin{array}{cc}\mathbf{A}=\left[{\mathbf{a}}_1,\cdots ,{\mathbf{a}}_n\right]\in {\mathbb{R}}^{d\times n},& \bar{\mathbf{A}}=\left[\bar{\mathbf{a}},\cdots ,\bar{\mathbf{a}}\right]=\mathbf{A}\frac{1}{n}{11}^{\top }.\end{array}$$

$$\{\begin{array}{c}\begin{array}{rl}& {\mathbf{u}}_i^{(t+1)}=\beta {\mathbf{u}}_i^{(t)}+{\mathbf{g}}_i^{(t)}\\ & {\mathbf{x}}_i^{(t+1)}={\mathbf{x}}_i^{(t)}-\eta {\mathbf{u}}_i^{(t+1)},\end{array}\end{array}$$

$${\mathbf{x}}_i^{(t+1)}={\mathbf{x}}_i^{(t)}-\eta {\mathbf{g}}_i^{(t)}+\beta \left({\mathbf{x}}_i^{(t)}-{\mathbf{x}}_i^{(t-1)}\right),$$

$$\{\begin{array}{c}\begin{array}{rl}& {\mathbf{u}}_i^{(t+1)}=\beta {\mathbf{u}}_i^{(t)}+{\mathbf{g}}_i^{(t)}\\ & {\mathbf{v}}_i^{(t+1)}=\beta {\mathbf{u}}_i^{(t+1)}+{\mathbf{g}}_i^{(t)}\\ & {\mathbf{x}}_i^{(t+1)}={\mathbf{x}}_i^{(t)}-\eta {\mathbf{v}}_i^{(t+1)}.\end{array}\end{array}$$

$${\mathbf{x}}_i^{(t+1)}={\mathbf{x}}_i^{(t)}-\eta {\mathbf{g}}_i^{(t)}+\beta \left({\mathbf{x}}_i^{(t)}-\eta {\mathbf{g}}_i^{(t)}-{\mathbf{x}}_i^{(t-1)}+\eta {\mathbf{g}}_i^{(t-1)}\right).$$

$$\{\begin{array}{c}\begin{array}{rl}& {\mathbf{u}}_i^{(t),\tau +1}={\mathbf{x}}_i^{(t),\tau }-\eta {\mathbf{g}}_i^{(t),\tau }\\ & {\mathbf{v}}_{i,}^{(t),\tau +1}={\mathbf{x}}_i^{(t),\tau }-\alpha \eta {\mathbf{g}}_i^{(t),\tau }\\ & {\mathbf{x}}_i^{(t),\tau +1}={\mathbf{u}}_i^{(t),\tau +1}+\beta \left({\mathbf{v}}_i^{(t),\tau +1}-{\mathbf{v}}_i^{(t),\tau }\right),\end{array}\end{array}$$

$${\mathbf{x}}_i^{(t+1)}=\sum _{j=1}^n{w}_{ij}{\mathbf{x}}_j^{(t),K}.$$

$$\frac{1}{KT}\sum _{t=0}^{T-1}\sum _{\tau =0}^{K-1}\mathbb{E}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })\Vert }^2$$

$$\le \frac{2\left(f({\mathbf{x}}_0)-f^{\ast }\right)}{\tilde{\eta }KT}+\frac{2{\beta }^2{\hat{\eta }}^2{L}^2{G}^2}{n(1-\beta {)}^4}+\frac{2L^2{C}_1}{n^2(1-Q_1)}$$

$$+\frac{L}{n}\left({\sigma }^2+2\tilde{\eta }{G}^2+3{\sigma }^2\tilde{\eta }\right),$$

$$0\le \eta \le \min \left\{\frac{\rho }{12\lambda },\frac{1-\beta }{2},\frac{3+\beta }{2\sqrt{3}\lambda (1+\alpha \beta )},\frac{1+\beta }{2\sqrt{3}\lambda \alpha \beta }\right\}\frac{1}{L}$$

$$\{\begin{array}{c}\begin{array}{rl}& (1+\alpha \beta )(1-\lambda )\le \frac{1}{2\sqrt{2}},\\ & \rho \le \frac{48{\lambda }^2{L}^2}{(1-\lambda {)}^2},\\ & 4{\beta }_0^2\left(1-\frac{\rho }{4}\right)<1,\\ & 8(1-\rho )(1+\alpha \beta {)}^2<1,\end{array}\end{array}$$

$$\frac{1}{T}\sum _{t=0}^{T-1}\mathbb{E}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t)})\Vert }^2$$

$$\le \frac{2\left(f({\mathbf{x}}_0)-f^{\ast }\right)}{\tilde{\eta }T}+\frac{8{\beta }^2{\hat{\eta }}^2{L}^2}{n(1-\beta {)}^4}\left({\sigma }^2+3{\zeta }^2+4G^2\right)$$

$$+\frac{12L^2}{n}\left(\frac{Q_3}{T}+\frac{(2-Q_2)Q_4}{(1-Q_2)T}+\frac{V_{\max }}{1-Q_2}+\frac{C_2}{1-Q_2}\right)$$

$$+\frac{4{\lambda }^2{\sigma }^2+16{\zeta }^2}{n}+8(1-{\lambda }^2){ϵ}^2+\frac{12{\beta }^2{\hat{\eta }}^2{L}^2{ϵ}^2}{(1-\beta {)}^4}$$

$${\mathbf{z}}_i^{(t),\tau +1}={\mathbf{z}}_i^{(t),\tau }-\frac{\eta }{1-\beta }{\mathbf{g}}_i^{(t),\tau }$$

$${\mathbf{c}}_i^{(t),\tau +1}=\beta {\mathbf{c}}_i^{(t),\tau }+\left(\alpha -\alpha \beta -1\right)\eta {\mathbf{g}}_i^{(t),\tau }.$$

$${\mathbf{z}}_i^{(t),\tau +1}$$

$$={\mathbf{x}}_i^{(t),\tau }+\frac{\beta }{1-\beta }\left({\mathbf{x}}_i^{(t),\tau }-{\mathbf{x}}_i^{(t),\tau -1}+\alpha \eta {\mathbf{g}}_i^{(t),\tau -1}\right)-\frac{\eta }{1-\beta }{\mathbf{g}}_i^{(t),\tau }$$

$$={\mathbf{x}}_i^{(t),\tau }+{\mathbf{y}}_i^{(t),\tau }-\frac{\eta }{1-\beta }{\mathbf{g}}_i^{(t),\tau }.$$

$${\mathbf{c}}_i^{(t),\tau +1}$$

$$=\beta \left({\mathbf{x}}_i^{(t),\tau }-{\mathbf{v}}_i^{(t),\tau }\right)+\left(\alpha -\alpha \beta -1\right)\eta {\mathbf{g}}_i^{(t),\tau }.$$

$${\mathbb{E}}_{t,\tau }f({\bar{\mathbf{z}}}^{(t),\tau +1})-f({\bar{\mathbf{z}}}^{(t),\tau })$$

$$-⟨\nabla f({\bar{\mathbf{z}}}^{(t),\tau }),{\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]⟩$$

$$\stackrel{(a)}{\le }\frac{1}{2a}{\Vert \nabla f({\bar{\mathbf{z}}}^{(t),\tau })-\nabla f({\bar{\mathbf{x}}}^{(t),\tau })\Vert }^2-\frac{1-a}{2}{\Vert {\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2-\frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })\Vert }^2+\frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })-{\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2$$

$$\stackrel{(b)}{\le }\frac{L^2}{2a}{\Vert {\bar{\mathbf{z}}}^{(t),\tau }-{\bar{\mathbf{x}}}^{(t),\tau }\Vert }^2-\frac{1-a}{2}{\Vert {\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2-\frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })\Vert }^2+\frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })-{\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2,$$

$${\Vert {\bar{\mathbf{z}}}^{(t),\tau }-{\bar{\mathbf{x}}}^{(t),\tau }\Vert }^2=\frac{{\beta }^2}{(1-\beta {)}^2}\Vert \underset{{\bar{\mathbf{c}}}^{(t),\tau }}{\underbrace{{\bar{\mathbf{x}}}^{(t),\tau }-{\bar{\mathbf{v}}}^{(t),\tau }}}{\Vert }^2.$$

$${\Vert {\bar{\mathbf{z}}}^{(t),\tau }-{\bar{\mathbf{x}}}^{(t),\tau }\Vert }^2$$

$$\stackrel{(b)}{=}\frac{{\beta }^2{\hat{\eta }}^2{s}^2}{(1-\beta {)}^2}{\Vert \sum _{j=0}^{l-1}\frac{{\beta }^{l-1-j}}{s}{\bar{\mathbf{g}}}^j\Vert }^2$$

$$\stackrel{(c)}{\le }\frac{{\beta }^2{\hat{\eta }}^2{s}^2}{(1-\beta {)}^2}\sum _{j=0}^{l-1}\frac{{\beta }^{l-1-j}}{s}{\Vert {\bar{\mathbf{g}}}^j\Vert }^2$$

$$\stackrel{(d)}{\le }\frac{{\beta }^2{\hat{\eta }}^2}{(1-\beta {)}^3}\sum _{j=0}^{l-1}{\beta }^{l-1-j}{\Vert {\bar{\mathbf{g}}}^j\Vert }^2,$$

$$-⟨\nabla f({\bar{\mathbf{z}}}^{(t),\tau }),{\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]⟩$$

$$\le \frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })-{\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2+\frac{{\beta }^2{\hat{\eta }}^2{L}^2}{(1-\beta {)}^3}\sum _{j=0}^{l-1}{\beta }^{l-1-j}{\Vert {\bar{\mathbf{g}}}^l\Vert }^2-\frac{1}{2}{\Vert \nabla f({\bar{\mathbf{x}}}^{(t),\tau })\Vert }^2-\frac{1}{4}{\Vert {\mathbb{E}}_{t,\tau }\left[{\bar{\mathbf{g}}}^{(t),\tau }\right]\Vert }^2.$$

$${\mathbb{E}}_{t,\tau }\left[\Vert {\bar{\mathbf{g}}}^{t,\tau }{\Vert }^2\right]\le {\Vert {\mathbb{E}}_{t,\tau }[{\bar{\mathbf{g}}}^{t,\tau }]\Vert }^2+\frac{{\sigma }^2}{n}$$

$${\mathbb{E}}_{t,\tau }f({\bar{\mathbf{z}}}^{(t),\tau +1})-f({\bar{\mathbf{z}}}^{(t),\tau })$$

Full text

1 Introduction

Distributed machine learning (DML) has emerged as an important paradigm in large-scale machine learning Wan et al. (2022); Zhang et al. (2022); Qu et al. (2022). In terms of how to aggregate the model parameters/gradients among workers, researchers classify the system architecture into two main classes: parameter server (PS) and decentralized. The former is generally considered as the centralized paradigm where the central server acts as a coordinator for convenience, while the latter allows communication in a peer-to-peer fashion over an underlying topology, which could guarantee the model consistency across all workers with better scalability.

Meanwhile, multiple complementary studies Fang et al. (2018); Yu et al. (2019); Hsieh et al. (2020) have focused on the issues of DML mainly based on the following two key aspects. $\bullet$ The property of non-convex objectives is quite complicated in deep learning, in particular in distributed scenarios Karimireddy et al. (2020); Lian et al. (2017). Although some standard theoretical results have been obtained for convex models Tao et al. (2022); Deng and Gao (2021); Tao et al. (2021), much less is applicable in non-convex settings since they may be lossy and cause serious obstacles (e.g., high computation complexity and poor generalization) Ghadimi et al. (2015); Mai and Johansson (2020). $\bullet$ It is well known that heterogeneity in the data is one of key challenges in distributed training, resulting in a slow and unstable convergence as well as poor model generalization. There still exists a gap between the disappointing empirical performance and the degree of data heterogeneity Shang et al. (2022); Lin et al. (2021); Esfandiari et al. (2021). Unfortunately, there are currently no existing works attempting to improve real-world decentralized training from a comprehensive perspective by taking both non-convexity and data heterogeneity into account. Thus, it is non-trivial to handle these challenges, which significantly hinder the development of real-life applications.

Motivated by the momentum’s effects on optimal convergence complexity and empirical evaluation successes  Koloskova et al. (2019); Yu et al. (2019); Han and Gao (2021); Lin et al. (2021), we propose UMP, a Unified, Momentum-based Paradigm in the decentralized learning without considering the communication overhead throughout the paper. It consists of two algorithms named D-SUM and GT-DSUM. The former one D-SUM explores the potential of momentum by maintaining and scaling the momentum buffer to sharpen the loss landscape significantly and overcomes the restrictions of non-convexity, leading to better model performance and faster convergence rate in the non-convex settings. Our latter algorithm GT-DSUM also aims to mitigate the impact of data heterogeneity on the discrepancy of local model parameters by introducing the gradient tracking (GT) technique Di Lorenzo and Scutari (2016). The core insight is that the variance between workers is decreasing while the local gradient asymptotically aligns with the global optimization direction independent on the heterogeneity of the data. GT-DSUM accelerates decentralized learning achieving better generalization performance under both non-convex and different degrees of non-IID.

This paper makes the following main contributions:

• •

We propose a unified momentum-based paradigm UMP with two algorithms for dealing with non-convex and the degree of non-IID simultaneously. Moreover, a variety of algorithms with the momentum technique could be obtained by specifying the parameters of our base algorithms.

• •

We design the first algorithm D-SUM, which achieves good model performance, demonstrating its applicability in terms of efficacy and efficiency. We also provide its convergence result under the non-convex cases.

• •

Our second one GT-DSUM, which is robust to the distribution drift problem by applying the GT technique, is being further developed. We rigorously prove its convergence bound in smooth, non-convex settings.

• •

We additionally conduct extensive experiments to evaluate the performance of UMP on common models, datasets, and dynamic real-world settings. Experimental results demonstrate that D-SUM and GT-DSUM improve the model accuracy by up to $35.8\%$ and $57.6\%$ respectively under different non-IID degrees compared with the well-known decentralized baselines. GT-DSUM performs better than D-SUM on model generalization across training tasks suffering from data skewness.

2 The Unified Paradigm: UMP

In this section, we first begin with the notation and revisit two momentum approaches: the heavy ball (HB) method Polyak (1964) and Nesterov’s momentum Nesterov (1983). Inspired by them, we generalize a unified momentum-based paradigm with two algorithms D-SUM and GT-DSUM, which could cover the above two classical methods and other momentum-based variants, aiming to address issues on non-convexity and data heterogeneity in real-world decentralized learning applications. Finally, we provide the convergence result that they could converge almost to a stationary point for general smooth, non-convex objectives.

2.1 Notation and Preliminary

To better demonstrate the applicable effect in real-world complex scenarios, we consider a decentralized setting with a network topology where $n$ workers jointly deal with an optimization problem. Assume that for every worker $i$, it holds its own datasets drawn from $\mathcal{D}_{i}$ distribution, which corresponds to data heterogeneity. Let $f_{i}:\mathbb{R}^{d}\to\mathbb{R}$ be the training datasets loss function of worker $i$ and can be given in a stochastic form $\mathbb{E}_{\xi_{i}\sim\mathcal{D}_{i}}\left[\nabla F_{i}(\mathbf{x},\xi_{i})\right]=\nabla f_{i}(\mathbf{x})$, where $F_{i}(\mathbf{x},\xi_{i})$ is the per-data loss function related with the mini-batch sample $\xi_{i}\sim\mathcal{D}_{i}$. Then, we formulate the empirical risk minimization with sum-structure objectives:

[TABLE]

Among workers, there is an underlying topology graph $\mathbf{W}\in\mathbb{R}^{n\times n}$, which is convenient to encode the communication between arbitrary two workers, i.e., we let $w_{ij}=0$ if and only if worker $i$ and $j$ are not connected.

Definition 1** (Consensus Matrix  Koloskova et al. (2021)).**

A matrix with non-negative entries $\mathbf{W}\in\left[0,1\right]^{n\times n}$ that is symmetric ($\mathbf{W}=\mathbf{W}^{\top}$), and doubly stochastic ($\mathbf{W1}=\mathbf{1},\mathbf{1}^{\top}\mathbf{W}=\mathbf{1}$), where $\mathbf{1}$ denotes the all-one vector in $\mathbb{R}^{n}$.

Throughout the paper, we use the notation $\mathbf{x}_{i}^{(t),\tau}$ to denote the sequence of model parameters on worker $i$ at the $\tau$-th local update in epoch $t$. For any vector $\mathbf{a}_{i}\in\mathbb{R}^{d}$, we denote its model averaging $\bar{\mathbf{a}}=\frac{1}{n}\sum_{i=1}^{n}\mathbf{a}_{i}$. Let $\left\|\cdot\right\|$, $\left\|\cdot\right\|_{F}$ denote the $l_{2}$ vector norm and Frobenius matrix norm, respectively.

For ease of presentation, we apply both vector and matrix notation whenever it is more convenient. We denote by a capital letter for the matrix form combining by $\mathbf{a}_{i}$ as follows,

[TABLE]

The introduction of a momentum term is one of the most common modifications, which is viewed as a critical component for training the state-of-the-art deep neural networks Qu et al. (2022); Lin et al. (2021). Corresponding to its empirical success, momentum attempts to enhance the convergence rate on non-convex objectives by setting the optimized searching direction as the combination of stochastic gradient and historical directions.

The HB method (i.e., also known as Polyak’s momentum) is first proposed for the smooth and convex settings, written as

[TABLE]

where $\mathbf{u}_{i}^{(t)}$, $\mathbf{g}_{i}^{(t)}$ are denoted as the momentum buffer, and the stochastic gradient of worker $i$ at epoch $t$, respectively. $\eta$ presents the learning rate. The momentum variable $\beta$ adjusts the magnitude of updating direction provided by the past information estimation with the stochastic gradient, indicating the direction of the steepest descent. Equivalently, (3) can be also updated below

[TABLE]

when $t\geq 1$. Holding the past gradient values, this style of update can have better stability to some extent and enables improvement compared with some vanilla SGD methods Cutkosky and Mehta (2020).

Another kind of technique called Nesterov’s shows that choosing with suitable parameters, the extrapolation step can be accelerated from $\mathcal{O}\left(\frac{1}{t}\right)$ to $\mathcal{O}\left(\frac{1}{t^{2}}\right)$, which is the optimal rate for the smooth convex problems. Concretely, its update step is described as follows

[TABLE]

The model parameters are updated by introducing the momentum vector $\mathbf{u}_{i}$ and extra auxiliary $\mathbf{v}_{i}$ sequences. Compared with (3), through decaying the momentum buffer $\mathbf{u}_{i}^{(t)}$, it effectively improves the rate of convergence without causing oscillations. Similarly, the above steps can be written as

[TABLE]

Based on (4) and (6), it is not difficult to observe that the former could evaluate the gradient and add momentum simultaneously, while the latter applies momentum after evaluating gradients, which intuitively causes more computation cost. Meanwhile, leveraging the idea of HB momentum, Nesterov’s acceleration brings us closer to the minimum (i.e., $\mathbf{x}^{\ast}$) by introducing an additional gradient descent rule by adding the subtracted gradients $\eta(\mathbf{g}_{i}^{(t-1)}-\mathbf{g}_{i}^{(t)})$ for general convex cases. The above two basic momentum-based approaches are firstly investigated in convex settings, showing their advantage compared with the vanilla SGD. However, there is still a shortage of a comprehensive analysis of momentum-based SGD under non-convex conditions in common real-world scenarios.

2.2 D-SUM Algorithm

In this section, we present UMP and its first algorithm D-SUM, which is employed in decentralized training under non-convex cases.

Under each epoch, workers first perform $K$ local updates using different optimizers (i.e., SGD, Adam  Kingma and Ba (2015), etc.) with or without momentum. In this paper, we mainly focus on the momentum-based SGD variants, which are demonstrated in (3), and (5) for example. From a comprehensive view, we apply the key update of the stochastic unified momentum (SUM) is according to

[TABLE]

where $\alpha\geq 0$, and $\beta\in\left[0,1\right)$. $\mathbf{a}_{i}^{(t),\tau}$ ($\mathbf{a}_{i}$ could be the instance for $\mathbf{x}_{i}$, $\mathbf{u}_{i}$, $\mathbf{v}_{i}$, and $\mathbf{g}_{i}$) is denoted as the related variables for worker $i$ after $\tau$ local updates in epoch $t$. After $K$ local steps, worker $i$ communicates with its neighbors according to the communication pattern $\mathbf{W}$ for exchanging their local model parameters. We call this synchronization operation as gossip averaging which can be compactly written as

[TABLE]

To present the difference between vanilla SGD and stochastic unified momentum in (7), we summarize the training procedure in Algorithm 1. The specific algorithm instance is obtained by tuning the hyperparameters $\alpha$, $\beta$, $\eta$, and $K$. We cover the basic Heavy Ball method (4) and Nesterov’s momentum (6) when setting $\alpha=0$, and $\alpha=1$, respectively. Besides, when $K=1$, it reduces to the standard mini-batch SGD with momentum acceleration. Specially, we update the auxiliary variable sequences $\left\{\mathbf{v}_{i}\right\}$ for any worker $i$ by using the same gossip synchronization as in (8) interpreted as a restart in the next training epoch to simplify theoretical analysis.

However, there is no theoretical or empirical analysis to demonstrate that the momentum gets rid of heterogeneity which degrades the distributed deep training due to the discrepancies between local activation statistics Hsieh et al. (2020). Not only taking non-convex functions into account, but we also incorporate a technique that is agnostic to data heterogeneity, gradient tracking into D-SUM to alleviate the impact of heterogeneous data in decentralized training for better model generalization in the following.

2.3 GT-DSUM Algorithm

In this subsection, we go further the fact that heterogeneity hinders the local momentum acceleration  Lin et al. (2021) and provides our second algorithm in UMP, termed GT-DSUM, which aims to generalize the consensus model parameters better and alleviate the impact of heterogeneous data by applying the gradient tracking technique.

Taking the discrepancies between workers’ local data partition into account, GT introduces an extra worker-sided auxiliary variable $\mathbf{y}_{i}^{(t)},\forall i$ aiming to asymptotically track the average of $\nabla f_{i}$ assuming the local accurate gradients are accessible at any epoch $t$. Intuitively, GT is agnostic to the heterogeneity, while $\mathbf{y}_{i}^{(t)}$ is approximately equivalent to the global gradient direction along with the epoch $t$ increases. Inspired by this, we introduce GT into D-SUM, yields GT-DSUM. Concretely, we normalize the applied gradient $\mathbf{m}_{i}^{(t),\tau}$ using the mini-batch gradient $\mathbf{g}_{i}^{(t),\tau}$, and the $\mathbf{y}_{i}^{(t)}$ with the dampening factor $\lambda$ to highlight the necessity of local updates. The detailed algorithm is described in Algorithm 2. Within local updates, the model parameters are updated on line 5 with D-SUM but using a normalization term $\mathbf{m}_{i}^{(t),\tau}$. Line 7 and 8 are the same as the basic D-SUM procedures in Algorithm 1. For GT-DSUM, we apply the difference of two consecutive synchronized models shown in line 9 to update the gradient tracker variable in line 10 using the gossip-liked style Xin et al. (2021b, a). Especially, when $K=1$, $\lambda=1$ and $\beta=0$, the Algorithm 2 can be reduced to the original GT algorithm  Koloskova et al. (2021) instance.

Since Algorithm 1 and 2 employ multiple consensus steps from parameters exchanging which significantly increase communication cost, we apply the communication compression technique GRACE Xu et al. (2021) to trade off between model generalization and communication overhead in Section 3.

2.4 Theoretical Analysis

In what follows, we present the convergence analysis of two algorithms in the UMP for general non-convex settings. The detailed proof is in Appendix. Firstly, we state our assumptions throughout the paper.

Assumption 1** ($L$-smooth).**

For each function $f_{i}:\mathbb{R}^{d}\to\mathbb{R}$ is differentiable, and there exists a constant $L>0$ such that for each $\mathbf{x},\mathbf{x}^{\prime}\in\mathbb{R}^{d}:\left\|\nabla f(\mathbf{x})-\nabla f(\mathbf{x}^{\prime})\right\|\leq L\left\|\mathbf{x}-\mathbf{x}^{\prime}\right\|$.

Assumption 2** (Bounded variances).**

We assume that there exists $\sigma>0$ and $\zeta>0$ for any $i,\mathbf{x}\in\mathbb{R}^{d}$ such that $\mathbb{E}_{\xi_{i}\sim\mathcal{D}_{i}}\left\|\nabla F_{i}(\mathbf{x},\xi_{i})-\nabla f_{i}(\mathbf{x})\right\|^{2}\leq\sigma^{2}$, and $\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla f_{i}(\mathbf{x})-\nabla f(\mathbf{x})\right\|^{2}\leq\zeta^{2}$.

Assumptions 1 and 2 are standard in general non-convex objective literature Lin et al. (2021); Yu et al. (2019); Koloskova et al. (2020) in order to ensure the basis of loss functions continuous and the limited influence of heterogeneity among distributed scenarios. Noted that when $\zeta=0$, we have $\nabla f_{i}(\mathbf{x})=\nabla f(\mathbf{x})$, i.e., it reduces to the case of IID data distribution across all participating workers. The third common assumption is to assume the stochastic gradients are uniformly bounded which is stated as follows.

Assumption 3** (Bounded stochastic gradient).**

We assume that the second moment of stochastic gradients is bounded for any $i,\mathbf{x}\in\mathbb{R}^{d},\mathbb{E}_{\xi_{i}\sim\mathcal{D}_{i}}\left\|\nabla F_{i}(\mathbf{x},\xi_{i})\right\|^{2}\leq G^{2}$.

Assumption 4**.**

The mixing matrix is doubly stochastic by Definition 1. Further, define $\bar{\mathbf{Z}}=\mathbf{Z}\frac{1}{n}\mathbf{11}^{\top}$ for any matrix $\mathbf{Z}\in\mathbb{R}^{d\times n}$. Then, the mixing matrix satisfies $\mathbb{E}_{\mathbf{W}}\left\|\mathbf{ZW}-\bar{\mathbf{Z}}\right\|^{2}_{F}\leq(1-\rho)\left\|\mathbf{Z}-\bar{\mathbf{Z}}\right\|^{2}_{F}$.

In Assumption 4, we assume that $\rho:=1-\max\left\{\left|\lambda_{2}(\mathbf{W})\right|,\left|\lambda_{n}(\mathbf{W})\right|\right\}^{2}>0$, where let $\lambda_{i}(\mathbf{W})$ denote the $i$-th largest eigenvalue of the mixing matrix $\mathbf{W}$ with $-1\leq\lambda_{n}(\mathbf{W})\leq\cdots\leq\lambda_{2}(\mathbf{W})\leq\lambda_{1}(\mathbf{W})\leq 1$. For example, the value of $\rho$ is commonly used when $\rho=1$ for the full-mesh (complete) communication topology.

2.4.1 Convergence Analysis of D-SUM

We now state our convergence result for D-SUM (red highlight) in Algorithm 1. The detailed proof is presented in Appendix B

Theorem 1**.**

Considering problem (1) under the above mentioned assumptions, we denote $\beta_{0}=\max\left\{1+\beta,1+\alpha\beta\right\}$, for all $T\geq 1$ and $K\geq 1$ in Algorithm 1 with learning rate $\eta\leq\frac{\rho}{5L}$ and parameters satisfy $\frac{4-\rho}{2}<\frac{1}{\beta_{0}^{2}}$, we have

[TABLE]

where $\tilde{\eta}=\frac{\eta}{1-\beta}$, $\hat{\eta}=\left((1-\beta)\alpha-1\right)\eta$, $Q_{1}=2\beta_{0}^{2}(1-\frac{\rho}{4})$, and $C_{1}=24\eta^{2}\beta_{0}^{2}\zeta^{2}/\rho+4\left(1-\rho\right)\left(1+2\alpha\beta+2\alpha^{2}\beta^{2}\right)\eta^{2}\sigma^{2}$.

Remark 1**.**

Theorem 1 proposes a non-asymptotic convergence bound of D-SUM for general neural network since the second term i.e., $\mathcal{O}\left(\frac{L^{2}\hat{\eta}^{2}}{n}\right)$ generates from the core SGD step in (7). Intuitively, there exists an appropriate $\alpha$ for achieving the optimal training performance in practice, which has been observed in the single node case Yan et al. (2018). In Section 3, we perform related experiments to confirm this speculation.

2.4.2 Convergence Analysis of GT-DSUM

Next theorem is the convergence result of GT-DSUM in Algorithm 2 when $K=1$ with a fixed communication topology among workers for convenience, and the detailed proof is in Appendix C. Based on the GT is addressed with the issue on how to apply the mini-batch gradient estimates to track the global optimization descent direction, we define the following proposition to clarify this illustration.

Proposition 1** (Gradients averaging tracker Di Lorenzo and Scutari (2016)).**

We assume a loose constraint that the auxiliary variables $\mathbf{y}_{i}^{(t)}$ are considered as the tracker of the average $\frac{1}{n}\sum_{j=1}^{n}\nabla f_{j}(\mathbf{x}_{i}^{(t)})$, which means for any epoch $t$, we have $\mathbb{E}\left\|\mathbf{y}_{i}^{(t)}-\frac{1}{n}\sum_{j=1}^{n}\nabla f_{j}(\mathbf{x}_{i}^{(t)})\right\|^{2}\leq\epsilon^{2}$.

Theorem 2**.**

Consider problem (1) under the listed specific assumptions, we denote $\beta_{0}=\max\left\{1+\alpha\beta,1+\beta\right\}$, and set $T\geq 1$ in Algorithm 2 without multiple local steps (i.e., $K=1$) with learning rate $\eta$ chosen as

[TABLE]

and parameters satisfy

[TABLE]

we have

[TABLE]

where $\tilde{\eta}=\frac{\eta}{1-\beta}$ and $\hat{\eta}=(\alpha-\alpha\beta-1)\eta$. In addition, $C_{2}=\frac{\beta_{0}^{2}\left(1-\frac{\rho}{2}\right)}{L^{2}}(\sigma^{2}+\zeta^{2})(192\lambda^{2}L^{2}+\rho)$, $Q_{2}\triangleq\min\left\{4\beta_{0}^{2}\left(1-\frac{\rho}{4}\right),8(1-\rho)(1+\alpha\beta)^{2}\right\}$, $Q_{3}=6(\zeta^{2}+\sigma^{2})$, $Q_{4}=\beta_{0}^{2}\left(1-\frac{\rho}{2}\right)(\sigma^{2}+\zeta^{2})\left(48+\frac{\rho}{L^{2}}+192\lambda^{2}\right)$. Furthermore, we define $V_{\max}\triangleq\max_{0\leq t\leq T-1}\left\{\frac{1}{n}\left(\mathbb{E}\left\|\mathbf{X}^{(t)}-\bar{\mathbf{X}}^{(t)}\right\|_{F}^{2}+\mathbb{E}\left\|\mathbf{Y}^{(t)}-\bar{\mathbf{Y}}^{(t)}\right\|_{F}^{2}\right)\right\}.$

Remark 2**.**

The fourth term on the right-hand side of the Theorem 2, i.e., $\mathcal{O}\left(\frac{1}{nT}+\frac{\beta_{0}^{2}}{nT}+\frac{1}{n\beta_{0}^{2}}\right)$ comes from the the additional GT step for searching global optimal descent estimation in line $10$, Algorithm 2. However, this term can be dominant when $\alpha$ scales due to its higher order. Clearly, $\beta_{0}=1+\alpha\beta$ when $\alpha>1$, and it performs the convergence rate $\mathcal{O}\left(\frac{\beta_{0}^{2}}{nT}\right)$, leading to a significant deterioration from convergence perspective if $\alpha$ rises. Hence, we will show the impact of $\alpha$ in Section 3.

3 Evaluation

Our main evaluation results demonstrate that D-SUM outperforms other methods in terms of model accuracy, and GT-DSUM achieves a higher performance under different levels of non-IID. All experiments are executed in a CPU/GPU cluster, equipped with Inter(R) Xeon(R) Gold 6126, 4 GTX 2080Ti cards, and 12 Tesla T4 cards. We used Pytorch and Ray Moritz et al. (2018) to implement and train our models.

3.1 Experiment Methodology

Baselines. We consider the following three decentralized methods with momentum, which are described as follows: $\bullet$ Local SGD Stich (2018) periodically averages model parameters among all worker nodes. Compared with the vanilla SGD, each node independently runs the single-node SGD with Heavy Ball momentum. $\bullet$ QG-DSGDm Lin et al. (2021) mimics the global optimization direction and integrates the quasi-global momentum into local stochastic gradients without causing extra communication costs. It empirically mitigates the impact on data heterogeneity. $\bullet$ SlowMo Wang et al. (2020) performs a slow, periodical momentum update through an All-Reduce pattern (model averaging) after multiple SGD steps. For simplicity, we use the common mini-batch gradient as the local update direction.

Datasets and models. We study the decentralized behaviors on both computer vision (CV) and natural language processing (NLP) tasks, including MNIST, EMNIST, CIFAR10, and AG NEWS. For all CV tasks, we train different CNN models. For NLP, we train an RNN, which includes an embedding layer, and a dropout layer, followed by a dense layer. The model description is shown in Appendix D.

Hyperparameters. For all algorithms with different benchmarks, the setting deploys $10$ workers by default. In our experiments, we set the local mini-batch size as $256$ for CIFAR10 and $128$ for the rest, and the number of local updates is set as $K=10$. To illustrate the challenge of data heterogeneity in decentralized deep training, we adopt the Dirichlet distribution value Lin et al. (2021) to control different levels of non-IID degree, for the case with non-IID $=0.1,1,10$; the smaller the value is, the more likely the workers hold samples from only one class of labels (i.e., non-IID $=0.1$ can be viewed as an extreme data skewness case). Besides, we set the scalar $\alpha$, momentum $\beta$, normalized parameter $\lambda$ as $2$, $0.9$, and $0.8$ respectively by default. Among choices of $\mathbf{W}$ considered in practice, we pre-construct a dynamic topology changing sequence varying from full-mesh to ring by the popular Metropolis-Hastings rule Koloskova et al. (2021) i.e., $w_{ij}=w_{ji}=\min\left\{\frac{1}{\text{deg}(i)+1},\frac{1}{\text{deg}(j)+1}\right\}$ for any $i,j$, $w_{ii}=1-\sum_{j=1}^{n}w_{ij}$. The learning rate $\eta$ is fine-tuned via a grid search on the set $\left\{10^{-2},10^{-1.5},10^{-1},10^{-0.5}\right\}$ for each algorithm and dataset.

Performance Metrics. We examine the effects of different momentum variants on decentralized deep learning, including

• •

Model generalization is measured by the proportion between the amount of the correct data by the model and that of all data in the test dataset. We report the averaged model performance of local models over test samples.

• •

Effect of different hyperparameters is explored by tuning their values to study the properties of D-SUM and GT-DSUM.

• •

Scalability is a crucial property while handling tasks in a distributed situation.

3.2 Evaluation results

Performance with compared baselines. In Table 1, we can see that our proposed algorithms outperform all other baselines across different levels of data skewness. For CIFAR10 and AG NEWS, the performance of our algorithms and benchmarks: GT-DSUM $>$ D-SUM $>$ SlowMo $>$ QG-DSGDm $>$ Local SGD w/ momentum. Our proposed algorithms outperform other benchmarks on model generalization and demonstrate that GT technique effectively mitigates the negative impact caused by data heterogeneity. As the non-IID level increases, GT-DSUM achieves a higher accuracy than Local SGD w/ momentum up to $57.6\%$ on CIFAR10.

Effect of local update. The number of local updates $K$ is one of the most important parameters since it influences the final model generalization and training time. As usual, the number of $K$ is set less than $20$ Reddi et al. (2020); Qu et al. (2022). Hence, we present the comparison with $K\in\left\{1,5,10,15,20\right\}$. We make two observations from the results in Figure 1. Firstly, our algorithms have better performance than Local SGD w/ momentum regardless $K$. For CIFAR10 and AG NEWS, shown in Figure 1(c), 1(d), we observe that GT-DSUM always keep competitive when the number of local updates increases. Among them, it improves accuracy $53.8\%$ than Local SGD w/ momentum when $K=5$. Secondly, we can find see that the workers may not guarantee to improve the model generalization substantially by increasing the number of local updates $K$. Besides, all benchmarks perform worst when $K=1$, while when $K$ is too large, performance may be degraded because workers’ local models drift too far apart in distributed optimizations Qu et al. (2022).

Effect of $\beta$. The momentum term $\beta$ is further investigated via grid search on the set $\left\{0,01,0.1,0.5,0.9,0.99\right\}$ as different strategies for handling algorithms. The third set of simulations evaluates the performance of model accuracy on different $\beta$, which are depicted in Figure 2. We have a key observation from those results. Regardless of datasets or models, evaluation results with a greater $\beta$ (e.g., 0.9, 0.99) trend to outperform with a smaller one (e.g., 0.01, 0.1). In addition, it can be observed that the testing accuracy monotonically increases with $\beta$ on CIFAR10 and AG NEWS. We note that GT-DSUM reaches a higher model accuracy compared with Local SGD w/ momentum when $\beta$ is increasing. Among all tasks, it improves test accuracy up to $56.8\%$ than Local SGD w/ momentum.

Sensitivity of $\alpha$. One of the most important parameters $\alpha$ is varying from [math] to $15$ to analyze its influence on the convergence performance. Two observations are as follows according to Table 2. Firstly, we note that different optimal values of $\alpha$ are always found when D-SUM is evaluated on various datasets with the same level of data heterogeneity (e.g., non-IID $=1$). It is hard and time-consuming to determine the optimum $\alpha$ due to different characteristics of datasets and models. Secondly, as $\alpha$ increases, GT-DSUM makes a significant degradation on model performance. This phenomenon verifies the analysis detailed in Remark 2, which indicates that GT-DSUM requires a stricter constraint on $\alpha$ than D-SUM to ensure model validity. Empirically, there is still a gap between the vulnerable property of $\alpha$ to the momentum-based optimizer and the robustness endowing with a superior performance.

Scalability. We finally train on different numbers of workers compared with Local SGD w/ momentum when non-IID $=1$. We evaluate this by extending the scale by adjusting the number of devices $n$ training on $4$, $10$, $16$, and $32$ workers. Results are shown in Figure 3. When the number of participating workers increases, the advantage of our schemes is readily apparent since our method GT-DSUM consistently reaches a higher model accuracy compared to the Local SGD w/ momentum in this non-IID case.

4 Conclusion

In this paper, we propose a unified momentum-based paradigm UMP with two algorithms D-SUM and GT-DSUM. The former obtains good model generalization, dealing with the validity under non-convex cases, while the latter is further developed by applying the GT technique to eliminate the negative impact of heterogeneous data. By deriving the convergence of general non-convex settings, these algorithms achieve competitive performance closely related to a critical parameter $\alpha$. Extensive experimental results show our UMP leads to at most $57.6\%$ increase in improvement of accuracy.

Appendix A Prerequisite

For giving the theoretical analysis of the convergence results of all proposed algorithms, we first present some preliminary facts as follows:

• •

**Fact 1: ** For any random vector $\mathbf{a}$, it holds for $\mathbb{E}\left\|\mathbf{a}\right\|^{2}=\mathbb{E}\left\|\mathbf{a}-\mathbb{E}\left[\mathbf{a}\right]\right\|^{2}+\left\|\mathbb{E}\left[\mathbf{a}\right]\right\|^{2}$.

• •

**Fact 2: ** For any $a>0$, we have $\pm\left\langle\mathbf{a},\mathbf{b}\right\rangle\leq\frac{1}{2a}\left\|\mathbf{a}\right\|^{2}+\frac{a}{2}\left\|\mathbf{b}\right\|^{2}$.

• •

**Fact 3: ** $\left\langle\mathbf{a},\mathbf{b}\right\rangle=\frac{1}{2}\left\|\mathbf{a}\right\|^{2}+\frac{1}{2}\left\|\mathbf{b}\right\|^{2}-\frac{1}{2}\left\|\mathbf{a}-\mathbf{b}\right\|^{2}$.

• •

**Fact 4: ** For given two vectors $\mathbf{a}$ and $\mathbf{b},\forall a>0$, we have $\left\|\mathbf{a}+\mathbf{b}\right\|^{2}\leq(1+a)\left\|\mathbf{a}\right\|^{2}+(1+\frac{1}{a})\left\|\mathbf{b}\right\|^{2}$.

• •

**Fact 5: ** For arbitrary set of $n$ vectors $\left\{\mathbf{a}_{i}\right\}_{i=1}^{n}$, we have $\left\|\sum_{i=1}^{n}\mathbf{a}_{i}\right\|^{2}\leq n\sum_{i=1}^{n}\left\|\mathbf{a}_{i}\right\|^{2}$.

• •

**Fact 6: ** Suppose $\left\{b_{i}\right\}_{i=1}^{n}$, and $\left\{\mathbf{a}_{i}\right\}_{i=1}^{n}$ are a set of non-negative scalars and vectors, respectively. We define $s=\sum_{i=1}^{n}b_{i}$. Then according to Jensen’s inequality, we have $\left\|\sum_{i=1}^{n}b_{i}\mathbf{a}_{i}\right\|^{2}=s^{2}\left\|\sum_{i=1}^{n}\frac{b_{i}}{s}\mathbf{a}_{i}\right\|^{2}\leq s^{2}\cdot\sum_{i=1}^{n}\frac{b_{i}}{s}\left\|\mathbf{a}_{i}\right\|^{2}=s\cdot\sum_{i=1}^{n}b_{i}\left\|\mathbf{a}_{i}\right\|^{2}$.

The inequalities of Fact 4 also hold for the sum of two matrices $\mathbf{A},\mathbf{B}$ in Frobenius norm.

Proposition 2**.**

One step of gossip averaging with the mixing matrix $\mathbf{W}$ defined in the Definition 1 preserves the averaging of the iterates, i.e., $\mathbf{XW}\frac{\mathbf{11}^{\top}}{n}=\mathbf{X}\frac{\mathbf{11}^{\top}}{n}$.

Note our schemes have the following observation on the role of momentum, we now state a basic lemma:

Lemma 1**.**

Let introduce an auxiliary variable $\mathbf{y}_{i}^{(t),\tau}=\frac{\beta}{1-\beta}\left(\mathbf{x}_{i}^{(t),\tau}-\mathbf{v}_{i}^{(t),\tau}\right)$, we define $\mathbf{z}_{i}^{(t),\tau}\triangleq\mathbf{x}_{i}^{(t),\tau}+\mathbf{y}_{i}^{(t),\tau}$ and $\mathbf{c}_{i}^{(t),\tau}\triangleq\frac{1-\beta}{\beta}\mathbf{y}_{i}^{(t),\tau}$. Then we have

[TABLE]

and

[TABLE]

Proof.

For (9), starting from the definition of $\mathbf{v}_{i}^{(t),\tau}$ in (7),

[TABLE]

Similarly for (10),

[TABLE]

∎

Appendix B Proof of D-SUM

Since the smoothness of $f$, it follows that

[TABLE]

where $\mathbb{E}_{t,\tau}$ denotes a conditional expectation over the randomness in the $\tau$-th local updates under epoch $t$, conditioned on all past random variables. According to the described factors , on the right hand side for (12):

[TABLE]

where $(a)$ follows from the combination of Fact 2 and Fact 3; $(b)$ follows by the smoothness in Assumption 1. Since we assume that for any $i$ in the initial stage $\mathbf{x}_{i}^{(0)}=\mathbf{v}_{i}^{(0)}=\mathbf{0}$, based on the definition of $\mathbf{z}^{(t),\tau}$, and $\mathbf{y}^{(t),\tau}$, it can be shown by averaging

[TABLE]

Applying the recursion of (10) in Lemma 1,

[TABLE]

where we omit the aspects of epoch (i.e., $t$), local updates (i.e., $\tau$) and replace them with a more general term: iteration (i.e., $l$) in $(a)$; we define $s=\sum_{j=0}^{l-1}\beta^{l-1-j}$ in $(b)$; $(c)$ follows by the Fact 6 and Jensen’s inequality; $(d)$ follows because $s=\frac{1-\beta^{l}}{1-\beta}<\frac{1}{1-\beta}$ since $\beta\in[0,1)$. Substituting (15) into (13), and we set $a=\frac{1}{2}$, which yields

[TABLE]

Moreover, for the second term in (12), we apply the identified Fact 1 for any vector and Assumption 2, then

[TABLE]

Then plug (16) and the above inequality into (12),

[TABLE]

where $\tilde{\eta}=\frac{\eta}{1-\beta}$. Taking the total expectation, and summing from $\tau=0$ to $K-1$, we have

[TABLE]

Summing from $t=0$ to $T-1$ and dividing both side by $KT$,

[TABLE]

We now bound the upper bound of $T_{1}$:

[TABLE]

where $(a)$ follows because of the Fact 4 by setting $a=1$; $(b)$ follows from the Assumption 2. Then we estimate the $T_{2}$:

[TABLE]

where $(a)$ follows from the complexity of $\left\|\cdot\right\|^{2}$ and Jensen’s inequality; $(b)$ follows by the Fact 4 ; $(c)$ follows by applying $\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left\|\nabla f_{i}(\mathbf{x})-\nabla f(\mathbf{x})\right\|^{2}\leq\zeta^{2}$ in Assumption 2. Finally, we bound the term $T_{3}$:

[TABLE]

where $(a)$ follows by noting that $\sum_{u=j+1}^{TK-1}\beta^{u-1-j}=\frac{1-\beta^{TK-1-j}}{1-\beta}\leq\frac{1}{1-\beta}$. Substitute (21), (22), and (23) into (20), which yields

[TABLE]

Using that $\sum_{i=1}^{n}\left\|\mathbf{a}_{i}\right\|^{2}=\left\|\mathbf{A}\right\|_{F}^{2}$ where $\mathbf{A}:=\left[\mathbf{a}_{1},\cdots,\mathbf{a}_{n}\right]$, we now try to bound the consensus error between the nodes’ parameters and its averaging. We first reiterate the update scheme of (7) in a matrix form regardless of epoch $t$ and local update $\tau$ and denote $l$ as the index of update iteration:

[TABLE]

For averaged parameters which are performed model averaging across all nodes, we can also simply the updates since $\mathbf{W}$ is doubly stochastic, which is described as follows:

[TABLE]

According to the above two equations, we have

[TABLE]

where $(a)$ follows by applying Assumption 4; $(b)$ follows because we add the expectation term of $\mathbf{G}$ and $\bar{\mathbf{G}}$, so that $\mathbf{G}-\bar{\mathbf{G}}=\left(\mathbf{G}-\mathbb{E}\left[\mathbf{G}\right]\right)-\left(\bar{\mathbf{G}}-\mathbb{E}\left[\bar{\mathbf{G}}\right]\right)+\left(\mathbb{E}\left[\mathbf{G}\right]-\mathbb{E}\left[\bar{\mathbf{G}}\right]\right)$, which satisfies the condition of Assumption 2 generalizing the constant $\Delta$; $(c)$ follows from the Fact 4 by setting $a=1$. Here we use the contractivity of the matrix $\mathbf{W}$ and Young’s inequality. We can further proceed as

[TABLE]

where $(a)$ follows by applying the Fact 4 and sets $a=\frac{\rho}{2}$; $(b)$ follows because positive $\rho\leq 1$, and using Fact 4 as well as Assumption 2. By choosing the learning rate $\eta\leq\frac{\rho}{5L}$ ensures that $6\eta^{2}L^{2}\leq\frac{\rho^{2}}{4}$, we have two cases since $a\geq 0$

• •

Case one: $\alpha\in\left[0,1\right)$, we define $\hat{\beta}=1+\beta$, then

[TABLE]

• •

Case two: $\alpha\in\left[1,\infty\right)$, we define $\tilde{\beta}=1+\alpha\beta$, just replace $\hat{\beta}$ term by $\tilde{\beta}$.

Here we denote $\beta_{0}\triangleq\max\left\{\hat{\beta},\tilde{\beta}\right\}$. Since $\mathbf{x}_{i}^{(0)}=\bar{\mathbf{x}}^{(0)}=\mathbf{x}_{0}$, we get $\frac{1}{n}\mathbb{E}\|\mathbf{X}^{(0)}-\bar{\mathbf{X}}^{(0)}\|_{F}^{2}=\mathbf{0}$. Furthermore, we can easily obtain $\frac{1}{n}\mathbb{E}\|\mathbf{X}^{(1)}-\bar{\mathbf{X}}^{(1)}\|_{F}^{2}=24\eta^{2}\beta_{0}^{2}\zeta^{2}/\rho+\Delta$. We observe that when $l\geq 1$, $\frac{1}{n}\mathbb{E}\|\mathbf{X}^{(l)}-\bar{\mathbf{X}}^{(l)}\|_{F}^{2}\leq\frac{C_{1}l}{n\left(1-Q_{1}\right)},$ where we denote $Q_{1}=2\beta_{0}^{2}(1-\frac{\rho}{4})$ ensuring that $Q_{1}<1$, and $C_{1}=24\eta^{2}\beta_{0}^{2}\zeta^{2}/\rho+\Delta$. Substituting the above inequality into (24), for any $\bar{\mathbf{g}}$, $\mathbb{E}\left\|\bar{\mathbf{g}}\right\|^{2}\leq\frac{G^{2}}{n}$ by Assumption 3, and $\bar{\mathbf{z}}^{(0)}=\bar{\mathbf{x}}^{(0)}=\mathbf{x}_{0}$ by definition, rearranging terms yields the Theorem 1.

Appendix C Proof of GT-DSUM

Next, we provide a rigorous proof of GT-DSUM under non-convexity. Here we consider a special case where $K=1$ with a fixed consensus matrix. Then we construct the matrix form of Algorithm 2 as follows:

[TABLE]

Proof sketch. we try to bound the consensus distance (Lemma 2) between the worker’s parameters and its averaging. During this step, we perform a propagation step which brings the parameters of the workers closer to each other. Moreover, we also perform additional gradient tracking (Lemma 3) and their accumulation steps (Lemma 4) which move the distance away from each other. After that, we could immediately apply Lemma 4 into the the single-step update progress in (45).

Lemma 2** (Consensus distance change).**

Given above assumptions in Section 2, let $\beta_{0}=\max\left\{1+\alpha\beta,1+\beta\right\}$, and the update rule generated by (29) using learning rate satisfy $\eta\leq\frac{\rho}{12\lambda L}$ when $t\geq 1$,

[TABLE]

Proof.

Starting from (29) and all consensus matrices satisfy Proposition 2, we have

[TABLE]

where $(a)$ follows by applying Assumption 4; $(b)$, $(c)$ follows by the Fact 4 by choosing $\rho/2$, and $1$ respectively. Then we try to bound the distance between $\mathbf{M}^{(t)}$ and $\bar{\mathbf{M}}^{(t)}$,

[TABLE]

where $(a)$ follows from (29); $(b)$ follows by applying the Fact 4; $(c)$ follows by Fact 5 with the vector set $\left\{\mathbf{G}-\mathbb{E}\left[\mathbf{G}\right],\mathbb{E}\left[\bar{\mathbf{G}}\right]-\bar{\mathbf{G}},\mathbb{E}\left[\mathbf{G}\right]-\mathbb{E}\left[\bar{\mathbf{G}}\right]\right\}$; $(d)$ follows from the Fact 4 and Assumption 2. Since the positive scalar $\rho\leq 1$, substitute (31) into (30) on the condition that $\eta\leq\frac{\rho}{12\lambda L}$ ensures that $36\lambda^{2}L^{2}\eta^{2}\leq\rho^{2}/4$, completing the proof. ∎

Lemma 3** (Gradient tracker distance change).**

Given the assumptions in Section 2, when $t\geq 1$, let the learning rate satisfy $\eta\leq\min\left\{\frac{3+\beta}{2\sqrt{3}\lambda L(1+\alpha\beta)},\frac{1+\beta}{2\sqrt{3}\lambda L\alpha\beta}\right\}$ and follows from the assumption of hyperparameter that $(1+\alpha\beta)(1-\lambda)\leq\frac{1}{2\sqrt{2}}$, which yields

[TABLE]

Proof.

According to the update scheme in (29), we can get

[TABLE]

and

[TABLE]

Then, based on (33) and (29), we have

[TABLE]

The inequality holds for Assumption 4 and Fact 4. Since $\mathbf{W}\prec\mathbf{I}$, $-\mathbf{I}\prec\mathbf{W}$ as well as the product of two doubly stochastic matrices is still doubly stochastic, we have $2\mathbf{I}-(1+\beta)\mathbf{W}\prec(3+\beta)\mathbf{I}$; $\beta\mathbf{W}^{2}-\mathbf{I}\prec(\beta+1)\mathbf{W}^{2}\prec(\beta+1)\mathbf{I}$, we can continue

[TABLE]

where the inequality follows by applying the basic inequality $\left\|\sum_{i=1}^{j}\mathbf{A}_{i}\right\|_{F}^{2}\leq j\sum_{i=1}^{j}\left\|\mathbf{A}_{i}\right\|_{F}^{2}$ for matrices of the same dimension with $j=4$. Plug (31) into(35) under the condition that $(\mathit{i})$ $\eta\leq\frac{3+\beta}{2\sqrt{3}\lambda L(1+\alpha\beta)}$; $(\mathit{ii})$ $\eta\leq\frac{1+\beta}{2\sqrt{3}\lambda L\alpha\beta}$; $(\mathit{iii})$ $(1+\alpha\beta)(1-\lambda)\leq\frac{1}{2\sqrt{2}}$ for ease of presentation. ∎

Lemma 4** (Distance step improvement).**

When $t\geq 1$, using learning rate $\eta\geq\max\left\{\frac{2\sqrt{2}(3+\beta)}{\beta_{0}},\frac{2\sqrt{2}(1+\beta)}{\beta_{0}-1}\right\}$ and $\rho\leq\frac{48\lambda^{2}L^{2}}{(1-\lambda)^{2}}$, satisfy

[TABLE]

Proof.

Adding the results of Lemma 2 and 3 gives the result following from our assumption of learning rate $\eta$ and the hyperparameters $\rho,\lambda,L$. ∎

We now state our convergence results for GT-DSUM in Algorithm 2. Similar to proof process of Theorem 1, with the smoothness of $f$,

[TABLE]

where $\tilde{\eta}=\frac{\eta}{1-\beta}$. For term $T_{1}$, we adopt the same derivation process as (15), which is also suitable for the circumstance with no multiple local updates, indicating that

[TABLE]

where $\hat{\eta}=((1-\beta)\alpha-1)\eta$. We can further obtain that

[TABLE]

where $(a)$ follows from (29) by model averaging; we omit the coefficients in $(b)$; $(c)$ follows by applying the Fact 4. For $T_{4}$,

[TABLE]

where $(a)$ follows the definition of $\bar{\mathbf{g}}_{i}$; $(b)$ follows by the Jensen’s inequality and the Fact 4; and $(c)$ follows because Assumption 2 and 3. Estimate $T_{5}$,

[TABLE]

where $(a)$ follows because $\bar{\mathbf{y}}^{(t)}=\frac{1}{n}\sum_{i=1}^{n}\mathbf{y}_{i}^{(t)}$; $(b)$ follows by applying the Fact 5; $(c)$ follows by the Proposition 1 and Assumption 2. Substituting (37), (38), (39), and (40) into the $T_{1}$ in (36), which yields

[TABLE]

We can omit $T_{2}$ in (36) on the assumption that $\eta\leq\frac{1-\beta}{2L}$ ensures that $\frac{\tilde{\eta}^{2}L}{2}-\frac{\tilde{\eta}}{4}<0$. Then we estimate the bound of $T_{3}$ in (36),

[TABLE]

where $(a)$ is follows by the definition of $\bar{\mathbf{m}}^{(t)}$; $(b)$ follows because of the Fact 4; $(c)$ and $(d)$ follow by applying the Fact 4, Jensen’s inequality and Assumption 2. For $T_{6}$, we can estimate as follows

[TABLE]

where $(a)$ follows from the Fact 4; $(b)$ follows because the Assumption 1 for the first term on the right hand of the inequality; $(c)$ follows by applying the Fact 4 with Jensen’s inequality; $(d)$ follows because the Assumption 2 and Proposition 1. Plugging (43) into (42) yields

[TABLE]

Plugging (41) and (44) into (36), which yields

[TABLE]

the last inequality holds because $\sum_{l=0}^{t-1}\beta^{t-1-l}\leq\frac{1}{1-\beta}$, and we add the non-negative term $\mathbb{E}\left\|\bar{\mathbf{Y}}^{(t)}-\mathbf{Y}^{(t)}\right\|_{F}^{2}$.

Proposition 3**.**

Let $\left\{V_{t}\right\}_{t\geq 0}$ be a non-negative sequence and $C_{2}\geq 0$ be some constant such that $\forall t\geq 1,V_{t+1}\leq Q_{2}V_{t}+Q_{2}V_{t-1}+C_{2}$, where $Q_{2}\in(0,1)$. Then the following inequality holds if $T\geq 1$,

[TABLE]

Summing $t$ over from [math] to $T-1$,

[TABLE]

where $V_{\max}=\max_{0\leq t\leq T-1}\left\{V_{t}\right\}$.

Summing (45) over $t$ from [math] to $T-1$, then dividing both sides by $2/\tilde{\eta}$ and rearranging terms. Finally, applying Lemma 4 to Proposition 3. Concretely, we consider $C_{2}=\frac{\beta_{0}^{2}\left(1-\frac{\rho}{2}\right)}{L^{2}}(\sigma^{2}+\zeta^{2})(192\lambda^{2}L^{2}+\rho)$, $V_{t}\triangleq\mathbb{E}\left\|\mathbf{X}^{(t)}-\bar{\mathbf{X}}^{(t)}\right\|_{F}^{2}+\mathbb{E}\left\|\mathbf{Y}^{(t)}-\bar{\mathbf{Y}}^{(t)}\right\|_{F}^{2}$, $Q_{2}\triangleq\min\left\{4\beta_{0}^{2}\left(1-\frac{\rho}{4}\right),8(1-\rho)(1+\alpha\beta)^{2}\right\}$ ensuring that $\left\{\begin{matrix}\begin{aligned} &4\beta_{0}^{2}\left(1-\frac{\rho}{4}\right)<1\\ &8(1-\rho)(1+\alpha\beta)^{2}<1\end{aligned}\end{matrix}\right.$. Since $\mathbf{x}_{i}^{(0)}=\bar{\mathbf{x}}^{(0)}=\mathbf{0}$, and $\mathbb{E}\left\|\mathbf{Y}^{(0)}-\bar{\mathbf{Y}}^{(0)}\right\|_{F}^{2}=\mathbb{E}\left\|\mathbf{G}^{(0)}-\bar{\mathbf{G}}^{(0)}\right\|_{F}^{2}\leq\frac{6L^{2}}{n}\mathbb{E}\left\|\mathbf{X}^{(0)}-\bar{\mathbf{X}}^{(0)}\right\|+6(\zeta^{2}+\sigma^{2})=6(\zeta^{2}+\sigma^{2})$, thus $V_{0}\leq Q_{3}=6(\zeta^{2}+\sigma^{2})$. Furthermore, we assume that $\forall i,\mathbf{x}_{i}^{(-1)}=\mathbf{y}_{i}^{(-1)}=\mathbf{0}$, which is extended at an initial stage (i.e., $t=0$) for Lemma 2, (3), and (4), thus $V_{1}\leq Q_{4}=\beta_{0}^{2}\left(1-\frac{\rho}{2}\right)(\sigma^{2}+\zeta^{2})\left(48+\frac{\rho}{L^{2}}+192\lambda^{2}\right)$. The proof is completed.

Appendix D Experimental Setup

D.1 Dataset and Model Description

MNIST is a 10-class handwritten digits image classification dataset with $70,000$ $28$ $\times$ $28$ examples, $60,000$ of which are training datasets, the remaining $10,000$ are test datasets. Its extended version, EMNIST consists of images of digits and upper and lower case English characters, which includes $62$ total classes. CIFAR10 is labeled subsets of the 80 million images dataset, sharing the same $60,000$ input images with $10$ unique labels. For NLP, AG NEWS is a 4-class classification dataset on categorized news articles, containing $120,000$ training samples and $7,600$ testing samples. An overall description is given in Table 3.

Appendix E Addtional Evaluations

E.1 Model Generalization

Figure 4, 5, 6 and 7 present the experimental results on the model training process for different benchmarks. We note that the superiority of our proposed methods is better reflected in the convergence acceleration. For example, both D-SUM and GT-DSUM require about $50$ epochs to reach convergence for LeNet over MNIST, which reduces the number of training epochs by $40\%$. Switching to large datasets, i.e., CIFAR10 and AG NEWS, our proposed algorithms converge faster than other baselines with respect to the training epochs.