Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers

1 INTRODUCTION . Despite the impressive success that deep neural networks have achieved in a wide range of challenging tasks , the inference in deep neural networks is highly memory-intensive and computationintensive due to the over-parameterization of deep neural networks . Network pruning ( LeCun et al . ( 1990 ) ; Han et al . ( 2015 ) ; Molchanov et al . ( 2017 ) ) has been recognized as an effective approach to improving the inference efficiency in resource-limited scenarios . Traditional pruning methods consist of dense network training followed with pruning and fine-tuning iterations . To avoid the expensive pruning and fine-tuning iterations , many sparse training methods ( Mocanu et al. , 2018 ; Bellec et al. , 2017 ; Mostafa & Wang , 2019 ; Dettmers & Zettlemoyer , 2019 ) have been proposed , where the network pruning is conducted during the training process . However , all these methods suffer from following three problems : Coarse-grained predefined pruning schedule . Most of the existing pruning methods use a predefined pruning schedule with many additional hyperparameters like pruning a % parameter each time and then fine-tuning for b epochs with totally c pruning steps . It is non-trivial to determine these hyperparameters for network architectures with various degrees of complexity . Therefore , usually a fixed pruning schedule is adopted for all the network architectures , which means that a very simple network architecture like LeNet-300-100 will have the same pruning schedule as a far more complex network like ResNet-152 . Besides , almost all the existing pruning methods conduct epoch-wise pruning , which means that the pruning is conducted between two epochs and no pruning operation happens inside each epoch . Failure to properly recover the pruned weights . Almost all the existing pruning methods conduct ” hard ” pruning that prunes weights by directly setting their values to 0 . Many works ( Guo et al. , 2016 ; Mocanu et al. , 2018 ; He et al. , 2018 ) have argued that the importance of network weights are not fixed and will change dynamically during the pruning and training process . Previously unimportant weights may tend to be important . So the ability to recover the pruned weights is of high significance . However , directly setting the pruned weights to 0 results in the loss of historical parameter importance , which makes it difficult to determine : 1 ) whether and when each pruned weight should be recovered , 2 ) what values should be assigned to the recovered weights . Therefore , existing methods that claim to be able to recover the pruned weights simply choose a predefined portion of pruned weights to recover and these recover weights are randomly initialized or initialized to the same value . Failure to properly determine layer-wise pruning rates . Modern neural network architectures usually contain dozens of layers with a various number of parameters . Therefore , the degree of parameter redundancy is very different among the layers . For simplicity , some methods prune the same percentage of parameters at each layer , which is not optimal . To obtain dynamic layer-wise pruning rates , a single global pruning threshold or layer-wise greedy algorithms are applied . Using a single global pruning threshold is exceedingly difficult to assess the local parameter importance of the individual layer since each layer has a significantly different amount of parameter and contribution to the model performance . This makes pruning algorithms based on a single global threshold inconsistent and non-robust . The problem of layer-by-layer greedy pruning methods is that the unimportant neurons in an early layer may have a significant influence on the responses in later layers , which may result in propagation and amplification of the reconstruction error ( Yu et al. , 2018 ) . We propose a novel end-to-end sparse training algorithm that properly solves the above problems . With only one additional hyperparameter used to set the final model sparsity , our method can achieve dynamic fine-grained pruning and recovery during the whole training process . Meanwhile , the layerwise pruning rates will be adjusted automatically with respect to the change of parameter importance during the training and pruning process . Our method achieves state-of-the-art performance compared with other sparse training algorithms . The proposed algorithm has following promising properties : • Step-wise pruning and recovery . A training epoch usually will have tens of thousands of training steps , which is the feed-forward and back-propagation pass for a single mini-batch . Instead of pruning between two training epochs with a predefined pruning schedule , our method prunes and recovers the network parameter at each training step , which is far more fine-grained than existing methods . • Neuron-wise or filter-wise trainable thresholds . All the existing methods adopt a single pruning threshold for each layer or the whole architecture . Our method defines a threshold vector for each layer . Therefore , our method adopts neuron-wise pruning thresholds for fully connected and recurrent layer and filter-wise pruning thresholds for convolutional layer . Additionally , all these pruning thresholds are trainable and will be updated automatically via back-propagation . • Dynamic pruning schedule . The training process of deep neural network consists of many hyperparameters . The learning rate is perhaps the most important hyperparameter . Usually , the learning rate will decay during the training process . Our method can automatically adjust the layer-wise pruning rates under different learning rates to get the optimal sparse network structure . • Consistent sparse pattern . Our algorithm can get a consistent layer-wise sparse pattern under different model sparsities , which indicates that our method can automatically determine the optimal layer-wise pruning rates given the target model sparsity . 2 RELATED WORK . Traditional Pruning Methods : LeCun et al . ( 1990 ) presented the early work about network pruning using second-order derivatives as the pruning criterion . The effective and popular training , pruning and fine-tuning pipeline was proposed by Han et al . ( 2015 ) , which used the parameter magnitude as the pruning criterion . Narang et al . ( 2017 ) extended this pipeline to prune the recurrent neural networks with a complicated pruning strategy . Molchanov et al . ( 2016 ) introduced first-order Taylor term as the pruning criterion and conduct global pruning . Li et al . ( 2016 ) used ` 1 regularization to force the unimportant parameters to zero . Sparse Neural Network Training : Recently , some works attempt to find the sparse network directly during the training process without the pruning and fine-tuning stage . Inspired by the growth and extinction of neural cells in biological neural networks , Mocanu et al . ( 2018 ) proposed a pruneregrowth procedure called Sparse Evolutionary Training ( SET ) that allows the pruned neurons and connections to revive randomly . However , the sparsity level needs to be set manually and the random recovery of network connections may provoke unexpected effects on the network . DEEP-R proposed by Bellec et al . ( 2017 ) used Bayesian sampling to decide the pruning and regrowth configuration , which is computationally expensive . Dynamic Sparse Reparameterization ( Mostafa & Wang , 2019 ) used dynamic parameter reallocation to find the sparse structure . However , the pruning threshold can only get halved if the percentage of parameter pruned is too high or get doubled if that percentage is too low for a certain layer . This coarse-grained adjustment of the pruning threshold significantly limits the ability of Dynamic Sparse Reparameterization . Additionally , a predefined pruning ratio and fractional tolerance are required . Dynamic Network Surgery ( Guo et al. , 2016 ) proposed pruning and splicing procedure that can prune or recover network connections according to the parameter magnitude but it requires manually determined thresholds that are fixed during the sparse learning process . These layer-wise thresholds are extremely hard to manually set . Meanwhile , Fixing the thresholds makes it hard to adapt to the rapid change of parameter importance . Dettmers & Zettlemoyer ( 2019 ) proposed sparse momentum that used the exponentially smoothed gradients as the criterion for pruning and regrowth . A fixed percentage of parameters are pruned at each pruning step . The pruning ratio and momentum scaling rate need to be searched from a relatively high parameter space . 3 DYNAMIC SPARSE TRAINING . 3.1 NOTATION . Deep neural network consists of a set of parameters { Wi : 1 ≤ i ≤ C } , where Wi denotes the parameter matrix at layer i and C denotes the number of layers in this network . For each fully connected layer and recurrent layer , the corresponding parameter is Wi ∈ Rco×ci , where co is the output dimension and ci is the input dimension . For each convolutional layer , there exists a convolution kernel Ki ∈ Rco×ci×w×h , where co is the number of output channels , ci is the number of input channels , w and h are the kernel sizes . Each filter in a convolution kernel Ki can be flattened to a vector . Therefore , a corresponding parameter matrix Wi ∈ Rco×z can be derived from each convolution kernel Ki ∈ Rco×ci×w×h , where z = ci × w × h. Actually , the pruning process is equivalent to finding a binary parameter mask Mi for each parameter matrix Wi . Thus , a set of binary parameter masks { Mi : 1 ≤ i ≤ C } will be found by network pruning . Each element for these parameter masks Mi is either 1 or 0 . 3.2 THRESHOLD VECTOR AND DYNAMIC PARAMETER MASK . Pruning can be regarded as applying a binary mask M to each parameter W . This binary parameter mask M preserves the information about the sparse structure . Given the parameter set { W1 , W2 , · · · , WC } , our method will dynamically find the corresponding parameter masks { M1 , M2 , · · · , MC } . To achieve this , for each parameter matrix W ∈ Rco×ci , a trainable pruning threshold vector t ∈ Rco is defined . Then we utilize a unit step function S ( x ) as shown in Figure 2 ( a ) to get the masks according to the magnitude of parameters and corresponding thresholds as present below . Qij = F ( Wij , ti ) = |Wij | − ti 1 ≤ i ≤ co , 1 ≤ j ≤ ci ( 1 ) Mij = S ( Qij ) 1 ≤ i ≤ co , 1 ≤ j ≤ ci ( 2 ) With the dynamic parameter mask M , the corresponding element in mask Mij will be set to 0 if Wij needs to be pruned . This means that the weight Wij is masked out by the 0 at Mij to get a sparse parameter W M . The value of underlying weight Wij will not change , which preserves the historical information about the parameter importance . For a fully connected layer or recurrent layer with parameter W ∈ Rco×ci and threshold vector t ∈ Rco , each weight Wij will have a neuron-wise threshold ti , where Wij is the jth weight associated with the ith output neuron . Similarly , the thresholds are filter-wise for convolutional layer . Besides , a threshold matrix or a single scalar threshold can also be chosen . More details are present in Appendix A.2 . 3.3 TRAINABLE MASKED LAYERS . With the threshold vector and dynamic parameter mask , the trainable masked fully connected , convolutional and recurrent layer are introduced as shown in Figure 1 , where X is the input of current layer and Y is the output . For fully connected and recurrent layers , instead of the dense parameter W , the sparse product W M will be used in the batched matrix multiplication , where denote Hadamard product operator . For convolutional layers , each convolution kernel K ∈ Rco×ci×w×h can be flatten to get W ∈ Rco×z . Therefore , the sparse kernel can be obtained by a similar process as fully connected layers . This sparse kernel will be used for subsequent convolution operation . In order to make the elements in threshold vector t trainable via back-propagation , the derivative of the binary step function S ( x ) is required . However , its original derivative is an impulse function whose value is zero almost everywhere and infinite at zero as shown in Figure 2 ( b ) . Thus the original derivative of the binary step function S ( x ) can not be applied in back-propagation and parameter updating directly . Some previous works ( Hubara et al . ( 2016 ) ; Rastegari et al . ( 2016 ) ; Zhou et al . ( 2016 ) ) demonstrated that by providing a derivative estimation , it is possible to train networks containing such binarization function . A clip function called straight-through estimator ( STE ) ( Bengio et al. , 2013 ) was employed in these works and is illustrated in Figure 2 ( c ) . Furthermore , Xu & Cheung ( 2019 ) discussed the derivative estimation in a balance of tight approximation and smooth back-propagation . We adopt this long-tailed higher-order estimator H ( x ) in our method . As shown in Figure 2 ( d ) , it has a wide active range between [ −1 , 1 ] with a non-zero gradient to avoid gradient vanishing during training . On the other hand , the gradient value near zero is a piecewise polynomial function giving tighter approximation than STE . The estimator is represented as d dx S ( x ) ≈ H ( x ) = 2− 4|x| , −0.4 ≤ x ≤ 0.4 0.4 , 0.4 < |x| ≤ 1 0 , otherwise ( 3 ) With this derivative estimator , the elements in the vector threshold t can be trained via backpropagation . Meanwhile , in trainable masked layers , the network parameter W can receive two branches of gradient , namely the performance gradient for better model performance and the structure gradient for better sparse structure , which helps to properly update the network parameter under sparse network connectivity . The structure gradient enables the pruned ( masked ) weights to be updated via back-propagation . The details about the feed-forward and back-propagation of the trainable masked layer are present in Appendix A.3 . Therefore , the pruned ( masked ) weights , the unpruned ( unmasked ) weights and the elements in the threshold vector can all be updated via back-propagation at each training step . The proposed method will conduct fine-grained step-wise pruning and recovery automatically .

Summary: This paper proposes an algorithm for training networks with sparse parameter tensors. This involves achieving sparsity by application of a binary mask, where the mask is determined by current parameter values and a learned threshold. It also involves the addition of a specific regularizer which encourages the thresholds used for the mask to be large. Gradients with respect to both masked-out parameters, and with respect to mask thresholds, are computed using a "long tailed" variant of the straight-through-estimator.