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

Authors: rayen  dhahri
alexander  immer
betrand  charpentier
stephan  günnemann
vincent  fortuin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:28:55

Abstract: Neural network sparsification is a promising avenue to save computational time and memory costs, especially in an age where many successful AI models are becoming too large to naively deploy on consumer hardware. While much work has focused on different weight pruning criteria, the overall sparsifiability of the network, i.e., its capacity to be pruned without quality loss, has often been overlooked. We present Sparsifiability via the Marginal likelihood (SpaM), a sparsification framework that highlights the effectiveness of using the Bayesian marginal likelihood in conjunction with sparsity-inducing priors for making neural networks more sparsifiable. Our approach implements an automatic Occam's razor that selects the most sparsifiable model that still explains the data well, both for structured and unstructured sparsification. In addition, we demonstrate that the pre-computed posterior precision from the Laplace approximation can be re-used to define a cheap pruning criterion, which outperforms many existing (more expensive) approaches. We demonstrate the effectiveness of our framework, especially at high sparsity levels, across a range of different neural network architectures and datasets.

Main paper: 
 introduction: 
 The availability of large datasets and powerful computing infrastructure has fueled the growth of deep learning, enabling the training of increasingly complex neural networks (NNs). While catalyzing performance gains across various domains, such as image recognition [1] and text generation [2] , this development has amplified the challenge of over-parameterization [3, 4] and raised concerns about the increase in model size and computational cost. Over-parameterized neural networks present significant deployment challenges, particularly in hardware-constrained environments [5, 6] . This has sparked the research field of NN sparsification or pruning, where the goal is to remove a (potentially large) number of parameters from a trained network to make it smaller and ultimately cheaper to apply [7, 8] . However, most existing research in this domain has focused on the question of finding better pruning criteria, that is, scoring functions that decide which parameters to prune away [9, 4, 10] . This ignores the challenge that many trained networks are not inherently sparsifiable, i.e., they resist Figure 1 : Overview of our proposed SpaM method. We start by training the network to maximize the marginal likelihood using the Laplace approximation, while simplifying the Hessian computation through either the KFAC or a diagonal approximation. We can then use our precomputed posterior precision as a pruning criterion (OPD). For the case of unstructured pruning, we compute thresholds to achieve different target sparsities, compute the mask, and apply it, while for the structured approach, we aggregate the score per layer for easier weight transfer, compute the mask, and then delete the masked structures to obtain a smaller model. effective pruning, regardless of the chosen criterion. Indeed, standard training of NNs does not encourage sparsifiability at all, so it should not be surprising that such trained NNs would use all of their parameters to some degree to fit the data.Our work tackles this problem by modifying the training process itself, showing that more sparsifiable networks can be achieved through Bayesian model selection using the marginal likelihood [11, 12] in conjunction with an adequate prior that will induce such sparsifiability. We call this Sparsifiability via the Marginal likelihood, or SpaM. Our approach implements an automatic Occam's razor [13] , guiding the training process towards models that are faithfully fitting the data using only a small subset of their available parameters, such that the remaining ones can be pruned afterwards. This is achieved by optimizing thousands of prior hyper-parameters to adaptively regularize weight magnitudes. We make use of recent advances in Laplace inference for Bayesian neural networks (BNNs) [12, 14] , allowing us to approximate the marginal likelihood [11] efficiently.Once trained, we can use any pruning criterion to more effectively sparsify these networks. Notably, the pre-computed posterior precision of the Laplace approximation obtained from the marginal likelihood training readily translates into a powerful pruning criterion, which we call Optimal Posterior Damage (OPD), similar to the popular Optimal Brain Damage [OBD ; 3] . Since it reuses existing computations, it is cheaper than many existing criteria in practice, and it often performs on par or even better.Extensive empirical evaluations demonstrate the strength of our SpaM approach and the derived OPD pruning criterion in both unstructured and structured sparsification tasks across various datasets and architectures. Moreover, they show that our framework strikes a compelling balance between performance and computational cost.We make the following contributions:• We propose SpaM, a novel approach to improve the sparsifiability of neural networks during training, using Bayesian model selection with the Laplace-approximated marginal likelihood, which works well in structured and unstructured pruning scenarios and with different pruning criteria.• We provide evidence-based recommendations for prior selection within the SpaM framework, showing that some priors can improve sparsifiability significantly better than others.• We present OPD, a cheap pruning criterion similar to the popular OBD, which performs comparably or better than many other criteria in practice and conveniently reuses computations performed in the context of SpaM. 
