Quantized Adam with Error Feedback

Introduction

Recently, deep neural networks (LeCun et al., 2015; Goodfellow et al., 2016) achieve high performances in many applications, such as computer vision (Krizhevsky et al., 2012; He et al., 2016), natural language processing (Devlin et al., 2018), speech recognition (Amodei et al., 2016), reinforcement learning (Mnih et al., 2015; Silver et al., 2016), etc. However, a huge deep neural network contains millions of parameters, so its training procedure requires a large amount of training data (Deng et al., 2009; Wu et al., 2019), which may not be stored in a single machine. In addition, due to some privacy issues (McMahan et al., 2021; Yang et al., 2019), all the training data cannot be sent to a single machine but can be stored in different devices. Therefore, how to accelerate the training process by using multiple machines over large-scale data or distributed data has already been a hot topic in both industrial and academic communities (Kraska et al., 2013; Li et al., 2014; Xing et al., 2015; Liu et al., 2017).

An efficient approach to tackle this problem is to develop distributed training algorithms for the huge neural networks (Dean et al., 2012). Most of the distributed algorithms can be summarized into two categories: one is the parameter-server (Smola and Narayanamurthy, 2010) model (or called centralized model) shown in Fig. 2, and the other is the decentralized model (Lian et al., 2017) shown in Fig. 2. For the centralized model in Fig. 2, there are one parameter server and multiple workers. In an update iteration, all workers report the update vectors to the parameter server. After gathering all the update vectors, the parameter server will update the parameters and send the parameters to all workers. While for the decentralized model in Fig. 2, there are $n$ nodes working simultaneously. In each update iteration, each worker computes its update vector respectively and communicates with its neighbors, and then updates its own parameters. When we use a distributed training algorithm such as distributed stochastic gradient descent (Li et al., 2014) in either the centralized model or the decentralized model, plenty of update vectors have to be communicated among different devices. Then, a communication issue emerges for huge networks.

To accelerate the distributed training process of huge deep learning models, we propose a new distributed adaptive stochastic gradient method with gradient quantization, weight quantization, and error-feedback in the parameter server model, as shown in Fig. 2. In what follows, we elaborate on each component used in the proposed method:

Quantization. Note that both gradient quantization and weight quantization are introduced in the proposed method to reduce the communication cost among the workers and the parameter server. Specifically, weight quantization is performed on the parameter server and the quantized weights are then broadcast to all the workers. Weight quantization is introduced because of the consideration of limited storage in edge devices. Meanwhile, gradient quantization is performed on each worker and then the quantized gradients are reported to the server. Thanks to the double quantization schemes, the communication cost can be largely reduced. In addition, for some resource-limited devices, storage is another issue. Weight quantization can also be used to reduce the deep neural network model size efficiently (Han et al., 2015; Zhou et al., 2016; Rastegari et al., 2016). Especially, in federated learning, a distributed device may be smartphones or Internet of things devices, which may encounter both the storage issue and the communication issue. Thus, the weight quantization and gradient quantization schemes can jointly solve these two issues.

Adaptive learning rate. To ease the labor of tuning learning rate, we also adopt the adaptive learning rate as (Kingma and Ba, 2014; Hinton et al., 2012; Duchi et al., 2011; Zou et al., 2018; Reddi et al., 2019; Chen et al., 2018) in the proposed method. Here, the adaptive learning rate is calculated by a similar definition to those in RMSProp (Hinton et al., 2012) and Adam (Kingma and Ba, 2014), except that the noisy gradients are estimated with quantized weights. Moreover, to guarantee the convergence of the proposed method, we set the exponential moving average parameter in estimating the adaptive learning rate the same as that used in Zou et al.(Zou et al., 2019).

Error-feedback. In the proposed method, an error-feedback technique is leveraged to reduce the bias introduced by gradient quantization. The error-feedback technique is also performed on each worker. Actually, the error-feedback technique is motivated by Karimireddy et al. (Karimireddy et al., 2019) by introducing an additional term as the compensation term for the quantized gradient. However, due to the introduced adaptive learning rate and momentum, the compensation term is slightly different from that in Karimireddy et al. (Karimireddy et al., 2019). To the best of our knowledge, this is the first work that simultaneously employs the adaptive learning rate and the error-feedback technique.

Besides, we establish the convergence rate of the proposed algorithm. In the stochastic nonconvex setting, we show that the distributed adaptive stochastic gradient method with gradient quantization and error-feedback converges to the first-order stationary point, and that the distributed adaptive stochastic gradient method with weight quantization and error-feedback converges to the point related to the quantized level under both the single-worker and multi-worker modes. At last, we apply the proposed distributed adaptive method to train deep learning models, such as LeNet (LeCun et al., 1998) on the MNIST dataset (LeCun et al., 1998) and ResNet-101 (He et al., 2016) on the CIFAR100 dataset (Krizhevsky et al., 2009), respectively. The experimental results demonstrate the effectiveness of weight quantization, gradient quantization, and the error-feedback technique working in concert with distributed adaptive stochastic gradient method. Here, we summarize our contributions in three-fold:

