E FFECTIVE M ODEL S PARSIFICATION BY S CHEDULED G ROW - AND -P RUNE M ETHODS
A BSTRACT
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%.
Code available at: https://github.com/boone891214/GaP .
I NTRODUCTION
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; 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 [Figure] 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 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.
[Figure]
M ETHODOLOGY
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. C YCLIC G A P
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 Θ ∈ R M 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. [Figure]
In the C -GaP, all layers of the model { 1 , 2 , . . .
, L } are divided into κ subsets S 0 , S 1 , . . .
, S κ − 1 , where κ is an integer between 1 and L , as shown in Line 2.
Let Θ S and m S denote the subset of trainable weights and their masks.
Consequently, Θ is divided into κ partitions, denoted by Θ S 0 , . . .
, Θ S κ − 1 .
For the i -th partition Θ S i , we denote its mask by m S i .
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 Θ S j with a pruning ratio r .
In the mean time, the next sparse partition is grown to dense (Line 7).
It sets the mask m S
i to be all ones such that all weights Θ S
i 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).
2. P ARALLEL G A P
In this section, we introduce the parallel GaP ( P -GaP) as a flexible and highly efficient complement to the C -GaP sparse training method when sufficient training resources are available, as described in
Figure 3.
Unlike the C -GaP method where partitions are grown and pruned serially in one machine, the P -GaP method distributes the model to κ machines.
Each grows a different partition to dense and trains the model in parallel.
Then, the trained dense partitions are combined to a dense model and pruned to a sparse one again.
As a result, the P -GaP method utilizes κ more machines, but its training time is also κ times shorter than the C -GaP method.
Since each server trains the model independently, the inter-server communication only occurs when the dense parts are combined to synchronize their learned weights once every few epochs.
Thus its impact to performance is negligible.
Each pair of distribution and combining form one round in the P -GaP. Unlike C -GaP, one round of P -GaP consists of only one step.
The detailed P -GaP method is shown in Algorithm 2 and illustrated in Figure 4 (b) in Appendix D.
Some of the differences between the parallel GaP and the conventional distributed training are listed as follows.
In P -GaP, data communicates across different server nodes every T epochs of training.
It is far less frequent than conventional distributed training, which requires communications and synchronization in every mini-batch.
Therefore, the data bandwidth between the nodes remains minimal.
In P -GaP, each node can use the single-node hyper-parameters to train the sparse models.
It does not require excessively large batch-size to fully utilize the computing resource.
Different from data parallelism, model parallelism, and pipeline parallelism, the P -GaP explores another dimension of parallelism when training sparse models, which we call it mask parallelism .
This implies that masks in different partitions are less correlated and can be searched separately.
[Table]
E XPERIMENTAL R ESULTS
This section evaluates the proposed C -GaP and P -GaP sparse training methods on different machine learning tasks.
Since our goal is to search for the optimal sparse models to speed up inference, we report the best model accuracy obtained using the GaP methods for all models.
We compare our
GaP method with the state-of-the-arts, as well as our implemented prune-from-dense results.
If not otherwise specified, all prune-from-dense are implemented using magnitude-based iterative pruning schedule as described in GMP (Zhu & Gupta, 2017; Gale et al., 2019).
More ablation experiments on partition number, partition scheme, different sparsity granularity using GaP, and FC layer pruning are shown in Appendix A.
All of our experimental results are trained and inferenced using PyTorch in the machines with
NVIDIA-V100 GPUs.
The cyclic GaP is trained in one training node and the parallel GaP is trained in κ training nodes, where κ is the number of model partitions.
In the pruning stage of the GaP method, the weights with the lowest magnitudes are pruned to zero.
Note that in our experiments, only the weights for convolution, matrix-vector, and matrix-matrix multiplications are made sparse.
The weights related to biases and batch-normalization are kept dense, as they are critical for model quality but only contribute to a tiny fraction of the inference FLOPs.
For uniform sparsity, we prune individual layer separately by sorting all weights in the layer based on their magnitudes, and prune away any weight whose magnitude is below the percentile described as the sparsity level.
Thus, all layers receive the same sparsity.
For non-uniform sparsity, the pruning process is similar, but the weights in the entire model are sorted together.
Thus, the sparsity level in different layers may differ, with the more important layers pruned less than the trivial layers.
We list the implementation details and hyper-parameter settings for all tasks in Appendix B.
3. I MAGE CLASSIFICATION : R ES N ET -50 ON I MAGE N ET
Table 1 compares the accuracy between using the GaP methods and the previous works.
The ResNet50 is divided into four partitions by the last three downsampling stages.
For non-uniform sparsity, all layers are sparsified.
For the uniform sparsity, the first convolutional layer with 7 × 7 kernels is kept dense, the same as in Evci et al.
The fully connected (FC) layer only contributes 0.5% of
FLOPs, thus it is also kept dense (sparse FC layer has similar accuracy with dense FC layer, please check the ablation results in Appendix A).
To ensure fair comparison, we perform the following demonstrations in addition to the original baseline results.
We include the SNIP (Lee et al., 2019) and SET (Mocanu et al., 2018) results in uniform sparsity that are reproduced in Evci et al.
We also implement the RigL and
RigL 5 × (Evci et al., 2020) using our code base on PyTorch (please refer to Appendix B.2 for details).
In Table 1, our reproduction achieves equivalent or higher accuracy than the one reported in the paper.
Based on that, we extend the RigL training time by 12 × to match the GaP method.
Additionally, SET shares similar method with RigL, thus we also implement and extend SET training by 12 × .
We observe that the improvement over RigL 5 × (Evci et al., 2020) is 1.3% (77.9% vs. 76.6%) at uniform 80% sparsity, and 0.6% at 90% uniform sparsity.
The C -GaP is also better than our reimplementation of the prune-from-dense method using ADMM (Ren et al., 2019).
For non-uniform sparse ResNet-50 model, the improvement over RigL 5 × (Evci et al., 2020) is 1.0% and 1.5% at 80% and 90% sparsity, respectively.
Note that the 80% sparse ResNet-50 model can almost match the accuracy of the dense model reported by NVIDIA (2020a).
We also observe that, when we extend the training time for the SOTA methods such as SET and RigL to match the GaP training efforts, their accuracies still lag behind.
As an example, the model at 80% non-uniform sparsity trained using the
C -GaP method outperforms the re-implementation of the ERK RigL 12 × model with 0.7% accuracy increase and 21% FLOPS reduction (2.7 GFLOPS vs. 3.4 GFLOPS).
The 90% non-uniform sparsity model achieves 1.1% accuracy increase at the same FLOPS.
Comparing the C -GaP and P -GaP pruning methods with same training epochs, we find that CGaP achieves 0.4% and 0.2% higher accuracy at 80% and 90% uniform sparsity, respectively.
We conjecture the reason is that P -GaP converges slower than C -GaP.
In C -GaP, the weight values in the sparse partitions are continuously trained to get better model accuracy when the dense partition is explored, while the P -GaP method only keeps the dense partition and discards the sparse partitions when combining all dense partitions to the updated model.
3. O BJECT DETECTION : SSD ON COCO-
We divide the SSD network into 4 partitions with three in the backbone and one in the detection head.
We train the model using cyclic-partitioned C -GaP with a uniform weight sparsity of 90% on all layers except for the last output layer, which is kept dense.
We report the accuracy of the best model searched within 40 GaP steps.
As shown in Table 2, in the mean average precision category mAP@[.5:.95], the sparse model obtained using our C -GaP method exceeds the dense model by 0.
mAP (25.9 vs 25.2).
It also exceeds the best sparse model iteratively pruned from the dense model by 1.6 mAP (25.9 vs 24.3).
[Table]
3. 3D OBJECT PART SEGMENTATION : P OINT N ET ++ ON S HAPE N ET
We divide the PointNet++ model into 4 partitions with three in the backbone and one in the segmentation head.
We apply the C -GaP and P -GaP methods with uniform sparsity of 80% and 90% on all layers, respectively.
We report the accuracy of the best model searched within 40 GaP steps.
Tablecompares them with the dense model and the best sparse model pruned from the dense model using the ADMM algorithm.
The results show that on the class and instance mAP categories, the pruned model using the C -GaP method easily beats the model pruned from dense at both sparsities.
It even beats the dense model at 80% sparsity.
The pruned model using the P -GaP method also beats the model pruned from dense at 80% sparsity and is not far behind at 90% sparsity.
[Table]
3. T RANSLATION : T RANSFORMER ON WMT-14 E N -D E DATASET
In this part, we evaluate the C -GaP and the P -GaP methods on the translation task based on Transformer (Vaswani et al., 2017).
We train the Transformer models on the WMT-14 En-De dataset and evaluate the SacreBLEU score (Post, 2018).
We equally divide the Transformer models into three or six partitions with the decoder containing 2 × partitions of the encoder.
We apply the C -GaP and P -GaP methods with uniform sparsity of 80% and 90% on all layers, respectively.
We report the best SacreBLEU scores on the validation dataset in Table 4 within 30 GaP steps.
The models trained using both C -GaP and P -GaP methods significantly improve over the sparse models obtained by pruning from the dense counterparts.
They exceed the dense model quality at 80% sparsity, and the model using three-partition C -GaP method even outperforms the dense model at 90% sparsity.
[Table]
T HEORETICAL J USTIFICATION
We now provide the convergence guarantee for the cyclic-partitioned C -GaP algorithm.
We let
F (Θ) =
E x ∼ D f ( x ; Θ) be the loss function of the deep learning task where D is the data distribution.
In addition, we let m q be the mask at the start of the q -th round, Θ q be the learned model weights after q − 1 rounds, and Θ ( i ) q , i = 1 , · · · , κ be the learned weights in the q -th round after the i -th partition is grown and trained for T epochs.
Proposition 1 shows that C -GaP converges to a neighborhood around the stationary solution at rate O (1 / √ Q ) when learning rate is set properly.
Due to the space limitation, we put its proof in Appendix E.
Proposition 1 ( C -G A P CONVERGENCE ) .
Suppose the loss function F (Θ) is partition-wise Lsmooth, the sampled stochastic gradient is unbiased and has bounded variance, and the relative error introduced by each mask is bounded by δ 2 ∈ [0 , 1) , i.e., ∥ Θ q − Θ q ⊙ m q ∥ 2 ≤ δ 2 ∥ Θ q ∥ 2 .
If the learning rate γ = 1 / (4 κLT √ Q ) , the sparse models generated by the C-GaP method after Q rounds will converge as follows:
[Equation]
where G is a constant related to the gradient noise and the initialized model weights.
We make several remarks for Proposition 1 as follows.
Remark 1.
If there is no pruning in the training process, then it holds that δ 2 = 0 .
Substituting it into equation 1, we find that C -GaP will converge exactly to a stationary solution, i.e., E ∥∇ F (Θ Q ) ∥ 2 →as Q increases to infinity.
This implies C -GaP is also an effective algorithm even for dense training.
Remark 2.
When a sparse mask is utilized, it holds that δ 2 ̸ = 0 .
In this scenario, the mask-induced error will inevitably influence the accuracy of the trained model.
A smaller mask-induced error will lead to a more accurate model that C -GaP can converge to.
Remark 3.
In our proofs, we find that cycling through all partitions so that all weights can be trained per round is critical to guarantee the convergence of C -GaP.
Note that previous works update weights either greedily (e.g., RigL
(Evci et al., 2020)) or randomly (e.g., SET (Mocanu et al., 2018) and DSR (Mostafa & Wang, 2019)).
It may take numerous training steps for these algorithms to have each weight explored and updated.
This may explain why C -GaP achieves better accuracy than RigL,
SET, and DSR (see Table 1).
This intuition is consistent with (G¨urb¨uzbalaban et al., 2019;
Ying et al., 2018) which theoretically establish that sampling data cyclically converges faster than uniformly randomly in SGD training.
R ELATED W ORKS
Pruning from pre-trained dense models : Among various pruning algorithms, one-shot pruning (LeCun et al., 1990; Xu et al., 2021) zeros out a given percentage of the trained weights with the smallest magnitude.
Based on the magnitude-based one-shot pruning, the gradual magnitude pruning (GMP) (Zhu & Gupta, 2017; Gale et al., 2019) proposes an iterative pruning method that progressively prune models to the target pruning ratio.
However, those methods suffer from significant accuracy drop and relatively long training time.
Besides magnitude-based pruning, other approaches (Wen et al., 2016; He et al., 2017; Zhang et al., 2021a) adapt mathematics-oriented regularization-based algorithms to generate sparsity.
Later works (Zhang et al., 2018b; Ren et al., 2019; Ma et al., 2020a;b; Niu et al., 2020; 2021; Rumi et al., 2020; Ma et al., 2021a; Zhang et al., 2021b; Huang et al., 2021) utilizes dynamic regularization penalty such as ADMM to solve the pruning problem as an optimization problem and maintaining high accuracy.
Several methods (Zhou et al., 2021; Srinivas et al., 2017; Molchanov et al., 2017) work in probability space.
Other works such as HRank (Lin et al., 2020a),
SCOP (Tang et al., 2020), DMCP (Guo et al., 2020), MetaPruning (Liu et al., 2019a) use complicated rules to generate the sparsity distribution in the channel level.
Since pre-training a dense model from which to select critical weights consumes large memory space, our work differs substantially from them such that we do not rely on a pre-trained dense model.
Pruning at an early stage : A new trend of exploring sparsity at an early stage (Lee et al., 2019;
Wang et al., 2020; Tanaka et al., 2020; Wimmer et al., 2020; van Amersfoort et al., 2020; Ma et al., 2021b; Liu et al., 2021) has emerged to embrace the promising sparse training paradigm.
SNIP (Lee et al., 2019) finds the sparse masks based on the saliency score of each weight that is obtained after training the dense model for only a few iterations.
GraSP (Wang et al., 2020) prunes weights based on preserving the gradient flow in the model.
EarlyBird
(You et al., 2020) conducts a retraining process after searching sub-network topology for a few epochs.
PruneTrain (Lym et al., 2019) integrates structured pruning in the pretraining process.
However, pruning at an early stage fails to achieve acceptable accuracy on ImageNet (Deng et al., 2009).
We recognize the underlying reason is that the single-shot pruning criterion is based on the information from a few mini-batch of training samples, which cannot accurately represent the rich features of large-scale dataset.
The proposed GaP methods repeatedly update the sparse masks and hence overcome their drawbacks.
Sparse mask exploration : To satisfy the condition of low computation cost as well as low memory footprint, many works have incorporated the training-from-sparse-initialization paradigm, and optimize the network by sparse mask exploration.
DeepR (Bellec et al., 2018) proposes to train a sparse network at initialization, and dynamically adjust the sparse topology by pruning the weights that flip the sign and randomly growing new weights during training.
However, DeepR is primarily demonstrated with sparsification of fully-connected layers and applied to relatively small and shallow networks.
Sparse Evolutionary Training (SET) (Mocanu et al., 2018) and MEST (Yuan et al., 2021a) uses layer-wise magnitude-based pruning and random growth at the end of each training epoch.
DSR (Mostafa & Wang, 2019) and STR (Kusupati et al., 2020) designs a dynamic reparameterization method that allows weights to be re-distributed across layers by providing an effective and trainable way to allocate the sparsity across all layers.
Similarly, DNW (Wortsman et al., 2019) trains a sparse model and keeps dense gradients to dynamically adjust the network connections.
SNFS (Dettmers &
Zettlemoyer, 2019) develops sparse momentum to identify the importance of weights.
RigL
(Evci et al., 2020) and NeST (Dai et al., 2019) propose to use magnitude-based pruning and gradientflow-based growth that update sparse model topology during training.
Yuan et al.
(2021b) uses a continuous relaxation and optimization method to dynamically update the structural masks of DNN models.
All of the mentioned works update the sparse masks based on heuristics and there is little or no training at the mask update stage.
Our work differs from them in that the mask updating rule of the GaP methods is more optimized by training a subset of layers into dense and then pruned.
S OCIETAL I MPACT AND L IMITATIONS
Societal Impact.
The scheduled GaP methods target to obtain the best quality of the pruned model under a runtime budget.
Based on the researches in Facebook (Hazelwood et al., 2018; Wu et al., 2019), many models are inferenced many billions of times per day, which makes the cost of pruning such models a tiny fraction of the inference cost in one day.
Thus, a better pruning algorithm resulting to a smaller model may save a lot more computation on the inference side, which is of high value for the community and society to achieve Green AI (Schwartz et al., 2020).
On the other hand, our method cannot prevent the possible malicious usage, which may cause negative societal impact.
Limitations.
In this work, we focus on the unstructured sparsity to demonstrate the algorithm-level innovations of our proposed GaP methods, and we only preliminarily include a block sparsity example in Table 9 of Appendix A.3 to demonstrate its applicability to structured sparsity and its potential to improve the inference speed.
More detailed analysis is still a work in progress.
The scheduled GaP methods require more training time than conventional pruning methodology, which may potentially limit its applicability when the training compute resources are limited.
Thus, the scheduled GaP methods are mostly beneficial to models whose benefit of the reduced inference time and/or improved accuracy out weigh the cost of the moderately longer training time.
The key hyper-parameter settings such as partition numbers or partition strategies are determined heuristically.
In this work, we try to divide the models to partitions with similar parameter counts or computation.
We also intuitively group consecutive layers to the same partition, based on our hypothesis that adjacent layers are tighter correlated than more distant layers.
Our analysis in Proposition 1 cannot be directly extended to cover P -GaP since the analysis is built on the partition-wise cyclic updating structure of C
-GaP. We will leave the analysis for P -GaP as a future work.
The prediction results from the sparse models are generally biased comparing with the dense models they are pruned from (Hooker et al., 2019).
This is a general problem on sparse model architectures, and our GaP methodology cannot reduce such biases.
C ONCLUSIONS
Sparsity based model compression is gaining traction in research and industry to reduce the memory footprint and inference runtime.
However, conventional compression approaches result to noticeable accuracy loss during the pruning process.
In this paper, we propose a scheduled grow-and-prune (GaP) methodology to obtain sparse models.
It divides a model into several partitions, grows and prunes each partition cyclically or in parallel.
The experimental results show that this method preserves the model quality far better than previous works.
In all four experimented tasks, the GaP method matches or even beats the dense solution at 80% sparsity.
In addition, the GaP method does not require a pre-trained dense model and can continuously searching sparsity masks with new data, showing its flexibility and practical potential.  A CKNOWLEDGEMENT
This work was supported by Alibaba Group through Alibaba Research Intern Program, and was partly supported by the National Science Foundation CCF-1937500.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of NSF.
A A BLATION S TUDY
In this section
[Table]
we present some ablation study on the proposed GaP methods.
A. N UMBER OF PARTITIONS FOR G A P
The number of partitions in GaP methods play a key role in obtaining better results.
As compared in
Table 5 and Table 6, one partition is not sufficient to preserve the best quality.
Too many partitions may reduce quality as well (as shown in the model with six partitions in Table 6).
Empirically, a small number of equal partitions that naturally follow the structures of the models usually perform well.
For example, ResNet-50 model is composed of four blocks, each sub-samples from the previous block.
Thus a four-partition GaP is a natural choice.
Transformer model contains an encoder and a decoder, with the decoder twice the number of attention blocks as the encoder.
It is naturally divided to three partitions.
[Table] [Table]
A. P ARTITION THE MODEL IN A RANDOM MANNER
In the main text, each partition of the model consists of contiguous layers and they are cyclically grown and pruned.
We also explore a random partition strategy for each step in the C -GaP for sparse ResNet-50 and Transformer models.
Unlike the cyclic partition, a random partition chooses a uniformly random subset of layers to grow in each step of the GaP process.
On the results for
ResNet-50 described in Table 7, the random partition is slightly worse than the cyclic partition but still better than all previous existing solutions.
On the results for Transformer described in Table 8, however, the accuracy gap is much larger.
At 90% sparsity, Transformer model can achieve 27.
SacreBLEU score in uniform sparsity with 3 partitions using C-GaP. When a random 3-partition is applied to a 90% sparsity Transformer, the SacreBLEU score is reduced to 27.0 and 26.9 for uniform and non-uniform sparsity, respectively.
We conjecture that different model structures may result to different accuracy gaps.
For ResNet-50, the correlations between neighboring layers may be small so each layer may be optimized independently.
However, for Transformer, the inter-layer weight correlation is stronger (e.g., the W key , W query , W value in the same attention block have very strong correlation since they process the same input sequence for the self-attention computation)
This analysis shows that the partition mechanism affects the final model accuracy.
Our scheduled
GaP methods minimize the chance that the neighboring layers belong to different partitions, and have achieved superior accuracy.
[Table] [Table]
A. G ENERALIZATION TO OTHER SPARSE GRANULARITY
We use unstructured sparsity throughout the paper.
In this section, we demonstrate that the GaP method is also applicable to other coarse-grained sparse granularity such as block-wise sparsity.
The unit for the block-wise sparsity in this experiment is 1 × 8 , i.e., if the weight matrix is viewed as many 1 × 8 blocks, then the sparsity ratio of the matrix is defined as the ratio of all-zero blocks.
It is often believed that the accuracy would be impacted severely when the block size is 8 or greater.
Table 9 compares the SacreBleu scores of the sparse Transformer models obtained by applying the 3-partition cyclic C -GaP and the prune-from-dense methods.
It is observed that the GaP method improves over the prune-from-dense method significantly (27.21 vs. 26.21).
This further validates the benefits of mask updating strategy brought by the GaP method.
[Table]
A. P RUNING THE FULLY CONNECTED LAYER IN THE R ES N ET -50 MODEL
In the main context, the last fully connected (FC) layer of the ResNet-50 model is kept dense in the uniform sparsity distribution since it contributes to less than 1% of the FLOPs.
In Table 10, we perform additional experiments to supplement Table 1 by pruning the last FC layer using the C -GaP method and comparing them with Evci et al.
The results in Table 10 and Table 1 indicate that the accuracy of the ResNet-50 models with sparse and dense FC layers are similar, and both of them outperform the state-of-the-art results in Evci et al.
Please note that models with the non-uniform sparsifying distribution in Table 1 already have the last FC layer pruned, thus the experiment setup is the same as the ones in Evci et al.
[Table]
B E XPERIMENT S ETUP
Most training scripts and hyper-parameters are inherited from NVIDIA (2020a) and Yan (2019).
Identical hyper-parameters are applied to each GaP step and the final fine-tune stage.
For example, the initial learning rate of each step is restored to the same value.
This hyper-parameter restoration avoids pre-determining the total number of epochs before training and enables the GaP methods to continuously search for sparse masks.
B. I MAGE CLASSIFICATION : R ES N ET -50 ON I MAGE N ET
In the image classification task, we use standard data augmentation, a batch size of 2048, cosine annealing learning rate schedule, SGD optimizer with a momentum of 0.875, and a weight decay of 3.05e-5.
The learning rate is scheduled with a linear warm-up for 2 epochs before reaching the initial learning rate of 2.048.
Each GaP step with non-uniform and uniform sparsity includes 30 of training, respectively.
The final fine-tuning includes 150 epochs.
For C -GaP, we train for 28 steps (i.e., 7 rounds with 4 partitions), and for P -GaP, we train for 32 steps (i.e., 8 rounds with 4 partitions).
After the GaP step, we prune the dense partition(s) and fine-tune the model.
For the prune-from-dense method, we use iterative ADMM pruning.
We first pretrain a dense model using 250 epochs, and perform ADMM regularization training for 250 epochs.
Then we prune the model to 80% sparsity and finetune for another 250 epochs.
Thus, the total number of epochs for 80% sparsity is 750 epochs.
Starting from the 80% sparse model, we iteratively prune to 90% sparsity by firstly training the 80% sparse model with ADMM regularization for 250 epochs, and then prune to 90% sparsity and finetune 250 epochs.
It takes a total of 1250 epochs to train a 90% sparsity model.
B. O UR I MPLEMENTATION D ETAILS OF R IG L
In this section, we provide the implementation details of the RigL (Evci et al., 2020).
For the experiments with 100 training epochs, Most of the important hyper-paramters can be found in Evci et al.
We adopt the settings in our PyTorch code base and the accuracy of RigL in 80% and 90% sparsity can be reproduced, both in uniform and ERK sparsity distributions, respectively.
As shown in Table 1, our implementation has slightly higher accuracy, which indicates that our code base in PyTorch is comparable with the original implementation in Evci et al.
(2020), and can be used for reproducing and extending the training of RigL without causing fairness issue.
Based on the above implementation, we reproduce the RigL results with 500 training epochs.
However,
Evci et al.
(2020) doesn’t provide detailed hyper-parameter settings for 500-epoch training.
We perform extensive experiments and find out that using the same step learning rate decay scheduler can not reproduce the accuracy claimed in the original RigL paper, regardless of the learning rate decay factor or decay epochs.
We change the step learning rate decay scheduler to the cosine annealing learning rate scheduler (Loshchilov & Hutter, 2017), and add mixup technique (Zhang et al., 2018a) with the hyper-parameter 0.2 to improve the generalization ability of the network.
We find that the 500-epoch accuracy can be reproduced in these settings.
We conjecture that the cosine annealing learning rate schedule and the mixup technique reconcile the overfitting problem that easily occurs in long period training.
We inherit the RigL training recipe used in our 500-epoch training, and extend the training to 1200 epochs.
The reproduced results are shown in Table 1.
B. O BJECT DETECTION : SSD ON COCO-
To train SSD on COCO-2017 dataset, we use an input size of 300 × 300 pixels, a batch size of 64, step learning rate schedule, SGD optimizer with a momentum of 0.9, and a weight decay of 5e-4.
The learning rate is 2.6e-3 with 300 iterations of linear warm-up.
Each GaP step includes 32 epochs of training and the final fine-tuning includes 65 epochs.
For the prune-from-dense method, we report the mAP of 24.3 by iteratively pruning.
We first pretrain a dense model using 65 epochs, then we prune iteratively to 50% → 70% → 80% → 87.5% → 90%.
For each pruning iteration, we train the sparse model for 65 epochs.
It takes 390 epochs to get the final 90% sparsity model.
Please note that directly pruning a 90% sparsity model only results to 21.
mAP, which is significantly lower than the 24.3 mAP using the iterative pruning method.
The model pruned using C -GaP method can achieve 25.9 mAP.
B. 3D OBJECT PART SEGMENTATION : P OINE N ET ++ ON S HAPE N ET
When training PointNet++ (Qi et al., 2017; Yan, 2019) model on ShapeNet (Yi et al., 2016) dataset, we use a batch size of 32, step learning rate schedule with an initial learning rate of 0.001, SGD optimizer with a momentum of 0.9, and a weight decay of 1e-4.
We pretrain a dense model usingepochs.
Each step of GaP includes 25 epochs of training.
The final fine-tuning includes 250 epochs.
[Figure]
Apply random sparse mask [Figure] model ( [Figure]
[Figure]
P
[Figure]
Divide into partitions Grow part-0, prune part-( - 1) and keep other parts sparse Grow part-1, prune partand keep other parts sparse Grow part-( - 1) , prune part-( -2) and keep other parts sparse
Finetunemodel
[Figure]
[Figure]
[Figure]
Finetune
[Figure]
Apply random sparse mask on model
Divide into partitions and copy to machines
Grow part-[i] on machine-[i], and train the model on all machines.
i {0, 1, 2, ... - 1}
FinetuneModel? Prune and ﬁnetune model Output: pruned DNN model A randomly pruned model
Yes
NoRound+ ∈
Round=
[Figure]
κ κ
[Figure]
Collect and combine part-[i] into one model
[Figure] B. [Figure] T RANSLATION : T RANSFORMER ON WMT-14 E N -D E D ATASET
We use a batch size of 5120, inverse square-root scheduler with an initial learning rate of 0.000846, and Adam optimizer.
Each step in GaP includes 10 epochs of training, and the final fine-tuning includes 40 epochs.
We pretrain a dense model using 40 epochs.
For prune from dense method, we use one-shot magnitude-based pruning to prune Transformer directly to 80% and 90% sparsity.
C T HE D IFFERENCE B ETWEEN S CHEDULED G A P AND R ANDOM S PARSE M ASK E XPLORATION
In our scheduled GaP algorithm, we intend to use all information in a model.
i.e., explore all weights efficiently.
One may argue that when trained sufficiently long, all weights will be explored even for random mask exploration.
However, random sparse mask exploration and our scheduled GaP methods are like sampling operations with replacement and without replacement (Basu, 1958), and their required training time to get the same weight coverage differs a lot .
If we consider a very small model with only 10 weights, and each time we only train one weight and keep the remaining weights the same.
When we apply the scheduled GaP algorithm (without replacement), each weight is guaranteed to be trained after 10 steps (one step indicates training one weight).
This is because the scheduled GaP alternates the weight to train.
The already trained weight is guaranteed not to be re-trained again until all the remaining weights are trained.
On the other hand, if a random weight is selected to train in each step (with replacement), it may take 29 steps to be fairly confident that all weights are trained.
That is almost 3 times more training time than the scheduled GaP. Please note, in our example, we only have 10 weights.
A real neural network contains millions of weights and it is more difficult to guarantee a full exploration.
D A DETAILED WORKFLOW OF THE CYCLIC G A P AND PARALLEL G A P
In this section, we show the general flow of parallel GaP in Algorithm 2, and demonstrate the detailed workflow of the proposed cyclic GaP and parallel GaP method in Figure 4 as a supplement to the
Algorithm 1 and Algorithm 2.
E P ROOF OF P ROPOSITIONE. P ROBLEM FORMULATION
Training deep neural networks can be formulated into minimizing the following problem:
[Equation]
[Equation]
where Θ ∈ R d is the model parameter to be learned, and f ( x ; Θ) is a smooth non-convex cost function in terms of Θ .
The random variable x denotes the data samples that follow the distribution
D .
E x ∼ D [ · ] denotes the expectation over the random variable x .
E. N OTATION
We first clarify the notions used in the convergence analysis.
To characterize how GaP updates model parameters in a partition-wise manner, we introduce matrices
U i ∈
R d × d i , i = 1 , · · · , κ , for which the identity matrix I d ∈ R d × d can be represented as
[Equation]
• Θ S i
=
U Ti Θ ∈ R
d i is the model parameter at the i -th partition (i e., partition S i ).
• m S i ∈
R d i is the mask for the i -th partition.
• ∇ f ( x ; Θ) ∈ R d and ∇ F (Θ) ∈
R
d are the complete stochastic and accurate gradients in terms of Θ , respectively.
• ∇ i f
( x ; Θ) =
U
Ti ∇ f ( x ; Θ) ∈ R
d i is the stochastic gradient for the i -th partition • ∇ i F (Θ) =
U Ti ∇ F (Θ) ∈
R
d i is the gradient for the i -th partition.
• Θ ( i ) q,t ∈
R
d denotes the complete model parameter in which its i -th partition (i
e., [Θ ( i ) q,t ] S i ∈
R
d i ) is newly updated at the q -th round and t -th inner iteration while the other partitions remain unchanged.
• m ( i )
q ∈ R
d is a complete mask vector for the i -th partition in the q -th round.
• x ( i ) q,t is the data sampled at q -th round and t -th inner iteration to update the i -th partition.
E. A LGORITHM REFORMULATION
In this subsection we reformulate the cyclic GaP ( C -Gap) in Algorithm 1 in a way that can present the convergence analysis more easily.
With the notations in Sec.
E.2, the C -Gap Algorithm 1 can be reformulated as
Algorithm 3: C -GaP Algorithm (A math-friendly version) 1 Initialize Θ (1)1 , 0 and m (0)1 randomly.
Initialize pruning ratio r .
2 for round q = 1 , 2 , ..., Q do for partition i = 1 , 2 , ..., κ do
Generate a mask m ( i ) q = GaPMask (Θ ( i ) q, 0 , m ( i − 1) q , r ) ; for t = 1 , ..., T do
Update Θ ( i ) q,t =
Θ ( i ) q,t − 1 − γU i ∇ i f ( x ( i ) q,t − 1 ; Θ ( i ) q,t − 1 ⊙ m ( i ) q )
Update Θ ( i +1)
q,=
Θ ( i ) q,T
Update Θ (1) q +1 , 0 =
Θ ( κ +1) q,T and m (0)
q +1 = m ( κ ) q ; 9 Output: Θ ( κ ) Q,T ⊙ m ( κ ) Q .
In GaPMask (Θ ( i ) q, 0 , m ( i − 1) q , r ) , the partition of [ m ( i ) q ] S i is updated as:
Prune ( i -1)-th partition: update [ m ( i ) q ] S i − 1 ← ArgPruneTo ([Θ ( i ) q,t ] S i − 1 , r ) for pruning ratio r
Grow i -th partition: Let [ m ( i ) q ] S i =
1 d i ∈ R
d i
Other j -th partition: [ m ( i ) q ] S j = [ m ( i − 1) q ] S j (for all j ̸ = i and j ̸
= i − 1 ) and m ( i ) q
= [ m [Equation]
E. A SSUMPTIONS
We now introduce several assumptions on the cost function and the gradient noise, which are standard in the literature.
Assumption 1 (S MOOTHNESS ) .
We assume the cost function F (Θ) is partition-wise L -smooth, i.e.,
[Equation]
The above assumption implies that
[Equation]
or equivalently,
[Equation]
Assumption 2 (G RADIENT NOISE ) .
We assume for any k , t , and i that
[Equation]
where σ > 0 is a constent.
Moreover, we assume the data sample x ( i ) q,t is independent of each other for any k , t , and i .
This assumption implies that the stochastic partition-gradient is unbiased and has bounded variance.
Assumption 3 (M ASK - INCURRED E RROR ) .
It is assumed that
[Equation]
where δ ∈ (0 , 1) . [Equation]
[Equation] Note that [ m ( i ) q ] S i = 1 d i , and only the i -th partition in Θ ( i ) q,t is got updated during iterations ( q, 1) , · · · , ( q, T ) , we have
[Equation]
E. C ONVERGENCE ANALYSIS
Now we are ready to establish the convergence property of C -GaP in Algorithm 3.
The arguments to establish the convergence are strait-forward:
• First, we need to establish the function value E [ F (Θ ( i ) q,t − 1 )] will decrease to E [ F (Θ ( i ) q,t )] , up to an error term caused by the gradient noise and mask pruning, after each inner-iteration t when updating the i -th partition (See Lemma 1).
• Based on the above fact, we next prove E [ F (Θ ( i ) q, 0 )] will decrease to E [ F (Θ ( i +1) q, 0 )] when T -steps training of the i -th partition is completed (See Lemma 2).
• Finally, with the second fact, we show that E [ F (Θ (1) q, 0 )] will decrease to E [ F (Θ (1) q +1 , 0 )] after the k -th outer round (See Lemmas 3 and 4).
Since the function value decreases for each round k , we can prove that C -GaP will converge to a stationary solution after sufficiently large Q rounds (See Theorem 1).
Next we present the detail analysis.
Lemma 1 (D ESCENT L EMMA AFTER EACH I NNER -I TERATION ) .
Under Assumptions 1-3, it holds for each k , i , and t that
[Equation]
Remark 4.
It is observed in Lemma 1 that the function value will decrease by γ 3 E ∥∇ i F (Θ ( i ) q,t − 1 ) ∥for each inner-iteration t .
But such decrement suffers from two error terms: one is caused by stochastic gradient noise, and the other is by inexact pruning.
Proof.
Since F (Θ) is L -smooth (see Assumption 1)
[Equation]
it holds that
[Equation]
With Assumption
[Equation]
it is easy to verify that
[Equation]
equation 7 =
[Equation]
[Equation]
Using a similar way
[Equation]
we can prove
[Equation]
with equation
[Equation]
that
[Equation]
By taking the expectation over equation 12 and substituting equation 13 and equation 14 into equation
[Equation]
we achieve
[Equation] [Equation]
Note that
[Equation]
and, similarly,
[Equation]
Substituting equation 16 and equation 17 into equation
[Equation]
and by setting γ ≤ 1 / (6 L )
[Equation]
we have
E
[Equation]
Lemma 2 (D ESCENT L EMMA AFTER EACH PARTITION UPDATE ) .
Under Assumptions 1-3, it holds for each k and i that
[Equation]
Proof.
Summing the inequality over equation 11 for t = [Equation] [Equation] T
[Equation] we achieve γ T
[Equation] [Equation]
which completes the proof
[Equation]
Lemma 3 (D ESCENT L EMMA AFTER EACH ROUND ) .
Under Assumptions 1-3, it holds for each k that
[Equation]
Proof.
Summing the inequality in equation 19 over i = 1 , · · · , κ , we have
[Equation]
On the other hand
[Equation]
by the updating rule of Θ ( i ) q,
[Equation]
we have
[Equation]
Substituting equation 23 into equation
[Equation]
we achieve
[Equation]
(24) which completes the proof.
To finish the convergence proof, we still need to establish the relation between the decrement (cid:80) κi
=
[Equation]
Tt =
[Equation]
E ∥∇ i F (Θ ( i ) q,t − 1 ) ∥ 2 and E ∥∇ F (Θ (1) q, 0 ) ∥ so that
Lemma 4 (D ESCENT L EMMA II
A FTER EACH R OUND ) .
Under Assumptions 1-3, and learning rate γ ≤ 1 4 κLT , it holds for each k that
[Equation]
Proof.
In the following we lower bound the term (cid:80) κi = [Equation] Tt
[Equation] E ∥∇ i F ( [Equation] ∥ 2 in equation 21.
Recall Algorithm 3 that
[Equation]
With these relations, it holds for i ≥ 1 and t ≥ 1 that
[Equation]
[Equation]
where (a) holds because Θ ( i ) q, 0 − Θ ( i − 1) q,and Θ ( j ) q, 0 − Θ ( j − 1) q,are orthogonal to each other when i ̸
= j , (b) holds because of equation 26 and equation 27, (c) holds because of Assumption 2, and (d) holds because of the Jensen’s inequality.
With the above inequality, it holds for t ∈ [1 , T ] that
[Equation]
equation
[Equation]
equation
[Equation]
[Equation]
Summing the above inequality over t = 0 , 1 , · · · , T , we have
[Equation]
Summing the above inequality over i =
[Equation] [Equation]
κ
[Equation]
we have
T [Equation] i
[Equation]
Combining equation 24 and equation [Equation]
we have
[Equation] [Equation]
where (a) and (b) use the facts that
[Equation]
Then the inequality equation 32 becomes
[Equation]
Finally we establish the convergence of the C -Gap algorithm as follows:
Theorem 1 (C-GAP CONVERGENCE ) .
Under Assumptions 1-3, if learning rate is set as γ =4 κLT √ Q , it holds that
[Equation]
where G = 96 κL E [ F (Θ (1)1 , 0 )] + σ 2 is a constant.
Proof.
Since
[Equation]
we thus have
[Equation]
[Equation] Summing the above inequality over q =
[Equation]
[Equation]
Q and taking its average
[Equation]
we achieve
[Equation]
If we let γ = 1 4 κLT √ Q , the above inequalities becomes
[Equation]
which completes the proof
[Equation]