NTK-SAP: Improving neural network pruning by aligning training dynamics
Pruning neural networks before training has received increasing interest due
to its potential to reduce training time and memory. One popular method is to
prune the connections based on a certain metric, but it is not entirely clear
what metric is the best choice. Recent advances in neural tangent kernel (NTK)
theory suggest that the training dynamics of large enough neural networks is
closely related to the spectrum of the NTK. Motivated by this finding, we
propose to prune the connections that have the least influence on the spectrum
of the NTK. This method can help maintain the NTK spectrum, which may help
align the training dynamics to that of its dense counterpart. However, one
possible issue is that the fixed-weight-NTK corresponding to a given initial
point can be very different from the NTK corresponding to later iterates during
the training phase. We further propose to sample multiple realizations of
random weights to estimate the NTK spectrum. Note that our approach is
weight-agnostic, which is different from most existing methods that are
weight-dependent. In addition, we use random inputs to compute the
fixed-weight-NTK, making our method data-agnostic as well. We name our
foresight pruning algorithm Neural Tangent Kernel Spectrum-Aware Pruning
(NTK-SAP). Empirically, our method achieves better performance than all
baselines on multiple datasets.
Yite Wang, Dawei Li, Ruoyu Sun
# NTK-SAP: Improving neural network pruning by aligning training dynamics

###### Abstract