We propose a distributed variant of the adaptive stochastic gradient method to train deep learning models. The proposed approach exploits gradient quantization, weight quantization, and the error-feedback technique to accelerate the training process.

We establish the convergence rates of the proposed distributed adaptive stochastic gradient algorithms with weight quantization, gradient quantization, and error-feedback in the nonconvex stochastic setting under the single-worker and multi-worker environments, which are far different from the stochastic gradient setting because adaptive learning rate is introduced into the algorithms.

We apply the proposed algorithms to train deep learning models including LeNet and ResNet-101. The experiments demonstrate the efficacy of the proposed algorithms.

Related Works

In this section, we enumerate several works that are most related to this work. We split the related works into two categories: distributed quantized algorithms and adaptive learning rate.

The quantization functions can be divided into two categories: unbiased quantization functions and biased quantization functions. For unbiased quantization functions, Wen et al. (Wen et al., 2017) showed that with an unbiased ternary quantization function, the distributed stochastic gradient descent algorithm can almost surely converge to a minimum point. Jiang et al. (Jiang and Agrawal, 2018) showed with an unbiased quantization function, the centralized distributed stochastic gradient descent algorithm can converge with convergence rate $\mathcal{O}(1/\sqrt{T})$. Besides, Hou et al.(Hou et al., 2018) showed that in the stochastic convex setting, with gradient quantization solely, the algorithm they proposed will converge to the optimal solution, while with weight quantization the algorithm will converge to the point near the optimal solution which is related to the weight quantization level. However, they can only deal with the unbiased quantization function, which limits the use of both algorithms and theorems.

For biased quantization functions, the main issue is to eliminate the biased error during optimization. A common technique to tackle this issue is error-feedback, where each worker stores the error of the quantization and adds the error term to the next communication before quantization. Based on the decentralized model in Fig. 2, Tang et al. (Tang et al., 2019) and Koloskova et al. (Koloskova et al., 2019) showed that distributed stochastic gradient descent with quantized communication and error-feedback can converge to a stationary point in the nonconvex setting with convergence rate $\mathcal{O}(1/\sqrt{T})$. Based on the centralized model in Fig. 2, Zhou et al. (Zhou et al., 2016) and Wu et al. (Wu et al., 2018) showed that few bits or integer networks can be trained empirically. Zheng et al. (Zheng et al., 2019) showed the convergence of the algorithm with a block quantization function in the nonconvex setting.

Among the above-mentioned algorithms, Hou et al. (Hou et al., 2018) is the most related work to our proposed algorithm. However, their proposed algorithms do not adopt unbiased quantization on gradients. Moreover, they do not incorporate momentum acceleration terms into their algorithm to accelerate its piratical performance. In addition, the convergence analysis in Hou et al. (Hou et al., 2018) is merely restricted to the stochastic convex setting, which makes their algorithm heuristic when it is applied to train deep learning models. By contrast, the convergence rates of our algorithms are established in the more difficult nonconvex setting. In this work, we first extend the error-feedback technique to adaptive stochastic gradient method (Adam) and then establish its convergence in the nonconvex setting, and we compare the most related works in Table 1.

2. Adaptive Learning Rate

Adaptive learning rate, as a popular optimization technique for training deep learning models, has attracted much attention. Numerous papers have studied the convergences of adaptive stochastic gradient methods, such as AdaGrad (Duchi et al., 2011; McMahan and Streeter, 2010), RMSprop (Hinton et al., 2012), Adam (Kingma and Ba, 2014), and AMSGrad (Reddi et al., 2019). Besides the counterexample of divergence when using the Adam algorithm in the convex case in (Reddi et al., 2019), various works have proposed different conditions to make Adam-type methods converge to first-order stationary points. For example, (Ward et al., 2019; Li and Orabona, 2019; Zou et al., 2018) establish global convergence of AdaGrad in the nonconvex setting; Reddi et al. (Reddi et al., 2019) check the difference between learning rates of two adjacent iterations and proposes a new variant called AMSGrad; Chen et al. (Chen et al., 2018) establish the convergence of AMSGrad in the nonconvex setting; Basu et al. (Basu et al., 2018) show that Adam converges when a full-batch gradient is used; Zhou et al. (Zhou et al., 2018) check the independence between gradient square and learning rate to ensure the convergence for the counterexamples in (Reddi et al., 2019), and Zou et al. (Zou et al., 2019) check the parameter setting to give a sufficient condition to guarantee the convergences of both Adam and RMSProp. Also, Reddi et al. (Reddi et al., 2020) introduce distributed stochastic adaptive gradient methods in the centralized model and Nazari et al. (Nazari et al., 2019) introduce a decentralized adaptive gradient method. In this paper, we propose a distributed variant of Adam method by incorporating quantization and error-feedback techniques. We show that the proposed method converges to a saddle point with quantized update vectors, and will be close to a saddle point when we quantize the weights of a certain network.

