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 work presents 1) Semi-Relaxed Quantization (SRQ), a method that targets learning low-bit neural networks, 2) DropBits, a method that performs dropout-like regularization on the bit width of the quantizers with an option to also automatically optimise the bit-width per layer according to the data, and 3) quantised lottery ticket hypothesis. SRQ is an extension of Relaxed Quantization (RQ), which is prior work, in two ways; firstly the authors replace the sampling from the concrete relaxation during training to deterministically selecting the mode (which is non-differentiable) and, secondly, they propose a specific straight-through gradient estimator (STE) than only propagates the gradient backwards for the elements that were selected in the forward pass. DropBits is motivated from the perspective of reducing the bias of the STE gradient estimator by randomly dropping grid points associated with a specific bit-width and then renormalising the SRQ distribution over the grid. This essentially induces stochasticity in the sampling distribution for the quantised value (which was removed before by selecting the mode in SRQ). The authors further extend DropBits in a way that allows for learning the drop probabilities for each bit-width, thus allowing for learning mixed-precision networks. Finally the authors postulate the quantised lottery ticket hypothesis, which refers to that “one can find the learned bit-width network which can perform better than the network with the same but fixed bit-widths from scratch”.