Pruning neural networks before training has received increasing interest due to its potential to reduce training time and memory. One popular method is to prune the connections based on a certain metric, but it is not entirely clear what metric is the best choice. Recent advances in neural tangent kernel (NTK) theory suggest that the training dynamics of large enough neural networks is closely related to the spectrum of the NTK. Motivated by this finding, we propose to prune the connections that have the least influence on the spectrum of the NTK. This method can help maintain the NTK spectrum, which may help align the training dynamics to that of its dense counterpart. However, one possible issue is that the fixed-weight-NTK corresponding to a given initial point can be very different from the NTK corresponding to later iterates during the training phase. We further propose to sample multiple realizations of random weights to estimate the NTK spectrum. Note that our approach is weight-agnostic, which is different from most existing methods that are weight-dependent. In addition, we use random inputs to compute the fixed-weight-NTK, making our method data-agnostic as well. We name our foresight pruning algorithm Neural Tangent Kernel Spectrum-Aware Pruning (NTK-SAP). Empirically, our method achieves better performance than all baselines on multiple datasets. Our code is available at [https://github.com/YiteWang/NTK-SAP](https://github.com/YiteWang/NTK-SAP).

## 1 Introduction

The past decade has witnessed the success of deep neural networks (DNNs) in various applications. Modern DNNs are usually highly over-parameterized, making training and deployment computationally expensive. Network pruning has emerged as a powerful tool for reducing time and memory costs. There are mainly two types of pruning methods: post-hoc pruning (Han et al., 2015; Renda et al., 2020; Molchanov et al., 2019; LeCun et al., 1989; Hassibi and Stork, 1992) and foresight pruning (Lee et al., 2018; Wang et al., 2020; Alizadeh et al., 2022; Tanaka et al., 2020; de Jorge et al., 2020; Liu and Zenke, 2020). The former methods prune the network after training the large network, while the latter methods prune the network before training. In this paper, we focus on foresight pruning.

SNIP (Lee et al., 2018) is probably the first foresight pruning method for modern neural networks. It prunes those initial weights that have the least impact on the initial loss function. SNIP can be viewed as a special case of saliency score based pruning: prune the connections that have the least "saliency scores", where the saliency score is a certain metric that measures the importance of the connections. For SNIP, the saliency score of a connection is the difference of loss function before and after pruning this connection. GraSP (Wang et al., 2020) and Synflow (Tanaka et al., 2020) used two different saliency scores. One possible issue of these saliency scores is that they are related to the initial few steps of training, and thus may not be good choices for later stages of training. Is there a saliency score that is more directly related to the whole training dynamics?
### Our Methodology

Recently, there are many works on the global optimization of neural networks; see, e.g., Liang et al. (2018, 2019, 2021); Venturi et al. (2018); Safran and Shamir (2017); Li et al. (2022); Ding et al. (2022); Soltanolkotabi et al. (2019); Sun et al. (2020a); Lin et al. (2021, 2021); Zhang et al. (2021) and the surveys Sun et al. (2020b); Sun (2020). Among them, one line of research uses neural tangent kernel (NTK) (Jacot et al., 2018) to describe the gradient descent dynamics of DNNs when the network size is large enough. More specifically, for large enough DNNs, the NTK is asymptotically constant during training, and the convergence behavior can be characterized by the spectrum of the NTK. This theory indicates that the spectrum of the NTK might be a reasonable metric for the whole training dynamics instead of just a few initial iterations. It is then natural to consider the following conceptual pruning method: prune the connections that have the least impact on the NTK spectrum.

There are a few questions on implementing this conceptual pruning method.

First, what metric to compute? Computing the whole eigenspectrum of the NTK is too time-consuming. Following the practice in numerical linear algebra and deep learning (Lee et al., 2019; Xiao et al., 2020), we use the nuclear norm (sum of eigenvalues) as a scalar indicator of the spectrum.

Second, what "NTK" matrix to pick? We call the NTK matrix defined for the given architecture with a random initialization as a "fixed-weight-NTK", and use "analytic NTK" (Jacot et al., 2018) to refer to the asymptotic limit of the fixed-weight-NTK as the network width goes to infinity. The analytic NTK is the one studied in NTK theory (Jacot et al., 2018), and we think its spectrum may serve as a performance indicator of a certain architecture throughout the whole training process.1 However, computing its nuclear norm is still too time-consuming (either using the analytic form given in Jacot et al. (2018) or handling an ultra-wide network). The nuclear norm of a fixed-weight-NTK is easy to compute, but the fixed-weight-NTK may be quite different from the analytic NTK. To resolve this issue, we notice a less-mentioned fact: the analytic NTK is also the limit of the expectation (over random weights) of fixed-weight-NTK2, and thus it can be approximated by the expectation of fixed-weight-NTK for a given width. The expectation of fixed-weight-NTK shall be a better approximation of analytic NTK than a single fixed-weight-NTK. Of course, to estimate the expectation, we can use a few "samples" of weight configurations and compute the average of a few fixed-weight-NTKs.

Footnote 1: Please see Appendix B.2 for more discussions from an empirical perspective.

One more possible issue arises: would computing, say, \(100\) fixed-weight-NTKs take \(100\) times more computation cost? We use one more computation trick to keep the computation cost low: for each mini-batch of input, we use a fresh sample of weight configuration to compute one fixed-weight-NTK (or, more precisely, its nuclear norm). This will not increase the computation cost compared to computing the fixed-weight-NTK for one weight configuration with \(100\) mini-batches. We call this "_new-input-new-weight_" (INNW) trick.

We name the proposed foresight pruning algorithm **N**eural **T**angent **K**ernel **S**pectrum-**A**ware **P**runing **(**NTK-SAP**). We show that NTK-SAP is competitive on multiple datasets, including CIFAR-10, CIFAR-100, Tiny-ImageNet, and ImageNet. In summary, our contributions are:

* We propose a theory-motivated foresight pruning method named NTK-SAP, which prunes networks based on the spectrum of NTK.
* We introduce a multi-sampling formulation which uses different weight configurations to better capture the expected behavior of pruned neural networks. A "_new-input-new-weight_" (INNW) trick is leveraged to reduce the computational cost, and may be of independent interest.
* Empirically, we show that NTK-SAP, as a data-agnostic foresight pruning method, achieves state-of-the-art performance in multiple settings.

## 2 Related work and background

**Pruning after training (Post-hoc pruning).** Post-hoc pruning can be dated back to the 1980s (Janowsky, 1989; Mozer and Smolensky, 1989) and they usually require multiple rounds of train-prune-retrain procedure (Han et al., 2015; LeCun et al., 1989). Most of these pruning methods use
magnitude (Han et al., 2015), Hessian (LeCun et al., 1989; Hassibi and Stork, 1992; Molchanov et al., 2019) and other information (Dai et al., 2018; Guo et al., 2016; Dong et al., 2017; Yu et al., 2018) to define the scoring metrics which determine which weights to prune. Some other recent works (Verma and Pesquet, 2021) identify sparse networks through iterative optimization. Most post-hoc pruning methods intend to reduce the inference time, with the exception of lottery ticket pruning (Frankle and Carbin, 2018; Renda et al., 2020) which prunes the networks for the purpose of finding a trainable sub-network.

**Pruning during training.** Another class of pruning algorithms gradually sparsifies a dense neural network during training. Some works use explicit \(\ell_{0}\)(Louizos et al., 2017) or \(\ell_{1}\)(Wen et al., 2016) regularization to learn a sparse solution. Other works introduce trainable masks (Liu et al., 2020; Savarese et al., 2020; Kang and Han, 2020; Kusupati et al., 2020; Srinivas et al., 2017; Xiao et al., 2019) or use projected gradient descent (Zhou et al., 2021) to specifically learn the connections to be remained. DST methods (Mocanu et al., 2018; Bellec et al., 2017; Mostafa and Wang, 2019; Dettmers and Zettlemoyer, 2019; Liu et al., 2021b; Evici et al., 2020; Liu et al., 2021c;a) train the neural network with a fixed parameter count budget while gradually adjust the parameters throughout training. Moreover, another line of work (You et al., 2019) tries to identify the winning ticket at early epochs. These methods can reduce some training time, but often do not reduce the memory cost of training.

**Pruning before training (foresight pruning).** SNIP (Lee et al., 2018) is arguably the first foresight pruning method. Later, GraSP (Wang et al., 2020) uses Hessian-gradient product to design the pruning metric. To address the layer collapse issue of SNIP and GraSP at high sparsity ratios, Synflow (Tanaka et al., 2020) prunes the network iteratively in a data-independent manner. FORCE and Iterative SNIP (de Jorge et al., 2020) also intend to improve the performance of SNIP at high sparsity ratios. Lee et al. (2019) uses orthogonal constraint to improve signal propagation of networks pruned by SNIP. Neural tangent transfer (Liu and Zenke, 2020) explicitly lets subnetworks learn the full empirical neural tangent kernel of dense networks. PropsPr (Alizadeh et al., 2022) also mentions the "shortsightedness" of existing foresight pruning methods and proposes to use meta-gradients to utilize the information of the first few iterations.

## 3 General pipeline for foresight pruning methods

In this section, we review the formulation of pruning-at-initialization problem and revisit several existing foresight pruning methods. Here we consider Synflow (Tanaka et al., 2020), GraSP (Wang et al., 2020), SNIP (Lee et al., 2018) and Iterative SNIP (de Jorge et al., 2020). When given a dataset \(\mathcal{D}=(\mathbf{x}_{i},\mathbf{y}_{i})_{i=1}^{n}\) and an untrained neural network \(f(\cdot;\mathbf{\theta})\) parameterized by \(\mathbf{\theta}=\{\theta^{1},\theta^{2},\cdots,\theta^{p}\}\in\mathbb{R}^{p}\), we are interested in finding a binary mask \(\mathbf{m}=\{m^{1},m^{2},\cdots,m^{p}\}\in\{0,1\}^{p}\) which minimizes the following equation:

\[\min_{\mathbf{m}}\mathcal{L}(\mathcal{A}(\mathbf{\theta}_{0},\mathbf{ m})\odot\mathbf{m};\mathcal{D}) =\min_{\mathbf{m}}\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}(f(\mathbf{x}_{i}; \mathcal{A}(\mathbf{\theta}_{0},\mathbf{m})\odot\mathbf{m}),\mathbf{y}_{i})\] (1) s.t. \[\mathbf{m}\in\{0,1\}^{p},\quad\|\mathbf{m}\|_{0}/p\leq d=1-k\]

