Revisiting Loss Modelling for Unstructured Pruning

1 INTRODUCTION . Neural networks are getting bigger , requiring more and more computational resources not only for training , but also when used for inference . However , resources are sometimes limited , especially on mobile devices and low-power chips . In unstructured pruning , the goal is to remove some parameters ( i.e . setting them to zeros ) , while still maintaining good prediction performances . This is fundamentally a combinatorial optimization problem which is intractable even for small scale neural networks , and thus various heuristics have been developed to prune the model either before training ( Lee et al. , 2019b ; Wang et al. , 2020 ) , during training ( Louizos et al. , 2017 ; Molchanov et al. , 2017 ; Ding et al. , 2019 ) , or in an iterative training/fine-tuning fashion ( LeCun et al. , 1990 ; Hassibi & Stork , 1993 ; Han et al. , 2015 ; Frankle & Carbin , 2018 ; Renda et al. , 2020 ) . Early pruning work Optimal Brain Damage ( OBD ) ( LeCun et al. , 1990 ) , and later Optimal Brain Surgeon ( OBS ) ( Hassibi & Stork , 1993 ) , proposed to estimate the importance of each parameter by approximating the effect of removing it , using the second order term of a Taylor expansion of the loss function around converged parameters . This type of approach involves computing the Hessian , which is challenging to compute since it scales quadratically with the number of parameters in the network . Several approximations have thus been explored in the literature ( LeCun et al. , 1990 ; Hassibi & Stork , 1993 ; Heskes , 2000 ; Zeng & Urtasun , 2019 ; Wang et al. , 2019 ) . However , state-ofthe-art unstructured pruning methods typically rely on Magnitude Pruning ( MP ) ( Han et al. , 2015 ) , a simple and computationally cheap criterion based on weight magnitude , that works extremely well in practice ( Renda et al. , 2020 ) . This paper revisits linear and diagonal quadratic models of the local loss landscape for unstructured pruning . In particular , since these models are local approximations and thus assume that pruning steps correspond to small vectors in parameter space , we propose to investigate how this locality assumption affects their performance . Moreover , we show that the convergence assumption behind OBD and OBS , which is overlooked and violated in current methods , can be relaxed by maintaining the gradient term in the quadratic model . Finally , to prevent having to compute second order information , we propose to compare diagonal quadratic models to simpler linear models . While our empirical study demonstrates that pruning criteria based on linear and quadratic loss models are good at preserving the training loss , it also shows that this benefit does not necessarily transfer to better networks after fine-tuning , suggesting that preserving the loss might not be the best objective to optimize for . Our contributions can be summarized as follows : 1 . We present pruning criteria based on both linear and diagonal quadratic models of the loss , and show how they compare at preserving training loss compared to OBD and MP . 2 . We study two strategies to better enforce locality in the pruning steps , pruning in several stages and regularising the step size , and show how they improve the quality of the criteria . 3 . We show that using pruning criteria that are better at preserving the loss does not necessarily transfer to better fine-tuned networks , raising questions about the adequacy of such criteria . 2 BACKGROUND : UNSTRUCTURED PRUNING . 2.1 UNSTRUCTURED PRUNING PROBLEM FORMULATION . For a given architecture , neural networks are a family of functions fθ : X → Y from an input space X to an output space Y , where θ ∈ RD is the vector that contains all the parameters of the network . Neural networks are usually trained by seeking parameters θ that minimize the empirical risk L ( θ ) = 1N ∑ i ` ( fθ ( xi ) , ti ) of a loss function ` on a training dataset D = { ( xi , ti ) } 1≤i≤N , composed of N ( example , target ) pairs . The goal of unstructured pruning is to find a step ∆θ to add to the current parameters θ such that ‖θ+∆θ‖0 = ( 1−κ ) D , i.e . the parameter vector after pruning is of desired sparsity κ ∈ [ 0 , 1 ] . While doing so , the performance of the pruned network should be maintained , so L ( θ + ∆θ ) should not differ much from L ( θ ) . Unstructured pruning thus amounts to the following minimization problem : minimize ∆θ ∆L ( θ , ∆θ ) def= |L ( θ + ∆θ ) − L ( θ ) | s.t . ‖θ + ∆θ‖0 = ( 1− κ ) D ( 1 ) Directly solving this problem would require evaluating L ( θ + ∆θ ) for all possible values of ∆θ , which is prohibitively expensive , so one needs to rely on heuristics to find good solutions . 2.2 OPTIMAL BRAIN DAMAGE CRITERION . Optimal Brain Damage ( OBD ) ( LeCun et al. , 1990 ) proposes to use a quadratic modelling of L ( θ + ∆θ ) , leading to the following approximation of ∆L ( θ , ∆θ ) : ∆LQM ( θ , ∆θ ) = ∣∣∣∣∣∂L ( θ ) ∂θ > ∆θ + 12∆θ > H ( θ ) ∆θ ∣∣∣∣∣ ( 2 ) where H ( θ ) is the Hessian of L ( θ ) . H ( θ ) being intractable , even for small-scale networks , its Generalized Gauss-Newton approximation G ( θ ) ( Schraudolph , 2002 ) is used in practice , as detailed in Appendix A.1 Then , two more approximations are made : first , it assumes the training of the network has converged , thus the gradient of the loss wrt θ is 0 , which makes the linear term vanish . Then , it neglects the interactions between parameters , which corresponds to a diagonal approximation of G ( θ ) , leading to the following model : ∆LOBD ( θ , ∆θk ) ≈ 1 2 Gkk ( θ ) ∆θ 2 k ⇒ sOBDk = 1 2 Gkk ( θ ) θ 2 k ( 3 ) sOBDk is the saliency of each parameter , estimating how much the loss will change if that parameter is pruned , so if ∆θk = −θk . Parameters can thus be ranked by order of importance , and the ones with the smallest saliencies ( i.e . the least influence on the loss ) are pruned , while the ones with the biggest saliencies are kept unchanged . This can be interpreted as finding and applying a binary mask m ∈ { 0 , 1 } D to the parameters such that θ + ∆θ = θ m , where is the element-wise product . 2.3 MAGNITUDE PRUNING CRITERION . Magnitude Pruning ( MP ) ( Han et al. , 2015 ) , is a popular pruning criterion in which the saliency is simply based on the norm of the parameter : sMPk = θ 2 k ( 4 ) Despite its simplicity , MP works extremely well in practice ( Gale et al. , 2019 ) , and is used in current state-of-the-art methods ( Renda et al. , 2020 ) . We use global MP as baseline in all our experiments . 1Although LeCun et al . ( 1990 ) uses H ( θ ) in the equations of OBD , it is actually G ( θ ) which was used in practice ( LeCun , 2007 ) . 2.4 OPTIMAL BRAIN SURGEON . Optimal Brain Surgeon ( OBS ) ( Hassibi & Stork , 1993 ) also relies on the quadratic model in Equation 2 to solve the minimization problem given in Equation 1 , but uses the Lagrangian formulation to include the constraint to the solution of the minimization problem . Since OBS requires to compute the inverse of H ( θ ) , several approximations have been explored in the literature , including diagonal , as in the original OBS , Kronecker-factored ( Martens & Grosse , 2015 ) as in ML-Prune ( Zeng & Urtasun , 2019 ) , or diagonal , but in an Kronecker-factored Eigenbasis ( George et al. , 2018 ) , as in EigenDamage ( Wang et al. , 2019 ) . While we use OBD in our demonstrations and experimental setup , everything presented in this paper can also be used in OBS-based methods . We leave that for future work . 3 REVISITING LOSS MODELLING FOR UNSTRUCTURED PRUNING . In this work , we investigate linear and diagonal quadratic models of the loss function and their performance when used for pruning neural networks . In our empirical study , we aim at answering the following questions : 1 . How do criteria based on weight magnitude , or linear or quadratic models compare at preserving training loss ( i.e . at solving the minimization problem in Equation 1 ) ? 2 . How does the locality assumption behind criteria based on linear and quadratic models affect their performances ? 3 . Do pruning criteria that are better at preserving the loss lead to better fine-tuned networks ? We now describe the linear and quadratic models we use , as well as the strategies to enforce locality of the pruning steps . 3.1 LINEAR AND QUADRATIC MODELS . In current training strategies , regularization techniques such as early stopping or dropout ( Srivastava et al. , 2014 ) are often used to counteract overfitting . In these setups , there is no reason to assume that the training has converged , implying that the linear term in the Taylor expansion should not be neglected . Thus , one can build a pruning criterion similar to OBD that includes the gradient term in the quadratic model from Equation 2 , leading to the following saliencies:2 ∆LQM ( θ , ∆θk ) ≈ ∣∣∣∣∣∂L ( θ ) ∂θk > ∆θk + 1 2 Gkk ( θ ) ∆θ 2 k ∣∣∣∣∣⇒ sQMk = ∣∣∣∣−∂L ( θ ) ∂θk θk + 12Gkk ( θ ) θ2k ∣∣∣∣ ( 5 ) Recall the constraint ∆θk ∈ { −θk , 0 } , hence the saliencies . This criterion generalizes OBD for networks that are not at convergence , and provides similar saliencies for networks that have converged . To avoid the computational cost associated with computing second order information , which is prohibitive for large scale neural networks , one can use a simpler linear model ( LM ) instead of a quadratic one to approximate ∆L ( θ , ∆θ ) , leading to the following approximation and saliencies : ∆LLM ( θ , ∆θ ) = ∣∣∣∣∣∂L ( θ ) ∂θ > ∆θ ∣∣∣∣∣ ⇒ sLMk = ∣∣∣∣∂L ( θ ) ∂θk θk ∣∣∣∣ ( 6 ) The saliencies of the linear model are very related to the criterion used in Single-shot Network Pruning ( Lee et al. , 2019b ) , as demonstrated by Wang et al . ( 2020 ) . 3.2 ENFORCING LOCALITY . One important point to keep in mind is that linear and quadratic models ( whether diagonal or not ) are local approximations , and are generally only faithful in a small neighbourhood of the current parameters . Explicitly showing the terms that are neglected , we have : ∆L ( θ , ∆θ ) = ∆LLM ( θ , ∆θ ) +O ( ‖∆θ‖22 ) = ∆LQM ( θ , ∆θ ) +O ( ‖∆θ‖32 ) ( 7 ) 2Note that this idea has been explored for OBS ( Singh & Alistarh , 2020 ) , as well as for structured prun- ing ( Molchanov et al. , 2019 ) So when approximating ∆L with ∆LLM we neglect the terms in O ( ‖∆θ‖22 ) , and when approximating ∆L with ∆LQM we neglect the terms in O ( ‖∆θ‖32 ) . Both approximations are thus only valid in a small neighbourhood of θ , and are extremely likely to be wrong when ‖∆θ‖2 is large . We list here different tricks to prevent this from happening . Performing the Pruning in several Stages ‖∆θ‖2 can be large when a large portion of the parameters is pruned at once . An easy fix typically used to mitigate this issue is to perform the pruning in several stages , re-estimating the model at each stage . The number of stages , which we denote by π , is typically overlooked ( e.g . both Zeng & Urtasun ( 2019 ) and Wang et al . ( 2019 ) use only 6 stages of pruning ) . Our experiments , in upcoming Section 5 , show that it has a drastic impact on the performances . Note that , without fine-tuning phases between the different pruning stages , this strategy violates the convergence assumption behind OBD and OBS , since after the first stage of pruning the network is no more at convergence . The sparsity at each stage can be increased following either a linear schedule , where each step prunes the same number of parameter , or an exponential schedule , where the number of parameters pruned at each stage gets smaller and smaller . The later is typically used in the literature ( Zeng & Urtasun , 2019 ; Wang et al. , 2019 ; Frankle & Carbin , 2018 ; Renda et al. , 2020 ) . We compare them in Section 5 . Constraining the Step Size As is often done when using quadratic models ( e.g . Nocedal & Wright ( 2006 ) ) , one can penalize the model when it decides to take steps that are too large , in order to stay in a region where we can trust the model . This can be done by simply adding the norm penalty λ2 ‖θk‖ 2 2 to the saliencies computed by any criterion ( Equations 3 , 5 or 6 ) , where λ is a hyper-parameter that controls the strength of the constraint : a small value of λ leaves the saliencies unchanged , and a large value of λ transforms the pruning criterion into MP ( Equation 4 ) . Other Considerations ‖∆θ‖2 can be large if θ is large itself . This is dependent on the training procedure of the network , but can be easily mitigated by constraining the norm of the weights , which can be done using L2 regularisation or weight decay . Since nowadays weight decay is almost systematically used by default when training networks ( e.g . He et al . ( 2016b ) ; Xie et al . ( 2017 ) ; Devlin et al . ( 2018 ) ) , we do not investigate this further .

Summary: The authors study the use of loss-modeling to maintain model quality when inducing unstructured sparsity in deep neural networks. They study a range of different approximations and modifications that can help improve the quality of the approximation (taking local steps, avoid large changes in weight magnitude, avoiding assumptions about convergence). The authors conduct a thorough empirical investigation that yields practical observations for the design of future pruning techniques.