background: 
 We use deep neural networks to model learning tasks with inputs x n ∈ R D and targets y n ∈ R C collected in a dataset D = {(x n , y n )} N n=1 of N pairs. A model is parameterized by weights θ ∈ R P , and maps from inputs to targets using the neural network function f θ (x). Assuming the data are i.i.d., we have a likelihood p(D|θ) = N n=1 p(y n |f θ (x n )). We minimize the negative log likelihood, which corresponds to common losses like the cross-entropy in classification. Additionally, regularization in the form of weight decay is commonly used and corresponds to a Gaussian prior on parameters p(θ) = N (θ; 0, diag(δ)) with diagonal precision. We employ Gaussian priors due to their analytical tractability and seamless integration with the Laplace approximation, which requires differentiability. This is essential for maintaining computational efficiency in our framework to approximate the marginal likelihood in a practical and scalable manner. Furthermore, Gaussian priors enable automatic relevance determination by allowing each parameter to have its own variance, facilitating the regularization process without introducing significant computational overhead.The marginal likelihood serves as the probabilistic foundation for model evaluation and selection. It provides an objective to optimize the tradeoff between data fit and model complexity, akin to the concept of Occam's razor [13, 15] , by quantifying how well a model M, with all its inherent uncertainties, explains the observed data:EQUATIONHowever, it requires computing an intractable integral over all neural network parameters.The Laplace approximation [LA, 16] provides a tractable and effective approximation to the marginal likelihood for deep learning [12] . It arises from a second-order Taylor approximation around an estimate of the mode, θ * , resulting inEQUATIONwhere P θ * is the posterior precision given by the Hessian of the negative log joint distribution, -∇ 2 θ log p(D, θ|M), evaluated at θ * . Defining H θ * as the Hessian of the negative log likelihood objective -∇ 2 θ log p(D|θ, M), the posterior precision decomposes as P θ * = H θ * + diag(δ). In practice, the Hessian of the negative log likelihood is often approximated by the positive semidefinite generalized Gauss-Newton [GGN, 17] ,EQUATIONwhich relies on the Jacobians of the neural network function and second derivative of the negative log likelihood at the output. Further, it is amenable to efficient structured approximations like diagonal or layer-wise variants [e.g., 18, 19] .Diagonal and block-diagonal GGN approximations are efficient, and therefore commonly used for Laplace approximations in deep learning [20, 14] . The diagonal LA is cheap in terms of storage and computation by only modeling the marginal variances of parameters. Kronecker-factored LA [KFAC LA, 20] instead relies on a block-diagonal approximation to the GGN of the parameters θ l in the lth layer,EQUATIONwhere the factors are given by the outer products of pre-activations and Jacobians w.r.t. the output of a layer, respectively [18, 19] . Here, A l and G l are the uncentered covariances of the respective layer inputs and output gradients. The top left of Figure 1 shows a comparison of both structures.The goal of the pruning procedure is to remove parameters from θ without affecting the quality of the model output f θ (x). While unstructured pruning consists in zeroing individual entries θ p of the weight matrices, structured pruning consists in deleting entire structured sets of parameters g, like rows or columns [21, 22] . The results of structured pruning enable smaller matrix multiplications that directly provide real-world efficiency gains on most hardware, including GPUs.Pruning procedures usually follow three steps: (1) We use a scoring function S(•) to evaluate the importance of each individual parameter S(θ p ) for unstructured pruning, or of a structured set of parameters S(g) for structured pruning. (2) We compute a binary mask m with the same dimensions as θ, which assigns 0 values to parameters whose unstructured or structured pruning scores are below a threshold T , and 1 otherwise. While the threshold T is determined based on the target sparsity across layers for global pruning, it is determined per layer for uniform pruning [21] . (3) We apply the mask on the weight matrix with element-wise multiplication, m • θ, to effectively remove the least important parameters. Alternatively, structured pruning enables us to directly remove rows or columns whose mask values are 0 to reduce weight matrix dimensions. 
