Reparameterization invariance in approximate Bayesian inference

We thank the reviewer for the excellent review and we are very grateful for a thorough engagement with our work and for pinpointing its highlights.

We acknowledge your point (under Weaknesses) about providing a bit more explanation of unconventional terms such as kernel (that in machine learning often means something entirely different). We will take this feedback to heart and update the manuscript accordingly.

Regarding the opening paragraph of Sec. 3, then we acknowledge that, in particular, the phrasing “pathological representations” is too vague and opens the door to misinterpretations. We suggest to rephrase the sentence “However, as we have argued, traditional approximate Bayesian inference does not correctly capture this, leading to pathological representations” into “However, as we have argued, traditional approximate Bayesian inference does not correctly capture this and assigns different probability measures to identical functions.” We believe such a change clarifies the text, but are happy to take further feedback.

Yes, $f_{lin}^w$ is exactly the mentioned Taylor expansion. The equation is given in line 107 of the submitted manuscript, but it would have been better to already introduce this notation around line 64. We will make this change.

We hope this answers your questions and concerns, and are happy to discuss further. Once again, we are grateful for the favorable review that exactly captures our intention, and we hope that you will continue to fight for the paper.

We thank the reviewer for their valuable review.

• “Although the proposed method is theoretically sound, its computational complexity is higher than simpler approximations. [...]”

While this is a valid comment we believe it stems from a slight misunderstanding of the main point that is being made with our experiments.

The main goal of this paper is to understand the uncertainties obtained from Laplace approximations and more generally to understand some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.

One way to think about the experiments is as follows: Observation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. Hypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace. Theory: The theory presented in the paper provides mathematical justification for this hypothesis. Experiment: We build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper.

For this reason, we voluntarily spend the extra compute to get as close as possible to the theoretical ideal. Specifically, for a $P$ parameters network, computing $S$ samples by approximating the diffusion with $T$ discrete time steps and a low-rank $K$ approximation of the GGN requires $O(KP + SP)$ in memory and $O(K^2PST)$ in time, i.e. it is $ST$ times as expensive as linearized Laplace.

We could have made the method more scalable by fixing low values for $K$ or $T$. However, the resulting experiments can no longer validate the theory due to the brutal approximations. This is also why we only include comparison with baselines in the appendix omitting them from the main paper, where we only compare the different predictive functions.

Of course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.

• “Similar to previous works in Linearized Laplace, the need for new computational pipelines to handle the GGN matrix and its induced structure is acknowledged but not fully addressed in the paper.”

Indeed, our paper concludes with a brief discussion of this computational challenge. We intended to acknowledge the inherent difficulty and to draw more attention to this problem. That being said, the paper does contribute advances in this direction as our implementation uses efficient matrix-free access to the GGN using automatic differentiation combined with iterative solvers to obtain samples from the Laplace approximation. This is a significant improvement over existing implementations. We deemphasize these computational aspects of our work and only discuss them implicitly because our main focus is on the theoretical explanation.

• “Although the experiments are comprehensive, they demonstrate results on standard datasets and relatively small to medium-sized neural networks. Testing the method on more diverse datasets and larger-scale models would strengthen the empirical validation.”

We have chosen to replicate the key experiments from the Laplace Redux paper [1] which is an important reference for Laplace approximations. Their largest experiments are also on CIFAR-10 with ResNets so we don’t believe the scale of our experiments is too unreasonable. Furthermore, scaling to larger datasets and models necessarily requires making some fairly brutal approximations, by doing subnetwork inference, KFAC, diagonal approximations, etc. We are not at all opposed to this but hopefully, it is clear from the previous response why such approximations are inappropriate in the context of theory-validation.

• “Although mentioned in related work, I feel that this paper needs comparative experiments with methods such as Riemannian Laplace and (Bergman et al., 2024) and Connectivity Laplace (Kim et al., 2023).”

As mentioned above, our main experimental goal is to validate the theory. Further note that the public code for the first reference is orders of magnitude slower than our code and it is infeasible to run on as large models as we do. The public repository for the second reference has no documentation and no code for reproducing the experiments in their paper. Consequently, it would be highly non-trivial to run them on our benchmarks.

• “As mentioned in Weakness section, I would recommend the author to evaluate the Riemannian Diffusion in other domain dataset & architectures (e.g., Transformers).”

To our knowledge, it's not standard practice to evaluate Laplace approximations on transformers. We are only aware of [2] which evaluates Laplace approximations on transformers and they perform Laplace approximations only on the low-rank adaptation weights. Hopefully from the previous response, it is clear why such approximations are not appropriate for our experiments given our purpose.

We thank the reviewer for their efforts and we urge them to reconsider their assessment of the experiments in light of the additional context we have provided above.

[1] Daxberger, Erik, et al. "Laplace redux-effortless Bayesian deep learning." Advances in Neural Information Processing Systems 34 (2021)

