Reparameterization invariance in approximate Bayesian inference

Current approximate posteriors in Bayesian neural networks (BNNs) exhibit a crucial limitation: they fail to maintain invariance under reparameterization, i.e. BNNs assign different posterior densities to different parametrizations of identical functions. This creates a fundamental flaw in the application of Bayesian principles as it breaks the correspondence between uncertainty over the parameters with uncertainty over the parametrized function. In this paper, we investigate this issue in the context of the increasingly popular linearized Laplace approximation. Specifically, it has been observed that linearized predictives alleviate the common underfitting problems of the Laplace approximation. We develop a new geometric view of reparametrizations from which we explain the success of linearization. Moreover, we demonstrate that these reparameterization invariance properties can be extended to the original neural network predictive using a Riemannian diffusion process giving a straightforward algorithm for approximate posterior sampling, which empirically improves posterior fit.

Paper

Similar papers

Peer review

Reviewer v5yb8/10 · confidence 4/52024-07-03

Summary

The paper studies the effects of the invariance of Bayesian neural networks under reparametrization on their approximate posterior inference. The primary effect is the ambiguity on the uncertainty of the inferred function as it can be represented by multiple reparameterizations, each of which may be assigned a different uncertainty estimate. The paper takes Laplace approximation as a case and studies its properties when used for Bayesian neural net inference using linear algebraic and differential geometric tools. The paper derives an algorithm from its key theoretical outcomes and shows it to outperform its counterparts in various downstream uncertainty quantification tasks.

Strengths

* This is a very high-calibre paper, a complete piece of work, that studies an original and significant problem with impressive theoretical rigor. * The presentation is extremely clear. Figures 1 and 2 are truly helpful to the reader to grasp the visual intuition behind the studied problem and the proposed solution. * The compilation of theoretical tools such as the a parameter manifold on the Generalized Gauss-Newton (GGN) approximation and a diffusion defined on it are original, advanced, and elegant. * The reported numerical results are both comprehensive and strongly in favor of the central claims.

Weaknesses

* Overall this is excellent piece of work with a mature write-up. Just its presentation may be improved tiny bit. For instance, certain terms are used very early on in the paper such as the introduction and Figure 1 in their meanings a bit unusual for the broad ML research readership, such as a "kernel" and "diffusion". While their allocation in the particular context is appropriate, reading experience may be improved if their meanings are made more explicit, at least verbally, in their first use.

Questions

* How does the argumentation in the first paragraph of Section 3 come together with the main message of the paper? Is the problem not instead that the posterior assigns separate masses to different weights that represent the same functional form? How exactly does this produce a pathology in the posterior besides the need to detect multiple equivalent representations and add their masses? Why do we actually need to design the related degrees of freedom? Overall I buy the argument and see the problem, but my point is that probably this bridge paragraph may be misleading the reader. * What does f_lin^w precisely mean in Eq 4? Is it the first-order Taylor approximation of f?

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The paper briefly discusses its limitations in the final paragraph of the Conclusion section. The potential negative societal impacts of the studied work are not discussed probably because they are thought to be inapplicable.

Reviewer gWds5/10 · confidence 3/52024-07-11

Summary

The paper addresses the problem of maintaining invariance under reparameterization of Bayesian Neural Networks (BNNs). This issue undermines the reliability of BNNs since different parameterizations of identical functions produce different posterior densities. Through theoretical analysis and empirical validation, the authors demonstrate the effectiveness of their approach by extending a geometric view (in Linearized Laplace) and using a Riemannian diffusion process to achieve reparameterization invariance.

Strengths

* Proposed Riemannian diffusion process incorporates the properties of the linearized Laplace approximation into the original neural network prediction, resulting in a better posterior fit. * The method leverages advanced mathematical tools from differential geometry, offering a sophisticated approach to maintaining reparameterization invariance. * In both in-distribution and out-of-distribution scenarios, the proposed method consistently outperforms other approaches. * The paper is well-organized and well-written, with each section building logically on its predecessor.

Weaknesses

* Although the proposed method is theoretically sound, its computational complexity is higher than simpler approximations. This could limit its practicality for very large-scale neural networks. * 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. * 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.

Questions

* 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 in Weakness section, I would recommend the author to evaluate the Riemannian Diffusion in other domain dataset & architectures (e.g., Transformers).

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are mentioned in "Weakness" and "Questions".

Reviewer gWds2024-08-13

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

Authorsrebuttal2024-08-13

