Adaptive Gradient Quantization for Data-Parallel SGD

Introduction

Stochastic gradient descent (SGD) and its variants are currently the method of choice for training deep models. Yet, large datasets cannot always be trained on a single computational node due to memory and scalability limitations. Data-parallel SGD is a remarkably scalable variant, in particular on multi-GPU systems . However, despite its many advantages, distribution introduces new challenges for optimization algorithms. In particular, data-parallel SGD has large communication cost due to the need to transmit potentially huge gradient vectors. Ideally, we want distributed optimization methods that match the performance of SGD on a single hypothetical super machine, while paying a negligible communication cost.

A common approach to reducing the communication cost in data-parallel SGD is gradient compression and quantization . In full-precision data-parallel SGD, each processor broadcasts its locally computed stochastic gradient vector at every iteration, whereas in quantized data-parallel SGD, each processor compresses its stochastic gradient before broadcasting. Current quantization methods are either designed heuristically or fixed prior to training. Convergence rates in a stochastic optimization problem are controlled by the trace of the gradient covariance matrix, which is referred as the gradient variance in this paper . As Fig. 1 shows, no fixed method can be optimal throughout the entire training because the distribution of gradients changes. A quantization method that is optimal at the first iteration will not be optimal after only a single epoch.

In this paper, we propose two adaptive methods for quantizing the gradients in data-parallel SGD. We study methods that are defined by a norm and a set of quantization levels. In Adaptive Level Quantization (ALQ), we minimize the excess variance of quantization given an estimate of the distribution of the gradients. In Adaptive Multiplier Quantization (AMQ), we minimize the same objective as ALQ by modelling quantization levels as exponentially spaced levels. AMQ solves for the optimal value of a single multiplier parametrizing the exponentially spaced levels.

We propose two adaptive gradient quantization methods, ALQ and AMQ, in which processors update their compression methods in parallel.

We establish an upper bound on the excess variance for any arbitrary sequence of quantization levels under general normalization that is tight in dimension, an upper bound on the expected number of communication bits per iteration, and strong convergence guarantees on a number of problems under standard assumptions. Our bounds hold for any adaptive method, including ALQ and AMQ.

We improve the validation accuracy by almost $2\%$ on CIFAR-10 and $1\%$ on ImageNet in challenging low-cost communication setups. Our adaptive methods are significantly more robust to the choice of hyperparameters.Open source code: http://github.com/tabrizian/learning-to-quantize

2 Related work

Adaptive quantization has been used for speech communication and storage . In machine learning, several biased and unbiased schemes have been proposed to compress networks and gradients. Recently, lattice-based quantization has been studied for distributed mean estimation and variance reduction . In this work, we focus on unbiased and coordinate-wise schemes to compress gradients.

Alistarh et al. proposed Quantized SGD (QSGD) focusing on the uniform quantization of stochastic gradients normalized to have unit Euclidean norm. Their experiments illustrate a similar quantization method, where gradients are normalized to have unit $L^{\infty}$ norm, achieves better performance. We refer to this method as QSGDinf or Qinf in short. Wen et al. proposed TernGrad, which can be viewed as a special case of QSGDinf with three quantization levels.

Ramezani-Kebrya et al. proposed nonuniform quantization levels (NUQSGD) and demonstrated superior empirical results compared to QSGDinf. Horváth et al. proposed natural compression and dithering schemes, where the latter is a special case of logarithmic quantization.

There have been prior attempts at adaptive quantization methods. Zhang et al. proposed ZipML, which is an optimal quantization method if all points to be quantized are known a priori. To find the optimal sequence of quantization levels, a dynamic program is solved whose computational and memory cost is quadratic in the number of points to be quantized, which in the case of gradients would correspond to their dimension. For this reason, ZipML is impractical for quantizing on the fly, and is in fact used for (offline) dataset compression. They also proposed an approximation where a subsampled set of points is used and proposed to scan the data once to find the subset. However, as we show in this paper, this one-time scan is not enough as the distribution of stochastic gradients changes during the training.

