BMRS: Bayesian Model Reduction for Structured Pruning

We thank the reviewer for their time and thoughtful review! We are happy that they noted that the paper tackles and important problem, it is well-written and clearly and comprehensively addresses a gap in the literature, and that the experiments are thoughtfully designed. We address their concerns and questions below.

W1. Differences in performance advantage for different settings

The reviewer is referring to the continuous pruning experiments. For MNIST, Fashion-MNIST, and CIFAR10 with an MLP and LeNet, BMRS clearly outperforms the next best baseline in terms of compression vs. accuracy, while for ResNet50 and ViT on CIFAR10 and Tiny-ImageNet, SNR provides slightly higher compression at comparable accuracy. We explain this as follows:

SNR is dependent on selecting a pruning threshold, which we set at a fixed value (1.0) as used in previous work [1]. We show that this results in radically different pruning behavior for different model sizes and datasets. In contrast, BMRS$_{\mathcal{N}}$ does not use any explicit threshold, and its solution is dependent on how good our variational approximation is. We chose to end training at 200 epochs for the larger models, but with further training we expect the compression gap to close.

Additionally, we show two settings of $p_{1}$ for $BMRS_{\mathcal{U}}$ (8 and 4), but with lower values of $p_{1}$ we would expect to see a higher compression rate (e.g., see Figure 4).

W2. Runtime information

This is an excellent point! We have included the average training runtime for a sample of datasets and models in the additional supplement. We will update these numbers for all experiments in the final version. Also note that this is the average total training time, which can vary depending on the load on our compute cluster and which nodes the models are run on. We can provide fair training and inference benchmarking in the final version of the paper.

Q1. Improve Figure 1

Thank you for the suggestions! We have incorporated the changes to Figure 1 and included it in the additional PDF.

Q2. No pruning threshold for SNR

That is correct, and it is an oversight on our part; the pruning threshold is set to 1.0 for SNR in order to compare with previous work. This is mentioned at line 275 as the threshold is only used for continuous pruning, but we can additionally mention this at line 243.

References

[1] K. Neklyudov, D. Molchanov, A. Ashukha, and D. P. Vetrov. Structured bayesian pruning via log-normal multiplicative noise. In Advances in Neural Information Processing Systems (NeurIPS), pages 6775–6784, 2017.

We thank the reviewer for their time and engagement with the paper. We are happy that they found the approach interesting and well motivated, the writing clear, the math sound, the connection to floating point format interesting, and that we derive a variety of pruning algorithms. We address their weaknesses below:

W1. Figure clarity

We will revise the figures to be more clear, including the suggestion for Figure 1 which has now been updated in the rebuttal PDF.

W2. No ImageNet

The ResNet and ViT that we use in the paper are pre-trained on ImageNet, and then further fine-tuned on Tiny-ImageNet, which consists of 100,000 ImageNet images from 200 classes. While not as large scale as ImageNet, we do not expect to see radically different compression vs. accuracy if we perform further fine-tuning on the full ImageNet dataset. That being said, we are happy to perform experiments and include results on ImageNet in the final version of the paper, which could not be completed within the rebuttal period due to the size of ImageNet.

W3. Compare pruning characteristics when applied to complex structures

This is indeed an interesting point, and well worth looking into. The scope of our work is focused on pruning criteria for Bayesian structured pruning, while this question is slightly broader (relating to the Bayesian pruning literature in general). We can point this out in the paper as a useful future direction to explore.

W4. How many filters are required to learn a given dataset?

This is also an interesting research question! Figures 2, 5, and 7 (the post-training pruning experiments) show us this empirically. The knee-points of the SNR curve show us where accuracy begins to drop for higher compression. We see here that BMRS gives us a lower bound on the number of prunable structures in a network i.e. slightly higher compression can be achieved without significant drops in accuracy, but the accuracy will quickly decrease with further compression.

It should also be noted that different sparsity inducing priors will lead to different answers to this question. For example, [1] introduce a hierarchical prior design to induce both unstructured and structured sparsity, which can lead to a more sparse solution (and thus fewer parameters required to learn a given dataset), but is more complex to train and BMR cannot be directly used for group sparsity.

W5. More Baselines

We agree on the importance of relevant baselines. In our case, the baselines are different pruning criteria for multiplicative noise with the log Uniform prior, so we include L2 norm, SNR, and a no-pruning baseline.

