Differentiable Fine-grained Quantization for Deep Neural Network Compression
Abstract
Neural networks have shown great performance in cognitive tasks.
When deploying network models on mobile devices with limited resources, weight quantization has been widely adopted.
Binary quantization obtains the highest compression but usually results in big accuracy drop.
In practice, 8-bit or 16-bit quantization is often used aiming at maintaining the same accuracy as the original 32-bit precision.
We observe different quantization schemes have different accuracy impact on different layers.
Thus judiciously selecting different precision for different layers/structures can potentially produce more efﬁcient models compared to traditional quantization methods by striking a better balance between accuracy and compression rate.
In this work, we propose a ﬁne-grained quantization approach for deep neural network compression by relaxing the search space of quantization bitwidth from discrete to a continuous domain.
The proposed approach applies gradient descend based optimization to generate a mixed-precision quantization scheme that outperforms the accuracy of traditional quantization methods under the same compression rate.
Introduction
State-of-the-art neural networks have demonstrated promising performance in tasks such as image classiﬁcation and object detection [1][2][3][4].
These network models are designed for high accuracy with less consideration in the computational cost and inference delay.
Thus deploying them on resource-constrained platform such as mobile phones is usually inefﬁcient or even infeasible.
Even the recently proposed MobileNet [5], which is customized for mobile platforms, has a relatively large 4.24M parameters.
Extensive studies have been carried out in developing network models for resource constraint platforms.
Quantization is one of the most popular approaches [6][7].
For example, Bai et al.
recently proposed a proximal operators for training low-precision deep neural netwokrs [8].
Courbariaux et al.
proposed a radical binary representation of the inputs, weights, and activations [6].
Rastegari et al.
[7] theoretically analyzed the binary network and introduced a scaling scheme for their XNOR-Net—a network based on [6] with higher accuracy.
Despite the signiﬁcant improvement on inference accuracy, none of the above networks is able to achieve comparable accuracy as the full-precision counterparts.
We observe that different layers may have different accuracy sensitivity of quantization, thus a ﬁne-grained quantization for each layer has the penitential to preserve accuracy under the same compression rate (deﬁned as the ratio of original model size and compressed model size) compared to traditional course-grained quantization that uses the same quantization for the entire model.
The dotted line in Figure 1 shows the trade-off between accuracy and compression rate in traditional quantization for VGG-16.
Our goal is to push the trade-off between accuracy and compression rate into the shaded region of Figure 1 to achieve better compression efﬁciency, i.e., higher accuracy under the same compression rate.
To achieve this, we propose a ﬁne-grained quantization approach that relaxes the search space of quantization bitwidth from discrete to continuous domain and applies gradient descend optimization to generate best quantization scheme for each layer, i.e., applies lower bit for less quantization sensitive layers while preserving high bit precision for quantization sensitive layers.
Our experimental results show that the proposed approach outperforms the accuracy of traditional quantization methods under the same compression rate.
Proposed Approach
In this section, we propose a methodology to judiciously determine the best quantization scheme for each layer based on each layer’s accuracy sensitivity of quantization.
For easy description, we only use two-level quantization: binary and 8-bit quantization as an example to illustrate our approach and conduct a preliminary evaluation.
It is straightforward to extend to more quantization levels.
We relax the discrete variables to a continuous domain as the ﬁner granularity of which can provide more accurate indication in quantization searching.
We adopt gradient descent based searching algorithm as it is fast and can be easily deployed in different machine learning frameworks.
We use Softmax function to relax the search space from discrete to continuous.
We denote the output of layer i with continuous relaxation as q i .
For example, binary and 8-bit quantization represented as q i 0 and q i 1 .
Softmax ( q i 0 , q i 1 ) can be translated as the probability of binary and 8-bit quantization, respectively.
Thus q i can be computed as
[Equation]
where B is the batch normalization operation.
The output q i is used as the input of the following layer.
The search space for a network with k layer is α = { α 0 , α 1 , ..., α k − 1 } .
To explore the trade-off between different quantization schemes, we model the target objective function as
[Equation]
where
[Equation]
Here G represents the model size, L is the cross entropy loss, θ is the expected maximum loss, w denotes the weights of the model, and α represents the coefﬁcient of either quantization method (binary or 8-bit) in a certain layer.
In our model, (3) is the constraint for optimization problem (2).
We can rewrite the above as a bi-level optimization problem:
[Equation]
To solve this bi-level optimization problem, we adopt the approximate algorithm in [9].
First, we retrain the network to ﬁnd the weights that result in the minimal loss on the training set.
Then the
Lagrange multiplier problem is solved by ﬁxing the weights.
As shown in Algorithm 1, solving the
Lagrange multiplier problem starts with maximizing the target function w.r.t.
λ : if L ( α, w ∗ ) − θ ≤ 0 , λ approaches 0; otherwise λ approaches inﬁnite.
Here θ is a tunable hyperparameter representing the tolerance of accuracy drop.
Larger θ tolerates less accuracy drop but may also result in smaller compression rate.
While smaller θ can potentially achieve a higher compression rate, it may cause larger accuracy drop.
Our setting of θ is using the (expected or target) loss in full precision model.
Finally we minimize the target function w.r.t.
α .
Once obtaining the hyperparameter set α with the best trade-off, we retrain the quantization and ﬁne tune the quantized weights to generate the ﬁnal network model.
Experimental Evaluation
We evaluate our proposed methodology on a pretrained 2-layer depth- [Equation] separable convolution neural network using MNIST data set as well as VGG-16 neural network model using CIFAR-10 data set.
For each model, we compare our approach with the following baselines: 32-bit ﬂoating point (full-precision) model, 8-bit ﬁxed precision model, and binary ﬁxed precision model.
As shown in Table 1, the results of MNIST experiment suggest that our algorithm is capable of ﬁnd a quantization scheme that achieves 28x compression rate while keeping the accuracy drop less than 0.5%. In
CIFAR-10 experiment, we set θ as 0.6 for VGG-16. Comparing to whole binary quantization, our approach obtains a compression rate that is very close to binary quantization while gaining 1.5% more accuracy.
Figure 2 shows the memory consumption of our model and the original 32-bit full precision model.
The memory usage is dramatically decreased especially at the middle layers.
It is worth mentioning that our method is orthogonal to weight pruning.
Combing with state-of-the-art pruning methods [10][11] which achieve approximately 30x compression rate, the overall compression rate can be up to approximately 900x.
Algorithm 1: Differentiable quantization
Initialization; while not converged do
Update weights w by descending ∇ w L train ( w, α ) ; if L valid ( w, α ) − θ ≤ 0 then λ ← 0 ; else λ ← inf ; end
Update probability α by descending ∇ α ( G ( α ) + λ ( L valid ( w ∗ , α ) − θ )) ; end
[Table]
ourresultbaseline the mixed precision model generated by our algorithm.
The width of a rectangle denotes the size (i.e., memory consumption) of the corresponding layer.
In this paper, we propose a differentiable mixed-precision search method for compressing deep neural networks efﬁciently. Unlike the traditional quantization methods, our approach relaxes quantization bitwidths to a continuous domain and combined with loss function.
Deep neural networks can be either quantized from the start of training phase or from a pretrained model using our proposed methodology.
Moreover, our approach ensures quantized model remain a similar accuracy while being compressed up to 30X.
The proposed methodology is not tied into any speciﬁc neural network topology, so it can potentially be extended to mixed-precision quantization of different neural network architectures, such as RNN and LSTM.
We are currently working on providing more quantization options for each layer.
For example, each layer can be quantized to x bits, and x ∈ { 1 , 2 , ..., 32 } .
These new quantization options drastically increase the search space.
Therefore, We plan to design a predictor combined with autoencdoer-decoder architecture to expedite the search process of layer-wise quantization.
Acknowledgement
This work is supported in part by the following grants: National Science Foundation CCF-1756013,
IIS-1838024, 1717657 and Air Force Research Laboratory FA8750-18-2-0057.