GQ-Net: Training Quantization-Friendly Deep Networks

1 INTRODUCTION . Neural network quantization is a technique to reduce the size of deep networks and to bypass computationally and energetically expensive floating-point arithmetic operations in favor of efficient integer arithmetic on quantized versions of model weights and activations . Network quantization has been the focus of intensive research in recent years ( Rastegari et al. , 2016 ; Zhou et al. , 2016 ; Jacob et al. , 2018 ; Krishnamoorthi , 2018 ; Jung et al. , 2018 ; Louizos et al. , 2019 ; Nagel et al. , 2019 ; Gong et al. , 2019 ) , with most works belonging to one of two categories . The first line of work quantizes parts of the network while leaving a portion of its operations , e.g . computations in the first and last network layers in floating point . While such networks can be highly efficient , using bitwidths down to 5 or 4 bits with minimal loss in network accuracy ( Zhang et al. , 2018 ; Jung et al. , 2018 ) , they may be difficult to deploy in certain practical settings , due to the complexity of extra floating point hardware needed to execute the non-quantized portions of the network . Another line of work aims for ease of real world deployment by quantizing the entire network , including all weights and activations in all convolutional and fully connected layers ; we term this scheme strict quantization . Maintaining accuracy under strict quantization is considerably more challenging . While nearly lossless 8-bit strictly quantized networks have been proposed ( Jacob et al. , 2018 ) , to date state-of-the-art 4 bit networks incur large losses in accuracy compared to full precision reference models . For example , the strict 4-bit ResNet-18 model in Louizos et al . ( 2019 ) has 61.52 % accuracy , compared to 69.76 % for the full precision model , while the strict 4-bit MobileNet-v2 model in Krishnamoorthi ( 2018 ) has 62.00 % accuracy , compared to 71.88 % accuracy in full precision . To understand the difficulty of training accurate low-bitwidth strictly quantized networks , consider a common training procedure which begins with a pre-trained network , quantizes the model , then applies fine-tuning using straight-through estimators ( STE ) for gradient updates until the model achieves sufficient quantized accuracy . This process faces two problems . First , as the pre-trained model was not initially trained with the task of being subsequently quantized in mind , it may not be “ quantization-friendly ” . That is , the fine-tuning process may need to make substantial changes to the initial model in order to transform it to an accurate quantized model . Second , fine-tuning a model , especially at low bitwidths , is difficult due to the lack of accurate gradient information provided by STE . In particular , fine-tuning using STE is done by updating a model represented internally with floating point values using gradients computed at the nearest quantizations of the floating point values . Thus for example , if we apply 4 bit quantization to floating point model parameters in the range [ 0 , 1 ] , a random parameter will incur an average round-off error of 1/32 , which will be incorporated into the error in the STE gradient for this parameter , leading to possibly ineffective fine-tuning . To address these problems , we propose GQ-Net , a guided quantization training algorithm . The main goal of GQ-Net is to produce an accurate and quantization-friendly full precision model , i.e . a model whose quantized version , obtained by simply rounding each full precision value to its nearest quantized point , has nearly the same accuracy as itself . To do this , we design a loss function for the model which includes two components , one to minimize error with respect to the training labels , and another component to minimize the distributional difference between the model ’ s outputs and the outputs of the model ’ s quantized version . This loss function has the effect of guiding the optimization process towards a model which is both accurate , by virtue of minimizing the first loss component , and which is also similar enough to its quantized version due to minimization of the second component to ensure that the quantized model is also accurate . In addition , because the first component of the loss function deals only with floating point values , it provides accurate gradient information during optimization , in contrast to STE-based optimization which uses biased gradients at rounded points , which further improves the accuracy of the quantized model . Since GQ-Net directly produces a quantized model which does not require further fine-tuning , the number of epochs required to train GQ-Net is substantially less than the total number of epochs needed to train and fine-tune a model using the traditional quantization approach , leading to significantly reduced wall-clock training time . We note that GQ-Net ’ s technique is independent of and can be used in conjunction with other techniques for improving quantization accuracy , as we demonstrate in Section 4.3 . Finally , we believe that the guided training technique we propose can also be applied to other neural network structural optimization problems such as network pruning . We implemented GQ-Net in PyTorch and our codebase and trained models are publicly available 1 . We validated GQ-Net on the ImageNet classification task with the widely used ResNet-18 and 1An anonymous codebase has been submitted to OpenReview . The GitHub repository will be made public after the review process . compact MobileNet-v1/v2 models , and also performed a thorough set of ablation experiments to study different aspects of our technique . In terms of quantization accuracy loss compared to reference floating point models , GQ-Net strictly quantized using 4-bit weights and activations surpasses existing state-of-the-art strict methods by up to 2.7× , and also improves upon these methods even when they use higher bitwidths . In particular , 4-bit GQ-Net applied to ResNet-18 achieves 66.68 % top-1 accuracy , compared to 61.52 % accuracy in Louizos et al . ( 2019 ) and a reference floating point accuracy of 69.76 % , while on MobileNet-v2 GQ-Net achieves 66.15 % top-1 accuracy compared to 62.00 % accuracy in Krishnamoorthi ( 2018 ) and a reference floating point accuracy of 71.88 % . Additionally , GQ-Net achieves these results using layer-wise quantization , as opposed to channel-wise quantization in Krishnamoorthi ( 2018 ) , which further enhances the efficiency and practicality of the technique . 2 RELATED WORKS . Neural network quantization has been the subject of extensive investigation in recent years . Quantization can be applied to different part of neural networks , including weights , activations or gradients . Courbariaux et al . ( 2015 ) , Hou et al . ( 2016 ) , Zhou et al . ( 2017 ) , Hou & Kwok ( 2018 ) and other works quantized model weights to binary , ternary or multi-bit integers to reduce model size . Wei et al . ( 2018 ) quantized activations of object detection models for knowledge transfer . Alistarh et al . ( 2016 ) , Hou et al . ( 2019 ) quantized model gradients to accelerate distributed training . Another line of work quantizes both weights and activations to accelerate model inference by utilizing fix-point or integer arithmetic . These works include Courbariaux et al . ( 2016 ) , Rastegari et al . ( 2016 ) , Gysel et al . ( 2016 ) , Krishnamoorthi ( 2018 ) , Choi et al . ( 2018 ) , Zhang et al . ( 2018 ) , Jung et al . ( 2018 ) . A large set of methods have been proposed to improve training or fine-tuning for network quantization . Straight through estimators ( Bengio et al. , 2013 ) ( STE ) propagate gradients through non-differentiable operations with the identity mapping . Other training methods “ soften ” nondifferentiable operations to similar differentiable ones in order for gradients to pass through , then gradually anneal to piecewise continuous functions by applying stronger constraints . This line of works include Louizos et al . ( 2019 ) , Gong et al . ( 2019 ) , Bai et al . ( 2018 ) . Some works regard quantization as a stochastic process that produces parameterized discrete distributions , and guides training using gradients with respect to these parameters Soudry et al . ( 2014 ) , Shayer et al . ( 2018 ) . Another line of works does not require fine tuning , and instead re-calibrates or modifies the original network to recover accuracy using little or even no data He & Cheng ( 2018 ) , Nagel et al . ( 2019 ) , Meller et al . ( 2019 ) . Several recent works have focused on quantizing all parts of a network , typically in order to support deployment using only integer arithmetic units and avoiding the cost and complexity of additional floating point units . Gysel et al . ( 2016 ) proposed performing network inference using dynamic fixed-point arithmetic , where bitwidths for the integer and mantissa parts are determined based on a model ’ s weight distribution . Jacob et al . ( 2018 ) ; Krishnamoorthi ( 2018 ) proposed the quantization training and deployment algorithm behind the Tensorflow-Lite quantization runtime , which generates strictly quantized networks that can be easily implemented in hardware . Louizos et al . ( 2019 ) proposed a training method for strictly quantized models based on annealing a smooth quantization function to a piecewise continuous one . There has also been recent work on using parameterized quantizers which are optimized during quantization training . Choi et al . ( 2018 ) introduced learnable upper bounds to control the range of quantization . Zhang et al . ( 2018 ) proposed quantizers with a learnable basis which an be executed using fixed-point arithmetic . Jung et al . ( 2018 ) proposed to optimize weight scaling and quantization ranges jointly from task losses . 3 GQ-NET . In this section we describe the architecture of our proposed GQ-Net and then discuss components of the architecture which can be tuned to improve performance . 3.1 GQ-NET ARCHITECTURE . The major components of GQ-Net include the following , and are illustrated in Figure 1 : 1 . An L-layer neural network hW ( · ) with all computations performed using full precision floating point arithmetic . Here W = { W1 , . . . , WL } denotes the parameter ( weights ) of the model , with Wi , i ∈ 1 . . . L being the weights in layer i and expressed in floating point . 2 . The quantized model ĥW , Q ( · ) built from hW ( · ) . Here Q = { Qw1 , . . . , QwL , Qa0 , . . . , QaL } is a set of quantizers , i.e . mappings from floating point to ( scaled ) integer values ; the quantizers may be parameterized , and we describe how to optimize these parameters in Section 3.2 . Qwi quantizes weights Wi and Q a i quantizes activations in layer i . Let x0 denote an input to hW . To construct output ĥW , Q ( x0 ) of the quantized network , we proceed layer by layer . We first quantize the weights in layers i = 1 , . . . , L as ŵi = Qwi ( wi ) , and also quantize the input by setting x̂0 = Q a 0 ( x0 ) . we compute the quantized activations x̂i in layer i iteratively for i = 1 , . . . , L using x̂i = Qai ( x̃i ) , where x̃i = gi ( ŵi ∗ x̂i−1 ) , and gi ( · ) denotes the nonlinearity function in layer i and ∗ denotes convolution . Note that since ŵi and x̂i−1 are quantized , x̃i can be computed using integer or fixed point arithmetic . 3 . Next , we construct a loss function L incorporating both the training loss Lf of the full precision model hW and a loss Lq capturing the difference between hW and the quantized model ĥW , Q . L = ωfLf + ωqLq ( 1 ) Here ωf , ωq ∈ R are parameters capturing the relative importance of training loss versus distributional loss . In this paper , we focus on image classification networks , and thus we set Lf to be the cross-entropy loss between outputs from hW and the training labels . In addition , we set Lq = DKL ( σ ( hW ( · ) ) ||σ ( ĥW , Q ( · ) ) ) , where σ denotes the softmax function , to be the KL divergence between distributions σ ( hW ) and σ ( ĥW , Q ) on each input . Hence , minimizing the second term in L corresponds to pushing the floating point and quantized models to behave similarly to each other . Since the weight parameters W appear in both terms in L , the two terms can give conflicting signals for updating W during the optimization of L , causing the optimization to be unstable . We discuss how to deal with this problem in Section 3.2 . To train GQ-Net , we successively take mini-batches of training samples and labels and use them to compute L during the forward pass and propagate gradients with respect to W and the parameters of Q during the backward pass in order to minimize L. After L has converged sufficiently , we take the quantized weights in ĥW , Q ( · ) as the quantized model .

Summary: This work introduces GQ-Net, a novel technique that trains quantization friendly networks that facilitate for 4 bit weights and activations. This is achieved by introducing a loss function that consists of a linear combination of two components: one that aims to minimize the error of the network on the training labels of the dataset and one that aims to minimize the discrepancy of the model output with respect to the output of the model when the weights and activations are quantized. The authors argue that this has the effect of “guiding” the optimization procedure in finding networks that can be quantized without loss of performance. For the discrepancy metric the authors use the KL divergence from the predictive distribution of the floating point model to the one of the quantized model. The authors then propose several extra techniques that boost the performance of their method: 1. scheduling the weighting coefficients of the two loss terms (something which reminisces iterative pruning methods), 2. stopping the gradient of the floating point model w.r.t. the second loss term, 3. learning the parameters of the uniform quantizer, 4. alternating optimization between the weights and the parameters of the quantizers and 5. using separate batch normalization statistics for the floating point and quantized models. The authors then evaluate their method on Imagenet classification using ResNet-18 and Mobilenet v1 / v2, while also performing an ablation study about the extra tricks that they propose.