Riemannian Laplace approximations for Bayesian neural networks

Summary

This paper develops a Riemannian Laplace approximation, which is a Laplace approximation that takes into account the Riemannian geometry of the loss surface. The contributions of this paper are as follows: i) showing that such a loss-aware Laplace approximation is better able to capture the true posterior (and predictive); ii) presenting the Riemannian geometry framework for the Laplace approximation, with the Hessian in both the normal and tangential space; iii) a practical algorithm for efficiently integrating the required ODE; iv) experimental evidence on several commonly used datasets.

Strengths

The paper is for the most part very well written, it is for the most part easy to follow (main idea and background). And the authors give several examples (including figures) to help the readers further. The experimental section is decent, having both a toy example that helps gain intuition and a quantitative evaluation on several standard datasets, where the method is benchmarked against vanilla Laplace approximation and the MAP estimate. I think for the purpose of this paper, the various alternative approximate inference algorithms would not be necessary to compare to, as this is a direct extension of Laplace.

Weaknesses

Even though the paper is generally very well written, I did, however, struggle with some parts of the main section. As a reviewer who is very familiar with Bayesian neural networks and the Laplace approximation but less so with Riemannian geometry, I would prefer to have a more in-depth background section on Riemannian geometry. I think the Laplace approximation section can be shortened if space is needed (e.g. tricks of the trade and strengths and weaknesses could be shortened). The main weakness is the scalability of the method (see limitations). Solving the ODE takes a very long time. According to Fig. (4), this is in the order of tens of seconds for a mini-batch. Even for very small neural networks. The models used in the experimental section are tiny, e.g. single hidden layer networks with only 50 hidden units. For the CNN, the authors mention that they use 2 conv layers, but I didn't find the kernel size. It is also not clear to me if and which approximation the authors use for the Hessian. Can the complexity be reduced e.g. with a diagonal Hessian approximation? How would this compare to a diagonal or Kronecker-Factorised approximation of the standard LA?

Questions

In Figure 4, why does the NLL have a minimum and not strictly decrease with larger mini-batch size? Shouldn't the estimate become more and more precise? Is the wall-clock time really on the order of tens of seconds? What is the model size? For the CNN, what is the kernel size? Which approximation to the Hessian do you use for the Riemannian Laplace approximation (e.g. Diagonal or Kronecker-Factorized)?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The method has been applied to very small neural networks and the authors argue that this is because the method is computationally very expensive. While the authors acknowledge this in the limitations, it potentially limits the applicability to realistic neural network model sizes. It would be important to know more precisely what the computational complexity is (also as a function of the model size). Furthermore, I would like to see a comparison in wall clock time for the entire approximation compared to standard Laplace approximation, not only for different mini-batch sizes, but also for different model sizes, including large model sizes.

Summary

The paper presents a Laplace approximation for Bayesian neural networks that adapts the covariance to the local geometry of the loss, effectively overcoming the quadratic approximation of the loss. The authors report competitive performance with the standard Laplace approximation (both Monte Carlo sampled and linearized) on regression and MNIST-scale classification problems as well as a reduced reliance on tuning the precision of the prior. The approach is explained clearly and makes a lot of sense (at least with my rather superficial understanding of Riemannian geometry), it is applicable in more general probabilistic models and I would expect it to lead to various follow-up works. While methodologically this is a very nice paper, I feel like it is let down by the empirical evaluation. The method is only tested on UCI and MNIST-scale datasets, which are hardly relevant for deep learning these days. The authors mention the computational cost of their method, but only discuss the reasons superficially without providing exact benchmark figures to give a sense of where the main bottlenecks arise. Given the apparent computational limitations of the method, an experiment with a non-NN model could have strengthened the paper. All things considered, the strong methodological contribution outweighs the unconvincing empirical evaluation for me, so I would lean towards acceptance, although I wish could have given the paper a much higher rating.

Strengths

* The core idea makes a lot of sense, is applicable beyond inference in neural networks and seems to work well for the experiments that are considered. It has the potential to address the rather restrictive approximation of a quadratic loss in the Laplace approximation. * I am confident that the paper will inspire various pieces of follow-up work. * The paper is well-structured and -written. * Effective use of illustrative examples throughout.

Weaknesses

