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: This study takes use of pruning-rewiring-pruning method to get a high unstructured sparse model. Broadly, this is one dropout regularization method. With step-by-step growing and pruning, most of the weights are close to $0$ and will be pruned in the next step. Experiments show that this training method can improve the sparsity and performance for a wide range of architectures and applications,