Main Results

Throughout this paper, we consider the following stochastic nonconvex optimization:

Gradient $\small\nabla{f}$ is $L$-Lipschitz continuous, i.e., $\small\|\nabla\!f\left(x\right)\!-\!\nabla\!f\left(y\right)\|\!\leq\!L\|x\!-\!y\|$. Moreover, the noisy gradient estimation $g_{t}$ is upper bounded and unbiased, i.e., $\small E[g_{t}]\!=\!\nabla\!f\left(x_{t}\right)$ and $\small\|g_{t}\|\!\leq\!G$.

Below, we introduce the gradient quantization operator $\small Q_{g}\left(\cdot\right)$ and weight quantization operator $\small Q_{x}\left(\cdot\right)$ that satisfy the following assumptions, respectively.

Let $\small Q_{g}\left(\cdot\right)$ be the gradient quantization operator defined by Definition 1. We assume that there exists a constant $\small\delta_{g}\geq 0$ such that the inequality holds $\small\|g-Q_{g}\left(g\right)\|\leq(1-\delta_{g})\|g\|$.

Let $\small Q_{x}\left(\cdot\right)$ be the weight quantization operator defined by Definition 1. First, we assume that the noisy gradient estimation at point $x_{t}$ is an unbiased estimation of $\small\nabla f\left(Q_{x}\left(x_{t}\right)\right)$, i.e., $\small E[g_{t}]=\nabla f\left(Q_{x}\left(x_{t}\right)\right)$. In addition, we assume that there exists $\small\delta_{x}\geq 0$ such that $\small\|x-Q_{x}\left(x\right)\|\leq\delta_{x}$.

Assumption 1 is commonly used in analyzing adaptive stochastic type methods (Chen et al., 2018; Reddi et al., 2019; Duchi et al., 2011). Especially, for the gradient quantization and weight quantization, the Lipschitz continuity conditions are used for bounding the error term introduced by quantization. For the weight quantization, the unbiased estimation condition $\small E[g_{t}]=\nabla\!f\!\left(Q_{x}\left(x_{t}\right)\right)$ is used, which has also been used in Hou et al. (Hou et al., 2018). All the detailed proof procedures are placed in Section 4.

In this subsection, we first present the quantized Generic Adam with weight quantization, gradient quantization, and the error-feedback technique working on a single machine. Then, to show the influence on convergence related to gradient quantization or weight quantization, we establish its convergence rate with either gradient quantization or weight quantization.

Algorithm 1 unifies weight quantization, gradient quantization, and the error-feedback technique into Adam, in which $\small Q_{x}\left(\cdot\right)$ denotes the weight quantization operator, $\small Q_{g}\left(\cdot\right)$ denotes the gradient quantization operator, and $e_{t}$ denotes the error-feedback term. In addition, to establish the convergence rate of Algorithm 1 in the nonconvex setting, we make the following assumptions on momentum parameter $\beta_{t}$, exponential moving average parameter $\theta_{t}$, and base learning rate $\alpha_{t}$.

Assume that momentum parameter $\beta_{t}$, exponential moving average parameter $\theta_{t}$, and base learning rate $\alpha_{t}$ satisfy $\beta_{t}\!\in\![0,\beta]$ with $0\!<\!\beta\!<\!1$, $\theta_{t}\!=\!1\!-\!{\theta}/{t}$, and $\small\alpha_{t}\!=\!{\alpha}/{\sqrt{t}}$, respectively. Furthermore, we denote $\small\gamma\!=\!\beta/\theta^{\prime}$ and $\small C_{1}\!=\!\prod_{j=1}^{N}\!\frac{\theta_{j}}{\theta^{\prime}}$ with $\small N\!=\!\max\{j|\theta_{j}\!<\!\theta^{\prime}\}$ and $\theta^{\prime}$ satisfies $\beta^{2}\!<\!\theta^{\prime}<1$.

The above assumption on the hyperparameters is used to establish the convergence of adaptive stochastic type gradient method like Zou et al. (Zou et al., 2019). In this paper, we use a simplified setting for momentum parameter $\beta_{t}$, exponential moving average parameter $\theta_{t}$, and the base learning rate $\alpha_{t}$ to simplify the convergence analysis, compared with the sufficient condition in Zou et al. (Zou et al., 2019).

Let $\small Q_{x}\left(x\right)=x$. The quantized generic Adam reduces to be generic Adam with gradient quantization and error-feedback. Below, we present the convergence rate of Algorithm 1 in the single-machine mode.

