S$^3$: Sign-Sparse-Shift Reparametrization for Effective Training of Low-bit Shift Networks

1 Introduction . While deep neural networks ( DNNs ) have achieved widely success in various tasks , the training and inference of DNNs usually cost prohibitive resources due to the fact that they are often overparametrized and composed of computational costly multiplication operations for both forward and backward propagation [ 28 ] . To enable the application feasibility of DNNs in resource-constrained scenarios , substantial efforts have been made to reduce the computation complexity of neural networks while preserving their accuracy . Among different proposed approaches for this purpose , low-bit neural networks with binary weights [ 6 , 22 ] or ternary weights [ 15 ] are recently designed to replace the expensive operations like multiplication with cheaper ones , e.g. , replacing multiplication with sign flip operation during inference . We focus on low-bit shift neural networks [ 11 , 9 , 30 ] that replace multiplication with the bit-shift operation . Following the conventional linear algebra notations , we denote scalars by x , vectors in bold small letters x , and matrices by bold capital letters X . Traditional neural networks require computing the inner product w > x = ∑ i wixi , where w is the weight vector and x is the feature vector . In comparison , shift neural networks limit the values of the weight vectors to a set of discrete values wshift ∈ { 0 } ∪ { ±2p } . In this case , multiplication can be replaced with bit-shift operations , as multiplying wi = 2p is equavilent to shifting p places to the left ( if p is positive ) or to the right ( if p is negative ) . The bit-shift operation alone only covers discrete neural networks with positive weights , therefore , most shift neural networks further add a sign flip operation to the shift to allow wi = { ±2p } . Inference with bit-shift operations is highly hardware friendly and can save up to 196× energy cost on FPGA and 24× on ASIC compare to their multiplication counterparts [ 28 ] . Several works aim to further improve the memory efficiency of shift neural networks by reducing the bit-width of the weights . Zhou et al . [ 30 ] propose to fine-tune pre-trained full-precision weights 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . with a progressive power-of-two weight quantization strategy . They suggest to keep a portion of weights unquantized during fine-tuning to improve model performance . DeepShift [ 9 ] is the state-ofthe-art technique for training low-bit shift networks from random initialization . Initialization with pre-trained full-precision weights provides a significant performance gain for DeepShift . All latest techniques heavily rely on the initialization of full-precision pre-trained weights , which implies some optimization difficulties exist in the current training diagram of low-bit shift networks . However , although extremely memory efficient and hardware friendly , low-bit shift networks suffer from a significant accuracy drop on large datasets compared to their full-precision counterpart , and the performance of them is sensitive to weight initialization . To address these problems , in this work , we analyze the current training diagram of low-bit shift neural networks , and find that the accuracy degradation and the weight initialization sensitivity of the low-bit shift networks are caused by the design flaw of the weight quantizer during training . Specifically , training shift neural networks is different from training full-precision networks . The gradient-based optimization algorithms are designed for optimizing continuous variables , but the weights of shift neural networks are discrete values . To learn the discrete weights with a gradient-based approach , it is necessary to exploit weight re-parameterization . One of the most widely adopted discrete weight re-parameterization technique is using a quantizer function to map a continuous parameter w ∈ R to a discrete weight w ∈ { v1 , v2 , · · · , vk } with k possible values . As the weight quantizers are generally non-differentiable functions , a gradient approximator , such as straight-through estimator [ 3 ] , is utilized to approximate their derivatives during training . Our analysis shows the quantizer leads to a severe gradient vanishing and weight sign freezing . In this paper , we design an effective training strategy to combat the gradient vanishing and weight sign freezing problem in low-bit shift networks . We first impose a dense weight priori during training , which equals maximizing the ` 0 norm of the discrete weights . However , optimizing ` 0 norm during training is a combinatorial problem , instead we propose S3 re-parameterization and regularize the sparsity parameter . S3 is a new technique which re-parameterizes the discrete weights of shift networks in a sign-sparse-shift , a 3-dimensional augmented parameter space to disentangle the roles of quantization values , and hopefully train more effectively thanks to more orthogonal axes in the augmented space . Although our method introduces extra parameters during training , its memory occupation under low-bit is slightly higher than the full-precision network since the additional memory overhead depends on the weight size . In contrast , the activation size dominates the training memory occupation , especially using a large batch size . We evaluate our approach on the ImageNet dataset . While all previous methods require at least 5-bit weight representation to achieve the same performance as the full-precision neural networks on large datasets such as ImageNet , our experimental results show that our proposed method surpasses all previous methods , pushes this boundary further to 3-bits . Besides , our approach requires no complicated weight initialization or training strategy . Moreover , we define two indices of weight dynamics , named weight sign variation rate ( WSVR ) and weight low-value rate ( WLVR ) , and compare the trend of these two indices during the training process of different methods . Experiments show that the weight dynamics of the ternary weights trained with our proposed technique better align with the weight dynamics of full-precision weights than trained with a traditional quantizers . This desirable dynamics is caused by our re-parametrization , and is the key to efficient training . We hope that this opens the possibility for methods similar as S3 to be used in training low-bit networks . 2 Related Works . Different approaches have been proposed to reduce the computational complexity of neural networks . A type of the popular multiplication-free neural networks is low-bit neural networks with binary weights [ 6 , 22 ] or ternary weights [ 15 ] . However , they suffer from under-fitting on large datasets , leading to an accuracy gap compared with their full-precision counterparts . Other approaches include replacing multiplication operations with addition operations [ 5 , 25 , 26 ] or bit-shift operations [ 11 , 9 , 30 ] . Although these approaches achieve a low accuracy drop on large datasets , they require higher weight representation bit-width . Other follow-up works try to improve the accuracy of multiplicationfree neural networks by replacing multiplication with both addition and bit-shift operation [ 28 ] , sum of binary bases [ 17 , 29 ] , or sum of shift kernels [ 16 ] . However , their computational cost is high as they use multiple operations per kernel . Training shift neural networks with discrete weight values is a challenging task . Many weight reparameterization approaches have been proposed to solve the challenge . One way is using stochastic weight [ 6 , 18 , 23 ] , but these methods suffer from the slow computation of sampling . Another way is utilizing a quantizer function [ 6 , 22 , 15 ] to map or threshold continuous weights to discrete values . Since the thresholding function is non-differentiable , the backward gradients across the quantizer are approximated by a gradient approximator . The Straight-Through-Estimator ( STE ) [ 3 ] is a popular choice for quantizer gradient approximation . There are also works study the vanishing gradient problem ( VGP ) that appears frequently in deep neural network training . [ 1 , 10 ] point out that the zero values in full-precision weights lead to VGP when training highly pruned neural networks , which eventually jeopardizes the accuracy of the model . [ 8 , 30 ] suggest to maintain a portion of the full-precision weights during quantized network training for improving the performance of the model . [ 30 ] shows that applying progressive quantization on low and high value weights differently can further improve the performance of the trained model . Quantizing high-value weights first and keep more low-value weights remains in full-precision during training helps to achieve a better shift neural network compared with randomly applying progressive quantization strategy on weights regardless their values . This observation coincide with our analysis , and we argue that the performance gain of progressive quantization methods come from the gradient flow preserved by maintaining a portion of full-precision weight during training . However , the study of the negative impacts caused by weight quantizers are still missing . To the best of our knowledge , we are the first to study the VGP caused by weight quantizer design . 3 Training Low-bit Networks . We first introduce the general form of the weight quantizers for discrete weight training , and then explain how quantization function leads to training difficulties . In low-bit networks , including shift networks , the weights are discrete values while usually the gradient-based optimization algorithms are designed for optimizing continuous variables . Specifically , quantizer function is one of the most widely-used approach for training discrete weight neural networks . It maps a continuous weight w to one of the k discrete weight values Q ( w ) ∈ { v1 , v2 , · · · , vk } . Several quantizers of low-bit networks [ 15 , 31 , 27 ] and shift networks [ 9 , 30 ] are in the following form of staircase function : Q ( w ) = v1 w < t1 v2 t1 ≤ w < t2 ... vk tk ≤ w ( 1 ) However , training a quantized neural network with a quantizer function in the form of Equation ( 1 ) is challenging due to two issues , vanishing gradient ( VGP ) and weight sign freezing . We discuss the vanishing gradient issue in section 3.1 and the weight sign freezing issue in section 3.2 . 3.1 Gradient vanishing . To clarify further why quantization and VGP are entangled , suppose a ReLU-activated fully-connected layer can be generally formulated as : [ xl+10 , xl+11 ] = ReLU ( [ W l00 W l 01 W l10 W l 11 ] [ xl0 xl1 ] ) ( 2 ) The gradient update toward xl0 can be calculated as : ∂Loss ∂xl0 = ∂Loss ∂xl+10 ∂xl+10 ∂ ( W l00x l 0 +W l 01x l 1 ) W l00 + ∂Loss ∂xl+11 ∂xl+11 ∂ ( W l10x l 0 +W l 11x l 1 ) ︸ ︷︷ ︸ Derivative of ReLU W l10 ( 3 ) And the derivative of ReLU is , ∂xl+11 ∂ ( W l10x l 0 +W l 11x l 1 ) = { 1 W l10x l 0 +W l 11x l 1 > 0 0 W l10x l 0 +W l 11x l 1 ≤ 0 ( 4 ) By observing Equation ( 3 ) , we can see that the gradient information ∂Loss ∂xl+11 propagating from xl+11 to xl0 vanishes in two distinct situations : i ) the corresponding derivative of ReLU is zero . ( i.e . the neuron is not activated ) , and ii ) the corresponding weight W l10 is zero . Both scenarios lead to a gradient information loss , and ultimately cause inferior performance even after training convergence . The situation i ) is widely known as the Dying ReLU issue caused by the sparse derivative of ReLU . Several ReLU-variants [ 20 , 12 ] with non-sparse derivatives are developed as remedies . The VGP caused by the quantizer is similar to the dying ReLU issue . As the first step , we motivate S3 technique on a ternary network and then generalize it further for higher bit shift networks . The ternary quantizer function ( 1 ) contributes to situation ii ) by mapping a portion of non-zero full-precision weights w to the exact zero value wter = Q ( w ) . This leads to a smaller ` 0 norm for the discrete weight term wter compared to its full-precision counterpart w , ‖Q ( w ) ‖0 ≤ ‖w‖0 . Shift networks exchange the costly multiplication operation with cheap shift operation , so inherently it is a nonuniform quantized network and faces the same problem . The VGP caused by the quantizer is the main optimization difficulty in training shift neural networks as well . This issue appears more severely in low-bit networks which include fewer quantized range with higher sparsity rate .

Summary: The paper suggests a new training technique for neural networks, where each parameter is broken down into multiple parameters, representing its sign, sparsity, and magnitude (or shift). They explain why this new technique is useful for overcoming vanishing gradients and sign-freezing. They show that this training can produce a network with 3 bits weights, with the same quality as the standard fp32 network.