[2] Yang, Adam X., et al. "Bayesian low-rank adaptation for large language models." arXiv:2308.13111 (2023)

Thanks for the detailed response from the authors.

The authors claim I misunderstood the paper, but I don't believe that's true. Linearized Laplace outperforms (sampled) Laplace, and there are already other theoretical approaches to the problem. (See Theorem 4.1 in [1])

In addition, since the experiment I requested is applied to a very simple transformer, I don't think it is impossible.

As a result, I will maintain the current score.

[1] Kim, SungYub, Kyungsu Kim, and Eunho Yang. "GEX: A flexible method for approximating influence via Geometric Ensemble." Advances in Neural Information Processing Systems 36 (2023).

Thanks for the excellent review, which we will reply to in parts.

• “The paper provides, to my knowledge, the first theoretical justification of why Laplace approximation sufferes from underfitting: the approximate posterior covariance is not reparameterization invariance.”

This is correct. Additionally, we are the first to provide a complete characterization of continuous reparameterizations of a neural network. Certain reparameterizations such as scale reparameterizations for ReLUs are already known but we provide a complete characterization of these in terms of the kernel manifold. Furthermore, to the best of our knowledge, we are the first to consider data-dependent reparameterizations rather than only global reparameterization (a subset of the former).

• “The proposed method, as mentioned by the authors as well, suffers from more expensive computation. It would be good to report the time metric for readers to properly assess the practicability of the method.”

For a $P$ parameters network, computing $S$ samples by approximating the diffusion with $T$ discrete time steps and a low-rank $K$ approximation of the GGN requires $O(KP + SP)$ in memory and $O(K^2PST)$ in time, i.e. it is $ST$ times as expensive as linearized Laplace.

However, we want to emphasize that our main goal here is to understand the uncertainties obtained from Laplace approximations and more generally to understand some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.

One way to think about the experiments is as follows:

Observation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. Hypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace. Theory: The theory presented in the paper provides mathematical justification for this hypothesis. Experiment: In addition to the theoretical justification, we build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper.

For this reason, we voluntarily pay the extra computational factor $T$ to get as close as possible to the theoretical ideal. For example, we could have made the method more scalable by fixing low values for $K$ or $T$. While such approximations might be more practical and scalable, the resulting experiments can no longer validate the theory due to the brutal approximations. This is also why we only include comparison with baselines in the appendix omitting them from the main paper, where we only compare the different predictive functions.

Of course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.

Further discussion of these issues and time complexity is available in Appendix E.1.

• “The datasets considered are a bit outdated and the networks considred seem to be quite small [...]”

For the experiments, we have chosen to replicate the key experiments from the Laplace Redux paper [1] which is an important reference for Laplace approximations. Their largest experiments are also on CIFAR-10 with ResNets so we don’t believe the scale of our experiments is too unreasonable.

Secondly, we are certainly aware of the recent works that do Laplace approximation on larger models (see references in the paper). But they all invariably rely on fairly brutal approximations such as KFAC, last-layer, or LoRa. We are not at all opposed to this but hopefully, it is clear from the previous response why such approximations are inappropriate in the context of our experiments: these aim for theory-validation rather than the usual focus on performance and scalability.

• “The method is tailored to improving Laplace approximation, and there is no discussion related to any other approximate inference techiques (e.g. variational inference).”

While it’s true that our primary focus is on Laplace approximations, we believe that the tools developed here are more generally applicable. The theory of reparameterizations can find applicability in various topics such as continual learning and in studying the loss landscape. The issues considered here are also relevant to other Bayesian methods. It's not possible to explicate this fully but we have a brief discussion of this direction in appendix E.4. Another possible direction is Variational Inference with the diffusion posterior as a variational family. Secondly, we believe that the underfitting of Sampled Laplace is an important open problem in the community, which should not be understated. For example, this has also been discussed in a recent position paper [2]

We thank the reviewer for their efforts and we urge them to reconsider their assessment of the experiments in light of the additional context we have provided above.

[1] Daxberger, Erik, et al. "Laplace redux-effortless Bayesian deep learning." Advances in Neural Information Processing Systems 34 (2021): 20089-20103.

[2] Papamarkou, Theodore, et al. "Position paper: Bayesian deep learning in the age of large-scale AI." arXiv preprint arXiv:2402.00809 (2024)

We thank the reviewer for their valuable review and questions, which will further improve the paper.

• “The lack of explanation of background knowledge, such as quotient space and Riemannian manifold, makes the paper difficult to understand.”

We acknowledge this issue, which is due to the space constraints of a conference paper. Because of the additional allowed page, we will be happy to extend the background section to cover the necessary knowledge in Riemannian geometry to benefit future readers.

• “Although it appears effective compared to Laplace and Linearized Laplace, the performance improvement seems marginal when compared to other baselines such as SWAG and Last-layer Laplace.”

While this is a valid comment we believe it stems from a slight misunderstanding of the main point of our experiments.

