Scalable Methods for 8-bit Training of Neural Networks

Introduction

Deep Neural Networks (DNNs) achieved remarkable results in many fields making them the most common off-the-shelf approach for a wide variety of machine learning applications. However, as networks get deeper, using neural network (NN) algorithms and training them on conventional general-purpose digital hardware is highly inefficient. The main computational effort is due to massive amounts of multiply-accumulate operations (MACs) required to compute the weighted sums of the neurons’ inputs and the parameters’ gradients.

Much work has been done to reduce the size of networks. The conventional approach is to compress a trained (full precision) network using weights sharing, low rank approximation, quantization, pruning or some combination thereof. For example, Han et al., 2015 successfully pruned several state-of-the-art large-scale networks and showed that the number of parameters can be reduced by an order of magnitude.

Since training neural networks requires approximately three times more computation power than just evaluating them, quantizing the gradients is a critical step towards faster training machines. Previous work demonstrated that by quantizing network parameters and intermediate activations during the training phase more computationally efficient DNNs could be constructed.

Researchers have shown that 16-bit is sufficient precision for most network training but further quantization (i.e., 8-bit) results with severe degradation. Our work is the first to almost exclusively train at 8-bit without harming classification accuracy. This is addressed by overcoming two main obstacles known to hamper numerical stability: batch normalization and gradient computations.

The traditional batch normalization implementation requires the computation of the sum of squares, square-root and reciprocal operations; these require high precision (to avoid zero variance) and a large dynamic range. It should come as no surprise that previous attempts to use low precision networks did not use batch normalization layers or kept them in full precision . This work replaces the batch norm operation with range batch-norm (range BN) that normalizes inputs by the range of the input distribution (i.e., $\max(x)-\min(x)$). This measure is more suitable for low-precision implementations. Range BN is shown analytically to approximate the original batch normalization by multiplying this range with a scale adjustment that depends on the size of the batch and equals to $(2\cdot\ln(n))^{-0.5}$. Experiments on ImageNet with Res18 and Res50 showed no distinguishable difference between accuracy of Range BN and traditional BN.

The second obstacle is related to the gradients quantization. Given an upstream gradient $g_{l}$ from layer $l$, layer $l-1$ needs to apply two different matrix multiplications: one for the layer gradient $g_{l-1}$ and the other for the weight gradient $g_{W}$ which are needed for the update rule. Our analysis indicates that the statistics of the gradient $g_{l}$ violates the assumptions at the crux of common quantization schemes. As such, quantizing these gradients constitutes the main cause of degradation in performance through training. Accordingly, we suggest to use two versions of layer gradients $g_{l}$, one with low-precision (8-bit) and another with higher-precision (16-bit). The idea is to keep all calculations with $g_{l}$ that does not involve a performance bottleneck at 16 bits, while the rest at 8 bits. As the gradients $g_{W}$ are required only for the weight update, they are computed using the 16 bits copy of $g_{l}$. On the other hand, the gradient $g_{l-1}$ is required for the entire backwards stream and as such it is computed using the corresponding 8-bit version of $g_{l}$. In most layers of the DNN these computations can be performed in parallel. Hence $g_{W}$ can be computed at high precision in parallel with $g_{l-1}$, without interrupting the propagation of $g_{l}$ to lower layers. We denote the use of two different arithmetic precision operations in the differentiation process as "Gradients Bifurcation".

Previous Work

While several works have shown that training at 16-bit is sufficient for most networks, more aggressive quantization schemes were also suggested . In the extreme case, the quantization process used only one bit which resulted in binarized neural networks (BNNs) where both weights and activations were constrained to -1 and 1. However, for more complex models and challenging datasets, the extreme compression rate resulted in a loss of accuracy. Recently, Mishra et al. showed that this accuracy loss can be prevented by merely increasing the number of filter maps in each layer, thus suggesting that quantized neural networks (QNNs) do not possess an inherent convergence problem. Nevertheless, increasing the number of filter maps enlarge quadratically the number of parameters, which raises questions about the efficiency of this approach.

