Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients

1 INTRODUCTION . Pruning at initialization—where we remove weights from a model before training begins—is a recent and promising area of research that enables us to enjoy the benefits of pruning at training time , and which may aid our understanding of training deep neural networks . Frankle & Carbin ( 2019 ) provide empirical evidence for the existence of sparse sub-networks that can be trained from initialization and achieve accuracies comparable to the original network . These “ winning tickets ” were originally found in an iterative process where , in each iteration , the network is trained to full convergence followed by pruning a subset of the weights by magnitude . The values of the remaining weights are then rewound to their value at initialization , and the process is repeated iteratively until the desired sparsity level is achieved . This process , known as Lottery Ticket Rewinding ( LTR ) , is very compute-intensive and is prone to failures . For instance , Frankle et al . ( 2020 ) show better results by rewinding weights not all the way back to initialization , but to early stages of training instead . LTR is especially prone to failure for more difficult problems ( e.g. , training on ImageNet ) , where we must rewind weights to their state several epochs into training . A recent line of work proposes alternative practical solutions to identify these sub-networks before training begins , without the cost of retraining the network iteratively Lee et al . ( 2018 ) ; Wang et al . ( 2020 ) ; de Jorge et al . ( 2021 ) ; Tanaka et al . ( 2020 ) . This class of methods uses gradients to assess the importance of neural network weights . These gradients are often known as Synaptic Saliencies and are used to estimate the effect of pruning a single parameter in isolation on various objectives , typically the loss function . This objective is not so different from classical pruning-at-convergence methods , but the gradients for a well-trained model are small ; therefore these methods must inspect higher-order metrics such as the Hessian to estimate the pruning effect ( LeCun et al. , 1990 ; Hassibi & Stork , 1993 ) . Pruning at initialization is desirable because the benefits of pruning ( in terms of memory and speed ) can be reaped during training , rather than only at inference/deployment time . However , the performance of prune-at-init methods remains poor : the degradation in accuracy is still significant compared to training the full model and LTR , making these methods impractical for many real-world problems ( Frankle et al. , 2021 ) . In this paper , we identify a fundamental limitation in the objective formulation of current methods , namely that saliency criteria do not take into account the fact that the model is going to be trained after the pruning step . If our aim was to simply prune a subset of weights without affecting the loss , then these saliency criteria are estimating the correct objective . However , this estimate does not take into account that we are going to train the weights after we prune them . We need a metric that captures the trainability of the weights during the optimization steps , rather than a single myopic estimate . Many methods attempt to overcome this by pruning gradually and/or adding training steps between iterative pruning steps ( Zhu & Gupta , 2018 ; You et al. , 2020 ; de Jorge et al. , 2021 ) . Although this approach has been shown to be effective , it is expensive and cumbersome in practice and ultimately is an indirect approximation to the trainability criteria we are looking to incorporate into our objective . In this paper , we propose Prospect Pruning ( ProsPr ) , a new pruning-at-init method that learns from the first few steps of optimization which parameters to prune . We explicitly formulate our saliency criteria to account for the fact that the network will be trained after pruning . More precisely , ProsPr uses meta-gradients by backpropagating through the first few model updates in order to estimate the effect the initial pruning parameters have on the loss after a few gradient descent steps . Effectively this enables us to account for both higher-order effects of pruning weights on the loss , as well as the trainability of individual weights . Similar to other methods we apply pruning to initialization values of weights and train our models from scratch . In summary , our contributions are : • We identify a key limitation in prior saliency criteria for pruning neural networks—namely that they do not explicitly incorporate trainability-after-pruning into their criteria . • We propose a new pruning-at-init method , ProsPr , that uses meta-gradients over the first few training steps to bridge the gap between pruning and training . • We show empirically that ProsPr achieves higher accuracy compared to existing pruningat-init methods1 . Unlike other methods , our approach is single shot in the sense that the pruning is applied to the network initial weights in a single step . 2 BACKGROUND . In this section we review the key concepts that our method builds upon . We delay comparisons to other pruning techniques in the literature to Section 5 . Classic post-training pruning methods aim to identify and remove network weights with the least impact on the loss ( LeCun et al. , 1990 ; Hassibi & Stork , 1993 ) . They typically use the Taylor expansion of the loss with respect to parameters to define a saliency score for each parameter : δL ≈ ∇θL > δθ + 12δθ > H δθ , where H = ∇2θL is the Hessian matrix . When the network has converged , the first-order term in the expansion is negligible , and hence these methods resort to using H. Lee et al . ( 2018 ) introduce SNIP , and show that the same objective of minimizing the change in loss can be used at initialization to obtain a trainable pruned network . At initialization , the first-order gradients ∇θ in the local quadratic approximation are still significant , so higher-order terms can be ignored . Hence the computation of the parameter saliencies can be done using backpropagation . The Taylor expansion approximates the effect of small additive perturbations to the loss . To better approximate the effect of removing a weight , Lee et al . ( 2018 ) attach a multiplicative all-one mask to the computation graph of each weight . This does not change the forward-pass of the network , but it enables us to form the Taylor expansion around the mask values , rather than the weights , to estimate the effect of changing the mask values from 1 to 0 . More specifically , SNIP computes the saliency scores according to : sj = |gj ( w , D ) |∑m k=1 |gk ( w , D ) | , ( 1 ) with gj ( w , D ) = ∂L ( c w , D ) ∂cj , ( 2 ) 1Our code and models are publicly available at https : //anonymous-during-review where m is the number of weights in the network , c ∈ { 0 , 1 } m is the pruning mask ( initialised to 1 above ) , D is the training dataset , w are the neural network weights , L is the loss function , and is the Hadamard product . These saliency scores are computed before training the network , using one ( or more ) mini-batches from the training set . The global Top-K weights with the highest saliency scores are retained ( cj = 1 ) , and all other weights are pruned ( cj = 0 ) , before the network is trained . Our method , to be introduced in Section 3 , also relies on computing the saliency scores for each mask element , but uses a more sophisticated loss function to incorporate the notion of trainability . 3 OUR METHOD : PROSPR . In this section we introduce our method , Prospect Pruning ( ProsPr ) . We note that for the problem of pruning at initialization , the pruning step is immediately followed by training . Therefore , pruning should take into account the trainability of a weight , instead of only its immediate impact on the loss before training . In other words , we want to be able to identify weights that are not only important at initialization , but which may be useful for reducing the loss during training . To this end , we propose to estimate the effect of pruning on the loss over several steps of gradient descent at the beginning of training , rather than the changes in loss at initialization . More specifically , ProsPr models how training would happen by performing multiple ( M ) iterations of backpropagation and weight updates—like during normal training . We can then backpropagate through the entire computation graph , from the loss several steps into training , back to the original mask , since the gradient descent procedure is itself differentiable . Once the pruning mask is computed , we rewind the weights back to their values at initialization and train the pruned network . The gradient-of-gradients is called a meta-gradient . This algorithm is illustrated visually in Figure 1 . The higher-order information in the meta-gradient includes interactions between the weights during training . When pruning at initialization , our ultimate goal is to pick a pruned model , A , which is more trainable than an alternative pruned model B . That means we want the loss L ( ŷA , y ) to be lower than L ( ŷB , y ) at convergence ( for a fixed pruning ratio ) . Finding the optimal pruning mask is generally infeasible since the training horizon is long ( i.e. , evaluation is costly ) and the space of possible pruning masks is large . Unlike other methods that must compute the saliency scores iteratively , we can use the meta-gradients to compute the pruning mask in one shot . This picks a line in loss-space , which more closely predicts the eventual actual loss . This is because it smooths out over more steps , and takes into account interactions between weights in the training dynamics . Crucially , in the limit of large M , the match to the ultimate objective is exact . 3.1 SALIENCY SCORES VIA META-GRADIENTS . We now introduce ProsPr formally . After initialising the network weights randomly to obtain winit , we apply a weight mask to the initial weights , w0 = c winit . ( 3 ) This weight mask contains only ones , c = 1 , as in SNIP ( Lee et al. , 2018 ) , and represents the connectivity of the corresponding weights . We then sample M+1 batches of data Di ∼ Dtrain ( i ∈ { 0 , . . . , M } ; M ≥ 1 ) for the pruning step , and perform M weight updates2 , w1 = w0 − α∇w0L ( w0 , D0 ) ( 4 ) ... wM = wM−1 − α∇wM−1L ( wM−1 , DM−1 ) . ( 5 ) Then , we compute a meta-gradient that backpropagates through these updates . Specifically , we compute the gradient of the final loss w.r.t . the initial mask , ∇c L ( wM , DM ) . ( 6 ) Using the chain rule , we can write out the form of the meta-gradient beginning from the last step : ∇cL ( wM , D ) = ∇wML ( wM , D ) ( ∇cwM ) , ( 7 ) repeating for each step until we reach the zero ’ th step whose gradient is trivial , = ∇wML ( wM , D ) ( ∇wM−1wM ) . . . ( ∇w0w1 ) ( ∇cw0 ) ( 8 ) = ∇wML ( wM , D ) ( ∇wM−1wM ) . . . ( ∇w0w1 ) ( ∇c ( c winit ) ) ( 9 ) = ∇wML ( wM , D ) [ M∏ m=1 ( ∇wm−1wm ) ] winit . ( 10 ) In practice , we can compute the meta-gradients by relying on automatic differentiation software such as PyTorch ( Paszke et al. , 2019 ) . However , care must be taken to ensure that weights at each step are kept in memory so that the entire computation graph , including gradients , is visible to the automatic differentiation software . The saliency scores are now given by sj = |gj ( w , D ) |∑m k=1 |gk ( w , D ) | , ( 11 ) with gj ( w , D ) = ∂L ( wM , D ) ∂cj , ( 12 ) where wM is a function of c. Equation ( 12 ) stands in contrast to SNIP , where the saliency is computed using the loss at c ·winit rather than wM . The saliency scores are then used to prune the initial weights winit : the ones with the highest saliency scores are retained ( cj = 1 ) , and all other weights are pruned ( cj = 0 ) . Finally , the network is trained with the pruned weights ŵinit . Algorithm 1 summarises the proposed method , ProsPr .

Summary: This work studies the problem of pruning neural networks at initialization. It first identifies that the saliency score defined by the existing method SNIP has room for improvement. Specifically, the authors propose a method named prospect pruning to take into account the sequence of weight updates to determine the pruning mask. The experimental results on Tiny ImageNet and CIFAR show that the proposed method achieves better performance than existing methods of pruning at initialization.