Probabilistic Neural Pruning via Sparsity Evolutionary Fokker-Planck-Kolmogorov Equation

We sincerely appreciate Reviewer Y7sb's time, effort, insightful feedback, and constructive suggestions, which are valuable in helping us refine our work. Below, we address the concerns raised and provide additional discussions and justifications to clarify our contribution.

W1. Using mini-batch gradient in SFPK.

Response to W1. We thanks the Reviewer-Y7sb for his valuable comment and we will add the following discussions into the Appendix A.1.

"Specifically, we would like to emphasize that the mini-batch gradients involved in Algorithm 2 is crucial for mask particle simulation. In theory, the SFPK dynamic should be initialized at $\pi_0^*$, the optimal mask distribution within $M_0^{\rho}$. Since $\rho > 0$, $\pi_0^*$ is closed to, but not identical to, the delta distribution centered on the fully dense mask. In practice, we approximate $\pi_0^*$ by running SFPK with stochastic mini-batch gradient from the delta distribution at the fully dense model. During the SFPK particle simulation process, the stochasticity of the mini-batch gradients enhances robustness and stability in sampling, while also encouraging the exploration of diverse mask particles."

W2. Performance of SNIP and SynFlow in Table 1.

Response to W2. We have thoroughly checked our SNIP and SynFlow codes and confirm they match the official release with recommended hyperparameters (see Section 13.1 of [1]). Our results shows that SNIP and SynFlow fails to achieve satisfactory results in the one-shot pruning without retraining scenario.

We confirm that SynFlow does not consistently work well, likely due to:

1. As a pruning-at-initialization (PAI) method, SynFlow is less effective for post-training pruning (PTP) as its data-agnostic SynFlow score is less informative than data-dependent pruning scores when the data and a well trained weight are accessible.
2. As noted in Section 12 in [1], the SynFlow score can be overly large and numerically unstable. While rescaling weights can mitigate this in PAI, it is not applicable to PTP and may harm the pretrained weights. Hence, SynFlow fails to generalize to new architectures (e.g., ResNet-18/20, WRN-19/20, VGG-16/16bn) and weights (e.g., randomly initialized / pretrained) with the recommended hyperparameter.

We ran PAI experiments on VGG16 CIFAR-100 at 90% sparsity using our code and the official SynFlow code. Our code achieved 56.41% accuracy, while the official code achieved 56.34%, showing that our SynFlow code is correct. This can be verified by running the following lines in the official SynFlow Github repo:

python main.py --experiment singleshot --model vgg16 --dataset cifar100 --model-class lottery --optimizer momentum --train-batch-size 128 --pre-epochs 0 --post-epochs 160 --lr 0.1 --lr-drops 60 120 --lr-drop-rate 0.1 --weight-decay 0.0001 --pruner synflow --compression 1.0 --prune-epochs 100 --mask-scope global

We also confirm that SNIP does not consistently perform well, likely because the Taylor expansion score it uses becomes less effective when a significant number of parameters are removed simultaneously. Furthermore, we conjecture that its limited performance stems from SNIP's greedy pruning strategy, which prevents it from regrowing prematurely pruned weights, potentially leading to suboptimal results.

[1]. Hidenori Tanaka, Daniel Kunin, et al. Pruning neural networks without any data by iteratively conserving synaptic flow. (NeurlPS 2020)

W3. Definition of distributional transition $\mathcal{T}[\cdot]$.

Response to W3. In particular, the closed-form expression of the distributional transition $\mathcal{T}[\pi_t]$ is given by $- \nabla \cdot[T(\cdot; \pi_t) \pi_t]$ as proven in Proposition 3. Therefore, it is difficult to introduce the formal definition of $\mathcal{T}[\cdot]$ until Proposition 3.

To enhance clarity, we include a footnote after Equation (2) stating, "The formal expression of the distributional transition $\partial_t\pi_t =\mathcal{T}[\pi_t]$ will be introduced in Proposition 3."

W4. Definition of polarizer.