Let $\small\{x_{t}\}$ be the point generated by Algorithm 1 with $\small Q_{x}\left(x\right)=x$. In addition, let $x_{\tau}^{T}$ represent random variable $x_{\tau}$ with $\tau$ taking from $\small\{1,2,\ldots,T\}$ with the same probability. If Assumptions 1, 2, 4 further hold, the convergence result of Algorithm 1 holds as follows:

where $\small C=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}\right)$, $\small C^{\prime}=\frac{2\sqrt{G^{2}+\epsilon d}C_{3}}{\left(1-\beta\right)\alpha}$, and $\small C_{3}=\frac{1}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{L\left(2-\delta_{g}\right)G^{2}\alpha^{2}}{\epsilon\delta_{g}}+C_{2}\theta\right)$.

This theoretical result shows that with gradient quantization and error-feedback the proposed algorithm can converge to the first-order stationary point in the nonconvex setting. In addition, the convergence rate is of the same order as the original Adam in the nonconvex setting (Zou et al., 2019). Besides, paying attention to the constant $C_{3}$, it can be seen that the constant factor appearing in the convergence rate is related to the quantized level.

This result shows Algorithm 1 with gradient quantization and error feedback technique can convergence in the same order as some popular method such as stochastic gradient descent and vanilla Adam.

1.2. Weight Quantization

In this subsection, we set $\small Q_{g}\left(g\right)=g$ in Algorithm 1. The proposed quantized generic Adam reduces to generic Adam with the weight quantization. In this situation, to establish the convergence rate of Algorithm 1 are given below.

Let $\{x_{t}\}$ be the point generated by Algorithm 1 with $\small Q_{g}\left(g\right)=g$. In addition, let $x_{\tau}^{T}$ represent random variable $x_{\tau}$ with $\tau$ taking from $\small\{1,2,\ldots,T\}$ with the same probability. If Assumptions 1, 3, 4 further hold, the convergence result of Algorithm 1 holds as follows:

where $\small C_{5}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}\right)$, $\small C_{6}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{LG^{2}\alpha^{2}}{\epsilon}+C_{2}\theta\right)$, $\small C_{7}=\frac{8\delta_{x}\sqrt{G^{2}+\epsilon d}LG}{\left(1-\beta\right)\sqrt{\epsilon C_{1}}\left(1-\sqrt{\gamma}\right)}$, and $\small C_{2}$ is defined in Theorem 3.1.

This theoretical result shows that with weight quantization the algorithm will converge to the point related to the quantized level, and when we don’t use quantization the proposed algorithm will converge to the first-order stationary point by setting $\delta_{x}=0$ directly. In addition, weight quantization on stochastic type gradient methods has already been considered in Khaled et al.(Khaled and Richtárik, 2019), in which the authors also showed that weight quantized SGD converges to a point near the global optimum. However, the analysis of weight quantized SGD in Khaled et al. (Khaled and Richtárik, 2019) is merely restricted to the strongly convex setting.

This result shows Algorithm 1 with weight quantization can convergence to the point near the stationary point due to quantization, but the speed to near stationary is in the same order as some popular method such as stochastic gradient descent and vallina Adam.

2. Multi-Worker Analysis

In this subsection, we extend Algorithm 1 to the multi-worker setting via the parameter server model. Below, we use Algorithm 2 to represent the iteration schemes of the distributed quantized generic Adam algorithm in the parameter server, and Algorithm 3 to represent the iteration schemes in all workers, respectively. Here, we assume that all the workers work independently.

Note that communicated information $\hat{x}_{t}$ and $\delta_{t}^{i}$ between the server and works is all quantized in order to improve the communication efficiency. The weight quantization procedure is performed on the server, while the gradient quantization and error-feedback procedures are performed on the workers. Below, we establish the convergence rates of distributed Adam with weight quantization, gradient quantization, and error-feedback in Algorithms 2-3 in the parameter server model.

Let $\{x_{t}\}$ be the point generated by Algorithms 2-3. In addition, let $\small\hat{x}^{T}_{\tau}$ be the random variable $\hat{x}_{\tau}$ with $\tau$ taking from $\small\{1,2,\ldots,T\}$ with the same probability. If Assumptions 1-4 hold and the iterates $\small\|x_{t}\|\leq D$ are upper bounded, the convergence result of Algorithms 2-3 holds as follows:

where $\small C_{8}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}\right)$, $\small C_{9}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{L\left(2-\delta_{g}\right)G^{2}\alpha^{2}}{\epsilon\delta_{g}}+C_{2}\theta\right)$, $\small C_{10}=\frac{4\sqrt{G^{2}+\epsilon d}\delta_{x}LG}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)\sqrt{\epsilon}\left(1-\beta\right)}$, and $\small C_{2}$ is defined in Theorem 3.1.