where \(k\) and \(d=1-k\) are the target sparsity and density, \(\mathbf{\theta}_{0}\) is an initialization point which follows a specific weight initialization method \(\mathcal{P}\) (e.g. Xavier initialization), \(\odot\) denotes the Hadamard product, and \(\mathcal{A}\) denotes a training algorithm (e.g. Adam) which takes in mask \(\mathbf{m}\) as well as initialization \(\mathbf{\theta}_{0}\), and outputs trained weights at convergence \(\mathbf{\theta}_{\text{final}}\odot\mathbf{m}\).

Since directly minimizing Equation (1) is intractable, existing foresight pruning methods assign each mask and corresponding weight at initialization \(\theta_{0}^{j}\) with a saliency score \(S(\theta_{0}^{j})\) in the following generalized form:

\[S(m^{j})=S(\theta_{0}^{j})=\frac{\partial\mathcal{I}}{\partial m^{j}}=\frac{ \partial\mathcal{I}}{\partial\theta_{0}^{j}}\cdot\theta_{0}^{j} \tag{2}\]

where \(\mathcal{I}\) is a function of parameters \(\mathbf{\theta}_{0}\) and corresponding mask \(\mathbf{m}\). Intuitively, Equation (2) measures the influence of removing connection \(m^{j}\) on \(\mathcal{I}\). Once the saliency score for each mask is computed, we will only keep the top-\(d\) masks. Specifically, the masks with the highest saliency scores will be retained (\(m^{j}=1\)), while other masks will be pruned (\(m^{j}=0\)).
**SNIP.** SNIP (Lee et al., 2018) defines \(S_{\text{SNIP}}(m^{j})=\left|\frac{\partial\mathcal{L}(\mathbf{\theta}_{0}\odot\mathbf{m};D )}{\partial\theta_{0}^{j}}\cdot\theta_{0}^{j}\right|\), where \(D\subset\mathcal{D}\) is a subset of the whole training set. SNIP prunes the whole network in one round.

**Iterative SNIP.** The saliency score of Iterative SNIP3 (de Jorge et al., 2020) remains the same as the original SNIP. However, multiple rounds of pruning are performed.

Footnote 3: We do not consider FORCE in this work since we find empirically that FORCE is less stable than Iterative SNIP. See Appendix F for a detailed comparison.

**GraSP.** GraSP (Wang et al., 2020) defines \(S_{\text{GraSP}}(m^{j})=-\left[H(\mathbf{\theta}_{0}\odot\mathbf{m};D)\frac{\partial \mathcal{L}(\mathbf{\theta}_{0}\odot\mathbf{m};D)}{\partial\theta_{0}}\right]^{j} \cdot\theta_{0}^{j}\), where \(H\) is the Hessian matrix. Additionally, GraSP is also a single round foresight pruning method.

**Synflow.** Synflow (Tanaka et al., 2020) uses \(S_{\text{Syaffow}}(m^{j})=\left|\frac{\partial\sum_{k}f(1;|\mathbf{\theta}_{0}| \odot\mathbf{m})_{k}}{\partial|\theta_{0}^{j}|}\cdot|\theta_{0}^{j}|\right|\) as saliency score, where \(f_{k}\) denotes the \(k\)-th output logit, \(\mathbb{1}\) denotes an input full of \(1\)s, and \(f(\cdot;|\theta_{0}^{j}|)\) means replacing all the parameters of the neural network by their absolute values. Similar to Iterative SNIP, Synflow is also an iterative pruning method but in a data-agnostic manner.

