Bit-wise Training of Neural Network Weights

We propose an algorithm where the individual bits representing the weights of a neural network are learned . This method allows training weights with integer values on arbitrary bit-depths and naturally uncovers sparse networks , without additional constraints or regularization techniques . We show better results than the standard training technique with fully connected networks and similar performance as compared to standard training for residual networks . By training bits in a selective manner we found that the biggest contribution to achieving high accuracy is given by the first three most significant bits , while the rest provide an intrinsic regularization . As a consequence we show that more than 90 % of a network can be used to store arbitrary codes without affecting its accuracy . These codes can be random noise , binary files or even the weights of previously trained networks . 1 INTRODUCTION . Many challenging areas of computer science have found very good solutions by using powerful techniques such as deep neural networks . Their applications range now from computer vision , speech recognition , natural language processing , game playing engines , natural sciences such as physics , chemistry , biology and even to automated driving . Their success is largely due to the increase in computing power of dedicated hardware which supports massive parallel matrix operations . This enabled researchers to build ever growing models with intricate architectures and millions or even billions of parameters , with impressive results . However , despite their effectiveness , many aspects of deep neural networks are not well understood . One such aspect is why over-parameterized models are able to generalize well . One of the important avenues of research towards a better understanding of deep learning architectures is neural network sparsity . Frankle & Carbin ( 2019 ) showed a simple , yet very effective magnitude based pruning technique capable of training neural networks in very high sparsity regimes while retaining the performance of the dense counterparts . This sparked new interest in parameter pruning and a large body of work on the topic has since been published . The techniques for weight pruning can be broadly categorized as follows : pruning after training , before training and pruning during training . The work of Frankle & Carbin ( 2019 ) falls in the first category because the method relies on removing the weights which reach small magnitudes after they have been trained . In the second kind of approach , such as ( Lee et al. , 2019 ; Wang et al. , 2020 ) , neural networks are pruned before training in order to avoid expensive computations at training time . The end goal is to remove connections such that the resulting network is sparse and the weights are efficiently trainable after the pruning procedure . The third kind of approach is to use dynamical pruning strategies ( Dai et al. , 2019 ; Mostafa & Wang , 2019 ) which train and remove weights at the same time . The main goal behind these pruning strategies is to find sparse neural networks which can be trained to large degrees of accuracy . However , it has been shown by Zhou et al . ( 2019 ) that there exist pruning masks which can be applied to an untrained network such that its performance is far better than chance . Furthermore , Ramanujan et al . ( 2019 ) developed an algorithm for finding good pruning masks for networks with fixed , random weights . Theoretical works ( Malach et al. , 2020 ; Orseau et al. , 2020 ) even proved that within random neural networks there exist highly efficient subnetworks , which can be found just by pruning . Orseau et al . ( 2020 ) advance the hypothesis that the main task of gradient descent is to prune the networks while on the second place is the fine-tuning of the weights . 2 MOTIVATION . A key issue we want to emphasize is that , in all these works , the way in which the networks are pruned in practice is by forcing them , through some criteria , to set a fraction of the weights to zero . Since it has been shown that sparse networks perform as well as their dense counterpart , or sometimes even better , the natural question that arises is : why doesn ’ t gradient descent itself prune the weights during training ? Why hasn ’ t pruning been spontaneously observed in practice ? One possible explanation is that , at least for classification tasks , the usual cross–entropy loss without additional regularization techniques are not well suited for this . Other factors such as the stochasticity of the data batches , optimization algorithm , weights initialization etc . might also play a role . However , we approach this question from a different perspective . We hypothesize that an important reason for weights not being set to zero is because this is a particular state where the bits representing a weight must all equal zero . This is highly unlikely since weights are usually represented on 32 bits . The probability of a single weight being set to exactly zero is 2−31 , the sign bit not playing a role . Therefore the chances that a significant number of weights is set to zero decreases very rapidly . If weights would be represented on a lower bit depth , then the chance that the optimizer sets them to zero should increase . In order to test the degree to which this hypothesis is true we experiment with neural networks for image classification where , instead of training the weights themselves , we train the individual bits representing the weights . This might allow gradient descent to reach stable states where all bits in a set of weights are zero and the loss function is around a local minimum . If our hypothesis is true then we expect a strong dependency between sparsity and bit-depth . By encoding weights on arbitrary precision we also touch upon the topic of network quantization and show that particular cases of this training technique result in algorithms developed in previous works which we will describe in Section 8 . Moreover , we show that weight quantization naturally leads to weight pruning and sparse networks without additional constraints such a regularizations , additional loss terms , architectural changes or other tricks usually involved in engineering low bit quantized networks . 3 BINARY DECOMPOSITION . We approximate weights on k bits by using the sign and magnitude representation , due to its simplicity . A weight tensor of a layer l can be decomposed as : θlk = ( k−2∑ i=0 ali · 2i+αl ) · ( −1 ) a l k−1 ( 1 ) with al ∈ { 0 , 1 } representing the binary coefficients and k the number of bits . The summation encodes the magnitude of the number while the second factor encodes the sign : this way we obtain numbers in a symmetric interval around zero . We add a negative constant αl to the exponent in order to allow the representation of fractional numbers ( see Table 1 ) . Additionally , this term controls the magnitude of the numbers and , therefore , the width of the weights distribution within a layer . Choosing αl < −k + 1 the weights are guaranteed to be less than 1 in magnitude . In order to constrain a to take binary values we use auxiliary floating point variables x ∈ R ( virtual bits ) passed through a unit step function : a = H ( x ) = 0 if x ≤ 0 , otherwise 1 . The weight initialization for the k bit training technique is as follows : for a fully connected layer the weight matrix is expanded into a 3D tensor of shape ( k , nl−1 , nl ) with k representing the number of bits and nl−1 , nl the number of nodes in the previous and current layer , respectively . Figure 1 illustrates a simple example of a ( 3 , 4 , 3 ) bit-tensor . For convolutional layers , where a weight tensor is in higher dimension , the procedure is analogous to the fully connected case and the bit-tensor is now of shape ( k , sx , sy , nl−1 , nl ) with sx , sy representing the kernel sizes in x and y direction . The value for each bit is chosen randomly with equal probability of being either 0 or 1 ( x ≤ 0 in the first case and x > 0 in the second ) . We ensure the weights sampled in this manner are not initialized at exactly zero , because this would mean pruning the network from the start and invalidate our hypothesis . Hence we obtain a uniform weight distribution without zeros . We adopt the Kaiming He ( He et al. , 2015a ) initialization technique for each layer ’ s weights , which means the standard deviation is √ 2/nl−1 , where nl−1 is the number of nodes in the previous layer . We have determined αl algorithmically via a simple binary search such that this condition is fulfilled for the weight distribution of each layer . This term is a fixed parameter in each layer and depends only on the structure of the network . The virtual bits , x , are chosen from a normal distribution which also satisfies the Kaiming He condition on its variance . For the particular situation where k = 2 the weights have only two values and the standard deviation is exactly 2α . Ramanujan et al . ( 2019 ) refer to this distribution as the Signed Kaiming Constant . During training , the feed-forward step is performed as usual , with the weights being calculated according to Eq . ( 1 ) . The backpropagation phase uses the straight through estimator ( STE ) ( Hinton , 2012 ; Bengio , 2013 ) for the step function introduced in the weight ’ s binary decomposition . The derivative of a hard threshold function such as the Heaviside step function is zero everywhere except at zero ( more specifically it is the Dirac delta function ) . Since the values of the weights are passed through this step function are almost never exactly zero , the gradients during backpropagation will almost always be zero . This situation leads to a stagnant network which never updates its weights and never learns . To avoid this , during the backpropagation phase the gradient of the step function is replaced by the gradient of a different function which is non-zero on a domain larger than for the step function . Such functions are usually referred to as proxy functions and can take many forms . Yin et al . ( 2019 ) ; Shekhovtsov & Yanush ( 2020 ) provide in-depth discussions on the properties of STEs . Throughout this work we adopt the method first proposed by Hinton ( 2012 ) which treats the gradient of a hard threshold function as if it were the identity function . This method has been shown to work very well in practice ( Liu et al. , 2018 ; Bulat et al. , 2019 ; 2021 ; Bethge et al. , 2019 ; Alizadeh et al. , 2019 ) Notice that in Eq . ( 1 ) the additive constant αl can be factored out of the sum . The resulting weights are in the form θlk = 2 αl · Θlk , where Θlk contains only integer numbers on k bits . The ReLU activation function has the property that σ ( α · x ) = α · σ ( x ) for any α > 0 . It can be shown that for a ReLU network of depth L , scaling the weights of each layer by a factor αl , with l ∈ [ 0 , 1 , . . . L−1 ] representing the layer index , is equivalent to scaling just a single layer with α = ∏L−1 l=0 αl , including the input layer . This means that we can gather all factors αl into a single α , scale the input images with that factor and train the network with just integer numbers represented on k bits . At inference time , for classification tasks α is not relevant because the output nodes are all scaled by the same coefficient and argmax ( α · x ) = argmax ( x ) for any α > 0 .

Summary: This paper proposes to directly train the bit values of each parameter in a neural network, instead of directly optimizing the floating point value of each weight. By varying the number of bits that are allowed to be optimized, the authors show that with less bits the network will become automatically sparser. This method has many interesting applications, including fixing some bits to be a message and only training the rest.