Regularized KL-Divergence for well-defined function space variational inference in BNNs

Tristan Cinquin, Robert Bamler

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Bayesian deep learning, Bayesian neural network, Variational inference in function space, Gaussian process, Divergences

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We address the problem of infinite KL divergence in function space variational inference in BNNs, with a new method based on generalized variational inference.

Abstract: Bayesian neural networks (BNN) promise to combine the predictive performance of neural networks with principled uncertainty modeling important for safety-critical applications and decision making. However, uncertainty estimates depend on the choice of prior, and finding informative priors in weight space has proven difficult. This has motivated variational inference (VI) methods that pose priors directly on the function generated by the BNN rather than on weights. In this paper, we point out that function space VI is ill-posed since the standard objective function (ELBO) is negative infinite for many interesting priors (Burt et al., 2020). Therefore, we argue that a more general framework is needed to perform function space inference in BNNs, and we propose a simple solution using generalized variational inference (Knoblauch et al., 2019) with the regularized KL divergence (Quang, 2019). Experiments show that our inference method accurately approximates the true Gaussian process posterior on synthetic and small real-world data sets, and provides competitive uncertainty estimates for regression and out-of-distribution detection compared to BNN baselines with both function and weight space priors.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 2545