End-to-end Quantized Training via Log-Barrier Extensions

1 INTRODUCTION . As state-of-the-art deep learning models for vision , language understanding and speech grow increasingly large and computationally burdensome ( He et al. , 2017 ; Devlin et al. , 2018 ; Karita et al. , 2019 ) , there is increasing antithetical demand , motivated by latency , security and privacy concerns , to perform training and inference in these models on smaller devices at the edge rather than in server farms in the cloud . Model quantization has emerged as a promising approach to enable deployment of deep learning models on edge devices that reduce energy , latency and storage requirements by performing floating-point computation in low precision ( less than 32 bits ) . There are two primary strategies for quantization : Post-training approaches quantize the parameters of a model trained in full precision post-hoc , and tend to suffer a heavy penalty on accuracy since their inference graph differs substantially from training ( Jacob et al. , 2018 ) . Quantization-aware training ( QAT ) ( Bhuwalka et al. , 2020 ) combats this discrepancy by simulating quantization during training , so that model parameters are learned that will work well when inference is performed in low precision . In this work , we focus on the latter setting , suitable for fully quantized training on low-precision ( e.g . 8-bit ) devices . Though QAT results in quantized models that perform largely on par with their non-quantized counterparts , current state-of-the-art QAT methods ( Wu et al. , 2018 ; Wang et al. , 2018 ; Bhuwalka et al. , 2020 ) are not suitable for training on fully low-precision hardware because they employ fake quantization , meaning each operation is executed using 32- or 16-bit floating point arithmetic , and its output is quantized to lower precision , e.g . int8 . This results in two key incompatibilities with fully low-precision training , and consequently deployment on real low-precision hardware . First , existing QAT approaches assume perfect sums in inner product operations , which means that the accumulators used to compute matrix multiplies ( the acc row in Table 1 ) must be higher precision than the values being multiplied ( other bit-precision rows in Table 1 ) . This is to avoid losing res- olution in low-precision additions , also known as swamping ( Wang et al. , 2018 ) 1 . Second , QAT commonly leverages dynamic quantization ranges per-layer , meaning the mapping between highand low-precision values varies by layer , carefully tuned as a function of the network architecture , optimization dynamics and data during training . While this practice results in higher quantized inference accuracy , it is also a challenge to low-precision training , since it is unclear how to tune those ranges when training on new data in the absence of high-precision arithmetic . These incompatibilities present a substantial hurdle to quantized training in practice . For example , an automotive electronics manufacturer may want to deploy a machine learning model on its 8-bit door lock or power window controller to adaptively fit the users ’ habits . In this scenario , existing approaches for quantized training would fail ( Sakr et al. , 2019 ) . In response , we propose a new approach for fully quantized training of neural network models , inspired by the barrier method from convex optimization ( Boyd & Vandenberghe , 2004 ) . Log Barrier Tail-bounded Quantization ( LogBTQ ) utilizes a log barrier extension loss ( Kervadec et al. , 2019 ) to constrain the output of the network , encouraging all model parameters and activations to stay within the same predefined range . The log barrier function itself is a smooth approximation of the indicator function , which is ideal for selecting the weights that are within the range of quantization ( see Figure 1 , left ) . By fixing a single quantization range throughout the network at the beginning of training , our approach both obviates the need for dynamic ranges , and the limits of the range are set so as to alleviate overflow2 in matrix multiply accumulations . We combine the log barrier extension loss with an L1 regularization term ( Hoffer et al. , 2018 ) to further reduce the total magnitude of parameters and activations in the model . To allow for gradients , which tend form a peaky distribution near extremely small values ( Zhou et al. , 2016 ; Jain et al. , 2020 ) , to be quantized using the same range as the rest of the network , we also adopt the nonlinear µ-law algorithm from audio applications ( Deng & Doroslovacki , 2006 ) to construct a new MU8 codebook that better deals with “ swamping ” issues compared to the standard IEEE Float Standard . Experiments show that our approach achieves competitive results compared to state-of-art full-precision models on the MNIST , CIFAR-10 and ImageNet classification benchmarks , despite our models being trained end-to-end using only 8 bits of precision . 1Swamping : Accumulation of floating-point numbers , where the small magnitude value is ignored ( or truncated ) when it is added to the large magnitude sum . 2Overflowing : for the fixed-point accumulation where the accumulated value wraps around to the small value when it exceeds the largest value representable by the given accumulation precision . 2 BACKGROUND AND RELATED WORK . 2.1 POST-TRAINING QUANTIZATION . There was been a recent surge of interest in quantization research . In 2020 alone , there were a number of important developments in post-training quantization . Rusci et al . ( 2020 ) ; Jain et al . ( 2020 ) ; Esser et al . ( 2020 ) ; Uhlich et al . ( 2020 ) proposed learning-based approaches for determining the quantization ranges of activation and weights at low precision . Stock et al . ( 2020 ) advocates preserving the quality of the reconstruction of the network outputs rather than its weights . They all show excellent performance compared to full-precision models after quantization . Sakr & Shanbhag ( 2019 ) presented a detailed analysis of reduced precision training for a feedforward network that accounts for both the forward and backward passes , demonstrating that precision can be greatly reduced throughout the network computations while largely preserving training quality . Our work share the same intuition of preferring small predetermined dynamic range ( PDR ) and small clipping rate3 . However , Sakr & Shanbhag ( 2019 ) ’ s approach requires the network first be trained to convergence at full 32 bit precision , which is a significant limitation . In this paper , we focus on training rather than inference on low-precision hardware , therefore , we do not assume access to a full-precision high-performing model as a starting point . 2.2 QUANTIZATION-AWARE TRAINING . Pioneering works in this domain ( Zhou et al. , 2016 ; Courbariaux et al. , 2015 ) looked at quantizing model weights , activations , gradients to lower precision to accelerate neural network training . The terminology quantization aware training ( QAT ) was first introduced by Jacob et al . ( 2018 ) . QAT incoporates quantization error as noise during training and as part of the overall loss , which the optimization algorithm tries to minimize . Hence , the model learns parameters that are more robust to quantization , but QAT is not meant to be performed entirely in low precision , it aims to learn parameters that will work well for low-precision inference . More recently , several works further pursued the goal of enabling fully low-precision training ( Wu et al. , 2018 ; Wang et al. , 2018 ; Das et al. , 2018 ; Sun et al. , 2019 ) . As shown in Table 1 , most existing work employs fake quantization , resorting to higher precision values to compensate for the swamping issue , especially during gradient accumulation . Mixed-precision quantization ( Das et al. , 2018 ; Wang et al. , 2018 ; Zhang et al. , 2020a ) , which quantizes a neural network using multiple bit precisions across layers , still relies on higher-precision gradients to preserve model accuracy . This means it is difficult , if not impossible , to implement these approaches on low-bit ( e.g . 8-bit ) hardware . Most similar to our work , Sun et al . ( 2019 ) claim it is possible to do every step in low precision , but the quantization range for the layers in their work is very carefully chosen empirically , which presents great difficulty if we were to train models from scratch on low-precision hardware . Their method also requires a copy of the quantization error ( residual ) in FP16 ( 1-6-9 ) ( hence 8/16 in Ta- 3see Appendix B of Sakr & Shanbhag ( 2019 ) explaining PDR ; refer to their Criterion 2 about clipping rate . ble 1 ) . In addition to the 9-bit mantissa , the exponent bit in their floating point format would need to be manually modified to store the residual due to its small value . In this paper , we propose a new quantization scheme : log-barrier tail-bounded quantization ( LogBTQ ) that can perform fully end-to-end low precision training , suitable for deployment on lowprecision hardware . Our major contributions are the following : 1 . We apply a log barrier extension loss to soft-threshold the values of network weights and activations to constrain all the values to be small . Our quantization scheme also enables global fixed-range quantization which together significantly alleviates the overflow issue caused by large numbers and dynamic range . 2 . We add an L1 loss term to encourage sparsity and further reduce overflow . 3 . We propose µ-law quantization ( MU8 ) instead of INT8 , FP8 ( 1-4-3 ) or FP8 ( 1-5-2 ) to construct a more accurate codebook that better compensates for the peaky concentration of network parameters around small values . 3 LOG BARRIER TAIL-BOUNDED QUANTIZATION ( LOGBTQ ) . The overall diagram of our quantization scheme is shown in Figure 2 ( left ) . Figure 2 ( right ) shows the backward pass , where we quantize everything at each layer and all operations including input x , weights w , activations a , errors e , and gradients g ( including the gradient accumulation step ) . We denote all these values as the set Z = { x , w , a , e , g } . In this work , different from previous works ( Sakr & Shanbhag , 2019 ; Zhang et al. , 2020b ) that used adaptive quantization range , we adopt a globally fixed quantization range for every element z ∈ Z , and set z ∈ [ −2 , 2 ] . We do not need to adjust the range and precision during training as in other quantization work that relies on layer-wise dynamic ranges . This would greatly reduce the overhead for implementation on hardware . 3.1 CONSTRAINED FORMULATION . Let D = { I1 , ... IN } denote the labeled set of N training images , and f denote the neural network model , θ here denotes all the parameters of the neural network including weights w. For task , such as image classification , we are usually solving such an optimization problem : min θ L ( fθ ( I ) ) where L is the loss function of our neural network training objective . In this work , we use the typical cross-entropy loss , and since we are interested in constraining the quantization threshold , we are effectively performing constrained optimization in such a form : minimize θ L ( fθ ( I ) ) subject to |θn| ≤ u , n = 1 , . . . , N. ( 1 ) With u our desired barrier ( perturbation ) . In practice , we set u = 0.1 to ensure we can represent as much information as possible within our quantization range ( Figure 1 , left ) . This setting is also explained further in Section 3.3 .

Summary: The paper is addressing an important and challenging problem of end-to-end training of deep nets in fixed-point, in this case, with 8-bit precision. A good solution to this problem can have a major impact on the deployability of deep nets on embedded hardware. The basic idea is to introduce an additional term (the log-barrier constraint) in the loss function to constrain the allowable range over which model parameters are allowed to take values. The authors use of mu-encoding to assign non-uniform quantization levels to minimize the quantization error. The main results are in Table 2 showing that the method eliminates overflow in practice and allows quantized networks to approach the accuracy of full-precision networks on the MNIST, CIFAR-10 and ImageNet.