All You Need is a Good Functional Prior for Bayesian Deep Learning

A prior over neural-network weights induces a distribution over functions. Choosing the variance of each weight offers indirect control over the behavior of those functions. We instead specify a function prior with a Gaussian process and fit a weight prior whose sampled functions resemble it.

Specifying a prior over functions

A Gaussian-process prior expresses assumptions through a mean function and a covariance kernel. These objects describe how outputs vary across inputs, making properties such as correlation and smoothness accessible at the level of the prediction problem.

A Bayesian neural network also defines a function prior: sample weights $w$ from $p_\psi(w)$, then evaluate the resulting network $f_w$. The parameters $\psi$ describe the weight prior. Our goal is to choose $\psi$ so that this induced function distribution matches the chosen GP prior.

Matching the two function distributions

We compare both priors through their function values at a collection of inputs. Write $P_{\mathrm{GP}}$ for the distribution of GP evaluations and $P_{\mathrm{NN},\psi}$ for the corresponding distribution from the neural network. The prior-fitting objective is

$$\min_\psi W_1\!\left(P_{\mathrm{GP}},P_{\mathrm{NN},\psi}\right),$$

where $W_1$ is the Wasserstein-1 distance. Its dual formulation introduces a critic that distinguishes sampled function values from the two distributions. We alternate between fitting this critic and updating the weight-prior parameters to reduce the estimated discrepancy.

The network architecture and the chosen prior family determine how closely the target behavior can be represented. The input locations used for matching determine where the two function distributions are compared.

Our Bayesian autoencoder study also selects weight priors by matching output distributions, using reference images as the target.

Posterior inference with the fitted prior

After fitting the prior, we use it in Bayesian inference with the observed training data. The likelihood updates the network's weight distribution, and predictions integrate over posterior weight samples obtained with scalable Markov chain Monte Carlo.

We compare the fitted priors with alternative weight priors and approximate Bayesian methods, including settings with convolutional neural networks. The fitted priors improve predictive performance in these comparisons. They transfer the GP's specified function behavior into a neural-network parameterization that supports posterior inference and prediction.