Dynamic Model Pruning with Feedback

1 INTRODUCTION . Highly overparametrized deep neural networks show impressive results on machine learning tasks . However , with the increase in model size comes also the demand for memory and computer power at inference stage—two resources that are scarcely available on low-end devices . Pruning techniques have been successfully applied to remove a significant fraction of the network weights while preserving test accuracy attained by dense models . In some cases , the generalization of compressed networks has even been found to be better than with full models ( Han et al. , 2015 ; 2017 ; Mocanu et al. , 2018 ) . The sparsity of a network is the number of weights that are identically zero , and can be obtained by applying a sparsity mask on the weights . There are several different approaches to find sparse models . For instance , one-shot pruning strategies find a suitable sparsity mask by inspecting the weights of a pretrained network ( Mozer & Smolensky , 1989 ; LeCun et al. , 1990 ; Han et al. , 2017 ) . While these algorithms achieve a substantial size reduction of the network with little degradation in accuracy , they are computationally expensive ( training and refinement on the dense model ) , and they are outperformed by algorithms that explore different sparsity masks instead of a single one . In dynamic pruning methods , the sparsity mask is readjusted during training according to different criteria ( Mostafa & Wang , 2019 ; Mocanu et al. , 2018 ) . However , these methods require fine-tuning of many hyperparameters . We propose a new pruning approach to obtain sparse neural networks with state-of-the-art test accuracy . Our compression scheme uses a new saliency criterion that identifies important weights in the network throughout training to propose candidate masks . As a key feature , our algorithm not only evolves the pruned sparse model alone , but jointly also a ( closely related ) dense model that is used in a natural way to correct for pruning errors during training . This results in better generalization properties on a wide variety of tasks , since the simplicity of the scheme allows us further to study it from a theoretical point of view , and to provide further insights and interpretation . We do not require tuning of additional hyperparameters , and no retraining of the sparse model is needed ( though can further improve performance ) . Contributions .. • A novel dynamic pruning scheme , that incorporates an error feedback in a natural way Sec . 3 and finds a trained sparse model in one training pass . Sec . 5 • We demonstrate state-of-the-art performance ( in accuracy and sparsity ) , Sec . 5 ourperforming all previously proposed pruning schemes . Sec . 5 • We complement our results by an ablation study that provides further insights . Sec . 6 and convergence analysis for convex and non-convex objectives . Sec . 4 2 RELATED WORK . Previous works on obtaining pruned networks can ( loosely ) be divided into three main categories . Pruning after training . Training approaches to obtain sparse networks usually include a three stage pipeline—training of a dense model , one-shot pruning and fine-tuning—e.g. , ( Han et al. , 2015 ) . Their results ( i.e. , moderate sparsity level with minor quality loss ) made them the standard method for network pruning and led to several variations ( Guo et al. , 2016 ; Carreira-Perpinán & Idelbayev , 2018 ) . Pruning during training . Zhu & Gupta ( 2017 ) propose the use of magnitude-based pruning and to gradually increase the sparsity ratio while training the model from scratch . A pruning schedule determines when the new masks are computed ( extending and simplifying ( Narang et al. , 2017 ) ) . He et al . ( 2018 ) ( SFP ) prune entire filters of the model at the end of each epoch , but allow the pruned filters to be updated when training the model . Deep Rewiring ( DeepR ) ( Bellec et al. , 2018 ) allows for even more adaptivity by performing pruning and regrowth decisions periodically . This approach is computationally expensive and challenging to apply to large networks and datasets . Sparse evolutionary training ( SET ) ( Mocanu et al. , 2018 ) simplifies prune–regrowth cycles by using heuristics for random growth at the end of each training epoch and NeST ( Dai et al. , 2019 ) by inspecting gradient magnitudes . Dynamic Sparse Reparameterization ( DSR ) ( Mostafa & Wang , 2019 ) implements a prune– redistribute–regrowth cycle where target sparsity levels are redistributed among layers , based on loss gradients ( in contrast to SET , which uses fixed , manually configured , sparsity levels ) . Sparse Momentum ( SM ) ( Dettmers & Zettlemoyer , 2019 ) follows the same cycle but instead using the mean momentum magnitude of each layer during the redistribute phase . SM outperforms DSR on ImageNet for unstructured pruning by a small margin but has no performance difference on CIFAR experiments . Our approach also falls in the dynamic category but we use error compensation mechanisms instead of hand crafted redistribute–regrowth cycles . Pruning before training . Recently—spurred by the lottery ticket hypothesis ( LT ) ( Frankle & Carbin , 2019 ) —methods which try to find a sparse mask that can be trained from scratch have attracted increased interest . For instance , Lee et al . ( 2019 ) propose SNIP to find a pruning mask by inspecting connection sensitivities and identifying structurally important connections in the network for a given task . Pruning is applied at initialization , and the sparsity mask remains fixed throughout training . Note that Frankle & Carbin ( 2019 ) ; Frankle et al . ( 2019 ) do not propose an efficient pruning scheme to find the mask , instead they rely on iterative pruning , repeated for several full training passes . Further Approaches . Srinivas et al . ( 2017 ) ; Louizos et al . ( 2018 ) learn gating variables ( e.g . through ` 0 regularization ) that minimize the number of nonzero weights , recent parallel work studies filter pruning for pre-trained models ( You et al. , 2019 ) . Gal et al . ( 2017 ) ; Neklyudov et al . ( 2017 ) ; Molchanov et al . ( 2017 ) prune from Bayesian perspectives to learn dropout probabilities during training to prune and sparsify networks as dropout weight probabilities reach 1 . Gale et al . ( 2019 ) extensively study recent unstructured pruning methods on large-scale learning tasks , and find that complex techniques ( Molchanov et al. , 2017 ; Louizos et al. , 2018 ) perform inconsistently . Simple magnitude pruning approaches achieve comparable or better results ( Zhu & Gupta , 2017 ) . 3 METHOD . We consider the training of a non-convex loss function f : Rd → R. We assume for a weight vector w ∈ Rd to have access to a stochastic gradient g ( w ) ∈ Rd such that E [ g ( w ) ] = ∇f ( w ) . This corresponds to the standard machine learning setting with g ( w ) representing a ( mini-batch ) gradient of one ( or several ) components of the loss function . Stochastic Gradient Descent ( SGD ) computes a sequence of iterates by the update rule wt+1 : = wt − γtg ( wt ) , ( SGD ) for some learning rate γt . To obtain a sparse model , a general approach is to prune some of the weights of wt , i.e. , to set them to zero . Such pruning can be implemented by applying a mask m ∈ { 0 , 1 } d to the weights , resulting in a sparse model w̃t : = m wt , where denotes the entry-wise ( Hadamard ) product . The mask could potentially depend on the weights wt ( e.g. , smallest magnitude pruning ) , or depend on t ( e.g. , the sparsity is incremented over time ) . Before we introduce our proposed dynamic pruning scheme , we formalize the three main existing types of pruning methodologies ( summarized in Figure 1 ) . These approaches differ in the way the mask is computed , and the moment when it is applied.1 Pruning before training . A mask m0 ( depending on e.g . the initialization w0 or the network architecture of f ) is applied and ( SGD ) is used for training on the resulting subnetwork f̃ ( w ) : = f ( m0 w ) with the advantage that only pruned weights need to be stored and updated2 , and that by training with SGD a local minimum of the subnetwork f̃ ( but not of f—the original training target ) can be reached . In practice however , it remains a challenge to efficiently determine a good mask m0 and a wrongly chosen mask at the beginning strongly impacts the performance . 1The method introduced in Section 2 typically follow one of these broad themes loosely , with slight variations in detail . For the sake of clarity we omit a too technical and detailed discussion here . 2When training on f̃ ( w ) , it suffices to access stochastic gradients of f̃ ( w ) , denoted by g̃ ( w ) , which can potentially be cheaper be computed than by naively applying the mask to g ( w ) ( note g̃ ( w ) = m0 g ( w ) ) . Pruning after training ( one-shot pruning ) . A dense model is trained , and pruning is applied to the trained model wT . As the pruned model w̃T = mT wT is very likely not at a local optimum of f , fine-tuning ( retraining with the fixed mask mT ) is necessary to improve performance . Pruning during training ( incremental and dynamic pruning ) . Dynamic schemes change the mask mt every ( few ) iterations based on observations during training ( i.e . by observing the weights and stochastic gradients ) . Incremental schemes monotonically increase the sparsity pattern , fully dynamic schemes can also reactivate previously pruned weights . In contrast to previous dynamic schemes that relied on elaborated heuristics to adapt the mask mt , we propose a simpler approach : Dynamic pruning with feedback ( DPF , Algorithm 1 ) . Our scheme evaluates a stochastic gradient at the pruned model w̃t = mt wt and applies it to the ( simultaneously maintained ) dense model wt : wt+1 : = wt − γtg ( mt wt ) = wt − γtg ( w̃t ) . ( DPF ) Applying the gradient to the full model allows to recover from “ errors ” , i.e . prematurely masking out important weights : when the accumulated gradient updates from the following steps drastically change a specific weight , it can become activated again ( in contrast to incremental pruning approaches that have to stick to sub-optimal decisions ) . For illustration , observe that ( DPF ) can equivalently be written as wt+1 = wt − γtg ( wt + et ) , where et : = w̃t −wt is the error produced by the compression . This provides a different intuition of the behavior of ( DPF ) , and connects it with the concept of error-feedback ( Stich et al. , 2018 ; Karimireddy et al. , 2019 ) .3 We illustrate this principle in Figure 2 and give detailed pseudocode and further implementation details in Appendix A.1 . The DPF scheme can also be seen as an instance of a more general class of schemes that apply ( arbitrary ) perturbed gradient updates to the dense model . For instance straight-through gradient estimators ( Bengio et al. , 2013 ) that are used to empirically simplify the backpropagation can be seen as such perturbations . Our stronger assumptions on the structure of the perturbation allow to derive non-asymptotic convergence rates in the next section , though our analysis could also be extended to the setting in ( Yin et al. , 2019 ) if the perturbations can be bounded .

Summary: This work proposes a simple pruning method that dynamically sparsifies the network during training. This is achieved by performing at fixed intervals magnitude based pruning for either individual weights or entire neurons. While similar methods have been explored before, this work proposes a slight twist; instead of updating the weights of the model by following the gradient of the parameters of the dense model, they update the parameters of the dense model according to the gradients of the sparse model. Essentially, this corresponds to a variant of the straight-through estimator [1], where in the forward pass we evaluate the compressed model, but in the backward pass we update the model as if the compression didn’t take place. The authors argue that this process allows for ``feedback” in the pruning mechanism, as the pruned weights still receive gradient updates hence they can be ``re-activated” at later stages of training. They then provide a convergence analysis about the optimization procedure with such a gradient, and show that for strongly convex functions the method converges in the vicinity of the global optimum, whereas for non-convex functions it converges to the neighbourhood of a stationary point. Finally, the authors perform extensive experimental evaluation and show that their method is better than the baselines that they considered.