In Algorithms 2-3, both the gradient quantization and weight quantization schemes are applied. We also show that the proposed algorithms converge to a point near the saddle point of problem (1) up to a constant. It is noted that the constant is affected by both the gradient quantized level $\delta_{g}$ and the weight quantized level $\delta_{x}$. In addition, the limit point of the generate iterates will be influenced merely by the weight quantized level. Once gradient quantization and weight quantization reduce to identity mappings, Algorithms 2-3 reduce to the distributed Adam in the parameter server model and Theorem 3 provides their convergence rates.

To close this section, we give several comments on the proposed Algorithms 1-3. Different from distributed Adam where each worker transmits gradient to the parameter server and the parameter server calculates learning rate and update vector, we calculate the learning rates and update vector in local. Therefore, the error feedback technique can be applied to the adaptive algorithm. However, the proof will be complicated due to $N$ different learning rates being involved in the algorithm, and the following section will give a detailed proof of the above theorems.

Proof Details

In this section, we provide the detailed proof procedures of Theorem 3.1, Theorem 3.3, Theorem 3.5 and the related corollaries of the main theorems.

Before providing the detailed proof of Theorem 3.1, we first denote several useful notations. Then, we provide several lemmas that are used to split the main proof of Theorem 3.1 for better readability.

Given two positive sequences $\{a_{i}\}_{i=1}^{n}$ and $\{b_{i}\}_{i=1}^{n}$, it holds that

where the second inequality is the arithmetic inequality with positive numbers $a_{i}b_{j}$ and $a_{j}b_{i}$. ∎

By using Notation 1 and the iteration scheme of Theorem 3.1, for all $t\geq 1$ the following inequality holds:

By using the definition of $m_{t}$ in Theorem 3.1, it directly holds that $\small m_{t}=\sum_{i=1}^{t}\beta^{t-i}\left(1-\beta\right)g_{i}$. Let $\small\Theta\left(t,i\right)=\prod_{j=i+1}^{t}\theta_{j}$ for $i<t$, and $\small\Theta\left(i,i\right)=1$. According to the definition of $v_{t}$, it holds that $\small v_{t}=\sum_{i=1}^{t}\left(\prod_{j=i+1}^{t}\theta_{j}\right)\left(1-\theta_{i}\right)g_{i}^{2}=\sum_{i=1}^{t}\Theta\left(t,i\right)\left(1-\theta_{i}\right)g_{i}^{2}$.

Let $\tau$ be randomly chosen from $\small\{1,2,\cdots,T\}$ with equal probabilities $\small p_{\tau}=\frac{1}{T}$. We have the following estimate:

Note that $\small\|\hat{v}_{t}\|_{1}=\theta_{t}\|v_{t-1}\|_{1}+\left(1-\theta_{t}\right)\|\sigma_{t}\|^{2}$ and $\small\|g_{t}\|\leq G$. It is straightforward to prove $\small\|v_{t}\|_{1}\leq G^{2}$. Hence, we have $\small\|\hat{v}_{t}+\epsilon\|_{1}\leq G^{2}+\epsilon d$.

Then, by using the definition of $x_{\tau}$, we obtain

By using Notation 1, the following inequality holds:

By using the definition of $m_{t}$, it holds $\small\|m_{t}\|^{2}\leq G^{2}$.

Then, $\small\|\Delta_{t}\|^{2}=\|\frac{\alpha_{t}m_{t}}{\sqrt{v_{t}+\epsilon}}\|^{2}\leq\frac{G^{2}}{\epsilon}\alpha_{t}^{2}$ by using the definition of $\small\Delta_{t}$.

Therefore, $\small\sum_{t=1}^{T}\|\Delta_{t}\|{{}^{2}}\leq\frac{G^{2}}{\epsilon}\sum_{t=1}^{T}\frac{\alpha^{2}}{t}$. ∎

By the iteration scheme of Algorithm 1, it holds that

By the definition of noisy term $e_{t}$ and $\small\Delta_{t}$, it holds

By the definition of $\small M_{k}$, it holds that

To split $\small M_{t}$, first we introduce the following two equalities. Using the definitions of $v_{t}$ and $\hat{v}_{t}$, we obtain

In addition, it is not hard to check that the following equality holds:

where the equalities hold according to the following inequities, respectively,

where the equality holds according to $\small\left|\frac{\left(1-\beta\right)\sigma_{t}}{\sqrt{v_{t}+\epsilon}+\sqrt{\hat{v}_{t}+\epsilon}}\right|\leq\left|\frac{\left(1-\beta\right)\sigma_{t}}{\sqrt{\hat{v}_{t}+\epsilon}}\right|\leq\frac{1-\beta}{\sqrt{1-\theta_{t}}}.$

