Introduction to Posterior Inference

We should keep in mind that for all measurements, models, and discretizations we have uncertainties. For measurements and observations, this typically appears in the form of measurement errors. Model equations equations, on the other hand, usually encompass only parts of a system we're interested in (leaving the remainder as an uncertainty), while for numerical simulations we inherently introduce discretization errors. So a very important question to ask here is how we can be sure that an answer we obtain is the correct one. From a statisticians viewpoint, we'd like to know the posterior probability distribution, a distribution that captures possible uncertainties we have about our model or data.

Uncertainty

This admittedly becomes even more difficult in the context of machine learning: we're typically facing the task of approximating complex and unknown functions. From a probabilistic perspective, the standard process of training an NN here yields a maximum likelihood estimation (MLE) for the parameters of the network. However, this MLE viewpoint does not take any of the uncertainties mentioned above into account: for DL training, we likewise have a numerical optimization, and hence an inherent approximation error and uncertainty regarding the learned representation. Ideally, we should reformulate the learning process such that it takes its own uncertainties into account, and it should make posterior inference possible, i.e. learn to produce the full output distribution. However, this turns out to be an extremely difficult task.

This is where so-called Bayesian neural network (BNN) approaches come into play. They allow for a form of posterior inference by making assumptions about the probability distributions of individual parameters of the network. This gives a distribution for the parameters, with which we can evaluate the network multiple times to obtain different versions of the output, and in this way sample the distribution of the output.

Nonetheless, the task remains very challenging. Training a BNN is typically significantly more difficult than training a regular NN. This should come as no surprise, as we're trying to learn something fundamentally different here: a full probability distribution instead of a point estimate. (All previous chapters "just" dealt with learning such point estimates, and the tasks were still far from trivial.)

Introduction to Bayesian Neural Networks

In order to combine posterior inference with Neural Networks, we can use standard techniques from Bayesian Modeling and combine them with the Deep Learning machinery. In Bayesian Modeling, we aim at learning distributions over the model parameters instead of these fixed point estimates. In the case of NNs, the model parameters are the weights and biases, summarized by $\theta$, for a neural network $f$. Our goal is therefore to learn a so-called posterior distribution $p(\theta |D)$ from the data $D$, which captures the uncertainty we have about the networks weights and biases after observing the data $D$. This posterior distribution is the central quantity of interest here: if we can estimate it reasonably well, we can use it to make good predictions, but also assess uncertainties related to those predictions. For both objectives it is necessary to marginalize over the posterior distribution, i.e. integrate it out. A single prediction for input $x_i$ can for example be obtained via

$$\widehat{y_i}=\int f(x_i;\theta )p(\theta |D)d\theta$$

Similarly, one can for instance compute the standard deviation in order to assess a measure of uncertainty over the prediction $\widehat{y_i}$.

Prior distributions

In order to obtain the required posterior distribution, in Bayesian modeling one first has to define a prior distribution $p(\theta )$ over the network parameters. This prior distribution should encompass the knowledge we have about the network weights before training the model. We for instance know that the weights of neural networks are usually rather small and can be both positive and negative. Centered normal distributions with some small variance parameter are therefore a standard choice. For computational simplicity, they are typically also assumed to be independent from another. When observing data $D$, the prior distribution over the weights is updated to the posterior according to Bayes rule:

$$p(\theta |D)=\frac{p(D|\theta )p(\theta )}{p(D)}\text{ . }$$

This is, we update our a-priori knowledge after observing data, i.e. we learn from data. The computation required for the Bayesian update is usually intractable, especially when dealing with non-trivial network architectures. Therefore, the posterior $p(\theta |D)$ is approximated with an easy-to-evaluate variational distribution $q_{\varphi }(\theta )$, parametrized by $\varphi$. Again, independent Normal distributions are typically used for each weight. Hence, the parameters $\varphi$ contain all mean and variance parameters $\mu ,\sigma$ of those normal distributions. The optimization goal is then to find a distribution $q_{\varphi }(\theta )$ that is close to the true posterior. One way of assessing this closeness is the KL-divergence, a method used widely in practice for measuring the similarity of two distributions.

