commit 67c36c33f35fc14721b35c2bcd5c3cde326b40bb
parent 8762eb8db062f64465208efb4fbfb611b052f32c
Author: Silas Brack <s174433@student.dtu.dk>
Date: Wed, 8 Feb 2023 18:01:51 +0100
Pretty much everything except intro + abstract done.
Diffstat:
11 files changed, 144 insertions(+), 383 deletions(-)
diff --git a/chapters/ablation_experiments.tex b/chapters/ablation_experiments.tex
@@ -7,6 +7,8 @@
In this project, we propose performing the Laplace approximation by only storing the computational graph for the implicit Hessian-vector product, thereby avoiding the explicit computation and storage of the quadratically-scaling Hessian matrix.
How significant of an effect does this have, though?
To test this, we perform Hessian-vector products \(\bm K \bm v\) and inverse Hessian-vector products \(\bm K^{-1} \bm v\), where \(\bm K \in \reals^{N \times N}\), and compare their performance.
+Specifically, we apply a manual instantiation and inverse (or solve) of the Hessian matrix with a random vector \(\bm v\) and compare this to the performance of the JAX efficient Hessian-vector products (which we will call HVPs) and GGN-vector products (which we will call GVPs).
+In reality, JAX computes these products implicitly by performing two Jacobian-vector products (for the HVP) and one Jacobian-vector product (for the GVP).
To compare performance with regard to memory usage and runtime, we perform these products for an increasing problem dimensionality \(N \in \{10^1, 10^2, \ldots, 10^9\}\).
It is difficult to benchmark memory consumption in JAX, since the just-in-time compiler optimises computations using XLA.
@@ -36,7 +38,7 @@ We then measure the total wall-clock time for computing this Hessian product by
Furthermore, since we are looking to compute some variation of the matrix inverse of \(\bm K\), and inverting a high-dimensional matrix can also be very expensive, we also compare the performance of the \emph{inverse} Hessian-vector product \(\bm K^{-1} \bm v\).
To compute the inverse for the efficient implicit HVP and GVP functions, we use the conjugate gradient method, which only requires the use of matrix-vector multiplications to compute this inverse product.
-To compute the explicit inverse, we use the \(\texttt{scipy.linalg}\) package, which uses the LAPACK library to compute the inverse.
+To compute the explicit inverse, we use the \(\texttt{jax.scipy.linalg}\) package, which performs the inverse using LAPACK (on CPU) or cuSOLVER (on GPU).
These experiments were run on a single NVIDIA A100 GPU.
% RESULTS
@@ -77,7 +79,7 @@ However, it is not clear why this is the case.
\begin{marginfigure}
\centering
\includegraphics{ablation/mnist_ablation.pdf}
- \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST. We can see that sampling fails categorically for \(Q < 5\) and succeeds categorically for \(Q \geq 10\). SHOULD I CHANGE THRESHOLD TO 0.01 OR SO?}
+ \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST. We can see that sampling fails unambiguously for \(Q < 5\) and succeeds for \(Q \geq 10\).}
\label{fig:mnist-ablation}
\end{marginfigure}
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -24,9 +24,6 @@ By analysing quantile--quantile plots of our posterior samples and visualising o
We then benchmark the performance of Hessian-vector products and estimate that our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
This is the first method for performing the Laplace approximation while only accessing the Hessian implicitly.
-BLEH, STILL NOT TOO HAPPY WITH MY ABSTRACT.
-I FEEL LIKE I SHOULD MAKE IT MORE CONCISE AND FOCUS MORE ON THE RESULTS.
-
% Our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
% We compute the full Laplace approximation on both a sine function and MNIST.
% To do this, we demonstrate that the samples generated by our approximate sampling procedure are correctly distributed and that training the marginal likelihood using our method is feasible and produces meaningful results.
diff --git a/chapters/appendix.tex b/chapters/appendix.tex
@@ -1,7 +1,7 @@
-\chapter{Algorithms}
\setchapterstyle{lines}
+\chapter{Algorithms}
% \setchapterpreamble[u]{\margintoc}
-
+\vspace{-3em}
\begin{algorithm}[h!]
\caption{Multi-Shift Minimum Residual (msMINRES)}\label{alg:msminres}
\Input{$\bm K \succ 0$, $\bm b$, $\bm P \succ 0$, $t_1, \ldots, t_Q$, $J > 0$}
@@ -52,8 +52,8 @@
}
\end{algorithm}
-\chapter{Additional Results}
\setchapterstyle{lines}
+\chapter{Additional Results}
\begin{table*}[h]
\centering
@@ -121,3 +121,21 @@
\caption[Chi-squared samples plots for ablation on MNIST.]{Chi-squared sample plots for visualising normality of 200 posterior samples from the full Laplace approximation on MNIST for a varying number of quadrature points \(Q\). Samples appear correctly distributed for \(Q \geq 5\), incorrectly distributed for \(Q < 4\), and borderline for \(Q = 4\).}
\label{fig:ablation-chisq}
\end{figure*}
+
+\begin{figure*}
+ \begin{subfigure}[h]{0.45\textwidth}
+ \centering
+ \includegraphics[scale=0.9]{sin_posterior_predictive_samples.pdf}
+ \vspace{-0.7cm}
+ \caption{\(\alpha=70\)}
+ \end{subfigure}
+ \begin{subfigure}[h]{0.45\textwidth}
+ \centering
+ \includegraphics[scale=0.9]{sin_posterior_predictive_samples_bad.pdf}
+ \vspace{-0.7cm}
+ \caption{\(\alpha=1\)}
+ \end{subfigure}
+ \vspace{-0.3cm}
+ \caption[Laplace predictive posterior samples for a sine curve for varying prior precision.]{Laplace posterior predictive samples for a sine curve for different prior precisions \(\alpha\).}
+ \label{fig:sin-posterior-predictive-samples}
+\end{figure*}
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -28,7 +28,7 @@ Overall, while we have developed novel methods for training and inference in the
% What were our limitations?
As mentioned in the previous section, we have not yet demonstrated the practicality of our method for full Laplace approximation, particularly for large-scale problems.
-Additionally, hyperparameter selection typically requires performing a grid search over a range of hyperparameters, which is computationally expensive.
+Hyperparameter selection typically requires performing a grid search over a range of hyperparameters, which is computationally expensive.
As such, the practicality of our method depends on the ability to perform marginal training to optimise hyperparameters via backpropagation.
Additionally, marginal training is prohibitively slow for training convolutions, and we have not yet found a way to speed it up.
Furthermore, estimation of the rank of the posterior precision is important to guarantee the closeness of the upper bound on the log-determinant term in the marginal, and this is still an open problem.
@@ -38,7 +38,8 @@ As such, there is still much work to be done to understand the methods we have d
% Also diagonal raw values maybe, compare the two?
% Diagonal lets you know exactly which parameters contribute to the largest eigenvalues
-We also have not shown that the full Laplace approximation shows improved predictive performance or uncertainty calibration when compared to the diagonal or Kronecker-factored Laplace approximations.
+Since we have not yet been able to train hyperparameters on the log-marginal likelihood, we have not yet been able to perform the full Laplace pipeline, which includes training the model parameters and hyperparameters and sampling from the posterior.
+Because of this, we have not yet run experiments which attempt to demonstrate that the full Laplace approximation shows improved predictive performance or uncertainty calibration when compared to the diagonal or Kronecker-factored Laplace approximations.
Furthermore, the full Laplace approximation should be compared to other methods for Bayesian inference, such as variational approximations, deep ensembles, Monte Carlo dropout, and Markov chain Monte Carlo---both in terms of performance and runtime.
% RELATED TO ACTUAL USEFULNESS OF THE METHOD
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -20,28 +20,52 @@ We then sample from the posterior and visualise the posterior predictive samples
\label{fig:sin}
\end{marginfigure}
-In \cref{fig:sin} we show the posterior predictive distribution for a fully connected neural network modelling a sine curve and compare it to a neural network trained on the MSE loss, i.e., maximum likelihood for regression.
+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.
We vary the prior precision \(\alpha\) to show how the posterior predictive distribution changes.
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.
-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}.
+\cref{fig:sin-posterior-predictive-samples} in the appendix shows the posterior predictive samples for this experiment.
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.
Furthermore, for this value of \(\alpha\) the mean of the posterior predictive distribution is relatively close to the mean of the MSE-optimised network.
-For a very low posterior precision \(\alpha = 1\), the posterior predictive distribution has a high uncertainty throughout the whole function domain.
+For a very low prior precision \(\alpha = 1\), the posterior predictive distribution has a high uncertainty throughout the whole function domain.
This is because the prior uncertainty is very high, so the posterior uncertainty is also high.
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.
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.
-However, if the prior precision is too low, the posterior predictive distribution will have a high uncertainty throughout the whole function domain, and the mean of the posterior predictive distribution will not be close to that of the MSE-optimised network.
+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.
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.
+
+\begin{marginfigure}
+ \centering
+ \includegraphics{sin_training_alpha_small.pdf}
+ \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.}
+ \label{fig:sin-training}
+\end{marginfigure}
+
+\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.
+First, we train the model parameters on the MSE loss for 10 epochs, and then start training on the marginal.
+As per \cref{eq:laplace-log-marginal}, the Laplace-approximate negative log-marginal likelihood loss during training can be decomposed as
+\begin{align}\label{eq:laplace-negative-log-marginal}
+ -\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}.
+\end{align}
+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.
+We can see that all three terms (and therefore also the negative log-marginal loss itself) decrease during training.
+The negative log-likelihood term dominates the overall decrease in the negative log-marginal.
+Since we're optimising the prior precision hyperparameter \(\alpha\) with the marginal likelihood, we also plot its variation throughout training in \cref{fig:sin-training} (bottom).
+We see that \(\alpha\) decreases linearly during training.
+In practice, if we keep training, we find that \(\alpha\) reaches a minimum value below 1.
+As we showed in \cref{fig:sin}, this is a poor prior precision, since it leads to an unreasonably high posterior predictive uncertainty.
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.
+Because of this, we choose a fixed prior precision in our other experiments.
-Furthermore, we can see from \cref{fig:chisq-sin} that the sample square Mahalanobis distances appear chi-squared distributed, so the samples themselves are multivariate normal with the Laplace posterior precision.
+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}.
+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.
+\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.
We deliberately choose an extremely large number of posterior samples (1000) to demonstrate the accuracy of the posterior samples.
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.
-\begin{figure}
+\begin{figure}[b]
\centering
\includegraphics{chisq_sin.pdf}
\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.}
@@ -54,7 +78,6 @@ We now attempt to use our approximate sampling method on a more realistic proble
To demonstrate that our approximate sampling method can be used on a more complex network, we train a convolutional neural network on the posterior distribution to compute the post-hoc Laplace approximation.\sidenote{Currently, the inclusion of convolutions in the network causes a major slowdown in the training procedure, so we maximise the posterior during training (i.e., post-hoc Laplace) instead of the marginal likelihood. We are working on a fix.}
Unfortunately, evaluating the quality of the posterior predictive distribution for MNIST is not as straightforward as for the sine curve.
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 it is a 10-dimensional vector and the input images are \(28 \times 28\) pixels).
-Because of this, it is difficult to choose an appropriate prior precision \(\alpha\) for the problem, since we would have to either learn \(\alpha\) by maximising the marginal likelihood or choose \(\alpha\) via cross-validation.
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}, where we sampled 1000 posterior samples.
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.\sidenote{In reality, for large-scale problems, we would not be able to store even 200 posterior samples, but between 5 and 20.}
We can see these results with \(\alpha = 0.1\) in \cref{fig:chisq-mnist}.
@@ -67,21 +90,21 @@ We can see these results with \(\alpha = 0.1\) in \cref{fig:chisq-mnist}.
\end{figure}
\cref{fig:chisq-mnist} shows that sampling succeeds for values of \(\alpha\) as low as 0.1, a value which is much smaller (and more ill-conditioned) than the value of \(\alpha\) we would likely use in practice.
-As such, we can conclude that our approximate sampling method would be able to sample from the posterior distribution for a wide range of values of \(\alpha\) on the MNIST problem.
-
-We also show the spectrum of the posterior precision matrix for the MNIST problem in \cref{fig:spectrum-mnist} for a CNN with a \(3 \times 3\) filter and a \(5 \times 5\) filter.
-The spectrum of the posterior precision matrix for the \(3 \times 3\) filter has two distinct clusters, with the first cluster being over an order of magnitude larger than the second.
-The first cluster (which dominates the spectrum) consists of nine eigenvalues, and so the precision has an effective rank close to 9.
-For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, with the spectrum appearing smooth (i.e., not having any distinct clusters).
-This distribution of eigenvalues suggests that the effective rank of the precision is over 50, rather than 25 which would be plausible when considering that a \(3 \times 3\) filter has rank 9.
-We do not yet understand this phenomenon and do not have a good explanation for why the posterior precision matrix for the \(5 \times 5\) filter has such a large effective rank and a smooth spectrum, unlike for the \(3 \times 3\) filter.
-
-\begin{marginfigure}
- \centering
- \includegraphics{mnist_spectrum_3vs5.pdf}
- \caption[Spectrum of the posterior precision on MNIST.]{Spectrum of the Laplace posterior precision on MNIST with a CNN. For the \(3 \times 3\) filter, there are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so the precision has an effective rank close to nine. For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, which does not have any distinct clusters. This suggests that the effective rank of the precision is over 50.}
- \label{fig:spectrum-mnist}
-\end{marginfigure}
+As such, we can conclude that our approximate sampling method would be able to sample from the posterior distribution for a wide range of values of \(\alpha\) on MNIST.
+
+% We also show the spectrum of the posterior precision matrix for the MNIST problem in \cref{fig:spectrum-mnist} for a CNN with a \(3 \times 3\) filter and a \(5 \times 5\) filter.
+% The spectrum of the posterior precision matrix for the \(3 \times 3\) filter has two distinct clusters, with the first cluster being over an order of magnitude larger than the second.
+% The first cluster (which dominates the spectrum) consists of nine eigenvalues, and so the precision has an effective rank close to 9.
+% For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, with the spectrum appearing smooth (i.e., not having any distinct clusters).
+% This distribution of eigenvalues suggests that the effective rank of the precision is over 50, rather than 25 which would be plausible when considering that a \(3 \times 3\) filter has rank 9.
+% We do not yet understand this phenomenon and do not have a good explanation for why the posterior precision matrix for the \(5 \times 5\) filter has such a large effective rank and a smooth spectrum, unlike for the \(3 \times 3\) filter.
+
+% \begin{marginfigure}
+% \centering
+% \includegraphics{mnist_spectrum_3vs5.pdf}
+% \caption[Spectrum of the posterior precision on MNIST.]{Spectrum of the Laplace posterior precision on MNIST with a CNN. For the \(3 \times 3\) filter, there are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so the precision has an effective rank close to nine. For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, which does not have any distinct clusters. This suggests that the effective rank of the precision is over 50.}
+% \label{fig:spectrum-mnist}
+% \end{marginfigure}
\section{Discussion}
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -10,14 +10,14 @@ We analyse the construction of the Laplace approximate posterior and discuss its
\section[Optimisation in Deep Learning]{Optimisation in Deep Learning}[Optimisation]
\label{sec:optimisation-frequentist}
-Suppose a neural network is a real-valued function \(f: \reals^N \times \reals^D \rightarrow \reals^O\) parametrised in \(\bm\theta\) which maps an input \(\bm x\) to an output \(f(\bm x, \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
+Suppose a neural network is a real-valued function \(f: \reals^N \times \reals^D \rightarrow \reals^O\) parametrised in \(\bm\theta \in \reals^D\) which maps an input \(\bm x \in \reals^N\) to an output \(f(\bm x, \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
Our goal is to find the parameters \(\bm\theta^*\) which best model the observed data.
From a frequentist perspective, we can define a loss function on the model \(\mathcal{L}(\bm\theta): \reals^D \rightarrow \reals\) such that the optimal parameters \(\bm\theta^*\) minimise this loss function for the given model.
Often, we formulate the loss function as the negative log-likelihood of the data under the model, \(\mathcal{L}(\bm\theta) = -\log \lik\).
In this case, the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the likelihood of the data.\sidenote{Maximising the likelihood is equivalent to minimising the negative log-likelihood, since the log function is strictly monotonically increasing.}
-The likelihood function is problem-specific, and is defined as the probability of the data given the parameters \(\bm\theta\).
+The likelihood function is problem-specific, and is defined as the probability density of the data given the parameters \(\bm\theta\).
Typically, the likelihood is chosen to be a Gaussian distribution with mean \(f_{\bm\theta}(\bm x)\) and precision \(\rho\) (which is frequently chosen to be equal to one) for regression problems, and a categorical distribution with logits \(f_{\bm\theta}(\bm x)\) for classification problems.
These two cases correspond to the mean-squared error (MSE) and cross-entropy (CE) negative log-likelihood loss functions, respectively.
@@ -28,11 +28,10 @@ Due to the assumption of independence between observations, we can factorise the
\end{align}
This allows us to compute the loss as the sum of the losses for each individual observation.
This loss function is defined in an optimisation problem to find the parameters \(\bm\theta^*\) of the model which minimise this loss, thereby maximising the performance of the model on the data.
-% THIS LAST BIT A TAD REDUNDANT?
-\subsection[Stochastic Gradient Descent]{Stochastic Gradient Descent}[SGD]
+\subsection{Gradient Descent}
-Stochastic gradient descent (SGD) is a first-order optimisation algorithm which iteratively updates the parameters \(\bm\theta\) in the direction of the negative gradient of the loss function \(\mathcal{L}(\bm\theta)\).\sidenote{First-order optimisation methods only use the first derivative of the loss function.}
+Gradient descent is a first-order optimisation algorithm which iteratively updates the parameters \(\bm\theta\) in the direction of the negative gradient of the loss function \(\mathcal{L}(\bm\theta)\).\sidenote{First-order optimisation methods only use the first derivative of the loss function.}
The update rule is given by
%
\begin{align}
@@ -50,46 +49,12 @@ Using the chain rule on the composition of the loss into the loss as a function
%
where \(\nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)\) is the gradient of the neural network output with respect to its parameters at iteration \(t\) and \(\nabla_{\bm\theta} \log p(y_t \given \bm \theta)\) is the gradient of the loss function with respect to the model output.
Note that in this case, we perform one gradient update for each data point, which is why we iterate over the data points and update the parameters in the same step \(t\).
-Thus, SGD only requires the gradient of the loss function with respect to the parameters of the neural network, which can be computed efficiently using automatic differentiation for a similar cost to a single forward pass through the network.
-This makes SGD a popular choice for training neural networks.
-However, SGD is prone to getting stuck in local minima, and can be slow to converge.
-It doesn't take into account the curvature of the loss function, and so can overshoot the minimum.
-
-% \subsection[Natural Gradient Descent]{Natural Gradient Descent}[The Natural Gradient]
-% % https://agustinus.kristia.de/techblog/2018/03/14/natural-gradient/
-
-% Natural Gradient Descent is an approximate second-order optimisation method. It has an interpretation as optimizing over a Riemannian manifold using an intrinsic distance metric, which implies the updates are invariant to transformations such as whitening.
-% By using the positive semi-definite (PSD) Gauss-Newton matrix to approximate the (possibly negative definite) Hessian, NGD can often work better than exact second-order methods.
-
-% \textcite{rattray1998natural} introduced the natural gradient as a way to optimize a function \(f\) parametrised by \(\bm\theta\) by following the direction of the steepest descent in the Fisher information metric.
-% It can be interpreted as
-% % The algorithm can be seen in \cref{natural-gradient}.
-% %
-% % \begin{algorithm}
-% % \caption{Natural Gradient Descent} \label{natural-gradient}
-% % \begin{algorithmic}[1]
-% % \For {$i = 1$ to $T$}
-% % \State Calculate \(\mathcal{L}(\bm\theta)\)
-% % \State Calculate the gradient of the loss \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta)\)
-% % \State Calculate the Fisher information matrix \(\mathcal{F}(\bm\theta)\)
-% % \State Calculate the natural gradient \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta) \mathcal{F}(\bm\theta)^{-1}\)
-% % \State Update the parameters \(\bm\theta \gets \bm\theta - \eta \nabla_{\bm\theta} \mathcal{L}(\bm\theta) \mathcal{F}(\bm\theta)^{-1}\)
-% % \EndFor
-% % \State \Return {\(\bm\theta\)}
-% % \end{algorithmic}
-% % \end{algorithm}
-% %
-% In practice, deep learning models have millions (or billions) or parameters, and the Fisher information matrix is consequently often too large to be computed and stored explicitly.
-% This therefore limits the applicability of the natural gradient.
-% Instead, it is often approximated by the diagonal of the Hessian.
-
-% % This brings us to the Adam optimizer, which is a popular choice for optimizing deep neural networks.
-% \cite{kingma2014adam}
-% \cite{martens2015optimizing}
-% \cite{botev2017practical}
-
-% For a recent, detailed discussion of the natural gradient, see \textcite{martens2020new}.
-% \cite{wu2019logan}
+Thus, gradient descent only requires the gradient of the loss function with respect to the parameters of the neural network, which can be computed efficiently using automatic differentiation for a similar cost to a single forward pass through the network.
+However, gradient descent is prone to getting stuck in local minima, and can be slow to converge.
+In practice, we cannot compute the gradient of the loss function with respect to the parameters of the neural network for the entire dataset at once, since this would require storing the entire dataset in memory and computing the gradient over it.
+Instead, we compute the gradient of the loss function with respect to the parameters of the neural network for a subset of the data at each iteration (a batch), and then update the parameters based on this unbiased estimate of the gradient.
+This is known as stochastic gradient descent (SGD), and is the most common form of gradient descent used in practice.
+SGD helps avoid getting stuck in local minima, but, like gradient descent, doesn't take into account the curvature of the loss function, and so can overshoot the minimum.
\subsection{Adam}
\label{sec:adam}
@@ -102,7 +67,6 @@ The moving averages are weighted by the parameters \(\beta_1\) and \(\beta_2\) r
The learning rate is then scaled by the ratio of the first and second moments.
The update rule for Adam is thus given by
\begin{align}
- % \bm\theta_{t+1} ={} & \bm\theta_t - \eta_t \bm m_t \odot \frac{1}{\sqrt{\bm v_t} + \varepsilon},
\bm\theta_{t+1} ={} & \bm\theta_t - \eta_t \bm m_t \odot \left(\sqrt{\bm v_t} + \varepsilon\right)^{-1},
\\ \bm m_t ={}& (1 - \beta_1^t) \left[ \beta_1 \bm m_{t-1} + (1 - \beta_1) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
\\ \bm v_t ={}& (1 - \beta_2^t) \left[ \beta_2 \bm v_{t-1} + (1 - \beta_2) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \odot \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
@@ -119,11 +83,11 @@ The use of the second moment is inspired by optimisation methods such as Newton'
These optimisation methods all identify a single set of parameters as the optimal solution.
In the next section, we will discuss how to learn a distribution over parameters instead.
-\section{Bayesian Deep Learning}
+\section{Maximum a Posteriori Estimation}
\label{sec:bayesian-deep-learning}
Bayesian deep learning is a framework for deep learning that allows for modelling the parameters as random variables instead of as single optima in the loss landscape.
-From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and a likelihood function \(\lik\) such that the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the posterior distribution \(p(\bm\theta \given \bm y)\).
+From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and a likelihood function \(\lik\) such that, under maximum a posteriori estimation, the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the posterior distribution \(p(\bm\theta \given \bm y)\).
This way, we can quantify the uncertainty in our model parameters \(\bm\theta\) by computing the posterior distribution \(p(\bm\theta \given \bm y)\).
It is obtained from Bayes' theorem, which allows for uncertainty quantification by
\begin{align}\label{eq:bayes-theorem}
@@ -136,13 +100,6 @@ However, since the posterior distribution is proportional to the product of the
\begin{align}\label{eq:unnormalised posterior}
p(\bm \theta \given \bm y) \propto p(\bm y \given \bm \theta)\,p(\bm \theta) =: \tilde{p}(\bm \theta \given \bm y) .
\end{align}
-\marginnote{SHOULD I TALK ABOUT HOW WE MAKE PREDICTIONS? I.E., HOW WE COMPUTE THE POSTERIOR PREDICTIVE BY MC SAMPLING FROM THE POSTERIOR?}
-% From here, we can make predictions in a Bayesian setting by computing the posterior predictive distribution
-% \begin{align}\label{eq:posterior predictive}
-% p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
-% \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y), \nonumber
-% \end{align}
-% where \(y^*\) is the new data point, \(S\) is the number of Monte Carlo posterior samples \(\bm \theta_s\).\sidenote{DOES THIS HOLD IN REGRESSION?}
Since we want to determine the parameters which maximise the posterior distribution, we can typically use the same optimisation techniques as for frequentist learning, but with the negative log-posterior loss instead of the negative log-likelihood.
Furthermore, since the unnormalised posterior from \cref{eq:unnormalised posterior} is proportional to the likelihood and the prior, the exact posterior itself need not be tractable.
@@ -159,6 +116,12 @@ However, the posterior distribution itself is typically intractable, so approxim
We use the posterior distribution to quantify the uncertainty in our model parameters.
Often, all that is needed is a set of samples from the posterior distribution.
+From here, we can make predictions in a Bayesian setting by computing the posterior predictive distribution and using Monte Carlo sampling as
+\begin{align}\label{eq:posterior predictive}
+ p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
+ \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y), \nonumber
+\end{align}
+where \(y^*\) is the new data point and \(S\) is the number of Monte Carlo posterior samples \(\bm \theta_s\).
However, in some cases, it is useful to have a closed-form expression for the posterior distribution.
There are several methods for approximating the posterior distribution, which we will discuss in the following sections.
Some, like Markov chain Monte Carlo (MCMC)~\cite{gelman1995bayesian}, are exact methods, which can be used to sample from the posterior distribution directly (though they are often computationally expensive).
@@ -176,16 +139,15 @@ In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mack
However, instead of finding the optimal Gaussian distribution locations and scales by maximising the ELBO, the LA finds the location by computing the MAP solution \(\bm \theta_{\textsc{map}}\) and the scale by approximating the log-posterior with a second degree Taylor expansion around this solution (\(\bm \theta_0 = \bm \theta_{\textsc{map}}\)) and determining the curvature via its Hessian matrix \(\bm \Lambda = - \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_\textsc{map}}\).
Since the Taylor expansion is performed around the MAP solution, the first order derivative is zero, and the expansion is given by
\begin{align}
- \ln p(\bm \theta \given \bm y)
- \approx{} & \ln p(\bm \theta_{\textsc{map}} \given \bm y) + \nonumber
- \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \left( \left.\nabla^2_{\bm \theta} \ln p(\bm \theta \given \bm y) \right|_{\bm \theta_\textsc{map}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) \label{eq:laplace-taylor}
+ \log p(\bm \theta \given \bm y)
+ \approx{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nonumber
+ \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \left( \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_\textsc{map}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) \label{eq:laplace-taylor}
\\ \Rightarrow \tilde{p}(\bm \theta \given \bm y) \approx{}& p(\bm \theta_{\textsc{map}} \given \bm y) \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right),
\end{align}
where \(\tilde{p}(\bm \theta \given \bm y)\) corresponds to the unnormalised posterior (\cref{eq:unnormalised posterior}).
Normalising it by the normalisation constant \(p(\bm y)\) gives the normalised posterior\marginnote{This normalisation constant is computed as the marginal likelihood for the Laplace approximation in \cref{sec:maximising-the-evidence}.}
\begin{align}\label{eq:laplace-normal}
p(\bm\theta \given \bm y) \approx{} & \sqrt{\frac{\det (\bm \Lambda)}{(2 \pi)^D}} \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)
- % \\={}& \sqrt{\frac{\det (- \bm H)}{(2 \pi)^D}} \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_0)\T (- \bm H) (\bm \theta - \bm \theta_0)\right) \nonumber
\\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}). \nonumber
\end{align}
The Laplace approximation can simply be trained as an MAP solution, where the Hessian is computed to obtain the posterior precision at inference time.
@@ -196,7 +158,7 @@ We approach the latter, as well as post-hoc Laplace, in this work.
We have now shown that approximating the log-posterior with a second degree Taylor expansion around the \(\bm\theta_{\textsc{map}}\) corresponds to approximating the posterior with a Gaussian distribution given by \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\) where \(\bm \Lambda = - \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_{\textsc{map}}}\).
However, this approximation only holds when the posterior precision \(\bm \Lambda\) is positive definite.
-For BNNs, this is not guaranteed, which means that the normal distribution assumed by the Laplace approximation may not be valid.
+For BNNs, this is not guaranteed as the loss function may not be convex with respect to the model parameters, which means that the normal distribution assumed by the Laplace approximation may not be valid.
In the next sections, we will go into the computation of the Hessian for the negative log-posterior loss function from \cref{sec:bayesian-deep-learning} and show how the generalised Gauss-Newton approximation can be used to obtain a positive definite Hessian approximation, even when the actual Hessian itself is not positive definite.
\subsection{The Hessian}
@@ -211,6 +173,7 @@ We are interested in the Hessian of the loss function \(\mathcal{L}(\bm\theta)\)
For a neural network with \(D\) parameters, the Hessian is therefore a \(D \times D\) square matrix.
Furthermore, if all the neural network's second partial derivatives are continuous, then the Hessian is symmetric.
In this case, the Hessian is generally dominated by the block-diagonal~\cite{martens2015optimizing}.
+However, the Hessian will only be positive definite if the loss function is a convex function in the model parameters \(\bm \theta\).
Typically, the loss function is considered as the negative log-likelihood or negative log-posterior.
In this framing, if considering, e.g., the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we get
@@ -227,7 +190,7 @@ In the next section, we will discuss how to obtain a positive definite Hessian a
\subsection{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
\label{sec:ggn}
-From the equation of the Hessian we can apply the chain rule twice and the product rule once to obtain a simpler expression,
+From the expression of the posterior precision as the Hessian of the negative log-posterior in \cref{eq:negative-log-posterior-hessian} we can apply the chain rule twice and the product rule once to obtain a simpler expression,
%
\begin{align}\label{eq:ggn-whole-hessian}
\nabla^2_{\bm \theta} \mathcal{L}
@@ -241,9 +204,7 @@ More precisely, we obtain
\begin{align}
\nabla^2_{\bm \theta} \mathcal{L}_{ij} = \sum_{n=1}^O \frac{\partial^2}{\partial \theta_i \theta_j} f_n \cdot \frac{\partial}{\partial f_n} \mathcal{L} + \sum_{n,m=1}^O \frac{\partial}{\partial \theta_i} f_n \cdot \frac{\partial^2}{\partial f_n f_m} \mathcal{L} \cdot \frac{\partial}{\partial \theta_j} f_m,
\end{align}
-where \(f_n\) is the n\textsuperscript{th} output of the neural network function \(f(\bm x, \bm \theta)\).
-% Mention that it's evaluated?
-% The first term is the GGN approximation of the Hessian of the loss function.
+where \(f_n\) is the \(n\)\textsuperscript{th} output of the neural network function \(f(\bm x, \bm \theta)\).
Assume now we will approximate \(f\) using a first-order Taylor expansion.
In this way, we are linearising our neural network function.
@@ -261,7 +222,6 @@ By combining \cref{eq:ggn-whole-hessian,eq:nn-taylor} we then obtain the general
\approx{} & \nabla_{\bm\theta} f\T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm\theta} f,
\end{align}
which corresponds to the second term in \cref{eq:ggn-whole-hessian}.
-% In this approximation, we have \(\nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0\).
This means that the GGN approximation is an approximation to the Hessian of the composition of two functions, \(\mathcal{L}(\bm \theta) \equiv (\mathcal{L} \circ f) (\bm \theta)\) where we linearise the inner function \(f\) around a point \(\bm \theta_0\).
For a given negative log-likelihood loss \(\mathcal{L}\), we thus have the GGN approximation of the Hessian
\begin{align}\label{eq:negative-log-likelihood-ggn}
@@ -274,18 +234,16 @@ For a normal prior with precision \(\alpha\) on \(\bm \theta\), we obtain the GG
\end{align}
where we have defined \(\bm J_i = J_{\bm \theta} f(\bm x_i, \bm \theta)\) and \(\bm H_i = \nabla^2_f \mathcal{L}(y_i, f(\bm x_i, \bm \theta))\).
-For mean square error (MSE) loss, we have one output variable (\(O = 1\)), so \(f: \reals^N \times \reals^D \rightarrow \reals\).
-We then get the simple expression for the GGN approximation of the Hessian for MSE loss,
+For mean square error (MSE) loss (e.g. for regression), we obtain the simple expression the GGN approximation of the Hessian
%
\begin{align}%\label{eq:mse-hessian}
- \bm H_i \approx & {}\, (y_i - f(\bm x_i, \bm \theta))^2 = 2,
+ \bm H_i \approx & {}\, (\bm y_i - f(\bm x_i, \bm \theta))^2 = 2 \identity,
\\ \bm \Lambda \approx&{}\, 2 \sum_{i=1}^N \bm J_i\T \bm J_i + \alpha \identity.
\end{align}
If the Hessian of the loss function with respect to the model outputs \(\bm H_i\) is positive semidefinite, the GGN approximation is positive semidefinite, since each term in the sum \(\bm J_i\T \bm H_i \bm J_i\) is positive semidefinite.\sidenote{This is equivalent to the loss function being a convex function of the model parameters \(\bm \theta\).}
Furthermore, this is the case for most common loss functions, such as MSE loss and cross-entropy loss.
For the case of the normal prior, the GGN approximation of the posterior precision will be positive definite, since the eigenvalues of \(\alpha \identity\) are all \(\alpha > 0\) and the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) are all greater than or equal to zero.
-% FISHER INFORMATION MATRIX
It is also interesting to note that, for our definition of the loss, the GGN approximation of the Hessian is equivalent to the Fisher information matrix, which is a measure of the curvature of the log-likelihood function with respect to the model parameters.
Here, the Fisher information matrix is given by
\begin{align}
@@ -319,18 +277,18 @@ The diagonal of this matrix is given by the element-wise square of the gradient
Notice that this is the second-order term used in the Adam optimiser, which we discussed in \cref{sec:adam}.
The choice of application of the generalised Gauss-Newton approximation is not based on having a reduced computational cost---both the Hessian and the GGN approximation require the same number of forward and backward passes.
-However, the GGN approximation is better behaved than the exact Hessian, thereby making it more suitable for practical applications.
+However, the GGN approximation is better behaved than the exact Hessian since it is guaranteed to be positive semidefinite for positive definite loss functions, thereby making it more suitable for practical applications.
This is because the GGN approximation is designed to be positive semi-definite, while the exact Hessian can be vulnerable to negative curvature.
-This feature can limit the use of the exact Hessian for optimisation algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
+This feature can limit the use of the exact Hessian for algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
\subsection{The Spectrum of the GGN Approximation}[The Spectrum of the GGN]
\label{sec:spectrum-ggn}
When we construct the posterior precision for the Laplace approximation, we compute the posterior precision as in \cref{eq:laplace-precision-ggn}.
For methods which we will introduce later, we are interested in the spectrum (i.e., the set of eigenvalues) of the posterior precision.
-The eigenvalues of this precision will thus be the sum of the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) and of \(\alpha \identity\).
+The eigenvalues of this precision will thus be the sum of the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) and of \(\alpha \identity\), where \(\alpha > 0\) is the prior precision.
The first term, being the positive semi-definite GGN matrix, will have non-negative eigenvalues.
-The second term, being a diagonal matrix with entries \(\alpha\), will have eigenvalues that are all equal to \(\alpha > 0\), where \(\alpha\) is the prior precision.
+The second term, being a diagonal matrix with entries \(\alpha\), will have eigenvalues that are all equal to \(\alpha\).
Therefore, the eigenvalues of the posterior precision will be the sum of the eigenvalues of the GGN matrix and \(\alpha\) and will thus all be positive.
Thus, the prior precision \(\alpha\) guarantees the positive definiteness of the posterior precision.
For some applications, it is desirable to have a weak prior, so that the regularisation effect of the prior is not too strong.
@@ -359,6 +317,17 @@ If the \(O\) non-zero eigenvectors are identical across the GGN matrices, then t
However, if these eigenvectors are not identical, then as the largest eigenvalue of the sum of the GGN increases and the number of non-zero eigenvalues of the sum of the GGN increases, the non-zero spectrum of the sum of the GGN matrices will get steeper.
For these reasons, as we add more data points, the posterior precision will become more ill-conditioned.\sidenote{A matrix is considered ill-conditioned when the ratio of its largest eigenvalue to its lowest eigenvalue, its \emph{condition number} (\(\kappa \coloneqq \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}\)), is very large.}
+Empirically, we find that, for MNIST, the effective rank of the posterior precision is significantly less than \(N \cdot O\) (see \cref{fig:spectrum-mnist}).
+This suggests that the eigenvectors of the posterior precision are relatively similar across different observations, though not identical.
+However, this seems to greatly depend on the structure of the neural network (e.g., the filter size) in ways that are not yet clear.
+
+\begin{marginfigure}
+ \centering
+ \includegraphics{mnist_spectrum_3vs5.pdf}
+ \caption[Spectrum of the posterior precision on MNIST.]{Spectrum of the Laplace posterior precision on 100 MNIST observations with a CNN. Note the log scale and the fact that we're zoomed in on the largest eigenvalues, as can be seen from the eigenvalue indices on the x-axis. For the \(3 \times 3\) filter, there are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so the precision has an effective rank close to nine. For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, which does not have any distinct clusters. This suggests that the effective rank of the precision is over 50, which is still significantly lower than the number of observations \(N = 100\) times the number of outputs \(O = 10\).}
+ \label{fig:spectrum-mnist}
+\end{marginfigure}
+
\section{Limitations}
One limitation of the Laplace approximation is that it can only approximate one mode of the posterior distribution (the MAP).
@@ -373,252 +342,6 @@ This is because the full Laplace approximation requires the computation of the p
Thus, in order to scale to large models and large datasets without approximating the posterior precision with, e.g., a diagonal matrix, we need to find a method to perform the Laplace approximation without instantiating the posterior precision.
During training, we typically only need to train the model in a deterministic way (to find the MAP), and then we can use the Laplace approximation to sample from the approximate posterior post-hoc.
Thus, for inference, we need to determine a method for sampling from the Gaussian approximate posterior without explicitly computing the posterior precision.
-However, it is also possible to perform the Laplace approximation by training a model on its marginal likelihood (i.e., the likelihood of the data under the model).
+We have so far only considered the Laplace approximation computed by maximising posterior (i.e., MAP estimation), but it is also possible to perform the Laplace approximation by training a model on its marginal likelihood (i.e., the likelihood of the data under the model).
This would reduce overfitting and allow us to optimise hyperparameters via backpropagation.
However, we must find a differentiable way to compute the Laplace approximate log-marginal likelihood without explicitly storing the posterior precision.
-
-% \section[The Fisher Information Matrix]{The Fisher Information Matrix}[Fisher Information]
-% \label{sec:fisher-information}
-% % https://agustinus.kristia.de/techblog/2018/03/11/fisher-information/
-
-% Consider the gradients of the log-likelihood optimisation objective, given by the expression \(\nabla_{\bm\theta} \log\lik\).
-% We can show that the expectation of these gradients is zero:
-% \begin{align*}
-% \E{\lik}{\nabla_{\bm\theta} \log\lik} = & {} \E{\lik}{\nabla \log \lik}
-% \\ =&{} \int \nabla \log \lik \lik \, d \bm y
-% \\ =&{} \int \frac{\nabla \lik}{\lik} \lik \, d \bm y % as per the chain rule
-% \\ =&{} \int \nabla \lik \, d \bm y % as per the Leibniz integral rule
-% \\ =&{} \nabla \int \lik \, d \bm y
-% \\ =&{} \nabla 1
-% \\ =&{} 0
-% \end{align*}
-
-% Let us now analyse the covariance of these gradients.
-% It is given by
-% %
-% \begin{align}\label{eq:gradient-covariance}
-% \mathbb{E}_\lik \Big[ & \left(\nabla \log\lik - \E{}{\nabla \log\lik}\right) \nonumber
-% \\ & \cdot \left(\nabla \log\lik - \E{}{\nabla \log\lik}\right)\T\Big] \nonumber
-% \\ ={}& \E{\lik}{\left(\nabla \log\lik - 0\right) \left(\nabla \log\lik - 0\right)\T} \nonumber
-% \\ ={}& \E{\lik}{\nabla \log\lik \nabla \log\lik\T} \nonumber
-% \\ \equiv{}& \mathcal{F}(\bm \theta)
-% \end{align}
-% %
-% where we have defined the Fisher information matrix \(\mathcal{F}(\bm \theta)\) as the expectation of the outer product of the score function under our model.
-% This product is usually intractable.
-% It can, however, be approximated empirically by using samples from an empirical distribution \(q(\bm y)\) which is close to the true likelihood \(\lik\) \cref{eq:empirical-fisher}.
-% %
-% \begin{equation}\label{eq:empirical-fisher}
-% \mathcal{F} \approx \frac{1}{N} \sum^N_{n=1} \nabla \log p(y_n \given \bm \theta) \nabla \log p(y_n \given \bm \theta)\T.
-% \end{equation}
-% %
-% Note, however, that this matrix is distinct from the Fisher information matrix~\cite{martens2020new, kunstner2019limitations}, since
-
-% Another non-trivial property of the Fisher information matrix is that it is equivalent to the negative expectation of the Hessian of the log-likelihood.
-% %
-% \begin{align*}
-% H_{\bm\theta} \left(\log \lik\right) = & {} \nabla^2_{\bm\theta} \log \lik
-% \\ =&{} \nabla_{\bm\theta} \nabla_{\bm\theta} \log \lik
-% \\ =&{} \nabla_{\bm\theta} \frac{\nabla_{\bm\theta} \lik}{\lik}
-% \\ =&{} \frac{}{}
-% \end{align*}
-% %
-% The Fisher information matrix can therefore be interpreted as the curvature of the log-likelihood.
-
-% RELATIONSHIP WITH GAUSS-NEWTON APPROXIMATION
-% for common loss functions, the GGN is equivalent to the Fisher information matrix~\cite{martens2020new}
-% \cite{kunstner2019limitations}
-
-% RELATIONSHIP WITH KL DIVERGENCE
-% % https://agustinus.kristia.de/techblog/2018/03/14/natural-gradient/
-% % https://en.wikipedia.org/wiki/Fisher_information#Relation_to_relative_entropy
-
-% LIMITATIONS OF FISHER INFORMATION (MAYBE ESPECIALLY EMPIRICAL?)
-% \cite{kunstner2019limitations}
-
-% \section[The Woodbury matrix identity]{The Woodbury matrix identity}[Woodbury]
-% \label{sec:woodbury}
-
-% In many methods, such as Natural Gradient Descent~\cite{martens2020new} and the Laplace approximation~\cite{daxberger2021laplace} (explained in \cref{sec:laplace}), we must invert the Hessian.
-% Since, in neural networks, the number of parameters is often very large, this inversion is computationally expensive and/or cannot be stored in memory.
-% Furthermore, the Hessian matrix is often ill-conditioned, which can lead to numerical instability.
-% We typically approximate the Hessian matrix with the Gauss-Newton approximation, which is a low-rank approximation of the Hessian matrix.
-% This also leads to many methods being inefficient (TODO: explain this more, cite, conditioning, iterative methods, etc.).
-% The Woodbury matrix identity~\cite{woodbury1950inverting} is a formula that allows us to invert a matrix that is composed as the sum of a low-rank matrix and a diagonal matrix.
-% If we have a matrix \(\bm P\) that is the sum of a low-rank matrix \(\bm L\T \bm L\) and a diagonal matrix \(\alpha \identity\), then we can invert \(\bm P\) using the Woodbury matrix identity as
-% % Theorem 1.1 ([18, Thm. 1.35]). Let U, V ∈ Cn×k with k ≤ n and assume that V ∗U is nonsingular. Let f be defined on the spectrum of A = αIn + U V ∗, and if k = n let f be defined at α. Then (1.2) f (A) = f (α)In + U (V ∗U )−1(f (αIk + V ∗U ) − f (α)Ik )V ∗.
-% \begin{align}
-% \bm P ={} & \bm L\T \bm L + \alpha \identity
-% \\ \bm P^{-1} ={}& \alpha^{-1} \identity - \alpha^{-2} \bm L\T \left( \identity + \alpha^{-1} \bm L \bm L\T \right)^{-1} \bm L.
-% \end{align}
-% We can then compute the matrix-vector product and inverse matrix-vector product as
-% \begin{align}
-% \bm P \bm v ={} & \bm L\T ( \bm L \bm v ) + \alpha \bm v
-% \\ \bm P^{-1} \bm v ={}& \alpha^{-1} \bm v - \alpha^{-2} \bm L\T ( \underbrace{\identity + \alpha^{-1} \bm L \bm L\T}_{k \times k} )^{-1} ( \bm L \bm v )
-% \end{align}
-% such that we do not need to instantiate the whole matrix.
-% Since \(\identity + \alpha^{-1} \bm L \bm L\T\) is framed as an outer product \(\bm L \bm L\T\), we only need to compute the inverse of a \(k \times k\) matrix, which is much easier than computing the inverse of a \(n\times n\) matrix, since \(k \ll N\).
-% Furthermore, the condition number of this matrix is given by the ratio of the largest to the smallest singular value of \(\bm L \bm L\T\), which is much smaller than the condition number of \(\bm P\) itself (since the smallest singular value of \(\bm L\T \bm L\) is zero).
-
-% % For an arbitrary function \(f\),
-% % \begin{align}
-% % f\left(J\T J + \alpha \identity\right) ={} & f(\alpha) \identity + J\T \left(J J\T\right)^{-1} \left( f\left( J J\T + \alpha \identity \right) - f(\alpha) \identity \right) J \nonumber
-% % \\ f\left(J\T J + \alpha \identity\right) v ={}& f(\alpha) v + J\T \left(J J\T\right)^{-1} \left( f\left( J J\T + \alpha \identity \right) - f(\alpha) \identity \right) J v
-% % \end{align}
-
-% % \begin{align}
-% % \left(J\T J + \alpha \identity\right)^{-1} ={} & \alpha^{-1} \identity + J\T \left(J J\T\right)^{-1} \left( \left( J J\T + \alpha \identity \right)^{-1} - \alpha^{-1} \identity \right) J
-% % \end{align}
-
-
-% % % https://mathoverflow.net/questions/121411/expectation-of-square-root-of-binomial-r-v
-% % \begin{align}
-% % \E{}{\sqrt{X}} \approx \sqrt{\E{}{X}} \left( 1 - \frac{\Var{X}}{8 \E{}{X}^2} \right)
-% % \end{align}
-
-% \section[Krylov Subspace Methods]{Krylov Subspace Methods}[Krylov Methods]
-
-% A Krylov subspace of a square matrix \(\bm K \in \reals^{N \times N}\) and a vector \(\bm b \in \reals^{N}\) is a subspace spanned by the set of images of \(\bm b\) under the first \(j\) powers of \(\bm K\).
-% \begin{align}
-% \mathcal{K}_j( \bm K, \bm b) = \mathrm{span}\left\{ \bm b, \bm K \bm b, \bm K^2 \bm b, \ldots, \bm K^{j-1} \bm b \right\}
-% \end{align}
-
-% Because of this, matrix-vector products alone are enough to construct a Krylov subspace.
-% Krylov subspace methods leverage this fact to estimate functions of matrices multiplied with vectors, \(f(\bm K) \bm b\).
-% These methods therefore identify approximate solutions which are contained within the \(j\)\textsuperscript{th} Krylov subspace, i.e., \(f(\bm K) \bm b \approx \bm c_j \in \mathcal{K}_j( \bm K, \bm b)\).
-
-% These methods are always exact after \(N\) iterations, however are often almost exact after \(J \ll N\) iterations~\cite{pleiss2020fast}.
-% One major advantage Krylov subspace methods have is that they only require computation of matrix-vector products \(\bm K \bm b\), thereby not requiring direct instantiation of the matrix \(\bm K\).
-% As such, if allocation of memory to store a matrix is infeasible, but the computational graph for this matrix can be determined, then methods such as the conjugate gradient method or the minimum residual method can be used to approximate the inverse of this matrix~\cite{giordano2018covariances, nocedal1999numerical}.
-
-% \subsection{The Lanczos Algorithm}
-% \label{sec:lanczos}
-
-% One application known as the Lanczos algorithm adapts the power iteration method, which computes the greatest eigenvector of a square matrix \(\bm K\).
-% The power iteration method involves iteratively computing the normalised product of \(\bm K\) with a random vector, i.e.,
-% \begin{align}
-% \bm b_{j+1} = \frac{\bm K \bm b_{j}}{\norm{\bm K \bm b_{j}}},
-% \end{align}
-% for some number of iterations \(J\), such that \(\lim_{j \to \infty} \bm b_j\) equals the eigenvector corresponding to the largest eigenvalue of \(\bm K\).
-% This thereby yields a set of vectors which is proportional to \(\{ \bm b_0, \bm K \bm b_0, \bm K^2 \bm b_0, \ldots, \bm K^{j-1} \bm b_0 \}\) which is clearly a spanning set for a Krylov subspace.\sidenote{This is because the norm of the Lanczos vectors \(\zeta_j \coloneqq \norm{\bm K \bm b_{j}}\) is a scalar, which means that \(\bm b_j = \bm K \bm b_{j-1} \zeta_{j-1}^{-1} = \zeta_{j-1}^{-1} \bm K \bm b_{j-1}\), so \(\bm b_j = Z \bm K \bm b_{j-1}\) where \(Z = \prod_{i=0}^{j-1} \zeta_{i}^{-1} \).}
-% However, the power iteration method only uses the last vector \(\bm b_J\), thereby effectively ``wasting'' the other vectors \(\bm b_1, \ldots, \bm b_{J-1}\).
-% To make use of the entire Krylov subspace \(\mathcal{K}_j(\bm K, \bm b_0)\), the Lanczos algorithm combines the power iteration method with the Gram–Schmidt process, along with the assumption that \(\bm K\) is Hermitian, to produce an orthonormal basis from the Krylov subspace.
-% This method then converts an eigendecomposition problem for \(\bm K\) into an eigendecomposition problem for a tridiagonal matrix of rank \(J\) whose maximum and minimum eigenvalues will be close to those of \(\bm K\).
-
-% \subsection[The Conjugate Gradient Method]{The Conjugate Gradient Method}[The Conjugate Gradient]
-
-% Further adjustments to the Lanczos algorithm yield the conjugate gradient method, a Krylov subspace method used to compute the inverse of a matrix by solving a system of equations as
-% \begin{align*}
-% \bm x = \bm K^{-1} \bm b,
-% \end{align*}
-% where \(\bm K\) is a positive semi-definite matrix.\marginnote{A positive semi-definite matrix is square.}
-% The algorithm itself can be seen in \cref{alg:cg}.
-
-% % EXPLANATION
-% % \emph{When even calculating or instantiating Hηη is prohibitively time-consuming, one can use conjugate gradient algorithms to approximately compute H−1 ηη gη (Wright and Nocedal, 1999, Chapter 5). The advantage of conjugate gradient algorithms is that they approximate H−1 ηη gη using only the Hessian-vector product Hηη gη, which can be computed efficiently using automatic differentiation without ever forming the full Hessian Hηη. See, for example, the hessian vector product method of the Python autograd package (Maclaurin et al., 2015). Note that a separate conjugate gradient problem must be solved for each column of gη| , so if the parameter of interest g (θ) is high-dimensional it may be faster to pay the price for computing and inverting the entire matrix Hηη . See 5.3.2 for more discussion of a specific example. In Theorem 2, we require η0∗ to be at a true local optimum. Otherwise the estimated sensitivities may not be reliable (e.g., the covariance implied by Eq. (14) may not be positive definite). We find that the classical MFVB coordinate ascent algorithms (Blei et al. (2016, Section 2.4)) and even quasi-second order methods, such as BFGS (e.g., Regier et al., 2015), may not actually find a local optimum unless run for a long time with very stringent convergence criteria. Consequently, we recommend fitting models using second-order Newton trust region methods. When the Hessian is slow to compute directly, as in Section 5, one can use the conjugate gradient trust region method of Wright and Nocedal (1999, Chapter 7), which takes advantage of fast automatic differentiation Hessian-vector products without forming or inverting the full Hessian.}
-
-% \begin{marginfigure}
-% \begin{algorithm}[H]
-% \caption{Conjugate Gradient (CG)}\label{alg:cg}
-% \Input{\(\bm K \succ 0\), \(\bm b\), \(\bm P \succ 0\), \(J > 0\)}
-% \Output{\(\bm x = \bm K^{-1} \bm b\), \(\bm\alpha\), \(\bm\beta\)}
-% \(\bm x_0 \gets \bm 0\)\;
-% \(\bm r_0 = \bm b - \bm K \bm x_0\)\;
-% \(\bm p_0 = \bm z_0 = \bm P \bm r_0\)\;
-% \(\gamma_0 = \bm r_0 \bm z_0\)\;
-% \For{\(j \gets 1\) \KwTo \(J\)}{
-% \(\bm y_{j} \gets \bm A \bm p_{j-1}\)\;
-% \(\alpha_j = \gamma_{j-1} / (\overline{\bm p_{j-1}} \bm y_{j})\)\;
-% \(\bm x_{j} = \bm x_{j-1} + \alpha_j \bm p_{j-1}\)\;
-% \(\bm r_{j} = \bm r_{j-1} - \alpha_j \bm y_{j}\)\;
-% \(\bm z_{j} = \bm M \bm r_{j}\)\;
-% \(\gamma_{j} = \overline{\bm r_{j}} \bm z_{j}\)\;
-% \(\beta_{j} = \gamma_{j} / \gamma_{j-1}\)\;
-% \(\bm p_{j} = \bm z_{j} + \beta_{j} \bm p_{j-1}\)\;
-% }
-% \(\bm x \gets \bm x_{J}\)\;
-% \end{algorithm}
-% \end{marginfigure}
-
-% \subsection[The Minimum Residual Method]{The Minimum Residual Method}[The Minimum Residual]
-% \label{sec:msminres}
-
-% Another method which is derived from the Lanczos algorithm is the minimum residual method (MINRES).
-% Given a certain number of iterations \(J\), the MINRES algorithm constructs a Krylov subspace by exploiting the three-vector recurrence (like conjugate gradient does) and solves a linear system \(\bm K \bm c = \bm b\), where \(\bm K\) is a square matrix (though not necessarily a symmetric one), by finding the least-squares solution which lies within this Krylov subspace of dimension \(J\), i.e., by solving the constrained optimisation problem characterised by
-% \begin{align}
-% \bm c^* = \argmin_{\bm c \in \mathcal{K}_J(\bm K, \bm b)} \norm{\bm K \bm c - \bm b}_2.
-% \end{align}
-
-% This method then
-
-% The MINRES algorithm can then be trivially extended to solve for multiple shifts, as seen in \cref{alg:msminres}.
-
-% % \begin{algorithm}
-% % \caption{Multi-Shift Minimum Residual (msMINRES)}\label{alg:msminres}
-% % \Input{$K \succ 0$, $b$, $P \succ 0$, $t_1, \ldots, t_Q$, $J > 0$}
-% % \Output{$c_1 = (K + t_1)^{-1} b, \ldots, c_Q = (K + t_Q)^{-1} b$}
-% % $q_1 \gets b / \norm{b}_2$\;
-% % $v_1 \gets K q_0$\;
-% % $\delta_1 \gets \norm{b}_2, \delta_0 \gets 1$\;
-% % \For{$q \gets 1$ \KwTo $Q$}{
-% % $c_1^{(q)} \gets 0$\;
-% % $d_1^{(q)}, d_0^{(q)} \gets 0$\;
-% % $\gamma_2^{(q)} \gets \norm{b}_2$\;
-% % $\eta_1^{(q)} \gets 1, \eta_0^{(q)} \gets 0$\;
-% % }
-% % \For{$j \gets 2$ \KwTo $J$}{
-% % $q_j \gets v_j / \delta_j$\;
-% % $v_j \gets K q_j - \delta_j q_{j-1}$\;
-% % $\gamma_j \gets q_j v_j$\;
-% % $v_j \gets \gamma_j q_j$\;
-% % $\delta_j \gets \norm{v_j}$\;
-% % \For{$q \gets 1$ \KwTo $Q$}{
-% % $\varepsilon_j^{(q)} \gets \delta_{j-1} \left( \delta_{j-2} / \sqrt{\delta_{j-2}^2 + \eta_{j-2}^{(q)2}} \right)$\;
-% % $\zeta_j^{(q)} \gets \delta_{j-1} \left( \eta_{j-2}^{(q)} / \sqrt{\delta_{j-2}^2 + \eta_{j-2}^{(q)2}} \right)$\;
-% % $\eta_j^{(q)} \gets (\gamma_j + t_q) \left( \eta_{j-1}^{(q)} / \sqrt{\delta_{j-1}^2 + \eta_{j-1}^{(q)2}} \right) + \zeta_j^{(q)} \left( \delta_{j-1} / \sqrt{\delta_{j-1}^2 + \eta_{j-1}^{(q)2}} \right)$\;
-% % $\zeta_j^{(q)} \gets \zeta_j^{(q)} \left( \eta_{j-1}^{(q)} / \sqrt{\delta_{j-1}^2 + \eta_{j-1}^{(q)2}} \right) + (\gamma_j + t_q) \left( \delta_{j-1} / \sqrt{\delta_{j-1}^2 + \eta_{j-1}^{(q)2}} \right)$\;
-% % $\eta_j^{(q)} \gets \eta_j^{(q)} \left( \eta_{j}^{(q)} / \sqrt{\delta_{j}^2 + \eta_{j}^{(q)2}} \right)$\;
-% % $\varphi_j^{(q)} \gets \varphi_{j-1}^{(q)} \left( \delta_{j-1} / \sqrt{\delta_{j-1}^2 + \eta_{j-1}^{(q)2}} \right) \left( \eta_{j}^{(q)} / \sqrt{\delta_{j}^2 + \eta_{j}^{(q)2}} \right)$\;
-% % $d_j^{(q)} \gets \left( q - \zeta_j^{(q)} d_{j-1}^{(q)} - \varepsilon_{j}^{(q)} d_{j-2}^{(q)} \right) / \eta_{j}^{(q)}$\;
-% % $c_j^{(q)} \gets c_{j-1}^{(q)} + \varphi_{j}^{(q)} + d_{j}^{(q)}$\;
-% % }
-% % }
-% % \end{algorithm}
-
-% \section{Jacobi Elliptic Functions}\label{sec:jacobi-elliptic}
-
-% Complete elliptic integral of the first kind is given by
-% \begin{align}
-% \mathcal{K}(p) = \int_0^{\pi/2} \left( 1 - m \sin(t)^2 \right)^{-1/2}\,dt,
-% \end{align}
-% where \(p = 1 - m\), and is defined in the domain \(0 < p \leq 1\).\sidenote{If \(p > 1\), then what should be computed is instead the identity \(\mathcal{K}(p) = \mathcal{K}(1/p) / \sqrt{p}\).}
-% This integral is approximated by
-% \begin{align}
-% \mathcal{K}(p) \approx P(p) - \log(p) Q(p)
-% \end{align}
-% where \(P\) and \(Q\) are tenth-order polynomials.
-
-% Jacobian elliptic functions \(\sn(u \given m)\), \(\cn(u \given m)\), and \(\dn(u \given m)\)
-% undefined outside the range \(0 \leq u \leq 1\)??? can't be right
-% % These functions are periodic, with quarter-period on the real axis equal to the complete elliptic integral ellipk(m).
-% % Relation to incomplete elliptic integral: If u = ellipk(phi,m), then sn(u|m) = sin(phi), and cn(u|m) = cos(phi). The phi is called the amplitude of u.
-% % Computation is by means of the arithmetic-geometric mean algorithm, except when m is within 1e-9 of 0 or 1. In the latter case with m close to 1, the approximation applies only for phi < pi/2.
-
-% Honestly don't understand any of these whatsoever right now
-
-% \begin{marginfigure}
-% % \input{images/ellipk.tex}
-% \includegraphics[scale=1]{ellipk.pdf}
-% \caption[Comparison of SciPy and JAX \texttt{ellipk} implementations.]{Comparison of SciPy and JAX \texttt{ellipk} implementations. For values for which the complete elliptic integral \(K(p)\) is well-defined, both implementations yield the same results. Furthermore, due to its static nature, JAX implementation is compatible with \emph{just-in-time} compilation.}
-% \labfig{fig:ellipk}
-% \end{marginfigure}
-
-% \begin{figure}
-% \centering
-% % \input{images/ellipj.tex}
-% \includegraphics[scale=1]{ellipj.pdf}
-% \caption[Comparison of SciPy and JAX \texttt{ellipj} implementations.]{Comparison of SciPy and JAX \texttt{ellipj} implementations. For values for which the functions are well-defined, both implementations yield the same results. Functions plotted are $\sn(u \given m = 0.7)$, $\cn(u \given m = 0.7)$, and $\dn(u \given m = 0.7)$ Furthermore, due to its static nature, JAX implementation is compatible with \emph{just-in-time} compilation.}
-% \label{fig:ellipj}
-% \end{figure}
-
-% % \section{The Mahalanobis Distance and the Chi-Squared Distribution}
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -23,7 +23,7 @@ Because of this, the covariance matrix is given by the inverse of this Hessian,
\end{align}
where \(\bm \Lambda = \nabla^2_{\bm\theta} \log p(\bm \theta \given \bm y)\).
In \cref{eq:sample-reparam}, the square root of a matrix yields another matrix which, when multiplied by itself, is equal to the original matrix.
-However, since a Gaussian distribution is invariant to rotation, in this special case, we just need to obtain a matrix which returns the original matrix \emph{up to a rotation}.\sidenote{That is, if we have a vector of standard normal samples \(\bm \varepsilon \sim \normal(\bm 0, \identity)\), then, for an orthogonal matrix \(\bm R\) (i.e., a rotation matrix), \(\bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\). We then have that we can sample \emph{non}-standard normal samples as \(\bm \Lambda^{-1/2} \bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\).}
+However, since an isotropic Gaussian distribution is invariant to rotation, in this special case, we just need to obtain a matrix which returns the original matrix \emph{up to a rotation}.\sidenote{That is, if we have a vector of standard normal samples \(\bm \varepsilon \sim \normal(\bm 0, \identity)\), then, for an orthogonal matrix \(\bm R\) (i.e., a rotation matrix), \(\bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\). We then have that we can sample \emph{non}-standard normal samples as \(\bm \Lambda^{-1/2} \bm R\T \bm \varepsilon \sim \normal(\bm 0, \bm \Lambda^{-1})\).}
We thus need to find a method which computes the inverse square root of the full \(D \times D\) posterior precision matrix up to a rotation, where \(D\) is the number of model parameters.
Some methods exist which compute the inverse square root of a matrix up to a rotation, such as the Cholesky decomposition.
@@ -191,7 +191,8 @@ Jacobian elliptic functions \(\sn(u \given m)\), \(\cn(u \given m)\), and \(\dn(
% Computation is by means of the arithmetic-geometric mean algorithm, except when m is within 1e-9 of 0 or 1. In the latter case with m close to 1, the approximation applies only for phi < pi/2.
To compute these values, we also use the elliptic modulus \(k = \kappa(\bm K)^{-1/2}\) (where \(\kappa(\bm K)\) denotes the condition number of \(\bm K\)) and \(u_q = (q - 1/2) / Q\).
These functions are not implemented in JAX, and SciPy uses the Cephes numerical library in C to compute them.
-As such, we have ported the Cephes Jacobi elliptical functions to JAX, and \cref{fig:ellipk,fig:ellipj} show the functions evaluated for the two implementations.
+Because of this, we ported the Cephes Jacobi elliptical functions to JAX, and \cref{fig:ellipk,fig:ellipj} show the functions evaluated for the two implementations.
+Our static JAX implementation allows for efficient \emph{just-in-time} compilation, which is not possible with the SciPy implementation.
We can thereby now compute the quadrature weights and shifts, as seen in \cref{alg:quad}.
% where \(\sn\), \(\cn\), and \(\dn\) are the Jacobi elliptic functions and the complete elliptic integral is \(\mathcal{K}^\prime(k) = \mathcal{K}(k^\prime) = \mathcal{K}(\sqrt{1 - k^2})\) (see \cref{sec:jacobi-elliptic}).
% To compute these values, we also use the elliptic modulus \(k = \kappa(\bm K)^{-1/2}\) (where \(\kappa(\bm K)\) denotes the condition number of \(\bm K\)), and \(u_q = (q - 1/2) / Q\).
@@ -326,7 +327,7 @@ The spectrum of this preconditioner would thus be similar to that of the true po
Since the lowest eigenvalue of the posterior precision is the prior precision \(\alpha\), if the ill conditioning is caused by this eigenvalue, then we can use few preconditioner observations to increase the lowest eigenvalue of the preconditioned system to be equal to one.
Alternatively, we can also instead estimate the eigenvectors and eigenvalues of \(\textstyle{N / B \sum^B_i} J_i\T J_i\).
-This would \emph{hopefully} have an eigenbasis that is closer to that of \(\textstyle{\sum^N_i} J_i\T J_i\) without needing to calculate the precision matrix over the whole dataset (we may need to guarantee all classes are represented in \(x_i, i \sim \mathcal{B}\)).
+This would \emph{hopefully} have an eigenbasis that is close to that of \(\textstyle{\sum^N_i} J_i\T J_i\) without needing to calculate the precision matrix over the whole dataset (we may need to guarantee all classes are represented in \(x_i, i \sim \mathcal{B}\)).\sidenote{Empirically, it seems likely that the eigenvectors of the posterior precision are relatively similar across different observations---see \cref{fig:spectrum-mnist} and \cref{sec:spectrum-ggn}.}
However, scaling the precondititioner up to the size of the entire dataset will lead to a preconditioner which is ill-conditioned, making inversion via conjugate gradient difficult.
As mentioned before, each msMINRES iteration performs one matrix-vector product and one preconditioner-vector product.
@@ -405,18 +406,18 @@ This would then require storing the factorisation \(\bm L \in \reals^{K \times D
As explained in \cref{sec:spectrum-ggn}, as we sum over multiple low-rank GGN matrices, we will likely find that the minimum non-zero eigenvalue gets progressively smaller compared to the largest.
The condition number will therefore get progressively worse.
-The preconditioning of \(\sum^N_{i=1} \bm J_i \bm H_i \bm J_i\T\) (which should be low-rank) will progressively become both \emph{less} low-rank and its condition number will progressively become higher.% MERGE WITH:
+The preconditioning of \(\sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) (which should be low-rank) will progressively become both \emph{less} low-rank and its condition number will progressively become higher.% MERGE WITH:
-Regarding the eigenvalues and eigenvectors, we consider a best-case scenario.
-Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors.
+Regarding the non-zero eigenvalues and eigenvectors, we consider a best-case scenario.
+Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors of the precision matrix.
If the eigenspectrum of this factorisation is steep enough, then the largest eigenvalue \(s_1\) will be significantly larger than the smallest eigenvalue \(s_K\) (i.e., there will be a steep drop-off from the largest to the smallest non-zero eigenvalue).
-Since the eigenvectors of each per-observation Hessian will not be the same, when we sum over them, the eigenvalues of the sum will not be the sum of the eigenvalues of the individual Hessians.
+Since the eigenvectors of each per-observation term \(\bm J_i\T \bm H_i \bm J_i\) will not be the same, when we sum over them, the eigenvalues of the sum will not be the sum of the eigenvalues of the individual Hessians.
However, some of the eigenvectors may be close to each other, and so the spectrum of the GGN sum may be approximately the sum of the GGN eigenvalues.
-However, the rank of the total Hessian will increase by the number of observations and so the number of non-zero eigenvalues consequently increase as well.
+However, the rank of the total Hessian will increase with the number of observations and so the number of non-zero eigenvalues consequently increase as well.
These eigenvalues will be smaller than the summed eigenvalues (because the eigenvectors may not be contained in the same subspace), but they will be larger than the smallest eigenvalue of the per-observation Hessian.
This will then lead to a steeper spectrum and a larger condition number of the total Hessian.
Thus, when the batch size \(B\) is large, the condition number of the sum of GGN Hessians will increase and the preconditioner will become less effective.
-How much the condition number increases with the batch size depends on the eigenspectrum of the per-observation Hessians (i.e., how close the eigenvectors are to each other).\sidenote{MAYBE I SHOULD REWRITE THIS PARAGRAPH, IT'S KINDA HARD TO EXPLAIN.}
+How much the condition number increases with the batch size depends on the eigenspectrum of the per-observation Hessians (i.e., how close the eigenvectors are to each other).
Unfortunately, if we want to use the pivoted Cholesky decomposition to approximate \(\sum^B_i \bm J_i\T \bm H_i \bm J_i\), we are largely trying to approximate the higher eigenvectors of the sum.
In order for this to be the case, we will want to oversample vectors from this matrix which correspond to the largest eigenvalues.
@@ -495,7 +496,6 @@ The null hypothesis of the test is that the samples are drawn from the same dist
This means that sample normality is rejected if the p-value of the test is less than the significance level of the test.
However, this test is very sensitive to the number of samples used to compute the statistic, and so a large number of samples will lead to a very low p-value even if the samples are drawn from the same distribution.
Since we can sample an arbitrarily large number of samples from the approximate posterior, this may lead to low p-values even when sampling succeeds.
-\marginnote{MAYBE ADD A LIMITATIONS / CONCLUSION SECTION AT THE END OF THIS CHAPTER?}
% \section{Limitations}
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -16,7 +16,7 @@ As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model
\end{align}
where \(p(\bm y)\) is known as the evidence. It's also known as the \emph{marginal likelihood}, since it's obtained by marginalising out \(\bm \theta\) from the likelihood, as per
\begin{align}
- p(\bm y) ={} & \int_{\bm \theta} p(\bm \theta \given \bm y)\,d\bm \theta.
+ p(\bm y) ={} & \int_{\bm \theta} p(\bm y \given \bm \theta)\,p(\bm \theta)\,d\bm \theta.
\end{align}
This term is typically intractable, since it involves integrating over the entire parameter space.
For this reason, we typically approximate the posterior by some tractable distribution, such as a Gaussian, and then use this to approximate the marginal likelihood.
@@ -73,12 +73,14 @@ Combining \cref{eq:laplace-log-marginal,eq:normal-log-likelihood,eq:normal-log-p
\\ & - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2 \nonumber
\\ & + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda, \nonumber
\end{align}
-where it can be seen that there are constant terms do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\).
+where it can be seen that there are constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\).
While these terms are not necessary for the optimisation problem, we choose to include them.
-In particular, the terms \(D /\,2 \log(2 \pi)\) are important for comparing models with different numbers of parameters.\sidenote{We note that these constant terms are not included in the computation of the \texttt{log\_marginal\_likelihood} function from the \texttt{laplace} Python library~\cite{daxberger2021laplace}. In spite of this, \textcite{immer2021scalable} use this library in their experiments, where they compare models with different numbers of parameters.}
-Furthermore, the other terms are important for interpreting the log-marginal as a likelihood.
+% In particular, the terms \(D /\,2 \log(2 \pi)\) are important for comparing models with different numbers of parameters.
+% \sidenote{We note that these constant terms are not included in the computation of the \texttt{log\_marginal\_likelihood} function from the \texttt{laplace} Python library~\cite{daxberger2021laplace}. In spite of this, \textcite{immer2021scalable} use this library in their experiments, where they compare models with different numbers of parameters.}
+% Furthermore, the other terms are important for interpreting the log-marginal as a likelihood.
+These terms allow us to interpret the log-marginal as a likelihood and do not significantly affect performance.
From \cref{eq:expanded-log-marginal}, we can see that the Laplace log-marginal likelihood requires computation of the log-determinant of the posterior precision matrix.
-This is a challenging task, as there is no closed-form expression for the log-determinant of a matrix without instantiating the matrix.
+This is a challenging task, as there is no \emph{cheap} closed-form expression for the log-determinant of a matrix without instantiating the matrix.
However, we can compute an upper bound on the log-determinant of the posterior precision matrix, which is useful for our optimisation problem which minimises the log-determinant term.
% \begin{tabular}{ll}
@@ -100,7 +102,7 @@ However, we can compute an upper bound on the log-determinant of the posterior p
\section{The Determinant Bound}
\label{sec:determinant-bound}
-We want to compute an upper bound on the log-determinant of the Hessian of the negative log-posterior \(\bm \Lambda \in \reals^{D \times D}\) such that we have \(\log \det \bm \Lambda \leq B_D(\mu_1, \mu_2, \beta)\).
+We want to compute an upper bound on the log-determinant of the posterior precision \(\bm \Lambda \in \reals^{D \times D}\) such that we have \(\log \det \bm \Lambda \leq B_D(\mu_1, \mu_2, \beta)\).
\textcite{bai1996bounds} provide one such bound,
\begin{align}\label{eq:bai-golub-bound}
B_D(\mu_1, \mu_2, \beta) \coloneqq{} & \begin{pmatrix} \log \beta & \log t \end{pmatrix}
@@ -108,7 +110,7 @@ We want to compute an upper bound on the log-determinant of the Hessian of the
\begin{pmatrix} \mu_1 \\ \mu_2 \end{pmatrix},
\end{align}
with \(\mu_1 = \Tr(\bm \Lambda)\), \(\mu_2 = \norm{\bm \Lambda}^2_F = \Tr\left(\bm \Lambda^2\right)\), \(t = \frac{\beta \mu_1 - \mu_2}{\beta D - \mu_1}\), where \(\beta\) is an upper bound on the largest eigenvalue of \(\bm \Lambda\).
-This bound involves approximating the log eigenvalues of \(\bm \Lambda\) with a second-order polynomial.
+This bound involves approximating the log eigenvalues of \(\bm \Lambda\) with a third-order polynomial.
Since \(B_D(\mu_1, \mu_2, \beta)\) is an upper bound on the log-determinant of \(\bm \Lambda\), we can use this to compute a lower bound on the log-marginal likelihood, as per \cref{eq:expanded-log-marginal}.
This is effective because we want to maximise the log-marginal.
@@ -125,14 +127,14 @@ The former can then be bounded by \cref{eq:bai-golub-bound} with an upper bound
We then compute the trace of this term as
\begin{align}\label{eq:laplace-trace-bound}
\Tr{\bm \Lambda} ={} & \sum_{k=1}^{K} (\lambda_k + \alpha) + (D - K) \alpha
- \\ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha
+ \\ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha \label{eq:laplace-trace-bound-2}
\\ ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha, \nonumber
\end{align}
and for the square trace we have
\begin{align}\label{eq:laplace-square-trace-bound}
\sum_{k=1}^{K} (\lambda_k + \alpha)^2 ={} & \Tr\left(\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)^2\right) + 2 \alpha \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha^2.
\end{align}
-We can bound the log-determinant of \(\bm \Lambda\) using \cref{eq:log-determinant-decomposition,eq:laplace-trace-bound,eq:laplace-square-trace-bound} as
+We can bound the log-determinant of \(\bm \Lambda\) using \cref{eq:log-determinant-decomposition,eq:laplace-trace-bound-2,eq:laplace-square-trace-bound} as
\begin{align}\label{eq:laplace-log-determinant-bound}
\log \det \bm \Lambda \leq{} & B_K\left(\mu_1, \mu_2, \beta\right) + (D - K) \log \alpha,
\\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,
@@ -172,16 +174,12 @@ Empirically, we find that, without the trace rearrangement in \cref{eq:trace-rea
\subsection{Scaling}
In standard negative log-likelihood training, the only term in the loss is \(-\log p(\bm y \given \bm \theta)\).
-For regression, the loss function which has a Bayesian interpretation as the negative log-likelihood is the sum of squared errors (SSE), though often the mean squared error is used instead, since computing the correct gradient just involves scaling the step size \(\eta\) up by a constant (specifically, the number of observations \(N\)).
-For negative log-posterior training, the loss is \(-\log p(\bm y \given \bm \theta) - \log p(\bm \theta)\), which is the sum of the negative log-likelihood and the negative log-prior.
-Since, for regression, the prior is usually a Gaussian, the negative log-prior is \(\alpha \norm{\bm \theta}^2\).
-This means that scaling the step size by \(N\) is no longer sufficient to scale the loss, since ratio of the prior regularisation and the likelihood is now different.
-As such, we need to multiply the prior precision \(\alpha\) by \(1 / N\) to scale the prior term in the loss to the SSE negative log-likelihood.
-However, for the log-marginal likelihood, scaling the terms is no longer possible, since scaling the loss by a constant \(c\) will not affect the gradients of the log-determinant term.
+For regression, the loss function which has a Bayesian interpretation as the negative log-likelihood is the sum of squared errors (SSE), though often the mean squared error is used instead, since computing the correct gradient just involves scaling the step size \(\eta\) up by the number of observations \(N\).
+For negative log-posterior training, the same scaling can be applied with a Gaussian prior by dividing the prior precision by \(N\).
+However, for the log-marginal likelihood, scaling the terms is no longer possible, since scaling the loss by a constant will not affect the gradients of the log-determinant term.
This means that it is important to compute the negative log-likelihood correctly as the SSE, since the gradients of the log-marginal likelihood are not invariant to scaling the log-marginal by a constant.
-
-For computing the posterior precision using a mini-batch, we thus have to scale the likelihood term by \(N / B\) to account for the fact that the likelihood is computed on a mini-batch of size \(B\) instead of the full dataset of size \(N\), yielding \(N / B \sum_{i=1}^B \bm J_i \bm H_i \bm J_i + \alpha \identity\).
-We compute the log-determinant of this mini-batched posterior precision as
+To compute the posterior precision using a mini-batch, we thus have to scale the likelihood term by \(N / B\) to account for the fact that the likelihood is computed on a mini-batch of size \(B\) instead of the full dataset of size \(N\), yielding \(N / B \sum_{i=1}^B \bm J_i \bm H_i \bm J_i + \alpha \identity\).
+The log-determinant of this mini-batched precision is equivalent to
\begin{align}
\log\det \left( \frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \alpha \identity \right)
={} & \log\det \left( \frac{N}{B} \left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \frac{B}{N} \alpha \identity \right) \right)\hspace{-0.05cm}
@@ -195,7 +193,7 @@ We do this by scaling the trace and square traces for \cref{eq:laplace-log-deter
\\ \Tr\left(\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right) ={}& \frac{N^2}{B^2} \Tr\left(\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right). \label{eq:laplace-square-trace-scaling}
\end{align}
By combining \cref{eq:laplace-log-determinant-bound,eq:bai-golub-bound,eq:hutchinson-trace,eq:hutchinson-square-trace,eq:trace-rearrangement,eq:laplace-trace-scaling,eq:laplace-square-trace-scaling}, we can therefore compute a differentiable upper bound on the log-marginal likelihood under the Laplace approximation using only Hessian-vector products without the need to instantiate the posterior precision.
-\marginnote{MAYBE ADD A LIMITATIONS / CONCLUSION SECTION AT THE END OF THIS CHAPTER?}
+% \marginnote{MAYBE ADD A LIMITATIONS / CONCLUSION SECTION AT THE END OF THIS CHAPTER?}
% \subsection{Limitations}
% Estimating the trace and square trace of the Hessian matrix is the most computationally expensive part of estimating the log-marginal under the Laplace approximation, since this involves computing the GGN-vector product.
diff --git a/figures/sin_training_alpha.pdf b/figures/sin_training_alpha.pdf
Binary files differ.
diff --git a/figures/sin_training_alpha_small.pdf b/figures/sin_training_alpha_small.pdf
Binary files differ.
diff --git a/main.tex b/main.tex
@@ -264,8 +264,7 @@
% The bibliography needs to be compiled with biber using your LaTeX editor, or on the command line with 'biber main' from the template directory
-\defbibnote{bibnote}{Here are the references in citation order.\par\bigskip} % Prepend this text to the bibliography
-\printbibliography[heading=bibintoc, title=Bibliography, prenote=bibnote] % Add the bibliography heading to the ToC, set the title of the bibliography and output the bibliography note
+\printbibliography[heading=bibintoc, title=Bibliography] % Add the bibliography heading to the ToC, set the title of the bibliography and output the bibliography note
%----------------------------------------------------------------------------------------
% APPENDIX