* Only small scale problems are considered in the experiments * This is exacerbated by lack of analysis of the computational cost. It is not really clear to me what specifically is preventing the method from being applied to larger networks and datasets (even something like CIFAR with ResNets would have been great). The discussion mentions scaling issues w.r.t. number of datapoints and parameters, but the relative behavior is not benchmarked at all and it is also not clear to me how much time the ODE solver actually spends e.g. calculating the Hessian-vector products vs computations independent of that. Given that this paper opens up a new direction for research, I think there should be much clearer pointers as to where exactly the current bottlenecks and limitations are and where improvements can realistically be achieved. * There is no experiment demonstrating the efficacy of the approach on a non-NN probabilistic model. Given the apparent computational cost of the method, this would have seemed like a rather natural experiment to include and neural networks don’t seem like a particularly good fit for the approach. * The explanation of why the method would work better with a mini-batch rather than the full dataset (section 3.3/fig 4) is not exactly clear and seems rather hand-wavy. * Lack of HMC ground-truth baselines for the regression problems

Questions

* What are the absolute runtimes of each method per sampled prediction in the expeirments? It would be great to get a better sense of this as a function of dataset/minibatch size and number of parameters/network size (I could see a synthetic experiment be illuminating here). * How much of the cost of the method lies in calculating the Hessian-vector products? These might take up a significant chunk of the total ODE compute time, so I’m wondering if explicit Hessian approximations (last-layer, KFAC, subset, ...) might improve scalability? * Could you elaborate on the discussion at the end of section 3.3? In particular I don’t follow how using the full dataset would over-regularize the geodesic as judging from the $N/B$ factor in the inline equation you seem to correctly rescale the mini-batch loss to match the full loss in expectation (note: I assume by ‘over-regularize’ you mean concentrating the samples around the mean, i.e. effectively reduce the entropy of the Gaussian approximate posterior). Any ideas for overcoming this over-regularization? * Could you comment on the results for the Riemannian Laplace approximation being noticeably better on the UCI datasets with a non-optimized precision? **Minor note**: the x axis labels in Fig 4 are cut off at the bottom

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Only discussed superficially and in a quite hand-wavy manner, I would have wanted to see concrete benchmarks and an analysis of how the compute time evolves with increasing dataset size and number of parameters respectively as well as some evidence from the literature that a tailor-made ODE solver could indeed allow the method to take the step from small conv nets to more modern architectures.

Summary

This paper presents a novel Laplace Approximation for Bayesian Neural Networks. A key insight is to examine the local loss landscape with a Riemannian metric, which is determined by the gradient of the log posterior. Using this metric and an exponential map, a Laplace Approximation technique is developed to draw posterior samples that fall into regimes with low negative log posterior. The paper also develops a sampling method, which relies on the 2nd order ODE solver. Several experiments are conducted. When compared to the standard Laplace Approximation, evidences are provided to illustrate the improvements.

Strengths

- the contribution provided by this work is original and novel to the best of my knowledge. - the paper is polished well generally. Despite that the materials are developed on differential geometry, intuitions are relatively provided well. - Laplace Approximation has been increasingly popular in recent years. Such extensions to incorporate Riemannian geometry could be relevant to the Bayesian Deep Learning community.

Weaknesses

One complaint about the paper is that, without referring to the appendix, it is difficult to comprehend the material fully. For example, in section 3.3, I wish that the connection between an ODE and Riemannian metric is difficult to understand directly. Differential geometry is not often thought in engineering courses at many universities. It may make sense to recap the essential concept in the main paper. How the method could be used for linearized Laplace Approximation is made very short. Another point for improvement is the choice of the baselines. It would make sense to include a deep ensemble and MC-dropout as a minimum. This could show how far the proposed Laplace Approximation can compete with popular methods in practice. In the experiments, the paper could improve on analyzing the computational complexity, in comparison to the standard Laplace Approximation. DaxBerger et al 2021 claim that the major benefit of Laplace Approximation is on simplicity, it would be great to project this paper's method more on the plateau of quality of uncertainty Vs computational complexity. While it is great for research, I think the methods based on differential geometry may have certain drawbacks among practitioners. The paper could be more convincing by providing the gains in uncertainty, but additional overhead due to the added complexity of the pipeline.

Questions

N/A

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There is a limitation section at the end of the paper.

Summary