> 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]) This paper (already referenced by us) does not compare Linearized and Sampled Laplace. Further, note that they make no theoretical links between linearization and reparametrization invariance and they study the Hessian instead of the GGN. > In addition, since the experiment I requested is applied to a very simple transformer, I don't think it is impossible. It is unclear to us what the exact experiment you suggest. Previously we interpreted your review to mean a standard transformer in an NLP task. We reemphasize that this is not a usual benchmark for evaluating Bayesian methods especially Laplace Approximations where standard regression and image classification tasks are far more popular. Hence it is not obvious what priors are appropriate, what dataset should be used, and what metrics should be reported. These are important questions and we do not aim to settle them here; Tristan et al. discuss some of the pitfalls and difficulties of such experiments. If the reviewer means a Vision Transformer (smaller version with only a few million parameters) on the standard image classification task, this is very much within reach. In fact, we even have code to run such experiments. However, doing so would require us to make some large approximations (limited Lanczos iterations and diffusion steps). The resulting experiment would, thus, not serve for theory validation, but we are happy to add them nonetheless if you see value therein. Cinquin, Tristan, et al. "Pathologies in priors and inference for Bayesian transformers." arXiv preprint arXiv:2110.04020 (2021). (edited)

Reviewer 12ke7/10 · confidence 3/52024-07-14

Summary

The authors provide theoretical analysis of the underfitting problem of Laplace approximation. Specifically, they show that the underfitting of Laplace approximation is due to the approximate posterior covariance is not invariant under reparameterization. Moreover, they propose a reparametrization invariant diffusion posterior to address undefitting. The method are evaluated over standard image classification benchmarks. The quality of the uncertainty for both in-distribution and out-of-distribution tasks are considered.

Strengths

1. 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. 2. Given the theoretical analysis, the idea of reparametrization invariant posterior is natural and proves to be effective empirically. 3. The paper is easy to follow and well organized.

Weaknesses

1. The proposed method, as mentioned by the authors as well, sufferes from more expensive computation. It would be good to report the time metric for readers to properly assess the practicability of the method. 2. The datasets considered are a bit outdated and the networks considred seem to be quite small, such that overall the performance is on the lower end (e.g. <90% acc. for CIFAR10). Recent BDL papers typically consider larger datasets (such as ImageNet) and deeper networks (e.g. [1]). Furthermore, last-layer Laplace seems to be competitive with or even outperform the proposed method in some experiments. 3. The method is tailored to improving Laplace approximation, and there is no discussion related to any other approximate inference techiques (e.g. variational inference). References [1] Antoran et al. Sampling-based inference for large linear models with application to linearised Laplace. ICLR 2023.

Questions

See weakness.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors do not address potential negative social impact since the paper is predominantly theoretical.

Reviewer AZGF5/10 · confidence 3/52024-07-14

Summary

This work questions why the linearized Laplace approximation can be effective for uncertainty estimation, whereas the vanilla Laplace approximation can lead to poor performance, such as underfitting. This work argues that this degradation is attributed to the fact that the random function derived from each parameterization of a DNN may not form an invariant function space. Specifically, since the null space $ \text{ker} (\text{GGN}_{w}) $ is not empty for over-parameterized NNs, sampling the random function containing elements of null space could result in an inconsistent random function. This can lead to under-fitting issues as observed in Laplace approximation. Therefore, the authors believe that sampling the random function on $ \text{im} (\text{GGN}_{w}) $, i.e., the linearized Laplace, could be more effective than the Laplace method. Furthermore, for a given weight $w$, they explore how to sample the same function with $ f(w) $, i.e., $ f(g(w)) = f(w) $, by sampling the weight on a specific manifold $ g(w) $. To this end, they employ the concept of quotient space to elaborately define the manifold and diffusion to sample $ w $ from the manifold $ g(w) $.

Strengths

* This work reasonably reveals why the linearized Laplace approximation can be more effective than the standard Laplace approximation. This explanation appears to be novel. * This work attempts to justify the above reasoning using a mathematical elaboration framework.

Weaknesses

* The lack of explanation of background knowledge, such as quotient space and Riemannian manifold, makes the paper difficult to understand. * 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.

Questions

* To confirm my understanding, does the Laplace diffusion denote that for a given $w$, the sample functions $ { w_t }$ are obtained by the update rule using $\text{GGN}^{+}$ as described in the SDE on the manifold? * Does $\text{GGN}^{+}$ mean the positive definite matrix of $\text{GGN}$, which is obtained by applying SVD on $\text{GGN}$ and then using the eigenvectors with positive eigenvalues? * Is the Laplace diffusion the post-hoc method, meaning that $w$ is first obtained by MAP inference, i.e., $w_{\text{MAP}} = R^{f}_{\mathcal{X}}(w)$, and then $w_t$ is obtained by Laplace diffusion with $\text{GGN}^{+}$? * 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 $w$, 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 $w$ in training procedure and obtains comparable performance, isn't it more important to focus on how to find good $w$ (for example, $ w_{\text{swa}} $ ) and explore the subspace of $w$? I am just curious about your opinion on this.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

See above weaknesses and limitations.

Authorsrebuttal2024-08-13

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

Reviewer v5yb2024-08-08

Keep score

Thanks for the clarification. Rephrasing the sentence in the suggested way will indeed do the trick. While I do agree with the other reviewers that the method has an computational overhead and experiments may be improved, I still think that the technical contribution is solid and valuable and the other issues are rather secondary. Hence, I tend to keep my score for now.

Reviewer AZGF2024-08-13

Thank you for your responses to my questions. 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. I am inclined to accept this work and will therefore maintain my score.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC