WAVEQ: GRADIENT-BASED DEEP QUANTIZATION OF NEURAL NETWORKS THROUGH SINUSOIDAL REGULARIZATION

1 INTRODUCTION . Quantization , in general , and deep quantization ( below eight bits ) ( Krishnamoorthi , 2018 ) , in particular , aims to not only reduce the compute requirements of DNNs but also reduce their memory footprint ( Zhou et al. , 2016 ; Judd et al. , 2016b ; Hubara et al. , 2017 ; Mishra et al. , 2018 ; Sharma et al. , 2018 ) . Nevertheless , without specialized training algorithms , quantization can diminish the accuracy . As such , the practical utility of quantization hinges upon addressing two fundamental challenges : ( 1 ) discovering the appropriate bitwidth of quantization for each layer while considering the accuracy ; and ( 2 ) learning weights in the quantized domain for a given set of bitwidths . This paper formulates both of these challenges as a gradient-based joint optimization problem by introducing an additional novel sinusoidal regularization term in the training loss , called WaveQ . The following two main insights drive this work . ( 1 ) Sinusoidal functions ( sin2 ) have inherent periodic minima and by adjusting the period , the minima can be positioned on quantization levels corresponding to a bitwidth at per-layer granularity . ( 2 ) As such , sinusoidal period becomes a direct and continuous representation of the bitwidth . Therefore , WaveQ incorporates this continuous variable ( i.e. , period ) as a differentiable part of the training loss in the form of a regularizer . Hence , WaveQ is a differentiable regularization mechanism , it piggy backs on the stochastic gradient descent that trains the neural network to also learn the bitwidth ( the period ) . Simultaneously this parametric sinusoidal regularizer pushes the weights to the quantization levels ( sin2 minima ) . By adding our parametric sinusoidal regularizer to the original training objective function , our method automatically yields the bitwidths for each layer along with nearly quantized weights for those bitwidths . In fact , the original optimization procedure itself is harnessed for this purpose , which is enabled by the differentiability of the sinusoidal regularization term . As such , quantized training algorithms ( Zhou et al. , 2016 ; Mishra et al. , 2018 ) that still use some form of backpropagation ( Rumelhart et al. , 1986 ) can effectively utilize the proposed mechanism by modifying their loss . Moreover , the proposed technique is flexible as it enables heterogeneous quantization across the layers . The WaveQ regularization can also be applied for training a model from scratch , or for fine-tuning a pretrained model . In contrast to the prior inspiring works ( Uhlich et al. , 2019 ; Esser et al. , 2019 ) , WaveQ is the only technique that casts finding the bitwidths and the corresponding quantized weights as a simultaneous gradient-based optimization through sinusoidal regularization during the training process . We also prove a theoretical result to provide an insight on why the proposed approach leads to solutions preserving the original accuracy during quantization . We evaluate WaveQ using different bitwidth assignments across different DNNs ( AlexNet , CIFAR-10 , MobileNet , ResNet-18 , ResNet20 , SVHN , and VGG-11 ) . To show the versatility of WaveQ , it is used with two different quantized training algorithms , DoReFa ( Zhou et al. , 2016 ) and WRPN ( Mishra et al. , 2018 ) . Over all the bitwidth assignments , the proposed regularization , on average , improves the top-1 accuracy of DoReFa by 4.8 % . The reduction in the bitwidth , on average , leads to 77.5 % reduction in the energy consumed during the execution of these networks . Finally , we apply WaveQ to Transformer DNNs citeppby augmenting their loss with WaveQ parametric sinusiodal regularization . In this case , the conventional stochastic gradient descent plus WaveQ regularization is used to quantize the big Transformer model from ( Ott et al. , 2018 ) for machine translation on the IWSLT14 German-English dataset ( IWS ) . For 5 , 6 , and 7-bit quantization , training with WaveQ yields 0.46 , 0.14 , 0.04 improved BiLingual Evaluation Understudy ( BLEU ) score , respectively . As a point of reference , the original big Transformer model from ( Ott et al. , 2018 ) improves the BLEU by only 0.1 over the state-of-the-art . Code available at https : //github.com/waveq-reg/waveq 2 JOINT LEARNING OF LAYER BITWIDTHS AND QUANTIZED PARAMETERS . Our proposed method WaveQ exploits weight regularization in order to automatically quantize a neural network while training . To that end , Section 2.1 describes the role of regularization in neural networks and then Section 2.2 explains WaveQ in more details . 2.1 PRELIMINARIES . Quantizer . We discuss how quantization of weight works . Consider a floating-point variable wf to be mapped into a quantized domain using ( b + 1 ) bits . Let Q be a set of ( 2k + 1 ) quantized values , where k = 2b − 1 . Considering linear quantization , Q can be represented as { −1 , −k−1k , ... , − 1 k , 0 , 1 k , ... , k−1 k , 1 } , where 1k is the size of the quantization bin . Now , wf can be mapped to the b-bit quantization Zhou et al . ( 2016 ) space as follows . wqo = 2× quantizeb ( tanh ( wf ) 2max ( |tanh ( Wf ) | ) + 1 2 ) − 1 ( 2.1 ) In Equation 2.1 , quantizeb ( x ) = 1 2b−1 round ( ( 2 b−1 ) x ) , wf is a scalar , Wf is a vector , and wqo is a scalar in the range [ −1 , 1 ] . Then , a scaling factor c is determined per layer to map the final quantized weight wq into the range [ −c , +c ] . As such , wq takes the form cwqo , where c > 0 , and wqo ∈ Q . Soft constraints through regularization and the loss landscape of neural networks . Neural networks ’ loss landscapes are known to be highly non-convex and it has been empirically verified that loss surfaces for large networks have many local minima that essentially provide equivalent test errors Choromanska et al . ( 2015 ) ; Li et al . ( 2018 ) . This opens up the possibility of adding soft constrains as extra custom objectives during the training process , in addition to the original objective ( i.e. , to minimize the accuracy loss ) . The added constraint could be with the purpose of increasing gener- alization or imposing some preference on the weights values . Regularization in action . Regularization effectively constrains weight parameters by adding a corresponding term ( regularizer ) to the objective loss function . A classical example is Weight Decay Krogh & Hertz ( 1991 ) that aims to reduce the network complexity by limiting the growth of the weights . This soft constraint is realized by adding a term , proportional to the L2 norm of the weights to the objective function as the regularizer that penalizes large weight values . WaveQ , on the other hand , uses regularization to push the weights to the quantization levels . For the sake of simplicity and clarity , Figure 1 ( a ) and ( b ) illustrate a geometrical sketch for a hypothetical loss surface ( original objective function to be minimized ) and an extra regularization term in 2-D weight space , respectively . For weight decay regularization ( Figure 1 ( a ) ) , the faded circle shows that as we get closer to the origin , the regularization loss is minimized . The point wopt is the optimum just for the loss function alone and the overall optimum solution is achieved by striking a balance between the original loss term and the regularization loss term . Similarly , Figure 1 ( b ) shows a representation of the proposed periodic regularization for a fixed bitwidth β . A periodic pattern of minima pockets are seen surrounding the original optimum point . The objective of the optimization problem is to find the best solution that is the closest to one of those minima pockets where weight values are nearly matching the desired quantization levels , hence the name quantization-friendly . 2.2 WAVEQ REGULARIZATION The proposed regularizer is formulated in Equation ( 2.2 ) where the first term pushes the weights to the quantization levels and the second correlated term aims to reduce the bitwidth of each individual layer heterogeneously . R ( w ; β ) = λw ∑ i ∑ j sin2 ( πwij ( 2 βi − 1 ) ) 2βi︸ ︷︷ ︸ Weights quantization regularization + λβ ∑ i βi︸ ︷︷ ︸ Bitwidth regularization ( 2.2 ) In Equation ( 2.2 ) , λw is the weights quantization regularization strength which governs how strongly weight quantization errors are penalized , and λβ is the bitwidth regularization strength . The parameter βi is proportional to the quantization bitwidth which is elaborated later in this section . Figure 2 ( a ) shows a 3-D visualization of our regularizer , R. Figure 2 ( b ) , ( c ) show a 2-D profile w.r.t weights ( w ) , while ( d ) shows a 2-D profile w.r.t the bitwidth ( β ) . Periodic sinusoidal regularization . As shown in Equation equation 2.2 , the first regularization term is based on a periodic function ( sinusoidal ) that adds a smooth and differentiable term to the original objective , Figure 2 ( b ) . The periodic regularizer induces a periodic pattern of minima that correspond to the desired quantization levels . Such correspondence is achieved by matching the period to the quantization step ( 1/ ( 2βi − 1 ) ) based on a particular number of bits ( βi ) for a given layer i . Learning the sinusoidal period . The parameter βi in ( Equation 2.2 ) controls the period of the sinusoidal regularizer . Thereby βi is directly proportional to the actual quantization bitwidth ( bi ) of layer i as follows : bi = dβie , and αi = 2bi/2βi ( 2.3 ) In Equation ( B.2 ) where αi ∈ R+ is a scaling factor . Note that bi ∈ Z is the only discrete parameter , while βi ∈ R+ is a continuous real valued variable , and d∗e . is the ceiling operator . While the first term in Equation equation 2.2 is only responsible for promoting quantized weights , the second term ( λβ ∑ i βi ) aims to reduce the number of bits for each layer i individually while the overall loss is aiming to maximize accuracy . As such , this term is a soft constraint that yields heterogeneous bitwidths for different layers . The main insight here is that βi , which also controls the period of the sinusoidal term , is a continuous valued parameter by definition . As such , βi acts as an ideal optimization objective constraint and a proxy to minimize the actual quantization bitwidth bi . Therefore , WaveQ avoids the issues of gradient-based optimization for discrete valued parameters . Furthermore , the benefit of learning the sinusoidal period is two-fold . First , it provides a smooth differentiable objective for finding minimal bitwidths . Second , simultaneously learning the scaling factor ( αi ) associated with the found bitwidth . Leveraging the sinusoidal properties , WaveQ learns the following two quantization parameters simultaneously : ( 1 ) a per-layer quantization bitwidth ( bi ) along with ( 2 ) a scaling factor ( αi ) through learning the period of the sinusoidal function . Additionally , by exploiting the periodicity , differentiability , and the local convexity profile in sinusoidal functions , WaveQ automatically propels network weights towards values that are inherently closer to quantization levels according to the jointly learned quantizer ’ s parameters bi , αi as follows . Mapping this to the quantizer ’ s design , for ( b+ 1 ) bits quantization , k = 2b − 1 , and c = α . These learned parameters ( b , α ) , as explained above , can be mapped to the quantizer parameters explained in Equation equation 2.1 . For ( b+1 ) bits quantization ( the extra bit is the sign bit ) : k = 2b − 1 , and c = α = 2b/2β ( 2.4 ) Bounding the gradients . The denominator in the first term of Equation equation 2.2 ( ∑ i ∑ j sin2 ( πwij ( 2βi−1 ) ) 2βi ) is used to control the range derivatives of the proposed regularization term with respect to β . This denominator is chosen to limit vanishing and exploding gradients during training . To this end , we compared three variants of equation equation 2.2 with different normalization defined , for k = 0 , 1 , and 2 , as : Rk ( w ; β ) = λw ∑ i ∑ j sin2 ( πwij ( 2 βi − 1 ) ) 2kβi + λβ ∑ i βi ( 2.5 ) Figure 3 ( a ) , ( b ) , ( c ) provide a visualization on how each of the proposed scaled variants affect the first and second derivatives . For Rk=0 and Rk=2 , there are regions of vanishing or exploding gradients . Only the regularization Rk=1 ( the proposed one ) is free of such issues . Setting the regularization strengths . The convergence behavior depends on the setting of the regularization strengths λw and λβ . Our proposed objective seeks to learn multiple quantization parameterization in conjunction . As such , the learning process can be portrayed as three phases ( Figure 2 ( e ) ) . In Phase ( 1 ) , the learning process optimizes for the original task loss E0 . Initially , the small λw and λβ values allow the gradient descent to explore the optimization surface freely . As the training process moves forward , we transition to phase ( 2 ) where the larger λw and λβ gradually engage both the weights quantization regularization and the bitwidth regularization , respectively . Note that , for this to work , the strength of the weights quantization regularization λw should be higher than the strength of the bitwidth regularization λβ such that a bitwidth per layer could be properly evaluated and eventually learned during this phase . After the bitwidth regularizer converges to a bitwidth for each layer , we transition to phase ( 3 ) , where we fix the learned bitwidths and gradually decay λβ while we keep λw high . The criterion for choosing λw and λβ is to balance the magnitude of regularization loss to be smaller than the magnitude of accuracy loss . The mathematical formula used to generate λw and λβ profiles can be found in the supplementary material . ( Figure 8 ) .

Summary: When training quantized neural networks, one typically first fixes the desired bitwidth $b$ (of weights and activations). While training, one maintains and updates full-bitwidth weights during backprop and "cheats" by using $b$-quantized versions of these full-precision weights during forward propagation. The quantization scheme used may vary.