Evidence lower bound

We cannot directly minimize the KL-divergence between approximate and true posterior $KL(q_{\varphi }(\theta )||p(\theta |D)$, because we do not have access to the true posterior distribution. It is however possible to show that one can equivalently maximize the so-called evidence lower bound (ELBO), a quantity well known from variational inference:

$$\mathcal{L}(\varphi )=E_{q_{\varphi }}[\log (p(D|\theta ))]-KL(q_{\varphi }(\theta )||p(\theta ))\text{ , }$$

The ELBO (or negative ELBO, if one prefers to minimize instead of maximize) is the optimization objective for BNNs. The first term is an expected log-likelihood of the data. Maximizing it means explaining the data as well as possible. In practice, the log-likelihood is often a conventional loss functions like mean squared error (note that MSE can be seen as negative log-likelihood for normal noise with unit variance). The second term is the negative KL-divergence between the approximate posterior and the prior. For suitable prior and approximate posterior choices (like the ones above), this term can be computed analytically. Maximizing it means encouraging the approximate network weight distributions to stay close to the prior distribution. In that sense, the two terms of the ELBO have opposite goals: The first term encourages the model to explain the data as well as possible, whereas the second term encourages the model to stay close to the (random) prior distributions, which implies randomness and regularization.

The expectation of the log-likelihood is typically not available in analytical form, but can be approximated in several ways. One can, for instance, use Monte-Carlo sampling and draw $S$ samples from $q_{\varphi }(\theta )$. The expectation is then approximated via $\frac{1}{S}{\sum }_{s=1}^S\log (p(D|{\theta }_s)$. In practice, even a single sample, i.e. $S=1$ can be enough. Furthermore, the expectation of the log-likelihood is typically not evaluated on the whole data set, but approximated by a batch of data $D_{batch}$, which enables the use of batch-wise stochastic gradient descent.

The BNN training loop

For $S=1$, one iteration of the training loop then boils down to

1. sampling a batch of data $D_{batch}$
2. sampling network weights and biases from $q_{\varphi }(\theta )$
3. forwarding the batch of data through the network according to the sampled weights
4. evaluating the ELBO
5. backpropagating the ELBO through the network and updating $\varphi$

If $S>1$, steps 2 and 3 have to be repeated $S$ times in order to compute the ELBO. In that sense, training a variational BNN is fairly similar to training a conventional NN: We still use SGD and forward-backward passes to optimize our loss function. Only that now we are optimizing over distributions instead of single values. If you are curious on how one can backpropagate through distributions, you can, e.g., read about it here, and a more detailed introduction to Bayesian Neural Networks is available in Y. Gals thesis, in chapters 2 & 3.

Dropout as alternative

Previous work has shown that using dropout is mathematically equivalent to an approximation to the probabilistic deep Gaussian process. Furthermore, for a specific prior (satisfying the so-called KL-condition) and specific approximate posterior choice (a product of Bernoulli distributions), training a neural network with dropout and L2 regularization and training a variational BNN results in an equivalent optimization procedure. In other words, dropout neural networks are a form of Bayesian neural networks.

Obtaining uncertainty estimates is then as easy as training a conventional neural network with dropout and extending dropout, which traditionally has only been used during the training phase, to the prediction phase. This will lead to the predictions being random (because dropout will randomly drop some of the activations), which allows us to compute average and standard deviation statistics for single samples from the data set, just like for the variational BNN case. It is an ongoing discussion in the field whether variational or dropout-based methods are preferable.

A practical example

As a first real example for posterior inference with variational BNNs, let's revisit the case of turbulent flows around airfoils, from Supervised training for RANS flows around airfoils. However, in contrast to the point estimate learned in this section, we'll now aim for learning the full posterior.