Robust Pruning at Initialization

1 INTRODUCTION . Overparameterized deep NNs have achieved state of the art ( SOTA ) performance in many tasks ( Nguyen and Hein , 2018 ; Du et al. , 2019 ; Zhang et al. , 2016 ; Neyshabur et al. , 2019 ) . However , it is impractical to implement such models on small devices such as mobile phones . To address this problem , network pruning is widely used to reduce the time and space requirements both at training and test time . The main idea is to identify weights that do not contribute significantly to the model performance based on some criterion , and remove them from the NN . However , most pruning procedures currently available can only be applied after having trained the full NN ( LeCun et al. , 1990 ; Hassibi et al. , 1993 ; Mozer and Smolensky , 1989 ; Dong et al. , 2017 ) although methods that consider pruning the NN during training have become available . For example , Louizos et al . ( 2018 ) propose an algorithm which adds a L0 regularization on the weights to enforce sparsity while Carreira-Perpiñán and Idelbayev ( 2018 ) ; Alvarez and Salzmann ( 2017 ) ; Li et al . ( 2020 ) propose the inclusion of compression inside training steps . Other pruning variants consider training a secondary network that learns a pruning mask for a given architecture ( Li et al . ( 2020 ) ; Liu et al . ( 2019 ) ) . Recently , Frankle and Carbin ( 2019 ) have introduced and validated experimentally the Lottery Ticket Hypothesis which conjectures the existence of a sparse subnetwork that achieves similar performance to the original NN . These empirical findings have motivated the development of pruning at initialization such as SNIP ( Lee et al . ( 2018 ) ) which demonstrated similar performance to classical pruning methods of pruning-after-training . Importantly , pruning at initialization never requires training the complete NN and is thus more memory efficient , allowing to train deep NN using limited computational resources . However , such techniques may suffer from different problems . In particular , nothing prevents such methods from pruning one whole layer of the NN , making it untrainable . More generally , it is typically difficult to train the resulting pruned NN ( Li et al. , 2018 ) . To solve this situation , Lee et al . ( 2020 ) try to tackle this issue by enforcing dynamical isometry using orthogonal weights , while Wang et al . ( 2020 ) ( GraSP ) uses Hessian based pruning to preserve gradient flow . Other work by Tanaka et al . ( 2020 ) considers a data-agnostic iterative approach using the concept of synaptic flow in order to avoid the layer-collapse phenomenon ( pruning a whole layer ) . In our work , we use principled scaling and re-parameterization to solve this issue , and show numerically that our algorithm achieves SOTA performance on CIFAR10 , CIFAR100 , TinyImageNet and ImageNet in some scenarios and remains competitive in others . In this paper , we provide novel algorithms for Sensitivity-Based Pruning ( SBP ) , i.e . pruning schemes that prune a weight W based on the magnitude of |W ∂L∂W | at initialization where L is the loss . Experimentally , compared to other available one-shot pruning schemes , these algorithms provide state-of the-art results ( this might not be true in some regimes ) . Our work is motivated by a new theoretical analysis of gradient back-propagation relying on the mean-field approximation of deep NN ( Hayou et al. , 2019 ; Schoenholz et al. , 2017 ; Poole et al. , 2016 ; Yang and Schoenholz , 2017 ; Xiao et al. , 2018 ; Lee et al. , 2018 ; Matthews et al. , 2018 ) . Our contribution is threefold : • For deep fully connected FeedForward NN ( FFNN ) and Convolutional NN ( CNN ) , it has been previously shown that only an initialization on the so-called Edge of Chaos ( EOC ) make models trainable ; see e.g . ( Schoenholz et al. , 2017 ; Hayou et al. , 2019 ) . For such models , we show that an EOC initialization is also necessary for SBP to be efficient . Outside this regime , one layer can be fully pruned . • For these models , pruning pushes the NN out of the EOC making the resulting pruned model difficult to train . We introduce a simple rescaling trick to bring the pruned model back in the EOC regime , making the pruned NN easily trainable . • Unlike FFNN and CNN , we show that Resnets are better suited for pruning at initialization since they ‘ live ’ on the EOC by default ( Yang and Schoenholz , 2017 ) . However , they can suffer from exploding gradients , which we resolve by introducing a re-parameterization , called ‘ Stable Resnet ’ ( SR ) . The performance of the resulting SBP-SR pruning algorithm is illustrated in Table 1 : SBP-SR allows for pruning up to 99.5 % of ResNet104 on CIFAR10 while still retaining around 87 % test accuracy . The precise statements and proofs of the theoretical results are given in the Supplementary . Appendix H also includes the proof of a weak version of the Lottery Ticket Hypothesis ( Frankle and Carbin , 2019 ) showing that , starting from a randomly initialized NN , there exists a subnetwork initialized on the EOC . 2 SENSITIVITY PRUNING FOR FFNN/CNN AND THE RESCALING TRICK . 2.1 SETUP AND NOTATIONS . Let x be an input in Rd . A NN of depth L is defined by yl ( x ) = Fl ( W l , yl−1 ( x ) ) +Bl , 1 ≤ l ≤ L , ( 1 ) where yl ( x ) is the vector of pre-activations , W l and Bl are respectively the weights and bias of the lth layer and Fl is a mapping that defines the nature of the layer . The weights and bias are initialized with W l iid∼ N ( 0 , σ2w/vl ) , where vl is a scaling factor used to control the variance of yl , and Bl iid∼ N ( 0 , σ2b ) . Hereafter , Ml denotes the number of weights in the lth layer , φ the activation function and [ m : n ] : = { m , m+ 1 , ... , n } for m ≤ n. Two examples of such architectures are : • Fully connected FFNN . For a FFNN of depth L and widths ( Nl ) 0≤l≤L , we have vl = Nl−1 , Ml = Nl−1Nl and y1i ( x ) = d∑ j=1 W 1ijxj +B 1 i , y l i ( x ) = Nl−1∑ j=1 W lijφ ( y l−1 j ( x ) ) +B l i for l ≥ 2 . ( 2 ) • CNN . For a 1D CNN of depth L , number of channels ( nl ) l≤L , and number of neurons per channel ( Nl ) l≤L , we have y1i , α ( x ) = nl−1∑ j=1 ∑ β∈kerl W 1i , j , βxj , α+β+b 1 i , y l i , α ( x ) = nl−1∑ j=1 ∑ β∈kerl W li , j , βφ ( y l−1 j , α+β ( x ) ) +b l i , for l ≥ 2 , ( 3 ) where i ∈ [ 1 : nl ] is the channel index , α ∈ [ 0 : Nl−1 ] is the neuron location , kerl = [ −kl : kl ] is the filter range , and 2kl + 1 is the filter size . To simplify the analysis , we assume hereafter that Nl = N and kl = k for all l. Here , we have vl = nl−1 ( 2k + 1 ) and Ml = nl−1nl ( 2k + 1 ) . We assume periodic boundary conditions ; so yli , α = y l i , α+N = y l i , α−N . Generalization to multidimensional convolutions is straightforward . When no specific architecture is mentioned , ( W li ) 1≤i≤Ml denotes the weights of the l th layer . In practice , a pruning algorithm creates a binary mask δ over the weights to force the pruned weights to be zero . The neural network after pruning is given by yl ( x ) = Fl ( δl ◦W l , yl−1 ( x ) ) +Bl , ( 4 ) where ◦ is the Hadamard ( i.e . element-wise ) product . In this paper , we focus on pruning at initialization . The mask is typically created by using a vector gl of the same dimension as W l using a mapping of choice ( see below ) , we then prune the network by keeping the weights that correspond to the top k values in the sequence ( gli ) i , l where k is fixed by the sparsity that we want to achieve . There are three popular types of criteria in the literature : •Magnitude based pruning ( MBP ) : We prune weights based on the magnitude |W | . • Sensitivity based pruning ( SBP ) : We prune the weights based on the values of |W ∂L∂W | where L is the loss . This is motivated by LW ≈ LW=0 +W ∂L∂W used in SNIP ( Lee et al . ( 2018 ) ) . •Hessian based pruning ( HBP ) : We prune the weights based on some function that uses the Hessian of the loss function as in GraSP ( Wang et al. , 2020 ) . In the remainder of the paper , we focus exclusively on SBP while our analysis of MBP is given in Appendix E. We leave HBP for future work . However , we include empirical results with GraSP ( Wang et al. , 2020 ) in Section 4 . Hereafter , we denote by s the sparsity , i.e . the fraction of weights we want to prune . Let Al be the set of indices of the weights in the lth layer that are pruned , i.e . Al = { i ∈ [ 1 : Ml ] , s.t . δli = 0 } . We define the critical sparsity scr by scr = min { s ∈ ( 0 , 1 ) , s.t . ∃l , |Al| = Ml } , where |Al| is the cardinality of Al . Intuitively , scr represents the maximal sparsity we are allowed to choose without fully pruning at least one layer . scr is random as the weights are initialized randomly . Thus , we study the behaviour of the expected value E [ scr ] where , hereafter , all expectations are taken w.r.t . to the random initial weights . This provides theoretical guidelines for pruning at initialization . For all l ∈ [ 1 : L ] , we define αl by vl = αlN where N > 0 , and ζl > 0 such that Ml = ζlN2 , where we recall that vl is a scaling factor controlling the variance of yl and Ml is the number of weights in the lth layer . This notation assumes that , in each layer , the number of weights is quadratic in the number of neurons , which is satisfied by classical FFNN and CNN architectures . 2.2 SENSITIVITY-BASED PRUNING ( SBP ) . SBP is a data-dependent pruning method that uses the data to compute the gradient with backpropagation at initialization ( one-shot pruning ) .We randomly sample a batch and compute the gradients of the loss with respect to each weight . The mask is then defined by δli = I ( |W li ∂L∂W li | ≥ ts ) , where ts = |W ∂L∂W | ( ks ) and ks = ( 1− s ) ∑ lMl and |W ∂L ∂W | ( ks ) is the kths order statistics of the sequence ( |W li ∂L∂W li | ) 1≤l≤L,1≤i≤Ml . However , this simple approach suffers from the well-known exploding/vanishing gradients problem which renders the first/last few layers respectively susceptible to be completely pruned . We give a formal definition to this problem . Definition 1 ( Well-conditioned & ill-conditioned NN ) . Let ml = E [ |W l1 ∂L∂W l1 | 2 ] for l ∈ [ 1 : L ] . We say that the NN is well-conditioned if there exist A , B > 0 such that for all L ≥ 1 and l ∈ [ 1 : L ] we have A ≤ ml/mL ≤ B , and it is ill-conditioned otherwise . Understanding the behaviour of gradients at initialization is thus crucial for SBP to be efficient . Using a mean-field approach , such analysis has been carried out in ( Schoenholz et al. , 2017 ; Hayou et al. , 2019 ; Xiao et al. , 2018 ; Poole et al. , 2016 ; Yang , 2019 ) where it has been shown that an initialization known as the EOC is beneficial for DNN training . The mean-field analysis of DNNs relies on two standard approximations that we will also use here . Approximation 1 ( Mean-Field Approximation ) . When Nl 1 for FFNN or nl 1 for CNN , we use the approximation of infinitely wide NN . This means infinite number of neurons per layer for fully connected layers and infinite number of channels per layer for convolutional layers . Approximation 2 ( Gradient Independence ) . The weights used for forward propagation are independent from those used for back-propagation . These two approximations are ubiquitous in literature on the mean-field analysis of neural networks . They have been used to derive theoretical results on signal propagation ( Schoenholz et al. , 2017 ; Hayou et al. , 2019 ; Poole et al. , 2016 ; Yang , 2019 ; Yang and Schoenholz , 2017 ; Yang et al. , 2019 ) and are also key tools in the derivation of the Neural Tangent Kernel ( Jacot et al. , 2018 ; Arora et al. , 2019 ; Hayou et al. , 2020 ) . Approximation 1 simplifies the analysis of the forward propagation as it allows the derivation of closed-form formulas for covariance propagation . Approximation 2 does the same for back-propagation . See Appendix A for a detailed discussion of these approximations . Throughout the paper , we provide numerical results that substantiate the theoretical results that we derive using these two approximations . We show that these approximations lead to excellent match between theoretical results and numerical experiments . Edge of Chaos ( EOC ) : For inputs x , x′ , let cl ( x , x′ ) be the correlation between yl ( x ) and yl ( x′ ) . From ( Schoenholz et al. , 2017 ; Hayou et al. , 2019 ) , there exists a so-called correlation function f that depends on ( σw , σb ) such that cl+1 ( x , x′ ) = f ( cl ( x , x′ ) ) . Let χ ( σb , σw ) = f ′ ( 1 ) . The EOC is the set of hyperparameters ( σw , σb ) satisfying χ ( σb , σw ) = 1 . When χ ( σb , σw ) > 1 , we are in the Chaotic phase , the gradient explodes and cl ( x , x′ ) converges exponentially to some c < 1 for x 6= x′ and the resulting output function is discontinuous everywhere . When χ ( σb , σw ) < 1 , we are in the Ordered phase where cl ( x , x′ ) converges exponentially fast to 1 and the NN outputs constant functions . Initialization on the EOC allows for better information propagation ( see Supplementary for more details ) . Hence , by leveraging the above results , we show that an initialization outside the EOC will lead to an ill-conditioned NN . Theorem 1 ( EOC Initialization is crucial for SBP ) . Consider a NN of type ( 2 ) or ( 3 ) ( FFNN or CNN ) . Assume ( σw , σb ) are chosen on the ordered phase , i.e . χ ( σb , σw ) < 1 , then the NN is ill-conditioned . Moreover , we have E [ scr ] ≤ 1 L ( 1 + log ( κLN2 ) κ ) +O ( 1 κ2 √ LN2 ) , where κ = | logχ ( σb , σw ) |/8 . If ( σw , σb ) are on the EOC , i.e . χ ( σb , σw ) = 1 , then the NN is well-conditioned . In this case , κ = 0 and the above upper bound no longer holds . The proof of Theorem 1 relies on the behaviour of the gradient norm at initialization . On the ordered phase , the gradient norm vanishes exponentially quickly as it back-propagates , thus resulting in an ill-conditioned network . We use another approximation for the sake of simplification of the proof ( Approximation 3 in the Supplementary ) but the result holds without this approximation although the resulting constants would be a bit different . Theorem 1 shows that the upper bound decreases the farther χ ( σb , σw ) is from 1 , i.e . the farther the initialization is from the EOC . For constant width FFNN with L = 100 , N = 100 and κ = 0.2 , the theoretical upper bound is E [ scr ] / 27 % while we obtain E [ scr ] ≈ 22 % based on 10 simulations . A similar result can be obtained when the NN is initialized on the chaotic phase ; in this case too , the NN is ill-conditioned . To illustrate these results , Figure 1 shows the impact of the initialization with sparsity s = 70 % . The dark area in Figure 1 ( b ) corresponds to layers that are fully pruned in the chaotic phase due to exploding gradients . Using an EOC initialization , Figure 1 ( a ) shows that pruned weights are well distributed in the NN , ensuring that no layer is fully pruned .

Summary: The theoretical analysis is clearly stated in an well-organized way and the derived sparsity bound is reasonable. With FFNN and CNN, a theorem is given to show that the model is trainable only when the initialization on Edge of Chaos (EOC) and also provided a rescaling method to make the pruned NN into EOC regime. With Resnet, it proves the pruning satisfies the EOC condition by default and further provides re-parameterization method to tackle exploding gradients. The experiments well support theoretical results for both FFNN/CNN and resNet.