Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients
Authors: Milad Alizadeh
Shyam A. Tailor
Luisa M Zintgraf
Joost van Amersfoort
Sebastian Farquhar
Nicholas Donald Lane
Yarin Gal
Venue: ICLR 2022 Poster (2022)
Keywords: pruning
lottery ticket hypothesis
pruning at initialization

Abstract: Pruning neural networks at initialization would enable us to find sparse models that retain the accuracy of the original network while consuming fewer computational resources for training and inference. However, current methods are insufficient to enable this optimization and lead to a large degradation in model performance. In this paper, we identify a fundamental limitation in the formulation of current methods, namely that their saliency criteria look at a single step at the start of training without taking into account the trainability of the network. While pruning iteratively and gradually has been shown to improve pruning performance, explicit consideration of the training stage that will immediately follow pruning has so far been absent from the computation of the saliency criterion. To overcome the short-sightedness of existing methods, we propose Prospect Pruning (ProsPr), which uses meta-gradients through the first few steps of optimization to determine which weights to prune. ProsPr combines an estimate of the higher-order effects of pruning on the loss and the optimization trajectory to identify the trainable sub-network. Our method achieves state-of-the-art pruning performance on a variety of vision classification tasks, with less data and in a single shot compared to existing pruning-at-initialization methods.

Published as a conference paper at ICLR 2022
PROSPECT PRUNING: FINDING TRAINABLE WEIGHTS AT INITIALIZATION USING META- GRADIENTS
Milad Alizadeh 112 Shyam A. Tailor Luisa Zintgraf Joost van Amersfoort
Sebastian Farquhar
Nicholas Donald Lane Yarin Gal
1ATMLL, Department of Computer Science, University of Oxford
2CaMLSys, Department of Computer Science and Technology, University of Cambridge
"WhiRL, Department of Computer Science, University of Oxford
Samsung Al Center, Cambridge
ABSTRACT
Pruning neural networks at initialization would enable us to find sparse models that retain the accuracy of the original network while consuming fewer compu- tational resources for training and inference. However, current methods are in- sufficient enable this optimization and lead to a large degradation in model performance. In this paper, we identify a fundamental limitation in the formula- tion of current methods, namely that their saliency criteria look at a single step at the start of training without taking into account the trainability of the network. While pruning iteratively and gradually has been shown to improve pruning per- formance, explicit consideration of the training stage that will immediately follow pruning has so far been absent from the computation of the saliency criterion. To overcome the short-sightedness of existing methods, propose Prospect Pruning (ProsPr), which uses meta- traddggtssents through the first few steps of optimization to determine which weights to prune. ProsPr combines estimate of the higher- order effects of pruning on the loss and the optimization trajectory to identify the trainable sub-network. Our method achieves state-of-the-art pruning performance on a variety of vision classification tasks, with less data and in a single shot com- pared existing gruning-at-iniing itiiiiiiaiiitoooonn methods. Our code is available online at
com/ mi ad/prospr
INTRODUCTION
Pruning at snnalliiaaion-hhere remove weights from model before training begins is 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 (019)) provide empirical evidence for the existence of sparse sub-networks that 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 trained 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 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.
compute-inennsive
A recent line 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. 202I); Tanaka et al. (2020). This class of methods uses gradients to assess
"Corresponding author. Contact mi add. izade Basehhcs... ac.uckuk

---

Published as a conference paper at ICLR 2022
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 different from classical pruning-at-conwen at-convergence methods, but the gradients for well-trained model are small; therefore these methods must inspect higher-order metrics such 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/deploymenn 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.) [202]). In this paper, we identify 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 al. 2021). Although this ap- proach has been shown to he effective, it expensive and cumbersome in practice and ultimately is an indirect approximation the trainability criteria we are looking to incorporate into our objective.
In this paper, we propose Prospect Pruning (ProsPr), a new pruning-at an- 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-g gradients by backpropagating through the first few model updates in order to estimate the effect the initial pruning parameters have on the loss after 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 key limitation in prior saliency criteria for pruning neural networks-namee that they do not explicitly incorporate trainabiliy--afer-rrummi
into their criteria.
We propose a new pruning-at-init method, ProsPr, that uses meta- 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 pruning- at-init methods. 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.
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 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: 8L ss VoL 50 + 500 H 80, where H V2LC is the Hessian matrix. When the network has converged, the first-order term in the expansion 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 V 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 -aass 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 to O. More specifically, SNIP computes the
2

---