The main goal of this paper is to understand the uncertainties obtained from Laplace approximations and more generally to shed light on some of the problems encountered in Bayesian deep learning. More specifically decomposing the uncertainty into its two fundamental components and how they interact with the choice of the predictive (neural network or the linearized neural network). Our purpose is not to advance the state-of-the-art of Bayesian methods.

We include the baselines in the appendix and omit them from the main paper for precisely this reason. We want to validate the theory in realistic settings while staying as close as possible to the theoretical ideal (these choices are also discussed in Sec. 7 and Appendix E.1).

One way to think about the experiments is as follows: Observation: Linearized Laplace underfits less (in-distribution) than Sampled Laplace. Hypothesis: This difference is due to improper accounting of the reparameterization properties in the Sampled Laplace. Theory: The theory presented in the paper provides mathematical justification for this hypothesis. Experiment: In addition to the theoretical justification, we build a method that is as close to Sampled Laplace as possible while also accounting for the reparameterization properties of the neural network. This is the Laplace diffusion which turns out to not underfit in-distribution, hence supporting the extensive theory developed in the paper.

Of course, we do not disregard the importance of building practical, performant, and scalable Bayesian methods. In future work, we intend to build on the here-presented theoretical insights to achieve this exact goal. However, understanding why Bayesian deep learning struggles (this paper) is a prerequisite. The main takeaway of this paper is not a particular method but a general principle: Think about reparameterization issues when building Bayesian methods.

To confirm my understanding, does the Laplace diffusion denote that for a given 𝑤, the sample functions 𝑤𝑡 are obtained by the update rule using GGN+ as described in the SDE on the manifold? Diffusion can be performed by simulating the SDE in the kernel manifold and the non-kernel manifold or both (using alternating steps, for example). In Figs. 2 and 7 we perform all 3 diffusions. For the larger experiments (Sec. 7), your description is correct. We only perform the diffusion in the non-kernel manifold. The reasons for this are discussed in Section 5.

• "Does GGN+ mean the positive definite matrix of GGN, which is obtained by applying SVD on GGN [...]?"

Essentially yes. However, we use a Lanczos decomposition of the GGN instead of SVD.

Is the Laplace diffusion the post-hoc method, meaning that 𝑤 is first obtained by MAP inference, i.e., 𝑤 MAP=𝑅𝑋𝑓(𝑤), and then 𝑤𝑡 is obtained by Laplace diffusion with GGN+?

Yes, it is a post-hoc method. As stated above we want to build a method as close to Sampled Laplace as possible which additionally is reparameterization invariant. Since Sampled Laplace is typically a post-hoc method it makes sense that Laplace diffusion should be as well.

• "Compared to the performance of SWAG described in attached appendix, the performance of the Laplace diffusion does not seem significantly improved. In this context, for a given 𝑤, is it important to sample the weight parameter according to the invariant Riemannian manifold? Rather, as considering that SWAG focuses on finding good neighborhood of 𝑤 in training procedure and obtains comparable performance, isn't it more important to focus on how to find good 𝑤(for example, 𝑤 swa) and explore the subspace of 𝑤? I am just curious about your opinion on this."

This is an interesting point of discussion, which we will approach in steps. First, note that while the mode (i.e. the MAP) of the Gaussian approximate posterior is very important, the reparametrization issue we discuss is mainly related to the covariance. SWAG restricts the covariance to a subspace. SWAG involves several heuristics that are hard to characterize theoretically, so here we consider a simplified version: Fix a mode parameter $w_m$, then perform one SGD step to find a new parameter $w_1$. Repeat this (starting from $w_m$) several times. Compute the empirical covariance of $w_1,w_2,..$ and use this for the Gaussian posterior approximation. It’s important to clarify that the kernel directions (i.e. 0-eigenvalue directions of the GGN) depend on the dataset. We can, thus, consider both the full-dataset-kernel directions and the batch-kernel directions. Notably, the full-dataset-kernel is always contained in the batch-kernel. The covariance obtained by simplified-SWAG is then guaranteed to be 0 along the full-dataset-kernel direction. This shows that the simplified-SWAG approximate posterior is locally reparametrization invariant, thereby demonstrating how the developed theory can help us understand existing methods and eventually develop new ones.

We thank the reviewer for their efforts and we encourage them to reconsider their assessment of the experiments in light of the additional context we have provided above.

During the rebuttal period, I was able to resolve my concerns and gain a deeper understanding of the importance of the reparametrization invariance that this work aims to highlight.

Thanks for the support. We very much agree that reparametrization invariance is important. When models grow in size (e.g. increasing depth), we spend a greater proportion of the parameters on reparametrizations. Ten years ago, models were sufficiently small that reparametrizations could practically be ignored, but today they are causing sufficiently many problems that we are forced to better understand them. We believe that reparametrizations are the reason why Bayesian approximations do not work well with contemporary models, even if they work for small models. This is why it's essential that we form an understanding of reparametrizations.