where the inequality holds according to $\small\sqrt{\hat{\eta}_{t}}\epsilon\leq\sqrt{\frac{\alpha_{t}\epsilon^{2}}{\sqrt{\hat{v}_{t}+\epsilon}}}\leq\sqrt{\frac{\alpha\epsilon^{3/2}}{\sqrt{\theta}}}.$

where the equalities hold according to $\small\sqrt{\hat{\eta}_{t}}\sqrt{\sigma_{t}^{2}+\epsilon}=\sqrt{\frac{\alpha_{t}\left(\sigma_{t}^{2}+\epsilon\right)}{\sqrt{\hat{v}_{t}+\epsilon}}}\leq\sqrt{\frac{\alpha\sqrt{G^{2}+\epsilon}}{\sqrt{\theta}}},$ and

Therefore, by combining the above upper estimations for the seven terms in Eq. (3), we obtain

Recalling the definition of $\small M_{t}$. For $\small M_{1}$, we have

where the last inequality holds by using $A^{1}_{t}$ and $A^{2}_{t}$.

Based on Notation (1) and the above lemmas, then we can prove Theorem 1. First, according to the gradient Lipschitz condition of $f$, it holds

Using the above lemmas and arranging the corresponding terms, we have

where $\small C_{3}=\frac{1}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{L\left(2-\delta_{g}\right)G^{2}\alpha^{2}}{\epsilon\delta_{g}}+C_{2}\theta\right)$, $\small C=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}\right)$, $\small C^{\prime}=\frac{2\sqrt{G^{2}+\epsilon d}C_{3}}{\left(1-\beta\right)\alpha}$, respectively. Hence, the proof is completed. ∎

For a fixed iteration $T$, let $\small\alpha_{t}=\frac{\alpha}{\sqrt{T}}$ and $\small\theta=1-\frac{\theta}{T}$. When $\small\alpha_{t}=\frac{\alpha}{\sqrt{T}}$, Lemma 4.4 will update to $\small\sum_{t=1}^{T}\|\Delta_{t}\|^{2}\leq\frac{G^{2}\alpha^{2}}{\epsilon}$. By the same proof in Lemma 4.6, we have

Based on the proof of Theorem 3.1, we have

By bounding $\small\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha\sqrt{T}}\left(f\left(x_{1}\right)-f^{*}+\frac{L\left(2-\delta_{g}\right)G^{2}\alpha^{2}}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)\epsilon\delta_{g}}+\frac{C_{2}\theta}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\right)$ by $\xi$, we obtain $\small T=O\left(\frac{1}{\xi^{2}}\right)$. ∎

2. Proof of Theorem 3.3

To prove Theorem 3.3, we first define some notations and provide several useful lemmas.

By using Notation 2, the following inequality holds:

Based on Notation 2, we have the following upper estimation for $\small M_{t}$,

Denoting the same notations $\small A^{1}_{t},A^{2}_{t},A^{3}_{t},A^{4}_{t},A^{5}_{t}$ in Lemma 4.6, we have

The rest six parts remain the same as Lemma 4.6. Then we have

where $\small C_{2}\!=\!\frac{5\alpha G^{3}\left(1\!-\!\beta\right)}{2\epsilon\sqrt{\theta}}\left(\frac{\beta}{\left(1\!-\!\beta\right)\sqrt{\theta_{1}C_{1}\left(1\!-\!\gamma\right)}}\!+\!1\right)^{2}\!+\!\frac{5\alpha G^{3}}{2\epsilon\sqrt{\theta}}\!+\!\frac{5\beta^{2}\alpha d\sqrt{\epsilon}}{2\sqrt{\theta}\left(1\!-\!\beta\right)\theta_{1}C_{1}\left(1\!-\!\gamma\right)}\!+\!\frac{5\alpha\sqrt{G^{2}+\epsilon}G^{2}\beta^{2}}{2\left(1\!-\!\beta\right)\sqrt{\theta}\theta_{1}C_{1}\left(1\!-\!\gamma\right)\epsilon}\!+\!\frac{5\alpha\sqrt{G^{2}+\epsilon}\beta^{2}d}{2\left(1\!-\!\beta\right)\sqrt{\theta}\theta_{1}C_{1}\left(1\!-\!\gamma\right)}$.

By the Lipschitz continuity of the gradient, we have

Arranging the corresponding terms suitably, we obtain

where $\small C_{5}$, $\small C_{6}$, and $\small C_{7}$ are defined as $\small C_{5}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}\right),C_{6}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{LG^{2}\alpha^{2}}{\epsilon}+C_{2}\theta\right),$ and $\small C_{7}=\frac{8\delta_{x}\sqrt{G^{2}+\epsilon d}LG}{\left(1-\beta\right)\sqrt{\epsilon C_{1}}\left(1-\sqrt{\gamma}\right)}$, respectively.

