Year: 2021

Abstract:
This paper proposes a new and simple way of training sparse neural networks. Our method is based on a differentiation of the forward and backward paths: the weights in the forward path are a thresholded version of the weights maintained in the backward path. This decoupling allows for micro-updates, produced by gradient descent, to stack up, leading to the possible re-activation of weights that were set to zero in earlier training steps. At the end of training, links with zero weights are pruned away. Additional critical specificities of our approach lie (i) in the progressive increase of the zeroed weight ratio along the training, and (ii) in the use of soft-thresholding rather than hard-tresholding to derive the forward-path weights from the ones maintained in the backward path. At constant accuracy, our approach reduces the number of training cycles to 1 compared to the state-of-the-art recursive pruning methods. At high pruning rates, it also improves the model accuracy compared to other single cycle pruning approaches (66.18% top-1 accuracy when training a ResNet-50 on ImageNet at 98% sparsity).

text: This paper proposes a new and simple way of training sparse neural networks. Our method is based on a differentiation of the forward and backward paths: the weights in the forward path are a thresholded version of the weights maintained in the backward path. This decoupling allows for micro-updates, produced by gradient descent, to stack up, leading to the possible re-activation of weights that were set to zero in earlier training steps. At the end of training, links with zero weights are pruned away.
Additional critical specificities of our approach lie (i) in the progressive increase of the zeroed weight ratio along the training, and (ii) in the use of soft-thresholding rather than hard-tresholding to derive the forward-path weights from the ones maintained in the backward path. At constant accuracy, our approach reduces the number of training cycles to 1 compared to the state-of-the-art recursive pruning methods. At high pruning rates, it also improves the model accuracy compared to other single cycle pruning approaches (66.18% top-1 accuracy when training a ResNet-50 on ImageNet at 98% sparsity).

heading: 1 INTRODUCTION
text: State-of-the-art neural networks are composed of a few million parameters, leading to a few billion computations per inference. To limit these amounts, sparse networks have been thoroughly investigated in the past few years (Frankle & Carbin (2019); Liu et al. (2019); Lee et al. (2019); Evci et al. (2019); Frankle et al. (2020); Renda et al. (2020)). Those networks reduce the inference complexity by setting most of their weight parameters to zero. Recent works have shown that this can be done without penalizing the prediction accuracy compared to non-sparse networks (cite).
Training methods leading to sparse networks generally build on multiple rounds, or cycles, of iterative gradient descent updates. As detailed in Section 2, each round applies recursively to the outcome of the previous round, after having pruned a (small) predefined fraction of the links, typically selected as the ones with smallest weights. A training cycle typically consists in a full training procedure, with a complete learning rate schedule and consequently a large number of weight updates. In contrast, our work proposes to train sparse networks in a single round of gradient updates, thereby dramatically reducing the training resources. Therefore, as an original contribution, our work adopts in the forward path a soft-thresholded version of the weights that were updated based on the back-propagated gradients, and updates the parameter of the soft-threshold operator to progressively increase the ratio of pruned (inactive??) weights along the training iterations. In that way, the pruning of a weight is never definitive, and weights that are pruned at some stage of the training process get the opportunity to become non-zero again when gradient feedback makes them relevant. Moreover, gradients can be back-propagated according to the (non-zero) weights maintained along the backward path, ensuring a dense propagation of feedback, without being hampered by the definitive pruning of some links. Those specificities appear to largely benefit the performance of the trained model, compared to previous works. In particular, our proposed method is on par in terms of prediction accuracy with state-of-the-art sparse model training approaches, while significantly reducing (up to one order of magnitude) the amount of training iterations.