## 4 Review of neural tangent kernel

Recent works (Jacot et al., 2018; Lee et al., 2019; Arora et al., 2019; Xiao et al., 2020) study the training dynamics of large DNNs under gradient descent. Formally, let \(\mathcal{D}=\{(\mathbf{x}_{i},\mathbf{y}_{i})\}_{i=1}^{n}\subset\mathbb{R}^{n _{0}}\times\mathbb{R}^{k}\) denote the training set and let \(\mathcal{X}=\{\mathbf{x}:(\mathbf{x},\mathbf{y})\in\mathcal{D}\}\), \(\mathcal{Y}=\{\mathbf{y}:(\mathbf{x},\mathbf{y})\in\mathcal{D}\}\) denote inputs and labels, respectively. Consider a neural network \(f(\cdot;\mathbf{\theta})\) parameterized by weights \(\mathbf{\theta}\). Let \(\mathcal{L}(\mathbf{\theta})=\sum_{(\mathbf{x},\mathbf{y})\in(\mathcal{X},\mathcal{ Y})}\ell(f(\mathbf{x};\mathbf{\theta}),\mathbf{y})\) denote the empirical loss where \(\ell\) is the loss function. Thus, \(f(\mathcal{X};\mathbf{\theta})\in\mathbb{R}^{kn\times 1}\) is the concatenated vector of predictions for all training inputs. Under continuous time gradient descent with learning rate \(\eta\), the evolution of the parameters \(\mathbf{\theta}_{t}\), where we emphasize the dependence of \(\mathbf{\theta}\) on time \(t\), and the training loss of the neural network \(\mathcal{L}\) can be expressed (Lee et al., 2019) as

\[\dot{\mathbf{\theta}}_{t}=-\eta\nabla_{\mathbf{\theta}_{t}}f(\mathcal{X}; \mathbf{\theta}_{t})^{T}\nabla_{f(\mathcal{X};\mathbf{\theta}_{t})}\mathcal{L} \tag{3}\] \[\dot{\mathcal{L}}=\nabla_{f(\mathcal{X};\mathbf{\theta}_{t})}\mathcal{ L}^{T}\nabla_{\mathbf{\theta}_{t}}f(\mathcal{X};\mathbf{\theta}_{t})\dot{\mathbf{\theta}}_{t} =-\eta\nabla_{f(\mathcal{X};\mathbf{\theta}_{t})}\mathcal{L}^{T}\hat{\mathbf{\Theta}}_ {t}(\mathcal{X},\mathcal{X})\nabla_{f(\mathcal{X};\mathbf{\theta}_{t})}\mathcal{L} \tag{4}\]

where we obtain the NTK at time \(t\):

\[\hat{\mathbf{\Theta}}_{t}(\mathcal{X},\mathcal{X};\mathbf{\theta}_{t})\triangleq\nabla _{\mathbf{\theta}_{t}}f(\mathcal{X};\mathbf{\theta}_{t})\nabla_{\mathbf{\theta}_{t}}f( \mathcal{X};\mathbf{\theta}_{t})^{T}\in\mathbb{R}^{kn\times kn}. \tag{5}\]

It is proven that when the width of the neural network becomes sufficiently large (Jacot et al., 2018; Arora et al., 2019), the fixed-weight-NTK converges to a deterministic matrix \(\mathbf{\Theta}(\mathcal{X},\mathcal{X})\), which we term analytic NTK. The analytic NTK remains asymptotically constant throughout training, i.e. \(\mathbf{\Theta}_{t}(\mathcal{X},\mathcal{X})\approx\mathbf{\Theta}_{0}(\mathcal{X}, \mathcal{X}),\ \forall t\geq 0\).

**The condition number of the NTK is related to neural network optimization.** The condition number of the NTK is known to be important in neural network optimization. Specifically, let the eigenvalue decomposition of the NTK at initialization be \(\mathbf{\Theta}_{0}(\mathcal{X},\mathcal{X})=U^{T}\Lambda U\) where \(\Lambda=\text{diag}(\lambda_{1},\cdots,\lambda_{kn})\) is the diagonal matrix containing all the eigenvalues of the NTK.

For large-width networks, the mean prediction of the neural network in the eigenbasis of NTK can be approximately described as \((U\mathbb{E}\left[f(\mathcal{X})\right])_{i}=(\mathbf{I}-e^{-\eta\lambda_{i}t} (U\mathcal{Y}))_{i}\)(Xiao et al., 2020). Therefore, when using the largest possible learning rate \(\eta\sim 2/\lambda_{1}\)(Lee et al., 2019), the convergence rate of the smallest eigenvalue is related to \(\lambda_{1}/\lambda_{kn}\), which is the condition number of the NTK. Empirically, the condition number of the NTK has been successfully used to identify promising architectures in the neural architecture search (NAS) field (Chen et al., 2021; Wang et al., 2022).

