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

We thank you very much for your detailed and helpful feedback. We really appreciate your suggestions and will incorporate them in the revision.

Weaknesses:

The more complicated KFAC with non-scalar prior approximation does not appear to be empirically [...] However, it is still an interesting and potentially useful approximation outside of this particular application.

Thank you for your comment and raising this point. You understood our motivation with this proposal. We would like to expand further by adressing your question below.

Nitpicks

Some details [...] "interleaved" training of the prior hyperparameters and network parameters, or what exactly is being optimised in the MAP comparisons.

The optimization is conducted after a burn-in phase (number of initial epochs) and at a specific frequency (after each x epochs) referred to as marglik_frequency in the hyperparameter sections. In is worth mentioning that after pruning (at test time), we just use the trained posterior mean (MAP solution under the optimized prior) and not the full Laplace predictive.

Some of the more interesting figures (in my opinion) are relegated to the appendix, [...] 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.

We thank you for the suggestions, it was challenging to fit all key results in the main body of the paper. We will utilize the space better in the revised version.

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)?

Proposition 3.1 came into place to bridge the gap and to emphasize the contribution of diagonal priors in making neural networks more prunable, as it was shown to perform exceptionally well when combining the diagonal approximation with a diagonal prior. It was shown that a diagonal prior in KFAC pushes the probability of the network compared to other priors like scalar and layerwise (the approximation results are referred to as parameter-wise in red in Figure B4, compared to the KFAC with solid lines). For our work, we do recommend the usage of the diagonal approximation with diagonal priors as we do not need the precision of KFAC compared to diagonal in offering a better approximation. However, other works, as you mentioned, might benefit from it all while pushing the probability of networks.

In the MAP experiments, are we using the same Laplace approximation that [...] how are the prior hyperparameters chosen? Please clarify if I appear to be misunderstanding this aspect.

For MAP, we use a standard training approach where we do not optimize the marginal log-likelihood, and thus generally do not use Laplace. In the case of using OPD with MAP, the inverse of the Hessian is computed to be able to perform Laplace post-hoc on a pre-trained network without the need for the network to be trained and optimized for the marginal likelihood.

[...] "optimize the Laplace approximation to the marginal likelihood after an initial burn-in phase with a certain frequency"?

Thank you for your insightful question. We will include an explanation in the revision or refer to appendix D.4. The approximation is conducted during training after an initial number of epochs (burn-in phase) and at a specific frequency (e.g., every 5 epochs) for computational considerations. For smaller architectures, the approximation can be used from the start of training and at each epoch, but for larger architectures, this would be challenging and require more ressources. The table D.4 shows the choice of the parameters for each experiment and explains their significance.

Does SpaM allow for improving structured sparsifiability during training as well as just unstructured? This wasn't entirely clear to me.

SpaM can be applied during training for the structured case, but it is more efficient to mask structures (zero them out) instead of compressing them (removing them) during training. Compression during training is inefficient due to the need to replace layers and copy weights. By masking structures, we can evaluate the smaller network before fully committing to removing any components, and this does not add much overhead, as we have everything pre-computed with each marginal likelihood update and only need to aggregate the score. After training, compression can be performed once.

Expand slightly more on the difference between uniform and global pruning.

Thanks, we will expand on this in the revised version.

Global pruning prunes across the entire network at a target percentage (e.g., 80%). This approach may result in some layers being pruned more heavily than others to achieve the overall target sparsity.

Uniform pruning, on the other hand, applies the same pruning percentage to each layer.

Remove some of the less relevant figures/discussion and include more of the experimental plots in the main text.

Thank you, we will adress your suggestion in the revised version.

Typo in Eq. 5. should be $\mathbf{Q}^T_A$?

Thank you for pointing this out, that's correct and we will add the subscript $A$.

In in Eq. 5, it would be clear to write [...] Please let me know if this is a misunderstanding on my part.

Thank you for the suggestion. We will add the identity for clarity.

Many of the figure references in the paper didn't seem to lead to the right place, [...] please double check.

