Differentiable Joint Pruning and Quantization for Hardware Efficiency

Abstract

We present a differentiable joint pruning and quantization (DJPQ) scheme. We frame neural network compression as a joint gradient-based optimization problem, trading off between model pruning and quantization automatically for hardware efficiency. DJPQ incorporates variational information bottleneck based structured pruning and mixed-bit precision quantization into a single differentiable loss function. In contrast to previous works which consider pruning and quantization separately, our method enables users to find the optimal trade-off between both in a single training procedure. To utilize the method for more efficient hardware inference, we extend DJPQ to integrate structured pruning with power-of-two bit-restricted quantization. We show that DJPQ significantly reduces the number of Bit-Operations (BOPs) for several networks while maintaining the top-1 accuracy of original floating-point models (e.g., 53$\times$ BOPs reduction in ResNet18 on ImageNet, 43$\times$ in MobileNetV2). Compared to the conventional two-stage approach, which optimizes pruning and quantization independently, our scheme outperforms in terms of both accuracy and BOPs. Even when considering bit-restricted quantization, DJPQ achieves larger compression ratios and better accuracy than the two-stage approach.

Y. Lu—Work done during internship at Qualcomm AI Research.

Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc.

Appendices

A Quantization scheme in DJPQ

Figure 9 illustrates the proposed quantization scheme. First, a non-linear function is applied to map any weight input x to $\tilde{x}$ (shown in blue curve). Then a uniform quantization is applied to $\tilde{x}$. The quantized value $x_q$ is shown in red.

Fig. 9.

Illustration of quantization scheme. The blue curve gives the nonlinear mapping function. The red curve corresponds to the quantization value. (Color figure online)

B Experimental details

1.1 B.1 Comparison of DJPQ with DQ

To show that joint optimization of pruning and quantization outperforms quantization only scheme such as DQ, we plot in Fig. 10 a comparison of weight and activation bit-width for DQ and DJPQ. The results are for ResNet18 on ImageNet. We plot the pruning ratio curve of DJPQ in both figures to better show the pruning effect in the joint optimization scheme. As seen in the figure, there is no big difference between weight bit-width for DQ and DJPQ. However, the difference between activation bit-width for the two schemes is significant. Layer 0 to 6 in DJPQ has much larger activation bit-width than those in DQ, while those layers correspond to high pruning ratios in DJPQ. It provides a clear evidence that pruning and quantization can well tradeoff between each other in DJPQ, resulting in lower redundancy in the compressed model.

1.2 B.2 DJPQ Results for MobileNetV2

Figures 11 and 12 show the optimized bit-width and pruning ratio distributions, respectively. It is observed that earlier layers tend to have larger pruning ratios than the following layers. And for many layers in MobileNetV2, DJPQ is able to quantize the layers into to small bit-width.

Fig. 10.

Comparison of bit-width distributions of DQ and DJPQ for ResNet18 on ImageNet. The top figure plots the weight bit-width for each layer, while the bottom plots the corresponding activation bit-width. The green curve is the pruning ratio of DJPQ in each layer.(Color figure online)

1.3 B.3 Experimental Setup

The input for all experiments are uniformly quantized to 8 bits. For the two-stage scheme including ‘VIBNet+fixed quant.’ and ‘VIBNet+DQ’, VIBNet pruning is firstly optimized at learning rate 1e–3 with SGD, with a pruning learning rate scaling of 5. The strength $\gamma$ is set to 5e–6. In DQ for the pruned model, $\beta$ is chosen to 1e–11 and the learning rate is 5e–4. The learning rate scaling for quantization is 0.05. All the pruning threshold $\alpha _{th}$ is chosen to 1e–3. The number of epochs is 20 for each of the stage.

Fig. 11.

Bit-width for MobileNetV2 on ImageNet with DJPQ scheme

Fig. 12.

Pruning ratio for MobileNetV2 on ImageNet with DJPQ scheme

For DJPQ experiments on VGG7, the learning rate is set to 1e–3 with an ADAM optimizer. The initial bit-width is 6. The strength $\gamma$ and $\beta$ are 1e–6 and 1e–9, respectively. The scaling of learning rate for pruning and quantization is 10 and 0.05, respectively. For DJPQ on ResNet18, we chose a learning rate 1e–3 with SGD optimization. The initial bit-width for weights and activations are 6 and 8. The scaling of learning rate for pruning and quantization are 5 and 0.05, respectively. The strength $\gamma$ and $\beta$ are 1e–5 and 1e–10. For DJPQ on MobileNetV2, the learning rate is set to 1e–4 for SGD optimization. The initial bit-width is 8. The strength $\gamma$ and $\beta$ are set to 1e–8 and 1e–11. The learning rate scaling for pruning and quantization are selected to be 1 and 0.005, respectively.