**The entire spectrum of the NTK is a better measure.** However, the smallest and the largest eigenvalues only provide very limited information on the optimization property of the DNNs. In the classical optimization theory, the convergence speed depends on all eigenvalues of the Hessians, as different eigenvalues control the convergence of different eigen-modes. Similarly, for neural networks, a few works (Kopitkov & Indelman, 2020; Su & Yang, 2019) have shown that the convergence speed depends on eigenvalues other than the smallest eigenvalue. In this sense, compared to only considering the condition number of the NTK, controlling the entire eigenspectrum may be a better way to ensure the satisfactory optimization performance of DNNs.
## 5 Our method

### Account for the effects of the entire eigenspectrum with trace norm

As discussed in Section 4, the eigenspectrum of the NTK captures the training dynamics of large DNNs, and thus we intend to use the NTK spectrum to guide foresight pruning. More specifically, we would like to keep the spectrum of the pruned network close to that of the large network, so that the pruned network may perform similarly to the large network.

As discussed in Section 3, the saliency-score-based method requires using a single metric to capture the optimization performance. This single metric is often a scalar, and thus we want to generate a scalar to capture the information of the whole spectrum. A very natural choice is to use the average value of all eigenvalues. In fact, Pennington et al. (2017, 2018); Shu et al. (2021) indeed analyze the optimization property of DNNs through the average eigenvalues of the NTK (or a related matrix called input-output Jacobian) explicitly or implicitly. Additionally, using the average eigenvalues have also a few good properties which make it more suitable than other metrics for pruning:

**Computational efficiency.** To begin with, the computational cost of computing all the eigenvalues would be too expensive. On the contrary, trace norm can be efficiently computed, which we will show in Section 5.2.

**Potentially preserve the entire spectrum.** As indicated in Section 3, the saliency score for each weight measures the change of the metric (here it is the trace norm of the NTK) when removing the weight. More precisely, removing weights with a smaller saliency score is less likely to induce significant change to the eigenvalue distribution. Hence, using the NTK trace norm as the indicator is beneficial in maintaining the entire spectrum.

### Approximating the trace norm of the NTK

**Nuclear norm as a proxy of the whole spectrum.** The NTK is a symmetric matrix. Hence, the trace of the NTK is equivalent to the squared Frobenius norm of parameter-output Jacobian, i.e.

\[\|\mathbf{\Theta}_{0}\|_{*}=\|\mathbf{\Theta}_{0}\|_{\text{tr}}=\|\nabla_{\mathbf{ \theta}_{0}}f(\mathcal{X};\mathbf{\theta}_{0})\|_{F}^{2}. \tag{6}\]

**Finite difference approximation.** Computing the Frobenius norm of the Jacobian matrix \(\nabla_{\mathbf{\theta}_{0}}f(\mathcal{X};\mathbf{\theta}_{0})\in\mathbb{R}^{kn\times p}\) is memory and computation heavy. More specifically, a naive solution that constructs the Jacobian matrix explicitly for a batch of training data may require at least \(k\) rounds of backward propagation, where \(k\) is as large as 1000 for ImageNet classification task. So approximation is needed.

One possible way is to lower-bound it using loss gradient with respect to weights \(\|\nabla_{\mathbf{\theta}_{0}}\mathcal{L}\|_{2}^{2}\). Such a method has been studied in FORCE (de Jorge et al., 2020) but found to be inferior to Iterative SNIP.

Alternatively, we use the finite difference expression \(\frac{1}{\epsilon}\mathbb{E}_{\Delta\mathbf{\theta}\sim\mathcal{N}(\mathbf{0},\mathbf{c })}\left(\|f(\mathcal{X};\mathbf{\theta}_{0})-f(\mathcal{X};\mathbf{\theta}_{0}+ \Delta\mathbf{\theta})\|_{2}^{2}\right]\) to approximate the Jacobian norm. Please see Appendix O for the justification.

**Approximating the training set with pruning set.** We use the following finite approximation expression to define a saliency score:

\[S_{\text{NTK-SAP}}(m^{j})=\left|\frac{\partial\mathbb{E}_{\Delta\mathbf{\theta} \sim\mathcal{N}(\mathbf{0},\mathbf{c})}\left[\|f(\mathbf{X}_{D};\mathbf{\theta}_{0}\odot \mathbf{m})-f(\mathbf{X}_{D};(\mathbf{\theta}_{0}+\Delta\mathbf{\theta})\odot\mathbf{ m})\|_{2}^{2}\right]}{\partial m^{j}}\right| \tag{7}\]

where we use the inputs of the pruning dataset \(\mathbf{X}_{D}\) to approximate the whole training inputs \(\mathcal{X}\). Our foresight pruning method will use the same framework shown in Section 3, which prunes the weights with the least saliency scores defined in Equation (7). Nevertheless, a few more tricks are needed to make this method practical, as discussed next.

### NTK-SAP: multi-sampling formulation

The last subsection computes the saliency score using the fixed-weight-NTK. We think the spectrum of analytic NTK may serve as a performance indicator of a certain architecture. While a single fixed-weight-NTK can be viewed as an approximation of the analytic NTK, the expectation
of fixed-weight-NTKs (expecation over random draw of weight configurations) shall be a better approximation. Further, to approximate the expectation, we sample a few independent realization of weight configurations and compute the average of their fixed-weight-NTKs. This helps explore the parameter space and better capture the expected behavior of the pruned neural networks.

