commit 4a5c2624aead9e4c12c24b2b17d16c210f5949c7
parent 02e043e7c86d815899acf77cc56644a8db81c614
Author: Silas Brack <s174433@student.dtu.dk>
Date: Thu, 9 Feb 2023 15:58:29 +0100
So so so close.
Diffstat:
7 files changed, 94 insertions(+), 102 deletions(-)
diff --git a/chapters/ablation_experiments.tex b/chapters/ablation_experiments.tex
@@ -18,7 +18,7 @@ This is not trivial, though, as it would require querying the memory usage from
Instead, we simply measure the value of \(N\) at which the program fails from an out-of-memory error instead, as a proxy for memory usage.
This experiment should then answer two questions:
\begin{enumerate}
- \item How much larger problems can we solve if we don't need to instantiate the whole Hessian?
+ \item How much larger problems can we solve if we do not need to instantiate the whole Hessian?
\item How much faster is it to compute the GVP versus, say, the HVP or the explicit Hessian-vector product?
\end{enumerate}
@@ -50,7 +50,7 @@ Additionally, while the wall clock time for computing the Hessian-vector product
This may be due to the explicit Hessian, at this size, being too large to fit in the cache of the GPU.
The largest number of parameters tested is 1\,B.
-Above 1\,B parameters, even implicit methods fail, although this failure occurs in the instantiation of a single random vector of that size (so it would fail on this hardware for any method, since we can't even store the \(\bm v\) nor the product \(\bm K \bm v\)).
+Above 1\,B parameters, even implicit methods fail, although this failure occurs in the instantiation of a single random vector of that size (so it would fail on this hardware for any method, since we can not even store the \(\bm v\) nor the product \(\bm K \bm v\)).
Notably, however, these implicit Hessian-vector product methods allowed for computation of matrix products for problems three orders of magnitude larger, and computation of solves for problems four orders of magnitude larger, than the explicit methods.
This is a significant improvement in scalability, and is a key reason why we use these methods in our experiments.
Additionally the implicit inverse products are significantly faster than explicit inversion, even for small problems.
@@ -63,7 +63,7 @@ Performance also appears identical regardless of whether the explicit inverse is
\label{fig:hessian-profile-ratios}
\end{marginfigure}
-However, due to the scale of the figure, it isn't clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product.
+However, due to the scale of the figure, it is not clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product.
We can then plot the speedup from going from the HVP to the GVP as a function of the number of model parameters.
This can be calculated as the ratio of the wall-clock time of the HVP to the wall-clock time of the GVP.
Thus, larger values suggest greater performance gains for the GVP, while values closer to one suggest there is no performance gain.
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -20,7 +20,7 @@ Finally, we showed that it is possible to perform the full Laplace approximation
We demonstrated on a simple sine wave regression problem and an MNIST classification problem that our method can be used to train the full Laplace approximation and yield predictive performance which rivals deterministic training.
However, we found that the sampling procedure is too slow too sample outside of a post-hoc setting, and may become a significant bottleneck for inference even under post-hoc Laplace on large-scale problems.
-Overall, while we have developed novel methods for training and inference in the full Laplace approximation, we have not yet demonstrated its improved performance over the diagonal Laplace approximation nor its practical application to large-scale problems.
+Overall, while we have developed novel methods for training and inference in the full Laplace approximation, we have yet to demonstrate its improved performance over the diagonal Laplace approximation nor its practical application to large-scale problems.
\section[Outlook and Future Developments]{Outlook and Future Developments}[Future Work]
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -10,13 +10,6 @@
% Specific applications of Laplace approximation in Bayesian neural networks.
% Why this problem is important.
-% The Laplace approximation is one such method which enables post-hoc estimation of the posterior distribution by assuming a normal distribution for the posterior and using the Hessian to approximate the posterior precision.
-% Additionally, since the Laplace approximation deduces uncertainty rather than learning it, it has many advantages over other Bayesian methods that attempt to learn variance.
-% Due to its construction, the Laplace approximation allows for training by maximising the marginal likelihood, which facilitates model selection and hyperparameter optimisation.
-% However, the Laplace approximation typically relies on crude approximations to the posterior precision to make it computationally feasible.
-% When we perform the Laplace approximation using this full posterior precision, we call this the full Laplace approximation.
-% Conventional full Laplace uses the full posterior precision matrix, which cannot be instantiated for all but toy problems, since it scales quadratically with the number of parameters.
-
Bayesian methods allow for estimates of uncertainty which enable more efficient usage of data (e.g., via active learning) and avoid overfitting.
They accomplish this by using Bayes' theorem to model a distribution on the model parameters, which is then used to make predictions.
From here we can sample from this distribution to make predictions, thereby yielding uncertainty estimates, thus enabling the assessment of model predictive confidence.
@@ -48,55 +41,59 @@ The Laplace approximation has been applied in, for example, the prediction of ea
The primary advantages of the Laplace approximation are that it is simple to implement, effective at out-of-distribution detection~\cite{miani2022laplacian} and can be performed post-hoc on a pre-trained model.
This last point is important, since it also means that the Laplace approximation can be as cheap to train as the optimization algorithms for deterministic neural networks, while only requiring a small amount of additional computation to compute the posterior samples for inference.
-In the Laplace approximation corresponds to performing a second-order Taylor expansion of the log-posterior around the \emph{maximum a posteriori} (MAP) estimate of the model parameters as
-\begin{align}
+The Laplace approximation corresponds to performing a second-order Taylor expansion of the log-posterior around the \emph{maximum a posteriori} (MAP) estimate of the model parameters as
+\begin{align*}
\log p(\bm \theta\,|\,\bm y)
- \approx{} & \log p(\bm \theta_{\textsc{map}}\,|\,\bm y) + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y) (\bm \theta - \bm \theta_{\textsc{map}}). \nonumber
-\end{align}
-This is equivalent to approximating the posterior as a Gaussian distribution, with mean and precision given by the MAP estimate and the Hessian of the log-posterior respectively, i.e.,
-\begin{align}
- p(\bm \theta\,|\,\bm y) \approx{} & \normal\left(\bm \theta\,\Big|\,\bm \theta_{\textsc{map}}, \left(- \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y)\right)^{-1}\right). \nonumber
-\end{align}
+ \approx{} & \log p(\bm \theta_{\textsc{map}}\,|\,\bm y) + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y) (\bm \theta - \bm \theta_{\textsc{map}}).
+\end{align*}
+This is equivalent to approximating the posterior as a Gaussian distribution with mean and precision given by the MAP estimate and the Hessian matrix of the log-posterior respectively, i.e.,
+\begin{align*}
+ p(\bm \theta\,|\,\bm y) \approx{} & \normal\left(\bm \theta\,\Big|\,\bm \theta_{\textsc{map}}, \left(- \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y)\right)^{-1}\right).
+\end{align*}
To compute the posterior distribution of the model parameters, we thus need to compute the Hessian of the negative log-posterior loss with respect to the model parameters.
-Typically, this is done by storing the Hessian matrix explicitly and computing the inverse of this matrix.
-Since the Hessian matrix is a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters, this is intractable for large models, which can contain millions or billions of parameters, thus limiting the applicability of full Laplace outside of toy problems of a few thousand parameters.
-To overcome this, we can use a crude approximation, such as only storing the Hessian matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.
-Many of these Hessian approximations have successfully been applied to the Laplace approximation~\cite{ritter2018scalable,ritter2018online}, though are limited by the quality of the approximation.
-To overcome this, we propose a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
+Typically, this is done by storing the posterior precision matrix explicitly and computing the inverse of this matrix.
+Since the full Hessian matrix is a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters, this is intractable for large models, which can contain millions or billions of parameters, thus limiting the applicability of full Laplace outside of toy problems of a few thousand parameters.
+To overcome this, we can use a crude approximation, such as only storing the posterior precision matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.
+Many of these approximations to the posterior precision have successfully been applied to the Laplace approximation~\cite{ritter2018scalable,ritter2018online}, though are limited by the quality of the approximation.
+To overcome this, we propose a method which performs the full Laplace approximation without requiring the precision to be stored explicitly.
\section{Large-Scale Laplace}
-We propose an implementation of the full Laplace approximation which doesn't require instantiation of the full posterior precision matrix.
-First, we would like to learn a posterior distribution over the model parameters by maximising either the marginal likelihood or the posterior distribution.
-Then, to perform inference using the Laplace approximation, we need to sample from the posterior distribution, which is approximated by a normal distribution with a precision given as the Hessian of the loss with respect to the model parameters.
-The standard procedure for sampling thus requires computing the covariance matrix as the inverse of this precision matrix, and then taking the square root of the covariance matrix.
-Then, we can sample from this posterior distribution to make predictions.
-To train using the negative log-marginal likelihood loss for the Laplace approximation, we must compute the log-determinant of the Hessian.
-However, for large models, which can contain millions or billions of parameters, we cannot instantiate the Hessian matrix, for reasons mentioned above.
-Instead, we would like to access the inverse square root of the Hessian without explicitly storing it.
-
-Our method will allow us to approximate the Hessian's inverse square root and log-determinant using only implicit Hessian-vector products (HVPs).
-These HVPs compute products of the Hessian matrix with a vector without explicitly storing the Hessian matrix itself.
-Since the HVP outputs a vector, the memory requirements are linear in the number of parameters, rather than quadratic, allowing us to perform the full Laplace approximation with significantly larger models.
-We use JAX~\cite{jax2018github}, a Python library for GPU-accelerated automatic differentiation and computational graph compilation, to efficiently compute Hessian-vector while only storing the implicit Hessian computational graph.
-JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the Hessian-vector product computation into efficient machine code.
-
-We will thus develop a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
-To do this, we will use JAX and XLA to compute Hessian-vector products without explicitly storing the Hessian matrix.
-We will then exploit these Hessian-vector products to compute the log-determinant and inverse square root of the Hessian.
-These quantities enable us to perform training and inference using the Laplace approximation.
-An overview of the Laplace approximation and our contributions is shown in \cref{fig:overview}.
+We propose an implementation of the Laplace approximation which does not require instantiation of the full posterior precision matrix.
+In general, the Laplace approximation has two steps.
+First, we would like to learn a posterior distribution over the model parameters by maximising either the posterior or the marginal likelihood.
+Maximising the posterior during training is trivial, since it has a cheap closed-form gradient.
+However, maximising the marginal likelihood is more difficult.
-\begin{marginfigure}[-15cm]
+\begin{marginfigure}
\centering
\input{figures/overview.tex}
- \caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Hessian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the Hessian. Our contributions allow us to forgo the typical crude approximations of the Hessian by performing the full Laplace approximation without explicitly storing the Hessian matrix.}
+ \caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Jacobian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the posterior precision. Our contributions allow us to forgo the typical crude approximations of the precision by performing the LA using the full posterior precision without instantiating it explicitly.}
\label{fig:overview}
\end{marginfigure}
-% \begin{figure}
-% \centering
-% \input{figures/overview.tex}
-% \caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Hessian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the Hessian. Our contributions allow us to forgo the typical crude approximations of the Hessian by performing the full Laplace approximation without explicitly storing the Hessian matrix.}
-% \label{fig:overview}
-% \end{figure}
+As we will prove later, the Laplace log-marginal likelihood is given by, ignoring constant terms,
+\begin{align*}
+ \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2 - \frac{1}{2} \alpha \norm{\bm \theta}^2 - \frac{1}{2} \log \det \bm \Lambda,
+\end{align*}
+where \(\bm \Lambda\) is the posterior precision.
+The first two terms are easy to compute (and differentiate), but computing the log-determinant of very large matrices is generally infeasible, since it typically has \(O(D^3)\) complexity or requires storing the matrix in memory.
+To overcome this, we thus need to estimate the log-determinant of the posterior precision matrix without instantiating it.
+
+The second step of the Laplace approximation is to perform inference and make predictions from the posterior predictive distribution.
+We estimate the posterior predictive via Monte Carlo sampling by sampling from the posterior distribution and averaging the model predictions over the posterior samples.
+Since the posterior distribution is approximated by a Gaussian, we sample from the posterior as
+\begin{align*}
+ \bm \varepsilon = \bm \theta_{\textsc{map}} + \bm \Lambda^{-1/2} \bm \varepsilon_0,
+\end{align*}
+which effectively only requires computing the inverse square root of the posterior precision matrix.
+This computation is also intractable for large models, since it usually involves computing the Cholesky decomposition of the precision, which suffers from the same performance issues as the computation of the log-determinant.
+
+As stated before, we will present a method which performs the full Laplace approximation without explicitly storing the posterior precision matrix.
+Our method approximates the posterior precision's inverse square root and log-determinant using only implicit Jacobian-vector products (JVPs).
+These JVPs compute products of the posterior precision matrix with a vector without storing the matrix itself in memory.
+Since this product outputs a vector, the memory requirements are linear in the number of parameters (rather than quadratic), allowing us to perform the full Laplace approximation with significantly larger models.
+We use JAX~\cite{jax2018github}, a Python library for GPU-accelerated automatic differentiation and computational graph compilation, to efficiently compute Jacobian-vector products while only storing the implicit computational graph.
+JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the JVP computation into efficient machine code.
+The computation of the log-determinant and inverse square root of the precision will thus enable us to perform training and inference using the Laplace approximation.
+An overview of the Laplace approximation and our contributions is shown in \cref{fig:overview}.
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -23,8 +23,8 @@ These two cases correspond to the mean-squared error (MSE) and cross-entropy (CE
Due to the assumption of independence between observations, we can factorise the likelihood and the negative log-likelihood loss as
\begin{align}\label{eq:negative-log-likelihood}
- p(\bm y \given \bm \theta) ={} & \prod_i^N p(y_i \given \bm \theta),
- \\ \mathcal{L}(\bm\theta) \equiv \mathcal{L}(\bm \theta, \bm y) \coloneqq -\log p(\bm y \given \bm \theta) ={}& - \sum_i^N \log p(y_i \given \bm \theta).
+ p(\bm y \given \bm \theta) ={} & \prod_i^N p(\bm y_i \given \bm \theta),
+ \\ \mathcal{L}(\bm\theta) \equiv \mathcal{L}(\bm \theta, \bm y) \coloneqq -\log p(\bm y \given \bm \theta) ={}& - \sum_i^N \log p(\bm y_i \given \bm \theta).
\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.
@@ -54,7 +54,7 @@ However, gradient descent is prone to getting stuck in local minima, and can be
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.
+SGD helps avoid getting stuck in local minima, but, like gradient descent, does not take into account the curvature of the loss function, and so can overshoot the minimum.
\subsection{Adam}
\label{sec:adam}
@@ -83,7 +83,7 @@ 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{Maximum a Posteriori Estimation}
+\section{Maximum a Posteriori Estimation}[MAP 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.
@@ -96,7 +96,7 @@ It is obtained from Bayes' theorem, which allows for uncertainty quantification
where \(p(\bm \theta \given \bm y)\) is the posterior distribution over the parameters, \(p(\bm \theta)\) is the prior distribution over parameters, \(p(\bm y \given \bm \theta)\) is the likelihood of the data given the parameters, and \(p(\bm y)\) is the marginal likelihood of the data.
The marginal likelihood is obtained by marginalising over the parameters, i.e., \(p(\bm y) = \int p(\bm y \given \bm \theta)\,p(\bm \theta)\,d\bm \theta\).
Due to this integral, the marginal likelihood is often intractable.
-However, since the posterior distribution is proportional to the product of the likelihood by the prior, we can obtain the unnormalised posterior from the likelihood and the prior\marginnote{Since the marginal likelihood is a normalisation constant and doesn't depend on the parameters, it does not affect the optimisation problem.}
+However, since the posterior distribution is proportional to the product of the likelihood by the prior, we can obtain the unnormalised posterior from the likelihood and the prior\marginnote{Since the marginal likelihood is a normalisation constant and does not depend on the parameters, it does not affect the optimisation problem.}
\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}
@@ -107,10 +107,10 @@ Similarly to the negative log-likelihood, we can factorise the negative log-post
\begin{align}\label{eq:negative-log-posterior}
\begin{split}
\mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
- \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}.
+ \\ ={}& \underbrace{-\sum_i^N \log p(\bm y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}.
\end{split}
% \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
- % \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}. \nonumber
+ % \\ ={}& \underbrace{-\sum_i^N \log p(\bm y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}. \nonumber
\end{align}
Here the regularisation term is given by the prior distribution \(p(\bm\theta)\).
Due to the black-box nature of neural networks, it is often difficult to define a prior distribution over the parameters using previous knowledge.
@@ -193,13 +193,13 @@ Typically, the loss function is considered as the negative log-likelihood or neg
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
\begin{align}\label{eq:negative-log-posterior-hessian}
\begin{split}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
- \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta)
- \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta).
+ \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(\bm y_i \given \bm \theta) - \log p(\bm \theta) \right)
+ \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta)
+ \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta).
\end{split}
- % \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
- % \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
- % \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
+ % \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(\bm y_i \given \bm \theta) - \log p(\bm \theta) \right)
+ % \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
+ % \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
\end{align}
In this construction, the first term is the Hessian of the negative log-likelihood, i.e., the Hessian of the loss function we are interested in.
For the common prior given by a zero-mean Gaussian \(\bm \theta \sim \normal(\bm 0, \alpha^{-1} \identity)\), the log-prior regularisation term is given by \(\nabla^2_{\bm\theta} \log p(\bm \theta) = - \alpha \identity\).
@@ -252,15 +252,15 @@ This means that the GGN approximation is an approximation to the Hessian of the
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}
\begin{split}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta)
- \\ \approx{}& \sum_{i=1}^N {\underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i\T}}\T \underbrace{\left(\nabla^2_f \mathcal{L}(y_i, f(\bm x_i, \bm \theta))\right)}_{\bm H_i} \underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i}.
+ \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \sum_i^N \nabla^2_{\bm\theta} \log p(\bm y_i \given \bm \theta)
+ \\ \approx{}& \sum_{i=1}^N {\underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i\T}}\T \underbrace{\left(\nabla^2_f \mathcal{L}(\bm y_i, f(\bm x_i, \bm \theta))\right)}_{\bm H_i} \underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i}.
\end{split}
\end{align}
For a normal prior with precision \(\alpha\) on \(\bm \theta\), we obtain the GGN approximation of the posterior precision
\begin{align}\label{eq:laplace-precision-ggn}
\bm \Lambda \approx \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity,
\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))\).
+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}(\bm y_i, f(\bm x_i, \bm \theta))\).
For mean square error (MSE) loss (e.g. for regression), we obtain the simple expression the GGN approximation of the Hessian
%
@@ -311,6 +311,16 @@ However, the GGN approximation is better behaved than the exact Hessian since it
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 algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
+In practice, as we explained in \cref{ch:intro}, we are looking for methods which access the product of the GGN matrix with a vector, rather than the GGN matrix itself.
+This then means we want to compute
+\begin{align}
+ \begin{split}
+ \bm \Lambda \bm v ={} & \left(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\right) \bm v
+ \\ ={}& \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i \bm v + \alpha \bm v,
+ \end{split}
+\end{align}
+such that we actually compute the product of the Jacobian of the neural network (as defined in \cref{eq:negative-log-likelihood-ggn}) by a vector \(\bm v\).
+
\subsection{The Spectrum of the GGN Approximation}[The Spectrum of the GGN]
\label{sec:spectrum-ggn}
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -4,7 +4,7 @@
We will now discuss how to perform inference by sampling from the Laplace posterior distribution.
We first give an overview of how we typically sample from the Laplace posterior and discuss the memory constraints of this sampling method when using the full posterior precision matrix.
-In \cref{sec:ciq}, we propose an algorithm for sampling from the Laplace posterior that doesn't require storing the full precision matrix, allowing us to scale sampling to larger models.
+In \cref{sec:ciq}, we propose an algorithm for sampling from the Laplace posterior that does not require storing the full precision matrix, allowing us to scale sampling to larger models.
Finally, we discuss how to mitigate some of the issues of our method (\cref{sec:preconditioning}) and propose a practical technique for determining whether sampling was successful (\cref{sec:sampling-evaluation}).
\section{Sampling from the Laplace Posterior}[Posterior Sampling]
@@ -189,18 +189,11 @@ This integral is approximated by
\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.
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.
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\).
-% We can thereby now compute the quadrature weights and shifts, as seen in \cref{alg:quad}.
\begin{figure}
\centering
@@ -291,7 +284,7 @@ If we have a rank-one preconditioner whose only eigenvector is the largest eigen
Thus, in the best-case scenario, you can then eliminate the largest eigenvalue of \(\bm \Lambda\) by using a rank-one preconditioner.
Thus, if we have a rank-\(K\) preconditioner, then we can eliminate at most the \(K\) challenging (largest / smallest) eigenvalues of \(\bm \Lambda\) by using a rank-\(K\) preconditioner.
The steeper the spectrum of \(\bm \Lambda\), the more effective this preconditioning will be (however, in this case, \(\bm \Lambda\) will probably also tend to be worse conditioned).
-% This analogy can be extended in a hand-wavy way to an imperfect preconditioner which doesn't exactly approximate the eigenvalue
+% This analogy can be extended in a hand-wavy way to an imperfect preconditioner which does not exactly approximate the eigenvalue
For a preconditioner with a condition number of \(\kappa^\prime\), we can at best obtain a preconditioned system with conditioning \(\kappa / \kappa^\prime\).
This is because we want the highest eigenvalue of the preconditioner to neutralise (i.e., align with) the highest eigenvalue of \(\bm \Lambda\), and the same for the lowest eigenvalues.
@@ -359,7 +352,7 @@ This identity can be used to compute the inverse of the GGN matrix with the prio
\begin{align}
\left( \alpha \identity + \bm J\T \bm H \bm J \right)^{-1} = \alpha^{-1} \identity - \alpha^{-2} \bm J\T \left( \bm H^{-1} + \alpha^{-1} \bm J \bm J\T \right)^{-1} \bm J,
\end{align}
-where \(\alpha\) is the prior precision and \(\bm H\) and \(\bm J\) are the Hessian and the Jacobian as per \cref{eq:laplace-precision-ggn} evaluated at a single observation.
+where \(\alpha\) is the prior precision and \(\bm H\) and \(\bm J\) are the Hessian and the Jacobian as per \cref{eq:negative-log-likelihood-ggn} evaluated at a single observation.
Here, to exploit the low-rank structure of the preconditioner, we only compute the GGN matrix over a single observation.
As such, for this preconditioner to be effect, we assume that the eigenvectors of the GGN matrix are constant across the dataset.
@@ -379,7 +372,7 @@ However, in the case of autoencoders, the loss can be the MSE loss, which has a
Thus, if \(\bm H\) is positive definite (which is necessary for the Laplace approximation) and easy to invert, then it should be possible to use the Woodbury matrix identity to compute the inverse of this preconditioner.
This preconditioner is fast to compute, but may not be effective for when the eigenvectors of the GGN matrix are not constant across the dataset.
-However, due to the assumption of stationary eigenvectors and the effectiveness of the sub-sampled GGN preconditioner, we didn't end up using this preconditioner.
+However, due to the assumption of stationary eigenvectors and the effectiveness of the sub-sampled GGN preconditioner, we did not end up using this preconditioner.
Additionally, this preconditioner cannot assume a higher rank than the number of outputs \(O\), since adding more observations to the preconditioner will increase its rank to some unknown value and prevent the application of the Woodbury matrix identity.
Ideally, we would like to be able to find a preconditioner for which we can choose the rank of the low-rank approximation \(\bm A \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) without simply computing the GGN on a single observation (\(\bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) for some observation \(i\)) as in \cref{sec:woodbury-preconditioner}, while still allowing us to apply the Woodbury matrix identity, unlike in the case of the sub-sampling preconditioner in \cref{sec:sub-sampling-preconditioner} (\(\sum^B_{i=1} \bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\)).
@@ -486,7 +479,7 @@ Samples \(\bm \varepsilon\) which are normally distributed but do not have a mea
\end{split}
\end{align}
%
-This has the advantage of only requiring computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it doesn't require instantiation of the matrix.
+This has the advantage of only requiring computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it does not require instantiation of the matrix.
For the Laplace approximation, this can easily be computed as $\bm \Lambda = - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y)$.
Since this precision matrix can be computed exactly, this calculation can be performed to evaluate the whether a set of samples \(\bm \varepsilon\) are drawn from a distribution \(\normal(\bm \mu, \bm \Lambda^{-1})\) while only implicitly accessing this precision matrix in the evaluation via matrix-vector products.
To do so, we can compute the squared Mahalanobis distance of each sample \(\bm \varepsilon_i\) from the mean \(\bm \mu\) and compare it to the chi-squared distribution with \(D\) degrees of freedom.
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -10,11 +10,11 @@ Finally, we discuss the implementation of online Laplace with mini-batching usin
\section{Maximising the Evidence}
\label{sec:maximising-the-evidence}
-As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model is equal to
+As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model is
\begin{align}
p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta)\,p(\bm \theta)}{p(\bm y)},
\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
+where \(p(\bm y)\) is known as the evidence. It's also known as the \emph{marginal likelihood}, as it's obtained by marginalising \(\bm \theta\) from the likelihood, as per
\begin{align}
p(\bm y) ={} & \int_{\bm \theta} p(\bm y \given \bm \theta)\,p(\bm \theta)\,d\bm \theta.
\end{align}
@@ -28,7 +28,7 @@ To further motivate training on the evidence term, \textcite{fong2020marginal} s
This result shows that the marginal likelihood is a good proxy for the true cross-validation performance of a model.
Because of this, it can be used for model selection, since it will choose the model that performs best on average across all permutations of the data.
-Additionally, the marginal likelihood can be used to optimise hyperparameters, since it can be used to compute the gradient of the log-marginal with respect to the hyperparameters.
+Additionally, unlike the posterior, the marginal likelihood can be used to optimise both model parameters and hyperparameters, since we can meaningfully minimise the log-marginal with respect to the hyperparameters (and compute the respective gradients).
This is useful, since it allows us to use gradient-based optimisation methods to find the hyperparameters that maximise the marginal likelihood, as has been commonly done in Gaussian Processes~\cite{williams2006gaussian,damianou2013deep}.
This provides an alternative to finding the hyperparameters by grid search, which is computationally expensive.
@@ -58,10 +58,10 @@ Taking the log of this approximate Laplace marginal, we obtain
\end{align}
For the normal likelihood and normal prior, we have
\begin{align}\label{eq:normal-likelihood}
- p(\bm y \given \bm \theta) ={} & \normal(\bm y \given f_{\bm\theta}(\bm x), \rho^{-1} \identity) = \prod_{i=1}^N \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity),
- \\ \log p(\bm y \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity) \label{eq:normal-log-likelihood}
- \\ ={}& \sum_{i=1}^N \left( - \frac{O}{2} \log(2 \pi) - \frac{1}{2} \log \det (\rho^{-1} \identity) - \frac{1}{2} \rho \norm{y_i - f_{\bm\theta}(\bm x_i)}^2\right) \nonumber
- \\={}& - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2, \nonumber
+ p(\bm y \given \bm \theta) ={} & \normal(\bm y \given f_{\bm\theta}(\bm x), \rho^{-1} \identity) = \prod_{i=1}^N \normal(\bm y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity),
+ \\ \log p(\bm y \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(\bm y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity) \label{eq:normal-log-likelihood}
+ \\ ={}& \sum_{i=1}^N \left( - \frac{O}{2} \log(2 \pi) - \frac{1}{2} \log \det (\rho^{-1} \identity) - \frac{1}{2} \rho \norm{\bm y_i - f_{\bm\theta}(\bm x_i)}^2\right) \nonumber
+ \\={}& - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2, \nonumber
\end{align}
and
\begin{align}\label{eq:normal-prior}
@@ -72,7 +72,7 @@ where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of
Combining \cref{eq:laplace-log-marginal,eq:normal-log-likelihood,eq:normal-log-prior}, we expand the Laplace approximation to the log-marginal likelihood as
\begin{align}\label{eq:expanded-log-marginal}
\begin{split}
- \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2
+ \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2
\\ & - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2
\\ & + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda,
\end{split}
@@ -84,19 +84,20 @@ While these terms are not necessary for the optimisation problem, we choose to i
% 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 \emph{cheap} closed-form expression for the log-determinant of a matrix without instantiating the matrix.
+This log-determinant term gets minimised in our optimisation, which favours models with small eigenvalues in the Hessian of the loss, i.e., points of low curvature in the loss landscape, as we argued earlier in this section.
+This is a challenging computation, 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}
% Minimise & Maximise
-% \\\addlinespace 1. \(\rho / 2 \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2\) & 1. \(N O / 2 \log \rho\)
+% \\\addlinespace 1. \(\rho / 2 \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2\) & 1. \(N O / 2 \log \rho\)
% \\ 2. \(1 / 2 \log \det \bm \Lambda\) & 2. \(D \log \alpha\)
% \\ 3. \(\alpha / 2 \norm{\bm \theta}^2\)
% \end{tabular}
% The \texttt{laplace} Python library provides a function to compute the log-marginal likelihood of the Laplace approximation. The function is called \texttt{log\_marginal\_likelihood} and is given by
% \begin{align}
-% \log p(\bm y) \mathrel{\stackrel{\textsc{la}}{\approx}}{} & \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2
+% \log p(\bm y) \mathrel{\stackrel{\textsc{la}}{\approx}}{} & \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2
% \\ & - \frac{1}{2} \left( \log \det \bm \Lambda - D \log \alpha + \alpha \norm{\bm \theta}^2 \right). \nonumber
% \end{align}
% As compared to the exact log-marginal from \cref{eq:expanded-log-marginal}, this implementation is obtained by omitting the constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\).
@@ -199,13 +200,4 @@ We do this by scaling the trace and square traces for \cref{eq:laplace-log-deter
\Tr\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) ={} & \frac{N}{B} \Tr\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) \label{eq:laplace-trace-scaling},
\\ \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?}
-% \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.
-
-% Training hyperparameters
-% Speed for convolutions
-% Scalability to huge networks hasn't been tested
-% Rank estimation
+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 Jacobian-vector products without the need to instantiate the posterior precision.
diff --git a/figures/overview.tex b/figures/overview.tex
@@ -10,7 +10,7 @@
\definecolor{color2}{rgb}{0.94, 0.97, 1.0}
\definecolor{kaogreen}{HTML}{91C889}
\node (nn) [box,fill=color1,minimum height=1.0cm,minimum width=3cm] {Neural Network};
- \node (max) [below=of nn,yshift=0.1cm] {\(\argmax_{\bm \theta}\)};
+ \node (max) [below=of nn,yshift=-0.3cm] {\(\argmax_{\bm \theta}\)};
\node (posterior) [box,fill=color2,minimum width=1.8cm,minimum height=2.1cm,below=of max,xshift=-1.1cm,yshift=0.5cm] {\(\log p(\bm \theta \given \bm y)\)};
\node (marginal) [box,fill=color2,minimum width=1.8cm,minimum height=2.1cm,below=of max,xshift=1.1cm,yshift=0.5cm] {\(\overbrace{\myrel{\text{\normalsize \(\log p(\bm \theta \given \bm y)\)}}{\text{\normalsize \colorbox{kaogreen}{$\log \det \bm \Lambda$}}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
% \node (marginal) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=2cm,yshift=0.5cm] {\(\overbrace{\log p(\bm \theta \given \bm y) \quad \text{\colorbox{kaogreen}{$\log \det \bm \Lambda$}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
