Progressive Skeletonization: Trimming more fat from a network at initialization

1 INTRODUCTION . The majority of pruning algorithms for Deep Neural Networks require training dense models and often fine-tuning sparse sub-networks in order to obtain their pruned counterparts . In Frankle & Carbin ( 2019 ) , the authors provide empirical evidence to support the hypothesis that there exist sparse sub-networks that can be trained from scratch to achieve similar performance as the dense ones . However , their method to find such sub-networks requires training the full-sized model and intermediate sub-networks , making the process much more expensive . Recently , Lee et al . ( 2019 ) presented SNIP . Building upon almost a three decades old saliency criterion for pruning trained models ( Mozer & Smolensky , 1989 ) , they are able to predict , at initialization , the importance each weight will have later in training . Pruning at initialization methods are much cheaper than conventional pruning methods . Moreover , while traditional pruning methods can help accelerate inference tasks , pruning at initialization may go one step further and provide the same benefits at train time Elsen et al . ( 2020 ) . Wang et al . ( 2020 ) ( GRASP ) noted that after applying the pruning mask , gradients are modified due to non-trivial interactions between weights . Thus , maximizing SNIP criterion before pruning might be sub-optimal . They present an approximation to maximize the gradient norm after pruning , where they treat pruning as a perturbation on the weight matrix and use the first order Taylor ’ s approximation . While they show improved performance , their approximation involves computing a Hessian-vector product which is expensive both in terms of memory and computation . ∗Correspondence to pau @ robots.ox.ac.uk †www.europe.naverlabs.com We argue that both SNIP and GRASP approximations of the gradients after pruning do not hold for high pruning levels , where a large portion of the weights are removed at once . In this work , while we rely on the saliency criteria introduced by Mozer & Smolensky ( 1989 ) , we optimize what this saliency would be after pruning , rather than before . Hence , we name our criteria Foresight Connection sEnsitivity ( FORCE ) . We introduce two approximate procedures to progressively optimize our objective . The first , which turns out to be equivalent to applying SNIP iteratively , removes a small fraction of weights at each step and re-computes the gradients after each pruning round . This allows to take into account the intricate interactions between weights , re-adjusting the importance of connections at each step . The second procedure , which we name FORCE , is also iterative in nature , but contrary to the first , it allows pruned parameters to resurrect . Hence , it supports exploration , which otherwise is not possible in the case of iterative SNIP . Moreover , one-shot SNIP can be viewed as a particular case of using only one iteration . Empirically , we find that both SNIP and GRASP have a sharp drop in performance when targeting higher pruning levels . Surprisingly , they perform even worse than random pruning as can be seen in Fig 1 . In contrast , our proposed pruning procedures prove to be significantly more robust on a wide range of pruning levels . 2 RELATED WORK . Pruning trained models Most of the pruning works follow the train – prune – fine-tune cycle ( Mozer & Smolensky , 1989 ; LeCun et al. , 1990 ; Hassibi et al. , 1993 ; Han et al. , 2015 ; Molchanov et al. , 2017 ; Guo et al. , 2016 ) , which requires training the dense network until convergence , followed by multiple iterations of pruning and fine-tuning until a target sparsity is reached . Particularly , Molchanov et al . ( 2017 ) present a criterion very similar to Mozer & Smolensky ( 1989 ) and therefore similar to Lee et al . ( 2019 ) and our FORCE , but they focus on pruning whole neurons , and involve training rounds while pruning . Frankle & Carbin ( 2019 ) and Frankle et al . ( 2020 ) showed that it was possible to find sparse sub-networks that , when trained from scratch or an early training iteration , were able to match or even surpass the performance of their dense counterparts . Nevertheless , to find them they use a costly procedure based on Han et al . ( 2015 ) . All these methods rely on having a trained network , thus , they are not applicable before training . In contrast , our algorithm is able to find a trainable sub-network with randomly initialized weights . Making the overall pruning cost much cheaper and presenting an opportunity to leverage the sparsity during training as well . Induce sparsity during training Another popular approach has been to induce sparsity during training . This can be achieved by modifying the loss function to consider sparsity as part of the optimization ( Chauvin , 1989 ; Carreira-Perpiñán & Idelbayev , 2018 ; Louizos et al. , 2018 ) or by dynamically pruning during training ( Bellec et al. , 2018 ; Mocanu et al. , 2018 ; Mostafa & Wang , 2019 ; Dai et al. , 2019 ; Dettmers & Zettlemoyer , 2020 ; Lin et al. , 2020 ; Kusupati et al. , 2020 ; Evci et al. , 2019 ) . These methods are usually cheaper than pruning after training , but they still need to train the network to select the final sparse sub-network . We focus on finding sparse sub-networks before any weight update , which is not directly comparable . Pruning at initialization These methods present a significant leap with respect to other pruning methods . While traditional pruning mechanisms focused on bringing speed-up and memory reduction at inference time , pruning at initialization methods bring the same gains both at training and inference time . Moreover , they can be seen as a form of Neural Architecture Search ( Zoph & Le , 2016 ) to find more efficient network topologies . Thus , they have both a theoretical and practical interest . Lee et al . ( 2019 ) presented SNIP , a method to estimate , at initialization , the importance that each weight could have later during training . SNIP analyses the effect of each weight on the loss function when perturbed at initialization . In Lee et al . ( 2020 ) , the authors studied pruning at initialization from a signal propagation perspective , focusing on the initialization scheme . Recently , Wang et al . ( 2020 ) proposed GRASP , a different method based on the gradient norm after pruning and showed a significant improvement for higher levels of sparsity . However , neither SNIP nor GRASP perform sufficiently well when larger compressions and speed-ups are required and a larger fraction of the weights need to be pruned . In this paper , we analyse the approximations made by SNIP and GRASP , and present a more suitable solution to maximize the saliency after pruning . 3 PROBLEM FORMULATION : PRUNING AT INITIALIZATION . Given a dataset D = { ( xi , yi ) } ni=1 , the training of a neural network f parameterized by θ ∈ Rm can be written as minimizing the following empirical risk : arg min θ 1 n ∑ i L ( ( f ( xi ; θ ) ) , yi ) s.t . θ ∈ C , ( 1 ) where L and C denote the loss function and the constraint set , respectively . Unconstrained ( standard ) training corresponds to C = Rm . Assuming we have access to the gradients ( batch-wise ) of the empirical risk , an optimization algorithm ( e.g . SGD ) is generally used to optimize the above objective , that , during the optimization process , produces a sequence of iterates { θi } Ti=0 , where θ0 and θT denote the initial and the final ( optimal ) parameters , respectively . Given a target sparsity level of k < m , the general parameter pruning problem involves C with a constraint ‖θT ‖0 ≤ k , i.e. , the final optimal iterate must have a maximum of k non-zero elements . Note that there is no such constraint with the intermediate iterates . Pruning at initialization , the main focus of this work , adds further restrictions to the above mentioned formulation by constraining all the iterates to lie in a fixed subspace of C. Precisely , the constraints are to find an initialization θ0 such that ‖θ0‖0 ≤ k 1 , and the intermediate iterates are θi ∈ C̄ ⊂ C , ∀i ∈ { 1 , . . . , T } , where C̄ is the subspace of Rm spanned by the natural basis vectors { ej } j∈supp ( θ0 ) . Here , supp ( θ0 ) denotes the support of θ0 , i.e. , the set of indices with non-zero entries . The first condition defines the sub-network at initialization with k parameters , and the second fixes its topology throughout the training process . Since there are ( m k ) such possible sub-spaces , exhaustive search to find the optimal sub-space to optimize ( 1 ) is impractical as it would require training ( m k ) neural networks . Below we discuss two recent approaches that circumvent this problem by maximizing a hand-designed data-dependent objective function . These objectives are tailored to preserve some relationships between the parameters , the loss , and the dataset , that might be sufficient to obtain a reliable θ0 . For the ease of notation , we will use θ to denote the dense initialization . SNIP Lee et al . ( 2019 ) present a method based on the saliency criterion from Mozer & Smolensky ( 1989 ) . They add a key insight and show this criteria works surprisingly well to predict , at initialization , the importance each connection will have during training . The idea is to preserve the parameters that will have maximum impact on the loss when perturbed . Let c ∈ { 0 , 1 } m be a binary vector , and the Hadamard product . Then , the connection sensitivity in SNIP is computed as : g ( θ ) : = ∂L ( θ c ) ∂c ∣∣∣∣ c=1 = ∂L ( θ ) ∂θ θ . ( 2 ) Once g ( θ ) is obtained , the parameters corresponding to the top-k values of |g ( θ ) i| are then kept . Intuitively , SNIP favors those weights that are far from the origin and provide high gradients ( irrespective of the direction ) . We note that SNIP objective can be written as the following problem : max c S ( θ , c ) : = ∑ i∈supp ( c ) |θi ∇L ( θ ) i| s.t . c ∈ { 0 , 1 } m , ‖c‖0 = k. ( 3 ) It is trivial to note that the optimal solution to the above problem can be obtained by selecting the indices corresponding to the top-k values of |θi ∇L ( θ ) i| . 1In practice , as will be done in this work as well , a subset of a given dense initialization is found using some saliency criterion ( will be discussed soon ) , however , note that our problem statement is more general than that . GRASP Wang et al . ( 2020 ) note that the SNIP saliency is measuring the connection sensitivity of the weights before pruning , however , it is likely to change after pruning . Moreover , they argue that , at initialization , it is more important to preserve the gradient signal than the loss itself . They propose to use as saliency the gradient norm of the loss ∆L ( θ ) = ∇L ( θ ) T∇L ( θ ) , but measured after pruning . To maximize it , Wang et al . ( 2020 ) adopt the same approximation introduced in LeCun et al . ( 1990 ) and treat pruning as a perturbation on the initial weights . Their method is equivalent to solving : max c G ( θ , c ) : = ∑ { i : ci=0 } −θi [ Hg ] i s.t . c ∈ { 0 , 1 } m , ‖c‖0 = k. ( 4 ) Where H and g denote the Hessian and the gradient of the loss respectively .

Summary: The paper finds that at extreme sparsities (>95%), existing approaches to pruning neural networks at initialization devolve to worse than random pruning. The paper posits that this degenerate behavior is due to the fact that weights are pruned in groups, though the saliency metrics only capture pointwise changes. The paper presents a modified saliency metric based on SNIP, allowing for calculating salience of partially pruned networks; this in turn allows for applying an iterative version of SNIP, as well as a variant of iterative SNIP that allows for rejuvenation. These pruning techniques are evaluated, showing that they maintain accuracy at high sparsities.