Reparameterization invariance in approximate Bayesian inference

2024-06-05

Neural Information Processing Systems

AI Summary

Bayesian neural networks (BNNs) suffer from parameterization-dependent approximate posteriors, violating reparameterization invariance and thereby decoupling parameter uncertainty from functional uncertainty—compromising uncertainty calibration and generalization. This work systematically investigates the issue within the linearized Laplace approximation framework. We propose the first reparameterization-invariant posterior sampling algorithm grounded in Riemannian diffusion processes; establish a geometric interpretation of linearized Laplace predictive success; and, for the first time, extend invariance guarantees from linearized predictions to predictions of the original nonlinear network. Experiments demonstrate substantial improvements in posterior approximation quality, yielding significantly more robust uncertainty calibration and generalization across multiple benchmarks.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks

Abstract

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.

Problem

• Fixes reparameterization invariance in Bayesian neural networks.
• Addresses underfitting in Laplace approximation via linearization.
• Extends invariance properties using Riemannian diffusion process.

Innovation

• Reparameterization invariance in BNNs
• Linearized Laplace approximation enhancement
• Riemannian diffusion for posterior sampling