ABS: Automatic Bit Sharing for Model Compression

1 INTRODUCTION . Deep neural networks ( DNNs ) have achieved great success in many challenging computer vision tasks , including image classification ( Krizhevsky et al. , 2012 ; He et al. , 2016 ) and object detection ( Lin et al. , 2017a ; b ) . However , a deep model usually has a large number of parameters and consumes huge amounts of computational resources , which remains great obstacles for many applications , especially on resource-limited devices with limited memory and computational resources , such as smartphones . To reduce the number of parameters and computational overhead , many methods ( He et al. , 2019 ; Zhou et al. , 2016 ) have been proposed to conduct model compression by removing the redundancy while maintaining the performance . In the last decades , we have witnessed a lot of model compression methods , such as network pruning ( He et al. , 2017 ; 2019 ) and quantization ( Zhou et al. , 2016 ; Hubara et al. , 2016 ) . Specifically , network pruning reduces the model size and computational costs by removing redundant modules while network quantization maps the full-precision values to low-precision ones . It has been shown that sequentially perform network pruning and quantization is able to get a compressed network with small model size and lower computational overhead ( Han et al. , 2016 ) . However , performing pruning and quantization in a separate step may lead to sub-optimal results . For example , the best quantization strategy for the uncompressed network is not necessarily the optimal one after network pruning . Therefore , we need to consider performing pruning and quantization simultaneously . Recently , many attempts have been made to automatically determine the compression configurations of each layer ( i.e. , pruning ratios , and/or bitwidths ) , either based on reinforcement learning ( RL ) ( Wang et al. , 2019 ) , evolutionary search ( ES ) ( Wang et al. , 2020 ) , Bayesian optimization ( BO ) ( Tung & Mori , 2018 ) or differentiable methods ( Wu et al. , 2018 ; Dong & Yang , 2019 ) . In particular , previous differentiable methods formulate model compression as a differentiable searching problem to explore the search space using gradient-based optimization . As shown in Figure 1 ( a ) , each candi- date operation is maintained as a separate path , which leads to a huge number of trainable parameters and high computational overhead when the search space becomes combinatorially large . Moreover , due to the non-differentiable quantizer and pruning process , the optimization of heavily compressed candidate networks can be more challenging than that in the conventional search problem . In this paper , we propose a simple yet effective model compression method named Automatic Bit Sharing ( ABS ) to reduce the search cost and ease the optimization for the compressed candidates . Inspired by recent single-path neural architecture search ( NAS ) methods ( Stamoulis et al. , 2019 ; Guo et al. , 2020 ) , the proposed ABS introduces a novel single-path super-bit to encode all effective bitwidths in the search space instead of formulating each candidate operation as a separate path , as shown in Figure 1 ( b ) . Specifically , we build upon the observation that the quantized values of a high bitwidth can share the ones of low bitwidths under some conditions . Therefore , we are able to decompose the quantized representation into the sum of the lowest bit quantization and a series of re-assignment offsets . We then introduce learnable binary gates to encode the choice of bitwidth , including filter-wise 0-bit for pruning . By jointly training the binary gates and network parameters , the compression ratio of each layer can be automatically determined . The proposed scheme has several advantages . First , we only need to solve the search problem as finding which subset of the super-bit to use for each layer ’ s weights and activations rather than selecting from different paths . Second , we enforce the candidate bitwidths to share the common quantized values . Hence , we are able to optimize them jointly instead of separately , which greatly reduces the optimization difficulty from the discontinuity of discretization . Our main contributions are summarized as follows : • We devise a novel super-bit scheme that encapsulates multiple compression configurations in a unified single-path framework . Relying on the super-bit scheme , we further introduce learnable binary gates to determine the optimal bitwidths ( including filter-wise 0-bit for pruning ) . The proposed ABS casts the search problem as subset selection problem , hence significantly reducing the search cost . • We formulate the quantized representation as a gated combination of the lowest bitwidth quantization and a series of re-assignment offsets , in which we explicitly share the quantized values between different bitwidths . In this way , we enable the candidate operations to learn jointly rather than separately , hence greatly easing the optimization , especially in the non-differentiable quantization scenario . • We evaluate our ABS on CIFAR-100 and ImageNet over various network architectures . Extensive experiments show that the proposed method achieves the state-of-the-art performance . For example , on ImageNet , our ABS compressed MobileNetV2 achieves 28.5× Bit-Operation ( BOP ) reduction with only 0.2 % performance drop on the Top-1 accuracy . 2 RELATED WORK . Network quantization . Network quantization represents the weights , activations and even gradients in low-precision to yield compact DNNs . With low-precision integers or power-of-two representations , the heavy matrix multiplications can be replaced by efficient bitwise operations , leading to much faster test-time inference and lower power consumption . To improve the quantization performance , current methods either focus on designing accurate quantizers by fitting the quantizer to the data ( Jung et al. , 2019 ; Zhang et al. , 2018 ; Choi et al. , 2018 ; Cai et al. , 2017 ) , or seek to approximate the gradients due to the non-differentiable discretization ( Ding et al. , 2019 ; Louizos et al. , 2019 ; Zhuang et al. , 2020 ) . Moreover , most previous works assign the same bitwidth for all layers ( Zhou et al. , 2016 ; Zhuang et al. , 2018a ; 2019 ; Jung et al. , 2019 ; Jin et al. , 2019 ; Li et al. , 2020 ; Esser et al. , 2020 ) . Though attractive for simplicity , setting a uniform precision places no guarantee on optimizing network performance since different layers have different redundancy and arithmetic intensity . Therefore , several studies proposed mixed-precision quantization ( Wang et al. , 2019 ; Dong et al. , 2019 ; Wu et al. , 2018 ; Uhlich et al. , 2020 ) to set different bitwidths according to the redundancy of each layer . In this paper , based on the proposed quantization decomposition , we devise an approach that can effectively learn appropriate bitwidths for each layer through gradient-based optimization . NAS and pruning . Neural architecture search ( NAS ) aims to automatically design efficient architectures with low model size and computational costs , either based on reinforcement learning ( Pham et al. , 2018 ; Guo et al. , 2019 ) , evolutionary search ( Real et al. , 2019 ) or gradient-based methods ( Liu et al. , 2019a ) . In particular , gradient-based NAS has gained increased popularity , where the search space can be divided into the multi-path design ( Liu et al. , 2019a ; Cai et al. , 2019 ) and single-path formulation ( Stamoulis et al. , 2019 ; Guo et al. , 2020 ) , depending on whether adding each operation as a separate path or not . While prevailing NAS methods optimize the network topology , the focus of this paper is to search optimal compression ratios for a given architecture . Moreover , network pruning can be treated as fine-grained NAS , which aims at removing redundant modules to accelerate the run-time inference speed , giving rise to methods based on unstructured weight pruning ( Han et al. , 2016 ; Guo et al. , 2016 ) or structured channel pruning ( He et al. , 2017 ; Zhuang et al. , 2018b ; Luo et al. , 2017 ) . Based on channel pruning , our paper further takes quantization into consideration to generate more compact networks . AutoML for model compression . Recently , much effort has been put into automatically determining either the optimal pruning rate ( Tung & Mori , 2018 ; Dong & Yang , 2019 ; He et al. , 2018 ) , or the bitwidth ( Lou et al. , 2019 ; Cai & Vasconcelos , 2020 ) of each layer via hyper-parameter search , without relying on heuristics . In particular , HAQ ( Wang et al. , 2019 ) employs reinforcement learning to search bitwidth strategies with the hardware accelerator ’ s feedback . Meta-pruning ( Liu et al. , 2019b ) uses meta-learning to generate the weight parameters of the pruned networks and then adopts an evolutionary search algorithm to find the layer-wise sparsity for channel pruning . More recently , several studies ( Wu et al. , 2018 ; Cai & Vasconcelos , 2020 ) have focused on using differentiable schemes via gradient-based optimization . Closely related methods . To further improve the compression ratio , several methods propose to jointly optimize pruning and quantization strategies . In particular , some works only support weight quantization ( Tung & Mori , 2018 ; Ye et al. , 2019 ) or use fine-grained pruning ( Yang et al. , 2020 ) . However , the resultant networks can not be implemented efficiently on edge devices . Recently , several methods ( Wu et al. , 2018 ; Wang et al. , 2020 ; Ying et al. , 2020 ) have been proposed to consider filter pruning , weight quantization , and activation quantization jointly . In contrast to these methods , we carefully design the compression search space by sharing the quantized values between different candidate configurations , which significantly reduces the search cost and eases the optimization . Compared with those methods that share the similarities of using quantized residual errors ( Chen et al. , 2010 ; Gong et al. , 2014 ; Li et al. , 2017b ; van Baalen et al. , 2020 ) , our proposed method recursively uses quantized residual errors to decompose a quantized representation as a set of candidate bitwidths and parameterize the selection of optimal bitwidth via binary gates . Our proposed ABS and Bayesian Bits ( van Baalen et al. , 2020 ) are developed concurrently that share a similar idea of quantization decomposition . Critically , our ABS differs from Bayesian Bits in several aspects : 1 ) The quantization decomposition in our methods can be extended to non-powerof-two bit widths ( i.e. , b1 can be set to arbitrary appropriate integer values ) , which is a general case of the one in Bayesian Bits . 2 ) The optimization problems are different . Specifically , we formulate model compression as a single-path subset selection problem while Bayesian Bits casts the optimization of the binary gates to a variational inference problem that requires more relaxations and hyperparameters . 3 ) Our compressed models with less or comparable BOPs outperform those of Bayesian Bits by a large margin on ImageNet ( See Table 2 ) . 3 PROPOSED METHOD . 3.1 PRELIMINARY : NORMALIZATION AND QUANTIZATION FUNCTION . Without loss of generality , given a convolutional layer , let x and w be the activations of the last layer and its weight parameters , respectively . First , for convenience , following ( Choi et al. , 2018 ; Bai et al. , 2019 ) , we can normalize x and w into scale [ 0 , 1 ] by Tx and Tw , respectively : zx = Tx ( x ) = clip ( x vx , 0 , 1 ) , ( 1 ) zw = Tw ( w ) = 1 2 ( clip ( w vw , −1 , 1 ) + 1 ) , ( 2 ) where the function clip ( v , vlow , vup ) = min ( max ( v , vlow ) , vup ) clips any number v into the range [ vlow , vup ] , and vx and vw are trainable quantization intervals which indicate the range of weights and activations to be quantized . Then , we can apply the following function to quantize the normalized activations and parameters , namely zx ∈ [ 0 , 1 ] and zw ∈ [ 0 , 1 ] , to discretized ones : D ( z , s ) = s · round ( z s ) , ( 3 ) where round ( · ) returns the nearest integer of a given value and s denotes the normalized step size . Typically , for k-bit quantization , the normalized step size s can be computed by s = 1 2k − 1 . ( 4 ) After doing the k-bit quantization , we shall have 2k − 1 quantized values . Specifically , we obtain the quantization Q ( w ) and Q ( x ) by Q ( w ) = T−1w ( D ( zw , s ) ) = vw · ( 2 ·D ( zw , s ) − 1 ) , ( 5 ) Q ( x ) = T−1x ( D ( zx , s ) ) = vx ·D ( zx , s ) , ( 6 ) where T−1w and T −1 x denote the inverse functions of Tw and Tx , respectively .

Summary: The paper describes the method to determine optimal quantization bit-width and pruning configuration for the neural network compression. Different from other approaches, the proposed method integrates multiple bit configurations (including pruning) into a single architecture, which is named “Super-bit”. The architecture uses binary gates to automatically select bit resolution. In addition, the super-bit model is differentiable and jointly trainable with parameters.