heading: 2 STATE OF THE ART
text: In this paper we concentrate our efforts on unstructured pruning. While pruning in a structured manner leads to an easier acceleration of CNN within existing libraries, the performance obtained at the same level of sparsity is often inferior (Li et al. (2017); He et al. (2019); Wang et al. (2020)). The work of Liu et al. (2019) also showed that the chosen structure was more important than the actual way the pruned filters were chosen, indicating that l1-based weight selection might not be the best way of choosing which filters to prune. Moreover, recent works (Kalchbrenner et al. (2018); Park et al. (2017); Gale et al. (2019)) managed to significantly speed up networks with unstructured sparsity on off-the-shelve hardware.
Naive pruning approaches remove links from the network in one shot at the end of an entire training cycle, based on their respective weight magnitude (LeCun et al. (1990); Han et al. (2015)). This however comes at the cost of a loss in prediction accuracy. In order to reduce this loss, the unpruned weights generally undergo a finetuning step (i.e. a shorter training cycle, with small learning rate), while the pruned ones stay at zero. Over the years, many variants have been proposed to improve this baseline approach. They can be differentiated based on the way they address the three following issues:
• Pruning criterion: How are the pruned weights selected ?
• Pruning and training interaction: When is the pruning implemented along the training iterations ?
• Pruning persistence: Do pruned weights get the opportunity to come back to non-zero values or not ?

heading: 2.1 PRUNING CRITERION
text: Obviously, weights with smaller magnitude should be pruned first. However, the magnitude criterion is generally affected by two other considerations.
On the one hand, it has been widely documented (Mocanu et al. (2018); Evci et al. (2019); Lee et al. (2021)) that the network prediction accuracy is more impacted when pruning occurs close to the input. More generally, the allocation of computational resources, and thus of non-zero weights, across layers remains an open question. Two heuristics have been recently proposed to address this issue in a satisfying manner:
1. Erdos-Rényi-Kernel (ERK) Evci et al. (2019) proposes to scale the global pruning ratio with a layer-wise pruning ratio, defined according to the number of neurons (n), the width (w) and the height (h) of the layer convolutional kernel. Formally, letting l denote the layer index, the scaling factor is defined to be 1− n l−1+nl+wl+hl
nl−1·nl·wl·hl , which results in a more aggressive pruning of layers with more parameters.
2. Layer-Aware-Pruning (LAMP) Lee et al. (2021) selects the weights to prune globally, but relies on a score that is assigned to each weight rather than on its magnitude. This score is computed on the sorted and flattened array of weights (Wlsorted) of each layer (l). Mathe-
matically, score(i,Wlsorted) = (W lsorted,i) 2∑ j≥i(W l sorted,j) 2
The ERK criterion has more rigid constraints and performs slightly worse in traditional one-cycle pruning than LAMP (Lee et al. (2021)). Hence, we use the latter in Section 4.3.3 as a comparison point to our method.

heading: 2.2 PRUNING AND TRAINING INTERACTION
text: Pruning should only be considered after some training iterations, to make sure only the links that do not significantly impact the prediction are removed from the network. Recent literature has however revealed that the best performance are not obtained simply by pruning a fully trained network. The recent discoveries in that respect can be summarized as follows :
Subnetworks. It has been discovered by Frankle & Carbin (2019) and Frankle et al. (2020) that so-called subnetworks, corresponding to weights that are deemed essential for the final prediction, are formed quite early during training. This phenomenon arises because the weights change more significantly during the first training iterations than in subsequent ones (Leclerc & Madry (2020); Frankle et al. (2020)). This is simply brought about by the gradient magnitude generally starkly decreasing after a few initial iterations, denoted Tcrit in the following, to progressively stabilize. Interestingly, recent works related to pruning have shown that it is advantageous to derive sparse networks from the weights available in the early stages of training, once the network has achieved this relative stability after this small number Tcrit of training iterations (You et al. (2019)). This is done recursively as described below.
Rewinding. After an initial training cycle, the weights with the lowest magnitudes are pruned as before, but the remaining ones are rewound to their values at Tcrit (Frankle et al. (2020)). A novel training cycle is then launched from Tcrit onwards: meaning that the learning rate (and optimizer parameters) are reset to their values at Tcrit, and that a full cycle of training iterations, with complete learning rate schedule, is run. The difference compared to the initial training cycle lies in the fact that the pruned weights that will stay at zero for the remainder of the restarted training cycle. Surprisingly, the sparse model obtained with this two-cycles training strategy appears to perform better than the reference baseline, obtained by finetuning the pruned weights. In addition, the work in (Frankle et al. (2020)) has shown that applying the pruning earlier during the first training cycle leads to worse performance, thereby justifying the cost of the preliminary training cycle. Later work by Renda et al. (2020) shows that the critical component of the rewinding step is the rewinding of the learning rate. In our method, the weights are constantly updated instead of being pruned or rewound at the end of a training cycle. However, as will be shown in Section 4, a rewind of the learning rate and weights can further boost the performance of our method, although coming at the cost of an additional training cycle.
Recursive Pruning with rewinding. As illustrated in our experiments (in Section 4), the prediction accuracy of the model obtained with the two-cycles training strategy sharply decreases when the pruning ratio increases. Implementing the pruning procedure recursively has been shown to effectively circumvent this issue: instead of pruning the whole ratio of weights after the initial training cycle, only a fraction of this ratio is actually pruned. The surviving weights are rewound, the training cycle resumes from Tcrit and, when it ends, an additional fraction of the weights are pruned. The remaining weights are then rewound, training restarts and so on. This procedure results in the state-of-the-art performance, but rapidly gets expensive in terms of training time. If we take a 20% step-size (in order to compare with Frankle & Carbin (2019)), 11 training cycles are necessary to reach ∼ 90% sparsity.