Response to W4. We add the formulation of mask polarizer in Section 4.1. And we refer the readers to the Appendix A.3 in the original paper of PSO [1.] for more discussions and implementation details of the mask polarizer.

[1.]. Zhanfeng Mo, Haosen Shi, and Sinno Jialin Pan. Pruning via sparsity-indexed ODE: a continuous sparsity viewpoint. (ICML 2023)

W5. $\lambda=0$ result in Figure 3.

Response to W5. We have added the $\lambda=0$ result in Figure 3 (left). In practice, SFPK with $\lambda=0$ still achieves decent performance compared to the iterative magnitude pruner. However, increasing $\lambda$ beyond 0 leads to further performance improvements.

Dear Reviewer Y7sb,

Thanks for your reply. We have included valid comparison on CIFAR-100, WRN-20 under the pruning-at-initialization setting in the revised manuscript.

As mentioned in both the Summary of rebuttal revision, Line 8 and our Response to Reviewer 3E6s-Q2, Part 2, we compare our SFPK against SNIP, SynFlow, GraSP, and Mag under the pruning-at-initialization (PaI) setting. Specifically, we randomly intialize a WRN-20, prune it to 95% sparsity, and retrain it for epochs to convergence. The average test accuracy among of final epochs is reported based on 3 trials.

As shown in Appendix C, Table 9, our SFPK illustrates better performance than other baselines. Furthermore, the performance of SFPK can be evidently improved when a well-trained checkpoint is provided. This evidence helps clarify the significance of the 'rational' constraint.

Due to the limited time and computation resource, the PaI experiment on ImageNet is still pending (i.e. the one-shot prune w/ fine-tuned version of Table 1). We would add the associated results to the final version as long as the experiments are completed.

Thank the authors for the additional experiments and discussions. I like this paper and appreciate the effort put into it, so I will increase my score from 6 to 8.

We sincerely appreciate Reviewer 3JmV's time, effort, insightful feedback, and constructive suggestions, which are valuable in helping us refine our work. Below, we address the concerns raised and provide additional discussions and justifications to clarify our contribution.

W1. Computational overhead w.r.t $n$ and $K$.

Response to W1. A detailed complexity analysis of SFPK is included in Appendix A.1.

In particular, the cost of applying SFPK pruner (with $n$ discretization steps and $K$ mask particles) equals to training the model for $n\times K$ mini-batches plus a marginal $O(nKd)$ inner-product computation.

W2. Comparison with iterative magnitude pruning.

Response to W2. In Table 2b and Table 2c in Section 5.2, we compared SFPK against two iterative magnitude pruning baselines, i.e. the Mag (iterative global magnitude pruning) and GMP + LS (iterative magnitude pruning with improved sparsity configuration) [1]. The experiment shows that SFPK outperforms iterative magnitude pruning variants on both MobileNet-V1 and ResNet-50, under ImageNet-1K.

[1]. Trevor Gale, Erich Elsen, et al. The state of sparsity in deep neural networks, 2019.

W3. Diversity of SFPK particles.

Response to W3. We conduct additional experiment to study the diversity of the simulated particle distribution of SFPK. Specifically, at each intermediate step of SFPK, we track the mean and standard deviation of the loss function (the energy) and the average sparsity deviation of each mask particle. The average relative mask deviation of the $i$-mask is computed by $1/(n-1)\sum_j\|\mathbf{m}^i_t - \mathbf{m}^j_t\|_2^2/ \|\mathbf{m}^i_t\|_2^2$. We visualize the mean and standard deviation (across $10$ mask particles) of both the loss value and the relative mask deviation.

As shown in Figure 5, Appendix C, as sparsity increases, the energy of each mask increases without drastic explosion. Meanwhile, the diversity of the masks progressively grows. This suggests that the SFPK mask distribution does not collapse and it generates diverse yet effective mask particles as sparsity decreases, allowing us to sample performant masks at the desired sparsity level.

We sincerely appreciate Reviewer TsAq time, effort, insightful feedback, and constructive suggestions, which are valuable in helping us refine our work. Below, we address the concerns raised and provide additional discussions and justifications to clarify our contribution.