The Laplace approximation offers a practical posterior but is limited due to the symmetry of the weight space it is parameterised in. The method proposed to improve posterior quality by adapting the posterior shape through a Riemannian metric that is determined by the log-posterior gradient.

Strengths

* Practically The community is interested in using Laplace approximations for Bayesian posterior estimation due to its simplicity and wide applicability. The proposed paper deals with improving the quality of such approximations, which would benefit many. * Methodologically Since the quality of the Laplace approximation is inherent to the space it is parameterised in, it is very interesting to see how Riemmanian geometry of the loss landscapes provides a better understanding of the approximations and provides avenues to improve posterior quality. * Results It is very promising to see that the Riemannian Laplace approximation allows for good posterior fits even when no linearization or prior tuning is being used.

Weaknesses

Overall, I think the promise of the paper of using the Riemmannian geometry of the loss surface to improve the posterior quality of Laplace approximation is very promising. My concerns mainly lie in the practicality of the approach in terms of scaling the method to larger models (deeper networks and models with more parameters) and larger datasets. Since the paper presents the method in the context of Bayesian deep learning, the scalability of the method is important. - Scalability to deeper networks The posterior samples from 'POSTERIOR - TWO LAYERS MODEL' of Fig. D.4 seem off. It would be good if the authors addresses how the method would perform for more complex model classes. Does the method break down in this case, or could this potentially be mitigated? - Scalability to models with more parameters As mentioned in the paper, there is a computational cost associated with the growing dimensionality of the parameter space because the number of necessary solver steps increases. I am worried that the method can not be applied to larger deeper NNs as the solutions that are found by ODE solver in practice for larger dimensions will not be of sufficient quality. Since deep neural networks typically consist of many more parameters than the models, this seems like a very big limitation. - Scalability to larger datasets The linear scaling in the number of data points. - Quantitative results and comparisons The method only considers very small models (e.g. single or 2 layer NNs) and small toy problem datasets. This small data regime would allow computing of (close to) exact posteriors, which would allow better quantitative assessment of the posterior found by Riemmanian Laplace approximation. Furthermore, it would be interesting to see larger model and data regimes.

Questions

a) What are the method's memory and computational budget limitations? b) How accurate is the ODE solver when the number of dimensions grows? If the quality of the solutions degrades in higher dimensions, can we expect the method to remain functional for larger models? c) What are current limitations to scale the method to larger models and datasets (e.g. resnets/transformers models at cifar/imagenet scale)? d) MacKay also considered notes that the choice of basis for a Laplace approximation is important in [1], which might be a relevant reference. What are the most important reasons to consider the proposed adaptation using the Riemannian metric over other potential changes of the basis? [1] MacKay, David JC. "Choice of basis for Laplace approximation." Machine learning 33 (1998): 77-86.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

As mentioned, I think being able to adapt the quality of Laplace approximations by considering geometrical aspects of the loss landscape is very interesting. My concerns are mainly in practicality of the approach.

Response

Thank you for your comments and providing runtime results and HMC references. I would still really love to see a quantitative non-NN comparison to the regular Laplace approximation. I unfortunately do not have a specific one in mind to suggest, but would have a look through the tutorials/examples of a couple of probabilistic programming frameworks that focus on MCMC for inference. I'm sure they will have comparisons where sampling works much better than VI with a Gaussian/Laplace, and it would be interesting to see to what extent using your Riemannian approach to adapt to the posterior closes the gap (assuming it does). I would also be curious how the method compares e.g. to normalizing flows in such a lower dimensional case. I understand the temptation of wanting to do neural nets first and foremost, but I think there is a really clear path for potential applications with more traditional probabilistic models, whereas BNNs will require more work on scalability. Both are interesting in terms of research of course, but for the impact of the paper it would, at least in my opinion, make a lot of sense to cover the former empirically. Overall and in light of the other reviews with there being a consensus for acceptance, I remain with my score.

Response to the Authors

I would like to thank the authors for the efforts. I have read other reviews as well as related responses. I stand by the current score -- I think clearly analyzing the computational complexity vs empirical gains in performance (my third point) is one missing point in the paper. This is also connected to some of the raised concerns by other reviewers. While authors discuss the computational complexity, it might help to include empirical results, especially for all the comparison results with standard laplace approximation. If the paper gets accepted, I also hope to see more baselines like MC dropout and deep ensembles.