Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks

Unstructured pruning reduces the memory footprint in deep neural networks ( DNNs ) . Recently , researchers proposed different types of structural pruning intending to reduce also the computation complexity . In this work , we first suggest a new measure called mask-diversity which correlates with the expected accuracy of the different types of structural pruning . We focus on the recently suggested N : M fine-grained block sparsity mask , in which for each block of M weights , we have at least N zeros . While N : M fine-grained block sparsity allows acceleration in actual modern hardware , it can be used only to accelerate the inference phase . In order to allow for similar accelerations in the training phase , we suggest a novel transposable fine-grained sparsity mask , where the same mask can be used for both forward and backward passes . Our transposable mask guarantees that both the weight matrix and its transpose follow the same sparsity pattern ; thus , the matrix multiplication required for passing the error backward can also be accelerated . We formulate the problem of finding the optimal transposable-mask as a minimum-cost flow problem . Additionally , to speed up the minimum-cost flow computation , we also introduce a fast linear-time approximation that can be used when the masks dynamically change during training . Our experiments suggest a 2x speed-up in the matrix multiplications with no accuracy degradation over vision and language models . Finally , to solve the problem of switching between different structure constraints , we suggest a method to convert a pre-trained model with unstructured sparsity to an N : M fine-grained block sparsity model with little to no training . A reference implementation can be found at https : //github.com/papers-submission/structured_transposable_masks . 1 Introduction . Deep neural networks ( DNNs ) have established themselves as the first-choice tool for a wide range of applications , including computer vision and natural language processing . However , their impressive performance comes at a price of extensive infrastructure costs — as state-of-the-art DNNs may contain trillions of parameters [ 10 ] and require thousands of petaflops [ 6 ] for the training process . For this reason , compression of DNNs training and inference process is a research topic of paramount ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . importance in both academia and industry . The main techniques of compression include quantization [ 2 , 34 ] , knowledge distillation [ 17 ] , and pruning [ 15 , 23 ] . Pruning DNNs is one of the most popular and widely studied methods to improve DNN resource efficiency . The different pruning methods can be categorized into two different groups : unstructured and structured pruning . While the former can achieve a very high compression ratio , it usually fails in reducing the computational footprint in modern hardware . In contrast , structured pruning methods , such as block [ 41 ] or filter [ 23 ] pruning , are more hardware friendly . Unfortunately , these methods usually fail to keep the original accuracy for high compression ratios [ 37 ] . Finding an optimal structured sparsity pattern is still an ongoing research topic . Recently , Nvidia [ 36 ] announced the A100 GPU , containing sparse tensor cores which are able to accelerate fine-grained sparse matrix multiplication . The sparse tensor cores in A100 enable a 2x acceleration of regular matrix multiplication in DNNs , Y =WX , where W and X are weight and input matrices , respectively . The only requirement is that W would have a fine-grained 2:4 sparsity structure , i.e . out of every four contiguous elements in W , two are pruned . Nvidia [ 36 ] suggested a two-fold scheme for pruning a pretrained dense model : ( a ) Define a fine-grained 2:4 fixed mask , and ( b ) retrain with the masked weights using original training schedule . Indeed , the Nvidia [ 36 ] approach is very appealing for the common case where a pretrained dense model is given . While the Nvidia [ 36 ] method works well on many models , a pretrained model is not always given . In those cases , one has to first train a dense model and only then try to prune it . To alleviate this demand , Zhou et al . [ 43 ] suggested a method that trains from scratch a model with N : M fine-grained mask , using a sparse-refined straight-through estimator ( SR-STE ) . Similarly to the quantization-awaretraining methods [ 18 ] , they maintain a dense copy of the weights and prune it in every iteration , passing the gradients using the straight-through estimator [ 5 ] . Since the mask dynamically changes while training , they suggest adding an extra weight decay on the masked ( i.e . pruned ) elements to reduce the mask changes during the training process . As opposed to Evci et al . [ 9 ] that aims to reduce memory footprint for sparse training from scratch , Zhou et al . [ 43 ] only eliminates the need to train a dense model before pruning it . Motivated by these promising results , our goal here is to answer three remaining questions : ( Fig . 1 ) 1 . How to rank different types of sparsity masks ? We suggest a new measure called “ mask diversity '' , which is the first to connect mask constraints and network accuracy ( Section 3 ) . 2 . Can fine-grained sparsity masks accelerate training ? We start by observing both the forward and the backward matrix-multiplications involving the weight matrix W . Since the backward pass requires using the transposed matrix WT : ∂Loss ∂X =WT · ∂Loss ∂Y , ( 1 ) and in general WT does not have an N : M fine-grained sparsity structure ( even if W has this structure ) , the methods suggested in [ 43 , 36 ] accelerate only the forward pass matrix-multiplication , 3 . Can we change the type of sparsity mask structure without re-training ? Different hardware devices can support different types of fine-grained sparsity masks . Therefore , we suggest the `` Adaprune '' method , which converts between types of sparsity masks ( even from unstructured masks ) without the need of re-training , and almost no degradation in accuracy ( Section 5 ) . 2 Related work . Pruning of neural networks weights has been extensively investigated , starting with classical methods in the late 1980s [ 20 , 31 , 32 , 21 ] and then amounting to dozens of papers published in recent years . Since DNNs are generally over-parameterized , pruning the weights reduces their memory footprint . In special cases , when the sparsity mask has a specific pattern , it has the potential to reduce computational footprint as well . The most common practice is to prune a pretrained dense model , so that it will be sparse at deployment . Since the accuracy of the pretrained dense model is known , one can tune the sparsity level to ensure comparable accuracy for the sparse model . Recently , a new line of research that aims to train sparse models from scratch [ 14 , 7 , 9 ] has emerged . The goal is to train models that can not fit into currently available hardware . Next , we briefly overview the structured , unstructured , and accelerating-sparse-training categories . Unstructured pruning . It removes individual elements of the matrix , aiming for high total sparsity , while being agnostic to the location of the pruned elements . Standard pruning methods are based on different criteria , such as magnitude [ 15 ] , approximate L0 regularization [ 26 ] , or connection sensitivity [ 22 ] . Recent methods [ 12 ] , suggested to train a dense network until convergence , extract the required mask ( “ winning ticket '' ) , and use the original training regime to re-train the active weights from their original initialization or final values [ 38 ] using the original training schedule . These methods are able to achieve over 80 % sparsity on ResNet50- ImageNet dataset [ 38 ] . Despite the high sparsity ratio that can be achieved with these methods , modern hardware can not efficiently utilize such a form of sparsity for reducing computational resources [ 36 ] . Structured pruning . It removes weights in specific location based patterns , which are more useful for hardware acceleration . Such methods can be applied at either the level of channels or layers . For example , Li et al . [ 23 ] remove the channels with the lower norm , Luo et al . [ 27 ] prune channels according to the effect on the activation of the following layer , and [ 41 ] split the filters into multiple groups , applying a group Lasso regularization . All these methods are natively supported in both hardware and software , as they effectively change the model structure by reducing channels or groups . Yet , no such method was able to achieve a reasonable accuracy with sparsity levels higher than 50 % . As observed by Liu et al . [ 25 ] , filter pruning of a pretrained dense over-parameterized model is rarely the best method to obtain an efficient final model . Thus , here the structured pruning serves mostly as a DNN architecture search for the optimal compact model [ 40 , 42 ] . Our work is most closely related to Zhou et al . [ 43 ] , which is the first work that attempted training with a fine-grained N : M structured sparsity mask , as explained above . Accelerating sparse training for model deployment . The most common approach for sparse model deployment requires a three steps process : ( a ) train a dense model ; ( b ) define a sparsity mask ( c ) fine-tune while enforcing the mask on the model ’ s weights . The lottery ticket hypothesis of Frankle & Carbin [ 12 ] demonstrates that we can avoid the dense training , i.e. , step ( a ) , had we known how to choose the appropriate mask . Since Frankle & Carbin [ 12 ] discovered the optimal mask ( winning ticket ) by applying dense training , the question of how to find the optimal mask without training remained open . Since setting a predefined fixed mask results in some accuracy degradation ( Gray et al . [ 14 ] that requires expanding the model size ) , several researchers [ 4 , 29 , 30 , 7 , 9 ] tried to enable dynamic mask changes during the training process . All methods focused on unstructured sparsity and aimed to enable training large models on hardware with memory limitations . In contrast to these approaches , Zhou et al . [ 43 ] does not aim to reduce the model memory footprint during training , but rather accelerate inference while avoiding dense pre-training . To that end , they keep a dense weight matrix , and in each iteration they re-calculate the mask . With the obtained pruned copy of the weights they perform the forward and backward pass and update the dense copy . Therefore , this method is most relevant when a pretrained dense model is not given , and one wishes to obtain a fine grained N : M sparse model for inference . While our work focuses on the setting of Zhou et al . [ 43 ] , we argue that it can easily combined with Evci et al . [ 9 ] work , as it aims to solve a different issue within the same problem .

Summary: The paper proposes a method to quickly find masks that allow DNN models to train with N:M sparsity. Prior works mainly leverage N:M sparsity for inference (i.e., forward pass), but training cannot directly benefit from N:M sparsity because model weights are transposed during backpropagation to calculate gradients. Therefore, the N:M sparsity constraints may no longer be satisfied during backward pass. The paper suggests that it is possible to find so-called N:M transposable masks, which allow both the forward pass and the backward pass of a DNN training to perform N:M sparse computation. One key challenge in this task is to find such transposable masks in a lightweight way since sparsity masks may need to dynamically change in order to get on-par accuracy as dense training. To quickly find masks that satisfy the N:M transposable mask constraints, the paper proposes to use a greedy-based approximation method. The authors show that such a method adds a modest amount of runtime overhead and scales almost linearly with respect to the input sizes. The paper conducts evaluations against some computer vision (e.g., ResNet) and NLP tasks (e.g., BERT fine-tuning tasks) to show that such a training regime can bring comparable accuracy as dense training.