Learning N:M  Fine-grained Structured Sparse Neural Networks From Scratch

1 INTRODUCTION . Deep neural networks ( DNNs ) have shown promising performances on various tasks including computer vision , natural language processing , speech recognition , etc . However , a DNN usually comes with a large number of learnable parameters , ranging from millions of to even billions of ( e.g. , GPT-3 ( Brown et al. , 2020 ) ) , making the DNN model burdensome and difficult to be applied to real-world deployments . Therefore , researchers began to investigate how to speed up and compress DNNs via various methods such as knowledge distillation ( Hinton et al. , 2015 ) , quantization ( Jacob et al. , 2018 ; Zhou et al. , 2017 ) , designing efficient model architectures ( Howard et al. , 2017 ) , and structured sparsity ( Wen et al. , 2016 ; Li et al. , 2016 ) . In this paper , we focus on the problem of sparsifying DNNs . Sparsity in DNNs can be categorized into unstructured sparsity and structured sparsity . Unstructured sparsity prunes individual weights at any location , which is fine-grained and can achieve extremely high compression ratio ( Han et al. , 2015 ; Guo et al. , 2016 ) . However , unstructured sparsity struggles to take advantage of vectorprocessing architectures , which increases latency due to dependent sequences of reads ( Nvidia , 2020 ) . Compared with unstructured sparsity , structured sparsity is more friendly to hardware , especially for block pruning ( Wang et al. , 2019 ) , kernel shape sparsity ( Tan et al. , 2020 ) or channel and filter pruning ( Li et al. , 2016 ; Wen et al. , 2016 ) . Although structured sparsity can speed up DNNs on commodity hardware , it hurts model performance more significantly than unstructured fine-grained sparsity . For example , ResNet-50 network generated by unstructured pruning can achieve a 5.96× compression ratio , with the same accuracy as the original network , but it can only achieve 1× com- ∗The first two authors equally contribute to this paper . pression in the case of structured sparsity ( Renda et al. , 2020 ) . Therefore , how to combine the unstructured sparsity and structured sparsity to accelerate DNNs on modern hardware ( e.g. , GPU ) becomes a challenging yet valuable problem . Recently , Nvidia Ampere A100 is equipped with the Sparse Tensor Cores to accelerate 2:4 structured fine-grained sparsity . Here , N : M sparsity indicates the sparsity of DNNs in which onlyN weights are non-zero for every continuousM weights . To the best of our knowledge , A100 is the first commodity sparse hardware , where the sparse tensor core can support several common operations including linear , convolutional , recurrent cells , transformer blocks , etc . Specifically , suppose a typical matrix multiplication X ×W in DNNs , X andW denote input tensor and parameter tensor respectively . The Dense Tensor Cores implementX16×32×W32×8 matrix multiplication by 2 cycles while the Sparse Tensor Cores only need 1 cycle if the parameter tensorW satisfies the 2:4 structured sparse pattern . Nvidia has proposed an ASP1 ( APEX ’ s Automatic Sparsity ) solution ( Nvidia , 2020 ) to sparsify a dense neural network to satisfy the 2:4 fine-grained structured sparsity requirement . The recipe contains three steps : ( 1 ) training a dense network until converge ; ( 2 ) pruning for 2:4 sparsity with magnitude-based single-shot pruning ; ( 3 ) repeating the original training procedure . However , ASP is computationally expensive since it requires training the full dense models from scratch and finetuning again . Therefore , we still lack a simple recipe to obtain a structured sparse DNN model consistent with the dense network without extra fine-tuning . This paper addresses this question : Can we design a simple yet universal recipe to learn N : M sparse neural networks from scratch in an efficient way ? It is difficult to find the optimal sparse architecture ( connections ) and optimal parameters ( Evci et al. , 2019b ) simultaneously during training sparse CNNs and Transformers although SET-MLP could easily outperform dense MLP ( Bourgin et al. , 2019 ) . There are two schemes to obtain such sparse models . One is a two-stage scheme , which discovers a sparse neural architecture by pruning a well-trained dense network and then uses the same or even greater computational resources to retrain the sparse models ( Nvidia , 2020 ; Evci et al. , 2019b ; Han et al. , 2015 ; Frankle & Carbin , 2018 ) . The other is a one-stage scheme , which adopts the dynamic method to alternatively optimize parameters and prunes network architectures based on different criteria ( Bellec et al. , 2017 ; Mocanu et al. , 2018 ; Mostafa & Wang , 2019 ; Evci et al. , 2019b ; Kusupati et al. , 2020 ; Dettmers & Zettlemoyer , 2019 ) . Compared with the two-stage scheme , the one-stage scheme can save training time and cost however usually obtains lower performance . To overcome the aforementioned trade-off between training cost and performance , we present a simple yet effective framework to train sparse neural networks from scratch . Specifically , we employ the magnitude-based pruning method ( Renda et al. , 2020 ; Gale et al. , 2019 ) during the forward process . Considering that the pruning operation is a non-differentiable operator ( a similar dilemma in model quantization ( Courbariaux et al. , 2016 ) ) , we extend the widely used Straight-through Estimator ( STE ) ( Bengio et al. , 2013 ) in model quantization to aid sparse neural network ’ s backpropagation . However , perturbations are introduced during the back-propagation ( Yin et al. , 2019 ; Bengio et al. , 2013 ) . Hence we define Sparse Architecture Divergence ( SAD ) to further analyze N : M sparse networks trained by STE methods so that we can identify the impact of perturbations on sparse neural networks training . Based on SAD analysis , to alleviate the negative impact , we propose a sparse-refined term mitigating the approximated gradients ’ influence . We also compare the performance of neural networks with different granularities of fine-grained structured sparsity ( i.e. , 1:4 , 2:4 , 2:8 , 4:8 ) and conduct thorough experiments on several typical deep neural networks with different N : M sparsity levels , covering image classification , detection , segmentation , optical flow estimation , and machine translation . Experimental results have shown that the models with our proposed structured sparsity can achieve neglectful performance drop and can even sometimes outperform the dense model . The main contributions of this paper are summarized as three-fold . ( 1 ) To the best of our knowledge , this is the first systematic study into training N : M structured sparse neural networks from scratch without performance drop . The N : M structured sparsity is a missing yet promising ingredient in model acceleration , which can be a valuable supplement with various compression methods . ( 2 ) We extend STE to tackle the problem of training N : M sparse neural networks . To alleviate the limitations of STE on sparsifying networks , we propose a sparse refined term to enhance the effectiveness 1https : //github.com/NVIDIA/apex/tree/master/apex/contrib/sparsity on training the sparse neural networks from scratch . ( 3 ) We conduct extensive experiments on various tasks with N : M fine-grained sparse nets , and provide benchmarks for N : M sparse net training to facilitate co-development of related software and hardware design . 2 RELATED WORK . Unstructured and Structured Sparsity . Sparsity of DNNs is a promising direction to compress and accelerate a deep learning model . Among all sparsity types , unstructured sparsity can achieve a significantly high compression ratios ( e.g . 13× ( Han et al. , 2015 ) and 108× ( Guo et al. , 2016 ) ) while ensuring decent accuracy by pruning . Many different pruning criterions and pruning methods are proposed for unstructured sparsity , e.g. , magnitude-based pruning ( Han et al. , 2015 ; Frankle & Carbin , 2018 ) , Hessian based heuristics ( LeCun et al. , 1990 ) , and pruning with connection sensitivity ( Lee et al. , 2018 ) . However , unstructured sparsity ’ s ability to accelerate is highly limited since it takes a lot of overhead to store the irregular non-zero index matrix . On the other hand , Wen et al . ( 2016 ) introduces the structural sparsity to speed up deep models on GPUs . Existing structural sparsity contains filter-wise sparsity ( Li et al. , 2016 ) , channel-wise sparsity ( Li et al. , 2016 ) , filtershape-wise sparsity . Different from existing sparsity patterns ( fine-grained unstructured sparsity and coarse-grained structured sparsity ) , this paper presents an N : M fine-grained structured sparsity , a sparsity type that has both high efficiency and lossless performance . One-stage and two-stage methods . There are mainly two types of techniques to obtain a sparse neural network , one-stage methods and two-stage ones . The two-stage method first prunes a trained dense neural network and then retrains fixed sparse network to recover its performance . Typical two-stage methods include single-shot pruning ( Lee et al. , 2018 ) and iterative pruning ( Han et al. , 2015 ; Guo et al. , 2016 ) . Later , the lottery ticket hypothesis ( Frankle & Carbin , 2018 ) shows that the sparse sub-network ( winning tickets ) can be trained from scratch with the same initialization while the winning tickets are discovered by dense training . Deep-Rewiring ( Bellec et al. , 2017 ) , on the other hand , is a typical one-stage method , which takes a Bayesian perspective and samples sparse network connections from a posterior , however is computationally expensive and challenging to be applied to large-scale tasks . Sparse Evolutionary Training ( Mocanu et al. , 2018 ) ( SET ) is proposed as a simpler scheme where weights are pruned according to the standard magnitude criterion used in pruning and growing connections in random locations . Dettmers & Zettlemoyer ( 2019 ) uses the momentum of each parameter as the criterion for growing weights and receives an improvement in test accuracy . GMP ( Gale et al. , 2019 ) trains the unstructured sparse net using variational dropout and l0 regularization from scratch , and shows that unstructured sparse architectures learned through pruning can not be trained from scratch to have the same testing performance as dense models do . Recently proposed state-of-the-art method STR ( Kusupati et al. , 2020 ) introduces pruning learnable thresholds to obtain a non-uniform sparse network . RigL ( Evci et al. , 2019a ) uses the magnitudebased method to prune and the periodic dense gradients to regrow connection . However , compared with training dense neural networks from scratch , to achieve the same performance , RigL needs 5× more training time.The most closely related work to ours may be DNW ( Wortsman et al. , 2019 ) which uses a fully dense gradient in the backward run to discover optimal wiring on the fly . 3 METHOD . 3.1 N : M FINE-GRAINED STRUCTURED SPARSITY Here we define the problem of training a neural network with N : M fine-grained structured sparsity . A neural network with N : M sparsity satisfies that , in each group of M consecutive weights of the network , there are at most N weights have non-zero values . Fig . 1 illustrates a 2:4 sparse network . Generally , our objective is to train an N : M sparse neural network as min S ( W , N , M ) L ( W ; D ) , ( 1 ) where D denotes the observed data , L represents the loss function , W = { W l : 0 < l 5 L } indicates the parameters of an L-layer neural network , and S ( W , N , M ) is the N : M sparse neural network parameters . 3.2 STRAIGHT-THROUGH ESTIMATOR ( STE ) ON TRAINING N : M SPARSE NETWORKS A straightforward solution for training an N : M sparsity network is to simply extend Straightthrough Estimator ( STE ) ( Bengio et al. , 2013 ) to perform online magnitude-based pruning and sparse parameter updating , which is depicted in Fig . 2 ( a ) . STE is widely used in model quantization ( Rastegari et al. , 2016 ) , since the quantized function is non-differentiable without STE and the networks optimized with STE has decent performance under careful settings ( Yin et al. , 2019 ) . In STE , a dense network is maintained during the training process . During the forward pass , we project the dense weightsW into sparse weights W̃ = S ( W , N , M ) satisfying N : M sparsity . Let w ⊂ W be a group of consecutive M parameters inW and w̃ ⊂ W̃ be the corresponding group in W̃ . The projection of w can be formulated as : w̃i = { wi if |wi| > ξ 0 if |wi| < ξ , for i = 1 , 2 , . . . , M ( 2 ) where ξ is the N -th largest value in w = { |w1| , |w2| , . . . , |wM | } . Intuitively , this projection function S ( · ) produces sparse parameters W̃ by setting N parameters that have the least significant absolute values to zero in each consecutive M -parameter group , while keeping the other parameters the same as before . The computation of an N : M sparse sub-network on-the-fly in the forward pass is illustrated in Fig . 1 . The projection function S ( · ) , which is non-differentiable during back-propagation , generates the N : M sparse sub-network on the fly . To get gradients during back-propagation , STE computes the gradients of the sub-network g ( W̃ ) = 5W̃L ( W̃ ; D ) based on the sparse sub-network W̃ , which can be directly back-projected to the dense network as the approximated gradients of the dense parameters . The approximated parameter update rule for the dense network ( see Fig . 2 ( a ) in Appendix ) can be formulated as Wt+1 ←Wt − γtg ( W̃t ) , ( 3 ) whereWt represents dense parameters at iteration t and γt indicates the learning rate .

Summary: The paper introduces a new sparse training algorithm (SR-STE) based on the straight through estimator which is specially designed for the hardware constraints of Nvidia A100 GPU. Auxiliary, in order to study better this algorithm, the paper introduces also a metric (SAD) to measure the changes in the sparse network topology during training. The contributions have real added value as they show that sparse neural networks can actually benefit of hardware designed to consider sparsity. The experiments on CNNs and Transformers support the claims.