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: In this paper, the authors propose a framework towards 4-bit auantization of CNNs. Specifically, during training, the proposed method contains a full precision branch supervised by classification loss for accurate prediction and representation learning, as well as a parameterized quantization branch to approximate the full precision branch. A quantization loss between the full precision branch and the quantization branch is defined to minimize the difference between activation distributions. The authors proposed a series of improvements, including alternative optimization, dynamic scheduling, detach and batch normalization to help boosting the performance to SOTA under 4-bit quantization.