Thank you. Indeed, some links to the appendix figures seem to be broken, where the reference is rendered correctly, but the hyperlinks seem to be broken due to a counter change. We solved this issue by setting hypertextnames to false in the hyperref package.

Thank you for your evaluation of our work and your constructive comment that shows a good understanding of the field and your readiness to increase the score if your concerns and questions are adressed. We would like to further explain our motivation behind using Laplace approximations.

Weaknesses:

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. [...] 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.

Thank you for your comment. The main reason not to go for the full Bayesian inference approaches is their known computational complexity and lack of scalability for large-scale problems [1] which is emphasized in the suggested work (Ghosh et al., 2019; Louizos et al., 2017; Cui et al., 2022; Polson & Ročková, 2018). With a full Bayesian approach, it is not possible to provide a solution that can be scaled across the different architectures considered in our work (e.g., ViT and GPT-2). Utilizing ASDL [2] to compute the second-order information needed for the Laplace approximation, we make use of modern scalable algebra engines and avoid the repeated sampling and complex variational optimization that tends to slow down training and reduce performance.

This makes our method easier to use for practitioners, since it scales better with large datasets and complex models and supports different types of layers. For instance, it made it possible for us to use our method starting from a "pre-trained" backbone (e.g., from huggingface, without adjustments to the architecture) and transfer learn or fine-tune on another task.

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.?

Yes, full (approximate) Bayesian inference methods are likely intractable for large-scale problems due to their computational demands and the complexity of approximating or sampling from high-dimensional posteriors, as laid out in our response above.

In our case, we wanted the method to be able to scale and not be restricted by the limitations that come with full Bayesian inference that would require downscaling the experiments and not being able to perform our approach on modern architectures like vision and language transformers. BNNs are usually challenging to implement and expensive to train. With the Laplace approximation, it is much easier to use BNNs [3] and, during inference, to use a single forward pass on the pruned architecture. This approach allowed us to even be able to perform OPD on pre-trained models (by computing the inverse Hessian needed for Laplace post-hoc) that are very expensive to train, like GPT-2, and as well to use SpaM on architectures like DistilBERT and ViT.

References

[1] Bai, J., Song, Q., & Cheng, G. (2020). Efficient variational inference for sparse deep learning with theoretical guarantee. Advances in Neural Information Processing Systems, 33, 466-476.

[2] Osawa, K., Ishikawa, S., Yokota, R., Li, S., & Hoefler, T. (2023). ASDL: A Unified Interface for Gradient Preconditioning in PyTorch.

[3] Daxberger, E., Kristiadi, A., Immer, A., Eschenhagen, R., Bauer, M., & Hennig, P. (2021). Laplace Redux – Effortless Bayesian Deep Learning. NeurIPS.

Methodological Clarity:

Computation/Storage challenges

We discuss computational and storage costs of SpaM in depth in the main text but also in appendices D.4 & D.7. A diagonal approximation costs as much to store as a parameter vector and KFAC costs roughly twice as much so we incur minimal storage overhead. Computationally, using the EF is as cheap as gradient computation while the GGN scales with the number of outputs. For the benefit provided, these are minor additional costs.

Integration of the method is unclear.

There seems to be a misunderstanding as we have already covered these details. In appendix B.1, we explain that we initially train the NN either with MAP (a) or SpaM (b). Once these networks are trained, we prune both models using the different pruning criteria at different sparsity levels. Thus, we can measure the effect of SpaM training and OPD separately and fairly.

Fairness of comparison [...].

We do not use the predictive posterior for inference or Bayesian model averaging (BMA) but simply use the point estimate network. During inference, the model trained with SpaM only uses a single forward pass over the pruned architecture, thus guaranteeing a fair comparison with the baselines. The posterior is solely used to estimate the marginal likelihood during SpaM training.

Experimental Design:

Computational and memory costs of Laplace approximation.

