Closed Form Variational Objectives For Bayesian Neural Networks with a Single Hidden Layer
This work provides a method for faster training and prediction in Bayesian neural networks, but it is incremental as it focuses on a specific network architecture and activation types.
The paper tackles the problem of computing variational objectives for Bayesian neural networks with piecewise polynomial activations, showing that closed-form solutions for the variational lower bound and predictive statistics are possible for single-layer networks with Normal likelihoods and structured Normal variational distributions, and also extends to approximate bounds for other likelihoods like softmax classification.
In this note we consider setups in which variational objectives for Bayesian neural networks can be computed in closed form. In particular we focus on single-layer networks in which the activation function is piecewise polynomial (e.g. ReLU). In this case we show that for a Normal likelihood and structured Normal variational distributions one can compute a variational lower bound in closed form. In addition we compute the predictive mean and variance in closed form. Finally, we also show how to compute approximate lower bounds for other likelihoods (e.g. softmax classification). In experiments we show how the resulting variational objectives can help improve training and provide fast test time predictions.