Few-sample Variational Inference of Bayesian Neural Networks with Arbitrary Nonlinearities

TL;DR

This paper introduces a method for efficient variational inference in Bayesian Neural Networks that uses only three samples to propagate moments through arbitrary nonlinearities, enabling flexible and physics-informed BNNs.

Contribution

We propose a simple, sample-efficient approach for moment propagation in BNNs with arbitrary nonlinearities, allowing broader network architectures and physics-informed priors.

Findings

Enables few-sample variational inference with arbitrary nonlinearities.

Introduces a new nonlinear activation function for physics-informed priors.

Demonstrates computational efficiency and flexibility in BNN inference.

Abstract

Bayesian Neural Networks (BNNs) extend traditional neural networks to provide uncertainties associated with their outputs. On the forward pass through a BNN, predictions (and their uncertainties) are made either by Monte Carlo sampling network weights from the learned posterior or by analytically propagating statistical moments through the network. Though flexible, Monte Carlo sampling is computationally expensive and can be infeasible or impractical under resource constraints or for large networks. While moment propagation can ameliorate the computational costs of BNN inference, it can be difficult or impossible for networks with arbitrary nonlinearities, thereby restricting the possible set of network layers permitted with such a scheme. In this work, we demonstrate a simple yet effective approach for propagating statistical moments through arbitrary nonlinearities with only 3 deterministic samples, enabling few-sample variational inference of BNNs without restricting the set of network layers used. Furthermore, we leverage this approach to demonstrate a novel nonlinear activation function that we use to inject physics-informed prior information into output nodes of a BNN.