Bayesian Neural Networks

Adapted from an appendix of my MS thesis.

Bayesian Neural Network

Deep neural networks (DNNs) are usually trained using a regularized maximum likelihood objective to find a single setting of parameters. However, large flexible models like neural networks can represent many functions, corresponding to different parameter settings, which fit the training data well, yet generalize in different ways. This phenomenon is known as underspecification [1].

Considering all of these different models together can lead to improved accuracy and uncertainty representation. This can be done by computing the posterior predictive distribution using Bayesian model averaging. The main challenges in applying Bayesian inference to DNNs are specifying suitable priors, and efficiently computing the posterior, which is challenging due to the large number of parameters and the large datasets [1].

Consider the generalized MLP with

$L-1$

hidden layers and a linear output as described in the companion MLP post [1].

The most common choice of priors is to use a factored Gaussian prior [1].

Initializing this model’s parameters at a particular random value is like sampling a point from this prior over functions. In the limit of infinitely wide neural networks, a neural network defines a Gaussian process (see the companion Gaussian process series) with a fixed kernel. Indeed, an MLP with one hidden layer, whose width goes to infinity, and which has a Gaussian prior on all the parameters, converges to a Gaussian process with a well-defined kernel. These kernels can be used in modeling non-stationary covariance structure [1].

Monte Carlo dropout (MCD) is a very simple and widely used method for approximating the Bayesian predictive distribution. Usually stochastic dropout layers are added as a form of regularization during training, and are turned off at test time. However, the idea of MCD is to also perform random sampling during testing [1].

More precisely, we drop out each hidden unit according to a

$\mathrm{Bernoulli}(p)$

distribution, and repeat this procedure

$K$

times to create

$K$

distinct models. We then create an equally weighted average of the predictive distribution for each of these models as shown below. One drawback of MCD is that it is slow at test time. However, this can be overcome by distilling the model’s predictions into a deterministic student network [1].

Another very simply approximation to the posterior is to only be Bayesian about the weights in the final layer. This is called the “Bayesian last layer” approximation. In more detail, let

$\mathbf{z}=f(\mathbf{x};\mathbf{\theta })$

be the predicted outputs of the model before any optional final nonlinearity. We assume this has the form

z=wL⊤​ϕ(x;θ)

, where

$\varphi (\mathbf{x})$

are the features extracted by the first

$L-1$

layers. Then, we can use standard techniques, such as Markov chain Monte Carlo (described in the companion MCMC series), to compute

p(wL​∣D)=N(μL​,ΣL​)

, given

$\varphi (\cdot )$

[1].

MCMC methods like Hamiltonian Monte Carlo (HMC) are generally considered to be the gold standard for posterior approximation, since they do not make strong assumptions about the form of the posterior. However, a significant limitation of standard MCMC procedures, including HMC, is that they require access to the full training set at each step. Stochastic gradient MCMC methods, such as Stochastic gradient Langevin dynamics (SGLD), operate instead using mini-batches of data, offering a scalable alternative [1].

Many conventional approximate inference methods, such as variational inference, focus on approximating the posterior

$p(\mathbf{\theta }|\mathcal{D})$

in a local neighborhood around one of the posterior modes. While this is often not a major limitation in classical machine learning, modern deep neural networks have highly multi-modal posteriors, with parameters in different modes giving rise to very different functions [1].

On the other hand, the functions in a neighborhood of a single mode may make fairly similar predictions. So using a local approximation to compute the posterior predictive will underestimate uncertainty and generalize more poorly. A simple alternative method is to train multiple models, and then to approximate the posterior using an equally weighted mixture of delta functions as follows where

$M$

is the number of models. This approach is called deep ensembles [1].

Once we have an approximation

$q(\mathbf{\theta }|\mathcal{D})$

of the parameter posterior

$p(\mathbf{\theta }|\mathcal{D})$

, we can use it to approximate the posterior predictive distribution [1].

We often approximate this integral using Monte Carlo where

θs∼q(θ∣D)

[1].

References

1. Kevin P. Murphy (2023) Probabilistic Machine Learning: Advanced Topics. MIT Press.