Quantized Adaptive Subgradient Algorithms and Their Applications

Stochastic Optimization (Lee:2018:DQA:3219819.3220075, ) is a popular approach in training large-scale machine learning models. We denote the objective function as $f({\bf x})$, depending on the model parameter ${\bf x}\in\chi$. The dataset could be split into $M$ workers, where each has $N$ data samples. We denote $f_{m,n}$ as the loss function associated with the $n^{th}$ sample on the $m^{th}$ worker, and $f({\bf x})$ is an average over all loss functions on all workers.

When training on a single machine, the stochastic optimization, like SGD, randomly chooses a subset of samples ${\bf{Z}}_{t}$ and calculates the gradient $f^{\prime}_{t}({\bf x}_{t},{\bf z}_{t})$, where ${\bf z}_{t}\in{\bf{Z}}_{t}$, to approximates the true gradient $f^{\prime}_{t}({\bf x}_{t})$ at $t^{th}$ iteration by

Similarly, in a distributed setting, each worker calculates the stochastic gradient $f^{\prime}_{t}({\bf x}_{t})^{(m)}$ based on its local data ${\bf{Z}}_{t}^{(m)}$. The parameter server averages the local gradients to get the synchronized gradient,

As previous work on gradient sparsification and quantization suggested, we could represent $f^{\prime}_{t}({\bf x}_{t})^{(m)}$ in a low-precision way. Since extreme gradient sparsification can also be seen as gradient quantization, this paper focuses on reducing communication costs through gradient quantization. We denote $Q(\cdot)$ as the quantization function and then the synchronized quantized gradient can be obtained by

The update formula for the synchronous distributed stochastic optimization with function $Q(\cdot)$ is

where $\eta$ is the learning rate.

General Framework. In the data parallelism distributed training scheme, the number of the exchanged gradients is equal to the number of model learnable parameters. But for a sparse model, we can avoid transmitting unnecessary gradient information, i.e. gradient corresponding to parameters that are always 0 before and after model update. To this end, each worker utilizes their local full precision gradients to obtain the model parameters after the local update. With the assumption that the model parameters change smoothly and be sparsity, we can transmit a small amount of the necessary gradient to save the communication cost. Furthermore, this small amount of the local gradients is quantized and encoded. After the global parameter server receives and decodes the quantized gradients, synchronous gradients are calculated by averaging the local sparse quantized gradients. The synchronous sparse gradients are then quantized and sent to the workers for model parameter update. The general framework of QCMD adagrad and QRDA adagrad is shown in Fig. 2.

Generate Sparse Model. Vanilla SGD is not particularly effective at producing a sparse model, even adding a subgradient of $l_{1}$ penalty to the gradient of the loss. Some other approaches such as proximal gradient descent (also named composite mirror descent) (singer2009efficient, ; duchi2010composite, ) and regularized dual averaging (xiao2010dual, ) introduce limited sparsity with $l_{1}$ regularization. Their adaptive subgradient extensions, CMD adagrad and RDA adagrad, (duchi2011adaptive, ) produce better accuracy vs. sparsity tradeoffs.

The original proximal gradient method (duchi2010composite, ) employs an immediate trade-off among the current gradient $f_{t}^{\prime}({\bf x}_{t})$, the regularizer $\phi$ and proximal function $\psi$. The proximal function $\psi$ aims to keep ${\bf x}$ staying close to ${\bf x}_{t}$ and sometimes be simply set to $\psi({\bf x})=\frac{1}{2}||{\bf x}-{\bf x}_{t}||^{2}_{2}$). It makes the model parameters meet the assumption of steady change. To achieve better bound of regret, the CMD adagrad adopts $\psi_{t}$ which varies with $t$. The update for CMD adagrad amounts to solving

Similarly, the RDA adagrad encompasses a trade off among a gradient-dependent linear term, the regularizer $\phi$ and a strongly convex term $\psi_{t}$ for well conditioned predictions. The update for RDA adagrad is shown as:

For both CMD adagrad and RDA adagrad, $\phi({\bf x})=\lambda||{\bf x}||_{1}$ encourages the sparsity of the model parameter. The Bregman divergence associated with $\psi_{t}({\bf x})=\frac{1}{2}{\bf x}^{\top}{\bf A}_{t}{\bf x}$ is defined as

${\bf A}_{t}$ is a diagonal matrix also named the adaptive learning rate matrix. Concretely, for some small fixed $\delta\geq 0$, the adaptive learning rate matrix is:

