Learning sparse DNNs with soft thresholding of weights during training

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 ) . 1 INTRODUCTION . 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 . 2 STATE OF THE ART . 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 ? 2.1 PRUNING CRITERION . 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 . 2.2 PRUNING AND TRAINING INTERACTION . 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 . 2.3 PRUNING PERSISTENCE . 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 .

Summary: 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.