Zhang et al. proposed LQ-Net, where weights and activations are quantized such that the inner products can be computed efficiently with bitwise operations. Compared to LQ-Net, our methods do not need additional memory for encoding vectors. Concurrent with our work, Fu et al. proposed to quantize activations and gradients by modelling them with Weibull distributions. In comparison, our proposed methods accommodate general distributions. Further, our approach does not require any assumptions on the upper bound of the gradients.

Preliminaries: data-parallel SGD

The update rule for full-precision SGD is given by {\bf w}_{t+1}={\bf P}_{\Omega}\big{(}{\bf w}_{t}-\alpha g({\bf w}_{t})) where ${\bf w}_{t}$ is the current parameter vector, $\alpha$ is the learning rate, and ${\bf P}_{\Omega}$ is the Euclidean projection onto $\Omega$. We consider data-parallel SGD, which is a synchronous and distributed framework consisting of $M$ processors. Each processor receives gradients from all other processors and aggregates them. In data-parallel SGD with compression, gradients are compressed by each processor before transmission and decompressed before aggregation . A stochastic compression method is unbiased if the vector after decompression is in expectation the same as the original vector.

Adaptive quantization

We define the variance of vector quantization to be the trace of the covariance matrix,

Let ${\bf v}$ be a random vector corresponding to a stochastic gradient and ${\bf h}$ capture the randomness of quantization for this random vector as defined above. We define two minimization problems, expected variance and expected normalized variance minimization:

The following theorem suggests that solving Eq. 3 is challenging in general; however, the sub-problem of optimizing a single level given other levels can be solved efficiently in closed form. Proofs are provided in Appendix B.

Performing the update rule above sequentially over coordinates $j$ is a form of coordinate descent (CD) that is guaranteed to converge to a local minima. CD is particularly interesting because it does not involve any projection step to the feasible set ${{\cal L}}$. In practice, we initialize the levels with either uniform levels or exponentially spaced levels proposed in . We observe that starting from either initialization CD converges in small number of steps (less than $10$).

2 Gradient descent

Computing $\nabla\Psi$ using Leibniz’s rule , the gradient descent (GD) algorithm to solve Eq. 3 is based on the following update rule:

3 AMQ: Exponentially spaced levels

4 Expected variance minimization

In this section, we consider the problem of minimizing the expected variance of quantization:

To solve the expected variance minimization problem, suppose that we observe $N$ stochastic gradients $\{{\bf v}_{1},\ldots,{\bf v}_{N}\}$. Let $F_{n}(r)$ and $p_{n}(r)$ denote the CDF and PDF of normalized coordinate conditioned on observing $\|{\bf v}_{n}\|$, respectively. By taking into account randomness in $\|{\bf v}\|$ and using the law of total expectation, an approximation of the expected variance in Eq. 9 is given by

The optimal levels to minimize Eq. 10 are a solution to the following problem:

Theoretical guarantees

One can alternatively design quantization levels to minimize the worst-case variance. However, compared to an optimal scheme, this worst-case scheme increases the expected variance by $\Omega(d)$, which is prohibitive in deep networks. We quantify the gap in Appendix E. Proofs are in appendices.

We consider a general adaptively quantized SGD (AQSGD) algorithm, described in Algorithm 1, where compression schemes are updated over the course of training.Our results hold for any adaptive method, including ALQ and AMQ. Many convergence results in stochastic optimization rely on a variance bound. We establish such a variance bound for our adaptive methods. Further, we verify that these optimization results can be made to rely only on the average variance. In the following, we provide theoretical guarantees for AQSGD algorithm, obtain variance and code-length bounds, and convergence guarantees for convex, nonconvex, and momentum-based variants of AQSGD.

The analysis of nonadaptive methods in can be considered as special cases of our theorems with fixed levels over the course of training. A naive adoption of available convergence guarantees results in having worst-case variance bounds over the course of training. In this paper, we show that an average variance bound can be applied on a number of problems. Under general normalization, we first obtain variance upper bound for arbitrary levels, in particular, for those obtained adaptively.

Theorem 3 provides a bound on the expected number of communication bits to encode the quantized stochastic gradients. As expected, the upper bound in Eq. 12 increases monotonically with $d$ and $s$.

We can combine variance and code-length upper bounds and obtain convergence guarantees for AQSGD when applied to various learning problems where we have convergence guarantees for full-precision SGD under standard assumptions.

On convex problems, convergence guarantees can be established along the lines of [17, Theorems 6.1].

In addition, AQSGD requires at most $\overline{N_{Q}}$ communication bits per iteration in expectation.

In Appendix H and Appendix I, we obtain convergence guarantees on nonconvex problems and for momentum-based variants of AQSGD under standard assumptions, respectively. Theoretical guarantees for levels with symmetry are established in Appendix J.

Experimental evaluation

In this section, we showcase the effectiveness of our adaptive quantization methods in speeding up training deep models. We compare our methods to the following baselines: single-GPU SGD (SGD), full-precision multi-GPU SGD (SuperSGD), uniform levels under $L^{\infty}$ normalization (QSGDinf) , ternary levels under $L^{\infty}$ normalization (TRN) , and exponential levels under $L^{2}$ normalization with exponential factor $p=0.5$ (NUQSGD) . We present results for the following variations of our proposed methods: ALQ and AMQ (with norm adjustments in Section 3.4), and their normalized variations ALQ-N and AMQ-N (Sections 3.1 and 3.3). We present full training results on ImageNet in Appendix K along with additional experimental details.

We compare methods in terms of the number of training iterations that is independent of a particular distributed setup. In Table 1, we present results for training ResNet-32 and ResNet-110 on CIFAR-10 , and ResNet-18 on ImageNet . We simulate training with $4$-GPUs on a single GPU by quantizing and dequantizing the gradient from $4$ mini-batches in each training iteration. These simulations allow us to compare the performance of quantization methods to the hypothetical full-precision SuperSGD.

All quantization methods studied in this section share two hyper-parameters: the number of bits ($\log_{2}$ of number of quantization levels) and a bucket size. A common trick used in normalized quantization is to encode and decode a high-dimensional vector in buckets such that each coordinate is normalized by the norm of its corresponding bucket instead of the norm of the entire vector . The bucket size controls the tradeoff between extra communication cost and loss of precision. With a small bucket size, there are more bucket norms to be communicated, while with a large bucket size, we lose numerical precision as a result of dividing each coordinate by a large number. In Section 5.1, we provide an empirical study of the hyperparameters.

Matching the accuracy of SuperSGD. Using only $3$ bits ($8$ levels), our adaptive methods match the performance of SuperSGD on CIFAR-10 and close the gap on ImageNet (bold in Table 1). Our most flexible method, ALQ, achieves the best overall performance on ImageNet and the gap on CIFAR-10 with ALQ-N is less than $0.3\%$. There is at least $1.4\%$ gap between our best performing method and previous work in training each model. To the best of our knowledge, matching the validation loss of SuperSGD has not been achieved in any previous work using only $3$ bits. Fig. 3 shows the test loss and Fig. 4 shows the average gradient variance where the average is taken over gradient coordinates. Our adaptive methods successfully achieve lower variance during training.

Comparison on the trajectory of SGD. Fig. 5 shows the average variance on the optimization trajectory of single-GPU without quantization. This graph provides a more fair comparison of the quantization error of different methods decoupled from their impact on the optimization trajectory. ALQ effectively finds an improved set of levels that reduce the variance in quantization. ALQ matches the variance of SuperSGD on Resnet-110 (Fig. 5(b)). In Figs. 5(b) and 5(c), the variance of QSGDinf is as high as TRN in the first half of training. This shows that extra levels ($8$ uniform levels) do not perform better unless designed carefully. As expected, the variance of SuperSGD is always smaller than the variance of SGD by a constant factor of the number of GPUs.

Negligible computational overhead. Our adaptive methods have similar per-step computation and communication cost compared to previous methods. On ImageNet, we save at least $60$ hours from $95$ hours of training and add only an additional cost of at most 10 minutes in total to adapt quantization. For bucket sizes $8192$ and $16384$ and $3$–$8$ bits used in our experiments, the per-step cost relative to SuperSGD ($32$-bits) is $21$–$25\%$ for ResNet-18 on ImageNet and $32$–$36\%$ for ResNet-50. That is the same as the cost of NUQSGD and QSGDinf without additional coding or pruning with the same number of bits and bucket sizes. The cost of the additional update specific to ALQ is $0.4$–$0.5\%$ of the total training time. In Section K.3, we provide tables with detailed timing results for varying bucket sizes and bits.

Fig. 7 shows quantization levels for each method at the end of training ResNet-32 on CIFAR-10. The quantization levels for our adaptive methods are more concentrated near zero. In Figs. 6(a) and 6(b), we study the impact of the bucket size and number of bits on the best validation accuracy achieved by quantization methods.

Adaptive levels are the best quantization methods across all values of bucket size and number of bits. ALQ and ALQ-N are the best performing methods across all values of bucket size and number of bits. The good performance of ALQ-N is unexpected as it suggests quantization for vectors with different norms can be shared. In practice, ALQ-N is easier to implement and faster to update compared to ALQ. We observe a similar relation between AMQ and AMQ-N methods. Adaptive multiplier methods show inferior performance to adaptive level methods as the bucket size significantly grows (above $10^{4}$) or shrinks (below $100$) as well as for very few bits ($2$). Note that there exists a known generalization gap between SGD and SuperSGD in ResNet-110 that can be closed by extensive hyperparameter tuning . Our adaptive methods reduce this gap with standard hyperparameters.

Bucket size significantly impacts non-adaptive methods. For bucket size $100$ and $3$ bits, NUQSGD performs nearly as good as adaptive methods but quickly loses accuracy as the bucket size grows or shrinks. QSGDinf stays competitive for a wider range of bucket sizes but still loses accuracy faster than other methods. This shows the impact of bucketing as an understudied trick in evaluating quantization methods.

Adaptive methods successfully scale to large number of GPUs. Table 2 shows the result of training CIFAR-10 on ResNet-32 using 16 and 32 GPUs. Note that with 32 GPUs, TRN is achieving almost the accuracy of SuperSGD with only 3 quantization levels, which is expected because TRN is unbiased and the variance of aggregated gradients decreases linearly with the number of GPUs.

Conclusions

To reduce communication costs of data-parallel SGD, we introduce two adaptively quantized methods, ALQ and AMQ, to learn and adapt gradient quantization method on the fly. In addition to quantization method, in both methods, processors learn and adapt their coding methods in parallel by efficiently computing sufficient statistics of a parametric distribution. We establish tight upper bounds on the excessive variance for any arbitrary sequence of quantization levels under general normalization and on the expected number of communication bits per iteration. Under standard assumptions, we establish a number of convergence guarantees for our adaptive methods. We demonstrate the superiority of ALQ and AMQ over nonadaptive methods empirically on deep models and large datasets.

Broader impact

This work provides additional understanding of statistical behaviour of deep machine learning models. We aim to train deep models using popular SGD algorithm as fast as possible without compromising learning outcome. As the amount of data gathered through web and a plethora of sensors deployed everywhere (e.g., IoT applications) is drastically increasing, the design of efficient machine learning algorithms that are capable of processing large-scale data in a reasonable time can improve everyone’s quality of life. Our compression schemes can be used in Federated Learning settings, where a deep model is trained on data distributed among multiple owners without exposing that data. Developing privacy-preserving learning algorithms is an integral part of responsible and ethical AI. However, the long-term impacts of our schemes may depend on how machine learning is used in society.

Acknowledgement

The authors would like to thank Blair Bilodeau, David Fleet, Mufan Li, and Jeffrey Negrea for helpful discussions. FF was supported by OGS Scholarship. DA and IM were supported the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement No 805223 ScaleML). DMR was supported by an NSERC Discovery Grant. ARK was supported by NSERC Postdoctoral Fellowship. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.www.vectorinstitute.ai/#partners

References

Appendix A CDF and its inverse

The probability density function (PDF) for $X\sim{\cal N}(\mu,\sigma^{2})$ is defined as

and the cumulative distribution function (CDF) defined as

The inverse of CDF for the normal distribution is given by

Various approximations of Eq. 14 and Eq. 15 are available in the literature.

A.2 Truncated normal distribution

The probability density function (PDF) of a truncated normal distribution that lies within the interval $(a,b)$ with $-\infty<a<b<\infty$ is defined as

where $p_{{\cal N}}(\cdot)$ is defined in Eq. 13 and the cumulative distribution function (CDF) is defined as

where $F_{{\cal N}}(\cdot)$ and $\Phi(\cdot)$ are defined in Eq. 14. Note that the mean and variance of a random variable with truncated normal distribution are not $\mu$ and $\sigma^{2}$ based on our notation. The mean and variance depend on the interval $(a,b)$, which is clear in contexts that we use.

The inverse of CDF for truncated normal distribution is given by

where \overline{y}=\big{(}F_{{\cal N}}(b)-F_{{\cal N}}(a)\big{)}y+F_{{\cal N}}(a) and $F_{{\cal N}}^{-1}(\cdot),~{}\Phi^{-1}(\cdot)$ are defined in Eq. 15.

Appendix B Expected normalized variance minimization

We prove Theorem 1 in two steps in Proposition 1 and Proposition 2.

Let $R$ denote a random variable with probability density function (PDF) $p$ and cumulative distribution function (CDF) $F$. To show that problem Eq. 3 is nonconvex, we first focus on the problem of optimizing two levels $\min_{a,b}Q(a,b)$ where

The function $Q(a,b)$ is nonconvex in general. It becomes convex if for all $0\leq a\leq b\leq 1$, we have

We can find the eigenvalues of $\nabla^{2}Q$ by solving $|\nabla^{2}Q-\lambda{\bf I}|=0$, which leads to the following quadratic equation:

We note that Eq. 20 is sufficient to guarantee $\nabla^{2}Q\succeq\mathbf{0}$. ∎

The sufficient condition Eq. 20 is satisfied if $R$ is uniformly distributed in the range $$.

We now solve the problem of optimizing a single level, i.e., $\min_{b}Q(b)$ where

We note that $Q$ is convex so we can find the closed-form optimal solution through satisfying the first order optimality condition. ∎

In the special case with $a=0$ and $c=1$, the optimal solution to minimize $Q(b)$ is given by

For the special case of a truncated normal, the inner integral is evaluated as

where $F_{{\cal T}}$ and $p_{{\cal T}}$, the CDF and PDF of the truncated normal, are defined in Appendix A.

B.2 Projected Gradient Descent

For the special case of a normal or truncated normal distribution, the gradient of the expected normalized variance used in Section 3.2 is:

B.3 Symmetric Levels

Let $r_{i}=|{\theta}_{i}|$ for $i=1,\ldots,d$. We have the following propositions.

The variance of quantization with symmetric levels is given by

Note that for symmetrical levels, we have

If PDF of normalized gradients is an even function,i.e., $p(-{\theta})=p({\theta})$ for ${\theta}\in$, the expected normalized variance in Eq. 3 can be rewritten as

For the case of symmetrical levels, the gradient of the expected normalized variance is given by

B.3.2 CD

where $F$ is the CDF of the normalized coordinate.

B.3.3 Exponentially spaced levels

Using Leibniz’s rule, we can compute the first order derivative:

In particular, in the special case of a normal or truncated normal distribution, we have

We can update $p$ efficiently by a gradient descent algorithm as we have a closed-form expression to find the gradient function.

Appendix C Expected variance minimization in Section 3.4

In the following, we provide the update rules and the analysis of computation complexity of ALQ, GD, and AMQ.

C.2 GD update

In the special case of (truncated) normal distribution, we have

C.3 AMQ (GD update with exponentially spaced levels)

In the special case of (truncated) normal distribution, we have

We can update $p$ efficiently by a gradient descent algorithm as we have a closed-form expression to find the gradient function.

C.4 Computational complexity and scalability

The number of iterations for ALQ method to converge is in the order of $O(s\log(1/\epsilon))$ where $\epsilon$ is the suboptimality gap of bisection search. The number of iterations for AMQ method to achieve a local minimum with gap $\epsilon$ is $O(1/\epsilon)$. The total number of gradient computations for GD method to achieve a local minimum with gap $\epsilon$ is $O(s/\epsilon)$. Note that processors can run our methods in parallel. The time complexity of these methods is independent of the number of samples, the number of processors, and the number of parameters. The extra computational overhead is negligible compared to costs of computation of stochastic gradients and communication. Furthermore, we do not need to optimize levels at each iteration. Our experimental results suggest that it is sufficient to optimize levels at the lr_scheduler iterations.

Appendix D Encoding

The DECODE function (for Algorithm 1) simply reads $b$ bits to reconstruct $\|{\bf v}\|$. Using $\Gamma^{-1}$, it decodes the index of the first coordinate, depending on whether the decoded entry is zero or nonzero, it may read one bit indicating the sign, and then proceeds to decode the next symbol. The process proceeds in rounds, mimicking the encoding process, finishing when all coordinates have been decoded. Note that we can improve coding efficiency by encoding blocks of symbols at the cost of increasing encoding/decoding complexity. In this paper, we focus on a simple lossless prefix coding scheme that encodes symbols separately.

for $j=1,\ldots,s$ where $F$ is the marginal CDF of normalized coordinates. In addition, we have

In the special case of truncated normal distribution, we have the symbol probabilities in closed-form:

Appendix E Variance gap

For any distribution where the gap between the expected variance of a normalized coordinate under an optimal quantization to minimize Eq. 3 and the worst-case one is lower bounded by some constant, the total gap is lower bounded by $\Omega(d)$. We quantify this gap for the special case of one level with truncated normal density.

Note that $\hat{b}=1/2$ minimizes the worst-case variance upper bound in Eq. 1 . In general, $b^{*}\neq\hat{b}$ depending on $\mu$ and $\sigma^{2}$. Without loss of generality, assume $b^{*}>\hat{b}$.

In the interval $[\hat{b},b^{*}]$, we have

In this interval, $Q$ is $\gamma$-strongly convex, i.e.,

Hence, the gap in the expected normalized variance under $b^{*}$ and $\hat{b}$ is lower bounded by:

Appendix F Proof of Theorem 2 (variance bound)

In our proofs, we use the following known lemma.

Note that Lemma 1 holds even when $q<1$ and $\|\cdot\|_{q}$ is merely a seminorm.

In the following, we derive a bound on $\sum_{r_{i}\in{\cal I}_{0}}(2^{-s}-r_{i})r_{i}$, which completes the proof.

We can solve Eq. 42 and obtain the optimal solution \theta^{*}=\big{(}\frac{1/p-1}{2/p-1}\big{)}^{p}. Substituting $\theta^{*}$ into Eq. 42, we obtain Eq. 41. ∎

Let ${{\cal S}_{j}}$ denote the coordinates of vector ${{\bf v}}$ whose elements fall into the ${(j+1)}$-th bin, i.e., ${{\cal S}_{j}\triangleq\{i:r_{i}\in[l_{j},l_{j+1}]\}}$ for ${j=0,\ldots,s}$.

Then, for any $0<p<1$ and $q\geq 2$, we have

where the third inequality holds as $\|{\bf v}\|_{p}\leq\|{\bf v}\|_{2}d^{1/p-1/2}$ using Lemma 1 and the last inequality holds as $\|{\bf v}\|_{q}\leq\|{\bf v}\|_{2}$ for $q\geq 2$.

For $q\geq 1$, we have $\|{\bf v}\|_{q}^{2-p}\leq\|{\bf v}\|_{2}^{2-p}d^{\frac{2-p}{\min\{q,2\}}-\frac{2-p}{2}}$, which gives

Appendix G Proof of Theorem 3 (code-length bound)

We first encode the norm $\|{\bf v}\|_{q}$ using $b$ bits where, in practice, we use standard 32-bit floating point encoding.

We send one bit for each nonzero entry of ${\bf h}$. Let ${{\cal S}_{j}\triangleq\{i:r_{i}\in[l_{j},l_{j+1}]\}}$ and ${d_{j}\triangleq|{\cal S}_{j}|}$ for ${j=0,\ldots,s}$. We have an upper bound on the expected number of nonzero entries as follows:

