experiments.tex (12830B)
1 \setchapterpreamble[u]{\margintoc} 2 \chapter{Experiments} 3 \labch{experiments} 4 5 In \cref{ch:ablation-experiments} we tested whether the approximate sampling scheme is able to produce correct posterior samples on an arbitrary network. 6 In this chapter, we use our custom bound on the marginal likelihood to train a Bayesian neural network on the marginal via online Laplace. 7 We then use our approximate sampling scheme to sample from the posterior and propagate the samples through the network to produce a predictive distribution. 8 In this way, we test the entire Bayesian neural network pipeline, from training to inference. 9 10 \section{The Sine Function} 11 12 We first use a simple sine function as a toy example to demonstrate the Bayesian neural network pipeline for regression. 13 First, we must train the network on the marginal likelihood to demonstrate that the model can learn a good posterior distribution. 14 We then sample from the posterior and visualise the posterior predictive samples to showcase the quality of the predictive distribution. 15 16 \begin{marginfigure} 17 \centering 18 \includegraphics{sin_posterior_predictive_bothalphas_clear.pdf} 19 \caption[Laplace predictive posterior for a sine curve.]{Laplace predictive posterior (with a confidence interval given by one standard deviation) for a multilayer perceptron modelling \(f(x) = \sin(5 x + 1)\) with parameters trained on the marginal likelihood for a fixed prior precision \(\alpha = 70\) (top) and \(\alpha = 1\) (bottom). We compute the marginal using our approximate lower bound and sample from the posterior using our approximate sampling scheme. The choice of prior precision has a large effect on the posterior uncertainty, but, when it is chosen appropriately, our approximate methods are able to learn a reasonable predictive distribution (i.e., more uncertainty in the regions with no data) with a similar mean prediction to the MSE-optimised network.} 20 \label{fig:sin} 21 \end{marginfigure} 22 23 In \cref{fig:sin} we show the posterior predictive distribution for a fully connected neural network modelling a sine curve \(f(x) = \sin(5 x + 1)\) and compare it to a neural network trained on the MSE loss, i.e., maximum likelihood for regression. 24 We vary the prior precision \(\alpha\) to show how the posterior predictive distribution changes. 25 We have the same prior precision for all weights and biases, with values \(\alpha = 70\) for the top figure and \(\alpha = 1\) for the bottom figure. 26 \cref{fig:sin-posterior-predictive-samples} in the appendix shows the posterior predictive samples for this experiment. 27 28 We can see from \cref{fig:sin} that for \(\alpha = 70\) the posterior predictive distribution has a low uncertainty in the regions with data, but a high uncertainty in the regions without data. 29 Furthermore, for this value of \(\alpha\) the mean of the posterior predictive distribution is relatively close to the mean of the MSE-optimised network. 30 31 For a very low prior precision \(\alpha = 1\), the posterior predictive distribution has a high uncertainty throughout the whole function domain. 32 This is because the prior uncertainty is very high, so the posterior uncertainty is also high. 33 Furthermore, because of the noise in the posterior samples, the mean of the posterior predictive distribution is not close to that of the MSE-optimised network. 34 35 We can thus see that maximising the marginal likelihood is able to learn a reasonable predictive distribution (i.e., more uncertainty in the regions with no data) with a similar mean prediction to the MSE-optimised network, when the prior precision is chosen appropriately. 36 However, if the prior precision is too low, the posterior predictive distribution will have an extremely high uncertainty and the mean of the posterior predictive distribution will not be close to that of the MSE-optimised network. 37 In these experiments, we have chosen the prior precision to be fixed, but in practice, we would choose the prior precision by maximising the marginal likelihood or by cross-validation. 38 39 \begin{marginfigure} 40 \centering 41 \includegraphics{sin_training_alpha_small.pdf} 42 \caption[Training a sine curve on the marginal.]{Loss curves (top) from training model parameters and prior precision with initialisation \(\alpha = 70\) (bottom) on the negative log-marginal likelihood loss \(- \log p(\bm y)\) with a 10 epoch warmup period of training on MSE loss. Loss curves show the individual terms that make up the negative log-marginal likelihood loss, as per \cref{eq:laplace-negative-log-marginal}. Negative log-marginal likelihood (in green) and negative log-likelihood (in orange) are both strongly smoothed by a moving average. All four terms can be seen to decrease during training, with the negative log-likelihood term dominating the decrease. We see that the prior precision \(\alpha\) decreases linearly during training.} 43 \label{fig:sin-training} 44 \end{marginfigure} 45 46 \cref{fig:sin-training} show the loss curves from training the model parameters and prior precision with initialisation \(\alpha = 70\) on the negative log-marginal likelihood loss \(- \log p(\bm y)\) for our sine function. 47 First, we train the model parameters on the MSE loss for 10 epochs, and then start training on the marginal. 48 As per \cref{eq:laplace-log-marginal}, the Laplace-approximate negative log-marginal likelihood loss during training can be decomposed as 49 \begin{align}\label{eq:laplace-negative-log-marginal} 50 -\log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & -\log p(\bm y \given \bm \theta) - \log p(\bm \theta) + \frac{1}{2} \log \det \bm \Lambda + \mathrm{const}. 51 \end{align} 52 These terms are shown in \cref{fig:sin-training} (top), where the negative log-likelihood \(-\log p(\bm y \given \bm \theta)\) and negative log-marginal \(- \log p(\bm y)\) terms are both strongly smoothed by a moving average. 53 We can see that all three terms (and therefore also the negative log-marginal loss itself) decrease during training. 54 The negative log-likelihood term dominates the overall decrease in the negative log-marginal. 55 Since we are optimising the prior precision hyperparameter \(\alpha\) with the marginal likelihood, we also plot its variation throughout training in \cref{fig:sin-training} (bottom). 56 We see that \(\alpha\) decreases linearly during training. 57 In practice, if we keep training, we find that \(\alpha\) reaches a minimum value below 1. 58 As we showed in \cref{fig:sin}, this is a poor prior precision, since it leads to an unreasonably high posterior predictive uncertainty. 59 However, training the hyperparameters on the marginal likelihood is a sensitive procedure, and so we have not yet managed to successfully optimise the prior precision. 60 Because of this, we choose a fixed prior precision in our other experiments. 61 62 In \cref{fig:chisq-sin} we show the chi-squared sample plot for the posterior samples on this sine regression task to demonstrate that the posterior samples are multivariate normal with the correct covariance, as we explain in \cref{sec:sampling-evaluation}. 63 We choose a low value for the prior precision of \(\alpha = 1\) to demonstrate that our sampling works even when the condition number of the posterior precision is high. 64 \cref{fig:chisq-sin} shows that the sample square Mahalanobis distances for these experiments appear chi-squared distributed, so the samples themselves are multivariate normal with the Laplace posterior precision. 65 We deliberately choose an extremely large number of posterior samples (1000) to demonstrate the accuracy of the posterior samples. 66 However, it would not be practical to compute this many posterior samples for larger problems, since we would have to either store them all in memory or compute them in batches. 67 68 \begin{figure}[b] 69 \centering 70 \includegraphics{chisq_sin_lower.pdf} 71 \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for visualising normality of 1000 posterior samples from the full Laplace approximation for a fully connected neural network modelling \(f(x) = \sin(5 x + 1)\), using CIQ. Since the sample square Mahalanobis distances appear chi-squared distributed, the samples themselves are multivariate normal with the correct (known) covariance.} 72 \label{fig:chisq-sin} 73 \end{figure} 74 75 \section{MNIST} 76 77 We now attempt to use our approximate sampling method on a more realistic problem, namely, the classification of handwritten digits from the MNIST dataset. 78 To demonstrate that our approximate sampling method can be used on a more complex network, we train a convolutional neural network on MNIST to compute the Laplace approximation. 79 Currently, the inclusion of convolutions in the network causes a major slowdown in our approximate log-marginal training procedure, so we maximise the posterior during training (i.e., performing post-hoc Laplace) instead of the marginal likelihood. 80 Because of this, these experiments have been performed with a fixed posterior precision \(\alpha\) which may not be the best choice for the posterior precision. 81 82 Unfortunately, evaluating the quality of the posterior predictive distribution for MNIST is not as straightforward as for the sine curve. 83 The reason for this is that we cannot easily visually determine the appropriate magnitude of the uncertainty in the posterior predictive distribution for MNIST (since the output is a 10-dimensional vector). 84 We instead only evaluate the quality of the posterior distribution by comparing the posterior samples to the theoretical posterior distribution, as we did for the sine function in \cref{fig:chisq-sin}. 85 For the sine function, we sampled 1000 posterior samples. 86 Since, in practice, we would not be able to store this many posterior samples for large-scale problems, we instead sample 200 parameter vectors from the posterior for MNIST, to demonstrate what these plots look like for a more realistic number of samples. 87 In reality, for large-scale problems, we would not be able to store even 200 posterior samples, but between 5 and 20. 88 We can see these results with \(\alpha = 0.1\) in \cref{fig:chisq-mnist}. 89 90 \cref{fig:chisq-mnist} shows that sampling succeeds for values of \(\alpha\) as low as 0.1. 91 This value of the prior precision is much lower than the value of \(\alpha\) we would likely use in practice. 92 Additionally, as we noted in \cref{sec:spectrum-ggn}, since \(\alpha\) is the lowest eigenvalue of the posterior precision, the condition number is inversely proportional to \(\alpha\). 93 Because of this, by sampling with a low value of \(\alpha\), we are sampling from the most ill-conditioned posterior precision we can expect to encounter. 94 As such, this experiment allows us to conclude that our approximate sampling method would successfully be able to sample from the posterior distribution for a wide range of values of \(\alpha\) on MNIST without being too ill-conditioned for our method. 95 96 \begin{figure} 97 \centering 98 \includegraphics{figures/chisq_mnist_lower.pdf} 99 \caption[Chi-squared sample plot for MNIST.]{Chi-squared sample plot for visualising normality of 200 posterior samples from the full Laplace approximation for a convolutional neural network trained on MNIST with prior precision \(\alpha = 0.1\), using CIQ. Though the sample square Mahalanobis distances do not appear \emph{exactly} chi-squared distributed, this is because of the relatively low sample size used to compute the empirical histogram and quantiles. Thus, the plots themselves strongly suggest that the samples are multivariate normal with the correct (known) covariance.} 100 \label{fig:chisq-mnist} 101 \end{figure} 102 103 \section{Discussion} 104 105 The choice of prior precision \(\alpha\) is non-trivial. 106 Furthermore, since the prior precision is the lowest eigenvalue of the posterior precision matrix, we must choose \(\alpha\) to be large enough to ensure that the posterior precision matrix is positive definite (i.e., the prior precision must counteract small numerical issues which may lead to some negative eigenvalues). 107 Additionally, the condition number of the posterior precision matrix is proportional to \(\alpha^{-1}\), so \(\alpha\) should be large enough to ensure that the posterior precision matrix is well-conditioned. 108 While we can ``choose'' \(\alpha\) to be a value large enough to reduce conditioning issues, this somewhat invalidates the Bayesian interpretation of the prior, which should theoretically be chosen based on prior knowledge of the problem, can also be learned from the data. 109 Choosing the prior in order to reduce ill-conditioning of our posterior precision means we sacrifice some of the benefits of the prior. 110 As such, we would like to learn \(\alpha\) from the data by maximising the marginal likelihood. 111 Current issues with training prevent us from reliably learning \(\alpha\) from the data across models of different sizes.