heading: 2.3 PRUNING PERSISTENCE
text: Above, once a weight has been pruned, it stays pruned all along the training cycle. Some works however have considered the possibility to re-activate a weight that has been set to zero. Works like Dettmers & Zettlemoyer (2019) and Evci et al. (2019) rely on computations on the gradient magnitudes to decide to re-activate weights. In contrast, Kusupati et al. (2020) learns a progressive layer-wise threshold via back-propagation, but only computes the backpropagated gradients on the unpruned weights. In contrast, our work allows for all gradients to flow through, updating even weights that didn’t take part in the forward computation.

heading: 3 OUR PRUNING METHOD
text: To train a sparse network in a single training cycle without penalizing its prediction accuracy, our work proposes to adopt a non-persistent soft-thresholded pruning strategy and to progressively increase the sparsity ratio of the network along the iterations of a training cycle. The remainder of this section presents and motivates the 3 main components of our method. Our ablation study in Section 4.3 will confirm our arguments experimentally.
Non-persistent pruning In contrast to most previous works, our approach proposes to maintain two versions of the weights. A dense version of the weights is maintained and updated in the backward path, while their thresholded version is considered in the forward path and to backpropagate the gradients, so as to mimick the impact of pruning. This differentiation of weights in the forward and backward paths is inspired by the straight-through-estimator used to train quantized neural networks (Hubara et al. (2016); Mellempudi et al. (2017)). Formally, we define the sparse weights Wsparse adopted in the forward path as a thresholded version of the dense weights Wdense updated in the backward path: Wparse = apply th(Wdense, th). A consequence of this definition of Wsparse is that a forward weight that has been zeroed at some point during training might become non-zero again if its corresponding backward weight magnitude sufficiently increases in subsequent training iterations. Hence, the pruning emulated in the forward path is non-persistent.
Soft Thresholding Most of the persistent pruning literature makes use of hard thresholding (Han et al. (2015); Frankle & Carbin (2019)). However, in presence of non-persistent pruning, the weights that are maintained in the backward path, are likely to cross the pruning threshold, and be subject to the hard thresholding function discontinuity. This might lead to abrupt changes in the forward path, which might be inconsistent with the update implemented in the backward path. Figure 1 depicts a case where the pruning threshold is large compared to the gradient magnitude. In this case, hard thresholding induces severe discontinuities in the forward weight update, while soft thresholding preserves a smooth and consistent evolution of this forward weight. Our experiments in Section 4.3.2 confirm that soft thresholding results in a more stable training.
Increasing the pruning ratio To mitigate the instabilities induced by pruning, and leave the network the time to adapt to the zeroing of some weights in the forward path, we progressively increase the sparsity ratio during training. This allows the remaining weights to more easily adapt to the removal of a majority of its connections and avoids the neural network from getting stuck in a local minimum early on. A slow increase of the sparsity ratio during training was already applied successfully by Kusupati et al. (2020) and Wang et al. (2020). The former used it to approximately regulate the sparsity of each layer independently, while the latter applied it to the regularization of small weights, forcing them closer to zero over time. We use this linearly increasing ratio to force an ever-increasing exact amount of sparsity into our network. This ratio will increase linearly during the first half of the training, which also corresponds to the highest learning rate plateau for the networks trained on Cifar-10. A pseudo-code for our method is provided in Appendix A.1.