We compared the cost of Laplace approximation used in SpaM (both using GGN and EF) with MAP and, as theoretically validated, it only adds minor overhead, especially in the diagonal EF variant.

Baseline methods are too simplistic.

We appreciate your feedback. To address this point and avoid any misunderstanding, we have addressed your comment in a detailed explanation in our general response.

Experiments are limited to CIFAR10. Additional datasets like CIFAR100 and ImageNet [...].

The experiments are not limited to CIFAR10. We have included CIFAR100 and IMDB as larger and different datasets, respectively.

Datasets and networks used are relatively simple.

We indeed cover in our paper a large variety of datasets and models, including MLP, CNN, MLP-Mixers, Transformers (Vision Transformers (ViT), Language Transformer (DistilBERT)) and show that our method can be applied to pre-trained models like GPT-2. The list of models and datasets are represented in the appendices D.1 and D.2, respectively.

Specific questions about the online pruning process.[...].

Appendix B.6 addresses these questions and explains how the pruning is conducted during training. Pruning occurs incrementally after a new computation of the marginal likelihood.

Novelty and Comparison:

The novelty is somewhat limited, [...].

We address the novelty of our approach with the general response.

Differences in uncertainty [...] are marginal. The presentation of results could be improved [...].

We acknowledge that the difference in uncertainty estimation between SpaM and MAP could be more clearly visualized and will adjust the y-axis scale accordingly (due to random pruning having high values). The difference is indeed significant at high sparsities for criteria like OPD, GraSP, and magnitude (see Figure B8 and Table B3).

Comparison with previous studies using Bayesian inference or variational methods for pruning..

Variational inference for deep neural networks usually slows down training and reduces performance in comparison to the MAP. Further, as found by Blundell et al. (2015) in their seminal Bayes-by-Backprop paper, it does not work in conjunction with prior hyperparameter updates and thus cannot be applied in the automatic regularization setting we are interested in. We will include a more thorough discussion in the related work section.

The scale of Experiments:

The experiments are considered limited in scale,[...]. Larger models and datasets should be used for validation.

We cover a large variety of models and datasets (vision and text):

Models: fully connected, LeNets, MLPMixer, (Wide) ResNets, ViTs, DistilBERT, GPT-2

Datasets: UCI breast cancer, MNIST, FashionMNIST, CIFAR10, CIFAR100, IMDB

[...]. The paper should position SpaM+OPD in relation to other pruning methodologies more clearly.

Based on Figure 2 which compares SpaM and MAP across different criteria, SpaM creates models that are inherently more sparsifiable. OPD (blue lines) consistently ranks among the top-performing criteria, maintaining near-unpruned performance even at 95% sparsity in vision tasks. GPT-2 results also highlight OPD's standalone effectiveness without SpaM.

Taken together, SpaM and OPD offer a powerful combination for achieving the ultimate pruning goal: high performance at high sparsity.

Questions

Previous Studies Comparison: [...] SpaM and previous studies using Bayesian inference or variational methods for pruning [...].

A similar comparison would down-scale our evaluation pipeline, as most VI methods and those relying on full (approximate) Bayesian inference are limited and likely intractable for large-scale models. Such approaches are typically constrained to smaller architectures like MLPs and CNNs [1-3].

Methods like MCMC are impractical for large-scale models due to high computational demands and numerous samples needed to estimate posterior distributions accurately. The Laplace approximation used in SpaM avoids sampling, making it computationally feasible.

[1] Molchanov et al. (2017). Variational Dropout Sparsifies Deep Neural Networks. ICML.

[2] Zhao et al. (2019). Variational Convolutional Neural Network Pruning. CVPR.

[3] Bai et al. (2020). Efficient Variational Inference for Sparse Deep Learning with Theoretical Guarantee. NeurIPS.

[...] The paper should explicitly explain the sequence of dense training followed by prune-after-training, [...].

We provide a detailed description of the training approach in the main text and more in detail in appendices B.1 and D.

Weaknesses

While using the marginal likelihood as the training objective in the context of network pruning seems original, it seems somewhat incremental.