In addition to the quantization of the forward pass, a growing interest is directed towards the quantization of the gradient propagation in neural networks. A fully quantized method, allowing both forward and backward low-precision operations will enable the use of dedicated hardware, with considerable computational, memory, and power benefits. Previous attempts to discretize the gradients managed to either reduce them to 16-bit without loss of accuracy or apply a more aggressive approach and reduce the precision to 6-8 bit with a noticeable degradation. Batch normalization is mentioned by as a bottleneck for network quantization and is either replaced by a constant scaling layer kept in full precision, or avoided altogether; this clearly has some impact on performance (e.g., AlexNet trained over ImageNet resulted with top-1 error of $51.6\%$, where the state of the art is near $42\%$) and better ways to quantize normalization are explicitly called for. Recently L1 batch norm with only linear operations in both forward and backward propagation was suggested by with improved numerical stability. Yet, our experiments show that with 8-bit training even L1 batch norm is prone to overflows when summing over many large positive values. Finally, Wen et al. focused on quantizing the gradient updates to ternary values to reduce the communication bandwidth in distributed systems.

We claim that although more aggressive quantization methods exist, 8-bit precision may prove to have a "sweet-spot" quality to it, by enabling training with no loss of accuracy and without modifying the original architecture. Moreover, we note that 8-bit quantization is better suited for future and even current hardware, many of which can already benefit from 8-bit operations . So far, to the best of our knowledge, no work has succeeded to quantize the activations, weights, and gradient of all layers (including batch normalization) to 8-bit without any degradation.

Range Batch-Normalization

For a layer with $n\times d-$dimensional input $x=(x^{(1)},x^{(2)},...,x^{(d)})$, traditional batch norm normalizes each dimension

where $\mu^{d}$ is the expectation over $x^{(d)}$, $n$ is the batch size, $C(n)=\frac{1}{\sqrt{2\cdot\ln(n)}}$ is a scale adjustment term, and $\text{range}(x)=\max(x)-\min(x)$.

The main idea behind Range BN is to use the scale adjustment $C(n)$ to approximate the standard deviation $\sigma$ (traditionally being used in vanilla batch norm) by multiplying it with the range of the input values. Assuming the input follows a Gaussian distribution, the range (spread) of the input is highly correlated with the standard deviation magnitude. Therefore by normalizing the range by $C(n)$ we can estimate $\sigma$. Note that the Gaussian assumption is a common approximation (e.g., Soudry et al. ), based on the fact that the neural input $x^{(d)}$ is a sum of many inputs, so we expect it to be approximately Gaussian from the central limit theorem.

We now turn to derive the normalization term $C(n)$. The expectation of maximum of Gaussian random variables are bounded as follows :

Since ${x^{(d)}-\mu^{d}}$ is symmetrical with respect to zero (centred at zero and assumed gaussian), it holds that $E[\max(\cdot)]=-E[\min(\cdot)]$; hence,

Therefore, by summing Equations 3 and 4 and multiplying the three parts of the inequality by the normalization term $C(n)$, Range BN in Eq. 2 approximates the original standard deviation measure $\sigma$ as follows:

Importantly, the scale adjustment term $C(n)$ plays a major role in RangeBN success. The performance was degraded in simulations when $C(n)$ was not used or modified to nearby values.

Quantized Back-Propagation

Following we used the GEMMLOWP quantization scheme as decribed in Google’s open source library . A detailed explanation of this approach is given in Appendix B. While GEMMLOWP is widely used for deployment, to the best of the authors knowledge this is the first time GEMMLOWP quantization is applied for training. Note that the activations maximum and minimum values were computed by the range BN operator, thus finding the normalization scale as defined in Appendix B does not require additional $O(n)$ operations.

Finally we note that a good convergence was achieved only by using stochastic rounding for the gradient quantization. This behaviour is not surprising as the gradients will serve eventually for the weight update thus unbiased quantization scheme is required to avoid noise accumulation.

starting from the last layer. Each layer needs to derive two sets of gradients to perform the recursive update. The layer activation gradients:

served for the Back-Propagation (BP) phase thus passed to the next layer,and the weights gradients

Similar to previous work , we used the straight-through estimator (STE) approach to approximate differentiation through discrete variables. This is the most simple and hardware friendly approach to deal with the fact that the exact derivative of discrete variables is zero almost everywhere.

When is quantization of neural networks possible?

This section provides some of the foundations needed for understanding the internal representation of quantized neural networks. It is well known that when batch norm is applied after a convolution layer, the output is invariant to the norm of the weight on the proceeding layer i.e., $BN(C\cdot W\cdot x)=BN(W\cdot x)$ for any given constant $C$. This quantity is often described geometrically as the norm of the weight tensor, and in the presence of this invariance, the only measure that needs to be preserved upon quantization is the directionality of the weight tensor. In the following we show that quantization preserves the direction (angle) of high-dimensional vectors when $W$ follows a Gaussian distribution.