heading: 4 EXPERIMENTS

heading: 4.1 EXPERIMENTAL SETUP AND TERMINOLOGY
text: Our experiments train ResNet-20 (He et al. (2015)) and VGG-13 (Simonyan & Zisserman (2014)) (as simplified by Liu et al. (2019)) on the Cifar-10 Krizhevsky (2009) dataset. Every training cycle uses the same hyperparameters as defined in Appendix A.2 regardless of the method. This follows the choices made in Frankle et al. (2020) and, in practice, no significant difference has been observed when tuning the parameters individually for each method. For each pruning method, each network is trained three times, with distinct initialization seeds. The same triplet of seeds is used
for all pruning methods, and all figures display the mean and variance of the accuracy achieved by the models obtained from those seeds as a function of the sparsity ratio.
To denote the methods compared in this section, we introduce the following terms:
1. OneCycle : the network is pruned after one training cycle
2. Rw (rewind): an additional training is applied to the sub-network pruned at the end of the initial (or previous, in case of recursive pruning) training cycle. The weights, learning rate and optimizer of this subnetwork are rewound to Tcrit.
3. Rec. Rw (recursive rewind) : the pruning/rewind/training procedure is applied recursively, by pruning 20% of the remaining weights at each recursive step.
4. SoftThresh (Ours): As explained in Section 3, an increasing fraction of weights is set to zero in the forward path, by soft-thresholding the weights maintained in the backward path. The soft-threshold operator parameter is defined to progressively increase the sparsity ratio of the weights along the full training cycle.

heading: 4.2 COMPARISON WITH BASELINES
text: Pruning baseline Recursive rewind corresponds to the SotA in terms of pruned network accuracy, while one-cycle with rewind corresponds to a less accurate but also less complex (in terms of number of training cycles) baseline. Figure 2 plots our method against those 2 baselines. Our method matches recursive rewind all the while needing exponentially less training cycles.
Prior work To the best of our knowledge, Kusupati et al. (2020) detains the SotA one-cycle pruning results with a ResNet-50 on ImageNet on all thresholds. Table 1 compares our method, which was trained with the same meta-parameters. Our approach exceeds the performance of Kusupati et al. (2020) at approximately equivalent pruning budgets.

heading: 4.3 ABLATION STUDY
text: This section considers variants of our method to answer the following 3 questions:
1. Are dense weights necessary for gradient computation?
2. Is soft-thresholding better than hard-thresholding in practice?
3. Does the pruning criterion, i.e. the selection of the soft-thresholding operator parameter for each layer, affect performance?

heading: 4.3.1 DENSE OR SPARSE?
text: As stated earlier in Section 3, we update the dense weights during back-propagation. To emphasize the importance of this dense update strategy, our results are compared to the sparse update of weights during the backward pass. In that case, only the weights whose magnitude is larger than the thresholding operator parameter are updated. In Figure 3, these 2 scenarios (respectively referred to
as persistent and Wsparse in BP) were inspected, and a clear trend emerges : updating only the active weights during back-propagation leads to a lower accuracy compared to updating all the weights (dense).

heading: 4.3.2 HARD OR SOFT?
text: In Section 3, an explanation was given as to why a soft-thresholded approach may lead to better results than a hard-thresholded one due to the removal of a discrepancy between forward and backward propagation. Figure 4 clearly confirms our intuition. We note that until ≈ 98% sparsity both strategies are relevant to threshold the weights: The hard- and soft-threshold performance after weight rewinding and re-training is nearly identical. However, the hard-threshold training barely manages to keep up with our baseline of one-cycle+rewind. This is caused by a highly fluctuating loss function during training, implying a great deal of instability. Moreover, this instability reaches a critical point when only ≈ 2% of the weights remain, after which the network becomes unable to learn anything due to a whole layer being pruned.