Thank you for the chance to expand more on the novelty offered with our approach and scalability compared to established works. In fact, our framework addresses existing scalability challenges, making it practical for real-world scenarios, and enables automatic relevance determination in deep learning for the first time.

Quality

The paper proposes a combination of two methods [...] Is training using the marginal [...] a simple maximum likelihood optimisation with, say, L1 regularisation? The authors don't answer this important question.

Thank you for your insightful question. We consistently compare MAP and marginal likelihood training throughout the paper (solid and dashed lines in the plots). MAP with Gaussian priors (L2 regularization) is our primary baseline, as it corresponds to the widely used standard weight decay. We ran an experiment with L1 for the rebuttal (see provided pdf), which gave worse results.

Exploring other prior distributions like the horseshoe prior or Laplace distribution, which have been shown to induce sparsity and improve performance in Bayesian neural networks [...] 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.

Thank you for the opportunity to clarify. We address your concerns about prior choice in the general response. Alternative prior families often require expensive inference, making them infeasible for large-scale models. The Laplace approximation, which we employ for scalability reasons, requires differentiability, excluding many of your suggested priors. We instead use the idea of relevance determination, that is regularization of individual weights or weight groups. With the Laplace approximation, this is scalable to networks at any scale, thanks to modern Hessian approximations.

Experimentally, I would have liked to see results for the different pruning criteria in the settings that were[...]. It is unclear if the paper's results are on par with the literature's.

Thank you for your valuable comments. While SNIP and GraSP are primarily used for pruning at initialization (PAI), their criteria based on weight contribution to loss and gradient preservation can be applied broadly for weight importance evaluation in various scheduling scenarios. However, using them as PAI often requires initial calibration passes for the network, especially at high sparsities, to achieve the desired pruning target. Additionally, compared to post-training methods, PAI can be less efficient, especially for large models that need retraining from scratch or rely on the availability of the original dataset (we have included PAI setup in our submitted code). In contrast, methods like SpaM can be more efficient for networks built on top of transfer learning, while OPD can be directly applied to pre-trained models by computing the inverse Hessian, as demonstrated in our GPT-2 experiments. Our main aim was to show the benefits of SpaM in a harmonized experimental setup, regardless of the specific pruning criterion.

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.

Thank you. We found that most methods, especially OPD, perform at almost identical performance to the baseline at 20% and hence did not report 0% sparsity so far. We will add a horizontal line with the baseline accuracy that will also make it easier to visualize the performance gap at higher sparsities.

Clarity

While the paper is overall very well-written, the terms in Eq. (4) lack definitions, and the appendix appears a little rough.

Thank you we will improve the definition in the revision and polish the appendix. $A_l$ and $G_l$ are the uncentered covariances of the layer inputs and output gradients, respectively. These are used as defined in the referenced works, for example [1].

[1] Martens, James, and Roger Grosse. "Optimizing Neural Networks with Kronecker-factored Approximate Curvature." 2020

The font size in the plots could be larger.

We will increase the font size to improve readability and accessibility.

Questions:

Additional experiments are probably infeasible for the rebuttal period, [...] experiment with other prior distribution families? Why did you choose a normal distribution rather than a sparsity-inducing family?

See our arguments regarding scalability above. Specifically, in the context of automatically learning regularization, the Gaussian group-wise prior, as in automatic relevance determination (ARD), has proven effective. Other priors, due to non-differentiability, could require intractable manual hyperparameter tuning.

The results for the online pruning approach in figure B6 confuse me a little. The online [...] for CIFAR-10, with LeNet on CIFAR-10 being particularly extreme [...]. Do you know what happens here?

The x-axis in the online pruning curves (section B.6) indirectly reflects training progress, not just sparsity. LeNet on CIFAR-10 is pruned starting from epoch 10 for 100 epochs, aiming for 99% sparsity. LeNet's curve trend on CIFAR-10 is therefore due to ongoing convergence during pruning, caused by dataset complexity relative to the architecture.