Semi-Relaxed Quantization with DropBits: Training Low-Bit Neural Networks via Bitwise Regularization

1 INTRODUCTION . Deep neural networks have achieved great success in various computer vision applications such as image classification , object detection/segmentation , pose estimation , action recognition , and so on . However , state-of-the-art neural network architectures require too much computation and memory to be deployed to resource-limited devices . Therefore , researchers have been exploring various approaches to compress deep neural networks to reduce their memory usage and computation cost . In this paper , we focus on neural network quantization , which aims to reduce the bit-width of a neural network while maintaining competitive performance with a full-precision network . It is typically divided into two groups , uniform and heterogeneous quantization . In uniform quantization , one of the simplest methods is to round the full-precision weights and activations to the nearest grid points : x̂ = αb xα + 1 2c where α controls the grid interval size . However , this naı̈ve approach incurs severe performance degradation on large datasets . Recent network quantization methods tackle this problem from different perspectives . In particular , Relaxed Quantization ( RQ ) ( Louizos et al. , 2019 ) employs Gumbel-Softmax ( Jang et al. , 2017 ; Maddison et al. , 2017 ) to force weights and activations to be located near quantization grids with high density . Louizos et al . ( 2019 ) notice the importance of keeping the gradient variance small , which leads them to use high Gumbel-Softmax temperatures in RQ . However , such high temperatures may cause a large quantization error , thus preventing quantized networks from achieving comparable performance to full-precision networks . To resolve this issue , we first propose Semi-Relaxed Quantization ( SRQ ) that uses the mode of the original categorical distribution in the forward pass , which induces small quantization error . It is clearly distinguished from Gumbel-Softmax choosing argmax among the samples from the concrete distribution . To cluster weights cohesively around quantization grid points , we devise a multi-class straight-through estimator ( STE ) that allows for efficient gradient-based optimization as well . As this STE is biased like a traditional one ( Bengio et al. , 2013 ) for the binary case , we present a novel technique , DropBits to reduce the distribution bias of the multi-calss STE in SRQ . Motivated from Dropout ( Srivastava et al. , 2014 ) , DropBits drops bits rather than neurons/filters to train low-bit neural networks under SRQ framework . In addition to uniform quantization , DropBits allows for heterogeneous quantization , which learns different bit-width per parameter/channel/layer by dropping redundant bits . DropBits with learnable bit-drop rates adaptively finds out the optimal bit-width for each group of parameters , possibly further reducing the overall bits . In contrast to recent studies ( Wang et al. , 2019 ; Uhlich et al. , 2020 ) in heterogeneous quantization that exhibit almost all layers possess at least 4 bits , up to 10-bit , our method yields much more resource-efficient low-bit neural networks with at most 4 bits for all layers . With trainable bit-widths , we also articulate a new hypothesis for quantization where one can find the learned bit-width network ( termed a ‘ quantized sub-network ’ ) which can perform better than the network with the same but fixed bit-widths from scratch . Our contribution is threefold : • We propose a new quantization method , Semi-Relaxed Quantization ( SRQ ) that introduces the multi-class straight-through estimator to reduce the quantization error of Relaxed Quantization for transforming continuous activations and weights to discrete ones . We further present a novel technique , DropBits to reduce the distribution bias of the multi-class straight-through estimator in SRQ . • Extending DropBits technique , we propose a more resource-efficient heterogeneous quantization algorithm to curtail redundant bit-widths across groups of weights and/or activations ( e.g . across layers ) and verify that our method is able to find out ‘ quantized sub-networks ’ . • We conduct extensive experiments on several benchmark datasets to demonstrate the effectiveness of our method . We accomplish new state-of-the-art results for ResNet-18 and MobileNetV2 on the ImageNet dataset when all layers are uniformly quantized . 2 RELATED WORK . While our goal in this work is to obtain an extremely low-bit neural network both for weights and activations , here we broadly discuss existing quantization techniques with various goals and settings . BinaryConnect ( Courbariaux et al. , 2015 ) first attempted to binarize weights to ±1 by employing deterministic or stochastic operation . To obtain better performance , various studies ( Rastegari et al. , 2016 ; Li et al. , 2016 ; Achterhold et al. , 2018 ; Shayer et al. , 2018 ) have been conducted in binarization and ternarization . To reduce hardware cost for inference , Geng et al . ( 2019 ) proposed softmax approximation via a look-up table . Although these works effectively decrease the model size and raise the accuracy , they are limited to quantizing weights with activations remained in full-precision . To take full advantage of quantization at run-time , it is necessary to quantize activations as well . Researchers have recently focused more on simultaneously quantizing both weights and activations ( Zhou et al. , 2016 ; Yin et al. , 2018 ; Choi et al. , 2018 ; Zhang et al. , 2018 ; Gong et al. , 2019 ; Jung et al. , 2019 ; Esser et al. , 2020 ) . XNOR-Net ( Rastegari et al. , 2016 ) , the beginning of this line of work , exploits the efficiency of XNOR and bit-counting operations . QIL ( Jung et al. , 2019 ) also quantizes weights and activations by introducing parametrized learnable quantizers that can be trained jointly with weight parameters . Esser et al . ( 2020 ) recently presented a simple technique to approximate the gradients with respect to the grid interval size to improve QIL . Nevertheless , these methods do not quantize the first or last layer , which leaves a room to improve power-efficiency . For ease of deployment in practice , it is inevitable to quantize weights and activations of all layers , which is the most challenging . Louizos et al . ( 2019 ) proposed to cluster weights by using GumbelSoftmax , but it shows drawbacks as we will discuss in Section 3.2 . Jain et al . ( 2019 ) presented efficient fixed-point implementations by formulating the grid interval size to the power of two , but they quantized the first and last layer to at least 8-bit . Zhao et al . ( 2020 ) proposed to quantize the grid interval size and network parameters in batch normalization for the deployment of quantized models on low-bit integer hardware , but it requires a specific accelerator only for this approach . As another line of work , Fromm et al . ( 2018 ) proposed a heterogeneous binarization given pre-defined bit-distribution . HAWQ ( Dong et al. , 2019 ) determines the bit-width for each block heuristically based on the top eigenvalue of Hessian . Unfortunately , both of them do not learn optimal bit-widths for heterogeneity . Toward this , Wang et al . ( 2019 ) and Uhlich et al . ( 2020 ) proposed a layer-wise heterogeneous quantization by exploiting reinforcement learning and learning dynamic range of quantizers , respectively . However , their results exhibit that almost all layers have up to 10-bit ( at least 4-bit ) , which would be suboptimal . Lou et al . ( 2020 ) presented a channel-wise heterogeneous quantization by exploiting hierarchical reinforcement learning , but channel-wise precision limits the structure of accelerators , thereby restricting the applicability of the model . 3 METHOD . In this section , we briefly review Relaxed Quantization ( RQ ) ( Louizos et al. , 2019 ) and propose Semi-Relaxed Quantization , which selects the nearest grid point in the forward pass to decrease the quantization error . To make it learnable and to cluster compressed parameters cohesively , SRQ expresses the nearest grid selection of the forward pass as the equivalent form , the combination of logistic distribution and argmax , and performs the backward pass on it . Then , we present DropBits technique to reduce the distribution bias of SRQ and its extension to heterogeneous quantization . 3.1 PRELIMINARY : RELAXED QUANTIZATION . Relaxed Quantization ( RQ ) considers the following quantization grids for weights : Ĝ = α [ −2b−1 , . . . , 0 , . . . , 2b−1 − 1 ] = : [ g0 , . . . , g2b−1 ] where b is the bit-width and a learnable parameter α > 0 for each layer controls a grid interval . When quantizing activations , the grid points in Ĝ start from zero since the output of ReLU is always non-negative . Then , x ( a weight or an activation ) is perturbed by noise as x̃ = x+ , which enables gradient-based optimization for non-differentiable rounding . The noise follows a distribution p ( ) = Logistic ( 0 , σ ) so that p ( x̃ ) is governed by Logistic ( x , σ ) where σ represents the standard deviation . Under p ( x̃ ) , we can easily compute the unnormalized probability of x̃ being quantized to each grid point gi in a closed form as below : πi = p ( x̂ = gi|x , α , σ ) = Sigmoid ( ( gi + α/2− x ) /σ ) − Sigmoid ( ( gi − α/2− x ) /σ ) , ( 1 ) where x̂ denotes a quantized realization of x̃ . Note that the cumulative distribution function of the logistic distribution is just a sigmoid function . Finally , given unnormalized categorical probability π = { πi } 2 b−1 i=0 for grid points Ĝ = { gi } 2b−1 i=0 , RQ discretizes x to x̂ = r ·Ĝ by sampling r = { ri } 2b−1 i=0 from the concrete distribution ( Jang et al. , 2017 ; Maddison et al. , 2017 ) with a temperature τ : ui ∼ Gumbel ( 0 , 1 ) , ri = exp ( ( log πi + ui ) /τ ) ∑2b−1 j=0 exp ( ( log πj + uj ) /τ ) , x̂ = 2b−1∑ i=0 rigi . ( 2 ) The algorithm of RQ is described in Appendix in detail . 3.2 SEMI-RELAXED QUANTIZATION - FIXING PITFALLS OF RQ . Although RQ achieves competitive performance with both weights and activations of neural networks quantized , the quantization probability modeling of RQ may still incur large quantization error , thereby yielding suboptimal performance . To be specific , Louizos et al . ( 2019 ) recommend high temperatures for the concrete distribution ( e.g . τ = 1.0 or 2.0 ) in ( 2 ) since exploiting low temperatures hinders networks from converging due to high variance of gradients . However , it turns out that the concrete distribution with such a high temperature is almost similar to the uniform distribution . As a concrete example , we consider 2-bit quantization with Ĝ = α [ −2 , −1 , 0 , 1 ] for a fixed scale parameter α > 0 , σ = α/3 , and we set τ to 1.0 as in Louizos et al . ( 2019 ) . Now , suppose that the original weight value is α/2 . As in Figure 1- ( b , d , e ) , ? can be sporadically quantized to below zero by RQ as the original categorical distribution has support for −α and −2α . It may be okay on average , but RQ computes only one sample in each forward pass due to computational burden , which can accidentally lead to very large quantization error for these particular sample . To avoid the counterintuitive-sample with large quantization error as seen in Figure 1- ( b , d , e ) , we propose ‘ Semi-Relaxed Quantization ’ ( SRQ ) which rather directly considers the original categorical distribution in Figure 1- ( g ) . To be concrete , for a weight or an activation x , the probability of x being quantized to each grid is ri = πi/ ∑2b−1 j=0 πj for i ∈ { 0 , · · · , 2b − 1 } with b-bit precision , where πi is computed as ( 1 ) . In such a manner , selecting a grid point for x can be thought of as sampling from the categorical distribution with categories Ĝ = { gi } 2 b−1 i=0 and the corresponding probabilities r = { ri } 2 b−1 i=0 as illustrated in Figure 1- ( g ) . Then , the grid point gimax with imax = argmaxi ri would be the most reasonable speculation due to the highest probability . SRQ therefore chooses the mode of the original categorical distribution , gimax and assign it to x̂ , entirely discriminated from Gumbel-Softmax which selects the argmax among samples from the concrete distribution . As a result , SRQ does not suffer from counterintutive-sample problem that RQ encounters at all . The last essential part for SRQ is to handle the non-differentiable argmax operator in computing imax . Toward this , we propose a multi-class straight-through estimator ( STE ) that allows for backpropagating through a non-differentiable categorical sample by approximating ∂L/∂rimax to ∂L/∂yimax and ∂L/∂ri to zero for i 6= imax , where L is the cross entropy between the true label and the prediction made by a quantized neural network as delineated in the previous paragraph and yimax is the imax-th entry of the one-hot vector y . The forward and backward passes of SRQ are summarized as follows . Forward : y = one hot [ argmax i ri ] , Backward : ∂L ∂rimax = ∂L ∂yimax and ∂L ∂ri = 0 for i 6= imax ( 3 ) Such a formulation brings two important advantages in network quantization . First of all , ( 3 ) makes the variance of gradient estimator become zero . Since SRQ always chooses the mode of the original categorical distribution ( i.e. , there is no randomness in the forward pass of SRQ ) , and the gradient of loss function L with respect to the individual categorical probabilities is defined as zero everywhere except for the coordinate corresponding to the mode , the variance of gradients in SRQ is indeed zero . The other advantage is that the backward pass ( 3 ) can cluster network weight parameters cohesively . Under the assumption that ri = πi , ∂L∂x is proportional to ∂πimax ∂x = 1 σ ( Sigmoid ( gimax +α/2−x σ ) Sigmoid ( − gimax +α/2−xσ ) −Sigmoid ( gimax−α/2−x σ ) Sigmoid ( − gimax−α/2−x σ ) ) 1 . When x is close to gimax , ∂πimax/∂x is nearly zero , so is ∂L/∂x . With an appropriate learning rate , x converges to gimax , which leads SRQ to cluster weights better than RQ as shown in Figure 2 . Although ∂L∂x is almost zero , α is still trained . After α is updated , there is a gap between x and α so that x can be trained . Hence , the network will continue to train until it finds the optimal α .

Summary: This paper deals with network quantization. It proposes Semi-Relaxed Quantization (SRQ) that uses a multi-class straight-through estimator to effectively reduce the bias and variance, along with a new regularization technique, DropBits that replaces dropout regularization to randomly drop the bits.  Extensive experiments are conducted to validate our method on various benchmark datasets and network architectures.