where ${\bf a}_{t,d}=||f^{\prime}({\bf x})_{1:t,d}||_{2}$, $f^{\prime}({\bf x})_{1:t}=[f^{\prime}({\bf x})_{1:t-1},f^{\prime}_{t}({\bf x}_{t})]$ denote the matrix obtained by concatenating the gradient sequence. $f^{\prime}({\bf x})_{1:t,d}$ is the $d^{th}$ row of this matrix and $diag(\cdot)$ converts a vector into a diagonal matrix. ${\bf I}$ is identity matrix.

To apply them to distributed learning, each worker needs to maintain a adaptive learning rate matrix locally. On the $m_{th}$ worker, the elements of the gradient $f^{\prime}({\bf x}_{t})^{(m)}$ are selected by a selected function:

where $\mathbb{I}^{(m)}$ is an indicator function that

$\hat{\bf x}_{t+1}^{(m)}$ is the local parameter that is updated based on the local full precision gradient $f^{\prime}_{t}({\bf x}_{t})^{(m)}$ and local adaptive learning rate matrix ${\bf A}_{t}^{(m)}$. Both the indicator and selected elements of the local gradients are sent to the global parameter server for decoding. The synchronous indicator is calculated by ”bitwise or” between the $M$ indicators.

Gradient quantization. Besides reducing the transmission amount of the corresponding gradient by generating a sparse model, we also hope to further reduce the communication cost by quantifying the selected elements of the gradient. The quantization function is define as $Q(\cdot)$. Each worker computes $Q(S(f^{\prime}_{t}({\bf x}_{t})^{(m)},\mathbb{I}^{m}))$. The global parameter server first decodes the quantized gradient through $S^{-1}$ and computes the synchronized quantized gradient based on Eq.(4). We furtherly quantize the synchronized gradient $SynQ_{t}$ to save more communication costs by

The server pushes the double quantized gradient ${\bf q}_{t}$ to every worker. Specially, the adaptive learning rate is calculated by the worker according to ${\bf q}_{t}$ sequence. So we construct a quantized gradient based adaptive learning rate matrix,

where ${\bf c}_{t,d}=||{\bf q}_{1:t,d}||_{2}$, ${\bf q}_{1:t}=[{\bf q}_{1:t-1},{\bf q}_{t}]$. ${\bf H}_{t}$ dynamically incorporates knowledge of the geometry of the data and the curvature of the loss function based on quantized gradient, which makes QCMD adagrad and QRDA adagrad achieve a good balance between model sparsity, accuracy and communication cost. Therefore, the objective function for QCMD adagrad becomes

The objective function for QRDA adagrad becomes

Solving Eq. (14), we have the update rule for QCMD adagrad:

The subscript $(\cdot)_{+}$ represents that we retain the value greater than 0. Solving Eq. (15), the update rule for QRDA adagrad becomes:

The overall QCMD adagrad and QRDA adagrad schemes are presented in Alg. 1.

In this section, we theoretically analyse that the quantization error introduced by the gradient quantization affects the convergence rate of regret for QCMD adagrad and QRDA adagrad.

Remark: Since $\mathbb{E}_{\bf q}||{\bf q}_{t}-f^{\prime}_{t}({\bf x}_{t})||^{2}_{\psi^{*}_{t}}$ is the quantization variance scaled by the adaptive learning rate and

We can simply regard $\mathbb{E}_{\bf q}||{\bf q}_{t}||^{2}_{\psi^{*}_{t}}$ as the error introduced by quantization for QCMD adagrad and regard $\mathbb{E}_{\bf q}||{\bf q}_{t}||^{2}_{\psi^{*}_{t-1}}$ as the error introduced by quantization for QRDA adagrad. Therefore, Proposition 1. and Proposition 2. show that gradient quantization introduces additional noise which affects the model convergence and sparsity.

Although gradient quantization can reduce the cost of gradient communication in distributed training, it also introduces additional errors, which affect the convergence of the model and the sparsity of parameters in QCMD adagrad and QRDA adagrad. As an unbiased gradient quantization method, TernGrad (wen2017terngrad, ) has already made a good balance between the encoding cost and accuracy of the general model. However, when it comes to the sparse model, a large quantization error still leads to slower convergence of the $l_{1}$ norm as a part of the objective function, which affects the sparsity of the model. In order to mitigate this problem, we apply the threshold quantization method to the QCMD adagrad and QRDA adagrad.