where the last inequality holds as $\|{\bf v}\|_{1}\leq\|{\bf v}\|_{q}d^{1-1/q}$ using Lemma 1. ∎

For each entry of ${\bf h}$, we send the associated codeword. The optimal expected code-length for transmitting one random symbol is within one bit of the entropy of the source. Hence, we need to transmit upto $d(H(L)+1)$ to transmit entries of ${\bf h}$ . Putting everything together, we have

Finally, note that the entropy of a source with $n$ outcomes is bounded above by $\log_{2}(n)$.

Appendix H AQSGD for smooth nonconvex optimization

On nonconvex problems, we can establish convergence guarantees in terms of convergence to a local minima for a smooth loss function along the lines of, e.g., [33, Theorem 2.1].

In addition, AQSGD requires at most $\overline{N_{Q}}$ communication bits per iteration in expectation.

Appendix I AQSGD with momentum

The update rule for full-precision unified momentum SGD (UMSGD) is given by

where ${\bf w}_{t}$ is the current parameter input and $\mu\in[0,1)$ is the momentum parameter. Note that the heavy-ball method and Nesterov’s accelerated gradient method are the special cases of UMSGD obtained by substituting $l=0$ and $l=1$ into Eq. 45, respectively.

The steps for data-parallel version of UMSGD are those in Algorithm 1 by replacing Algorithm 1 with an UMSGD update. We have convergence guarantees for adaptively quantized SGD with momentum (AQSGDM) along the lines of, e.g., [34, Theorem 1]. We first establish the convergence guarantees for convex optimization in the following theorem.

In addition, AQSGD requires at most $\overline{N_{Q}}$ communication bits per iteration in expectation.

On nonconvex problems, (weaker) convergence guarantees can be established for AQSGDM. In particular, AQSGDM is guaranteed to converge to a local minima for smooth general loss functions.

In addition, AQSGD requires at most $\overline{N_{Q}}$ communication bits per iteration in expectation.

Appendix J Theoretical guarantees for levels with symmetry

Following Proposition 3, the variance is given by

In the special case of truncated normal distribution, we have the symbol probabilities in closed-form:

Finally, we have the following bound on the expected number of communication bits per iteration for quantizing with symmetrical levels.

where $b$ is a constant and $L$ is a random variable with the probability mass function given by Proposition 8.

Appendix K Experimental details and additional experiments

In this section, we provide additional experiments for the methods evaluated in Section 5. In addition to baselines discussed in Section 5, we present results for ALQ-N, which minimizes the expected normalized variance using coordinate descent in Eq. 3, ALQ with norm adjustments in Section 3.4, ALQ adapted using gradient descent in Section 3.2 (ALQG, ALQG-N), AMQ-N, and AMQ. This section includes full ImageNet runs. Figs. 5(b), 3(b), 3(a) and 5(a) have also been extended to include all variations of the proposed algorithms and baselines.

An implementation challenge is that the value of the statistics, especially the variance, can become very small. This makes PDF and CDF calculations challenging. The challenge is that the value of PDF is very close to zero when it is far from the mean but not exactly zero. In order to overcome this challenge, we use histograms to model the distribution of gradients as a weighted sum of truncated normals. Another problem is the large number of statistics that are calculated. As presented in Section 3, we sample a number of gradients and then normalize the gradients. Then we split the gradients into buckets and calculate average, variance, and norm of each of the buckets. The number of means, variances, and norms can become very large with large networks and small bucket sizes. To reduce computational complexity of the algorithm, we sample uniformly from these values. This number of samples is equal to 20 for small networks such as ResNet-8 and networks trained on CIFAR-10; however, in experiments on ImageNet, we used 350 samples to achieve the desired accuracy.