shaving weights with occam's razor: 
 We identify sparsifiable neural networks by automatically regularizing (groups of) parameters to have small magnitudes, to facilitate pruning the least important ones, within a probabilistic framework. Specifically, we utilize priors that regularize parameters in potentially structured ways, leading to smaller magnitudes. To optimize the resulting prior hyperparameters, we employ the Bayesian marginal likelihood as a differentiable objective function, effectively implementing a Bayesian variant of Occam's razor that drives irrelevant parameters towards smaller absolute magnitudes. The regularized networks can then be pruned with any method. However, we additionally propose to reuse the computed posterior precision for sparsification as a cheap and effective criterion.To reduce the magnitude of parameters and make them more amenable to pruning, we introduce structured priors and show how to combine them with diagonal and KFAC Laplace approximations. While a scalar prior, corresponding to weight decay, is the most common, it suggests that all parameters in a neural network are equally relevant and favors a uniform magnitude of parameters, which is suboptimal for pruning [23, Sec. 3.6] .Instead of scalar priors, we regularize parameters with different strengths using layer-, unit-, and parameter-wise priors. Layer-wise priors regularize individual layers differently and have been shown to aid pruning and improve generalization [12, 24, 25, 26] . Unit-wise regularization has been used mostly in traditional statistics, for example, for group sparsity [27] , but recently also for channels or feature dimensions in neural networks [28, 29] .We consider different priors in the context of the Laplace approximation for marginal likelihood optimization and pruning: Scalar priors correspond to standard weight decay and are identical for all weights. Layer-wise priors provide a scalar regularizer δ l per layer that is stacked into a vector δ in line with the number of parameters per layer. Parameter-wise priors allow to specify δ p for each parameter p individually. We define unit-wise priors so that each unit, which denotes a channel for convolutional and a hidden neuron for fully-connected layers, has a regularization strength for incoming and outgoing weights separately. Thus, a weight θ p that connects unit i at layer l-1 with unit j in layer l has prior N (0, [δ l-1 ] i • [δ l ] j ), that is, each layer l with M l hidden units has a prior vector δ l ∈ R M l . A weight is thus regularized more strongly whenever both its in-and output neurons are. Our different priors are simple to combine additively with a diagonal Hessian approximation for the Laplace approximation (Equation (2)) but not with a KFAC structure. For that reason, so far, only scalar or layer-wise priors have been used for KFAC posterior approximations [14] . The main issue is that we have to preserve the Kronecker factors to keep the resulting memory cost low. For scalar or layer-wise priors, this can be achieved by an eigendecomposition of the individual factorsEQUATIONwhich means that the precision only needs to be added to the diagonal eigenvalues and no Kronecker product needs to be calculated for inversion or determinant calculation.To add a diagonal prior precision δ l to the KFAC of the lth layer, we derive an optimal approximation in the KFAC eigenbasis, so as to maintain the Kronecker-factored structure of the posterior: Proposition 3.1 (Diagonal Prior in KFAC Eigenbasis). Considering the Frobenius norm, the optimal diagonal perturbation of the KFAC eigenvalues Λ A ⊗ Λ B to add a diagonal prior precision is given by ΛA ⊗ Λ B + δ with mat( δ) = (Q T G ) 2 mat(δ)Q 2 Awhere the square is element-wise and mat(•) reshapes the vector to match the parameter shape used in KFAC. Thus, it can be computed efficiently without computing a Kronecker product.We provide the proof in Appendix A. The approach is similar to that of George et al. [30] , who correct KFAC's eigenvalues towards the diagonal Gauss-Newton, but solves the problem of adding a full-rank diagonal instead of a rank-1 outer product to the KFAC eigenbasis.To optimize the potentially millions of regularization parameters, for example, arising from a parameter-wise prior, we employ the marginal likelihood as a differentiable objective. Optimizing regularization parameters has the advantage that different (groups of) parameters will be regularized differently and, therefore, become easier to prune. While it would be intractable to optimize that many regularization parameters using validation-based forms of optimization, the marginal likelihood can be estimated and differentiated during training [12, 31, 32] .Automatically determining the relevance of parameter-groups (ARD) is a common approach in Bayesian learning that can lead to sparsity and smaller parameter magnitudes [33, 34] and has been used especially in linear models. The marginal likelihood provides an objective that automatically regularizes irrelevant parameter groups more to lower their magnitude. Therefore, it implements a Bayesian variant of Occam's razor, finding the simplest model that explains the data well [13] .Mathematically, all the prior parameters δ constitute the hyperparameters of the model M in the log marginal likelihood (Equation ( 1)) that we optimize interleaved with the neural network parameters. When optimizing the prior parameters, we use gradient ascentEQUATIONor adaptive optimizers like Adam [35] . We follow Immer et al. [12] and optimize the Laplace approximation to the marginal likelihood after an initial burn-in phase with a certain frequency. We describe the optimization process and the related hyperparameters in Appendix D.4While sparsity regularization learned by marginal likelihood training can be advantageously combined with any pruning criterion, like Single-shot Network Pruning [SNIP; 36], variants of Gradient Signal Preservation [GraSP; 37, 38, 39] , or magnitude pruning [7] , we further propose a new pruning criterion that uses our Laplace approximation and extends the unstructured Optimal Brain Damage (OBD) pruning criterion [3] . While OBD traditionally uses the Hessian (approximation) H θ of the loss, we propose to adapt it to use the posterior precision P θ , which additionally includes the prior precision δ. The importance score S(θ p ) for parameter θ p becomesEQUATION)where P pp denotes the posterior precision for the parameter θ p , extracted from the diagonal of the posterior precision matrix P θ . We call this novel posterior-based pruning criterion Optimal Posterior Damage (OPD). Intuitively, individual weights with high scores indicate certainty of the posterior distribution and a significant contribution to the model's functionality, as indicated by the magnitude.We also propose a structured version of OPD by aggregating the score over a set of parameters g, i.e.,EQUATION)In practice, the structured set of parameters g corresponds to all parameters along one dimension of the weight matrix inside a layer, in order to reduce the size of the matrix multiplications. Since subsequent layers might have significantly different weight matrix dimensions impacting the magnitude of the aggregated sum, we opt for uniform structured pruning to guarantee a fair pruning treatment across all layers. This means we prune each layer by the same target percentage, reducing the dimensions of each layer by the same proportion relative to the unpruned model. This approach contrasts with achieving a global target sparsity that varies across layers, which can make it more difficult to consistently compress and adjust the model's size.Moreover, as removing a full structure is more aggressive, we also apply gradual pruning during training. Finally, we omit pruning the final layer to mitigate overly strong impact on classification accuracy and computational stability [40] .When using our SpaM approach, the precomputed precision matrix from the Laplace approximation can be reused to compute OPD without computational overhead, in contrast to the other pruning criteria, which often require additional computations to be performed. Note that we will also show in our experiments that OPD additionally avoids the need for potentially expensive fine-tuning after pruning. Moreover, even in the case of maximum a posteriori (MAP) training, Laplace approximations of the inverse Hessian at θ * can be additionally computed to approximate OPD. Finally, the OPD criterion can not only be computed post-hoc after training, but even online during training. 
related work: 
 Laplace-approximated BNNs. From the early inception of Bayesian neural networks [41, 42] , the Laplace approximation was a popular inference method [16] . In recent years, it has undergone a renaissance [18, 19, 20, 14] , including critical work on using more scalable approximations for the associated marginal likelihood in the context of model selection [11, 12, 43] , which we use in our framework. To the best of our knowledge, we are the first to study the benefits of this Laplaceapproximated marginal likelihood in the context of sparsification of deep neural networks. However, similar methods that automatically quantify the relevance (ARD) of parameters have been derived and used for linear, bilinear, and kernel models [34, 44, 45] as an alternative to the Lasso. More recently, van der Ouderaa et al. [46] and Bouchiat et al. [47] used the ARD mechanism in deep learning to select layers and features by regularization, respectively.Pruning neural networks. Various pruning criteria have been proposed to determine the importance of model parameters. Many criteria prune based on the weight magnitude [7, 48, 49] but usually required additional fine-tuning to recover accuracy. Sun et al. [8] proposed to combine activation and weight norms for pruning without fine-tuning. Other approaches include pruning using firstorder information based on connectivity [36] or synaptic flow conservation [50] , or second-order information aiming at preserving gradient flow [37, 38, 39] . Recently, van der Ouderaa et al. [51] focused on pruning LLMs based on a second-order Taylor expansion. In contrast, OPD uses secondorder information provided by the posterior precision given by the Laplace approximation. Beyond pruning criteria, there have been many approaches to prune at initialization [36, 37, 50] , during training [52, 53] , and after training [7, 8] . In particular, multiple works proposed to leverage specific Figure 2 : Predictive performance as a function of sparsity level in unstructured pruning. We see that SpaM improves the performance over MAP training across most architectures, datasets, and pruning criteria, and that OPD often outperforms the other pruning criteria. Both of these effects are particularly visible at higher sparsity levels. The black star in each subfigure denotes the performance of the unpruned models, which is often identical to the performance of models pruned at 20% sparsity.training schemes promoting zero-invariant parameter groups for structured pruning [54, 55] . In contrast, SpaM induces sparsifiability during training, and is agnostic about the criterion. 
