Learned Threshold Pruning

1 INTRODUCTION . Deep neural networks ( DNNs ) have provided state-of-the-art solutions for several challenging tasks in many domains such as computer vision , natural language understanding , and speech processing . With the increasing demand for deploying DNNs on resource-constrained edge devices , it has become even more critical to reduce the memory footprint of neural networks and also to achieve power-efficient inference on these devices . Many methods in model compression Hassibi et al . ( 1993 ) ; LeCun et al . ( 1989 ) ; Han et al . ( 2015b ) ; Zhang et al . ( 2018 ) , model quantization Jacob et al . ( 2018 ) ; Lin et al . ( 2016 ) ; Zhou et al . ( 2017 ) ; Faraone et al . ( 2018 ) and neural architecture search Sandler et al . ( 2018 ) ; Tan & Le ( 2019a ) ; Cai et al . ( 2018 ) ; Wu et al . ( 2019 ) have been introduced with these goals in mind . Neural network compression mainly falls into two categories : structured and unstructured pruning . Structured pruning methods , e.g. , He et al . ( 2017 ) ; Li et al . ( 2017 ) ; Zhang et al . ( 2016 ) ; He et al . ( 2018 ) , change the network ’ s architecture by removing input channels from convolutional layers or by applying tensor decomposition to the layer weight matrices whereas unstructured pruning methods such as Han et al . ( 2015b ) ; Frankle & Carbin ( 2019 ) ; Zhang et al . ( 2018 ) rely on removing individual weights from the neural network . Although unstructured pruning methods achieve much higher weight sparsity ratio than structured pruning , unstructured is thought to be less hardware friendly because the irregular sparsity is often difficult to exploit for efficient computation Anwar et al . ( 2017 ) . However , recent advances in AI accelerator design Ignatov et al . ( 2018 ) have targeted support for highly efficient sparse matrix multiply-and-accumulate operations . Because of this , it is getting increasingly important to develop state-of-the-art algorithms for unstructured pruning . Most unstructured weight pruning methods are based on the assumption that smaller weights do not contribute as much to the model ’ s performance . These pruning methods iteratively prune the weights that are smaller than a certain threshold and retrain the network to regain the performance lost during pruning . A key challenge in unstructured pruning is to find an optimal setting for these pruning thresholds . Merely setting the same threshold for all layers may not be appropriate because the distribution and ranges of the weights in each layer can be very different . Also , different layers may have varying sensitivities to pruning , depending on their position in the network ( initial layers versus final layers ) or their type ( depth-wise separable versus standard convolutional layers ) . The best setting of thresholds should consider these layer-wise characteristics . Many methods Zhang et al . ( 2018 ) ; Ye et al . ( 2019 ) ; Manessi et al . ( 2018 ) propose a way to search these layer-wise thresholds but become quite computationally expensive for networks with a large number of layers , such as ResNet50 or EfficientNet . In this paper , we propose Learned Threshold Pruning ( LTP ) to address these challenges . Our proposed method uses separate pruning thresholds for every layer . We make the layer-wise thresholds trainable , allowing the training procedure to find optimal thresholds alongside the layer weights during finetuning . An added benefit of making these thresholds trainable is that it makes LTP fast , and the method converges quickly compared to other iterative methods such as Zhang et al . ( 2018 ) ; Ye et al . ( 2019 ) . LTP also achieves high compression on newer networks Tan & Le ( 2019a ) ; Sandler et al . ( 2018 ) ; Tan & Le ( 2019b ) with squeeze-excite Hu et al . ( 2018 ) and depth-wise convolutional layers Chollet ( 2017 ) . Our key contributions in this work are the following : • We propose a gradient-based algorithm for unstructured pruning , that introduces a learnable threshold parameter for every layer . This threshold is trained jointly with the layer weights . We use soft-pruning and soft L0 regularization to make this process end-to-end trainable . • We show that making layer-wise thresholds trainable makes LTP computationally very efficient compared to other methods that search for per-layer thresholds via an iterative pruning and finetuning process , e.g. , LTP pruned ResNet50 to 9.11x in just 18 epochs with 12 additional epochs of fine-tuning , and MixNet-S to 2x in 17 epochs without need for further finetuning . • We demonstrate state-of-the-art compression ratios on newer architectures , i.e. , 1.33× , 3× and 2× for MobileNetV2 , EfficientNet-B0 and MixNet-S , respectively , which are already optimized for efficient inference , with less than 1 % drop in Top-1 accuracy . • The proposed method provides a trace of checkpoints with varying pruning ratios and accuracies . Because of this , the user can choose any desired checkpoint based on the sparsity and performance requirements for the desired application . 2 RELATED WORK . Several methods have been proposed for both structured and unstructured pruning of deep networks . Methods like He et al . ( 2017 ) ; Li et al . ( 2017 ) use layer-wise statistics and data to remove input channels from convolutional layers . Other methods apply tensor decompositions on neural network layers , Denton et al . ( 2014 ) ; Jaderberg et al . ( 2014 ) ; Zhang et al . ( 2016 ) apply SVD to decompose weight matrices and Kim et al . ( 2015 ) ; Lebedev et al . ( 2014 ) apply tucker and cp-decompositions to compress . An overview of these methods can be found in Kuzmin et al . ( 2019 ) . These methods are all applied after training a network and need fine-tuning afterwards . Other structured methods change the shape of a neural network while training . Methods like Bayesian Compression Louizos et al . ( 2017 ) , VIBnets Dai et al . ( 2018 ) and L1/L0-regularization Srinivas et al . ( 2017 ) ; Louizos et al . ( 2018 ) add trainable gates to each layer to prune while training . In this paper we consider unstructured pruning , i.e . removing individual weights from a network . This type of pruning was already in use in 1989 in the optimal brain damage LeCun et al . ( 1989 ) and optimal brain surgeon Hassibi et al . ( 1993 ) papers , which removed individual weights in neural networks by use of Hessian information . More recently , Han et al . ( 2015a ) used the method from Han et al . ( 2015b ) as part of their full model compression pipeline , removing weights with small magnitudes and fine-tuning afterwards . This type of method is frequently used for pruning , and has recently been picked up for finding DNN subnetworks that work just as well as their mother network in Frankle & Carbin ( 2019 ) ; Zhou et al . ( 2019 ) . Another recent application of Han et al . ( 2015b ) is by Renda et al . ( 2020 ) where weight and learning-rate rewinding schemes are used to achieve competitive pruning performances . These methods , however , are very computationally extensive requiring many hundreds of epochs of re-training . Finally , papers such as Molchanov et al . ( 2017 ) ; Ullrich et al . ( 2017 ) apply a variational Bayesian framework on network pruning . Other methods that are similar to our work are Zhang et al . ( 2018 ) and Ye et al . ( 2019 ) . These papers apply the alternating method of Lagrange multipliers to pruning , which slowly coaxes a network into pruning weights with a L2-regularization-like term . One problem of these methods is that they are time-intensive , another is that they need manual tweaking of compression rates for each layer . In our method , we get rid of these restrictions and achieve comparable compression results , at fraction of the computational burden and without any need for setting per-layer pruning ratios manually . Kusupati et al . ( 2020 ) and Manessi et al . ( 2018 ) learn per-layer thresholds automatically using soft thresholding operator or a close variant of it . However they rely on L1 and/or L2 regularization , which as shown in section 3.2 , is inefficient when used in networks with batch-normalization Ioffe & Szegedy ( 2015 ) . He et al . ( 2018 ) use reinforcement learning to set layer-wise prune ratios for structured pruning , whereas we learn the pruning thresholds in the fine-tuning process . 3 METHOD . LTP comprises two key ideas , soft-pruning and soft L0 regularization , detailed in sections 3.1 and 3.2 , respectively . The full LTP algorithm is then presented in section 3.3 . 3.1 SOFT PRUNING . The main challenge in learning per-layer thresholds during training is that the pruning operation is not differentiable . More precisely , consider an N -layer DNN where the weights for the l-th convolutional or fully-connected layer are denoted by { wkl } , and let k index the weights within the layer . In magnitude-based pruning Han et al . ( 2015b ) the relation between layer l ’ s uncompressed weights and pruned weights is given by : vkl = wkl × step ( w2kl − τl ) , ( 1 ) where τl denotes the layer ’ s pruning threshold and step ( . ) denotes the Heaviside step function . We name this scheme hard-pruning . Since the step function is not differentiable , ( 1 ) can not be used to learn thresholds through back-propagation . To get around this problem , during training LTP replaces ( 1 ) with soft-pruning vkl , wkl · sigm ( w2kl − τl T ) , ( 2 ) where sigm ( . ) denotes the sigmoid function and T is a temperature hyper-parameter . As a result of ( 2 ) being differentiable , back-propagation can now be applied to learn both the weights and thresholds simultaneously . Defining soft-pruning as in ( 2 ) has another advantage . Note that if w2kl is much smaller than τl ( i.e. , τl − w2kl T ) , wkl ’ s soft-pruned version is almost zero and it is pruned away , whereas if it is much larger ( i.e. , w2kl − τl T ) , wkl ≈ vkl . Weights falling within the transitional region of the sigmoid function ( i.e. , |w2kl − τl| ∼ T ) , however , may end up being pruned or kept depending on their contribution to optimizing the loss function . If they are important , the weights are pushed above the threshold through minimization of the classification loss . Otherwise , they are pulled below the threshold through regularization . This means that although LTP utilizes pruning thresholds similar to previous methods , it is not entirely a magnitude-based pruning method , as it allows the network to keep important weights that were initially small and removing some of the unimportant weights that were initially large , c.f. , Figure 1 ( left ) . Continuing with equation ( 2 ) , it follows that ∂vkl ∂τl = −1 2 · σT ( wkl ) and ∂vkl ∂wkl = sigm ( w2kl − τl T ) + wkl · σT ( wkl ) , ( 3 ) with σT ( wkl ) , 2wkl T · sigm ( w 2 kl − τl T ) × ( 1− sigm ( w 2 kl − τl T ) ) . ( 4 ) The σT ( . ) function also appears in subsequent equations and merits some discussion . First note that σT ( wkl ) as given by ( 4 ) is the derivative of sigm ( ( w2kl − τl ) /T ) with respect to wkl . Since the latter approaches the step function ( located at w2kl = τl ) in the limit as T −→ 0 , it follows that the former , i.e. , σT ( wkl ) would approach a Dirac delta function , meaning that its value approaches zero everywhere except over the transitional region where it is inversely proportional to region ’ s width , i.e. , σT ( wkl ) ∼ 1 T , for |w2kl − τl| ∼ T. ( 5 )

Summary: The paper introduces a new type of soft threshold operator in conjunction with appropriate weight regularization that can be used in the context of neural network pruning to obtain sparse, performant networks from pre-trained, dense networks. The main idea is to replace the Heaviside step function that occurs in "hard threshold" pruning, which is non-differentiable, by a sigmoid function that can be differentiated and thus enables the efficient training/optimization of relevant pruning parameters. Pruning is hereby performed on a per-layer basis by training a regularized per-layer threshold.