Threshold quantization is an existing quantization method used for model quantization in (TWN, ). We apply it to gradient quantization since it produces less error than Terngrad(wen2017terngrad, ). In this section, we use ${\bf v}^{t}$ to represent the (stochastic) gradient in the $t_{th}$ iteration. Fig. 3 gives a brief explanation of threshold quantization, and more analysis is provided below.

$Q_{\triangle}(\cdot)$ is the threshold quantization function defined as below

where ${\bf v}$ is a ternary vector, $s$ is a non-negative scaling factor and $\triangle$ denotes the threshold. For the $i^{th}$ component of ${\bf v}$,

The error ${\epsilon}_{t}$ is defined as the difference between the full precision vector ${\bf v}^{t}$ and $Q_{\triangle}({\bf v}^{t})$:

In order to keep as more information of ${\bf v}^{t}$ as possible, the quantization method is required to minimize the Euclidean distance between ${\bf v}^{t}$ and $Q_{\triangle}({\bf v}^{t})$, i.e.,

Obviously, the above problem can be transformed to the following formulation:

where $I_{\Delta}=\{i||{\bf v}^{t}_{i}|>\Delta\}$ and $|I_{\Delta}|$ denotes the number of elements in $I_{\Delta}$. Thus, for any given $\Delta$, the optimal $s$ can be computed as follows,

The optimal $\Delta$ can be computed as follows,

To solve the above optimal threshold, we can rank the components of the gradient and treat them as potential thresholds. For each potential threshold, we calculate the corresponding objective value $\frac{1}{|I_{\Delta}|}(\sum_{i\in I_{\Delta}}|{\bf v}^{t}_{i}|)^{2}$ and take the potential threshold that maximizes the objective value as the optimal threshold. The computational complexity of this process is $O(d\log d)$ for $d$-dimensional gradients, where the major cost is the sorting over all elements of gradients.

We try to improve the computational efficiency of the threshold quantization procedure without harming the optimality of the coding described in Eq. (24). This improvement begins with the assumption that the gradient follows the Gaussian distribution. In case of ${\bf v}^{t}_{i}$ follows $N(0,\sigma^{2})$, Li and Liu (TWN, ) have given an approximate solution for the optimal threshold $\triangle^{*}$ by $0.6\sigma$ which equals $0.75\cdot\mathbb{E}(|{\bf v}^{t}_{i}|)\approx\frac{0.75}{d}\sum_{i=1}^{d}|{\bf v}_{i}^{t}|$. We also find in our experiment that most of the gradients satisfy the previous assumption. Fig. 4 shows the experimental result for training AlexNet (krizhevsky2012imagenet, ) on two workers. The left column visualizes the first convolutional layer and the right one visualizes the first fully-connected layer. The distribution of the original floating gradients is close to the Gaussian distribution for both convolutional and fully-connected layers. Based on this observation, we simply use $\frac{0.75}{d}\sum_{i=1}^{d}|{\bf v}_{i}^{t}|$ to approximate the optimal threshold to avoid the expensive cost of solving the optimal threshold $\Delta^{*}$ every iteration.

Encode. The gradient vectors need to be encoded after the threshold quantization. Specifically, we use 2 bits to encode ${\bf v}_{i}$, and one floating-point number to represent the scaling factor s. For the dense model, the overall communication cost is $(32+2d)$ bits, where $d$ denotes the parameter dimension. The communication cost of threshold quantization is the same as Terngrad. But for the sparse model, assume the number of non-zero parameters is $k$ $(k\ll d)$, the communication cost of QCMD adagrad and QRDA adagrad is $(32+d+2k)$ bits. Since we need at least $d$ bits to indicate which components of the parameter are non-zero.

In this section, we first conduct experiments on linear models to validate the effectiveness and efficiency of the proposed QCMD adagrad and QRDA adagrad for binary classification problems. After that, we use our proposed method to train convolutional neural networks, in order to validate the performance of them on non-convex problems.

Baselines. We compare QCMD adagrad and QRDA adagrad with several sparse model distributed optimization methods, including 32 bits Prox-gd(duchi2010composite, ), 32 bits CMD adagrad(duchi2011adaptive, ), 32 bits RDA adagrad(duchi2011adaptive, ) and their corresponding ternary variants(wen2017terngrad, ). ${\dagger}$ is to mark the methods that only the local gradients are quantized but the synchronous gradient.