experiments: 
 We conduct experiments on various datasets and models and outline our experimental setup in detail in Appendix D. We compare MAP training with our proposed SpaM approach with different priors, comparing our OPD pruning criterion with random pruning, magnitude pruning, SNIP [36] , GraSP [37, 38] , and SynFlow [50] . We show that SpaM improves pruning performance with different pruning criteria, especially at higher sparsities, and that our OPD criterion often outperforms the other criteria. This observation extends not only to predictive accuracy, but also uncertainty estimation. Moreover, we show that the choice of prior can play a significant role, and we introduce parameter-wise and unit-wise priors for the KFAC approximation. Finally, we show that SpaM and OPD also work in a structured pruning setting, leading to significant computational benefits. The code for our methods and experiments can be found at https: //github.com/fortuinlab/spam-pruning.We compare SpaM to MAP training with different pruning criteria, including OPD, across different models on tabular, vision, and language datasets. For SpaM in this unstructured pruning context, we use the diagonal Laplace approximation with a parameter-wise prior. Encouragingly, MAP and SpaM reach comparable performance during training, showing that the increased sparsifiability of SpaM comes at no additional cost in unpruned performance (see Figure B1 in the appendix).We see in Table 1 and Figure 2 that SpaM drastically improves the performance for many pruning criteria, especially magnitude pruning, GraSP, and OPD. We also see that OPD, despite being a cheap byproduct of our marginal likelihood computation, often outperforms the other pruning criteria, especially at higher sparsities. For instance, at 95 % pruning rate (i.e., with 20x fewer parameters), our combination of SpaM and OPD still retains almost the same performance as the unpruned model on vision tasks, while the other pruning criteria with MAP training have dropped to essentially unusable performance levels at this sparsity. Fine-tuning. We see in Figure B4 in the appendix that some of this performance difference can be remedied by costly fine-tuning of the networks after pruning, which however still does not allow the other methods to reach the full SpaM-OPD performance. Interestingly, in the case of OPD, this does not further improve its already near-optimal performance.Online pruning. Figure B5 in the appendix shows that our online version of SpaM, which uses the marginal likelihood and OPD during training to iteratively prune the network, reaches comparable performance levels to the post-hoc version, thus offering a computationally even more convenient way to effectively sparsify neural networks.Uncertainty estimation. Given that SpaM is a Bayesian method, it does not only offer high predictive accuracies but also calibrated uncertainty estimates. Indeed, we see in Figure 3 that the trends we have seen for accuracy also apply for negative log-likelihood, expected calibration error, and the Brier score. Again, SpaM improves the uncertainty estimates over MAP training, OPD outperforms most other criteria, and we achieve well-calibrated models up until very high sparsity levels. Note that the random baseline also achieves a low ECE at high sparsity levels because it essentially reverts to random guessing, which is a known weakness of the ECE metric [56] . Figure 5 : Similarly to unstructured pruning, we see in this experiment on structured pruning that SpaM (using a unit-wise prior) improves performance over MAP and that OPD mostly outperforms other pruning criteria, especially at higher sparsity levels. The black stars reflect the performance of the unpruned models. Figure 6 : Structured pruning with LeNet on FashionMNIST, using unit-wise priors. We see that our SpaM-OPD dominates the Pareto frontier, in terms of predictive performance as a function of computational time and memory cost, and is particularly competitive at lower costs.To understand the influence of the prior and Hessian approximation on performance in our proposed SpaM-OPD approach, we compare diagonal and KFAC approximations with scalar, layer-wise, unit-wise, and parameter-wise priors. Note regarding the latter two, that in this work, we are the first to implement them for the KFAC approximation, thus contributing to the general framework of Laplace-approximated BNNs [14] , independent of the pruning use case.We see in Figure 4 that our newly introduced unit-wise and parameter-wise priors for KFAC indeed outperform the others, especially at high sparsities. When comparing KFAC to the diagonal approximation, we see that KFAC often leads to slightly better performance at lower sparsity levels. However, we also see that the relatively simple choice of parameter-wise prior and diagonal Hessian approximation, as used in our previous experiments above, is a strong baseline across the board and can be recommended as a safe default option for unstructured pruning. Note that the unit-wise priors can be especially useful for structured pruning, as we will see in the following experiment. More detailed prior comparisons can be found in Appendix B.3.Here, we study the effect of SpaM and OPD in the more challenging setting of eliminating entire network structures, such as convolutional kernels. Studying different network architectures, we aim to generalize our unstructured pruning approach to the setting of structured pruning, where the structures can be freely defined depending on the use case.Encouragingly, we see in Figure 5 that our findings from the unstructured case transfer qualitatively also to the structured case, with SpaM-OPD outperforming the baselines at high sparsities. Crucially, while the sparsity patterns generated by unstructured pruning are more difficult to translate into computational benefits, structured pruning directly leads to computational savings on standard GPUs (see also Figure B11 in the appendix). We see in Figure 6 that SpaM-OPD dominates the Pareto frontier of the tradeoff between performance and computational cost at high sparsities (i.e., low costs), yielding 10x-20x savings in FLOPS and memory consumption with only minimal deterioration in performance. This positions our proposed framework as a potentially promising approach for the deployment of AI models in resource-constrained environments. 
