Effective Model Sparsification by Scheduled Grow-and-Prune Methods

Deep neural networks ( DNNs ) are effective in solving many real-world problems . Larger DNN models usually exhibit better quality ( e.g. , accuracy ) but their excessive computation results in long inference time . Model sparsification can reduce the computation and memory cost while maintaining model quality . Most existing sparsification algorithms unidirectionally remove weights , while others randomly or greedily explore a small subset of weights in each layer for pruning . The limitations of these algorithms reduce the level of achievable sparsity . In addition , many algorithms still require pre-trained dense models and thus suffer from large memory footprint . In this paper , we propose a novel scheduled grow-and-prune ( GaP ) methodology without having to pre-train a dense model . It addresses the shortcomings of the previous works by repeatedly growing a subset of layers to dense and then pruning them back to sparse after some training . Experiments show that the models pruned using the proposed methods match or beat the quality of the highly optimized dense models at 80 % sparsity on a variety of tasks , such as image classification , objective detection , 3D object part segmentation , and translation . They also outperform other state-of-the-art ( SOTA ) methods for model sparsification . As an example , a 90 % non-uniform sparse ResNet-50 model obtained via GaP achieves 77.9 % top-1 accuracy on ImageNet , improving the previous SOTA results by 1.5 % . All code will be publicly released . 1 INTRODUCTION . Deep neural networks ( DNNs ) have achieved great performance in many real-world scenarios . However , the large computation and memory requirements of deep neural networks discourage them from being applied to broader applications on resource-limited devices . Model compression via weight pruning is a popular research topic , where a large proportion of weights are set to zero , leading to significant reduction in both memory and computation . The introduction of sparse tensor cores in the NVIDIA A100 GPU ( NVIDIA , 2020b ) brings weight pruning into mainstream . Early works on weight pruning generally follow a prune-from-dense methodology ( Guo et al. , 2016 ; Yu et al. , 2018 ; Wen et al. , 2016 ; Liu et al. , 2019b ; Zoph & Le , 2017 ; Bender et al. , 2018 ; Pham et al. , 2018 ; Real et al. , 2019 ) , which usually requires 3 phases of training : pre-train a dense model , prune it to sparse , and fine-tune it . In such methodologies , the one-shot or iteratively pruning from a well-trained DNN can only remove weights , which lacks the flexibility of growing back weights that are considered unimportant early in the training process but showed to be significant later in training . On the other hand , early-stage pruning methods ( Lee et al. , 2019 ; Wang et al. , 2020 ; You et al. , 2020 ) avoid training dense models to converge . However , those sparse masks ( i.e. , the binary matrix in which a zero value removes the corresponding weight entry from the model ) are prematurely fixed , resulting in inferior model quality . Methods based on sparse mask exploration , such as DeepR ( Bellec et al. , 2018 ) , SET ( Mocanu et al. , 2018 ) , DSR ( Mostafa & Wang , 2019 ) , and RigL ( Evci et al. , 2020 ) maintain the target sparsity in all layers throughout the training process and selectively explore a small fraction of the weights periodically . The explored weights are chosen based on random or greedy heuristics , leading to limited coverage of sparse patterns and consequentially sub-optimal model quality . In order to overcome the shortcomings of the previous solutions , we propose a scheduled grow-andprune ( GaP ) methodology . It does not require a dense model any time during the training process , leading to largely-reduced memory footprints . Meanwhile , the sparse mask of a layer is updated after exploring all weights in the same layer , resulting in better mask-update efficiency . It targets to obtain the best quality sparse model under a runtime budget . Since some large models may be compressed below a target latency before being deployed in the wild to thousands of devices and inferenced billions of times per day ( Hazelwood et al. , 2018 ; Wu et al. , 2019 ) , spending a bit more time in training to find a better model is well rewarded . The scheduled GaP methodology divides a DNN into several partitions composed of contiguous layers . During model training , one of the sparse partitions is grown to dense while the rest remain sparse . After some epochs of training , the previously dense partition is pruned to sparse , and an alternate partition is grown to dense . This process is repeated so that all partitions are grown to dense and pruned back to sparse multiple times . If the partitions are grown to dense one by one sequentially , it is called the cyclic GaP method ( as shown in Figure 1 ) . If the model is replicated to multiple machines with different dense partitions , it is called the parallel GaP method . The GaP methods are carefully scheduled to explore all weights in a model efficiently , which is not guaranteed in the existing mask exploration methods . We illustrate their differences by an example in Figure 2 . In the scheduled GaP methods , all weights to be grown ( or to be pruned ) belong to the same partition . All weights in the model are explored when every partition is grown to dense and pruned to sparse . On the other hand , in the random or greedy exploration methods , the grown and pruned weights are distributed across all layers and they do not guarantee that all weights are explored , e.g , the connection between 1st input neuron and 3rd neuron in the first layer is not explored in Figure 2 ( b ) . Random sparse mask exploration and our scheduled GaP methods are like sampling operations with replacement and without replacement ( Basu , 1958 ) . The former requires a lot more samples to achieve the same weight coverage as the latter . Please see Appendix C for an illustrative example . The contribution of this paper is summarized as follows . • We propose two GaP methods such that they do not require to train dense models any time throughout the training process , reducing the training memory footprints . • Sparse models obtained by the proposed GaP methods outperform the SOTA sparsification algorithms on a wide variety of deep learning tasks , including 2D image recognition , object detection , 3D object part segmentation , and machine translation . • For all four tasks , the proposed GaP methods surprisingly generate sparse models that match or even exceed the dense counterparts at 80 % sparsity . • We provide theoretical guarantees for the GaP method and clarify its convergence property . In particular , we justify the influence of the mask-induced error . Algorithm 1 : C-GaP training flow . Input : An L-layer model with uninitialized weight Θ ; pruning ratio r. Output : An L-layer sparse model satisfying the target sparsity requirement . 1 Initialize f ( x ; Θ m ) with random weights Θ and random masks m that satisfy the sparsity requirement . 2 Divide all layers into κ partitions , denoted by Si , i ∈ { 0 , 1 , . . . , κ− 1 } . 3 step = 0 4 while step < K do 5 Partition to grow index i← step mod κ , partition to prune index j ← ( step− 1 ) mod κ . 6 Prune ΘSj and update mSj by [ ΘSj , mSj ] ← ArgPruneTo ( ΘSj , r ) . 7 Grow ΘSi by mSi ← ArgGrowTo ( ΘSi ) . 8 Train the model f ( x ; Θ m ) for T epochs . 9 step = step+ 1 . 10 Denote the final dense partition by Sd . 11 Prune ΘSd and update mSd by [ ΘSd , mSd ] ← ArgPruneTo ( ΘSd , r ) . 12 Fine-tune f ( x ; Θ m ) for T ′ epochs . 2 METHODOLOGY . In this section , we describe the scheduled grow-and-prune ( GaP ) methods in detail . The process starts from a randomly initialized model . First , its weights are pruned randomly to reach the target sparsity , i.e . a random sparse mask is applied to the weights . Then , the model is divided into κ partitions and each partition is grown and pruned separately . In the following two sub-sections , two variants of the GaP methods are described in detail . 2.1 CYCLIC GAP . As shown in Figure 1 , the cyclic GaP ( C-GaP ) method rotates the dense partitions among all κ partitions ( κ = 4 in Figure 1 ) . Starting from all partitions with random sparse masks , the first partition is grown to a dense one . All weights in the first partition and the masked weights in the remaining 3 partitions are trained for a few epochs ( Step 1 ) . Then , the first partition is pruned to sparse ( Step 2 ) . Next , we apply the same strategy to the second partition , then continue iterating over all κ partitions . When all layers are cyclically grown and pruned once after κ steps , we call it one round . This process is repeated K times . Algorithm 1 describes this process in more detail . In Line 1 , we use f ( x ; Θ ) to represent an L-layer deep neural network model , where Θ ∈ RM are M trainable weights and x are training samples . A sparse model is denoted as f ( x ; Θ m ) , where denotes element-wise multiplication . m is a binary mask m ∈ { 0 , 1 } M , where a zero in m means that the corresponding weight in Θ is fixed to be zero and a one in m means that the corresponding weight in Θ is free to be updated . In the C-GaP , all layers of the model { 1 , 2 , . . . , L } are divided into κ subsets S0 , S1 , . . . , Sκ−1 , where κ is an integer between 1 and L , as shown in Line 2 . Let ΘS and mS denote the subset of trainable weights and their masks . Consequently , Θ is divided into κ partitions , denoted by ΘS0 , . . . , ΘSκ−1 . For the i-th partition ΘSi , we denote its mask by mSi . Then the grow and prune process is iterated K times ( Line 4 ) . First , the indices of the partitions to grow and prune are calculated ( Line 5 ) . They are determined in a cyclic order . Then , the selected dense partition is pruned to sparse ( Line 6 ) . It prunes the partition ΘSj with a pruning ratio r. In the mean time , the next sparse partition is grown to dense ( Line 7 ) . It sets the mask mSi to be all ones such that all weights ΘSi are free to be updated during training . After that , the model is trained for T epochs ( Line 8 ) . After the entire grow-and-prune process is finished , the final dense partition is pruned to sparse ( Line 11 ) and the model is fine-tuned to get the final trained sparse model ( Line 12 ) .

Summary: Note: I have previously reviewed this paper for NeurIPS 2021 and I will be reusing my review from then and add the additional comments along the way based on how the paper has been updated along with the promised updates from the previous revisions.   This paper proposes a way to obtain sparse neural networks by repeatedly growing and pruning parts of the model in a scheduled fashion. Figure 1 gives a good idea of the whole technique. It basically divides the whole network into K parts and only one of these K parts is dense while the rest are sparse. The dense part keeps cycling ensuring over-parameterized learning when needed. The paper proposes that these scheduled routines help in improving the accuracy of the sparse neural network models compared to other sparsity-inducing techniques.  The paper has extensive experimental evaluation on tasks like classification, object detection, 3D object part segmentation, and Transformers. The authors also try to provide some theoretical explanations to these proposed ideas.  The paper has updated visualizations to make the algorithm much more clearer from the earlier version which is commendable.