For a fixed iteration $T$, let $\small\alpha_{t}=\frac{\alpha}{\sqrt{T}}$ and $\theta_{t}=1-\frac{\theta}{T}$. It is not hard to check that Lemma 4.7 updates to $\small\sum_{t=1}^{T}\|\Delta_{t}\|\leq\frac{\alpha G\sqrt{T}}{\sqrt{\epsilon}}$, and Lemma 4.8 updates to

where $\small C_{5}^{\prime}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}+\frac{LG^{2}\alpha^{2}}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)\epsilon}+\frac{C_{2}\theta}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\right)$, $\small C_{7}^{\prime}=\frac{4\delta_{x}\sqrt{G^{2}+\epsilon d}LG}{\left(1-\beta\right)\sqrt{\epsilon C_{1}}\left(1-\sqrt{\gamma}\right)}$.

Hence, by bounding $\small\frac{C_{5}^{\prime}}{\sqrt{T}}+C_{7}^{\prime}\leq C_{7}^{\prime}+\xi$, we obtain the desired result. ∎

3. Proof of Theorem 3.5

To prove Theorem 3, we first denote a few notations and provide several useful lemmas.

With Notation 3, we derive an upper estimation for $\small\hat{\Delta}_{t}$ as: $\small\sum_{t=1}^{T}\|\hat{\Delta}_{t}\|^{2}\leq\frac{G^{2}}{\epsilon}\sum_{t=1}^{T}\frac{\alpha^{2}}{t}.$

By using the definition of $\small\hat{\Delta}_{t}$, it is not hard to check that the following equations hold:

Let $e_{i}$ be the noisy term in Algorithm2 and $\small\hat{\Delta}_{t}$ be the term defined in Notation 3. Then it holds that

Using the definition of the noisy term $e_{t}$, the following holds:

Let $\tau$ be randomly chosen from $\small\{1,2,\cdots,T\}$ with equal probabilities $\small p_{\tau}=\frac{1}{T}$. We have the following estimate:

By the definition of $\small M_{t}$, we obtain its upper-estimation as follows:

Based on the similar proof of $\small\Delta_{t}-\frac{\beta\alpha_{t}}{\sqrt{\theta_{t}}\alpha_{t}}\Delta_{t-1}$ in Lemma 4.6, we define

Then, via the same proof as Lemma 4.6, the following equation holds

For the first term in Eq. (5), it holds that

The remain 6 terms in Eq. (5) are the same as Lemma 4.6. Thus, we have

By using the gradient Lipschitz continuity of $f$, it holds that

Taking summation over both sides of the above inequality, it holds that

where $\small C_{8}\!=\!\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1\!-\!\beta\right)\alpha}\left(f\left(x_{1}\right)\!-\!f^{*}\right)$,$\small C_{9}\!=\!\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1\!-\!\beta\right)\alpha\sqrt{C_{1}}\left(1\!-\!\sqrt{\gamma}\right)}\left(\frac{L\left(2\!-\!\delta_{g}\right)G^{2}\alpha^{2}}{\epsilon\delta_{g}}\!+\!C_{2}\theta\right)$, $\small C_{10}\!=\!\frac{8\sqrt{G^{2}+\epsilon d}\delta_{x}LG}{\sqrt{C_{1}}\left(1\!-\!\sqrt{\gamma}\right)\sqrt{\epsilon}\left(1\!-\!\beta\right)}$, respectively.

Given iteration $T$, let $\small\alpha_{t}\!=\!\frac{\alpha}{\sqrt{T}}$ and $\small\theta_{t}\!=\!1\!-\!\frac{\theta}{T}$. Lemma 4.12 updates to

where $\small C_{8}^{\prime}=\frac{2\sqrt{G^{2}+\epsilon d}}{\left(1-\beta\right)\alpha}\left(f\left(x_{1}\right)-f^{*}+\frac{1}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)}\left(\frac{L\left(2-\delta_{g}\right)G^{2}\alpha^{2}}{\epsilon\delta_{g}}+C_{2}\theta\right)\right)$, $\small C_{10}^{\prime}=\frac{4\sqrt{G^{2}+\epsilon d}\delta_{x}LG}{\sqrt{C_{1}}\left(1-\sqrt{\gamma}\right)\sqrt{\epsilon}\left(1-\beta\right)}$.

Hence, by bounding $\small\frac{C_{8}^{\prime}}{\sqrt{T}}+C_{10}^{\prime}\leq C_{10}^{\prime}+\xi$, we obtain the desired result. ∎

Experiments

In this section, we apply the proposed algorithms to train deep neural networks including VGG16 network (Simonyan and Zisserman, 2014) and ResNet-101 (He et al., 2016) on CIFAR10 (Krizhevsky et al., 2009) and CIFAR100 (Krizhevsky et al., 2009) datasets, respectively. Below, we first describe the implementation details and then show the experimental results.