heading: 4.3.3 L1 OR LAMP?
text: Criteria like LAMP (Lee et al. (2021)) try to mitigate instabilities by normalizing individual weights with respect to the other ones inside their layer. To the right in Figure 4 LAMP is shown to significantly improve the results of hard thresholding at high pruning ratio, getting rid of the noted instability.

heading: 5 WEIGHT INSIGHTS
text: The previous sections have numerically shown the importance of keeping the dense version of the gradient when updating the weights through gradient descent. By taking a look at the underlying weight distributions and zeroed weights, two additional observations stand out:
1. The weight distributions that result from our soft-thresholding method, are different from those generated by recursive pruning
2. During training our method encourages a high number of weights to flip between a zero (inactive) and non-zero (active) state. This high number of crossings leads to networks that generalize better.

heading: 5.1 WEIGHT DISTRIBUTIONS
text: The weight distributions for all the layers are depicted in Figure 5 at a sparsity ratio of 98.8% to more clearly see the emerging phenomena. The weight distributions generated by pruning methods (be it OneCycle+rewind or Recursive-rewind) differ heavily from those obtained with our soft-thresholded approach. Our initial training cycle generates a very peaky normal distribution. Rewinding the weights and only retraining the unpruned ones, puts us closer to distributions resembling the sum of 2 overlapping gaussians mirrored around zero. This is in opposition to pruning which generates a mostly normal distribution with some weights that are pushed to higher absolute values. This indicates that our method selects weights to set to zero in a fundamentally different way than traditional pruning.

heading: 5.2 THRESHOLD CROSSINGS
text: In Section 3, a plausible explanation was given for the use of growing sparsity. However, we can observe why this training regime is so beneficial: growing the sparsity ratio leads to more weights jumping from an inactive to an active state and vice-versa. This phenomenon will be referred to as threshold crossings. Figure 6 shows the number of times the weights will cross during different training phases that are each 40 epochs long. The weights that will end up staying inactive once the training completed, are separated to more clearly show the trend: a lot of crossings occur during the progressive growing of the sparsity ratio before stabilizing after 80 epochs.
This same experience was repeated using the whole sparsity ratio from Tcrit onwards in order to eliminate the high learning rate factor. Not allowing this ratio to progressively grow over time, leads first to a noticeably lower final accuracy and second we observe fewer threshold crossings during the early phases of the training process. This implies that the mask is chosen at Tcrit and doesn’t change much from there on out, moving our method closer to real pruning than is desirable. This observation points to the weights needing some time to adapt to the disappearance of their neighbours. This also consolidates our method’s starting point that applying fewer constraints on the weight updates, thanks to dense gradients, leads to the spontaneous emergence of weightadjustments.

heading: 6 CONCLUSION
text: We have proposed to apply soft-thresholding to the weights maintained in the network backward path, so as to define sparse weights in the forward path. The soft-thresholding operator parameter is changed along training to progressively increase the ratio of weights that are set to zero by the thresholding, thereby reaching high sparsity ratio in a single training cycle. This is in contrast to traditional persistent pruning methods that need many training cycles to obtain models whose prediction accuracy reaches what our method obtains in one shot. Our analysis has shown the importance of using soft-thresholding, and not hard-thresholding, to avoid sharp discrepancies in the evolution of the weights maintained in the forward and backward paths. The lesser constraints of these dense gradients on the weights allow for threshold crossings which validate the importance of our method. Overall, our method results in a new SotA for one-training-cycle accuracy on ImageNet with a ResNet-50 of 66.18% at 98% sparsity.
REPRODUCIBILITY STATEMENT
The source code for this paper has been added to the supplementary material of this submission and the code will be open-sourced after publication. Every training configuration has been safeguarded in order to minimize any possible misinterpretation of the manuscript. The network architectures used are standardized (ResNet, VGG) and so are the datasets (Cifar-10, ImageNet). The training mete-parameters were taken from previous in order to enable a better comparison accross the literature. Whenever it was computationally sound, experiments were run with 3 different seeds (shared across all experiments) in order to guarantee more consistent results.

