Tractable Function-Space Variational Inference in Bayesian Neural Networks

TL;DR

This paper introduces a scalable function-space variational inference method for Bayesian neural networks that improves uncertainty estimation and predictive performance, especially in safety-critical applications.

Contribution

It proposes a novel function-space variational inference approach that directly models the posterior over functions, enabling better prior incorporation and uncertainty quantification.

Findings

Achieves state-of-the-art uncertainty estimation and predictive accuracy.

Performs well on safety-critical medical diagnosis tasks.

Demonstrates scalability and effectiveness across various prediction tasks.

Abstract

Reliable predictive uncertainty estimation plays an important role in enabling the deployment of neural networks to safety-critical settings. A popular approach for estimating the predictive uncertainty of neural networks is to define a prior distribution over the network parameters, infer an approximate posterior distribution, and use it to make stochastic predictions. However, explicit inference over neural network parameters makes it difficult to incorporate meaningful prior information about the data-generating process into the model. In this paper, we pursue an alternative approach. Recognizing that the primary object of interest in most settings is the distribution over functions induced by the posterior distribution over neural network parameters, we frame Bayesian inference in neural networks explicitly as inferring a posterior distribution over functions and propose a scalable function-space variational inference method that allows incorporating prior information and results in reliable predictive uncertainty estimates. We show that the proposed method leads to state-of-the-art uncertainty estimation and predictive performance on a range of prediction tasks and demonstrate that it performs well on a challenging safety-critical medical diagnosis task in which reliable uncertainty estimation is essential.