Shaving Weights with Occam's Razor: Bayesian Sparsification for Neural Networks Using the Marginal Likelihood

Summary

This paper proposes to train neural networks that are more amenable to pruning, using a Bayesian approach. In particular, they specify separate priors over individual parameters rather than a shared global prior, which allows some parameters to be regularised more than others. They maximise a Laplace approximation to the marginal likelihood to optimise these priors, deriving an approximation that allows them to utilise richer block-diagonal KFAC precision matrices rather than simple diagonal ones. Furthermore, they propose a scoring criterion based on the posterior variance and magnitude of the weights for deciding which weights to prune. The authors found that their method for pretraining sparsifiable networks generally led to much better performance after pruning than traditional MAP training.

Strengths

- The proposed "SpaM" algorithm for pretraining sparsifiable networks in conjunction with their "OPD" criterion appears to be an effective method for pruning a variety of network types, based on the empirical results in the paper. - The experiments are extensive and relevant to the contributions of the paper - The approximation that allows KFAC to be used with non-scalar priors is useful to the KFAC literature in general, not just in the scope of network pruning. - The paper is well written and concepts are introduced in an intuitive fashion.

Weaknesses

- The more complicated KFAC with non-scalar prior approximation does not appear to be empirically better than using a diagonal precision matrix in the Laplace approximation, where the use of non-scalar priors is much simpler. However, it is still an interesting and potentially useful approximation outside of this particular application. ### Nitpicks - Some details may be unclear to people less familiar with the area, such as the "interleaved" training of the prior hyperparameters and network parameters, or what exactly is being optimised in the MAP comparisons. - Some of the more interesting figures (in my opinion) are relegated to the appendix, such as the fine-tuning figure, whilst the main paper has elements that feel superfluous or irrelevant, such as Figure 1, which seems unnecessary to explain the method, or a rather lengthy discussion of the use of marginal-likelihood in e.g. section 3.2.

Questions

### Questions - Do the authors believe that the disappointing performance of the KFAC with non-scalar priors was due to the extra approximation that was necessarily made (Proposition 3.1)? - In the MAP experiments, are we using the same Laplace approximation that was used to approximate the marginal log-likelihood, and if so, how are the prior hyperparameters chosen? Please clarify if I appear to be misunderstanding this aspect. - In lines 165-166, could you expand upon what is meant by "optimize the Laplace approximation to the marginal likelihood after an initial burn-in phase with a certain frequency"? - Does SpaM allow for improving structured sparsifiability **during training** as well as just unstructured? This wasn't entirely clear to me ### Suggestions - Expand slightly more on the difference between uniform and global pruning. - Remove some of the less relevant figures / discussion and include more of the experimental plots in the main text. - Typo in Eq. 5. $\mathbf Q^T$ should be $\mathbf Q^T_A$? - In Eq. 5, it would be clear to write $(\boldsymbol \Lambda_A \otimes \boldsymbol \Lambda_G + \delta \mathbf I)$ rather than $(\boldsymbol \Lambda_A \otimes \boldsymbol \Lambda_G + \delta)$ as the latter implies elementwise addition of $\delta$ rather than just the diagonal. Please let me know if this is a misunderstanding on my part. - Many of the figure references in the paper didn't seem to lead to the right place, e.g. on line 239. This might just be a problem with my machine, please double check.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors briefly discussed some generic limitations of Laplace approximations and additional computational cost of their SpaM training procedure over MAP. It would have been nice to comment on, for example, the remaining difficulty of structured pruning, since it is so much more beneficial than unstructured pruning.

Summary

This paper proposes to sparsify a neural network using Bayesian principles by optimizing the marginal likelihood (SpaM). Specifically, the authors use weight/node/layer-wise Gaussian priors over the weights and learn the corresponding scales by maximizing the marginal likelihood during training using Laplace approximation. Compared with MAP, where a shared regularization is used, SpaM optimizes the prior scale for each weight/node/layer to regularize weights adaptively. Moreover, a new important score, Optimal Brain Damage (OBD), is proposed for pruning weights, by using the approximated posterior. The effectiveness of this method is demonstrated with extensive experiments on various datasets and model architectures.

Strengths

The paper is novel and well-written. Although sparsity-inducing priors have been widely used to sparsify deep neural networks, optimizing the hyper-parameters in the prior using ML-II with Laplace approximation is novel and seems to be promising compared with using a single fixed scale to shave weights. The experiments are comprehensive, including results with different prior structures, pruning criteria, architectures (conv and transformer), and dataset domains (images and texts).

Weaknesses