Additional baselines that are comparable are:

• Expected value, $E[\theta]$
• Magnitude of gradient
• Hessian of gradients
• Magnitude of activation

Given the time constraints we have not been able to finish running all of these baselines, but we have submitted updated versions of Figure 2 and Table 1 with the finished results for $E[\theta]$ in the supplemental material (the rest will be included in the final version).

Q1. Advantage of variational pruning vs. classical methods

This depends on what the reviewer mean by "computational advantage"; the variational approach can lead to higher sparsification than classic approaches, which provides a net gain at inference time, but can be more costly to train due to the overhead of the variational parameters.

Q2. How many samples required for $\Delta F$?

One of the benefits of our proposed method is that the calculation of $\Delta F$ does not require any sampling; it is calculated analytically using the statistics of the variational distribution, original prior, and reduced prior, making it quite efficient (see Eq. 8 and Eq. 11).

Q3. Are there other choices of prior considered?

This depends on if the reviewer means the prior for the original model, or the reduced priors used to calculate $\Delta F$. The former is selected based on previous work that we build upon [2]. The latter are selected in order to be able to calculate $\Delta F$ analytically.

References

[1] D. Markovic, K. J. Friston, and S. J. Kiebel. Bayesian sparsification for deep neural networks with Bayesian model reduction. arXiv:2309.12095, 2023.

[2] K. Neklyudov, D. Molchanov, A. Ashukha, and D. P. Vetrov. Structured bayesian pruning via log-normal multiplicative noise. In Advances in Neural Information Processing Systems (NeurIPS), pages 6775–6784, 2017.

We thank the reviewer for their time reviewing the paper. We would like to address the points they make in their review:

Novelty

We respectfully disagree with the characterization of our method as a naive extension of variational dropout. We derive novel pruning criteria for a class of structured pruning models which are theoretically grounded, namely from the perspective of Bayesian model reduction (BMR). This yields a robust approach to structured pruning that performs well across our experiments/datasets without tuning of any threshold parameters. The different priors used in the paper are to attain different pruning criteria, as different realizations of BMR result from different priors; they are not different priors on the model parameters. In recent work, discussed in our paper, BMR has been successfully applied in the unstructured pruning case [1,2] but not in the structured case as far as we know.

Performance

Developing BMR for structured pruning is the main contribution of our work. Specifically, to 1) derive theoretically grounded pruning criteria for structured pruning and 2) empirically characterize the pruning characteristics of these criteria with respect to existing criteria. Our work serves as a step towards principled Bayesian structured pruning criteria via BMR. And when comparing with methods that require tuning of importance thresholds (L2, SNR) we show that we achieve reasonable trade-offs between compression and performance. It is not expected that we would attain higher compression than, e.g., Bayesian approaches which induce both unstructured and structured sparsity.

Baselines and Datasets

We would like to point out that we use MNIST, Fashion-MNIST, CIFAR10, and tiny-imagenet for our experiments (not only MNIST and CIFAR10 as the reviewer mentions).

Second, we compared to relevant baseline pruning criteria that have been used in previous work for the model class of Bayesian structured pruning that we study. Indeed in the paper referenced by the reviewer [3], their proposed pruning criteria are based on variations of $E_{q_{\phi}}[\theta]$ and SNR$(\theta)$; we compare to SNR$(\theta)$ as a baseline, and we have added $E_{q_{\phi}}[\theta]$ in the rebuttal PDF document, which will also be updated in the final version of the paper.

Missing citations

We thank the reviewer for pointing out the very recent work [3], which we will cite in our paper; the two other papers on variational pruning which the reviewer mentions are already discussed in our submission.

References

[1] J. Beckers, B. Van Erp, Z. Zhao, K. Kondrashov, and B. De Vries. Principled pruning of bayesian neural networks through variational free energy minimization. IEEE Open Journal of Signal Processing, 2024.

[2] D. Markovic, K. J. Friston, and S. J. Kiebel. Bayesian sparsification for deep neural networks with Bayesian model reduction. arXiv:2309.12095, 2023.

[3] J. Li, Z. Miao, Q. Qiu, and R. Zhang. Training bayesian neural networks with sparse subspace variational inference. CoRR, abs/2402.11025, 2024