Specifically, with \(R\) random weight configurations \(\mathbf{\theta}_{0,r}\stackrel{{\mathrm{iid}}}{{\sim}}\mathcal{P}\), we use the following stabilized version of the saliency score:

\[S_{\text{NTK-SAP}}(m^{j})=\left|\frac{\partial\frac{1}{R}\sum_{r=1}^{R}\mathbb{ E}_{\Delta\mathbf{\theta}\sim\mathcal{N}(\mathbf{\theta},\mathbf{\mathrm{I}})}\left[\|f( \mathbf{X}_{D};\mathbf{\theta}_{0,r}\odot\mathbf{m})-f(\mathbf{X}_{D};(\mathbf{\theta }_{0,r}+\Delta\mathbf{\theta})\odot\mathbf{m})\|_{2}^{2}\right]}{\partial m^{j}} \right|. \tag{8}\]

There is a major difference between the proposed saliency score and most existing foresight pruning scores: our score is weight-agnostic, i.e., it is not defined by a specific random draw of weight configuration. In other words, the score is mainly determined by the mask structure rather than specific weights.4

Footnote 4: Please refer to Appendix M for a discussion on whether a good pruning solution should be weight-agnostic.

**New-input-new-weight (NINW) trick.** It seems that using multiple weight configurations will cause a large increase in the computational cost. Nevertheless, by utilizing the fact that the training subset for pruning \(D=(\mathbf{x}_{i},\mathbf{y}_{i})|_{i=1}^{|D|}\subset\mathcal{D}\) is usually fed into neural networks in multiple batches, we further approximate Equation (8) with NINW trick as

\[S_{\text{NTK-SAP}}(m^{j})=\left|\frac{\partial\frac{1}{[D]}\sum_{i=1}^{|D|} \mathbb{E}_{\Delta\mathbf{\theta}\sim\mathcal{N}(\mathbf{\theta},\mathbf{\mathrm{I}})} \left[\|f(\mathbf{x}_{i};\mathbf{\theta}_{0,i}\odot\mathbf{m})-f(\mathbf{x}_{i};( \mathbf{\theta}_{0,i}+\Delta\mathbf{\theta})\odot\mathbf{m})\|_{2}^{2}\right]}{ \partial m^{j}}\right|. \tag{9}\]

**Replacing pruning set with random input.** Furthermore, we find that replacing the training subset \(D\) with standard Gaussian noise \(\mathcal{Z}\sim\mathcal{N}(\mathbf{0},\mathbf{\mathrm{I}})\) can generate comparable performance, which makes NTK-SAP purely data-agnostic. As for Gassuan perturbation, ideally, using multiple draws will indeed give a more accurate approximation. However, we empirically find that a single draw of Gaussian perturbation is good enough to generate satisfactory performance. Hence, by using one random draw of \(\Delta\mathbf{\theta}\) we have the final data-agnostic NTK-SAP with its saliency score

\[S_{\text{NTK-SAP}}(m^{j})=\left|\frac{\partial\frac{1}{[D]}\sum_{i=1}^{|D|} \left[\|f(\mathbf{z}_{i};\mathbf{\theta}_{0,i}\odot\mathbf{m})-f(\mathbf{z}_{i};( \mathbf{\theta}_{0,i}+\Delta\mathbf{\theta}_{i})\odot\mathbf{m})\|_{2}^{2}\right]}{ \partial m^{j}}\right| \tag{10}\]

where \(\mathbf{\theta}_{0,i}\stackrel{{\mathrm{iid}}}{{\sim}}\mathcal{P}\), \(\Delta\mathbf{\theta}_{i}\stackrel{{\mathrm{iid}}}{{\sim}}\mathcal{N}( \mathbf{0},\mathbf{\mathrm{\epsilon}}\mathbf{\mathrm{I}})\) and \(\mathbf{z}_{i}\stackrel{{\mathrm{iid}}}{{\sim}}\mathcal{N}(\mathbf{0},\mathbf{\mathrm{I}})\).

In addition, to potentially avoid layer collapse, we also iteratively prune neural networks like Synflow in \(T\) rounds.5 The major steps of the NTK-SAP algorithm are given in Algorithm 1.

Footnote 5: Please refer to Appendix N for a discussion on why NTK-SAP could potentially avoid layer collapse.

**Computational cost.** We compare the computational cost of NTK-SAP with Iterative SNIP as they are both iterative methods and use multi-batch inputs. Since the overhead of the reinitialization process is negligible, NTK-SAP needs roughly double the computational cost of Iterative SNIP due to two forward passes with the computational graph in the formula. However, one should note that NTK-SAP is a data-agnostic method that can be computed beforehand, and the computation time can be reduced with a smaller \(T\); thus, the computational overhead is reasonable. Effects of \(T\), a detailed comparison of pruning time, and FLOPs comparison can be found in Section 6.4, Appendix D.1, and Appendix D.2, respectively.

## 6 Experiments