conclusion: 
 We have shown that the Bayesian marginal likelihood, with its associated Occam's razor effect, can be used during training to select neural network models that are inherently more sparsifiable. Crucially, we have shown that this sparsifiability extends across different pruning criteria and enables large gains in performance and uncertainty estimation, especially at high sparsity levels. Conveniently, the computations needed for the marginal likelihood estimation using the Laplace approximation can be re-used to define a novel pruning criterion called OPD, which outperforms many existing (more expensive) criteria in our experiments. We have also presented guidelines for choosing priors within our framework and have shown that even in the challenging setting of structured pruning, our proposed SpaM approach can yield up to 20x savings in computational time and memory, with only small reductions in performance. Our work thus offers a promising path towards pruning large AI models at high sparsity levels for deployment on resource-constrained devices.Limitations. Our approach naturally inherits some limitations of the Laplace approximation, for instance, the fact that it only captures the local geometry of a single posterior mode or potential numerical instabilities in the Hessian computations when used with low-precision weights. Moreover, it accrues an additional computational cost compared to MAP training, which is then, however, amortized by the computational savings during the deployment of the sparsified model. 
 
 Appendix: Proposition A.1 (Diagonal Prior in KFAC Eigenbasis). Considering the Frobenius norm, the optimal diagonal perturbation of the KFAC eigenvalues Λ A ⊗ Λ B to add a diagonal prior precision is given by ΛA ⊗ Λ B + δ with mat( δ) = (Q T G ) 2 mat(δ)Q 2 Awhere the square is element-wise and mat(•) reshapes the vector to match the parameter shape used in KFAC. Thus, it can be computed efficiently without computing a Kronecker product.Proof. We prove this result in two steps. First, we show what the optimum looks like in terms of the Frobenius norm. Second, we show how to simplify the results to enable efficient computation without computing Kronecker products. We have a KFAC Hessian approximation A ⊗ B with A ∈ R Din×Din and B ∈ R Dout×Dout where the dimensionalities D • depend on the layer type [18] . In the case of a fully-connected layer, these are simply the dimensionality of the in-and output hidden representation. The same layer will have D in × D out parameters and thus the corresponding diagonal prior precision is given by δ ∈ R DinDout . For the Laplace approximation, the eigendecomposition of individual Kronecker factors is already computed asA = Q A Λ A Q TA and similarly for G as shown in Equation (5) . Recall also that diag(•) turns a vector into a diagonal matrix and extracts the diagonal entries of a matrix into a vector. We are interested in the Frobenius-optimal diagonal perturbation of the eigenvalues so as to maintain the efficiency structure of the KFAC, and, thus, the downstream Laplace approximation:arg min δ ∥(Q A ⊗ Q G )(Λ A ⊗ Λ G + diag ( δ))(Q T A ⊗ Q T G ) -(Q A ⊗ Q G )(Λ A ⊗ Λ G )(Q T A ⊗ Q T G ) + diag (δ)∥ 2 F = arg min δ ∥Λ A ⊗ Λ G + diag ( δ) -Λ A ⊗ Λ G + (Q T A ⊗ Q T G ) diag (δ)(Q A ⊗ Q G )∥ 2 F = diag ((Q T A ⊗ Q T G ) diag (δ)(Q A ⊗ Q G )), where we first multiplied the orthogonal bases from left and right and then realized that the values of δ need to be set to the entries of the prior δ projected into the basis.Naïvely, computing the optimum of δ would require expanding the Kronecker product above and lead to a potentially intractable complexity of O(D 2 in D 2 out ). However, it is possible to simplify it further to maintain efficient computation: For simplicity, consider the case without Kronecker factorization. We havediag (Q T diag (d)Q) = (Q T • Q T )d, where • is the element-wise Hadamard product. So we can express the diagonal of the matrix-matrix product as a matrix-vector product with the diagonal d as the vector. In the Kronecker-factored case, we need just one more simplification:diag((Q T A ⊗ Q T G ) diag (δ)(Q A ⊗ Q G )) = ((Q T A ⊗ Q T G ) • (Q T A ⊗ Q T G ))δ = ((Q T A • Q T A ) ⊗ (Q T G • Q T G ))δ = vec((Q T G • Q T G ) mat(δ)(Q A • Q A )) = vec(Q T G ) 2 mat(δ)Q 2 A, where we have used the mixed-product property of the Kronecker product and the properties for multiplying a Kronecker-product with a vector. The vec operator "flattens" a matrix, that is, turns a D out × D in matrix into a D out D in vector, and mat does the opposite. The final approximation δ can be computed efficiently in O(D 2 in + D 2 out ).Figure B1 illustrates that both MAP and SPAM achieve similar levels of performance throughout the training process. This observation underscores that SPAM's enhanced sparsifiability is achieved without compromising the unpruned performance. Furthermore, the comparable unpruned accuracies of SPAM and MAP models indicate that SPAM's sparsifiability benefits are not merely a result of higher baseline accuracies, but rather a distinct advantage offered by the SPAM methodology. The sparsification methods are performed on these models in a way that once the model is trained for a specific seed, we copy it and use it to perform the different sparsification methods; we repeat the steps for a minimum of 4 different seeds, ensuring the robustness of our findings. In addition, the model trained with SpaM only uses a single forward pass over the pruned architecture during inference, thus guaranteeing a fair comparison with the baselines. The posterior is solely used to estimate the marginal likelihood during SpaM training.In tables B1 and B2, we present our results comparing different methods using MAP and SpaM with various priors. Notably, SpaM with Diag LA and parameter-wise priors significantly outperforms MAP and other SpaM variants at high sparsity levels.Figure B2 and Figure B3 illustrate our findings when applying SpaM with various priors for both OPD and GraSP. Notably, Diag LA, using parameter-wise priors, excels in high-sparsity scenarios, even with complex models and datasets like ResNets. Furthermore, for MLPmixer, we observe that SpaM variants, employing parameter-wise priors and layerwise approaches, preserve baseline accuracy even at extreme sparsities of 99%.As shown in Section 5, networks trained using SpaM and parameter-wise priors were able to maintain a high accuracy at challenging sparsity levels up to 99%. Moreover, parameter-wise KFAC and unit-wise priors showed high performance for the OPD pruning approach. We show in Figure B3 that the combination of SpaM and these priors with the GraSP criterion yield qualitatively similar performance rankings as with our OPD criterion.As seen in Figure B4 , our proposed post-hoc pruning criterion, OPD, consistently demonstrates stable performance across diverse model architectures and datasets, achieving significant sparsity levels without the need for fine-tuning. It seamlessly operates either post-training or with pre-trained models, providing a highly flexible and versatile solution. Figure B5 shows that our online version of SpaM, which uses the marginal likelihood and OPD during training to iteratively prune the network, reaches comparable performance levels to the post-hoc version, thus offering a computationally even more convenient way to effectively sparsify neural networks. In the online approach, we prune jointly during SpaM training to reach a target sparsity; we perform the sparsity updates (dynamic masking) based on the marginal likelihood optimization parameters we refer to as n_epochs_burnin and marglik_frequency. Here, n_epochs_burnin specifies when we start the marginal likelihood optimization and marglik_frequency specifies after how many epochs we update the estimate. Pruning occurs only after a new marginal likelihood calculation, and the sparsity percentage is adjusted incrementally to reach the target by the training's end. The curve of OPD-Online reflects training progress (x-axis explicitly epochs / reached sparsity), which explains the curve of LeNet on CIFAR-10 that is still converging while pruning. At the same time, the one-shot post-training approach reflects converged models copied and pruned at different sparsity levels.We conducted experiments using L1 regularization, which is well-known for inducing sparsity in neural networks. To optimize performance, we performed an extensive search for the appropriate L1 regularization coefficient and applied various strengths during training.We found that SpaM consistently outperforms L1 regularization, achieving much higher levels of sparsity while maintaining the network's predictive performance as shown in Figure B6 . L1 regularization, while effective at inducing sparsity, often proved too aggressive, leading to networks that were excessively pruned, negatively affecting performance. In contrast, our method offers a key Figure B10 : Visualization of the sparsification effect on the model's layers for LeNet on FashionM-NIST sparsified at 95% using both the structured and unstructured approach. Blocks in black refer to masked filters, and columns refer to neurons pruned. We see that for unit-wise priors, a 95% sparsity yields more entire kernels and neurons being masked compared to parameter-wise priors. Disk size (Bytes) We see that through structured sparsification, we are able to obtain models that are still performant at a significantly reduced computational and memory cost, while unstructured pruning does not directly translate into computational benefits. We see that through structured sparsification, we are able to obtain models that are still performant at a significantly reduced computational and memory cost. At the same time, unstructured pruning does not directly translate into computational benefits.This leads to a network with fewer filters in convolutional layers and a reduced number of neurons in fully connected layers, resulting in a leaner and more efficient model.The process of compaction involves transferring the weights from the pruned model to a newly created, smaller architecture that is aligned with the dimensions of the retained active structures. This results in a denser, storage-and computation-optimized model.Algorithm 1 summarizes this entire process of structured pruning and model compaction.This approach transitions the model from a pruned state to a compact and optimized architecture. The final compressed model M compact not only retains essential predictive capabilities but is also further tuned for performance. The newly configured M compact is saved with updated parameters, ensuring efficient inference and ease of deployment, especially on resource-constrained edge devices.The reduced memory footprint and FLOPS of M compact are particularly beneficial for deployment on edge devices with limited computational resources. When models exceed the hardware limits, aggressive compression techniques like quantization may be required, which can compromise performance. Our method aims to significantly reduce the memory size of the model while minimizing performance trade-offs. The effectiveness of our approach in achieving this balance is explored in Section 5.Algorithm 1 outlines our structured pruning procedure, highlighting how we efficiently achieve a simpler model by transferring weights to a smaller one. [58] . LeNet on CIFAR-10 is not a very common benchmark for pruning; here, it is used to demonstrate how SpaM, and specifically SpaM-OPD, is able to prune at high percentages without a performance loss up to 80% in a model that struggles with representing the data's complexity, showing that our work extends beyond over-parametrized networks for the task at hand.• MLPMixer: The MLPMixer serves as a streamlined alternative to more complex models like CNNs and transformers. It relies solely on Multi-Layer Perceptrons (MLPs) for integrating inputs across spatial and channel dimensions [62] . We implement an MLPMixer with 2 blocks designed for MNIST. • ResNet with inplanes 64 and depth 18 for CIFAR-10: We modify the implementation of ResNet and incorporate fixup initialization and custom bias and scale parameters to align with the constraints of the ASDL backend [63] used for the Laplace computations in this work, which does not support batch normalization. • Wide ResNet: decreases depth compared to ResNet and increases the width of residual networks [64] with a depth of 16 and a widening factor of 4 (WRN16-4). We use fixup blocks to be able to utilize ASDL backend [63] . • Vision Transformer (ViT) [65] : unlike CNNs, which extract local features through filters and pooling layers, ViT breaks down images into fixed-size patches, treating each as a "token" in a sequence [65] . This allows it to leverage the Transformer architecture, initially designed for language processing, to analyze relationships between patches through self-attention mechanisms [66] . • DistilBERT: DistilBERT [67] is a smaller, faster, and cheaper version of BERT, achieved by leveraging knowledge distillation during the pre-training phase. This model retains 97% of BERT's language understanding capabilities while being 60% faster and 40% smaller. We use the pre-trained DistilBERT hosted in Hugging Face under (distilbert-base-uncased) [67] and tune it for sentiment analysis to classify reviews in the IMDB dataset [61] , which involves predicting the sentiment (positive or negative) of user reviews based on their textual content. • GPT-2: a large-scale transformer-based language model developed by OpenAI, with impressive text generation capabilities. Trained on a vast corpus of internet text [68] . In our study, we leverage the 124M parameter version of GPT-2, fine-tuning it on the IMDB dataset for sentiment analysis to assess its performance under different pruning conditions.For the computation of second-order information (e.g., Hessian, Fisher information) needed for the Laplace approximation, we utilize the ASDL Library [63] . We use the library in its version under https://github.com/kazukiosawa/asdl/tree/ 011a942b2698b9ec33b0c8c47c96bd49335e5d80. The ASDL Library is distributed under the MIT License, which allows for reuse with a few restrictions that we respect in our work.• Hessian Approximation: The choice between GGN and EF. GGN was initially employed for fully connected networks, LeNets, and ResNets. However, for complex architectures (WRNs, ViTs, DistilBERT), GGN's computational cost became prohibitive, exceeding MAP runtime by up to 20x and even more for casual modeling tasks. Switching to EF maintained pruning performance while closely matching MAP runtime, which is particularly beneficial as GGN scales linearly with the number of classes. We discuss further the cost in Appendix D.7. • n_epochs_burnin Dictates the number of epochs after which marginal likelihood optimization starts. If set superior to the number of training epochs, marginal likelihood is skipped, and the training is equivalent to MAP. • marglik_frequency Controls the frequency of marginal likelihood estimation. The default value of 1 signifies re-estimation after each epoch, while a value of 5 indicates approximation for every fifth epoch.We use these parameters to manage the computational cost of our experiments, where for small models like LeNets, FC Networks, the n_epochs_burnin is set to zero and marglik_frequency to one reflecting estimating each epoch since the start of the training. In contrast, for complex networks like emphasizing the impact of pruning on the training dynamics [37] . We replicate the GraSP implementation of Rachwan et al. [39] , where we consider the absolute value of the importance score initially proposed by Lubana and Dick [38] given by:I(Θ t ) = |Θ T t H L (Θ t )g L (Θ t )| (D2)Note that while the importance score was initially used before training, we propose to use this importance score as a one-shot criterion after the training process and show how SpaM can leverage the performance of GraSP.• Structured-SynFLow: We challenge the capabilities of SynFlow [50] , a data-agnostic pruning approach that prevents layer collapse that happens at high sparsities where layers are no longer able to perform at the model's predictive power. This typically occurs when the pruning algorithm, intentionally or inadvertently, removes a significant portion of weights or filters from a specific layer, effectively collapsing its functionality [50] . We push SynFlow to its limits through advanced structured pruning strategies, where we prune layers aggressively at the same target sparsity, which facilitates the compression process and resizing. By applying rigorous layer-specific filtering and neuron pruning, we aim to test the robustness and effectiveness of SynFlow in extreme sparsity structured scenarios. This approach not only benchmarks SynFlow's performance under stringent conditions but also explores its potential to maintain network functionality and accuracy in highly sparse neural network architectures. • Magnitude Pruning: Relies on the magnitude of weights for pruning, aiming to maintain model performance while reducing complexity [69] . After the success shown by Han et al. [69] , many methods adapted magnitude as a pruning criterion coupled with different scheduling [49, 70, 48] . • Random Pruning: Prune weights or structure randomly.For structured sparsification, contrasting with the unstructured approach, the process necessitates reshaping the weight matrices to effectively reduce model complexity. The steps include:1. One-shot structure masking based on aggregated importance scores. 2. Continue training for five epochs using the model from Step 1 for preliminary evaluation. 3. Implementing two software design approaches:• In-place layer replacement in the model with smaller ones fitting the non-masked regions. • Creating a new, flexible model initialized to match the dimensions of the non-masked areas, requiring repeated reading of the nonzero mask for state-dictionary and metadata alignment. 4. Transferring non-zero structures to smaller layers and tuning the model.Post structure removal, we extend the training phase to adapt the model weights and re-evaluate, ensuring seamless functionality once transferred to smaller layers. Particularly after significant structural reduction, our primary objective shifts to maximizing performance in the downsized model. This fine-tuning spans 5 or 10 epochs depending on the complexity of the original model's structure, which was initially trained for either 50 or 100 epochs.Instead of using the Generalized Gauss-Newton (GGN) approximation, which scales linearly with the number of classes, we can also use the Empirical Fisher (EF). For most architectures, using EF instead of GGN for SpaM does not add a very large computational overhead to MAP, as EF costs roughly as much as gradient computation. This is particularly beneficial as GGN scales linearly with the number of classes. The pruning results are not significantly affected by the choice of GGN or EF.The runtime of MAP and SpaM was close (roughly 1h and 20 minutes on A100s) for WRN-16 on CIFAR100 using SpaM (EF) with diagonal LA and parameter-wise priors (our recommended settings for pruning). For language transformers, specifically DistilBERT and GPT-2, SpaM with EF does result in a longer training time compared to MAP. However, this increase is considerably less than when using GGN, where a single epoch can take longer than the entire SpaM training with EF.In prior works [71] , it was found that GGN gives a better posterior predictive approximation, but we do not use it in this work. We find that EF works similarly well for pruning at a much lower cost. 


