DominoSearch: Find layer-wise fine-grained N:M sparse schemes from dense neural networks

1 Introduction . Modern deep neural networks ( DNNs ) achieve remarkable success in various tasks such as computer vision and language processing at a cost of memory footprint , computation , energy and storage . These have all become bottlenecks of deploying DNNs on commodity hardware for real-world applications . To address these challenges , many model compression techniques have been proposed to optimize DNN models . Neural network pruning is one promising technique orthogonal to quantization [ 1–3 ] , network architecture search [ 4–6 ] and knowledge distillation [ 7–9 ] . Recent innovations in hardware architectures ( e.g . Nvidia Ampere Sparse Tensor Core [ 10 ] ) and N : M fine-grained sparse DNN pruning algorithms [ 11 , 12 ] reveal a promising research line of neural ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . pruning that may achieve a high sparse ratio while maintaining a regular sparse structure . On the other hand , traditional unstructured sparsity can achieve a higher sparse ratio but can not maintain a regular sparse structure , which makes it difficult for acceleration [ 13 ] . Traditional structured sparsity can achieve regular sparse structure but may not achieve high compression ratio [ 11 ] . Figure 1 illustrates traditional unstructured/structured pruning and N : M fine-grained structured pruning . Applying uniform fine-grained N : M sparsity training , or fine-tuning to DNNs , will lead to a suboptimal solution because layers are treated equally . Modern DNNs tend to have a number of layers with a highly non-uniform distribution of redundancy . This well-known phenomenon has been exploited by recent unstructured pruning methods [ 14–17 ] that have achieved higher accuracy than the uniform counterparts , especially for relatively high sparsity . However , the existing methods mainly focus on unstructured pruning that treats the weights within a layer individually and does not have fine-grained position constraints , which can not be used for fine-grained N : M sparsity . As a consequence , a novel method for layer-wise fine-grained N : M sparsity is needed . Furthermore , sparsity acceleration relies on the underlying software-architecture system of the platform . We can make a reasonable assumption that there will be more but not arbitrary N : M schemes supported in future software-architecture designs . Although the flexibility of N and M is still an open system design research question , we assume a policy of fixed M and flexible N in this paper because it can provide fixed window size ( i.e . size = M ) for sparse weight matrix compression . Motivated by above-mentioned challenges , this work introduces an efficient framework to find a layerwise fine-grained N : M scheme with pre-defined N : M candidates which achieves higher accuracy compared to the uniform counterpart with equivalent complexity constraint ( e.g . model size or FLOPs ) . We propose DominoSearch , an iterative algorithm that finds layer-wise fine-grained sparse schemes from pre-trained dense weights . This is achieved using a magnitude-based criterion [ 14 , 18 ] for the pruning selection threshold , in combination with a weight penalty to drive values towards the threshold . DominoSearch also uses a layer-wise penalty factor to balance the heuristic layer-wise redundancy of parameters [ 15 ] and layer-wise computational complexity ( e.g . FLOPs ) . As indicated by the name Domino , the N value from the N : M ratio decrements during the search phase with the weight penalty acting as momentum . Due to the layer-wise sparsity , our framework is able to achieve state-of-the-art ( SOTA ) accuracymodel size trade-offs when compared with unstructured sparsity and accuracy-FLOPs trade-offs when compared with structured sparsity . Figure 2 demonstrates the advantages of layer-wise N : M sparsity under both scenarios on ResNet50 [ 19 ] . Details can be found in Table 2 . The major contributions of our paper can be summarized as follows : ( 1 ) We present DominoSearch - A novel and easy-to-use framework to search the layer-wise sparse schemes from pre-trained dense weights with a specified model complexity constraint . Our method does not rely on expensive design space exploration such as Network Architecture Search ( NAS ) or Reinforcement Learning ( RL ) . To the best of our knowledge , this is the first work on layer-wise fine-grained N : M sparsity . ( 2 ) We propose a layer-wise penalty factor to balance the parameters redundancy that does not depend on input tensor/image size and complexity that depends on input tensor/image size ( e.g . FLOPs ) . The computational complexity measurement FLOPs used in this paper can easily be replaced with latency , throughput or energy consumption when deployed on a specific platform . ( 3 ) We evaluate our solution on the large scale ImageNet dataset with both heavy ResNets [ 19 ] and compact RegNets [ 20 ] , and the smaller scale CIFAR100 dataset with ResNet56 . Results of ResNet50 on ImageNet show that layer-wise N : M sparsity searched by our framework achieves SOTA accuracy-compression trade-off under fine-grained N : M sparsity . 2 Related Work . Structure of sparsity : Traditionally , sparsity of DNNs is categorized into 1 ) unstructured sparsity and 2 ) structured sparsity . Unstructured sparsity can be achieved by weight pruning and reaches a relatively high sparse ratio without significant accuracy drop [ 14 , 16 , 18 , 21 , 22 ] . However , unstructured sparsity has a large execution time overhead on modern hardware due to the large overhead of handling the irregular non-zero weights in the matrix ; good speed-up can only be achieved when the sparse ratio is very high ( i.e. , > 95 % ) [ 13 , 23 ] . On the other hand , coarse-grained structured sparsity such as filter-wise sparsity and channel-wise sparsity [ 23 , 24 ] can maintain a hardware-friendly sparse structure but may not be able to achieve a high sparse ratio . In order to achieve a high compression ratio and a hardware-friendly sparse structure , various fine-grained structured sparsity patterns were proposed including Tile-wise [ 25 ] , intragroup sparsity structure [ 26 ] and N : M sparsity [ 10 ] . Our paper focuses on N : M fine-grained sparsity which has demonstrated success in getting accurate N : M sparse DNNs ( i.e . trained from scratch [ 11 ] or fine-tuned [ 27 ] from dense pretrained weights ) and usage in commodity hardware ( e.g . Nvidia Ampere Tensor Core ) . Figure 1 illustrates two traditional types ( structured/unstructured ) and fine-grained N : M sparsity . Layer-wise compression : Layer-wise compression is a well-known concept in the DNN optimization community . The non-uniform importance of different layers in DNNs is leveraged by mixed-precision quantization in which different bit-widths are allocated via NAS [ 28 , 29 ] or RL [ 3 ] . Recent discoveries demonstrate that by carefully selecting [ 14 , 15 ] or learning [ 16 ] the layer-wise sparsity , the sparse DNNs can achieve higher accuracy than their uniform counterparts . Evci et al . [ 15 ] propose a training scheme for unstructured sparse neural networks and present the Erdos-Renyi kernel ( ERK ) ( extension of [ 17 ] ) to heuristically select the layer-wise sparsity . Lee et al . [ 14 ] propose a layeradaptive magnitude-based pruning score to choose the layer-wise sparsity , focusing on unstructured pruning and the number of survival weights of each layer without taking the layer-wise computation complexity ( e.g . FLOPs ) into account . Kusupati et al . [ 16 ] present Soft Threshold Reparameterization ( STR ) to learn the layer-wise sparsity during training . However , the target model complexity constraint ( e.g . model size and FLOPs ) is implicitly controlled by two hyper-parameters ( λ and sinit ) and can not be specified , hence users have to tune λ and sinit and train multiple models to get the desired target model complexity . Our paper focuses on how to find the N : M scheme of each layer efficiently with the specified model complexity constraint . It avoids the need for a time consuming hyper-parameter search . 3 Method - DominoSearch . 3.1 Problem Formulation . Given a pretrained dense model with weight tensors w1 , . . . , wi , . . . , wL associated with L layers ( convolution layers + fully-connected layers ) , model complexity constraint Spt ( e.g . model size or FLOPs ) and candidate pool Np of sparse schemes ( e.g . [ 1:8,2:8,4:8,8:8 ] ) 2 . The problem can be formulated as : max S ( w , N ) Accuracy ( S ( w , N ) ) s.t . Complexity ( S ( w , N ) ) ≤ Spt N i ∈ Np ( 1 ) Where S ( · ) denotes the sparse model which can be obtained with dense weights w and sparse ratio N . Complexity ( · ) can be model size , FLOPs , real latency or even energy consumption depending on the use case and design objectives . The goal is to find optimal layer-wise schemes for the whole network . Due to the large discrete design space and the training/fine-tuning cost of evaluating a design point ( i.e . recovering the accuracy of the sparse model ) , it is usually not possible to handcraft the optimal solution or to perform an exhaustive search . Inspired by methods used in mixed-bit search [ 3 , 29 ] , which decouple mixed-bit policy searching and model accuracy recovering , we use a two-step solution : ( 1 ) find the layer-wise sparse schemes from the candidate pool Np ; ( 2 ) fine-tune/retrain sparse weights with searched sparsity ( i.e . N i ) to recover the accuracy . 3.2 Group-wise threshold and sparsity . We first emphasise the difference between layer-wise fined-grained sparsity and layer-wise unstructured sparsity . How to decide the layer-wise sparsity has been actively studied over recent years [ 14–16 ] . However , these methods mainly focus on unstructured sparsity whose layer-wise sparsity can be decided by a layer-wise threshold because the weights in a certain layer are treated individually . In contrast , fine-grained N : M sparsity treats weights from a group-wise view ( e.g . Figure 1 ( c ) ) . As a consequence , we define a group-wise threshold vector T i and group-wise sparsity vector N ig , where i is the layer index . The length of these two vectors is K = |w i| M , where || is cardinality ( i.e . the number of parameters ) , M is the group size and K is the number of groups . |wi| should be divisible by M . Group-wise threshold selection : Instead of using a small global value such as 11000 [ 22 ] to decide if the weights should be treated as pruned or not , we use a threshold vector T i that can represent the weight distributions of each group . In order to achieve this , we use the mean of the smallest ( absolute value ) M2 elements as threshold because it provides a group-wise feature in an efficient way and can mitigate the effect of outliers ( e.g . values like 0.0 ) . Note that we use M2 elements because it results in a sufficient search space and representativeness . Using more elements can result in a better representativeness because more elements are included , but results in a smaller search space because a larger threshold value will lead to a smaller starting value N ig and vice versa . Figure 3 shows an example of group-wise threshold selection . The larger variance of T i indicates that we need a group-wise threshold . On the other hand , a group-wise threshold will lead to group-wise sparsity . The column vector N ig encodes the number of kept weights of each group , from which we can infer the sparsity of this layer ( i.e . N : M = 3:4 ) . In practice , we can relax the requirement by using the Supermajority voting rule . Specifically , as long as the ratio of the number of x in N ig reaches a pre-defined voting ratio vr , the layer sparsity N i can be inferred as x . With above-mentioned group-wise definitions , we can reformulate wi and T i as wik and T ik , where k is the group index . We use the reformulated version in the next sections . 2Without losing generalization , we use N equal to power-of-two as search space .

Summary: The authors propose combining two existing techniques, N:M sparsity and per-layer (rather than uniform/global) pruning rates.  By allowing each layer to choose a particular N:M pair, the network can find a configuration that is not constrained to sharing a single N:M.  The solution is DominoSearch, an algorithm which starts with pretrained (dense) weights and determines a pruning threshold for each group of M weights based on the initial values.  Some number of iterations of fine-tuning adjust the weights gradually lower, in turn reducing each layer's chosen N.  Once some complexity target (either parameter or FLOP count) is hit, the network's layer-wise N:M are frozen, and the weights are finally fine-tuned.  A layer-wise penalty factor adjusts the overall complexity metric to account for layers which have relatively more or fewer FLOPs.  Empirical results show compelling reductions in FLOPs and parameters when compared to training with unstructured sparsity or uniform N:M sparsity.