commit 8762eb8db062f64465208efb4fbfb611b052f32c
parent c0eeb82fd776a42ddd7b18b2ef6a17d36bd126b9
Author: Silas Brack <s174433@student.dtu.dk>
Date: Tue, 7 Feb 2023 20:19:54 +0100
Implemented the Marco changes.
Diffstat:
5 files changed, 49 insertions(+), 47 deletions(-)
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -5,13 +5,13 @@
\section{Summary and Key Results}
In this project, we aimed to develop a scalable full Laplace approximation for Bayesian neural networks.
-In doing so, we developed a novel method for approximately sampling from the full Laplace posterior and methods for preconditioning the Hessian of the negative log posterior, which involved attempting to better understand the spectrum of the generalised Gauss-Newton Laplace posterior precision.
+In doing so, we developed a novel method for approximately sampling from the full Laplace posterior and methods for preconditioning the Hessian of the negative log-posterior, which involved attempting to better understand the spectrum of the generalised Gauss-Newton Laplace posterior precision.
To test the quality of our sampling procedure, we showed how to evaluate the quality of the samples it produces while only accessing the samples themselves and the desired posterior precision, enabling us and to choose a suitable preconditioner for it.
This evaluation method showed that, with an appropriate preconditioner, we can reliably produce samples which are Gaussian-distributed with the correct mean and covariance, as desired.
As such, we have shown that our approximate sampling method is a valid (albeit slow) method for sampling from the full Laplace posterior.
We also developed a novel method for training the full Laplace approximation by maximising the Laplace-approximate marginal likelihood of the data.
-To compute this evidence term, we developed a method for computing an upper bound on the log determinant of the Laplace posterior precision without storing the posterior precision itself.
+To compute this evidence term, we developed a method for computing an upper bound on the log-determinant of the Laplace posterior precision without storing the posterior precision itself.
This method allows us to compute a differentiable log-marginal likelihood bound, which we can train with standard gradient descent optimisation methods.
We showed that this method is a valid method for training the full Laplace approximation, and that it can be used to train the full Laplace approximation on a convolutional neural network.
However, we also found that the marginal training procedure to optimise the hyperparameters of the full Laplace approximation is not yet feasible, since the hyperparameters converge to unreasonable values.
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -62,11 +62,11 @@ First, we would like to learn a posterior distribution over the model parameters
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.
+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 allows us approximate the Hessian's inverse square root and log determinant using only implicit Hessian-vector products (HVPs).
+Our method allows us 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.
@@ -74,7 +74,7 @@ JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the
We have thus developed 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.
+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}.
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -15,7 +15,7 @@ Our goal is to find the parameters \(\bm\theta^*\) which best model the observed
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.}
+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\).
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.
@@ -153,7 +153,7 @@ Similarly to the negative log-likelihood, we can factorise the negative log-post
\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.
-Because of this, the prior is often chosen to be a Gaussian distribution with zero mean and a diagonal covariance matrix with some precision \(\alpha\), i.e., \(p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I)\).\sidenote{In the frequentist setting, this type of regularisation is known as weight decay or \(L_2\)-regularisation. Note that \(- 1 / N \log p(\bm y \given \bm \theta) - \gamma / 2 \norm{\bm \theta}^2\) is the typical formulation in this setting, where \(\gamma\) is the regularisation strength and is usually within the range of \(10^{-2}\) to \(10^{-4}\). In our setting, we are instead computing the log posterior \(-\log p(\bm y \given \bm \theta) - \alpha / 2 \norm{\bm \theta}^2\), and so we must multiply the regularisation strength by \(N\) to obtain \(\alpha\).}
+Because of this, the prior is often chosen to be a Gaussian distribution with zero mean and a diagonal covariance matrix with some precision \(\alpha\), i.e., \(p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I)\).\sidenote{In the frequentist setting, this type of regularisation is known as weight decay or \(L_2\)-regularisation. Note that \(- 1 / N \log p(\bm y \given \bm \theta) - \gamma / 2 \norm{\bm \theta}^2\) is the typical formulation in this setting, where \(\gamma\) is the regularisation strength and is usually within the range of \(10^{-2}\) to \(10^{-4}\). In our setting, we are instead computing the log-posterior \(-\log p(\bm y \given \bm \theta) - \alpha / 2 \norm{\bm \theta}^2\), and so we must multiply the regularisation strength by \(N\) to obtain \(\alpha\).}
The likelihood term is defined as in the frequentist setting (see \cref{sec:optimisation-frequentist}).
However, the posterior distribution itself is typically intractable, so approximations to the posterior are used instead.
@@ -173,7 +173,7 @@ In this project, we focus on the Laplace approximation (LA)~\cite{laplace1774mem
\label{sec:laplace}
In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI\@.
-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}}\).
+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)
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -96,13 +96,14 @@ it is then also trivial to compute the square root product by multiplying \(\bm
% ------------------------------------------
The first step of CIQ invokes Cauchy's integral formula.
-Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed circular contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as\marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
+Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as\marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
\begin{align}\label{eq:cauchy-integral-formula}
f(\bm K) ={} & \frac{1}{2 \pi i} \oint_\Gamma f(\tau) \left(\tau \identity - \bm K\right)^{-1}\,d\tau,
\end{align}
where \(\oint_\Gamma\) is the contour integral along the contour \(\Gamma\) parametrised in the complex variable \(\tau\).
Since, in our case, the eigenvalues of \(\bm K\) will all be real-valued, a circle centred on the real axis which encloses the minimum and maximum eigenvalues of \(\bm K\) will suffice.\marginnote{If a real-valued matrix is symmetric, its eigenvalues are real. In our case, \(\bm K\) will be the precision matrix obtained from the Laplace approximation.}
-To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).\sidenote{Krylov subspace methods are a class of iterative methods which estimate \(f(\bm K) \bm v\) for some matrix \(\bm K\) and vector \(\bm v\) by using a Krylov subspace. The Krylov subspace is a subspace spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), given by \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\left\{ \bm v, \bm K \bm v, \bm K^2 \bm v, \ldots, \bm K^{j-1} \bm v \right\}\).}
+To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).
+\sidenote{Krylov subspace methods are a class of iterative methods which estimate certain properties (e.g., the rank, eigenvalues, span) of some matrix \(\bm K\) by using a Krylov subspace for some vector \(\bm v\). The Krylov subspace is a subspace spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), given by \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\left\{ \bm v, \bm K \bm v, \bm K^2 \bm v, \ldots, \bm K^{j-1} \bm v \right\}\).}
% ------------------------------------------
% QUADRATURE
@@ -125,8 +126,9 @@ We can then apply Cauchy's integral formula to the function \(f(\bm K) = \bm K^{
\bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau \label{eq:ciq-integral-tau}
\\ \approx{}& \frac{1}{2 \pi i} \sum_{q=1}^Q \Tilde{w}_q \tau_q^{-1/2} \left( \tau_q \identity - \bm K \right)^{-1}, \nonumber %\label{eq:ciq-quadrature-tau}
\end{align}
-where \(\tau_1, \ldots, \tau_Q\) are the sampled quadrature points and \(\Tilde{w}_1, \ldots, \Tilde{w}_Q\) are the quadrature weights.
-Since we will be sampling a fixed number of quadrature points from the circle, to improve the closeness of the approximation of the integral, it would be beneficial to choose a closed circle which intersects the real axis \emph{at} the minimum and maximum eigenvalues, since this will avoid sampling quadrature points which lie outside the spectrum of \(\bm K\).
+where \(\tau_1, \ldots, \tau_Q \in \Gamma\) are the sampled quadrature points and \(\Tilde{w}_1, \ldots, \Tilde{w}_Q \in \reals\) are the quadrature weights.
+% Since we will be sampling a fixed number of quadrature points from the circle, to improve the closeness of the approximation of the integral, it would be beneficial to choose a closed circle which intersects the real axis \emph{at} the minimum and maximum eigenvalues, since this will avoid sampling quadrature points which lie outside the spectrum of \(\bm K\).
+Since we will be sampling a fixed number of quadrature points from the circle, the closer the contour is to the singular points (i.e., the eigenvalues), the better the approximation to the integral will be.
% ------------------------------------------
% JACOBI ELLIPTICAL FUNCTIONS
@@ -283,9 +285,10 @@ To do so, one could hypothetically attempt to compute the eigenvectors of \(\bm
If we have a rank-one preconditioner whose only eigenvector is the largest eigenvector of \(\bm \Lambda\), then the eigenvalue corresponding to this eigenvector becomes one and the second-largest eigenvalue becomes the largest (assuming that the second-largest eigenvalue is greater than one).
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, \(\Lambda\) will probably also tend to be worse conditioned).
+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
-For a preconditioner with a condition number of \(\kappa^\prime\), we can at best obtain a preconditioned system with conditioning \(\kappa / \kappa^\prime\).\sidenote{IS THIS STRICTLY TRUE?}
+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.
In CIQ, we will, in each iteration, perform one matrix-vector product \(\bm \Lambda \bm v\) and one preconditioner-vector product, as seen in \cref{alg:ciq}.%\sidenote{Because of this, to avoid the preconditioner becoming a performance bottleneck, the speed of the preconditioner-vector products should be significantly greater than that of the matrix-vector products.}
As such, if the preconditioner-vector product is similarly or more expensive to compute compared to the matrix-vector product, then preconditioning will significantly slow down the algorithm.
@@ -300,17 +303,18 @@ We will now discuss some potential choices for preconditioners to sample using C
One of the simplest preconditioners is similar to the second moment estimation used in Adam (\cref{sec:adam}).
If we calculate the GGN approximation of the Hessian by linearizing over the whole likelihood (as described in \cref{sec:practical-ggn}), we get
\begin{align}
- \bm P ={} & \nabla \nabla\T + \alpha \identity
- \\ \bm P^{-1} ={}& \beta \cdot \nabla \nabla\T + \alpha^{-1} \identity
- \\ \beta ={}& \left(\alpha \norm{\nabla}^2 + \norm{\nabla}^4\right)^{-1} - \left(\alpha \norm{\nabla}^2\right)^{-1},
+ \bm P ={} & \nabla \nabla\T + \alpha \identity \label{eq:adam-preconditioner}
+ \\ \bm P^{-1} ={}& \beta \cdot \nabla \nabla\T + \alpha^{-1} \identity \label{eq:adam-preconditioner-inverse}
+ \\ \beta ={}& \left(\alpha \norm{\nabla}^2 + \norm{\nabla}^4\right)^{-1} - \left(\alpha \norm{\nabla}^2\right)^{-1}, \label{eq:adam-preconditioner-beta}
\end{align}
where \(\nabla\) is the gradient of the loss with regard to the model parameters.
-The problem with using this approximation is that the rank of \(\bm P\) is 1.
+Since \cref{eq:adam-preconditioner} is in the form of a diagonal perturbation to an outer product, we have the closed-form inverse given by \cref{eq:adam-preconditioner-inverse,eq:adam-preconditioner-beta}.
+The problem with using this approximation is that the rank of the outer product is one, where this term determines the largest eigenvalue of the preconditioner.
This means that, best-case scenario, we can reduce the highest eigenvalue of the Hessian.
This occurs when the gradient corresponds to the direction of the highest precision (empirically, this may not be the case).
This would then leave the second-highest eigenvalue unaffected.
-If the rank of our precision matrix is greater than one, then this value seems likely to be similarly large to originally the largest eigenvalue.
-Therefore, the condition number of our problem will not be significantly reduced.
+% If the rank of our precision matrix is greater than one, then this value seems likely to be similarly large to originally the largest eigenvalue.
+% Therefore, the condition number of our problem will not be significantly reduced.
We therefore need to find a higher-rank approximation of the posterior precision.
\subsection{Sub-Sampling the Data}
@@ -325,24 +329,22 @@ Alternatively, we can also instead estimate the eigenvectors and eigenvalues of
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}\)).
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.
-% \[\textstyle{\sum^N_i} s_1 v_i / \alpha \leq \kappa \leq N s_1 / \alpha\]
-
As mentioned before, each msMINRES iteration performs one matrix-vector product and one preconditioner-vector product.
-However, even though the preconditioner is constructed in nearly the same way as the GGN matrix, the preconditioner is significantly cheaper to compute.
-For the subsampled GGN preconditioner, the preconditioner-vector product will then cost \(B / N \cdot t_{\mathrm{GGN}}\), where \(t_{\mathrm{GGN}}\) is the computational cost of performing the GGN-vector product.\sidenote{This analysis assumes that the computation of the GGN-vector product scales linearly in the number of observations used in the GGN matrix.}
+However, even though the preconditioner is constructed in nearly the same way as the sum of GGN matrices in the posterior precision, the preconditioner is significantly cheaper to compute. %, since we're only computing the GGN on a subset of the data.
+This is because the preconditioner-vector product will cost \(B / N \cdot t_{\mathrm{GGN}}\), where \(t_{\mathrm{GGN}}\) is the computational cost of performing the GGN-vector product.\sidenote{This analysis assumes that the computation of the GGN-vector product scales linearly in the number of observations used in the GGN matrix. Mathematically speaking, this is a valid assumption, due to the number of elementary operations that are performed in the algorithm. However, the efficient JAX XLA-compiled implementation of Jacobian-vector products can scale differently to na\"ive implementations, so the validity of the assumption is less clear in our case.}
Since we will have \(N \gg B\), then the cost of computing the preconditioner will be negligible compared to that of the GGN-vector product.
-This preconditioner can also be adjusted to cases where the smallest eigenvalue \(\alpha\) is not very small, but for which the largest eigenvalue is very large.
-In this case, we can compute this preconditioner with a larger value of \(\alpha\) and a large number of preconditioner observations.
-This will lead to a preconditioner which has a larger smallest eigenvalue than the true posterior precision, but whose largest eigenvalue is closer to that of the true posterior precision.
-This will lead to a preconditioner which can reduce the largest eigenvalues of the posterior precision.
-This may be useful, for instance, in problems for which the eigenvectors of the posterior precision are not constant across the dataset, but are instead highly variable.
-However, this preconditioner will be very expensive to compute, since it effectively requires computing the entire posterior precision.\sidenote{MEH Empirically, I found that this is useful for the sine problem, though, which seems to have a very constant posterior precision.}
+% This preconditioner can also be adjusted to cases where the smallest eigenvalue \(\alpha\) is not very small, but for which the largest eigenvalue is very large.
+% In this case, we can compute this preconditioner with a larger value of \(\alpha\) and a large number of preconditioner observations.
+% This will lead to a preconditioner which has a larger smallest eigenvalue than the true posterior precision, but whose largest eigenvalue is closer to that of the true posterior precision.
+% This will lead to a preconditioner which can reduce the largest eigenvalues of the posterior precision.
+% This may be useful, for instance, in problems for which the eigenvectors of the posterior precision are not constant across the dataset, but are instead highly variable.
+% However, this preconditioner will be very expensive to compute, since it effectively requires computing the entire posterior precision.\sidenote{MEH Empirically, I found that this is useful for the sine problem, though, which seems to have a very constant posterior precision.}
\subsection{Using the Woodbury Matrix Identity}[Woodbury]
\label{sec:woodbury-preconditioner}
-The Woodbury matrix identity is a useful identity for computing the inverse of a matrix which is the sum of two matrices of which one is low-rank and the other can easily be inverted.
+The Woodbury matrix identity~\cite{woodbury1950inverting} is a useful identity for computing the inverse of a matrix which is the sum of two matrices of which one is low-rank and the other can easily be inverted.
The identity is given by
\begin{align}
\left( \bm A + \bm B \bm C \bm D \right)^{-1} = \bm A^{-1} - \bm A^{-1} \bm B \left( \bm C^{-1} + \bm D \bm A^{-1} \bm B \right)^{-1} \bm D \bm A^{-1},
@@ -350,7 +352,7 @@ The identity is given by
where \(\bm A\) is a \(n \times n\) matrix and can easily be inverted, \(\bm B\) is a \(n \times m\) matrix, \(\bm C\) is a \(m \times m\) matrix, and \(\bm D\) is a \(m \times n\) matrix.
This identity can be used to compute the inverse of the GGN matrix with the prior precision as
\begin{align}
- \left( \alpha \identity + \bm J\T \bm H \bm J \right)^{-1} \bm v = \alpha^{-1} \identity - \alpha^{-2} \bm J\T \left( \bm H^{-1} + \alpha^{-1} \bm J \bm J\T \right)^{-1} \bm J,
+ \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.
@@ -366,7 +368,7 @@ In this way, we can increase the maximum eigenvalues of the preconditioner witho
Firstly, the Hessian of the loss with respect to the model outputs \(\bm H\) must be first inverted before the Woodbury matrix identity can be used.
To use the conjugate gradient method to invert this Hessian, we would have to perform nested conjugate gradient to also compute the inverse of the outer product term \(\bm H^{-1} + \alpha^{-1} \bm J \bm J\T\), which could be very expensive.
For regression tasks where the Hessian is diagonal, this is not a problem.
-For classification tasks, we can compute the inverse of the Hessian as a diagonal and outer product matrix, so this is also not a problem.
+For classification tasks using the cross-entropy loss, we can compute the inverse of the Hessian as a diagonal and outer product matrix, so this is also not a problem.
For models with a very large number of outputs, such as autoencoders, the matrix \(\bm H\) will be fairly large, and so computing its inverse manually can become expensive.
However, in the case of autoencoders, the loss can be the MSE loss, which has a diagonal Hessian, and so can be inverted easily.
@@ -403,7 +405,7 @@ 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 \(\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 \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:
Regarding the eigenvalues and eigenvectors, we consider a best-case scenario.
Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors.
@@ -421,6 +423,7 @@ In order for this to be the case, we will want to oversample vectors from this m
A common approach to sampling the largest eigenvalues is to sample vectors with probability proportional to the diagonal value~\cite{harbrecht2012low}.
However, it is not possible to generally and efficiently determine the diagonal of a matrix via matrix-vector products.
Often, it is necessary to multiply a vector of zeros with a value of one at the index at which to find the diagonal.
+This operation needs to be repeated as many times as there are elements in the diagonal with no shared computation between each element, which is very expensive.
Alternatively, we can approximate the diagonal of the GGN matrix as the element-wise product of the gradient with itself, similarly to the preconditioner described in \cref{sec:adam-preconditioner} and the GGN approximation computed in \cref{eq:fully-linear-ggn}, though this may not be the most accurate approximation.
In conclusion, to find the optimal parameters for the pivoted Cholesky preconditioner (the number of observations \(B\) and the rank of the pivoted Cholesky factorisation \(K\)), there are some rules of thumb to consider.
@@ -430,9 +433,8 @@ I've found that it's most effective to keep \(K\) low, such that the \(K \times
We thus need to choose \(K\) such that it is large enough to approximate enough large eigenvectors and small enough to ensure the condition number of the outer product is low.
Furthermore, you want enough observations \(B\) to be included in the preconditioner such that it's possible to accurately approximate the eigenvectors of the problem---it may be optimal to use the whole dataset, though this may worsen the conditioning of the outer product.
-Furthermore, the pivoted Cholesky preconditioner is not easy to implement in JAX.
-The pivoted Cholesky decomposition is not implemented in JAX, and so we would need to implement it manually.
-Since the pivoted Cholesky decomposition involves incrementally adding columns to the factorisation, and JAX is not friendly towards mutable variables, it is not straightforward to implement.
+Furthermore, the pivoted Cholesky decomposition is not implemented in JAX.
+Since the pivoted Cholesky decomposition involves incrementally adding columns to the factorisation, and JAX is not friendly towards mutable variables, this preconditioner is not straightforward to implement.
Additionally, the pivoted Cholesky decomposition traditionally uses the diagonal of the matrix to determine which columns to add to the factorisation, which is difficult without instantiating the matrix,we would need to approximate the diagonal of the GGN matrix.
\subsection{Other Preconditioners}
@@ -460,13 +462,13 @@ This means evaluating whether the samples are drawn from a normal distribution w
What we would like to fundamentally test is whether the CIQ algorithm converged correctly and thus whether the samples were generated successfully.
However, for cases when sampling fails to converge, this is typically caused by ill-conditioning of the posterior precision matrix.
Since we only have access to the precision matrix and computing the covariance involves inverting the precision matrix using, for example, the conjugate gradient method (which is itself sensitive to ill-conditioning), we cannot directly evaluate the covariance of the samples.
-Instead, we require a method to evaluate the quality of the samples without access to the covariance matrix.
+Instead, we require a method to evaluate the quality of the samples without access to the covariance matrix, which is what we will discuss in the next section.
\subsection{The Chi-Squared Distribution}[Chi-Squared Distribution]
-As per the defition of the chi-squared distribution, for \(D\)-dimensional standard normal samples \(\bm \varepsilon_0\), the sum of square deviations from the mean is distributed as a chi-squared distribution with \(D\) degrees of freedom, as per
+As per the definition of the chi-squared distribution, for \(D\)-dimensional standard normal samples \(\bm \varepsilon_0\), the sum of square deviations from the mean is distributed as a chi-squared distribution with \(D\) degrees of freedom, as per
\begin{align}
- \bm \varepsilon_0\T \bm \varepsilon_0 \sim \chi^2(D).
+ {\bm \varepsilon_0}\T \bm \varepsilon_0 \sim \chi^2(D).
\end{align}
Samples \(\bm \varepsilon\) which are normally distributed but do not have a mean of zero and a standard deviation of one can be standardised and their sum of square deviations be calculated by computing their squared Mahalanobis distances, which will be chi-squared distributed with \(D\) degrees of freedom, as per
%
@@ -475,9 +477,9 @@ Samples \(\bm \varepsilon\) which are normally distributed but do not have a mea
\\ \Rightarrow{}& \bm d\T \bm d \sim \chi^2(D). \nonumber
\end{align}
%
-This has the advantage of only requiring computation of the precision matrix \(\bm \Lambda\).
+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.
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 accessing this precision matrix in the evaluation.
+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.
We do this by comparing the histogram of these squared Mahalanobis distances to the appropriate PDF and by visualising the sample empirical percentiles against the theoretical percentiles in a quantile-quantile plot.
As such, this method can be used to evaluate whether the approximate normal samples obtained from CIQ are correctly distributed.
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -114,13 +114,13 @@ This is effective because we want to maximise the log-marginal.
Assume \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) has rank \(K \leq O \cdot N\).
Then, \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues greater than \(\alpha\), and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
-We then decompose the log determinant as
+We then decompose the log-determinant as
\begin{align}\label{eq:log-determinant-decomposition}
\log \det \bm \Lambda ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + \sum_{k=K+1}^D \log \alpha
\\ ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + (D - K) \log \alpha, \nonumber
\end{align}
where \(\lambda_1, \ldots, \lambda_K\) are the non-zero eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\).
-The first term is the log determinant of a positive definite matrix with eigenvalues \(\{\lambda_1 + \alpha, \ldots, \lambda_K + \alpha\}\), and the second term is a constant.
+The first term is the log-determinant of a positive definite matrix with eigenvalues \(\{\lambda_1 + \alpha, \ldots, \lambda_K + \alpha\}\), and the second term is a constant.
The former can then be bounded by \cref{eq:bai-golub-bound} with an upper bound on its eigenvalues \(\beta \geq \lambda_1 + \alpha\) as \(B_D\left(\sum_{k=1}^{K} (\lambda_k + \alpha), \sum_{k=1}^{K} (\lambda_k + \alpha)^2, \beta\right)\).
We then compute the trace of this term as
\begin{align}\label{eq:laplace-trace-bound}
@@ -132,19 +132,19 @@ 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,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,
\\ \mu_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}
For the upper bound on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\) as an upper bound.\sidenote{Since the matrix is positive semidefinite, all eigenvalues are non-negative and thus the trace is bigger than the largest eigenvalue.}
-From \cref{eq:laplace-log-determinant-bound}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log determinant of the posterior precision.
+From \cref{eq:laplace-log-determinant-bound}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log-determinant of the posterior precision.
% WHERE DOES THE BOUND COME FROM? WHICH UNCERTAINTY? THE UPPER BOUND ON THE EIGENVALUE?
\subsection{Hutchinson's Trace Estimator}[The Trace Estimator]
-To compute the upper bound on the log determinant of the posterior precision, we need to compute \(\mu_1\) and \(\mu_2\), which depend on the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\).
+To compute the upper bound on the log-determinant of the posterior precision, we need to compute \(\mu_1\) and \(\mu_2\), which depend on the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\).
To compute these, we can use Hutchinson's trace estimator~\cite{hutchinson1990stochastic}.
Let \(\bm \varepsilon_0 \sim \normal(\bm 0, \identity)\).
Since \(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} = \identity\), we compute \(\mu_1\) as