More specifically, for networks with $M$-bit fixed point representation, the angle is preserved when the number of quantization levels $2^{M}$ is much larger than $\sqrt{2\ln(N)}$, where $N$ is the size of quantized vector. This shows that significant quantization is possible on practical settings. Taking for example the dimensionality of the joint product in a batch with 1024 examples corresponding to the last layer of ResNet-50, we need no more than 8-bit of precision to preserve the angle well (i.e., $\sqrt{2\ln(3\cdot 3\cdot 2048\cdot 1024)}=5.7<<2^{8}$). We stress that this result heavily relays on values being distributed according to a Gaussian distribution, and suggests why some vectors are robust to quantization (e.g., weights and activations) while others are more fragile (e.g., gradients).

Given a vector of weights $W=(w_{0},w_{1},...,w_{N-1})$, where the weights follow a Gaussian distribution $W\sim N(0,\sigma)$, we would like to measure the cosine similarity (i.e., cosine of the angle) between $W$ and $Q(W)$, where $Q(\cdot)$ is a quantization function. More formally, we are interested in estimating the following geometric measure:

We next define the quantization function $Q(\cdot)$ using a fixed quantization step between adjacent quantified levels as follows:

We consider the case where quantization step $\Delta$ is much smaller than $\text{mean}(|W|)$. Under this assumption correlation between $W$ and quantization noise $W-Q(W)=(\epsilon_{0},\epsilon_{1},...,\epsilon_{N-1})$ is negligible, and can be approximated as an additive noise. Our model assumes an additive quantization noise $\bar{\epsilon}$ with a uniform distribution i.e., $\epsilon_{i}\sim\mathcal{U}[-\Delta/2,\Delta/2]$ for each index $i$. Our goal is to estimate the angle between $W$ and $W+\bar{\epsilon}$ for high dimensions (i.e., $N\to\infty$).

2 Angle preservation during quantization

In order to estimate the angle between $W$ and $W+\epsilon$, we first estimate the angle between $W$ and $\epsilon$. It is well known that if $\epsilon$ and $W$ are independent, then at high dimension the angle between $W$ and $\epsilon$ tends to $\dfrac{\pi}{2}$ i.e., we get a right angle triangle with $W$ and $\epsilon$ as the legs, while $W+\epsilon$ is the hypotenuse as illustrated in Figure 1-right. The cosine of the angle $\theta$ in that triangle can be approximated as follows:

Since $W$ is Gaussian, we have that $E(||W||)\cong\sqrt{N}\sigma$ in high dimensions . Additionally, in Appendix A we show that $E(||\bar{\epsilon}||)\leq\sqrt{N/12}\cdot\Delta$. Moreover, at high dimensions, the relative error made as considering $E||X||$ instead of the random variable $||X||$ becomes asymptotically negligible . Therefore, the following holds in high dimensions:

Finally, $E(\max(W))\leq\sqrt{2}\sigma\sqrt{\ln(N})$ when $W$ follows a Gaussian distribution , establishing the following:

Eq. 12 establishes that when $2^{M}>>\sqrt{\ln(N)}$ the angle is preserved during quantization. It is easy to see that in most practical settings this condition holds even for challenging quantizations. Moreover, this results highly depends on the assumption made about the Gaussian distribution of $W$ (transition from equation 11 to equation 12).

Experiments

We evaluated the ideas of Range Batch-Norm and Quantized Back-Propagation on multiple different models and datasets. The code to replicate all of our experiments is available on-line https://github.com/eladhoffer/quantized.pytorch.

To validate our assumption that the cosine similarity is a good measure for the quality of the quantization, we ran a set of experiments on Cifar-10 dataset, each with a different number of bits, and then plotted the average angle and the final accuracy. As can be seen in Figure 2 there is a high correlation between the two. Taking a closer look the following additional observations can be made: (1) During quantization the direction of vectors is better preserved with the forward pass compared to the backward pass; (2) validation accuracy follows tightly the cosine of the angle in the backward pass, indicating gradient quantization as the primary bottleneck; (3) as expected, the bound on $E(\cos(\theta))$ in Eq. 12 holds in the forward pass, but less so in the backward pass, where the Gaussian assumption tends to break. The histograms in Figure 2 further confirms that the layer gradients $g_{l}$ do not follow Gaussian distribution. These are the values that are bifurcated into low and high precision copies to reduce noise accumulation.

