DNN Quantization with Attention

1 Introduction and Related Work . In the last decade , Deep Neural Networks ( DNNs ) in general and Convolutional Neural Networks ( CNNs ) in particular became state-of-the-art in many computer vision tasks , such as image classification or segmentation , object detection and face recognition ( LeCun et al . ( 1998 ) ; Iandola et al . ( 2016 ) ; Simonyan & Zisserman ( 2014 ) ; Graham ( 2014 ) ; Szegedy et al . ( 2015 ) ) . However , to be the state-of-the-art , DNNs often contain a large number of trainable parameters and require considerable computational power . Therefore , due to their large power and memory consumption , implementing DNNs on embedded systems with limited resources can be a real challenge . To alleviate this problem , a large number of different network compression methods that reduce the resource requirements of DNNs has been proposed in the past . Among them are for example pruning , distillation or quantization methods . Pruning methods have been first introduced by LeCun et al . ( 1990 ) . They identify and remove the most insignificant DNN parameters and yield networks with a reduced memory footprint and a smaller computational complexity . As reported by Yamamoto & Maeno ( 2018 ) ; Ramakrishnan et al . ( 2020 ) ; He et al . ( 2020 ) , pruning methods can either remove single DNN parameters , intermediate inputs or even whole network layers that are , according to a specific criteria , irrelevant for a good network performance . Another line of work is distillation . As introduced by Hinton et al . ( 2015 ) , it aims at training a small student DNN to reproduce the output of a bigger teacher network . While distillation methods initially only matched the final outputs of the teacher and student networks , methods evolved to take into account intermediate representations ( Romero et al . ( 2014 ) ; Koratana et al . ( 2018 ) ) . Our work is about DNN quantization , where weights and activations are represented with a smaller number of bits n ăă 32 . Quantization reduces the memory footprint of DNNs because the number of bits that are required to store their parameters is reduced . However , it also reduces their computational complexity , because low bit operations can be implemented more efficiently on dedicated low precision hardware ( Merolla et al . ( 2014 ) ; Farabet et al . ( 2011 ) ; Cowan et al . ( 2020 ) ; Han et al . ( 2020 ) ; Hacene et al . ( 2018 ) ) . Following Gupta et al . ( 2015 ) , many works have experimentally demonstrated that neural networks do not lose a lot of performance when their parameters are restricted to a small set of possible values . For instance , Choi et al . ( 2018 ) introduced PArameterized Clipping acTivation ( PACT ) in combination with Statistics-Aware Weight Binning ( SAWB ) –a method that aims at uniformly quantizing both weights and activations to n bit . Learned Step Size Quantization ( LSQ ) presented by Esser et al . ( 2019 ) is a quantization method that learns the quantization steps during training . Unlike other methods , it scales the gradient during backpropagation to speed up the training . Gradient scaling is important especially at transition points . Nikolić et al . ( 2020 ) proposed Bit-Pruning , a quantization method to learn the number of bits that each layer requires to represent its parameters and activations . In the same vein , Differentiable Quantization of Deep Neural Networks proposed by Uhlich et al . ( 2019 ) ( DQDNN ) tries to combine the features of both LSQ and Bit-Pruning . They propose a quantization technique where both the number of bits and the quantization steps are learned simultaneous . Other more aggressive quantization methods proposed to use low-bit precision down to binarization ( resp . ternarization ) with only two ( resp . three ) possible values and one ( resp . two ) bit storage for each parameter and/or activation ( Hubara et al . ( 2016 ) ; Courbariaux et al . ( 2015 ) ; Li et al . ( 2016b ) ; Zhu et al . ( 2016 ) ; Li et al . ( 2016a ) ) . Zhou et al . ( 2017 ) observed that training quantized networks to low precision benefits from incremental training . Rather than quantizing all the weights at once , they are quantizing them incrementally in groups , with some training iterations between each quantization step . In practice , 50 % of the weights are quantized in the first step , then 75 % , 87.5 % and finally 100 % . Another method for incremental quantization is Binary-Relax ( BR ) ( Yin et al . ( 2018 ) ) . Rather than splitting the parameters into groups , it uses a weighted linear combination of quantized and full-precision parameters and adapts a strategy to push the weights towards the quantized state , by gradually increasing the scaling factor corresponding to the quantized parameters . However , their strategy how to move from full-precision to quantized parameters is handcrafted and may not be optimal . In comparison to most of the previously mentioned works , we do not propose an improved quantization method , but a way how to train a quantized DNN with any existing quantizer . We rely on the fact that the DNN performance can be increased if the network has the ability to learn other features in addition to its own parameters . In the context of DNN quantization , this has already been discussed by Elsken et al . ( 2019 ) ; Ramakrishnan et al . ( 2020 ) ; Hacene et al . ( 2019 ) ; Uhlich et al . ( 2019 ) ; Nikolić et al . ( 2020 ) . In this contribution , we introduce DNN Quantization with Attention ( DQA ) , an attention mechanism-based learnable approach for DNN quantization ( Vaswani et al . ( 2017 ) ) . As shown in Figure 1 , rather than quantizing with just a single bitwidth DQA performs quantization with multiple different bitwidths in parallel . At each stage of the training , DQA can select the quantizer with the optimal bitwidth by putting more attention weight on it . In particular , by starting with a uniform initialization of the attention weights at the beginning of the training , DQA can move smoothly from high precision to low precision quantization . We demonstrate in this paper that this learnable approach results in better DNN accuracy for the exact same complexity and number of bits . The outline of the paper is as follows . In Section 3 we introduce the proposed method . Section 4 presents experiments results and compares our method with other state-of-the-art approaches on challenging computer vision datasets . Finally , we conclude in Section 5 . 2 Background . In the following , we explain how differentiable quantization with attention ( DQA ) can be used to train quantized DNNs . For simplicity reasons , we only discuss how to train DNNs with quantized weights . However , DQA is general and can also be used to train DNNs with quantized weights and activations . In particular , we will also provide experimental results for this case in Section 4 . Training DNNs with quantized weights is challenging . Especially when considering low bitwidths n ăă 32 , quantization can cause a severe accuracy degradation if compared to the full precision networks . This is mainly caused by the capacity reduction and by additional optimization issues that go hand in hand with the quantization . In particular , quantization yields non-smooth loss functions with gradients that are zero almost everywhere . As discussed by Uhlich et al . ( 2019 ) , that effectively stops gradient backpropagation and therefore harms training . For these reasons , quantized DNNs are usually not trained with standard gradient based training procedures , but require some tricks that allow for gradient backpropagation and that stabilize the training . The most commonly used trick is to apply straight through gradient estimates ( STE ) and to ignore the quantization during backpropagation . STE yield non-zero gradients that are suited for DNN training . However , at the same time they introduce a mismatch between the forward ( FW ) and backward ( BW ) pass , what often causes training instabilities and oscillations . More specifically , mismatch means that the gradients are calculated at the position of the quantized parameters , assuming the original float32 loss surface . Note , that because of this FW/BW mismatch , training with STE does not necessarily converge to the optimum . This mismatch is for example problematic , if the gradient of the cost function changes signs within one quantization step . For this case , gradient descent with STE would start to oscillate near the quantization thresholds and would not converge to the optimum , even if the cost function is convex . Of course , this problem is most pronounced for low bitwidhts . A training procedure that can alleviate this problem to some degree is Binary-Relax ( BR ) . As proposed by Yin et al . ( 2018 ) , BR does not only apply STE , but also uses a linear combination of the quantized and the full-precision parameters . This effectively reduces the FW/BW mismatch at the beginning , while still enabling the DNN to use the exact low-bit quantized loss at the end of the training . 3 Methodology . DQA builds on a similar approach . As shown in figure 1 , let fpx ; Qpw ; nqq be the transfer function of a quantized DNN layer , where x P RD , w P RM and Qpw ; nq are the layer input , the full-precision weights and the quantized weights , respectively . Similar to the idea of Binary-Relax ( BR ) , DQA relaxes the quantization problem and combines different quantization schemes during training . More specifically , instead of using just one single Qpw ; nq , we propose to train a quantized DNN with a set of K different quantization functions that are combined linearly during training as follows : y “ fpx ; qq ( 1 ) q “ QTa ( 2 ) Q “ » — — — – Q1pw ; n1qT Q2pw ; n2qT ... QKpw ; nKqT fi ffi ffi ffi fl , ( 3 ) where q is the weighted sum of K quantized weight vectors , Q P RKˆM is a matrix whose row vectors are the quantized weight vectors and a P r0 , 1sK is the attention vector on the quantization functions . Note , that each row of Q is calculated , using a different quantization function Qkpw ; nkq and bitwidth nk P N. In particular , we assume that the quantization functions in Q are sorted by the bitwidth , i.e. , n1 ă n2 ă ... ă nK . In general , DQA is agnostic to the choice of the actual quantization method and can be used with any existing method like min-max , SAWB , binary or ternary quantization . In the following section , we review and define popular quantization methods that we used in our experiments . The attention a is computed from a soft attention vector α P RK , using a softmax function with temperature T P R ` , i.e. , a “ e α T řK k “ 1 e αk T , P RK . ( 4 ) In particular , a reflects the importance of the K quantization methods Qk . During training , the soft attention α is treated as a trainable parameter that is optimized in parallel to the weights w. Note that , increasing αk will also increase the corresponding attention weight ak and therefore the importance of Qkpw ; nkq . In this manner , the quantized DNN can learn which bitwidth should be used at which stage , during the training . DQA exponentially cools down the temperature T T pbq “ T p0qΨb . ( 5 ) Here , b “ 1 , 2 , ... , B is the batch index for batch-wise training , T p0q P R ` is the initial temperature and Ψ P r0 , 1r is the decay rate . Because of that schedule , DQA progressively moves from the full mixture of quantization functions at the beginning of the training to just one single quantization function at the end of training . Note , that BR can be seen as a special case of DQA , where q “ rQpw ; 2q , Qpw ; 32qsT , i.e . for the case that we only use two quantizers with n “ 2bit and n “ 32bit , and for the case that we use a fixed schedule to change the attention vector a . However , DQA has two advantages : 1 ) The way how we change a and move from high to low precision quantization is learned and data dependent . Hence , DQA can choose the optimal mixture of the quantizers at each training iteration . 2 ) As shown in Figure 2 , DQA gives a smoother transition from high to low-precision parameters . Here , we plot the absolute quantization error for a fixed temperature based schedule for a . In particular , we choose α “ r3 { 4 , 1 { 4sT and α “ r4 { 7 , 2 { 4 , 1 { 7sT for BR and DQA with 2 , 4 , 8bit quantization , respectively . Then , we start with a large temperature T Ñ 8 , for which we effectively take the average 1K řK k “ 1Qkpw ; nkq , and move towards 2bit quantization for T Ñ 0 . Note , that for the whole interval that we consider for T , DQA results in a lower quantization error , meaning that it also yields a smaller FW/BW mismatch . In general , training quantized DNNs with such a mixture of different weight quantizations and decaying T will not necessarily result in a quantized DNN that uses a low bitwidth . To enforce a low-bit quantized DNN , we therefore add to the loss function a separate regularizer for each layer rpαq “ λg Tapαq S , ( 6 ) where S is the number of weights in the whole network . Note , that the normalization by S makes the regularizer , and therefore the choice of λ , independent of the actual network size . g “ rg1 , g2 , ¨ ¨ ¨ , gKsT is a penalty vector , where gk is increasing with growing k. Because we assume , that the quantization functions Qkpw ; nkq are sorted by the bitwidth , i.e. , n1 ă n2 ă ... ă nK adding gTapαq penalizes if large attention values are assigned to quantizers with a large bitwidth . Hence , it helps the method to converge to the lowest-bit quantization . Algorithm 1 summarizes the DQA training . To quantize a given value x , we use min-max , Statistics-Aware Weight Binning SAWB , Binary Weight Network ( BWN ) or Ternary Weight Network ( TWN ) as defined by Nikolić et al . ( 2020 ) ; Choi et al . ( 2018 ) ; Rastegari et al . ( 2016 ) ; Li et al . ( 2016a ) respectively , and detailed in appendix A.1 Algorithm 1 DQA algorithm for a single network layer Inputs : Input vector x , initial softmax temperature T p0q , final softmax temperature T pBq , number of training iterations B , and layer transfer function f Output : Output tensor y ψ “ e log ˆ T pBq T p0q ˙ B ă 1 for each b “ 1 , 2 , ... , B do T pbq Ð T p0qψb α Ð αstdpαq a Ð softmaxpα { T pbqq q “ QTa ( linear combination ) y “ fpx , qq Update w and α via backpropagation . end for

Summary: This work presents a training method for low-bit network quantization. While training, it employs a multi-bitwidth paradigm in order to alleviate the nonsmooth optimization landscape with lower bitwidth. It uses a temperature parameter and a penalty term to force the network to gradually converge to the target low bit. Experiments are conducted on CIFAR10, CIFAR100, ImageNet Classification with ResNet 18 and MobilenetV2.