W1. Computation cost vs DST methods.

Response to W1. As shown in Appendix C, we compare the training and inference FLOPs of SFPK with 3 dynamic sparse training (DST) methods, i.e. RigL [1] (Table 14), SWAMP [2] (Table 12), and MEST [3] (Table 16). We reported the training FLOPs in units of $1e18$ and inference FLOPs in units $1e9$.

Specifically, Table 14 indicates that SFPK is less expensive than RigL in terms of training FLOPs.

[1]. Utku Evci, Trevor Gale, et al. Rigging the lottery: Making all tickets winners. (ICML 2020)

[2]. Moonseok Choi, Hyungi Lee, et al. Sparse weight averaging with multiple particles for iterative magnitude pruning. (ICLR 2024)

[3]. Geng Yuan, Xiaolong Ma, et al. Mest: Accurate and fast memory-economic sparse training framework on the edge. (NeurlPS 2021)

W2. Reason for using polarizer.

Response to W2. To analyze mask evolution under an infinitesimal $dt$-sparsity increment, we generalize discrete masks to continuous-valued and nearly-sparse soft masks. Accordingly, the hard sparsity measure $1 - \|m\|_0 / d$ is extended to a soft sparsity measure $1 - \|m\|_2 / d$.

To study the sparsity evolution of masks, an ideal polarization operator $P(m)$ is expected to preserve the sparsity of soft masks, satisfying the condition $1 - \|m\|_2 / d = 1 - \|P(m)\|_2 / d$. In this case, the soft sparsity of the nearly-sparse mask $P(m)$ can be evaluated directly by computing the $\ell_2$ norm of the soft, unpolarized mask $m$.

In contrast, standard softmax or sigmoid operators do "sparsify" the mask, but they fail to preserve the soft sparsity of the resulting "sparsified masks." This issue can be resolved through renormalization, i.e. setting $P(m) = \text{softmax}(m) \cdot \|m\|_2 / \|\text{softmax}(m)\|_2.$

Generally, designing better mask polarization operators for neural pruning is an orthogonal research direction to this paper, and we will postpone this to future work.

W3. Comparisons in Table 1.

Response to W3. In the post-training scenario (e.g. pruning a pretrained dense model), SFPK shares the same objective with Mag, SNIP, and WF: all of them are primarily designed to minimize the pruning-induced accuracy drop (i.e. pruning error).

We will consider removing the results of SynFlow from Table 1, as it was not designed to minimize the loss upon sparsification.

To minimize loss upon sparsification, our SFPK use a particle simulation approach, while Mag, SNIP, and WF minimize a Taylor-expansion-based pruning error bound, i.e. $L(\theta\odot m) - L(\theta)\approx \nabla L(\theta)^\top \delta\theta+1/2\ \delta\theta^\top \nabla^2L(\theta)\delta\theta$, where $\delta\theta\triangleq (\theta\odot m) - \theta$ is the mask perturbation. Then, Mag aims to minimize $\|\delta\theta\|$ and SNIP aims to minimize $\|\nabla L(\theta)^\top\delta\theta\|$, both neglecting the second-order term. In contrast, WF aims to minimize $\|\delta\theta^\top \nabla^2L(\theta)\delta\theta\|$, neglecting the first order term.

Empirically, our Table 1 shows SFPK is more beneficial in reducing pruning error. It validates that our SFPK theory provides a better approximation of the optimal masks than other baseline score-based methods (e.g., Mag, SNIP, and WF).

Q4. Pruning scope and sparsity calculation.

Response to Q4. We follow the pruning scope and sparsity calculation as in Section S7.1 of WoodFisher [4] (also in DPF [5] and STR [6]): only the weights in fully-connected and convolutional layers are pruned, and none of the batch-norm and bias parameters are pruned if presented. In this paper, we consistently evaluate the sparsity based on '1 - (# remaining parameters) / (# total parameters in convolutional and linear layers)' across different pruning and sparse training methods.

