Few-sample Variational Inference of Bayesian Neural Networks with Arbitrary Nonlinearities
This addresses the problem of efficient uncertainty estimation in BNNs for machine learning practitioners, offering a practical solution with incremental improvements in computational efficiency and flexibility.
The paper tackles the computational expense and layer restrictions in Bayesian Neural Networks (BNNs) by introducing a method for propagating statistical moments through arbitrary nonlinearities using only 3 deterministic samples, enabling few-sample variational inference without limiting network layers, and demonstrates a novel physics-informed activation function for output nodes.
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.