BSQ: Exploring Bit-Level Sparsity for Mixed-Precision Neural Network Quantization

1 INTRODUCTION . Numerous deep neural network ( DNN ) models have been designed to tackle real-world problems and achieved beyond-human performance . DNN models commonly demand extremely high computation cost and large memory consumption , making the deployment and real-time processing on embedded and edge devices difficult ( Han et al. , 2015b ; Wen et al. , 2016 ) . To address this challenge , model compression techniques , such as pruning ( Han et al. , 2015b ; Wen et al. , 2016 ; Yang et al. , 2020 ) , factorization ( Jaderberg et al. , 2014 ; Zhang et al. , 2015 ) and fixed-point quantization ( Zhou et al. , 2016 ; Wu et al. , 2019 ; Dong et al. , 2019 ) , have been extensively studied . Among them , fixed-point quantization works directly on the data representation by converting weight parameters originally in the 32-bit floating-point form to low-precision values in a fixed-point format . For a DNN model , its quantized version requires much less memory for weight storage . Moreover , it can better utilize fixed-point processing units in mobile and edge devices to run much faster and more efficiently . Typically , model compression techniques aim to reduce a DNN model size while maintaining its performance . The two optimization objectives in this tradeoff , however , have a contrary nature : the performance can be formulated as a differentiable loss function L ( W ) w.r.t . the model ’ s weights W ; yet the model size , typically measured by the number of non-zero parameters or operations , is a discrete function determined mainly by the model architecture . To co-optimize the performance and model size , some previous pruning and factorization methods relax the representation of model size as a differentiable regularization term R ( W ) . For example , group Lasso ( Wen et al. , 2016 ) and DeepHoyer ( Yang et al. , 2020 ) induce weight sparsity for pruning , and the attractive force regularizer ( Wen et al. , 2017 ) and nuclear norm ( Xu et al. , 2018 ) are utilized to induce low rank . The combined objective L ( W ) + αR ( W ) can be directly minimized with a gradient-based optimizer for optimizing the performance and model size simultaneously . Here , the hyperparameter α controls the strength of the regularization and governs the performance-size tradeoff of the compressed model . Unlike for pruning and factorization , there lacks a well-defined differentiable regularization term that can effectively induce quantization schemes . Early works in quantization mitigate the tradeoff exploration complexity by applying the same precision to the entire model . This line of research focuses on improving the accuracy of ultra low-precision DNN models , e.g. , quantizing all the weights to 3 or less bits ( Zhou et al. , 2016 ; Zhang et al. , 2018 ) , even to 1-bit ( Rastegari et al. , 2016 ) . These models commonly incur significant accuracy loss , even after integrating emerging training techniques like straight-through estimator ( Bengio et al. , 2013 ; Zhou et al. , 2016 ) , dynamic range scaling ( Polino et al. , 2018 ) and non-linear trainable quantizers ( Zhang et al. , 2018 ) . As different layers of a DNN model present different sensitivities with performance , a mixed-precision quantization scheme would be ideal for the performance-size tradeoff ( Dong et al. , 2019 ) . There have also been accelerator designs to support the efficient inference of mixed-precision DNN models ( Sharma et al. , 2018 ) . However , to achieve the optimal layer-wise precision configuration , it needs to exhaustively explore the aforementioned discrete search space , the size of which grows exponentially with the number of layers . Moreover , the dynamic change of each layer ’ s precision can not be formulated into a differentiable objective , which hinders the efficiency of the design space exploration . Prior studies ( Wu et al. , 2019 ; Wang et al. , 2019 ) utilize neural architecture search ( NAS ) , which suffers from extremely high searching cost due to the large space of mixed-precision quantization scheme . Recently , Dong et al . ( 2019 ) propose to rank each layer based on the corresponding Hessian information and then determine the relative precision order of layers based on their ranking . The method , however , still requires to manually select the precision level for each layer . Here , we propose to revisit the fixed-point quantization process from a new angle of bit-level sparsity : decreasing the precision of a fixed-point number can be taken as forcing one or a few bits , most likely the least significant bit ( LSB ) , to be zero ; and reducing the precision of a layer is equivalent to zeroing out a specific bit of all the weight parameters of the layer . In other words , the precision reduction can be viewed as increasing the layer-wise bit-level sparsity . By considering the bits of fixed-point DNN parameters as continuous trainable variables during DNN training , we can utilize a sparsity-inducing regularizer to explore the bit-level sparsity with gradient-based optimization , dynamically reduce the layer precision and lead to a series of mixed-precision quantization schemes . More specific , we propose Bit-level Sparsity Quantization ( BSQ ) method with the following contributions : • We propose a gradient based training algorithm for bit-level quantized DNN models . The algorithm considers each bit of quantized weights as an independent trainable variable and enables the gradient-based optimization with straight-through estimator ( STE ) . • We propose a bit-level group Lasso regularizer to dynamically reduce the weight precision of every layer and therefore induce mixed-precision quantization schemes . • BSQ uses only one hyperparameter , the strength of the regularizer , to trade-off the model performance and size , making the exploration more efficient . This work exclusively focuses on layer-wise mixed-precision quantization , which is the granularity considered in most previous works . However , the flexibility of BSQ enables it to explore mixedprecision quantization of any granularity with the same cost regardless of the search space size . 2 RELATED WORKS ON DNN QUANTIZATION . Quantization techniques convert floating-point weight parameters to low-precision fixed-point representations . Directly quantizing a pre-trained model inevitably introduces significant accuracy loss . So many of early research focus on how to finetune quantized models in low-precision configurations . As the quantized weights adopt discrete values , conventional gradient-based methods that are designed for continuous space can not be directly used for training quantized models . To mitigate this problem , algorithms like DoReFa-Net utilize a straight-through estimator ( STE ) to approximate the quantized model training with trainable floating-point parameters ( Zhou et al. , 2016 ) . As shown in Equation ( 1 ) , a floating-point weight element w is kept throughout the entire training process . Along the forward pass , the STE will quantize w to n-bit fixed-point representation wq , which will be used to compute the model output and loss L. During the backward pass , the STE will directly pass the gradient w.r.t . wq onto w , which enables w to be updated with the standard gradient-based optimizer . Forward : wq = 1 2n − 1 Round [ ( 2n − 1 ) w ] ; Backward : ∂L ∂w = ∂L ∂wq . ( 1 ) Early studies revealed that weights of different layers have different dynamic ranges . It is important to keep the dynamic range of each layer for maintaining the model performance , especially for quantized models . He et al . ( 2016b ) and Polino et al . ( 2018 ) propose to explicitly keep track of the dynamic range of each layer by scaling all the weight elements in a layer to the range of [ 0,1 ] at every training step , before applying the quantization STE . Other techniques , such as learnable nonlinear quantifier function ( Zhang et al. , 2018 ) and incremental quantization ( Zhou et al. , 2017 ) , are also useful in improving the performance of quantized models . However , it is still very difficult to quantize the entire DNN model to a unified ultra-low precision without incurring significant accuracy loss . Recent research shows that different layers in a DNN model contribute to the overall performance in varying extents . Therefore mixed-precision quantization scheme that assigns different precision to layers ( Wu et al. , 2019 ; Dong et al. , 2019 ) presents a better accuracy-compression tradeoff . The challenge lies in how to determine the quantization scheme , i.e. , the precision of each layer , as it needs to explore a large and discrete search space . Some works design quantization criteria based on concepts like “ noise gain ” ( Sakr & Shanbhag , 2018 ; 2019 ) to constraint the relationship between each layer ’ s precision and thus largely reduce the search space , yet those criteria are often heuristic , preventing these methods to reach ultra-low precision and find the optimal tradeoff point between model size and accuracy . Other works utilize neural architecture search ( NAS ) . For example , Wang et al . ( 2019 ) consider the precision assignment of each layer as an action and seek for the optimal design policy via reinforcement learning . Wu et al . ( 2019 ) combine all possible design choices into a “ stochastic super net ” and approximate the optimal scheme via sampling . However , the cost of NAS methods scales up quickly as the quantization search space grows exponentially with the number of layers . Common practices of constraining the search cost include limiting the precision choices or designing the quantization scheme in a coarse granularity . A recent line of research work attempts to rank layers based on their importance measured by the sensitivity or Hessian information . Higher precision will then be assigned to more important layers ( Dong et al. , 2019 ) . The exact precision of each layer , however , needs to be manually selected . So these methods can not adequately explore the whole search space for the optimal quantization scheme . 3 THE BSQ METHOD . BSQ aims to obtain an optimal mixed-precision quantization scheme through a single-pass training process of a quantized model . In this section , we first introduce how to convert a DNN model to the bit representation and propose a gradient-based algorithm for training the resulted bit-level model . A bit-level group Lasso regularizer is then proposed to induce precision reduction . In the end , we elaborate the overall training objective of BSQ and the dynamical precision adjustment procedure . 3.1 TRAINING THE BIT REPRESENTATION OF DNN . As illustrated in Figure 1 ( a ) , we convert a floating-point weight matrix W of a pretrained network to its bit representation through a pipeline of scaling , quantization and binary conversion . Similar to the practice in ( He et al. , 2016b ; Polino et al. , 2018 ) , we retain the dynamic range of W by scaling all the elements to the range of [ 0 , 1 ] before applying quantization . However , these prior works always scale the largest element to 1 to fully utilize all the quantized bins at every training step , which makes the dynamic precision reduction impossible . Instead , our method conducts the scaling only once , which is right before the bit representation training . Formally , before converting W to its bit representation , we first extract its dynamic range as W = s ·Ws , where s = max|W | is the scaling factor and Ws is the scaled weight matrix . The absolute value of any element ws in Ws is within the range of [ 0 , 1 ] . Now we apply an n-bit uniform quantization to the absolute value of ws such as wq = Round [ |ws|× ( 2n − 1 ) ] / ( 2n − 1 ) . Then wq can be exactly represented by a n-bit binary number as wq = [ ∑n−1 b=0 w ( b ) s 2b ] / ( 2n−1 ) , where w ( b ) s denotes the bth bit in the binary representation . Till this point , W in the floating-point form is replaced with W ≡ sign ( W ) sWq ≡ sign ( W ) s 2n − 1 n−1∑ b=0 W ( b ) s 2 b , ( 2 ) where denotes the element-wise Hadamard product . We consider the bit representation W ( b ) s where b ∈ [ 0 , n− 1 ] and the scaling factor s as independent trainable variables in the training process . Note that W ( b ) s is composed of binary values by definition and sign ( W ) is a discrete function . Neither of them can be directly trained with gradient descent . To mitigate the binary constraint of W ( b ) s , we adopt the STE proposed by Bengio et al . ( 2013 ) during the training process . As shown in Equation ( 1 ) , STE enables a quantized model to be trained with continuous floating-point weights . Specifically , the STE for the bit representation training is defined as : Forward : Wq = 1 2n − 1 Round [ n−1∑ b=0 W ( b ) s 2 b ] ; Backward : ∂L ∂W ( b ) s = 2b 2n − 1 ∂L ∂Wq . ( 3 ) STE relaxes the binary constraint and allows gradient updates for the elements in W ( b ) s . As illustrated in Figure 1 ( b ) , during the forward pass , s ·Wq will be used to reconstruct the model weight W and compute the loss , which demonstrates the performance of the current model after quantization . The gradient w.r.t . Wq from the back-propagation will be passed through the rounding function and updated on the continuous values of W ( b ) s . The proposed bit representation can therefore be trained with any gradient-based optimizer . The proposed bit representation training only leads to minimal computational and run-time memory overhead comparing to the normal back propagation procedure . From the memory consumption perspective , the bit representation training treats each bit as separated floating-point trainable variables , so a N -bit model in bit representation will have N times more parameters and gradients to be stored comparing to that of the baseline training . Though for actual run-time memory consumption , the hidden feature between each layer consumes a significantly larger memory than weights and gradients . As the bit representation does not affect the hidden features , the increase in trainable variables does not lead to significant increase in run-time memory consumption . From the perspective of computation cost , note that the gradient w.r.t . each W ( b ) s can be computed as the gradient w.r.t . the corresponding Wq scaled by a power of 2 . So under a N -bit scheme there will only be N additional scaling for each parameter comparing to the normal training . These additional computations are very cheap comparing to the floating-point operations involved in back propagation . So the proposed bit representation training only leads to minimal computational overhead comparing to a normal back propagation . We restrict the value of W ( b ) s within [ 0 , 2 ] throughout the training , so that the corresponding Wq has the chance to increase or decrease its precision in the “ precision adjustment ” step , which will be discussed in Section 3.3 . This is enforced by trimming W ( b ) s to 0 or 2 if it exceeds the range after a training step . To enable the dynamic update of sign ( W ) during training , we separate the positive and negative elements in Ws as Ws = ( Wp−Wn ) before quantization . Here Wp =Ws 1 ( Ws ≥ 0 ) contains all the positive elements and Wn = −Ws 1 ( Ws < 0 ) includes the absolute value of all the negative weight elements . Wp and Wn will be respectively converted to W ( b ) p and W ( b ) n by following the process in Equation ( 2 ) , so that W ( b ) s = W ( b ) p −W ( b ) n . Note that the replacement of W ( b ) s with W ( b ) p −W ( b ) n does not introduce any non-differentiable function . Therefore all elements in W ( b ) p and W ( b ) n can take continuous values between [ 0 , 2 ] and be trained with the bit representation STE in Equation ( 3 ) . As such , the original weight matrix W is converted into trainable variables W ( b ) p , W ( b ) n and s throughout the BSQ training process .

Summary: This paper basically proposed to learn the quantization bits (precision) in each layer. Specially, weights are constructed with binary representation as $W_s = \[W_s^1,...,W_s^b\]$. During training, $W_s^i$ is relaxed to $ \in \[0, 2\]$. And a group sparsity is imposed to all $W_s^i$ for all weights in a layer, leading to certain $W_s^i \to 0$, thus cancelling the bit allocation in $i$-th. Experimental results is promising.