Quantization Networks
Jiwei Yang, Xu Shen, Jun Xing, Xinmei Tian, Houqiang Li, Bing Deng, Jianqiang Huang, Xiansheng Hua
Introduction
Although deep neural networks (DNNs) have achieved huge success in various domains, their high computational and memory costs prohibit their deployment in scenarios where both computation and storage resources are limited. Thus, the democratization of deep learning hinges on the advancement of efficient DNNs. Various techniques have been proposed to lighten DNNs by either reducing the number of weights and connections or by quantizing the weights and activations to lower bits. As exemplified by ResNet , SqueezeNet and MobileNet , numerous efforts have been devoted to designing networks with compact layers and architectures. Once trained, these networks can be further compressed with techniques such as network pruning , weight sharing or matrix factorization .
Approaches for quantizing full-precision networks into low-bit networks can be roughly divided into two categories: approximation-based and optimization-based. Methods in the first category approximate the full-precision (32-bit) values with discrete low-bit (e.g. binary) values via step functions in the forward pass . Because the gradients of such approximations are saturated, additional approximations in the backward process are needed. As a consequence, the use of different forward and backward approximations causes a gradient mismatch problem, which makes the optimization unstable. To avoid the approximation of gradients, some methods formulate the quantization of neural networks as a discretely constrained optimization problem, where losses of the networks are incorporated . Unfortunately, optimization-based methods are only suitable for the quantization of weights. Moreover, the iterative solution of the optimization problem suffers from a high computational complexity during training.
Intuitively, if we can formulate the quantization operation as a simple non-linear function similar to the common activation functions (e.g., Sigmoid , ReLU or Maxout ), no approximation of gradients would be needed, and the quantization of any learnable parameters in DNNs, including activations and weights, can be learned straightforwardly and efficiently. Inspired by that, we present a novel perspective for interpreting and implementing quantization in neural networks. Specifically, we formulate quantization as a differentiable non-linear mapping function, termed quantization function. As shown in Fig. 1, the quantization function is formed as a linear combination of several Sigmoid functions with learnable biases and scales. In this way, the proposed quantization function can be learned in a lossless and end-to-end manner and works for any weights and activations in neural networks, avoiding the gradient mismatch problem. As illustrated in Fig. 2, the quantization is achieved via the continuous relaxation of the steepness of the Sigmoid functions during the training stage.
Our main contributions are summarized as follows:
In contrast to existing low-bit quantization methods, we are the first to formulate quantization as a differentiable non-linear mapping function, which provides a simple/straightforward and general/uniform solution for any-bit weight and activation quantization, without suffering the severe gradient mismatch problem.
We implement a simple and effective form of quantization networks, which could be learned in a lossless and end-to-end manner and outperform state-of-the-art quantization methods on both image classification and object detection tasks.
Related Work
In this paper, we propose formulating the quantization operation as a differentiable non-linear function. In this section, we give a brief review of both low-bit quantization methods and non-linear functions used in neural networks.
Approaches for quantizing full-precision networks into low-bit networks can be roughly divided into two categories: approximation-based and optimization-based. The first approach is to approximate the 32-bit full-precision values with discrete low-bit values in the forward pass of networks. BinaryConnect directly optimizes the loss of the network with weights replaced by sign(), and approximates the sign function with the “hard tanh” function in the backward process, to avoid the zero-gradient problem. Binary weight network (BWN) adds scale factors for the weights during binarization. Ternary weight network (TWN) introduces ternary weights and achieves better performance. Trained ternary quantization (TTQ) proposes learning both ternary values and scaled gradients for -bit weights. DoReFa-Net proposes quantizing -bit weights, activations and gradients using different widths of bits. Gradients are approximated by a custom-defined form based on the mean of the absolute values of full-precision weights. In , weights, activations, gradients and errors are all approximated by low-bitwidth integers based on rounding and shifting operations. Jacob et al. propose an affine mapping of integers to real numbers that allows inference to be performed using integer-only arithmetic. As discussed before, the approximation-based methods use different forward and backward approximations, which causes a gradient mismatch problem. Friesen and Domingos observe that setting targets for hard-threshold hidden units to minimize loss is a discrete optimization problem. Zhuang et al. propose a two-stage approach to quantize the weights and activations in a two-step manner. Lin et al. approximate full-precision weights with the linear combination of multiple binary weight bases. Zhang et al. propose an flexible un-uniform quantization method to quantize both network weights and activations. Cai et al. used several piece-wise backward approximators to overcome the problem of gradient mismatch. Zhou et al. proposed a decoupling step-by-step operation to efficiently convert a pre-trained full-precision convolutional neural network (CNN) model into a low-precision version. As a specific quantization, HashNet adopts a similar continuous relaxation to train the hash function, where a single tanh function is used for binarization. However, our training case (multi-bits quantization of both activations and weights in multi-layers) is much more complicated and challenging.
To avoid the gradient approximation problem, optimization-based quantization methods are recently proposed. They directly formulate the quantization of neural networks as a discretely constrained optimization problem . Leng et al. introduce convex linear constraints for the weights and solve the problem by the alternating direction method of multipliers (ADMM). Hou and Kwok directly optimize the loss function w.r.t. the ternarized weights using proximal Newton algorithm. However, these methods are only suitable for quantization of weights and such iterative solution suffers from high computational costs in training.
2 Non-Linear Functions in Deep Neural Networks
In neural networks, the design of hidden units is distinguished by the choice of the non-linear activation function for hidden units . The simplest form of a neural network is perceptron , where a unit step function is introduced to produce a binary output:
This form is similar to the binary quantization operation, i.e., discretize the continuous inputs into binary values. However, the problem is that it is not immediately obvious how to learn the perceptron networks .
To solve this problem, the sigmoid activation function is adopted in the early form of feedforward neural networks:
which has smooth and non-zero gradient everywhere so that the sigmoid neurons can be learned via back-propagation. When the absolute value of is very large, the outputs of a sigmoid function is close to a unit step function.
Currently, rectified linear units (ReLU) are more frequently used as the activation functions in deep neural networks:
The ReLU function outputs zero across half of its domain and is linear in the other half, which makes the DNNs easy to optimize.
A generalization of the rectified linear units is Maxout. Its activation function is defined as:
where {} and {} are learned parameters. The form of Maxout indicates that a complex convex function can be approximated by a combination of simple linear functions.
Quantization Networks
The main idea of this work is to formulate the quantization operation as a differentiable non-linear function, which can be applied to any weights and activations in deep neural networks. We first present our novel interpretation of quantization from the perspective of non-linear functions. Then, our simple and effective quantization function is introduced and the learning of quantization networks are given.
The quantization operation is mapping continuous inputs into discrete integer numbers, which is similar to the perceptron. Thus, from the perspective of non-linear mapping functions, a binary quantization operation can be formed of a unit step function. Inspired by the design of Maxout units, quantizing continuous values into a set of integer numbers can be formulated as a combination of several binary quantizations. In other words, the ideal low-bit quantization function is a combination of several unit step functions with specified biases and scales, as shown in Fig. 2(e):
where is the full-precision weight/activation to be quantized, is the quantized integer constrained to a predefined set , and is the number of quantization intervals. is the scale factor of inputs. is the standard unit step function. and are the scales and biases for the unit step functions, . The global offset keeps the quantized output zero-centered. Once the expected quantized integer set is given, , and offset can be directly obtained.
For example, for a 3-bit quantization, the output is restricted to , , , and . and are parameters to be learned. Because the step function is not smooth, it is not immediately obvious how we can learn a feedforward networks with Eq. (5) applied to activations or weights .
2 Training and Inference with Quantization Networks
Inspired by the advantage of sigmoid units against the perceptron in feedforward networks, we propose replacing the unit step functions in the ideal quantization function Eq. (5) with sigmoid functions. With this replacement, we can have a differentiable quantization function, termed soft quantization function, as shown in Fig. 2(c). Thus, we can learn any low-bit quantized neural networks in an end-to-end manner based on back propagation.
However, the ideal quantization function Eq. (5) is applied in the inference stage. The use of different quantization functions in training and inference stages may decrease the performance of DNNs. To narrow the gap between the ideal quantization function used in inference stage and the soft quantization function used in training stage, we introduce a temperature to the sigmoid function, motivated by the temperature introduced in distilling ,
With a larger , the gap between two quantization functions is smaller, but the learning capacity of the quantization networks is lower since the gradients of the soft quantization function will be zero in more cases. To solve this problem, in the training stage we start with a small to ensure the quantized networks can be well learned, and then gradually increase w.r.t. the training epochs. In this way, the quantized networks can be well learned and the gap between two quantization functions will be very small at the end of the training.
Forward Propagation. In detail, for a set of full-precision weights or activations to be quantized , the quantization function is applied to each independently:
where and are the scale factors of the input and output respectively. , where indicates the beginning of the input for the -th quantization interval except the first quantization interval, and the beginning of the first quantization interval is . The temperature controls the gap between the ideal quantization function and the soft quantization function. The gradual change from no quantization to complete quantization along with the adjustment of is depicted in Fig. 2.
The quantization function Eq. (7) is applied to every full-precision value that need to be quantized, just as applying ReLU in traditional DNNs. can be either a weight or an activation in DNNs. The output replaces for further computing.
where . we do not need to compute the gradients of , and offset , because their are directly obtained by . Our soft quantization function is a differentiable transformation that introduces quantized weights and activations into the network.
Training and Inference. To quantize a network, we specify a set of weights or activations and insert the quantization function for each of them, according to Eq. (7). Any layer that previously received as an input, now receives . Any module that previously used as parameters, now uses . The smooth quantization function allows efficient training for networks, but it is neither necessary nor desirable during inference; we want the specified weights or activations to be discrete numbers. For this, once the network has been trained, we replace the sigmoid function in Eq. (7) by the unit step function for quantization:
Algorithm 1 summarizes the procedure for training quantization networks. For a full-precision network with modules, where a module can be either a convolutional layer or a fully connected layer, we denote all the activations to be quantized in the -th module as , and denote all the weights to be quantized in the -th module as . All elements in share the same quantization function parameters . All elements in share the same quantization function parameters . We apply the quantization function module by module. Then, we train the network with gradually increased temperature .
Experiments
To compare with state-of-the-art methods, we evaluate our method on ImageNet (ILSVRC 2012). ImageNet has approximately million training images from thousand categories and thousand validation images. We evaluate our method on AlexNet (over-parameterized architectures) and ResNet-18/ResNet-50 (compact-parameterized architectures). We report our classification performance using Top-1 and Top-5 accuracies with networks quantized to Binary({0, 1}, 1 bit), Ternary({-1, 0, 1}, 2 bits), {-2, -1, 0, 1, 2} (denoted as 3 bits(2)), {-4, -2, -1, 0, 1, 2, 4 } (denoted as 3 bits(4)), and {-15, -14, , -1, 0, 1, , 14, 15 } (5 bits). All the parameters are fine-tuned from pretrained full-precision models.
All the images from ImageNet are resized to have pixels for the smaller dimension, and then a random crop of is selected for training. Each pixel of the input images is subtracted by the mean values and divided by variances. Random horizontal flipping is introduced for preprocessing. No other data augmentation tricks are used in the learning process. The batch size is set to . Following and , the parameters of the first convolutional layer and the last fully connected layer for classification are not quantized. For testing, images are resized to for the smaller side, and a center crop of is selected.
For our quantization function Eq. (7), to ensure all the input full-precision values lie in the linear region of our quantization function, the input scale is initialized as , where is the max absolute value of elements in and is the max absolute value of elements in . The output scale is initialized by , keeping the magnitude of the inputs unchanged after quantization.
Weight quantization: For binary quantization, only sigmoid function is needed; thus , , , and . For ternary quantization (), , and , ensuring that of the values in $n+1k\{c_{1},\dots,c_{n+1}\}b_{i}=\frac{c_{i}+c_{i+1}}{2}s_{\lfloor\frac{n}{2}\rfloor}=s_{\lfloor\frac{n}{2}\rfloor+1}=1b_{\lfloor\frac{n}{2}\rfloor}=-0.05,b_{\lfloor\frac{n}{2}\rfloor+1}=0.055\%$ are quantized to .
Activation quantization: Outputs of the ReLU units are used for activation quantization. It means that the block is Conv−-BN−-ReLU(-−Pooling)-−Quant in our method. The in Eq. (7) is set to because all activations are non-negative. For binary quantization(), only sigmoid function is needed, i.e. and . For two-bit quantization of activations (), and . is obtained by clustering as in weight quantization. We randomly sample 1000 samples from the dataset, and get the min/max activation values of the output layer by layer for ’s initialization .
The whole training process consists of phases. First, disable activation quantization and only train the quantization of weights. Second, fix the quantization of weights and only train the quantization of activations. Third, release quantization of both weights and activations until the model converges.
AlexNet: This network has five convolutional layers and two fully connected layers. This network is the mostly used benchmark for the quantization of neural networks. As in , we use AlexNet coupled with batch normalization layers. We update the model by stochastic gradient descent (SGD) with the momentum set to . The learning rate is initialized by and decayed by at epochs and respectively. The model is trained for at most epochs in total. The weight decay is set to . The temperature is set to and increased linearly w.r.t. the training epochs, i.e., . Gradients are clipped with a maximum L2 norm of .
The results of different quantization methods are shown in Table 1. denotes both weights and activations are binary quantized. As shown, our quantization network outperforms state-of-the-art methods in both weight quantization and activation quantization. Moreover, our quantization network is highly flexible. It is suitable for arbitrary bits quantizaion and can be applied for quantization of both weights and activation.
ResNet: The most common baseline architectures, including AlexNet, VGG and GoogleNet, are all over-parameterized by design for accuracy improvements. Therefore, it is easy to obtain sizable compression of these architectures with a small accuracy degradation. A more meaningful benchmark would be to quantize model architectures that are already with efficient parameters, e.g., ResNet. We use the ResNet-18 and ResNet-50 proposed in .
The learning rate is decayed by at epochs and , and the model is trained for at most epochs in total. The weight decay is set to . The temperature is set to and increased linearly w.r.t the training epochs (). The other settings are the same as these for AlexNet. The results of different quantization methods are shown in Table 2 and Table 3 for ResNet-18 and ResNet-50, respectively. We can see that the performance degradation of quantized models is larger than that on AlexNet. This is reasonable because that the parameters of the original model are more compact. It’s worth noting that even in such a compact model, our method still achieves lossless results with only bits. And as far as we know, we are the first to surpass the full-precision model on ResNet-18 with bits weight quantization.
2 Object Detection
In order to evaluate our quantization network on object detection task, we test it on the popular architecture SSD (single shot multibox detection) . The models are trained on Pascal VOC and train datasets, and are tested on Pascal VOC test dataset. We follow the same settings in and the input images are resized to . Except the final convolutional layers with kernels and the first convolution layer, parameters of all other layers in the backbone VGG16 are quantized.
We update the model by SGD with the momentum set to . The initial learning rate is set to for quantized parameters, for non-quantized parameters and decayed by at epochs and . Models are trained for epochs in total. The batch size is set to and the weight decay is . We increase the temperature by every epoch, i.e., . Gradients are clipped with maximum L2 norm of .
The results are given in Table 4. Here we compare our model with ADMM only because other baseline quantization methods did not report their performance on object detection task. As shown in Table 4, our model is slightly better than ADMM. This result is very promising since our method is much simpler and much more general than ADMM.
3 Ablation Experiments
In this section we discuss about the settings of our quantization network. All statistics are collected from the training process of Alexnet and ResNet-18 on ImageNet.
Configuration of Bias . Generally, the quantized values are pre-defined linearly (e.g., ) or logarithmically (e.g., ) with a scale factor . In this paper, we find that the distribution of full-precision parameters of pre-trained model is roughly subjected to Gaussian distribution, as shown in Fig. 3. It indicates that quantizing weights into linear or logarithmical intervals may not be the most suitable way. Thus, a non-uniform quantization (e.g. K-means clustering) is adopted to counterbalance this. So, we use the clustering centers to determine the boundaries of quantization intervals {}. The experimental results in Table 5 demonstrate the superior of non-uniform quantization over linear quantization. We also find that adaptive learning of biases during training does not show superiority over the fixed version. Therefore, we freeze the biases after initialization in all experiments.
Effect of layer-wise quantization. As shown in Fig. 3, the parameter magnitudes are quite different from layer to layer (full-precision network). Therefore, it is unsuitable and less efficient to use a shared quantization function across layers. We adopt layer-wise quantization in this paper, i.e., weights/activations from the same layer share the same quantization function and weights/activations from different layers use different quantization functions. Table 6 shows a comparison between shared quantization function across layers and layer-wise shared quantization function.
Effect of Temperature. As discussed in Section 3, the temperature controls the gap between the hard quantization function Eq. (12) in the inference stage and the soft quantization function Eq. (7) in the training stage. In order to investigate the effect of this gap to the performance of quantized network, we compare the testing accuracy of the models (trained with different s) when soft and hard quantization functions are adopted, as shown in Fig. 4. We can see that as the temperature increases, the difference between them is gradually reduced. Thus, gradually increasing temperature during training can achieve a good balance between model learning capacity and quantization gap.
4\{-4,+4\}. The gap between training and testing model converges when learning proceeds. Training from pre-trained model. In our training, the temperature parameter is increased linearly w.r.t. the training epochs. When training from scratch, the temperature may become quite large before the network is well-converged, and the saturated neurons will slow down the network training process and make the network stuck in bad minima. According to Table 7, training from a pre-trained model could greatly improve the performance compared to training from scratch.
Time-space complexity of the final model for inference. Table 8 shows the time-space complexities of the final quantization networks for inference based on VU9P FPGA evaluation. We can see that both time and space complexity are significantly reduced via low-bit quantization of Neural Networks.
Convergence for Temprature . The training process is very stable w.r.t. different s (shown in Fig. 5). The approximation of the final “soft” quantization function to a “hard” step function is determined by the final temperature, which is controlled by the maximum training epoch (). The increasing speed of temperature (e.g.10) controls the speed of convergence (or learning rate) from a “soft” to “hard” quantization (shown in Figure in our paper), and it is consistent with the learning progress of the backbone model. Practically, for differentbackbone models, we can tune in via performance on validation set as the way of learning rate for DL models.
Conclusion
This work focused on interpreting and implementing low-bit quantization of deep neural networks from the perspective of non-linear functions. Inspired by activation functions in DNNs, a soft quantization function is proposed and incorporated in deep neural networks as a new kind of activation function. With this differentiable non-linear quantization function embedded, quantization networks can be learned in an end-to-end manner. Our quantization method is highly flexible. It is suitable for arbitrary bits quantization and can be applied for quantization of both weights and activations. Extensive experiments on image classification and object detection tasks have verified the effectiveness of the proposed method.
Acknowledgements
This work was supported in part by the National Key R&D Program of China under contract No. 2017YFB1002203 and NSFC No. 61872329.