Published as a conference paper at ICLR 2022
saliency scores according to:
1g (w, D)1
8j
(I)
k=1 9 (w, D)
with
acic O w, D)
93 (w, D)
Dcg
where m is the number of weights in the network, с E {0, 1 the pruning mask (initialised to above), D is the training dataset, 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 (c, = 1), and all other weights are pruned (cg = 0), before the network is trained.
Our method, to be introduced in Section B] also relies computing the saliency scores for each mask element, but uses a more sophisticated loss function to incorporate the notion of trainability.
OUR METHOD: PROSPR
In this section we introduce our method, Prospect Pruning (ProsPr). We note that for the problem of pruning 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. 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 itself differentiable. Once the pruning mask is com- puted, we rewind the weights back to their values at initialization and train the pruned network. The pradientoff-gradient called mela-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 to pick pruned model, A. which is more trainable than an alternative pruned model B. That means we want the loss L(iA, y) to be lower than L(DB, y) at convergence (for 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 line 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 Winis we apply weight mask the initial weights,
Wo = 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.
3

---

Published as a conference paper at ICLR 2022
Backpropagate through entire computation graph with respect
WM-1
V 2i
War .D))
C Wo
Do)
W M-11
PM-1)
DM
WIN
Perform Mf weight updates with batches of data D;
WAN
Figure I: Visualization of computing saliency scores in method, ProsPr. By backpropagating through several gradient steps we capture higher-order information about the objective that we care about in practice, i.e., saliency of parameters during training and not just initialization.
We then sample M+1 batches of data D; ~ Diman (i E (O....,,My. MY; M 2 1) for the pruning step, and perform M weight updates
W1 Wo - a C(wo.
(4)
WM WM- aV WW WW
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,
Ve C(w Dm)
(6)
Using the chain rule, we write out the form of the meta-gradient beginning from the last step:
V.L(wM, D) = Vw C(war,
repeating for each step until we reach the zero" th step whose gradient is trivial,
Vw L(war, D) war .wM) ... (Wwo W (V)(V.
(8)
Wwn, L(was,
WM) ... (Vw, w1 (V)((V(CO))) Win)))
(9)
M
V C(was, D)
II (V WV..., Wm) Lm=1
Winit-
(10)
In practice, compute the meta-gradien meta-gradients by relying automatic differentiation software such as PyTorch (Paszke et al., (201)). However, care must taken to ensure that weights at each step kept in memory that the entire computation graph, including gradients, is visible to the automatic differentiation software. The saliency scores are now given by
19s (w, D)1
8, -k=1 19n(w, (w, D)I*
(11)
with
(w, D) aL(wM. D)
(12)
Oc,
where Wa function of c. Equation 12 stands contrast to SNIP, where the saliency is computed using the loss at Wini rather than W.r- The saliency scores are then used to prune the initial weights Winiti the ones with the highest saliency scores are retained (c, = 1), and all other weights are pruned (cj 0). Finally, the network is trained with the pruned weights Wiit-
Algorithm summarises the proposed method, ProsPr.
We formalise the weight updates using vanilla SGD here; in practice these may be different when using approaches such as momentum or BatchNorm (loffe & Szegedy 2015). Since our implementation relies on automatic differentiation in Py PyTochh (Paszke et al.] 2019) we can use any type of update, long it is differentiable the initial mask c.
4

---

Published as a conference paper at ICLR 2022
Algorithm ProsPr Pseudo-Code
1: Inputs: training dataset Dimin, number of initial training steps M, number of main training steps N (M & N), learning rate
2: Initialise: network weights Wisit
3:
P Initialise mask with ones Apply mask to initial weights
Wo Cinit 3 Winit
for O,.... .M- 1 do
6: DE ~ Ditrin
D- Sample batch of data
7: Wi+1
V C(w,, Da)
Update network weights
8:
end for
9: 93 (w D) = (L(WM, D),
D. Compute meta- graddet
Ia, 1.1
10: 8j Σr [W]
> Compute saliency scores
11: Determine the k-th largest element s, 8k-
1, cy 2
12: Gpnune
Set pruning mask
0, otherwise
13: We Cprie C Wint
b Apply mask to initial weights Winit
14: for L, N do
D Train pruned model
Wi+1 Wi V [ww,
16: end
3.2 FIRST-ORDER APPROXIMATION
Taking the meta-gradient through many model updates (Fquation 6 can be memory intensive: in the forward pass, all gradients of the individual update steps need to be retained in memory to then able to backpropagate all the way the initial mask. However, we only need to perform perform the pruning step CPU
few steps at the beginning of training which usually has access to more memory compared to GPU. We apply this approach in our own experiments, with overheads of around 30 seconds being observed for the pruning step.
in practice
Alternativel,, when the number of training steps needs to be large we can use the following first- order approximation. Using Equation (TO), the meta vreta-gradientiin is:
VeL(WM. DM) Vw L(WM, Dar)
(V Wm)
Winit:
(13)
Lm=1
writing wm in terms of Wm-1 following SGD,
M
V Da))
V Wm-1 (Wm-1 aV Wm-1 L( Dm))
Winit
mm=1
(14)
carrying through the partial derivative,
Vw C(w.. Dar)
II I- C..oaW,,,WWWm-1 Dm)
Winits
(15)
Lm=1
and finally dropping small terms for sufficiently small learning rates,
M I
Vwze L(WM, Dar)
Winits
(16)
Lm=1
Vwn C(WM. Da))
Winit-
(17)
In the second-to-lase step, drop the higher-order terms, which gives us first-order approximation of the meta-gradien
use steps for experiments on CIFAR-10. 10, CIFAR- 100 and TinyImageNet datasets
Note that this approximation also works for optimisers other than vanilla SGD (e.g., Adam, Adamw, Ad-
ahound), except that the which is dropped (r.h.s. of Equation Equation [5
looks slightly different.

---

Published as a conference paper at ICLR 2022
ResNet-20 (CIFAR-10)
VGG-16 (CIFAR-10)
94
90
Test Accurscy
Test Accuracy
93
85
92
80-
75
90
20.0
48.8
67.2
79.0
86.0
91.4
94.5
96.5
97.7
20.0
48.8
67.2
79.0
86.0
91.4
94.5
96.55
97.7
Sparsity (9))
Sparsity (%)
ResNet-
Prune as Initislization
Prune at Convergence
Test Test Accurscy
50-
ProsPr
Weight Rewinding
Magnitude
Maguitude
SNIP
45
GraSP
SyuFlow
40
Random
Unpruned Baseline
20.0
48.8
07.2
79.0
86.6
91.4
91.5
96.5
07.7
Sparsity (%)
Figure 2: Accuracy of ProsPr against other prune-at-init and prune-after-convergenee methods as
benchmarked by Frankle et al. (202I). The shaded areas denote the standard deviation of the runs.
With this approximation, we only need to save the initial weight vector Winit in memory and multiply it with the final gradient. This approximation can be crude when the Laplacian terms are large, but with sufficiently small learning rate becomes precise. The approximation allows us to take many more intermediate gradient-steps which can be beneficial for performance when the training dataset has many classes, as we will see in Section
EXPERIMENTS
We empirically evaluate the performance of our method, ProsPr, compared to various vision classi- fication baselines across different architectures and datasets. In supplementary sections show effectiveness of our method on image segmentation tasks (Appendix DD and when using self- supervised initialization (Appendix E We provide details of our hyper-parameters, experiment setup, and implementation details in Appendix A
4.1 RESULTS ON CIFAR AND TINY-IMAGENET
In recent work, Frankle et al. (202I) extensively study and evaluate different pruning- i--iiiiaaiizaiinnn methods under various effects such as weight re-initialization, weight shuffling, and score inversion. They report the best achievable results by these methods and highlight the gap between their perfor- mance and two pruning-at aataaaat-coceennnn methods, weight rewinding and magnitude pruning (Renda ct al. 2020 Frankle et al., 2020).
In Figure we evaluate ProsPr on this benchmark using ResNet-20 and VGG-16 on CIFAR-10, and R on Imag-em can be seen that ProsPr reduces the performance gap, espe- cially higher sparsity levels, and in some cases exceeds the accuracy of pruning-after after-conveeeence methods. Full results are also summarised in Appendix B
This is remarkable achievement: ProsPr is the first work to close the gap to methods that prune after training. Previous works that prune at the start have not been able to outperform methods that prune after training any settings. including smaller datasets such CIFAR-10 also important to note that other baselines that have comparable accuracies are all iterative methods. ProsPr is the only method that can do this in single shot after using only 3 steps batch-sizes of 512
ImageNe. It is

---

Published as a conference paper at ICLR 2022
Table 1: Test accuracies of VGG- 19 and ResNet-50 on ImageNet. First-Order ProsPR exceeds the results reported by de Jorge et al. (2021) in all configurations but one, where GraSP works best.
VGG-19
ResNet-50
Sparsity
90%
95%
90%
95%
Accuracy
Top-1
Top-5
Top-1
Top-5
Top-1
Top-5
Top-1
Top-5
Unpruned Baseline
73.1
91.3
75.6
92.8
ProsPr (ours)
70.75
89.9
66.1
87.2
66.86
87.88
59.62
82.82
FORCE
70.2
89.5
65.8
86.8
64.9
86.5
59.0
82.3
Iter-SNIP
69.8
89.5
65.9
86.9
63.7
85.5
54.7
78.9
GRASP-MB
69.5
89.2
67.6
87.8
65,4
86.7
46.2
66.0
SNIP-MB
68.5
88.8
63.8
86.0
61.5
83.9
44.3
69.6
Random
64.2
86.0
56.6
81.0
64.6
86.0
57.2
80.8
in the inner-loop before computing the metu-gradients. In total, we only use 4 batches data. We also do not do any averaging of scores by repeating the method multiple times.
The performance in these small datasets comes from the fact that ProsPr computes higher-order gradients. While there other iterative methods that work without any data, their effect is mostly a more graceful degradation at extreme pruning ratios as opposed to best accuracy at more practical sparsity levels. One example is SynFlow which is similar to FORCE but uses an all-one input tensor instead of samples from the training set (Tanaka et al. 2020
4.2 RESULTS ON IMAGENET DATASET
To evaluate the performance of ProsPr more difficult tasks we run experiments on the larger Im- ageNet dataset. Extending gradient-based pruning methods this dataset poses several challenges.
Number of classes In aynaptic-saliena methods, the mini batches must have enough examples from all classes in the dataset. |Wang al.) (2020) recommend using glas-balaacced mini- mini-batches sized ten times the number of classes. In datasets with few classes this is not issue and even single batch includes multiple examples per class. This one reason why methods like SNIP work with a single batch, and why kept the number of steps in ProsPr's inner loop fixed to only 3. ImageNet however has 1,000 classes, and using a single or a handful small batches is inadequate. Previous methods such FORCE, GraSP, SynFlow avoid this problem by repeating the algorithm with new data batches and averaging the saliency scores. In ProsPr we instead increase the number of updates before computing the meta-gradients, ensuring they flow through enough data. Computing meta-gradients through many steps however poses new challenges.
gradient stability issues when computing gradients deep loops. Gradient degradation problems, i.c., vanishing and exploding gradients, have also been observed in other fields that use meta-gradients such as Meta-Learning. Many solutions have been proposed to stabilize gradients when the length of loop increases beyond 4 or 5 steps, although this remains an open area of research (Antoniou et al., 2019).
Gradient degradation We start
Computation Complexity For ImageNet we must make the inner loop hundreds of steps deep achieve balanced data representation. addition to stability issues, backpropagating through hundreds of steps is very compute intensive.
Therefore for our experiments on ImageNet we use the first-order approximation of ProsPr (Sec 3.2 We evaluate ProsPr using ResNet-50 and VGG-19 architectures and compare against Nate-of-the-an methods FORCE and INeP introduced by de Jorge et al. [022]) We include multi-batch ver- sions of SNIP and GraSP (SNIP-MB and GraSP-MB) to provide fair comparison to iterative meth- ods, which partially prune several times during training, in terms of the amount data presented to the method. We 1024 steps with batch size of 256 (i.e. 262,144 samples) for ResNet-50. For VGG-19, a much larger model, and which requires more GPU memory we do 256 steps with batch size of 128. This is still far fewer samples than other methods. Force, for example, gradually prunes

---

Published as a conference paper at ICLR 2022
in 60 steps, where each step involves computing and averaging scores over 40 batches of size 256, ie. performing backpropagation 2400 times and showing 614,400 samples to the algorithm.
Table 1 shows results compared to the baselines reported by [de Jorge et al. (2021). First-order ProsPr exceeds previous results in all configurations except one, where it is outperformed by GraSP. Note the surprisingly good performance of random pruning of ResNets, which was also observed by de Jorge et al. (2021). This could be explained by the fact that VGG-19 is a much larger architecture with 143.6 million parameters, compared to 15.5 million in ResNet-50s.. More specifically the final three dense layers of VGG-19 constitute 86% of its total prunable parameters. The convolution layers of VGG constitute only 14% of the prunable weights. Pruning methods are therefore able to keep more of the convolution weights and instead prune extensively from the over- parametrized dense layers. ResNet architectures on the hand have a single dense classifier the end.
4.3 STRUCTURED PRUNING
We also evaluate ProsPr in the structured pruning setup where instead of pruning individual weights, entire channels (or columns of linear layers) removed. This is more restricted setup, however it offers memory savings and reduces the computational cost of training and inference.
Adopting ProsPr for structured pruning is as simple as changing the shape of the pruning mask in to have one entry per channel (or column of the weight matrix). We evaluate our method against 3SP, method that extends SNIP structured pruning (van Amersfoort et al.) 2020). Our results summarized in 2 which show accuracy improvements in all scenarios. In Appendix C we also evaluate wall-clock improvements training time as a result of structured pruning at initialization.
Test accuracies for structured pruning using VGG-19 on CIFAR-10 and CIFAR-100. ProsPr achieves better accuracy in all configurations.
Table 2:
Sparsity
Method
CIFAR- 10 Ace (%)
CIFAR- 100 Acc (%)
Unprumed Baseline
93.6
72.5
80%
ProsPr (ours)
93.61 + 0.01
72,29 + 0.11
3SP
93.4 1 0.03
69.9 1 0.14
3SP reinit
93.4 0.04
70.3 1 0.16
3SP rescale
93.3 0.03
70.5 0.13
Random
92.0 0.08
67.5 1 0.16
90%
ProsPr (ours)
93.64 0.24
71.12 + 0.26
3SP
93.1 = 0.04
68.3 + 0.12
3SP reinit
93.0 = 0.02
69.0 + 0.08
3SP rescale
93.0 0.06
69.2 : 0.11
Random
90.4 + 0.12
63.8 0.13
95%
ProsPr (ours)
93.32 + 0.15
68.03 + 0.38
3SP
92.5 0.12
63.2 0.52
3SP reinit
92.6 0.09
64.2 0.35
3SP rescale
92.5 0.06
63.5 0.63
Random
89.0 1 0.15
60.1 1 0.29
4.4 NUMBER OF META STEPS
Finally, we evaluate ProsPr when using a varying number of meta steps, which gives insight into whether using meta-gradients is beneficial. We repeated experiments from Section 4.3 but this time we vary the depth of training steps between 0 and 3. The results in Table 3 show that the final accuracy consistently increases of meta- -reaa-gaddsnns. We used the same data batch in all M training steps to isolate the effect of M, while in other experiments
increase the depth of the training, showing the effectiveness
new batch in every step.
In theory increasing the number of training steps should always help and match the ultimate objective (estimating the loss after many epochs of training) in the limit. However, in practice increasing the number steps beyond poses a lot of gradient stability issues (and is computationally expensive). These issues have been also identified in the meta-learning literature (Antoniou al. 2019
8

---

Published as a conference paper at ICLR 2022
Table 3: Evaluating the effect of meta steps (M) on structured pruning performance of VGG-19
(a) 90% Sparsity
(b) 95% Sparsity
M
CIFAR-I0 Acc
CIFAR-I00 Acc
M
CIFAR-10 Acc
CIFAR-100 Acc
0
93.1% + 0.04
68.3% + 0.12
0
92.5% + 0.12
63.20% + 0.52
93.3%
0.12
68.8% 0.18
1
92.9% 0.31
64.87% + 0.35
93.5% + 0.21
69.8% + 0.20
2
93.25% 0.24
67.12% = 0.42
93.6% + 0.19
71.0% + 0.21
3
93.29% 0.29
67.98% + 0.31
RELATED WORK
Pruning at initialization Several works extend the approach proposed by Lee et al. (0188) de Jorge et al. ) evaluate SNIP objective in a loop in which pruned parameters still receive gradients and therefore have chance get un-pruned. The gradual pruning helps avoid the layer- collapse issue, and their method, known as FORCE, achieves better performance at extreme sparsity levels. Tanaka al. (2020) provide theoretical justification for why iteratively pruning helps with the layer-collapse issue and propose data-free version of the method where an all-one input tensor used instead of real training data. Wang al. (2020) propose an alternative criterion to minimizing changes in the loss and instead argue for preserving the gradient flow. Their method, GraSP, keeps weights that contribute most to the norm of the gradients. yan Amersfoort et al. (2020) extends SNIP and GraSP to structured pruning make training and inference faster. They further augment the scores by their compute cost to push the pruning decision towards more FLOPS reduction.
Gradual pruning As discussed in Section [I in existing methods the training step has been ab- sent from the saliency computation step. As a workaround, many methods make their approaches waining-aware by applying pruning gradually and interleaving it with training: Zhu & Gupta (2018) proposed an exponential schedule for effectiveness in broader range of tasks. Frankle & Carbin (2019) show that weight rewinding
feiiin-duuiin-triinnng and Gale et al. (2019) showed its
achieves better results when done in multiple prune-retrain steps. Lym et al. (2019) continuously apply structured pruning via group-lasso regularization while the same time increasing batch
sizes. You et al. (2020) find pruned architectures after few epochs of training-and and-pruninn and monitoring distance metric.
Backpropagation through gradients, and its first-order approximation, is also used model-agnostic meta-learning literature (Finn et al. 2017: Zintgraf et al., 2019) where the objective is to find model that can be adapted to new data in few training steps. Similar to our setup, the meta-loss captures the trainability
Meta-Gradients
a model, but additionally, the meta-gradients are used to update the network's 's weights in second loop. In self-supervised learning setting, Xiioo et al. [222T) use meta-gradients to explicitly optimize a learn-to-ge learn-t--generalize regularization term in is also used regularize loss with a penalty on the gradients, for instance, to enforce Lipschitz continuity on the network (Gulrajani
nested meta-learning loops. Computing ofddients-of-rradies
al. 2017 control different norms of the gradients (Alizadeh al. 2020).
6 DISCUSSION
Although pruning at initialization has the potential to greatly reduce the cost of training neural networks, existing methods have not lived up to their promise. We argue that this is, part, because they do not account for the fact that the pruned network is going to be trained after it is pruned. We take this into account, using a saliency score that captures the effect of a pruning mask on the training procedure. As a result, our method is competitive not just with methods that prune before training, but also with methods that prune iteratively during training and those that prune after training. In principle, compressing neural networks at initialization has the potential to reduce energy and environmental costs of machine learning. Beyond our context, taking into account that methods which prune-at-converge could have further implications for these pruning methods as well (Molchanov et al. 2016 Wang et al. 2019
generally have be fine-tuned, it is possible that work

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
Milad Alizadeh is grateful for funding by the EPSRC (grant references EP/R512333/I) and Arm (via NPIF 2017 studentship). Shyam Tailor is supported by EPSRC grants EP/M50659X/1 and EP/S001530/1 (the MOA project) and the European Research Council via the REDIAL project (Grant Agreement ID: 805194). Luisa Zintgraf is supported by the 2017 Microsoft Research PhD Scholarship Program, and the 2020 Microsoft Research EMEA PhD Award. Joost van Amersfoort is grateful for funding by the EPSRC (grant reference FP/N50971 1/1) and Google-DeepMind. Se- bastian Farquhar is supported by the EPSRC via the Centre for Doctoral Training in Cybersecurity at the University of Oxford well Christ Church, University of Oxford.
REFERENCES
Milad Alizadeh, Arash Behboodi, Mart van Baalen, Christos Louizos, Tijmen Blankevoort, and Max Welling, Gradient & regularization for quantization robustness. In International Conference on Learning Representations, 2020.
Antreas Antoniou, Harrison Edwards, and Amos Storkey. How to train your maml, 2019.
Pau de Jorge, Amartya Sanyal, Harkirat Behl, Philip Torr, Grégory Rogez, and Puneet K. Dokania. Progressive skeletonization: Trimming more fat from a network at initialization. In International Conference Learning Representations, 2021.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp- 1126-1135. PMLR, 06-11 Aug 2017.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference Learning Representations, 2019.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis, 2020.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Pruning neural net- works initialization: Why are missing the mark? In International Conference on Learning Representations, 2021.
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint 19t.::990220997444 2019.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukeuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent- new approach to self- superviee learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 21271-21221 Curran Associates, Inc., 2020.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training wasserstein gans, 2017.
Babak Hassibi and David G Stork. Second order derivatives for network pruning: Optimal brain surgeon. In Advances in neural information processing systems, pp. 164-171, 1993.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 448 456, Lille, France, 07-09 Jul 2015. PMLR.
Tachee Jeong, Manasa Bollavaram, Elliott Delaye, and Ashish Sirasao. Neural network pruning for biomedical image segmentation. In Medical Imaging 2021: Image-Guided Procedures, Robotic Interventions, and Modeling, volume 11598, pp. 115981J. International Society for Optics and Photonics, 2021.
10

---

Published as a conference paper at ICLR 2022
Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in neural information processing systems, pp. 598-605, 1990.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based connection sensitivity. arXiv preprint arXiv: 1af::8880.0000 180.022400, 2018.
Sangkug Lym, Esha Choukse, Siavash Zangeneh, Wei Wen, Sujay Sanghavi, and Mattan Erez. Prunetrain. Proceedings of the International Conference for High Performance Computing, Net- working. Storage and Analysis, Nov 2019. doi: 10. 00.1145/99550003356155.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. In International Conference Learning Repre- sentations, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Alex Renda, Jonathan Frankle, and Michael Carbin. Comparing rewinding and fine-tuning in neural network pruning, 2020.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234 241. Springer, 2015.
Hidenori Tanaka, Daniel Kunin, Daniel I. Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. In H. Larochelle, M. Ranzato, R. Had- sell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, vol- ume 33, pp. 6377-6389. Curran Associates, Inc., 2020.
Joost van Amersfoort, Milad Alizadeh, Sehastian Farquhar, Nicholas Lane, and Yarin Gal. Single shot structured pruning before training, 2020.
Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang. Eigendamage: Structured pruning in the kronecker-factored eigenbasis. In International Conference on Machine Learning, pp- 655666577, 2019.
Chaoqi Wang. Guodong Zhang, and Roger Grosse. Picking winning tickets before training by
preserving gradient flow. In International Conference on Learning Representations, 2020.
Taihong Xiao, Xin-Yu Zhang, Haolin Jia, Ming-Ming Cheng, and Ming-Hsuan Yang. Semi- supervised learning with meta-gradient. In Arindam Banerjee and Kenji Fukumizu (eds.), Pro- ceedings of The 24th International Conference on Artificial Intelligence and Statisties, volume 130 of Proceedings of Machine Learning Research, pp. 73-81. PMLR, 13-15 Apr 2021.
Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Yingyan Lin, Zhangyang Wang, and Richard G. Baraniuk. Drawing early-bird tickets: Toward more efficient training of deep networks. In International Conference on Learning Representations, 2020.
Michael H. Zhu and Suyog Gupta. To prune, or not prune: Exploring the efficacy of pruning for model compression, 2018.
Luisa Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast con- text adaptation via meta-learning, In ICML 2019: Proceedings of the Thirty- Sixth International Conference on Machine Learning, June 2019.
11

---

Published as a conference paper at ICLR 2022
A EXPERIMENTAL SETUP
A.1 ARCHITECTURE DETAILS
We use standard VGG and ResNet models provided by torchvision throughout this work where possible. The ResNet-20 model, which is not commonly evaluated, was implemented match the version used by Frankle et al. (2021) so that we could compare using the benchmark supplied by this paper.
For smaller datasets, is common to patch models defined for ImageNet. Specifically, for ResNets, replace the first convolution with one x 3 filter size, and stride the first max-pooling layer replaced with an identity operation. For VGG, we follow the convention used by works such FORCE (de Jorge al. 2021). We do not change any convolutional layers, but change the classifier to use a single global average pooling layer, followed by a single fully-connected layer.
A.2 TRAINING DETAILS
For CIFAR-I0, CIFAR-I00 and TinylmageNet we perform 3 meta-steps to calculate our saliency criteria. We train the resulting models for 200 epochs, with initial learning rate 0.1; we divide the learning rate by 10 at epochs 100 and 150. Weight decay was set to 5 x 10 - Batch size for CIFAR- 10, CIFAR-106 100, and TinylmageNet was 256. For CIFAR-10 10 and CIFAR- 100 we augment training data by applying random cropping (32 32, padding 4), and horizontal flipping. For TinylmageNet we use the same procedure, with random cropping parameters set to x 64, padding 4.
For ImageNet we train models for 100 epochs, with an initial learning rate of 0.1; we divide the learning rate by 10 at epochs 30, 60 and 90. Weight decay was set to 1 x 10-4, Batch size was 256. We use the first order approximation to do pruning, and use 1024 steps for ResNet-50. For VGG-19 we use 2048 steps, but with batch size set to 128 (due to memory limitations, as our implementation only utilized single GPU for meta-training)) We apply random resizing, then crop the image to 224 x 224, with horizontal flipping.
A.3 IMPLEMENTATIONS
In addition to our code, the reader may find it useful to reference the following repos from related work. Our experiments were performed using code derived from these implementations:
htte:///gtttubb.cc//aeeeerrrrrr
bttees//gghhhbbbcco////lleeessss
httpes//ggteubb..oom//1a/a/////
B NUMBERS FROM FIGURE 2
Table 4: Numerical results for ResNet-20 on CIFAR-10
200
115
14
A8 2OL
AL aM
B113:
85-04
00
251402
71100
101-0
001403
Srea nowe w S01-E1 See-En acTenS STeen most mced Sc-011 I4-42 30--55
Table 5: Numerical results for VGG-16 on CIFAR-10
K55
200
052
011.11
63-41
01-31
-1
61.11
01101
01
03--11 15
02.02
11
62-33
02-2
83 10
82131
917481
804 meet mavet Weeet aren 61431 Nawel meeer
05401
01e33
12

---

Published as a conference paper at ICLR 2022
Table 6: Numerical results for ResNet-18 on TinyImageNet
286
01
714
704
63
44.
114
901
113
100
30 51
11.
111
41181
03:41
2010
11
05:31
04111
02131
01:11
015
Trads
5:8444
C WALL CLOCK TIME FOR STRUCTURE PRUNING AT INITIALIZATION
When pruning is done at convergence, the benefits of having a compressed model (in terms of memory saving and speed-up) can only be utilized inference/deploymmet time. However, with mruning-at-nntiialzzan
these benefits can be reaped during training as well. This is especially true in the case of structured pruning, where pruning results in weight and convolutional kernels with smaller dimensions (as opposed to unstructured pruning, where we end up with sparse weights with the original dimensions). This means that in addition memory savings, training take fewer operations which speeds up training. To evaluate the benefits of training at initialization in terms of speed improvements we measured the wall-time training time on an NVIDIA RTX 2080 Ti GPU for the architectures used in Section 4.3 (an additionally on ImageNet dataset). The results in Table show that structured pruning with ProsPr can significantly reduce the overall training time.
Table 7: Wall-time Training Speed-ups by Structured Pruning at Initialization
Dataset
Epochs
Model
Unpruned
80% pruned
90% pruned
95% pruned
CIFAR-100
200
ResNet-18
83.9 mins
60.76 mins
54.8 mins
46.6 mins
CIFAR-10 100
200
VGG-19
50.3 mins
45.8 mins
39.93 mins
38.8 mins
Tiny-Imagen
200
ResNet-18
9.79 hours
8 hours
5.4 hours
4.92 hours
Tine-ImageNet
200
VGG-19
5.75 hours
4.0 hours
3.38 hours
2.7 hours
ImageNet
90
R8sNet-18
73.7 hours
72.15 hours
65.9 hours
64.6 hours
D RESULTS ON SEGMENTATION TASK
An interesting, albeit less common, application for pruning models is within the context of segmen- tation. In recent paper Jeong al. (20T)) train and prune the U-Net (Ronneberger al. 2015) architecture on two image datasets from the Cell Tracking Challenge (PhC-C2DH-U373 and DIC- C2DH-HeLa). They use the classic multi-stey approach of gradually applying magnitude-pruningg interleaved with fine-tuning stages. To evaluate the flexibility of our method we used meta-gradients the beginning of training (on a randomly initialized U-Net), prune in single shot, and train the network once for the same number of epochs (50). We kept the training set-up the same the base- line by Jeong et al. (202T) (i.e., resizing images and segmentation maps to (256,256), setting aside 30% of training data for validation) and similarly aim to find the highest prune ratio that does not result in IOU degradation. We report intersection-oven ceereectioo-ove--uiioon (IOU) metric for the two datasets in
Tables 8 and g
Table 8: Mean-IOU on U373 validation
Table 9: Mean-IOU on Hel.a validation
Method
Prune Ratio
Mean IOU
Method
Prune Ratio
Mean IOU
Unpruned 95%
0.9371
Unpruned Jeong et al.
0.7514
Jeong et al.
0.9368
81.8% 90%
0.741
ProsPr
97%
0.9369
ProsPr
0.7491
These results show that our method works as well (or better) compared to this compute- expentiv baseline, in the sense that we can prune more parameters while keeping the IOU score the same.
13

---

Published as a conference paper at ICLR 2022
E SELF-SUPERVISED INITIALIZATION
To evaluate the robustness and consistency of our method against non-random initialization we ran experiments using BYOL to learn representations from unlabeled samples (Grill al., 2020). We used ResNet-18 as backbone and trained for 1000 epochs with an embedding size of 64, Unlike the vanilla ResNet-18 architecture used Section 4.3 we used the commonly-usee modified version of ResNet- 18 for smaller inputs (removing the first pooling layer and modifying the first convolutional layer to have kernel kernel size of 3. stride of 1, and padding size of 1). We then used this trained R as the initialization for our meta-gradient pruning method. After the pruning step, all layers were trained as before until convergence. All training sapee-parameters were kept as before. The results (final test accuracies for 95% pruning) are summarized in Table I0
Table 10: Comparing test accuracies (%) of ProsPr randomly initialized ResNet-18 18 and initial- ization from self-supervised learning (B (BY)
Dataset
Random Init
BYOL init
CIFAR-10
93.6 73.2
93.62 74.02
CIFAR-100
These results show the robustness of method for this particular self-supervised initialization. Starting from a learned representation can be challenging because these representations are much closer to weight values convergence, and therefore the magnitude of their gradients is significantly smaller than randomly initialized weights. However, this is less of a problem for meta meta-gradienn geeta-grdeeets as their magnitude is still significant due pack-propagation through training steps. This can be seen in Figure 3 which shows the L.2 norm of gradients of each layer of B BYOL-initialized ResNet- 18 for meta- mnta-gradients compared to normal gradients. It can be seen that meta meta-gradieets provide stronger signal compared to normal gradients.
10a
1. grodear
SNIP
PeaPr BYOL ini
10-1
10
10-3
17 19 20
Lanss
Figure 3: Comparing L2-norm of gradients and meta- ceta-gradieets in BBOOL ResNet-18
14


Reviews:
[{"review_id": "wwQYE7TDGMk", "reviewer": "Reviewer_9cLN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Strengths]\n- The proposed ProsPr is a simple and effective pruning method that prunes weights at initialization. The authors propose to use meta-gradients to compute saliency scores when pruning weights.\n\n- The higher-order temrs in meta-gradients can be further dropped (Equ. 16) such that saving the initial weights $w_{init}$ is enough when computing meta-gradients.\n\n- Experimental results show that ProsPr achieves state-of-the-art pruning performance.\n\n\n[Weaknesses]\n- Writing needs improvement. E.g., \"Many methods attempts...\", \"...the original, unpruned, model\" and \"Previous works that prune at the start have training have not been...\".\n\n- Although experimental results show that using estimation over several steps of gradient descent improves pruning performance, the connection between meta-gradients and motivation (i.e., the trainability of weights) is not strong and convincing enough.\n", "summary_of_the_paper": "This work focuses on weight pruning at initialization. In this paper, the authors point out an important problem that the pruned subnetwork at initialization is going to be trained and previous prune-at-init methods ignore this fact. As a result, these prune-at-init methods ignore the trainability of weights. This paper proposes to use meta-gradients through the first few steps of optimization to determine which weights to prune. Experimental results show that ProsPr (this paper) achieves state-of-the-art pruning performance.", "summary_of_the_review": "It is a good technical paper, but it requires better writing. Although the proposed ProsPr, which is implemented by meta-gradients, is effective, I think the novelty is limited.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "w8R_5b7dy0O", "reviewer": "Reviewer_fSju", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n+ This paper is well-written.\n\n+ the proposed method is easy to implement.\n\n+ The experimental results support the claim that the proposed method achieves higher accuracy compared to existing pruning-at-init methods.\n\nWeaknesses:\n\n1. According to the proposed first-order approximation, it seems that the proposed method is limited to apply with the optimization method SGD. There is a variety of models that are trained with other optimization methods, like Adam, Adamw, and Adabound. It would be helpful to provide a discussion on the direction of how to adapt the proposed method to work with generic (or specific) optimization methods.\n\n2. The first-order approximation (Eq. (17)) may not be well studied. The following questions are not clear: Is it sensitive to how $W_{init}$ is initialized (e.g., xavier_uniform, kaiming_normal, etc)? Is it robust to the sequence data? How does the number $M$ of initial training steps affect the accuracy and sparsity?\n\n3. An important step to the first-order approximation is to drop the higher-order terms of Eq. (15), i.e., $\\prod_{m=1}^{M} I - \\alpha \\nabla_{w_{m−1}}^{2} \\mathcal{L} (w_{m−1}; \\mathcal{D})$. A follow-up question is, what is the error bound (w.r.t. M) between the approximated meta-gradient and the ground-truth meta-gradient? For example, assume that the loss function meets the second-order necessary condition $\\nabla_{w_{m−1}}^{2} \\mathcal{L} (w_{m−1}; \\mathcal{D}) \\succeq 0$, as $M$ gets large enough, the operation $\\prod_{m=1}^{M} I - \\alpha \\nabla_{w_{m−1}}^{2} \\mathcal{L} (w_{m−1}; \\mathcal{D})$ could lead to vanishing and exploding gradients. However, the approximated gradient $\\nabla_{w_{M}} \\mathcal{L}( w_{m−1}; \\mathcal{D}) w_{init}$ seems to be still robust. How to interpret the discrepancy?\n\n4. Figure 2 shows that pruning at convergence achieves a better trade-off between accuracy and sparsity than pruning at initialization. So it is not clear that the advantage(s) of the techniques of pruning at initialization over the ones of pruning at initialization. I looked into the introduction and the related work of pruning at initialization. I didn't find the points related to this. It would be better to explicitly discuss the comparison to make the motivation clearer. If I missed something, please point it out.\n\n5. The algorithm presumes that $w_{init}$ is randomly initialized. It would be interesting to verify if the proposed method behaves consistently with $w_{init}$ that is learned from unlabeled samples, which is a common practice in self-supervised learning. \n\n6. The related work of meta-gradients misses a related work, that is, meta-gradient in semi-supervised learning [r1].\n\n7. This is not a critical comment, but I'd like to bring it to the discussion. Regarding the novelty, the proposed method relies on the saliency score (Eq. 1 and Eq. 11) defined in SNIP. This may make this work arguably look incremental. The authors present a discussion \"our method, to be introduced in Section 3, also relies on computing the saliency scores for each element in the mask but uses a more sophisticated loss function to incorporate the notion of trainability into the objective.\" Still, I think the work is a bit weak in terms of the novelty of the methodology.\n\nReferences:\n\n[r1] Xiao, Taihong, Xin-Yu Zhang, Haolin Jia, Ming-Ming Cheng, and Ming-Hsuan Yang. \"Semi-Supervised Learning with Meta-Gradient.\" In International Conference on Artificial Intelligence and Statistics, pp. 73-81. PMLR, 2021.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper is well-written and provides detailed derivations. As a result, the proposed method has the potential to apply to a wide range of real-world applications. On the other hand, I think the proposed method may not be well studied and the experiment can be improved by adding the results of using $w_{init}$ learned from unlabeled data. Therefore, my initial recommendation is \"marginally below the acceptance threshold\". I will go over the review comments by other reviewers and the responses by the authors and adjust my recommendation accordingly.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Tp35d64khp6", "reviewer": "Reviewer_D88v", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1) Since the current methods are insufficient to enable this optimization and lead to a large degradation in model performance, the authors proposed a new method to identify a fundamental limitation, namely that their saliency criteria look at a single step at the start of training without considering the trainability of the network. The paper tackles an important problem. \n2) The paper is well organized and easy to follow. \n3) It would be better if the authors can validate their method on more tasks, such as human segmentation and image denoising, etc.", "summary_of_the_paper": "This paper proposes Prospect Pruning (ProsPr) to handle the problem of short-sightedness of existing methods. The main idea is to use meta-gradients through the first few steps of optimization to determine which weights to prune.", "summary_of_the_review": "See the main review.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "CYiNd44BnL9", "reviewer": "Reviewer_MHq7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\nThe overall method is well explained and easy to follow.\n\nWeakness:\nThe performance of this work can be decoupled into two aspects:\n1) If the meta-gradients w.r.t. masks ($\\textbf{c}$ is vital, or single-step gradients (i.e. $M=1$) may also work.\n2) If dropping small terms (Eq. 15 to 16) is indeed negligible.\n\nI did not see ablation studies on these two aspects. Therefore I propose some further studies:\n1) If the single-step gradient is used (no meta-gradients) to update the mask with the same total number of iterations (e.g. 1024 steps with a batch size of 256 for ResNet-50), what will be the pruning results? This is actually more GPU-memory efficient since no extra computational graphs are stored. Both \"pruning by optimizing the mask\" (this work) and magnitude pruning could be studied.\n2) If we do not drop small terms in Eq. 15, what will be the pruning results? Smaller $M$ would be acceptable if storing computational graph is GPU-memory consuming.\n\nIn short, I think it is important to understand which part is vital to pruning: 1) optimizing the mask or magnitude pruning; 2) mata-gradients or just single-step gradients with more steps; 3) dropping or keeping the small terms.", "summary_of_the_paper": "This work proposed a new efficient pruning method, by leveraging both loss sensitivity (\"saliency score\") during a few initial training steps. The authors leveraged meta-gradients with appropriate approximations to stabilize and speed up the pruning.", "summary_of_the_review": "I generally think the novelty of this work is a bit trivial. Leveraging information from early training steps is a reasonable strategy, but is still straightforward.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is proposed to address neural network pruning at initialization with the help of meta-gradients considering the high-order relations between loss and optimization of trainable sub-network. The paper is well organized and written with the clear logic. The discussions of related works, as well as their limitations, are comprehensive. To verify the proposed method, the authors have tested it on various benchmarks with different settings. Overall, the meta-learning idea for model pruning is relatively new, which may bring more inspirations to the community.
Decision: Accept (Poster)