In this section, we empirically evaluate the effectiveness of our proposed method, NTK-SAP (red). We compare with two baseline methods random pruning (gray) and magnitude pruning (magenta), as
well as foresight pruning algorithms SNIP (blue) (Lee et al., 2018), Iterative SNIP (yellow) (de Jorge et al., 2020b), GraSP (green) (Wang et al., 2020) and Synflow (cyan) (Tanaka et al., 2020) across multiple architectures and datasets.

**Target datasets, models, and sparsity ratios.** We develop our code based on the original code of Synflow (Tanaka et al., 2020). On smaller datasets, we consider ResNet-20 (He et al., 2016) on CIFAR-10, VGG-16 (Simonyan and Zisserman, 2014) on CIFAR-100 (Krizhevsky et al., 2009) and ResNet-18 on Tiny-ImageNet. For ImageNet (Deng et al., 2009) experiments, we consider ResNet-18 and ResNet-50. For the training protocols, we roughly follow the settings in Frankle et al. (2020b). Details of training hyperparameters can be found in Appendix A. We divide the sparsity ratios into three ranges: trivial sparsity ratios (36.00%, 59.04% and 73.80%), mild sparsity ratios (83.22%, 89.30% and 93.12%) and extreme sparsity ratios (95.60%, 97.17% and 98.20%).

Figure 1: Performance of NTK-SAP against other foresight pruning methods. Results are averaged over 3 random seeds, and the shaded areas denote the standard deviation. We highlight SNIP and GraSP as they perform best among baselines on mild and extreme density ratios, respectively.
**Implementation details.** Following the original implementation, SNIP, GraSP, and Iterative SNIP are run by setting BatchNorm layers as train mode while Synflow metric is evaluated under evaluation mode. The number of training samples for pruning we use is ten times of the number of classes as suggested by the original implementations of SNIP and GraSP. For iterative methods, including Iterative SNIP and Synflow, we prune all the weights in 100 rounds in an exponential schedule following Tanaka et al. (2020); Frankle et al. (2020b). Further, following Frankle et al. (2020b), Kaiming normal initialization is used for all models. We prune networks using a batch size of 256 for CIFAR-10/100 and Tiny-ImageNet datasets and a batch size of 128 for ImageNet experiments.

For NTK-SAP, we use the same number of Gaussian noise batches as in Iterative SNIP and prunes neural network progressively in \(T=100\) rounds for CIFAR-100 and Tiny-ImageNet. Since the pruning dataset for CIFAR-10 only contains 1 batch of the training sample, we alternatively sample \(R=5\) independent initialization points and prune progressively in \(T=20\) rounds, keeping the overall computational cost roughly the same. More details can be found in Appendix A.

### Results on CIFAR and Tiny-ImageNet

In Figure 1 we evaluate NTK-SAP using ResNet-20 on CIFAR-10, VGG-16 on CIFAR-100 and ResNet-18 on Tiny-ImageNet. It can be observed that NTK-SAP is competitive for all sparsity ratios.

**Trivial and mild sparsity ratios.** For trivial sparsity ratios, all methods show good performance except GraSP, magnitude, and random pruning. On ResNet-20, all methods are similar, while random pruning and GraSP show slightly worse performance. On VGG-16 (ResNet-18), similar trend can be observed while the performance gap is larger. NTK-SAP shows consistent good performance, especially on ResNet-18 (TinyImageNet).

For mild sparsity ratios, NTK-SAP is the only method that performs well across all datasets and models. On ResNet-20, SNIP and Iterative SNIP take the lead with NTK-SAP while suffering from a performance drop starting at 93.12% sparsity. Magnitude pruning and GraSP have similar performance to NTK-SAP after 89.30%. On VGG-16 and ResNet-18, NTK-SAP and SNIP remain dominant, while NTK-SAP outperforms SNIP.

**Extreme sparsity ratios.** NTK-SAP's superior performance can be seen at extreme sparsity ratios, especially 98.20%. On the CIFAR-10 dataset, GraSP and magnitude pruning are competitive as in matching sparsity ratios. On VGG-16, SNIP first dominates with NTK-SAP and then becomes much inferior. At the sparsest level 98.20%, magnitude, Iterative SNIP, and GraSP show similar results, about 1 percentage point lower than NTK-SAP.

### Results on larger dataset

We further evaluate the performance of NTK-SAP on more difficult tasks: we run experiments on the ImageNet dataset using ResNet-18 and ResNet-50 with pruning ratios {89.26%, 95.60%}.6 We compare NTK-SAP against Synflow, SNIP, Iterative SNIP, GraSP, Magnitude, and Random pruning.

Footnote 6: We choose these sparsity ratios as they include matching and extreme sparsity ratios and close to {90%, 95%} as used by Iterative SNIP (de Jorge et al., 2020b) and ProsPr (Alizadeh et al., 2022).

The results are shown in Table 1. We find that SNIP, which performs reasonably well for smaller datasets, is no longer competitive. For ResNet-50, NTK-SAP beats other pruning methods with margins at least 0.5% and 1% for sparsity ratio 89.26% and 95.60%, respectively. For ResNet-18, NTK-SAP slightly outperforms magnitude pruning under 89.26% sparsity, and GraSP follows. For sparsity ratio 95.60 %, GraSP is better than magnitude pruning while NTK-SAP still takes the lead.

