Closed Form Variational Objectives For Bayesian Neural Networks with a Single Hidden Layer

TL;DR

This paper derives closed-form variational objectives for Bayesian neural networks with a single hidden layer and piecewise polynomial activation functions, enabling efficient training and prediction.

Contribution

It introduces a method to compute variational lower bounds and predictive distributions in closed form for single-layer Bayesian neural networks with ReLU activations.

Findings

Closed-form variational lower bounds improve training efficiency.

Fast computation of predictive mean and variance.

Effective for softmax classification with approximate bounds.

Abstract

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.