Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: #### Summary

The authors focus on the important problem of scalable approximate inference in Bayesian NNs. More specifically, they propose a method for scalable BNNs via a (full-covariance Gaussian) Laplace approximation on a (Wasserstein-based) pruned subnetwork within a deterministically-trained model. They include a theoretical analysis for a simple generalized linear model, and experiments on 1D regression, tabular regression, and larger-scale image classification with CIFAR-10 (using the dataset shift setup from Ovadia et al., (2019)). From the experiments, they show that their method generally outperforms comparable methods (including deep ensembles) on metric performance and on the ability to capture in-between uncertainty.

#### Strengths

- Scalable approximate inference for Bayesian models is an important research area.
- Expressive, diverse approximate posteriors is also an important area of research, especially given the limitations of mean-field VI and recent literature.
- The proposed method demonstrates better results for robustness to dataset shifts, as well as in-between uncertainty that mean-field VI misses.
- In general, the paper is well-written, clearly-motivated, includes both theoretical and empirical results, and adequately compares to, or discusses, relevant literature in the space.

#### Weaknesses

- The authors push on the idea of *scalable* approximate inference, yet the largest experiment shown is on CIFAR-10. Given this focus on scalability, and the experiments in recent literature in this space, I think experiments on ImageNet would greatly strengthen the paper (though I sympathize with the idea that this can a high bar from a resources standpoint).
- As I noted down below, the experiments currently lack results for the standard variational BNN with mean-field Gaussians. More generally, I think it would be great to include the remaining models from Ovadia et al. (2019). More recent results from ICML could also useful to include (as referenced in the related works sections).

#### Recommendation

Overall, I believe this is a good paper, but the current lack of experiments on a dataset larger than CIFAR-10, while also focusing on scalability, make it somewhat difficult to fully recommend acceptance. Therefore, I am currently recommending marginal acceptance for this paper.

#### Additional comments

- p. 5-7: Including tables of results for each experiment (containing NLL, ECE, accuracy, etc.) in the main text would be helpful to more easily assess
- p. 7: For the MNIST experiments, in Ovadia et al. (2019) they found that variational BNNs (SVI) outperformed all other methods (including deep ensembles) on all shifted and OOD experiments. How does your proposed method compare? I think this would be an interesting experiment to include, especially since the consensus in Ovadia et al. (2019) (and other related literature) is that full variational BNNs are quite promising but generally methodologically difficult to scale to large problems, with relative performance degrading even on CIFAR-10.