2 Experiment results on ImageNet dataset: Range Batch-Normalization

We ran experiments with Res50 on ImageNet dataset showing the equivalence between the standard batch-norm and Range BN in terms of accuracy. The only difference between the experiments was the use of Range BN instead of the traditional batch-norm. Figure 3 compares between the two. It shows equivalence when models are trained at high precision. We also ran simulations on other datasets and models. When examining the final results, both were equivalent i.e., 32.5% vs 32.4% for ResNet-18 on ImageNet and 10.5% vs 10.7% for ResNet-56 on Cifar10. To conclude, these simulations prove that we can replace standard batch-norm with Range BN while keeping accuracy unchanged. Replacing the sum of squares and square root operations in standard batch-norm by a few maximum and minimum operations has a major benefit in low-precision implementations.

3 Experiment results on ImageNet dataset: Putting it all together

We conducted experiments using RangeBN together with Quantized Back-Propagation. To validate this low precision scheme, we were quantizing the vast majority of operations to 8-bit. The only operations left at higher precising were the updates (float32) needed to accumulate small changes from stochastic gradient descent, and a copy of the layer gradients at 16 bits needed to compute $g_{W}$. Note that the float32 updates are done once per minibatch while the propagations are done for each example (e.g., for a minibatch of 256 examples the updates constitute less than 0.4% of the training effort). Figure 4 presents the result of this experiment on ImageNet dataset using ResNet18 and ResNet50. We provide additional results using more aggressive quantizations in Appedix F.

Discussion

In this study, we investigate the internal representation of low precision neural networks and present guidelines for their quantization. Considering the preservation of direction during quantization, we analytically show that significant quantization is possible for vectors with a Gaussian distribution. On the forward pass the inputs to each layer are known to be distributed according to a Gaussian distribution, but on the backward pass we observe that the layer gradients $g_{l}$ do not follow this distribution. Our experiments further assess that angle is not well preserved on the backward pass, and moreover final validation accuracy tightly follows that angle. Accordingly, we bifurcate the layer gradients $g_{l}$ and use it at 16-bits for the computation of the weight gradient $g_{W}$ while keeping the computation of next layer gradient $g_{l-1}$ at 8-bit. This enables the (slower) 16-bits computation of $g_{W}$ to be be done in parallel with $g_{l-1}$, without interrupting the propagation the layer gradients.

We further show that Range BN is comparable to the traditional batch norm in terms of accuracy and convergence rate. This makes it a viable alternative for low precision training. During the forward-propagation phase computation of the square and square root operations are avoided and replaced by $\max(\cdot)$ and $\min(\cdot)$ operations. During the back-propagation phase, the derivative of $\max(\cdot)$ or $\min(\cdot)$ is set to one where the coordinates for which the maximal or minimal values are attained, and is set to zero otherwise.

Finally, we combine the two novelties into a single training scheme and demonstrate, for the first time, that 8-bit training on a large scale dataset does not harm accuracy. Our quantization approach has major performance benefits in terms of speed, memory, and energy. By replacing float32 with int8, multiplications become 16 times faster and at least 15 times more energy efficient . This impact is attained for 2/3 of all the multiplications, namely the forward pass and the calculations of the layer gradients $g_{l}$. The weight gradients $g_{W}$ are computed as a product of 8-bit precision (layer input) with a 16-bit precision (unquantized version of $g_{l}$), resulting with a speedup of x8 for the rest of multiplications and at least x2 power savings. Although previous works considered an even lower precision quantization (up-to 1-bit), we claim that 8-bit quantization may prove to be more of an interest. Furthermore, 8-bit matrix multiplication is available as an off-the-shelf operation in existing hardware and can be easily adopted and used with our methods.

References

Appendix

By Jensen inequality the following holds true:

Given a uniform random variable in $\epsilon_{i}\sim U[-\Delta/2,\Delta/2]$, we next derive the expected value of its square as follows:

Note that $\Delta$ is a random variable. Hence, we substitute Eq 14 into Eq 13 using the following conditional expectation:

We can now establish $E(||\epsilon||)$ as follows:

