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: The papers addresses the problem of compression of neural networks. The paper builds upon binary-relax prior work. The precision is adapted during training with a mixture-based quantization method through temperature cooling and a set of an ``attention’’ vector a. The idea of the method, called DQA, is to progressively moves from a mixture of quantization functions (mixing high with low precision, for instance 32 bit with 2 bit training), to a single one (low precision) towards the end of the training. The paper states that the method can be used with several types of quantization methods. The evaluation is carried on computer vision architectures (ResNet18, MobileNet) for image recognition tasks (Cifar-10, Cifar-100, Imagenet ILSVRC 2012).