One should note the surprisingly good performance of magnitude pruning. We hypothesize that magnitude pruning chooses an appropriate layer-wise sparsity ratio since Kaiming normal initialization already gives weights different variances depending on the number of parameters of different layers. Thus, layers with more parameters will be more likely pruned.

### An overall comparison with SNIP and GraSP

In this section, we explicitly compare NTK-SAP with two competitive baseline methods in detail, i.e., SNIP and GraSP, since there is no state-of-the-art foresight pruning algorithms (Frankle et al., 2020b).
We highlight both baseline algorithms in Figure 1. Specifically, SNIP is considered competitive as it generally performs well, especially for trivial and mild sparsity ratios. However, SNIP shows poor performance for ImageNet. On the contrary, GraSP indeed shows promising performance on the ImageNet dataset and extreme sparsity ratios on smaller datasets, while it suffers from poor performance on trivial and mild sparsity ratios of smaller datasets. In this sense, NTK-SAP is the first foresight pruning algorithm that works consistently well on all datasets.

### More ablation studies

We finally evaluate the performance of NTK-SAP when varying the number of iterative steps (\(T\)), which can give a better understanding of how iterative pruning helps NTK-SAP achieve better performance at higher sparsity ratios. We repeat experiments of Section 6.1 with varying \(T\). Results show that increasing \(T\) consistently boosts the performance of NTK-SAP. Moreover, such an increase saturates when \(T\) is roughly \(1/5\) of the adopted value in Section 6.1. Hence, the computational cost of NTK-SAP can be further decreased if we use a smaller value of \(T\). Please refer to Appendix C for more details. We also conduct an ablation study on the effects of perturbation hyperparameter \(\epsilon\). Results show that NTK-SAP is robust to the choices of \(\epsilon\). See Appendix I for more details.

### The eigenspectrum of the fixed-weight-NTK after pruning

To visualize the effectiveness of our proposed NTK-SAP method, we show eigenspectrum of pruned subnetwork on CIFAR-10 dataset with ResNet-20 at sparsity ratios \(73.80\%\). Figure 2 shows that NTK-SAP successfully maintain the eigenspectrum of the fixed-weight-NTK compared to other methods. See Appendix G.1 for the eigenspectrum of the fixed-weight-NTK for other sparsity ratios.

## 7 Conclusion

Pruning at initialization is appealing compared to pruning-after-training algorithms since it circumvents the need for training to identify the masks. However, most of the existing foresight pruning methods do not capture the training dynamics after pruning. We argue that such an issue can be mitigated by using NTK theory. Hence, we propose NTK-SAP, which iteratively prunes neural networks without any data. Additionally, NTK-SAP pays attention to multiple random draws of initialization points. Empirically, our method is competitive in all datasets and models studied in this work. We consider that our algorithm can encourage more works to apply theoretical findings to empirical applications. Additionally, our novel way of combining multiple initialization configurations may further encourage future work to investigate the similarity of pruning and NAS just like Liu et al. (2018); Abdelfattah et al. (2021).

The limitations of the work include the relatively heavy computation nature of iterative pruning methods compared to single-shot pruning methods. We also want to emphasize the great significance of pruning methods in reducing the energy cost of training and applying deep learning models.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Network & \multicolumn{2}{c}{**ResNet-18**} & \multicolumn{2}{c}{**ResNet-50**} \\ \hline Sparsity percentage & 89.26\% & 95.60\% & 89.26\% & 95.60\% \\ \hline (Dense Baseline) & 69.98 & & 76.20 & \\ \hline Synflow & 57.30 & 45.65 & 66.81 & 58.88 \\ SNIP & 57.41 & 45.04 & 60.98 & 40.69 \\ Iterative SNIP & 52.97 & 37.44 & 52.53 & 36.82 \\ GraSP & 58.16 & 49.17 & 67.74 & 59.73 \\ Magnitude & 58.75 & 48.50 & 66.80 & 43.79 \\ Random & 54.48 & 42.80 & 65.30 & 56.53 \\ \hline \hline
**Our method** & **58.87** & **49.43** & **68.28** & **60.79** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Test performance of foresight pruning methods on the ImageNet dataset. Best results are in **bold**; second-best results are underlined.

Figure 2: Eigenvalue distribution of the fixed-weight-NTK on ResNet-20 (CIFAR-10) at sparsity ratio \(73.80\%\).
## 8 Acknowledgement

This work utilizes resources supported by the National Science Foundation's Major Research Instrumentation program, grant No.1725729 (Kindratenko et al., 2020), as well as the University of Illinois at Urbana-Champaign. This paper is supported in part by Hetao Shenzhen-Hong Kong Science and Technology Innovation Cooperation Zone Project (No.HZQSWS-KCCYB-2022046); University Development Fund UDF01001491 from the Chinese University of Hong Kong, Shenzhen; Guangdong Key Lab on the Mathematical Foundation of Artificial Intelligence, Department of Science and Technology of Guangdong Province.

We would like to thank the program chairs and the area chairs for handling our paper. We would also like to thank the reviewers for their time and efforts in providing constructive suggestions on our paper.