Besides, in all experiments we choose $\beta$ as 0.99, $\theta$ as 0.999, and $\epsilon$ as $1e-5$. Since it is natural to choose the exponential decay strategy on learning rate, we choose to reduce $\alpha$ by half every 50 epochs instead of using ${\alpha}/{\sqrt{t}}$ as (Reddi et al., 2019). We choose the starting learning rate as 0.001. The value is chosen based on grid search on the set $\{0.01,0.001,0.0001\}$ based on the accuracy of the full precision setting. For gradient quantization, we use function $\small Q_{g}(g)=\|g\|_{\infty}\arg\min_{\hat{g}\in\mathcal{G}^{d}}\|g/\|g\|_{\infty}-\hat{g}\|$, where $\small\mathcal{G}=\{-1,\cdots,-2^{-k_{g}},0,2^{-k_{g}},2^{-k_{g}+1},\cdots,1\}$. For weight quantization, we use function $\small Q_{x}(x)=0.5\times\arg\min_{\hat{x}\in\mathcal{X}}\|2x-\hat{x}\|$, where $\small\mathcal{X}=\left\{-1,\cdots,-\frac{1}{2^{k_{x}}},0,\frac{1}{2^{k_{x}}},\frac{2}{2^{k_{x}}},\cdots,1\right\}$.

We compared our method with TernGrad (Wen et al., 2017) and Zheng et. al(Zheng et al., 2019) for gradient quantization, where the learning rate in these two methods is 0.1 which generated by grid search in $\{0.1,0.05,0.01\}$. For weight quantization, we compare with the result which quantizes the final model directly named WQuan in tables.

2. Results Illustration

In this section, we will illustrate the results of training ResNet-101 on the CIFAR100 dataset and training VGG16 on the CIFAR10 dataset, respectively. Two tables show the test accuracy after 200 epochs of training, where the first column represents training methods, the second column represents the test accuracy, the third column represents bits required for gradient communication, and the last column represents the bits to save a model. As for the same method, we can set different $k_{x}$ and $k_{g}$, and we can get a different number of bits needed for gradients and weights.

Figure 4 and Table 2 show the result of training ResNet-101 on the CIFAR100 dataset. When the algorithm involves gradient quantization, we compare our algorithm with Zheng et al.(Zheng et al., 2019) and TernGrad (Wen et al., 2017) with a different number of communication bits, which has been shown in the left figure of Figure 4 and the first 9 lines in Table 2. It can be shown that even with gradient quantization, our algorithm can outperform TernGrad(Wen et al., 2017) and Zheng et al. (Zheng et al., 2019). The middle figure in Figure 4 shows the result of using weight quantization only. Row 10-12 in Table 2 shows the comparison results between quantizing weight during training and after training, which shows when quantizing weight during the training process can achieve higher test accuracy. The right figure in Figure 4 and the last 4 rows demonstrate the results of combining gradient quantization and weight quantization. It shows even though we shrink model size into 1/4 of its original size and gradient size into 1/16 of its original size, respectively, it can still give comparable results.

2.2. Results of Training VGG16 on the CIFAR10

Figure 4 and Table 3 show the result of training VGG16 on the CIFAR10 dataset. The left figure in Figure 4 and the first 9 rows in Table reftable2 show results of gradient quantization comparison among our algorithm, TernGrad(Wen et al., 2017) and Zheng et al.(Zheng et al., 2019). Our results and Zheng et al.(Zheng et al., 2019) can achieve similar performance, but TernGrad (Wen et al., 2017) gets worse performance due to noise to ensure unbiasedness. When considering weight quantization, shown in the middle figure of Figure 4 and row 10-13 in Table 3, performance is similar between quantizing during training or after training. Further, as it has shown in the last 4 rows with different sizes of the model and different gradient quantization, our method can still achieve high accuracy compared to the full precision version (first row in Table 3).

Conclusions

To accelerate the training process of deep learning models, we proposed distributed Adam with weight quantization, gradient quantization, and the error-feedback technique in the parameter-server model. Through capitalizing on the two schemes of weight quantization and gradient quantization, the communication cost between the server and works can be significantly alleviated. In addition, the proposed error-feedback technique can suppress the bias caused by the gradient quantization step, thereby making the proposed algorithms more efficient. We further established the convergence rates of the proposed algorithms in the nonconvex stochastic setting and showed that quantized Adam with the error-feedback technique converges to the neighborhood of a stationary point under both the single-worker and multi-worker modes. Moreover, we applied the proposed algorithms to train VGG16 on the CIFAR10 dataset and ResNet-101 on the CIFAR100 dataset, respectively. The experiments demonstrate the efficacy of the proposed algorithms.

References