One other understudied detail in quantizing is how bucketing is performed. In , gradient coordinates in each bucket do not exceed the layer size. It means that the gradient coordinates in a bucket do not contain gradient coordinates from the next layer even if the bucket size is not fully utilized. This leads to creation of under-sized buckets that can be problematic for quantization performance. Different tricks are employed to fix this problem. These tricks include transmitting biases or under-sized buckets in full-precision (not that typically biases are main sources of under-sized buckets). In our implementation, we normalize the buckets network-wise and do not consider the layer size as the bucket size boundary. We only transmit the last bucket in full precision if it is smaller than the specified bucket size.

Update Schedule. In the ImageNet and CIFAR-10 runs, adaptive level updates are scheduled at 100 and 2000 iterations only once and every 10K iterations. The reason for this schedule is changes in the gradient statistics over the course of training. As shown in Fig. 1, the average variance changes rapidly during the first iterations and then only changes at every learning rate schedule. In practice, we noticed accuracy degradation especially when the levels are not updated during the initial iterations where the average variance is rapidly changing.

Convergence of level updating. Fig. 8 shows the expected normalized variance (the objective in Eq. 3) and expected variance (the objective in Eq. 9) during one step of adapting levels. This figure shows that the objective function in Eq. 3 is nonconvex and different initializations lead to sub-optimal solutions. ALQG and ALQG-N refer to variations of ALQ using gradient descent instead of coordinate descent.

Hyperparameters used for training. Table 3 shows the hyperparameters used for training CIFAR-10 and ImageNet. These are conventional hyperparameters for training ResNet models. SuperSGD is able to replicate the accuracy reported in showing the correct setting for training.

Validation accuracy on full ImageNet run Table 4 shows the validation accuracy of full ImageNet runs on ResNet-18. Total number of iterations required for a full ImageNet run is 600K. This table shows that ALQ and ALQ-N are able to outperform QSGDinf by 1% on ImageNet.

Similar performance of normalized and unnormalized variations of AMQ and ALQ methods suggests that for given datasets and deep models, the distribution of normalized gradient coordinates can be represented by either of the forumulations in Section 3. In AMQ-N and ALQ-N, $\mu$ and $\sigma$ values for the truncated normal distribution is equal to the average of $\mu$ and $\sigma$ for individual buckets.

Figs. 10(a) and 10(b) show an interesting observation for NUQSGD. Although NUQSGD has worse performance in terms of the training loss and average variance compared to all other approaches, it is able to achieve better validation accuracy. This suggests that NUQSGD is able to generalize better in this specific setting. However, this pattern does not repeat when it comes to the ImageNet dataset.

Figs. 9, 10, 11 and 12 are extended versions of Figs. 3, 11 and 4 figures in the main body. The difference is that they contain more baselines compared to the figures in the main body. Fig. 9 contains the training loss for the experiments on CIFAR-10. It was not possible to include the same figure for ImageNet, because calculating the full training loss on ImageNet takes a very long time.

The expected variance, training loss, and the validation loss for the results presented in Table 2 are shown in Fig. 13. Although ALQ performs better in expected variance and training loss, it seems to have trouble when it comes to the validation loss for 32-GPUs. We suspect that this is due to the large total batch size used for these experiment that results in overfitting. The batch size for each GPU is 128.

K.2 Effect of Using Gradient Clipping

TRN introduced the idea of gradient clipping before quantization. Gradient clipping replaces the gradient coordinates that are far from the mean to reduce the gradient variance. The gradient coordinates that are very far from the mean can affect the normalization. In order to tackle this problem, they clip the gradients before quantization. The clipping process can be described using Eq. 49:

The constant $c$ used in TRN equals to 2.5. In order to investigate the effect of gradient clipping in ALQ and AMQ, we train a ResNet-8 on CIFAR-10 dataset for various bucket sizes. Fig. 14 shows the validation accuracy of the baselines and the algorithms we proposed. ALQ and ALQ-N always maintain better or equal accuracy compared to the other quantization schemes. It is also worth noting that the quantization is performed by each layer instead of a performing the quantization across the network without considering the layers.

K.3 Timing Overhead

In this section, we provide the timing results per step for training ResNet-18 (Table 6) and ResNet-50 (Table 7) on ImageNet with mini-batch size $512$. The training setup consists of $4$ AWS nodes with one V100 GPU on each. Network bandwidth programmatically constrained to 1GBit/s.