Decision: This paper introduces a framework called Sparsifiability via the Marginal likelihood (SpaM) that leverages Bayesian marginal likelihood with sparsity-inducing priors to enhance the sparsification of neural networks. All reviewers praised the originality and effectiveness of Bayesian marginal likelihood as a pruning criterion, and most reviewers also appreciated the thoroughness of the experiments performed. Some reviewers questioned the choice of Laplace approximation and Gaussian approximation and asked for comparison with previous works featuring other sparsity-inducing priors like horseshoe priors; the authors mainly defended from the perspective of computational feasibility and the fact that some of the prior works require more invasive changes to the architectures/training protocols, which largely addressed the reviewers' concerns. On this point, while the AC also does not see this as a major issue post-rebuttal and acknowledges the computational concerns, the AC still thinks that some empirical comparison against these prior works, even on smaller scale setups where these baselines are indeed applicable, could still be helpful to strengthen the arguments in the paper. 

Other concerns were raised on novelty (although the AC does not view this as a major weakness given that all reviewers acknowledged the originality and effectiveness) and comparison against approaches with type-II MLE, but the authors have answered the questions satisfactorily.

Post-rebuttal, all reviewers unanimously recommend acceptance and the AC concurs with this recommendation. The authors are encouraged to incorporate all requested changes and discussions into the camera-ready version of the paper.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: **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: 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?
limitations: limitations: The authors have adequately addressed the limitations of their approach.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.?
limitations: limitations: The authors mentioned the limitation coming from the Laplace approximation.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: ### 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.
limitations: 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.