As our denominator (which equals $25502912$) is larger than '# convolutional parameters' (which equals $23454912$), the sparsity metric adopted in our paper is generally higher and more restrictive than the metric '1 - (# remaining parameters) / (# convolutional parameters)' when given the same sparsity level.

Since we consistently use the sparsity metric '1 - (# remaining parameters) / (# total parameters in convolutional and linear layers)', "avoiding pruning linear layers" in SWAMP [7] is a specific choice of 'pruning implementation or configuration', which does not impact the calculation or definition of the sparsity metric (some methods prefer avoiding pruning some specific modules according to some heuristic rules). Additionally, the authors of SWAMP [7] also state that "we empirically checked that including the fully-connected layers as part of the prunable parameters does not affect the results we have obtained".

[4]. Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. (NeurlPS 2020)

[5]. Tao Lin, Sebastian U. Stich, et al. Dynamic model pruning with feedback. (ICLR 2020)

[6]. Aditya Kusupati, Vivek Ramanujan, et al. Soft threshold weight reparameterization for learnable sparsity. (ICML 2020)

[7]. Moonseok Choi, Hyungi Lee, et al. Sparse weight averaging with multiple particles for iterative magnitude pruning. (ICLR 2024)

Q5. Why termed as SFPK?

Response to Q5. We coined the proposed analysis framework as SFPK to highlight that it is a sparsity-indexed FPK equation, in contrast to the thermodynamic time-indexed FPK equation.

W5. Pruning ResNet-50 on ImageNet.

Response to W5. Following the setup in Table 2c, we prune ResNet-50 to the 95% sparsity level using SFPK with the same hyperparameters. Due to the limited computation resource, the reported accuracy is based on one trial, and we will add the results of $2$ extra trials into the final version. As shown in Table 10, Appendix C, our SFPK outperforms WF, STR, GMP, DNW, and RIGL+ERK.

In particular, the experiments in this paper are designed to validate the effectiveness of the proposed SFPK theory in approximating the sparsity evolution of the optimal mask distribution in neural pruning. Thus, we mainly focused on the pruning settings, where the mask and weights are not jointly optimized as in sparse training methods, e.g. RiGL, Spartan. We have added a reference to Spartan in the related work section. Extending SFPK to the sparse training scenario, which requires considering the joint evolution of optimal masks and weights, is deferred to future work.

We have also compared our SFPK-pruner against L1-spred [2] in Table 13 and Figure 6, Appendix C, a SOTA L1 pruning method proposed in 2023. For comprehensiveness, we further include the results of CrAM [3], another pruning method proposed in 2023, into our Table 2c.

To improve clarity, we revise the claim from "our SFPK-pruner demonstrates decent performance compared to SOTA baselines" to "our SFPK-pruner demonstrates decent performance compared to various competitive pruning baselines".

[2]. Liu Ziyin and Zihao Wang. spred: Solving l1 penalty with SGD. (ICML 2023).

[3] Peste, A., et al. CrAM: A Compression-Aware Minimizer. (ICLR 2023).

Q1. Typo in SFPK overhead.

Response to Q1. We apologize for this typo. The phrase "... that of only ONE epoch ..." should be "... that of only 0.3 epoch ...," as stated in the complexity analysis section in Appendix A.1.

Q2. SPFK with random initialization.

Response to Q2. We compare our SFPK against SNIP, SynFlow, GraSP, and Mag under the pruning-at-initialization setting. Specifically, we randomly intialize a WRN-20, prune it to 95% sparsity, and retrain it for $100$ epochs to convergence. The average test accuracy among of final $5$ epochs is reported based on 3 trials.

As shown in Appendix C, Table 9, our SFPK illustrates better performance than other baselines. However, the performance of SFPK can be evidently improved when a well-trained checkpoint is provided. This evidence helps clarify the significance of the 'rational' constraint.

Q3. Diversity of SFPK particles.