This paper is generally well written with extensive experiment, and I only have one following concern (I'm very happy to increase my score if this is addressed): Hyper-parameters in the prior often can be tackled in two ways as a Bayesian: 1. optimize the hyper-parameters with ML-II; 2. conduct Bayesian inference on the hyper-parameters by having hyper-priors on them. This paper focuses on the first approach, but the comparison with the second approach is missing. In fact, the second approach has been widely used [1-4] to sparsify deep neural nets; for example, the node/layer-wise horseshoe prior and spike-and-slab prior can also offer different regularization strengths on different weights as well as give a more structured sparsity with good theoretical guarantees. Moreover, ML-II is known to have the risk of overfitting compared with the full Bayesian approach. So, I believe it is important to compare ML-II with the full Bayesian inference. [1] Ghosh, Soumya, Jiayu Yao, and Finale Doshi-Velez. "Model selection in Bayesian neural networks via horseshoe priors." Journal of Machine Learning Research 20.182 (2019): 1-46. [2] Louizos, Christos, Karen Ullrich, and Max Welling. "Bayesian compression for deep learning." Advances in neural information processing systems 30 (2017). [3] Cui, Tianyu, et al. "Informative Bayesian neural network priors for weak signals." Bayesian Analysis 17.4 (2022): 1121-1151. [4] Polson, Nicholas G., and Veronika Ročková. "Posterior concentration for sparse deep learning." Advances in Neural Information Processing Systems 31 (2018).

Questions

Is optimizing scales with ML-II using Laplace approximation better than doing a full (approximate) Bayesian inference over scales with common sparsity-inducing priors (e.g., horseshoe) using VI or SGHMC, in terms of computation, accuracy/uncertainty given a sparsity level, etc.?

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors mentioned the limitation coming from the Laplace approximation.

Summary

The paper introduces a new pruning technique named SpaM, which leverages Laplace approximation and Bayesian marginal likelihood for approximating the posterior in Bayesian Neural Networks. This approach includes two main components: the Gaussian prior variance and OPD, a pruning method that utilizes the Laplace posterior precision.

Strengths

Strengths: - Clarity and Accessibility: The paper is well-written, making it accessible to readers. - Motivation and Methodology: The use of marginal likelihood for pruning neural network models is well-motivated. The proposed method is applicable to both unstructured and structured pruning scenarios. - Bayesian Justification: The method is grounded in Bayesian theory, presenting a new pruning criterion based on posterior approximation. - Simplicity: The approach is relatively simple compared to existing Bayesian pruning methods.

Weaknesses

Weaknesses: Methodological Clarity: - The computational and storage challenges associated with the Fisher matrix or generalized Gaussian Newton matrix in the context of Laplace approximation for Bayesian Neural Networks are not sufficiently addressed in the limitations section. - The integration of the method with existing pruning criteria is unclear. The process involving MAP solution, Laplace approximation, and OPD with the precision matrix needs clarification, especially in comparison with zero-shot pruning methods like GraSP, SNIP, and Random. - The fairness of comparing the proposed method with Monte Carlo sampling for BMA performance computation against MAP is questionable. It should be clarified if only a single pruned model or Laplace approximation is used. Experimental Design: - The computational and memory costs of Laplace approximation should be compared with other baselines. - The baseline methods used are too simplistic. Stronger baselines such as IMP, RigL, or DST should be included. - The experiments are limited to CIFAR10. Additional datasets like CIFAR100 and ImageNet would strengthen the findings. - The datasets and networks used are relatively simple, raising questions about the method's performance on more complex datasets and models. - Specific questions about the online pruning process need addressing, such as the timing of pruning, the validity of Laplace approximation during training, and optimization of prior parameters. Novelty and Comparison: - The novelty is somewhat limited, as the method relies on known ideas from the literature. - The differences in uncertainty estimation between SpaM and MAP are marginal. The presentation of results could be improved to better highlight these differences. - Previous studies using Bayesian inference or variational methods for pruning should be compared with SpaM in both related work and experiment sections to enhance understanding. The scale of Experiments: - The experiments are considered limited in scale, with MNIST and CIFAR being seen as insufficient to provide meaningful insights. Larger models and datasets should be used for validation. Performance Goal: - The ultimate objective in neural network pruning is to achieve performance comparable to the unpruned dense model at a given sparsity level. The paper should position SpaM+OPD in relation to other pruning methodologies more clearly.

Questions

Questions: - Previous Studies Comparison: A comprehensive comparison between SpaM and previous studies using Bayesian inference or variational methods for pruning would be beneficial. - Clarification of Procedures: The paper should explicitly explain the sequence of dense training followed by prune-after-training, as this is critical for reader understanding.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have sufficiently addressed the potential negative societal impacts of their work. For details on the limitations, refer to the Weaknesses and Questions sections.

Summary

The authors present a framework for assessing the sparsifiability of a parametric model, a measure of how many parameters can be pruned without severely affecting the modelling performance. In essence, the authors suggest training a Bayesian neural network (BNN) using the marginal likelihood estimated via the Laplace approximation. The marginal likelihood's automatic Occam's Razor ability to identify good trade-offs between model complexity and data fit will encourage sparsity, and the trained network can then be subjected to a pruning criterion of choice. While any pruning criterion can be used, the authors propose the Optimal Posterior Damage (OPD), which is computed as a cheap byproduct of their marginal likelihood approximation. This often outperforms more expensive approaches. The authors demonstrate the effectiveness of their proposed approach through several experiments covering performance at different sparsity levels, online pruning, uncertainty estimation, the influence of prior choice, and structured sparsification. Generally, their proposed approach and pruning criterion outperforms baselines.

Strengths

**Originality** 1. Using the marginal likelihood to encourage less complex networks appears original in the context of pruning, although a little incremental. 2. The proposed pruning criterion, Optimal Posterior Damage (OPD), and the structured priors for the KFAC approximation appear to be original contributions. **Quality** 1. Marginal likelihood optimisation is a well-known technique which has many theoretically and intuitively pleasing properties. 2. The OPD pruning criterion makes intuitive sense and seems remarkably powerful. 3. The experimental evaluation is quite extensive and includes repetitions over multiple seeds (and the resulting uncertainties over the reported metrics). I find the evaluation of different choices for the structure of the prior is particularly interesting, and the recommendation for a default configuration is great. **Clarity** 1. The paper is very well-written, and the plots are generally informative and easy to understand. **Significance** 1. The marginal likelihood idea is simple yet works well in practice, which is great for the potential significance of the approach. 2. The proposed pruning criterion seems bound to become a strong baseline whenever the training approach is used. 3. The structured priors for the KFAC approximation are a minor contribution, but will likely be used outside of the pruning literature too, thus making them quite significant.

Weaknesses

**Originality** 1. While using the marginal likelihood as the training objective in the context of network pruning seems original, it seems somewhat incremental. **Quality** 1. The paper proposes a combination of two methods, training a Bayesian network using (an approximation to) the marginal likelihood and a pruning criterion, but only the choice of pruning criterion is evaluated experimentally, not the training scheme. Is training using the marginal likelihood better than a simple maximum likelihood optimisation with, say, L1 regularisation? The authors don't answer this important question. 2. While the authors test the effect of different prior structures (which are great and interesting experiments!), they do not test the effect of different prior distribution families. They seem to use a normal distribution for all experiments, which is surprising since the normal distribution doesn't induce sparsity and has been shown experimentally to be quite a poor choice for BNNs. Some obvious choices would be to test the horseshoe prior or the Indian buffet process prior, which can both encourage sparsity (e.g., Ghosh et al., 2018; Kessler et al., 2021). Other interesting choices would be the Laplace distribution, for which the MAP solution with a diagonal prior would correspond to L1 regularisation, Student's t and the spike-and-slab prior. Taken together with the remark above, it seems like the paper is missing half of the experimental analysis. 3. Experimentally, I would have liked to see results for the different pruning criteria in the settings that were used in their original publications (e.g., same training scheme). It is unclear if the paper's results are on par with the literature's. 4. Minor weakness: it would have been helpful to see the performance of the unpruned networks (i.e., 0% sparsity) to understand the penalty caused by the sparsification fully. Comment: if the authors decide to include more experiments in a future revision, I think sections 1 through 3 could be significantly shortened without losing too much context. **Clarity** 1. While the paper is overall very well-written, the terms in Eq. (4) lack definitions, and the appendix appears a little rough. 2. The font size in the plots could be a little larger. **References** - Ghosh et al. (2018). "Structured variational learning of Bayesian neural networks with horseshoe priors." ICML. - Kessler et al. (2021). "Hierarchical Indian buffet neural networks for Bayesian continual learning." UAI.

Questions

1. Additional experiments are probably infeasible for the rebuttal period, but did you already experiment with other prior distribution families? Why did you choose a normal distribution rather than a sparsity-inducing family? 2. The results for the online pruning approach in figure B6 confuse me a little. The online approach seems to generally perform better for higher sparsity levels for CIFAR-10, with LeNet on CIFAR-10 being particularly extreme (test accuracy around 50% for 20% sparsity compared to 65% for 80% sparsity ). Do you know what happens here?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations of their approach.