Implementation Details. All the experiments are carried out in a distributed framework that the network bandwidth is 100MBps. For the linear models training, each worker only utilizes the CPUs. For the training of convolutional neural networks, each worker is allocated with 1 NVIDIA Tesla P40 GPU. The methods are evaluated on four publicly available datasets. news20 and rcv1 are text datasets with a high dimension of input features from LIBSVM (chang2001libsvm, ). MNIST (lecun1998gradient, ) is for handwritten digits classification problem and CIFAR10 (krizhevsky2009learning, ) is for image classification problem. Table 1 shows the details of datasets. For news20 and rcv1, the $\ell_{1}$ norm regularized logistic regression model are trained. As for multi-classification problems, we train LeNet for MNIST(lecun1998gradient, ) and AlexNet (krizhevsky2012imagenet, ) for CIFAR10. To generate a network with large sparsity, the batch normalization layer is added before each convolution layer and fully connected layer. The code is implemented via tensorflow. Experimental results are averaged over 5 runs with a random initialization seed.

Hyperparameter Discussion. Table 2 lists the related hyperparameters of the sparse model training in distributed settings, including the learning rate $\eta$, coefficient of regularization $\lambda$, batch size per worker and the number of the worker. The optimal choice of $\eta$ and $\lambda$ can vary somewhat from dataset to dataset. The structure and random initial of the network also affect them. In order to find the appropriate hyperparameters, we first select the best learning rate based on cross-validations without the regularization and then find the coefficient of regularization that maximizes the sparsity of the model without reducing the accuracy as much as possible. For ease of comparison, we set the same hyper parameters for corresponding ternary quantization methods. Fig. 5 shows the accuracy and model sparsity as the $\eta$ and $\lambda$ change on MINIST. The sparsity increase with the $\lambda$. A strong coefficient of regularization $\lambda$ enforces a highly sparse model which may deteriorate the accuracy. When the learning rate is high or $\lambda$ is large, it is easy to cause oscillation for QCMD adagrad. For QCMD adagrad, the model sparsity increases not only with the coefficient of regularization $\lambda$ but also the learning rate $\eta$. For QRDA adagrad, model sparsity is mainly affected by $\lambda$, while the learning rate has less effect.

We summarize three metrics including the accuracy, model sparsity and error of the quantized gradient in Table 3, and report the time consumption in Fig. 6. From these results, we draw several conclusions.

Firstly, the sparsity of model on the four datasets is 98.18%, 97.39%, 98.69% and 81.86% for QRDA adagrad, respectively. QRDA adagrad is easier to generate extremely sparse models than QCMD adagrad, since it uses the same evaluation criteria $\lambda$ for each dimension of the model parameters to determine whether to retain the parameter values. In the case of highly sparse model for QRDA adagrad, the accuracy decreases slightly compared with QCMD adagrad whose sparseness of model is relatively low. For example, on MNIST, the accuracy for threshold QRDA adagrad achieves 97.70%, which is $0.76\%$ lower than that of threshold QCMD adagrad.

Secondly, although the sparsity of the model generated by QCMD adagrad is limited. However, it still performs better than 32 bits Prox-gd in accuracy and sparsity, since the proximal function $\psi$ use the Bregman divergence to keep the parameter ${\bf x}_{t+1}$ close to ${\bf x}_{t}$ and it adaptively retains the parameter values for each dimension.

Thirdly, the noise introduced by the quantization method affects the convergence of accuracy and sparsity of the model. The one-way quantization (only the local gradients are quantized instead of both local gradients and the synchronous gradient) introduces less noise than the double quantization. The error of threshold quantization is smaller than Terngrad’s and it preserves sparsity and accuracy of the model both for QCMD adagrad and QRDA adagrad. Specifically, for threshold quantization, the accuracy and sparsity of the model for double quantized scheme are relatively close to 32 bits gradient scheme’s, while ternary quantization reduces the sparsity and accuracy of the model on these 4 datasets.

Lastly, double quantization reduces most of the communication cost. Fig. 6 decomposes the time consumption. For example, double quantized threshold QRDA can save $96.06\%$ of the total training time on CIFAR10. Threshold QRDA (only the local gradients are quantized) saves $48.03\%$ of the total training time on CIFAR10. The QRDA adagrad consumes the least communication cost since it generates a highly sparse model.