Response to Q3. We conduct additional experiment to study the diversity of the simulated particle distribution of SFPK. Specifically, at each intermediate step of SFPK, we track the mean and standard deviation of the loss function (the energy) and the average sparsity deviation of each mask particle. The average relative mask deviation of the $i$-mask is computed by $1/(n-1)\sum_j\|\mathbf{m}^i_t - \mathbf{m}^j_t\|_2^2/ \|\mathbf{m}^i_t\|_2^2$. We visualize the mean and standard deviation (across $10$ mask particles) of both the loss value and the relative mask deviation.

As shown in Figure 5, Appendix C, as sparsity increases, the energy of each mask increases without drastic explosion. Meanwhile, the diversity of the masks progressively grows. This suggests that the SFPK mask distribution does not collapse and it generates diverse yet effective mask particles as sparsity decreases, allowing us to sample performant masks at the desired sparsity level.

We sincerely appreciate Reviewer 3E6s's time, effort, insightful feedback, and constructive suggestions, which are valuable in helping us refine our work. Below, we address the concerns raised and provide additional discussions and justifications to clarify our contribution.

W1. Error bars in tables.

Response to W1. All the results presented in Table 1 and Table 2 are computed based on three trials. We postpone the full results in Table 5 and Table 6 of Appendix C.

W2. Cost comparison in oneshot pruning.

Response to W2. In Table 7, Appendix B, we report the actual and theoretical runtime complexities of WF, PSO, and our SFPK-pruner for the one-shot pruning experiment on ImageNet-1K with ResNet-50 and MobileNet-V1 at $80\%$ sparsity. Our results show that, in practice, SFPK outperforms WF while maintaining a comparable computational cost. As expected, the computational cost of SFPK is approximately $10$ times that of PSO, as both methods use $700$ discretization steps, while SFPK additionally simulates $10$ particles.

We also compare the theoretical complexity of our SFPK-pruner with that of PSO and WF. Let $m$ denote the data batch size, $d$ the model size, $K$ the discretization step for both PSO and SFPK-pruner, and $n$ the number of particles in the SFPK-pruner. The theoretical runtime complexity of the SFPK-pruner is $O(nK(d+m))$, the theoretical runtime complexity of PSO is $O(K(d+m))$, and the theoretical runtime complexity of WF is $O(md^2)$. As we will show, for large models where $d$ is extremely large, WF becomes significantly expensive due to the need to compute the Hessian matrix. As expected, the computational cost of SFPK is approximately $n$ times that of PSO, as both methods use $K$ discretization steps, while SFPK additionally simulates $n$ particles.

[1]. Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. (NeurlPS 2020)

W3. Additional oneshot pruning results.

Response to W3. Following the same configuration as described in Section 5.2, we perform one-shot unstructured and structured pruning experiments on the MobileNet-V1, ResNet-50, and DeiT-T architecture to provide a more comprehensive evaluation. As shown in Appendix B, Table 6, our SFPK exhibits competitive performance against other baselines across various scenarios.

Due to limited computational resources, we are unable to complete the one-shot pruning with tuning results during the rebuttal period. However, we will make every effort to include it in the final version. Nonetheless, the gradual pruning results shown in Table 2 have partially validated the effectiveness of our SFPK in the prune-and-retrain scenario.

W4. How rational $\theta^\*$ needs to be in SPFK?

Response to W4. To study how the rationality of $\theta^*$ affects the effectiveness of SFPK, we conduct an ablation study on the initialization of SFPK on WRN-32x4 CIFAR-100. Specifically, we first pretrain a dense WRN-32x4 from scratch for 300 epochs to convergence, and we save the intermediate model checkpoints $\theta_t^*$ at training epoch $t\in\{0, 50, 100, 150, 200, 300\}$. Then, we prune each $\theta_t^*$ to $90\%$ sparsity and retrain it for 100 epochs. We report the training accuracy of each setting based on 3 trials.

As shown in Appendix C, Table 8, the peformance of SFPK increases as the pretraining epochs increases. Specifically, the performance of SFPK retrains at a decent level even when the model is not pretrained to convergence. This empirical evidence suggests the robustness of SFPK w.r.t to the rationality of the model weight.