heading: A APPENDIX
text: A.1 PSEUDO-CODE
Listing 1: Pytorch-like pseudo-code for Recursive pruning with rewind a c t p r u n i n g = 0
W = Win i t Wmask = o n e s l i k e (W) t r a i n c y c l e s = 0 whi le True :
# Comple te T r a i n i n g c y c l e f o r epoch in range ( n e p o c h s ) :
f o r i n p u t s , t r u e o u t p u t s in d a t a s e t : o u t p u t s = NN(W * Wmask ) ( i n p u t s ) l o s s = l o s s f x ( t r u e o u t p u t s , o u t p u t s ) c o m p u t e g r a d i e n t s ( l o s s , W[Wmask==1] )
W = o p t i m i z e r . u p d a t e (W[Wmask==1] , l e a r n i n g r a t e ) l e a r n i n g r a t e . u p d a t e ( epoch )
a c t p r u n i n g = 1 − (1 − p r u n i n g s t e p ) ˆ t r a i n c y c l e s i f a c t p r u n i n g > f i n a l p r u n i n g :
break
# Prune more w e i g h t s Wmask = Wmask − prune (W[Wmask==1] , p r u n i n g s t e p )
Listing 2: Pytorch-like pseudo-code for our soft-thresholded sparse training with increasing sparsity W = Wini t s t a r t e p o c h = T c r i t end epoch = n e p o c h s / / 2 p r u n i n g d e l t a = f i n a l p r u n i n g / ( l e n ( d a t a s e t ) * ( end epoch − s t a r t e p o c h ) ) a c t p r u n i n g = 0 # Comple te T r a i n i n g c y c l e f o r epoch in range ( n e p o c h s ) :
f o r i n p u t s , t r u e o u t p u t s in d a t a s e t : t h = c o m p u t e p r u n e t h (W, a c t p r u n i n g )
# Do n o t compute g r a d i e n t s on p r u n i n g o f f s e t wi th n o g r a d i e n t s ( ) :
Wdel ta = s o f t t h r e s h o l d (W, t h ) − W
Wpruned = W + Wdelta o u t p u t s = NN( Wpruned ) ( i n p u t s ) l o s s = l o s s f x ( t r u e o u t p u t s , o u t p u t s ) c o m p u t e g r a d i e n t s ( l o s s , W)
W = o p t i m i z e r . u p d a t e (W, l e a r n i n g r a t e )
i f epoch >= s t a r t e p o c h and epoch < end epoch : a c t p r u n i n g += p r u n i n g d e l t a
l e a r n i n g r a t e . u p d a t e ( epoch )
A.2 HYPER-PARAMETER TUNING
The hyper parameters used for ResNet-20 and VGG-13 trained on Cifar-10 are displayed in Table 2. All run configurations are also available in the code submitted with the supplementary material.

Review summaries:
The paper uses three techniques to "sparsify" a dense network over training. These techniques are: non-persistent pruning, which backpropagates the gradients and updates a copy of the dense weights of the network; soft thresholding, which replaces the hard threshold typically used while prunning neural networks with the soft threshold operator; and increasing the pruning ratio linearly while training. The paper shows that these three techniques can achieve good accuracy at a high sparsity levels on ImageNet and CIFAR-10, when compared to other approaches. 

In this paper, the authors proposed a training algorithm for sparse deep neural network. Their key idea is to decouple the forward and backward paths, that is, the weights used in the forward path are a threshold version of the weights maintained in the backward path. In this way, some pruned weights have possibility to be re-activated. The authors also conducted some experiments to evaluate the performance of their methods. 

The paper proposes a single training cycle method to train sparse networks. It adopts a soft-threshold pruning strategy and progressively increase the sparsity ratio of the network along the training iterations. Experiments show its advantage.

The authors propose a method for unstructured model pruning that applies soft-thresholding in the forward pass, while using the straight-through estimator in the backward pass. The latter allows for accumulation of gradients and rewiring of earlier-pruned weights. The authors show SOTA performance for the high sparsity-regime on ResNet-50 for Imagenet, and include ablation experiments on CIFAR10 with ResNet-20.