Comments: title: Official comment
comment: Dear authors,

Thank you for the additional comments - the discussion on the L1 regularisation is insightful and really interesting. I'll keep this in mind and discuss further with the other reviewers.
Thank authors for addressing my concerns in the author response. I have increased my score to 5 as a result of the clarification. However, I still have some reservations about the practicality of the proposed method.
Thank you for the clarification.

I maintain my score of 8 (strong accept) which reflects my opinion that this is good work that has been carried out to an excellent standard.
Thank you for your rebuttal, which answered my questions from a motivation perspective. I have increased my score accordingly.
Thanks for the follow-up question.

> Does SpaM do anything to encourage weights to be amenable to structured pruning in particular, or does the question of "structured or unstructured" only affect the final pruning using OPD?

Apart from differently aggregated weights in OPD between structured and unstructured, SpaM can also encourage structured pruning by specifying priors that correspond to sensible groups, for example, rows or columns of weight matrices. One such prior would be the unit-wise prior defined in lines 130-134. One unit can correspond to a neuron in the fully-connected case or a filter in the convolutional case

> I personally think it would be useful to have a diagram entirely devoted to the SpaM training process.

Thanks for the suggestion. We will try to include this process visually or make a separate figure or algorithmic description for it in the next revision.
Thank you for raising your score. We are happy that we have addressed your concerns.
>Which of the plots in your original submission should I compare the rebuttal plot to? In particular, I'm looking for the MAP (L2 reg.) results for the same experiments.

