Efficient Neural Network Training via Forward and Backward Propagation Sparsification

1 Introduction . In the last decade , deep neural networks ( DNNs ) [ 35 , 11 , 38 ] have proved their outstanding performance in various fields such as computer vision and natural language processing . However , training such large-sized networks is still very challenging , requiring huge computational power and storage . This hinders us from exploring larger networks , which are likely to have better performance . Moreover , it is a widely-recognized property that modern neural networks are significantly overparameterized , which means that a fully trained network can always be sparsified dramatically by network pruning techniques [ 9 , 8 , 25 , 46 , 20 ] into a small sub-network with negligible degradation in accuracy . After pruning , the inference efficiency can be greatly improved . Therefore , a natural question is can we exploit this sparsity to improve the training efficiency ? The emerging technique called sparse network training [ 10 ] is closely related with our question , which can obtain sparse networks by training from scratch . We can divide existing methods into two categories , i.e. , parametric and non-parametric , based on whether they explicitly parameterize network structures with trainable variables ( termed structure parameters ) . Empirical results [ 24 , 34 , 44 , 23 ] demonstrate that the sparse networks they obtain have comparable accuracy with those ∗Equal contribution †Jointly with Google Research 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . obtained from network pruning . However , most of them narrowly aim at finding a sparse subnetwork instead of simultaneously sparsifying the computation of training by exploiting the sparse structure . As a consequence , it is hard for them to effectively accelerate the training process in practice on general platforms , e.g. , Tensorflow [ 1 ] and Pytorch [ 31 ] . Detailed reasons are discussed below : • Non-parametric methods find the sparse network by repeating a two-stage procedure that alternates between weight optimization and pruning [ 10 , 6 ] , or by adding a proper sparsityinducing regularizer on the weights to the objective [ 22 , 41 ] . The two-stage methods prune the networks in weight space and usually require retraining the obtained subnetwork from scratch every time when new weights are pruned , which makes training process even more time-consuming . Moreover , the computation of regularized methods is dense since the gradients of a zero-valued weights/filters are still nonzero . • All the parametric approaches estimate the gradients based on chain rule . The gradient w.r.t . the structure parameters can be nonzero even when the corresponding channel/weight is pruned . Thus , to calculate the gradient via backward propagation , the error has to be propagated through all the neurons/channels . This means that the computation of backward propagation has to be dense . Concrete analysis can be found in Section 3 . We notice that some existing methods [ 4 , 28 ] can achieve training speedup by careful implementation . For example , the dense to sparse algorithm [ 28 ] removes some channels if the corresponding weights are quite small for a long time . However , these methods always need to work with a large model at the beginning epochs and consume huge memory and heavy computation in the early stage . Therefore , even with such careful implementations , the speedups they can achieve are still limited . In this paper , we propose an efficient channel-level parametric sparse neural network training method , which is comprised of completely sparse ( See Remark 1 ) forward and backward propagation . We adopt channel-level sparsity since such sparsity can be efficiently implemented on the current training platforms to save the computational cost . In our method , we first parameterize the network structure by associating each filter with a binary mask modeled as an independent Bernoulli random variable , which can be continuously parameterized by the probability . Next , inspired by the recent work [ 47 ] , we globally control the network size during the whole training process by controlling the sum of the Bernoulli distribution parameters . Thus , we can formulate the sparse network training problem into a constrained minimization problem on both the weights and structure parameters ( i.e. , the probability ) . The main novelty and contribution of this paper lies in our efficient training method called completely sparse neural network training for solving the minimization problem . Specifically , to fully exploit the sparse structure , we separate training iteration into two parts , i.e. , weight update and structure parameter update . For weight update , the conventional backward propagation is used to calculate the gradient , which can be sparsified completely because the gradients of the filters with zero valued masks are also zero . For structure parameter update , we develop a new variance reduced policy gradient estimator ( VR-PGE ) . Unlike the conventional chain rule based gradient estimators ( e.g. , straight through [ 2 ] ) , VR-PGE estimates the gradient via two forward propagations , which is completely sparse because of the sparse subnetwork . Finally , extensive empirical results demonstrate that our method can significantly accelerate the training process of neural networks . The main contributions of this paper can be summarized as follows : • We develop an efficient sparse neural network training algorithm with the following three appealing features : – In our algorithm , the computation in both forward and backward propagations is completely sparse , i.e. , they do not need to go through any pruned channels , making the computational complexity significantly lower than that in standard training . – During the whole training procedure , our algorithm works on small sub-networks with the target sparsity instead of follows a dense-to-sparse scheme . – Our algorithm can be implemented easily on widely-used platforms , e.g. , Pytorch and Tensorflow , to achieve practical speedup . • We develop a variance reduced policy gradient estimator VR-PGE specifically for sparse neural network training , and prove that its variance is bounded . • Experimental results demonstrate that our methods can achieve significant speed-up in training sparse neural networks . This implies that our method can enable us to explore larger-sized neural networks in the future . Remark 1 . We call a sparse training algorithm completely sparse if both its forward and backward propagation do not need to go through any pruned channels . For such algorithms , the computational cost in forward and backward propagation cost can be roughly reduced to ρ2 ∗ 100 % , with ρ being the ratio of remaining unpruned channels . 2 Related Work . In this section , we briefly review the studies on neural network pruning , which refers to the algorithms that prune DNNs after fully trained , and the recent works on sparse neural network training . 2.1 Neural Network Pruning . Network Pruning [ 10 ] is a promising technique for reducing the model size and inference time of DNNs . The key idea of existing methods [ 10 , 8 , 46 , 20 , 27 , 13 , 48 , 40 , 43 , 32 , 16 ] is to develop effective criteria ( e.g , weight magnitude ) to identify and remove the massive unimportant weights contained in networks after training . To achieve practical speedup on general devices , some of them prune networks in a structured manner , i.e. , remove the weights in a certain group ( e.g. , filter ) together , while others prune the weights individually . It has been reported in the literature [ 8 , 25 , 46 , 20 ] that they can improve inference efficiency and reduce memory usage of DNNs by orders of magnitudes with minor loss in accuracy , which enables the deployment of DNNs on low-power devices . We notice that although some pruning methods can be easily extended to train sparse networks , they can not accelerate or could even slow down the training process . One reason is they are developed in the scenario that a fully trained dense network is given , and can not work well on the models learned in the early stage of training . Another reason is after each pruning iteration , one has to fine tune or even retrain the network for lots of epoch to compensate the caused accuracy degradation . 2.2 Sparse Neural Network Training . The research on sparse neural network training has emerged in the recent years . Different from the pruning methods , they can find sparse networks without pre-training a dense one . Existing works can be divided into four categories based on their granularity in pruning and whether the network structures are explicitly parameterized . To the best of our knowledge , no significant training speedups achieved in practice are reported in the literature . Table 1 summarizes some representative works . Weight-level non-parametric methods , e.g. , [ 6 , 10 , 48 , 29 , 30 ] , always adopt a two-stage training procedure that alternates between weight optimization and pruning . They differ in the schedules of tuning the prune ratio over training and layers . [ 10 ] prunes the weights with the magnitude below a certain threshold and [ 48 , 6 ] gradually increase the pruning rate during training . [ 30 , 5 ] automatically reallocate parameters across layers during training via controlling the global sparsity . Channel-level non-parametric methods [ 12 , 41 ] are proposed to achieve a practical acceleration in inference . [ 41 ] is a structured sparse learning method , which adds a group Lasso regularization into the objective function of DNNs with each group comprised of the weights in a filter . [ 12 ] proposes a soft filter pruning method . It zeroizes instead of hard pruning the filters with small ` 2 norm , after which these filters are treated the same with other filters in training . It is obvious that these methods can not achieve significant speedup in training since they need to calculate the full gradient in backward propagation although the forward propagation could be sparsified if implemented carefully . Parametric methods multiply each weight/channel with a binary [ 47 , 44 , 37 , 42 ] or continuous [ 24 , 26 , 19 , 18 ] mask , which can be either deterministic [ 24 , 42 ] or stochastic [ 47 , 44 , 26 , 37 , 19 , 18 ] . The mask is always parameterized via a continuous trainable variable , i.e. , structure parameter . The sparsity is achieved by adding sparsity-inducing regularizers on the masks . The novelties of these methods lie in estimating the gradients w.r.t structure parameters in training . To be precise , • Deterministic Binary Mask . [ 42 ] parameterizes its deterministic binary mask as a simple step function and estimates the gradients via sigmoid straight through estimator ( STE ) [ 2 ] . • Deterministic Continuous Mask . [ 24 ] uses the linear coefficients of batch normalization ( BN ) as a continuous mask and enforces most of them to 0 by penalizing the objective with ` 1 norm of the coefficients . [ 18 ] defines the mask as a soft threshold function with learnable threshold . These methods can estimate the gradients via standard backward propagation . • Stochastic Binary Mask . [ 44 , 37 ] model the mask as a bernoulli random variable and the gradients w.r.t . the parameters of bernoulli distributions are estimated via STE . [ 47 ] estimates the gradients via Gumbel-Softmax trick [ 15 ] , which is more accurate than STE . • Stochastic Continuous Mask . [ 26 , 19 ] parameterize the mask as a continuous function g ( c , ) , which is differentiable w.r.t . c , and is a parameter free noise , e.g. , Gaussian noise N ( 0 , 1 ) . In this way , the gradients can be calcuated via conventional backward propagation . Therefore , we can see that all of these parametric methods estimate the gradients of the structure parameters based on the chain rule in backward propagation . This makes the training iteration can not be sparsified by exploiting the sparse network structure . For the details , please refer to Section 3 .

Summary: Paper proposes a variational based pruning method where, with 2 structurally sparse forward passes and 1 sparse backward, the training can be significantly faster. Description is clear and theoretically sounds, equations are easy to follow, code is provided, results are encouraging. The paper might be a great benchmark and starting point for faster pruning-while-training with efficient inference.