##### Minor
- p. 6: In the phrase "for 'in-between' uncertainty", the first quotation mark on 'in-between' needs to be the forward mark rather than the backward mark (i.e., $`in-between'$).
- p. 7: s/out of sitribution/out of distribution/
- p. 8: s/expensive approaches 2) allows/expensive approaches, 2) allows/
- p. 8: s/estimates 3) is/estimates, and 3) is/
- In the references:
  - Various words in many of the references need capitalization, such as "ai" in Amodei et al. (2016), "bayesian" in many of the papers, and "Advances in neural information processing systems" in several of the papers.
  - Dusenberry et al. (2020) was published in ICML 2020
  - Osawa et al. (2019) was published in NeurIPS 2019
  - Swiatkowski et al. (2020) was published in ICML 2020
- p. 13, supplement, Fig. 5: error bar regions should be upper and lowered bounded by [0, 1] for accuracy.
- p. 13, Table 2: Splitting this into two tables, one for MNIST and one for CIFAR-10, would be easier to read.
review_title: Review

final_score: rating: 6: Marginally above acceptance threshold
review_content: The paper proposes to approximate the posterior distribution of a Bayesian neural network by an approximation that consists of a deterministic component. The authors select a sub network and infer approximate posterior distributions over the weights in the sub network. All other weights are estimated via MAP point estimation.  A sufficiently small sub-network allows high fidelity posterior approximations that do not make restrictive mean field assumptions to be tractable.

The paper is generally well written and easy to follow. The idea that BNNs have too many parameters for reliably inferring posterior distributions over all of them is a reasonable one. Splitting the posterior approximation into a deterministic and stochastic component to deal with this issue is interesting.  The experiments do indicate improvements over fully factorized posterior approximations. 

Concerns:
* The process of selecting the sub-network over which the full posterior distribution is inferred is crucial and herein lies my main concern with the approach presented in the paper. Minimizing the Wasserstein distance between the subnetwork posterior p(W_s | data) and the true posterior over all weights p(W | data) is sensible. However, since the true posterior is intractable, the authors instead appear to minimize the distance to an Laplace approximated posterior q(W | data). While this is fine for smaller models when the approximation can use a full covariance,  why does it make sense for larger models where the authors use diagonal approximations to the generalized Gauss Newton matrix (Section 5.3)? If the goal is to do better than such diagonal approximations then why treat the subnetwork with the lowest Wasserstein distance to such crude approximations as optimal? How much worse is the  random selection strategy on the tasks listed in section 5.3? 
* In the toy experiments, when the full covariance posteriors approximations  are available, using the deterministic + stochastic sub-network approximation does a little bit worse than the full covariance approximation but better than the diagonal approximation. If instead the diagonal posterior approximation is used to select the subnetwork in these models, does the resulting approximation still improve upon the diagonal approximation (as seems to be happening in the experiments in 5.3)? 
* The experiment in section 5.3 has other curious issues. Different methods appear to be using different priors (Gaussians with different precisions). How do we then know that the benefits reported stem from the proposed approximation rather than the differences in the model, especially since diagonal Laplace is using a Gaussian prior with a precision that is 80 times smaller?  What priors were used for deep ensembles and SWAG? 
* The grid search by which the prior precisions were selected needs more details. The presented numbers are surprisingly small suggesting that a-priori  with high probability we expect all weights to be zero, and hence the prior predictive functions to be all zero as well. This does not seem like a sensible prior. 
* At the very least there needs to be a discussion about sparse Bayesian deep learning techniques (and preferably an empirical comparison)[1, 2, 3], that use sparsity inducing priors to prune away weights and nodes from a larger network instead of the approach presented here. 
* (Minor) The conclusion of 5.3 that subnetwork posteriors are better calibrated are not supported by Figure 4. I would suggest moving the ECE and Brier score plots from the appendix to Figure 4. 

Overall, although I have several concerns they primarily stem from experimental issues in 5.3. Assuming that the authors are able to sufficiently address these in the rebuttal, I am leaning towards a borderline accept. 

[1] https://www.jmlr.org/papers/v20/19-236.html

[2] https://papers.nips.cc/paper/7372-posterior-concentration-for-sparse-deep-learning.pdf

[3] https://arxiv.org/pdf/2002.10243.pdf
review_title: Interesting ideas but various clarifications are needed

final_score: rating: 5: Marginally below acceptance threshold
review_content: The authors present a new method for Bayesian deep learning motivated by the difficulty of posterior inference in the "overparameterized" regime of deep neural network models. The proposed method provides a principled strategy for selecting a subset of the neural network's parameters (forming a so-called "subnetwork") for which a full-covariance approximate posterior can be computed. The authors use the well-studied Laplace approximation with the generalized Gauss-Newton Hessian approximation for the covariance. An empirical analysis is presented which attempts to assess the efficacy of the proposed method in prediction accuracy and uncertainty quantification.

The presented approach is novel and appears to be a promising contribution to the study of Bayesian neural networks. As a fellow Bayesian, I applaud the authors' efforts. Unfortunately, the paper has a number of significant weaknesses which I detail below. The authors' experimental results appear to me to not sufficiently support some of their claims. There are also a number of formatting issues. As such, I cannot recommend accepting this work in its current state. I would be happy to revisit my rating after revision and discussion.

Pros:
- The basic idea behind the method is very interesting and constructive for the field
- Theoretical justification is solid
- Overall well written with only minor clarity issues

Cons:
- Experimental results lack thoroughness
- Results do not seem to adequately support authors' (somewhat bold) claims
- Lacking detail in certain areas of the method description


I will organize my comments for the authors by section.

#### Section 1
1. "In turn, we can apply more expensive, but more faithful, posterior approximations to just that subnetwork to achieve better uncertainty quantification than if we apply cheaper, but more crude, approximations to the full network."
This sentence is too long and too difficult to read. Please try rephrasing to make it more clear.

#### Section 2

2. It would be nice to see the subnetwork posterior fully defined in terms of the fixed weights, e.g:
$$p(W_s | y,X,W^*) \propto p(y|X,W_s,W^*)p(W_s)p(W^*) = p(W_s|y,X)p(W^*)$$
I don't meant to be pendantic here. It wasn't obvious to me (at first) that the delta functions were, in fact, a degenerate prior over the fixed weights, $W^*$; this would make it more clear.

3. It might be worth discussing here or later what implications the degenerate prior has for the subnetwork posterior. This seems to, at least, preclude any ideas about possibly applying gradient-optimized MAP directly to the subnetwork posterior.
Furthermore, what are the benefits of this approach as opposed to assigning a non-degenerate Gaussian prior with fixed variance to the remaining weights?

#### Section 4

4. "We now analyze the following procedure..."
Are you analyzing the procedure for selecting a subnetwork? Or the entire procedure described previously? This is not made clear.

5. In step 1, you specify the analytical solution for w_MAP, but in section 3 you describe MAP as being performed using stochastic gradient optimization. Which one are you using? Related to previous point, is this the same step as outlined in section 3? Or a separate step?

6. In general, this section needs work. It is not clear to me where you actually detail how the subnetwork is selected. You discuss approximating the optimal subnetwork w.r.t to the Wasserstein distance in equation 11, but this equation requires M_S to be already available. M_S is generated by "a (one-shot) procedure of choice"; I expected this choice to be clearly explained in this section, but this is not the case.

#### Section 5.1

7. It's a bit confusing that you say "50%, 97%, and 99% of model parameters" when really you mean that this percentage of weights were pruned. It would be visually more intuitive as well in Figure 2 if you named these 50%, 3%, 1%, since the posterior size is getting *smaller* with each one.

8. Just to clarify, homeoscedastic here is w.r.t to the sequence of data points? i.e. you have one variance rather than one per data point?

9. Please provide some explanation for the prior precision being set to 3.

10. Bolding most of the last sentence in section 5.1 is unnecessary and looks weird. I would suggest bolding individual words or phrases in the sentence, or just not bolding at all.

#### Section 5.2

11. 1e4 ->$10^4$

12. It would be helpful to provide a very brief summary of the three datasets as well as what the "gap variants" are (it's fairly straightforward, just creating "gaps" in the training data). I cannot find the "kin8nm" dataset on UCI, so at bare minimum this needs to be clarified.

13. Figure 3 is quite difficult to interpret at first glance. You should at least use a discretized perceptually uniform colormap to make more visually apparent the pattern in increasing network size. I would also consider using a line plot rather than scatter plot here, but this is a matter of taste.

14. Why are you using the log likelihood as opposed to the full posterior probability? Since we're in a Bayesian setting here, it seems worth considering the priors.

15. While log likelihood can, in principle, serve as an indirect proxy for uncertainty calibration, there are important caveats to consider which make it unconvincing as a standalone measure (see "Pitfalls of In-Domain Uncertainty Estimation...", Ashukha et al. 2020). Furthermore, a quick scan of recent literature (cited in this work) confirms that most authors tend to use multiple methods of assessment (e.g. Brier score, accuracy calibration curves, etc), not just log likelihood. The raw likelihood scores also don't provide any interpretability in *how much better* one model is than another, i.e. what concrete impact does the difference have on uncertainty quantification? Thus, **it's unconvincing that here (and elsewhere) you rely entirely on log ilkelihood** to demonstrate the alleged superiority of your method.

16. It's noteworthy that 2/3 of the "gap" datasets produce inconsistent results, particularly considering that they are designed specifically to test out-of-sample uncertainty. This needs to be discussed.

17. typo: modelling -> modeling

18. In light of points 15 and 16, I do not think you have sufficient empirical basis to make the claim that "...it is better to perform subnetwork inference in a large model than full network inference in a small one". The results are simply too weak to support this. I suggest either running a more comprehensive experiment or dialing back the confidence of this claim.

#### Section 5.3

19. You mention "grid search" multiple times but do not provide any indication of the grid you searched over.

20. typo: sitribution -> distribution

21. Please specify the error metrics being used in figure 4 (and discussed in the text).

22. Once again, marginally higher likelihood scores are not, in my view, sufficient basis to claim that your method assigns "high uncertainty to out of distribution points" or that it is better calibrated. This is especially the case considering that there seems to be little to no improvement in the error. The case for improved "robustness" also appears to be equally unfounded on these grounds.
review_title: Good idea, solid theoretical foundation, but empirically weak

final_score: rating: 5: Marginally below acceptance threshold
review_content: The paper presents a new way for approximating posteriors in Bayesian DNN. The network is split into two subnets. One uses only point estimates while another one uses full (non-diagonal) Gaussian approximation. The structure of that subnet is found by taking largest second derivatives of Hessian of linearized DNN (the authors call it generalized Gauss-Newton (GGN) matrix). The authors show that under very specific conditions such choice correposnds to minimization of Wasserstein-2 distance between their approximation and the true posterior. In the experimental part they provide a set of explorative experiments showing that it may be better to use their approximation for inference in large newtork than both using standard (simple) approximations in large network and full Bayesian inference in small network. This is very nice methodologically and I welcome such demonstration but this can only be considered as a (good) proof of concept. The flagship experiment however looks very inconvincing (see below).

Pros.
1. Interesting idea of finding a better approximation for the posterior on subset of parameters.
2. Methodologically nice PoC.
3. Thorough comparison with alternative similar techniques such as SWAG.

Cons.
1. The authors claim that they theoretically characterize the descrepancy and derive optimal strategy (see contribution 3) but they (0) consider linearized approximation of DNN (they admit this); (1) do this ONLY for regression problem although their flagship experiment is on classification problems; (2) the method they derive is based on un-natural assumption that covariance matrix is diagonal (it the matrix is diagonal there is no way to approximate with a full submatrix anyway). I would not call it an optimal strategy - it rather looks like a reasonalbe heuristic.
2. My major concern is section 5.3. The authors claim that their method estimates uncertainty better than all baselines including deep ensembles (DE). I am afraid in its current form the comparison is not fair. They use only DE of size 5 while their method requires approx. (42K)*(42K) = 1756B of parameters that is enough to keep in memory about 160 initial networks. So it would be fair to compare aganist DE of size 160 since we know that larger DE estimate uncertainties better. It is not that surprizing that the suggested method outperforms other baselines since all of them require much less memory. So I would recomend the authors to compare (1) their current model aganist DE that requires similar amount of memory; (2) their reduced model (that requires approximately the same amount of memory as baselines) aganist other baselines to check wether the proposed algorithm may still estimate uncertainty better given the same memory budget.
review_title: Good PoC but the main experiment raises questions

Meta-review: This paper propose an approach to efficient Bayesian deep learning by applying Laplace approximations to sub-structures within a larger network architecture. In terms of strengths, scalable approximate Bayesian inference methods for deep learning models are an important and timely topic. The paper includes an extensive set of experiments with promising results. 

In terms of issues, the reviewers originally raised many concerns and the authors provided a large update to the paper. However, following that update and the discussion, several concerns remain. First, the reviewers noted that the originally submitted draft made claims about the optimality of the sub-network selection procedure that were incorrect due to the use of a diagonal approximation. The authors subsequently retracted these claims and re-focused on the idea that the subset selection approach is theoretically well-motivated heuristic that performs well empirically. Following the discussion, the reviewers continued to express concerns about the heuristic nature of this procedure. 

A second point has to do with scalability. The reviewers noted that the authors had only evaluated their approach on small data sets, leaving open the question of how scalable the method is. The authors responded by adding experiments on the same data sets using larger models, which does not squarely address the issue raised. Third, an additional point was raised regarding the lack of control of resource use in the experiments. The authors note that their approach can use more resources when available while many other methods can not. However, some methods including deep ensembles can also expand to use more resources, as can posterior ensembles produced using MCMC methods like SGLD and SGHMC. The authors need to consider quantifying space-performance and time-performance trade-offs in the same units for different approaches to satisfactorily address this issue. While the authors added one set of experiments looking at deep ensembles in isolation, their conclusions that performance saturates for these models at low ensembles sizes seems to be hasty in some cases (e.g., deep ensembles show continued improvement for large corruptions in Figure 5(right) despite the claim by the authors that the models saturate after 15 epochs). 

In summary, this appears to be a promising approach. While the authors made significant efforts to correct issues and address questions with the original draft, the majority view of the reviewers following discussion is that this paper requires additional work to more carefully expand on the revised results and to address the heuristic status of the sub-network selection approach.
Decision: Reject