The experimental plots to compare to would be Figure 2 bottom left (WRN on CIFAR-100) and Figure B6 top right corner.

>It's quite interesting that the L1 regularisation performs this badly, I think. Do you have an intuition for why this is? Perhaps L1 regularisation is simply too aggressive? To be clear, I'm very happy that your method works better, I'm just trying to understand the reason, if possible.

We have thoroughly searched for an L1 regularization coefficient but could not find a setting that led to better performance. Indeed, L1 regularization seems to be too aggressive and leads to an overly pruned network in the end. One advantage of our method in comparison to L1 regularization is that it applies different regularization per weight and adapts to the data instead of one global fixed parameter. Theoreticaly, Wipf and Nagarajan [1] find that the per-parameter learned regularization that we use (ARD) is related to a complex form of per-weight L1 regularization. Such a per-weight L1 regularization is, however, impossible to realize in practice as it has too many hyperparameters.

[1] Wipf, D., & Nagarajan, S. A new view of automatic relevance determination. NeurIPS 2007.
Thank you for your detailed rebuttal. You answered many of my questions satisfactorily.

> Does SpaM allow for improving structured sparsifiability during training as well as just unstructured?

Sorry, my question doesn't seem to have been clear. I meant "does SpaM do anything to encourage weights to be amenable to structured pruning in particular, or does the question of "structured or unstructured" only affect the final pruning using OPD?". I believe I'm specifically talking about when you pruning post-hoc, rather than online.

Thinking more about Figure 1, I personally think it would be useful to have a diagram entirely devoted to the SpaM training process (e.g. compute laplace approximation, do backprop for K steps, update laplace approximation etc.), rather than focusing on the actual pruning using OPD after, as the main idea of the paper to me seems to be encouraging weights to be sparsifiable during training, using clever priors.
Dear authors,

Thank you for your reply and for performing the additional experiment with L1 regularisation. Which of the plots in your original submission should I compare the rebuttal plot to? In particular, I'm looking for the MAP (L2 reg.) results for the same experiments. 

It's quite interesting that the L1 regularisation performs this badly, I think. Do you have an intuition for why this is? Perhaps L1 regularisation is simply too aggressive? To be clear, I'm very happy that your method works better, I'm just trying to understand the reason, if possible.

In any case, I have increased my score. You have nicely addressed my questions and concerns, and I also think my comment that your contribution seems incremental was perhaps too harsh.