Appendix B Quantization methods

Following we choose to use the GMMLOWP quantization scheme as decribed in Google’s open source library . Given an input tensor $x$, clamping values $[v_{\min},v_{\max}]$,and number of bits $M$ we set the output to be:

The clamping values for the weights and activations were defined as the input’s absolute maximum and minimum. Since the activations can have a high dynamic range which can be aggressively clamped as shown by we defined its clamping values to be the average of absolute maximum and minimum values of K chunks. This reduces the dynamic range Variance and allows smaller quantization steps.

it is important to note that a good convergence was achieved only by using stochastic rounding . This behaviour is not surprisings as the gradients serves eventually for the weight update thus unbias quantization scheme is required to avoid quantization noise accumulation.

Appendix C Additional Experiments

In this section we present our more aggressive quantization experiments of the Quantized Back-Propagation scheme. In the extreme case, QBP ternarizes the gradients and uses only 1-bit for the weights and activations. In this case, we refer to QBP networks as Ternarized Back-Propagation (TBP), in which all forward MACs operations can be replaced with XNOR and population count (i.e., counting the number of ones in the binary number) operations. To avoid significant degradation in test accuracy, we apply stochastic ternarization and increase the number of filter maps in a each convolution layer.

A well studied dataset is the CIFAR10 image classification benchmark first introduced by Krizhevsky . CIFAR10 consists of a training set of size 50K, and a test set of size 10K color images. Here, each images represents one of the following categories: airplanes, automobiles, birds, cats, deer, dogs, frogs, horses, ships and trucks. We trained a VGG-like network similar to the one suggested by Hubara et al. on the CIFAR10 dataset, with the same hyper-parameters used in the original work. We compared two variants: the original model BNN model and the BNN model trained with TBP (this paper), where ternarized gradients are used. The results shown in Table (2) demonstrate that if we inflate the convolutional filters by 3 we can achieve similar results as the BNN and full precision models achieved. This is in accordance with previous finding , that found that widening the network can mitigate accuracy drop inflicted by low precision training. To make sure this is not a unique case for BNN we also applied TBP on ResNet with depth of 18. As can be seen from Table (2), as before, inflating the network improves performance, until it is only 1% from the original performance, after inflating it by 5.

C.2 ImageNet

Next, we applied TBP to the more challenging ImageNet classification task introduced by Deng et al. . It consists of a training set of size 1.2M samples and a test set of size 50K. Each instance is labeled with one of 1000 categories including objects, animals, scenes, and even some abstract shapes. We report two error rates for this dataset: top-1 and top-5, as is typical done. Top-$k$ error rate represent the fraction of test images for which the correct label is not among the $k$ most probable labels predicted by the model.

We run several experiments on AlexNet inflated by 3. Similarly to previous work and to ease the comparison we kept first and last layer in full precision. With binarized weights, 4-bit activations and gradients, TBP converged to 53.3% top-1 accuracy and 75.84% top-5 accuracy. By using only 2bit activation TBP reached 49.6% top-1 accuracy and 73.1% top-5 accuracy. We are currently working on more advanced typologies such as ResNet-50 model . Results are summarized in Table (1).

C.3 Additional experiments

Ternarizing both stages results with completely MAC free training. However, our results show that without enabling at least 3bit precision for the update stage the model reaches only approximately 80% accuracy. This indicates that the ternarization noise is too high, and thus distorts the update gradients direction. If we stop the gradients ternarization once the accuracy ceases to increase, the convergence continues and the accuracy increases to the same accuracy as TBP. Thus, ternarizing the update stage can be used to accelerate TBP training of BNN networks by first training it with ternarized weights gradients and then, for the last couple of epochs, continue training with full precision weights gradients.

To alleviate the need for float MAC operation in the update phase we suggest to use gated XNOR operation multiple times, each time with different stochastic sample of the ternarized tensor and average the results. The algorithm is detailed in Algorithm (1) and results are given in Table (3). As expected the accuracy improves with the amount of sampling. To find the number of samples needed for each layer we adopted a similar geometrical approach as suggested by Anderson & Berg and measured the correlation coefficient ($R$) between the update gradients received with and without ternarization. Our experiments indicate that more samples are required for the first two convolution layer (12 samples) while the rest of the layers need approximately 6 samples. Using this configuration keeps the correlation coefficient above 0.7 and results with 87.5% accuracy.