commit 0e31212b9a128bac2882af690a4b3ca19381828a
parent 3910d50712d48a9eddac4a9fb77721ff0b9f5004
Author: Silas Brack <s174433@student.dtu.dk>
Date: Wed, 1 Feb 2023 16:32:22 +0100
Close to first draft.
Diffstat:
14 files changed, 1040 insertions(+), 680 deletions(-)
diff --git a/README.md b/README.md
@@ -12,3 +12,30 @@ More specifically, this project investigates the connections between the Laplace
We use a sampling-based training procedure, where a sample is first drawn from a Gaussian weight-posterior, a gradient step is performed on this sampled neural network, and the variance of the weight-posterior is updated with an approximate Hessian.
Approximating the Hessian is the most time-consuming and painful-to-engineer step of this training procedure.
In this project, we tap into the potential of modern machine learning frameworks to efficiently approximate the Hessian.
+
+# Setup
+
+This document is designed to be generated with `latexmk` and cleaned with `latexmk -C`.
+You can use either `pdflatex` or `xelatex` to compile the document, though I've used `pdflatex` throughout the process and can't guarantee that `xelatex` will work out of the box.
+If you want to use `xelatex`, you should update the `latexmkrc` file to use `xelatex` instead of `pdflatex`.
+I highly recommend using VSCode with the LaTeX Workshop extension to compile the document.
+I also lightly recommend using the LTeX extension and the stylecheck script to check for common errors, though this is more out of personal preference.
+If using the LaTeX Workshop extension, I would add this to your `settings.json` file:
+
+```json
+{
+ "latex-workshop.hover.preview.mathjax.extensions": ["boldsymbol"],
+ "latex-workshop.hover.preview.newcommand.newcommandFile": "mathjax.tex",
+ "latex-workshop.hover.preview.newcommand.parseTeXFile.enabled": true,
+ "latex-workshop.latex.clean.method": "cleanCommand",
+ "latex-workshop.latex.recipe.default": "latexmk (latexmkrc)",
+ "latex-workshop.latex.tools": [
+ {
+ "name": "latexmk_rconly",
+ "command": "latexmk",
+ "args": [],
+ "env": {}
+ },
+ ],
+}
+```
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -13,7 +13,7 @@ Introduce Bayesian inference and its applications in machine learning.
General advantages Bayesian neural networks have over other traditional deterministic neural networks (and disadvantages).
Specific applications where these advantages are important and why they're important.
Specific applications of Laplace approximation in Bayesian neural networks.
-Why this problem is importnat.
+Why this problem is important.
Bayesian inference is used to ..., but requires an approximation to the posterior to make it computationally tractable.
The Laplace is one such approximation which assumes a normal distribution for the posterior and uses the Hessian to approximate the posterior precision.
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -4,32 +4,48 @@
\section{Summary and Key Results}
-What results did we get? What did we learn? What did we propose?
-Remember to actually include some results and a quantitative look at the results that were obtained.
+In this project, we aimed to develop a scalable full-rank Laplace approximation for Bayesian neural networks.
+In doing so, we developed a novel method for approximately sampling from the full-rank 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 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-rank Laplace posterior.
+
+Furthermore,
+
+% Which quantitative results do I have which I can put here?
How did we get these results? What methods did we use?
-How did our findings contribute to the field? What changes?
+% How did our findings contribute to the field? What changes?
+Overall, while we have developed novel methods for training and inference in the full-rank Laplace approximation, we have not yet demonstrated its improved performance over the diagonal Laplace approximation nor its feasibility.
\section[Outlook and Future Developments]{Outlook and Future Developments}[Future Work]
% OUR ISSUES
-What were our limitations?
+% What were our limitations?
Fix marginal training, specifically hyperparameter optimisation
Fix convolution performance in marginal
-Rank estimation
-
-Understand Hessian structure and spectrum
-Also diagonal raw values maybe, compare the two?
-Diagonal lets you know exactly which parameters contribute to the largest eigenvalues
+Rank estimation, and (very closely related) determining if the eigenvectors are relatively stationary across different observations
-% RELATED TO ACTUAL USEFULNESS OF THE METHOD
-
-Perform a large-scale experiment.
-Propose one: resnet? dataset with millions of rows? (maybe not millions, but a lot)
+Furthermore, while we have improved the understanding of the spectrum of the Generalised Gauss-Newton Laplace posterior precision, we have not yet fully understood it.
+In particular, ...
+As such, it would be interesting to further investigate the spectrum of the Generalised Gauss-Newton Laplace posterior precision, in particular to better understand both the spectrum and the Hessian structure and their relationship to ...
-% WHERE THE PROBLEM SHOULD BE EXPLORED NEXT
+% Also diagonal raw values maybe, compare the two?
+% Diagonal lets you know exactly which parameters contribute to the largest eigenvalues
+We did not show that the full-rank Laplace approximation shows improved predictive performance or uncertainty calibration when compared to the diagonal or Kronecker-factored Laplace approximations.
+...
+Furthermore, the full-rank 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
+Furthermore, the feasibility of the full-rank Laplace approximation must be demonstrated on a large-scale problem.
+It is not clear how the full-rank Laplace approximation will scale with the number of parameters, data points, and model outputs.
+This is a factor both regarding the actual performance of the method and the post-hoc sampling time for the approximate posterior.
+As such, important experiments must be performed to determine the scalability of the full-rank Laplace approximation.
+This includes experiments on a large-scale problem, such as a large-scale image classification problem on a dataset with hundreds of thousands to millions of rows, with significantly higher resolution than the MNIST dataset, and more than 10 classes, such as for zero-shot learning.
+These experiments would thus also require models with many parameters, such as a ResNet.
+
+% WHERE THE PROBLEM SHOULD BE EXPLORED NEXT. NOT SURE THERES MUCH FOR ME TO SAY
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -28,8 +28,8 @@ This experiment will then answer two questions:
\label{fig:hessian-profile}
\end{marginfigure}
-We then compose the loss function \((g \circ f) (\bm x) := \sum^N_{i=1} (2 x_i + 0.5)^2\) as
-\(f(\bm x) := 2 \bm x + 0.5, g(\bm x) := \sum^N_{i=1} x_i^2\).
+We then compose the loss function \((g \circ f) (\bm x) \coloneqq \sum^N_{i=1} (2 x_i + 0.5)^2\) as
+\(f(\bm x) \coloneqq 2 \bm x + 0.5, g(\bm x) \coloneqq \sum^N_{i=1} x_i^2\).
This is equivalent to the sum-of-square-error loss \(g\) for a simple linear model \(f\).
Since \(f(x)\) is a linear function, the GGN matrix will exactly equal the actual Hessian.
We then initialise \(\bm x\) to a random vector and compute the matrix-vector product \(\bm K \bm v\) where \(\bm v\) is a random vector using three Hessian-vector products: the JAX efficient HVP and GVP functions and a manual Hessian instantiation and multiplication.
@@ -84,31 +84,40 @@ We then investigated the number of quadrature points \(Q\) required to determine
We vary \(q \in \{1, \ldots, 20\}\) to determine the number of quadrature points required for the Laplace approximation.
-We then compute the Mahalanobis distance between the samples and the true posterior mean, and compare this to the Chi-squared distribution with \(D\) degrees of freedom.
+We then compute the Mahalanobis distance between the samples and the true posterior mean, and compare this to the chi-squared distribution with \(D\) degrees of freedom.
We then compute the p-value of the Kolmogorov-Smirnov test for each value of \(q\) by comparing the empirical CDF of the Mahalanobis distances to the theoretical \(\chi^2_D\) CDF.
-This p-value is then used to determine whether the null hypothesis that the samples are multivariate normal can be rejected based on some significance threshold \(\alpha\).
+This p-value is then used to determine whether the null hypothesis that the samples are multivariate normal can be rejected based on some significance threshold.
Sample evaluation is explained in more detail in \cref{sec:sampling-evaluation}.
The results of this experiment are shown in \cref{tab:mnist-ablation,fig:mnist-ablation}.
\begin{margintable}
\centering
- \caption[Ablation over quadrature points for MNIST.]{Sampling time and p-values for Kolmogorov-Smirnov with varying quadrature points \(q\) on MNIST with prior precision \(\alpha = 0.1\). Lower p-value indicates a higher probability of rejecting normality.}
+ \caption[Ablation over quadrature points for MNIST.]{Sampling time and p-values for Kolmogorov-Smirnov with varying quadrature points \(q\) on MNIST with prior precision \(\alpha = 0.1\). Lower p-value indicates a higher probability of rejecting normality. Bold indicates sampling was successful based on visual inspection of Q--Q plots.}
\label{tab:mnist-ablation}
\include{tables/mnist_ablation.tex}
\end{margintable}
% RESULTS
-Based on \cref{fig:mnist-ablation}, sampling appears to be successful (based on a significance threshold of \(\alpha = 0.05\)) for \(q \in \{10, \ldots, 20\}\).
+Based on \cref{fig:mnist-ablation}, sampling appears to be successful (based on a significance threshold of \(0.05\)) for \(q \in \{10, \ldots, 20\}\).
% DISCUSSION
However, inspection of CHISQ PLOTS suggests that there is not a significant difference between these values and \(q = 5\).
In general, significance testing can be sensitive to the number of samples used, and, for large numbers of samples, these tests can excessively reject the null hypothesis.
However, the number of samples is generally fixed as a function of the number of parameters in the model and the amount of memory available.
As such, it is important to select an appropriate significance threshold for the test based on the number of posterior samples that are required.
-For 200 posterior samples, a more appropriate significance threshold of \(\alpha = 0.01\) or lower would reject sampling for \(q \leq 3\), since the cases where sampling completely fails have p-values that are below single floating point precision.
+For 200 posterior samples, a more appropriate significance threshold of \(0.01\) or lower would reject sampling for \(q \leq 3\), since the cases where sampling completely fails have p-values that are below single floating point precision.
Furthermore, the plots in CHISQ PLOTS suggest that the performance of the sampling method is effectively equal for \(q \in \{5, \ldots, 20\}\).
+In terms of the time required to compute the samples, the time required to compute the samples is XXXX in \(q\).
+However, in taking as much as 20 to 30 minutes to compute 200 samples of an MNIST model with 15\,000 parameters, the time required to compute the samples is not negligible.
+In practice, it is likely that only a small number of samples will be obtained (around 5-15), since for each sample it is necessary to store a parameter vector of size \(D\).
+However, for 10 samples, it still takes XXXX minutes.
+As such, the time required to compute the samples may be a significant factor in the overall time required to compute the Laplace approximation.
+In particular, for models with a very large number of parameters, the time required to compute the samples may be prohibitive when compared when the near-instant time required to sample from the diagonal Laplace approximation.
+Even if the sampling time scales sublinearly with the number of parameters, it would still prevent methods such as online Laplace, which require the Laplace approximation to be computed at each iteration, from being used in practice.
+However, this sampling time constraint is also significant for MCMC methods, which are the most common method for sampling from the posterior distribution of a BNN.
+
\section[Sampling for the Sine Function]{Sampling for the Sine Function}[The Sine Function]
The choice of \(\alpha\) is a non-trivial one
@@ -123,7 +132,7 @@ The choice of \(\alpha\) is a non-trivial one
\begin{figure}
\centering
\includegraphics{chisq_sin.pdf}
- \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for testing normality of Laplace approximation samples for a neural network modelling \(f(x) = \sin(x)\). The sample Mahalanobis distances are Chi-squared distributed, then the samples themselves are multivariate normal with correct (fixed) covariance. As can be seen by looking at the histogram and the q-q plot, samples are multivariate normal.}
+ \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for testing normality of Laplace approximation samples for a neural network modelling \(f(x) = \sin(x)\). The sample Mahalanobis distances are chi-squared distributed, then the samples themselves are multivariate normal with correct (fixed) covariance. As can be seen by looking at the histogram and the q-q plot, samples are multivariate normal.}
\label{fig:chisq-sin}
\end{figure}
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -2,20 +2,21 @@
\chapter{Introduction}
\labch{intro}
-% \section{The Problem}
\section{Bayesian Deep Learning}
+% Introduce Bayesian inference and its applications in machine learning.
+% General advantages Bayesian neural networks have over other traditional deterministic neural networks (and disadvantages).
+% Specific applications where these advantages are important and why they're important.
+% Specific applications of Laplace approximation in Bayesian neural networks.
+% Why this problem is important.
+
Bayesian methods allow for estimates of uncertainty which enable more efficient usage of data (e.g., via active learning) and avoid overfitting.
-Furthermore, these uncertainty estimates improve model interpretability and assessment of model predictive confidence.
+They accomplish this by modelling a distribution on the model parameters, which are then used to make predictions.
+These predictions will then have uncertainty estimates, thus improving model interpretability and enabling the assessment of model predictive confidence.
These methods have been successfully applied to a wide range of classical problems in statistics, though attempts to apply them to deep learning have had limited success.
-% are particularly useful in the context of deep learning,
-% where the models are often black-boxes and the data is often scarce.
-Introduce Bayesian inference and its applications in machine learning.
-General advantages Bayesian neural networks have over other traditional deterministic neural networks (and disadvantages).
-Specific applications where these advantages are important and why they're important.
-Specific applications of Laplace approximation in Bayesian neural networks.
-Why this problem is importnat.
+For deep learning, the Bayesian approach is often limited by the computational cost of Bayesian inference.
+This is because Bayesian inference requires computing the posterior distribution of the model parameters given the data, which, due to the nonlinear complexity of the neural network architecture, is often intractable.
Typically, the Bayesian approach to deep learning is to use a prior distribution over the model parameters and then compute the posterior distribution of the model parameters given the data, where the posterior is maximised to obtain the optimal model parameters for the available data.
This posterior distribution can then be sampled from to make predictions.
@@ -27,52 +28,56 @@ However, even in this case, the choice of prior precision is still difficult, of
Instead, maximising the marginal likelihood allows for optimisation of hyperparameters in the likelihood and prior, which can be done using gradient-based optimisation algorithms.
However, this is computationally infeasible, and so approximate Bayesian methods are used instead.
-% \section{The State of Research in the Area}
\section{Current Methods}
Currently, approximate Bayesian methods are either expensive to compute (Markov Chain Monte Carlo), are significantly more difficult to implement (such as variational inference), or simply perform poorly and are limited in their Bayesian interpretation (MC dropout).
-% Beyond these, there are also domain-specific models which aim to provide
+Of these, MC dropout and Markov Chain Monte Carlo are the most widely used, though both have significant drawbacks.
As such, there is demand for a method which exhibits the same computational cost as the optimization algorithms for deterministic neural networks while providing accurate posterior approximations and working out-of-the-box for any given architecture.
-The Laplace approximation~\sidecite{daxberger2021laplace,laplace1774memoire} is a simple yet theoretically well-supported posterior approximation suitable for Bayesian modeling.
-The Laplace approximation has been applied in, for example, the prediction of earthquake hypocenters~\sidecite{izzatullah2022laplace}.
-The primary advantages of the Laplace approximation are that it is simple to implement, effective at out-of-distribution detection~\sidecite{miani2022laplacian} and can be performed post-hoc on a pre-trained model.
+The Laplace approximation~\cite{daxberger2021laplace,laplace1774memoire} is a simple yet theoretically well-supported posterior approximation suitable for Bayesian modeling.
+The Laplace approximation has been applied in, for example, the prediction of earthquake hypocenters~\cite{izzatullah2022laplace}.
+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, to compute the posterior distribution of the model parameters, we need to compute the Hessian of the 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-rank Laplace outside of toy problems.
-To overcome this, we can use a crude approximation, such as only storing the Hessian matrix's diagonal~\sidecite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\sidecite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.%~\sidecite{}.
-Many of these Hessian approximations have successfully been applied to the Laplace approximation~\sidecite{ritter2018scalable,ritter2018online}, though are limited by the quality of the approximation.
+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.%~\cite{}.
+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-rank Laplace approximation without requiring the Hessian to be stored explicitly.
\section{Large-Scale Laplace}
-Large-Scale Laplace (LSL) is a method which approximates the Hessian of the loss with respect to the model parameters without requiring the Hessian to be stored explicitly.
-This is achieved by using a sampling-based training procedure, where a sample is first drawn from a Gaussian weight prior, and then the loss is evaluated on this sample.
-Sampling from the posterior requires computing the covariance matrix as the inverse of the Hessian of the loss with respect to the model parameters, and then taking the square root of this covariance matrix.
-However, for large models, which can contain millions or billions of parameters, this is intractable using conventional methods, since the Hessian will be a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters.
-
-\sidetextcite{pleiss2020fast} propose a contour integral quadrature (CIQ), a method which approximates the square root of a square matrix multiplied by a vector without requiring this matrix to be stored explicitly.
-While they use it for computing the square root of the kernel matrix in Gaussian processes, we use it for computing the square root of the covariance matrix in Bayesian neural networks using the Laplace approximation.
-To run CIQ, we therefore only require the ability to compute matrix-vector products.
-Since, to sample from the Laplace posterior distribution we need to compute the square root of the covariance matrix, we need be able to compute Hessian-vector products (i.e., matrix-vector products with the Hessian matrix) without explicitly storing the Hessian matrix.
-% he Hessian-vector product, which is a crucial step in the Laplace approximation.
-% This method is used to approximate the Hessian-vector product, which is used to update the posterior variance of the weights in a Bayesian neural network.
-% This method is used in conjunction with the Laplace approximation, which is a simple, yet principled, posterior approximation suitable for Bayesian modeling.
-% The Laplace approximation relies on access to the Hessian matrix of the loss with respect to the model parameters.
-% For large models, it is even intractable to store this matrix, and crude approximations (e.g. only storing the Hessian diagonal) are called upon.
-% We investigate storing the Hessian implicitly as means to avoid the fuzz.
+We are seeking an implementation of the full-rank Laplace approximation which doesn't require instantiation of the full posterior precision matrix.
+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.
+Sampling thus requires computing the covariance matrix as the inverse of the Hessian of the loss with respect to the model parameters, and then taking the square root of this covariance matrix.
+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 compute the inverse square root of the Hessian without explicitly storing it.
+% , and instead only storing the output of Hessian-vector products.
+Furthermore, we would like to perform training using the negative log marginal likelihood loss, which, for the Laplace approximation, requires computing the log determinant of the Hessian.
+To perform these computations we need to approximate these quantities without explicitly storing the Hessian matrix.
+One way to do this is to instead compute Hessian-vector products, which are matrix-vector products with the Hessian matrix.
+Fortunately, it is possible to compute Hessian-vector products while only storing the implicit Hessian computational graph.
+
+% \textcite{pleiss2020fast} propose a contour integral quadrature (CIQ), a method which approximates the product of a square root of a matrix with a vector without requiring this matrix to be stored explicitly.
+% Instead of instantiating the matrix, they approximate this square-root matrix-vector product by performing repeated matrix-vector products with this matrix, and then approximating the result using a quadrature rule.
+% While they use it for computing the square root of the kernel matrix in Gaussian processes, we use it for computing the inverse square root of the posterior precision in Bayesian neural networks using the Laplace approximation.
+% Since, to sample from the Laplace posterior distribution we need to compute the square root of the covariance matrix, we need be able to compute Hessian-vector products (i.e., matrix-vector products with the Hessian matrix) without explicitly storing the Hessian matrix.
+% Fortunately, it is possible to compute Hessian-vector products while only storing the implicit Hessian computational graph. (MAYBE SHOULD REMOVE THIS PARAGRAPH, TOO MUCH INTO METHODS)
% \section{JAX and XLA}
-JAX~\sidecite{jax2018github} is a Python library for GPU-accelerated automatic differentiation and compilation of computational graphs into efficient machine code using XLA~\sidecite{sabne2020xla}.
-% To perform this compilation, it uses XLA, a compiler for linear algebra~\sidecite{sabne2020xla}.
+JAX~\cite{jax2018github} is a Python library for GPU-accelerated automatic differentiation and compilation of computational graphs into efficient machine code using XLA~\cite{sabne2020xla}, a compiler for linear algebra.
JAX enables us to efficiently compute Hessian-vector products without explicitly storing the Hessian matrix by compiling the Hessian-vector product computation with XLA.
-% JAX is designed to be a drop-in replacement for NumPy, and is compatible with most of the Python scientific stack. JAX is built on top of XLA, a domain-specific compiler for linear algebra. XLA is able to compile Python code into efficient machine code, and is used by Google to accelerate TensorFlow and PyTorch. JAX is able to compile Python code into efficient machine code, and is used by Google to accelerate TensorFlow and PyTorch.
-% It is designed to be a drop-in replacement for NumPy, and is compatible with most of the Python scientific stack. JAX is built on top of XLA, a domain-specific compiler for linear algebra. XLA is able to compile Python code into efficient machine code, and is used by Google to accelerate TensorFlow and PyTorch. JAX is able to compile Python code into efficient machine code, and is used by Google to accelerate TensorFlow and PyTorch.
% \section{Research Objectives}
+We would thus like to develop a method which performs the full-rank 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 use these Hessian-vector products to compute the log determinant and inverse square root of the Hessian.
+These quantities will be used to perform training and inference using the Laplace approximation.
+% SHOULD THIS BE LONGER? TALK ABOUT EVALUATION SOMEWHERE?
+
% This should be discussed in a bit more detail in a thesis, since there
% are certain objectives discussed in the project plan.
% In an article, a quick sentence which summarises the rest of the
diff --git a/chapters/literature.tex b/chapters/literature.tex
@@ -6,8 +6,8 @@
% How does adam work?
-\sidecite{sutskever2013importance}
-\sidecite{kingma2014adam}
+\cite{sutskever2013importance}
+\cite{kingma2014adam}
% Integrate with the previous section
Note that Adam approximates the second moment of the gradient of the loss, that this corresponds to the variance of this gradient, and that this is the diagonal of the Fisher information matrix, as we defined in \cref{sec:fisher-information} (\cref{eq:gradient-covariance})!
@@ -16,21 +16,21 @@ Note that Adam approximates the second moment of the gradient of the loss, that
% \section{K-FAC}
% Another alternative to approximating the Fisher information matrix is by calculating the Kronecker-factored Approximate Curvature (K-FAC).
-% \sidecite{martens2015optimizing}
+% \cite{martens2015optimizing}
% \section{Noisy K-FAC}
-% \sidecite{zhang2018noisy}
+% \cite{zhang2018noisy}
% \section{Gauss-Newton Optimization}
-% \sidecite{botev2017practical}
+% \cite{botev2017practical}
% \section{Variational Adaptive-Newton}
-% \sidecite{khan2017variational}
+% \cite{khan2017variational}
% \section{Natural Gradient Variational Inference}
-% Natural gradient VI (VOGN)~\sidecite{osawa2019practical} approximates the posterior as
+% Natural gradient VI (VOGN)~\cite{osawa2019practical} approximates the posterior as
% %
% \begin{equation*}
% p(\bm \theta) \approx \normal(\bm 0, \bm I / \delta)
@@ -43,7 +43,7 @@ Note that Adam approximates the second moment of the gradient of the loss, that
\section{Laplace Approximation}\label{sec:laplace}
-In the Laplace approximation (LA)~\sidecite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI.
+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}}\).
Since the Taylor expansion is performed around the MAP solution, the first order derivative is zero, and the expansion is simply given by
\begin{align}
@@ -81,14 +81,14 @@ minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees po
% % https://martiningram.github.io/vi-with-good-covariances/
% % https://github.com/martiningram/jax_advi/blob/main/jax_advi/lrvb.py#L49
-% \sidecite{giordano2018covariances}
+% \cite{giordano2018covariances}
% INLA: Laplace for nested models (hierarchical models)
-% \sidecite{rue2017bayesian}
+% \cite{rue2017bayesian}
% \section{Online Laplace}
-% \sidecite{miani2022laplacian}
+% \cite{miani2022laplacian}
% \section{Bayes’ Theorem}
@@ -183,8 +183,8 @@ minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees po
% \end{align}
% where \(\bm{\theta}^{(s)} \sim \normal(\bm{\mu}, \bm{\sigma})\).
-% The use of the mean-field approximation for BNNs has been found to be unreliable~\sidecite{wu2018deterministic} and, for increasingly wide neural networks, converges toward the prior~\sidecite{coker2021wide}.
-% However, for \emph{deep} BNNs, the mean-field assumption may be reasonable~\sidecite{farquhar2020liberty}.
+% The use of the mean-field approximation for BNNs has been found to be unreliable~\cite{wu2018deterministic} and, for increasingly wide neural networks, converges toward the prior~\cite{coker2021wide}.
+% However, for \emph{deep} BNNs, the mean-field assumption may be reasonable~\cite{farquhar2020liberty}.
% \subsection{Full-Rank Variational Approximation}
@@ -224,7 +224,7 @@ minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees po
% \subsection{Deep Ensembles}
-% In deep ensembles~\sidecite{lakshminarayanan2017simple}, \(M\) neural networks are trained with different initialisations.
+% In deep ensembles~\cite{lakshminarayanan2017simple}, \(M\) neural networks are trained with different initialisations.
% In this way, multiple local MAP solutions \(\bm\theta^{(m)}\) are obtained and equally considered, such that \(p(\bm \theta = \bm \theta^{(m)} \given \bm y) = 1 / M\) for every \(\bm \theta^{(m)} \in \{\bm\theta^{(1)}, \ldots, \bm\theta^{(M)}\}\) (and 0 elsewhere).
% To make predictions with a deep ensemble, we use the posterior predictive given by
% \begin{align}
@@ -238,13 +238,13 @@ minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees po
% Stochastic weight averaging estimates the final point estimate of the weights of a neural network \(\bm \theta_\mathrm{SWA}\) as the average of the model parameters \(\bm \theta\) at the end of each epoch obtained via gradient descent after convergence has been reached.
-% For SWAG~\sidecite{maddox2019simple}, the standard deviation of the model parameters \(\bm \sigma_\mathrm{SWA}\) is also calculated.
+% For SWAG~\cite{maddox2019simple}, the standard deviation of the model parameters \(\bm \sigma_\mathrm{SWA}\) is also calculated.
% From the SWA estimates of the model parameter mean and standard deviations, the posterior is then estimated as
% \begin{align}\label{eq:swag}
% p(\bm \theta \given \bm y) \approx{} & \normal(\bm \theta_\mathrm{SWA}, \bm \sigma_\mathrm{SWA}).
% \end{align}
-% For MultiSWAG~\sidecite{wilson2020bayesian}, \(M\) SWAG estimates are obtained from different initialisations, as in deep ensembles.
+% For MultiSWAG~\cite{wilson2020bayesian}, \(M\) SWAG estimates are obtained from different initialisations, as in deep ensembles.
% Then, a Gaussian mixture model (GMM) is formed as a combination of each Gaussian SWAG posterior estimate \(\normal(\bm \theta_\mathrm{SWA}^{(m)}, \bm \sigma_\mathrm{SWA}^{(m)})\) where each component is equally weighed, as
% \begin{align}\label{eq:multiswag}
% p(\bm \theta \given \bm y) \approx{} & \frac{1}{M} \sum_{m=1}^M \normal(\bm \theta_\mathrm{SWA}^{(m)}, \bm \sigma_\mathrm{SWA}^{(m)}).
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -5,7 +5,7 @@
\section{Introduction}
In many methods, such as the Laplace approximation (\cref{sec:laplace}), we sample from a multivariate normal distribution parameterised by a mean vector \(\bm \mu\) and a covariance matrix \(\bm \Sigma\).
-To sample from a Gaussian distribution, you can apply the \emph{reparametrisation trick} to compute a vector of samples as
+To sample from a Gaussian distribution, you can apply the reparametrisation trick to compute a vector of samples as
\begin{align}\label{eq:sample-reparam}
\bm \varepsilon ={} & \bm \mu + \bm \Sigma^{1/2} \bm \varepsilon_0,
\end{align}
@@ -28,7 +28,7 @@ However, the precision matrix we are interested in is not itself low-rank but a
This is discussed in more detail in \cref{sec:laplace,sec:preconditioning}.
% As such, these are not adequate methods for computing the inverse square root using only matrix-vector products.
-Furthermore, methods which do not require computation of the square root of the covariance matrix typically attempt to exploit properties such as a block structure or sparsity of the covariance matrix~\sidecite{vono2022high}.
+Furthermore, methods which do not require computation of the square root of the covariance matrix typically attempt to exploit properties such as a block structure or sparsity of the covariance matrix~\cite{vono2022high}.
However, we do not have a good understanding of the structure of the Laplace covariance matrix, and thus these methods are not suitable for our problem.
% Talk about what we do know about the structure? Wrt the sum of low-rank GGNs
When sampling from high-dimensional Gaussian distributions, it is common to use Markov chain Monte Carlo (MCMC) sampling methods such as Gibbs sampling~\cite{vono2022high}.
@@ -38,8 +38,8 @@ For this, we need a method to compute the inverse square root using only matrix-
\section[Contour Integral Quadrature]{Contour Integral Quadrature}[Contour Integral Quad]
-\sidetextcite{pleiss2020fast} propose one such method.
-Contour integral quadrature (CIQ) attempts to solve the problem of computing the inverse-square root of a matrix, up to a rotation, by approximating the expression \(f(\bm K) \bm \varepsilon_0 = \bm K^{-1/2} \bm \varepsilon_0\) where \(\bm K\) is a positive semi-definite matrix using Cauchy's integral formula.
+\textcite{pleiss2020fast} propose one such method.
+Contour integral quadrature (CIQ) attempts to solve the problem of computing the inverse square root of a matrix, up to a rotation, by approximating the expression \(f(\bm K) \bm \varepsilon_0 = \bm K^{-1/2} \bm \varepsilon_0\) where \(\bm K\) is a positive semi-definite matrix using Cauchy's integral formula.
To approximate this integral, it exploits the Lanczos algorithm, Gaussian quadrature, and the msMINRES algorithm, methods which only require the computation of matrix-vector products \(\bm K \bm v\).
There are then essentially three steps to CIQ, namely:
\begin{enumerate}
@@ -89,13 +89,13 @@ it is then also trivial to compute the square root product by multiplying \(\bm
% CAUCHY INTEGRAL FORMULA
% ------------------------------------------
-The first step invokes Cauchy's integral formula.
+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.}
\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}
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 (\cref{sec:lanczos}) to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).
+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\}\).}
% ------------------------------------------
% QUADRATURE
@@ -113,7 +113,7 @@ To estimate these values, we can use the Lanczos algorithm (\cref{sec:lanczos})
\label{fig:ciq}
\end{marginfigure}
-We can then apply this to the function \(f(\bm K) = \bm K^{-1/2}\), apply a change of variable, and then approximate this contour integral using the quadrature rule with \(Q\) quadrature points (see \cref{fig:ciq}), as per
+We can then apply Cauchy's integral formula to the function \(f(\bm K) = \bm K^{-1/2}\), apply a change of variable, and then approximate this contour integral using the quadrature rule with \(Q\) quadrature points (see \cref{fig:ciq}), as per
\begin{align}
\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} \label{eq:ciq-quadrature-tau}
@@ -126,8 +126,8 @@ Since we will be sampling a fixed number of quadrature points from the circle, t
% ------------------------------------------
We could sample these quadrature points from the circle uniformly, which corresponds to the regular trapezoid quadrature rule.
-However, the convergence for this quadrature is linear with regard to the condition number of \(\bm K\)~\sidecite{pleiss2020fast,hale2008computing}.
-Since the precision matrix (which will be computed from the Laplace approximation as the GGN matrix) can often have a very low rank and poor conditioning, uniform sampling would then require a very large number of quadrature points \(Q\) and would therefore be inadequate.
+However, the convergence for this quadrature is linear with regard to the condition number of \(\bm K\)~\cite{pleiss2020fast,hale2008computing}.
+Since the precision matrix (which will be computed from the Laplace approximation as the GGN matrix) can often have a very low rank and poor conditioning, uniform sampling would then require a very large number of quadrature points \(Q\) and would therefore be inadequate~\cite{pleiss2020fast}.
Since our GGN precision matrix \(\bm \Lambda = \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) is often low-rank, we would instead prefer to oversample quadrature points which lie close to the minimum eigenvalue.
This is accomplished by applying a change of variable from \(\tau\) to \(\sigma = \tau^{1/2}\) and sampling the quadrature points from \(\sigma\).
@@ -145,22 +145,54 @@ From here, the inverse square root vector product is given by
We can then solve the \(Q\) systems of equations to obtain the solutions \(\bm c_q = \left(\tau_q \identity - \bm K \right)^{-1} \bm v\) for \(q \in 1, \ldots, Q\).
Since the integrand is symmetric with respect to the real axis, only the imaginary component of \(\Gamma_\sigma\) is required.\sidenote{This is because the circle with its centre on the \(x\)-axis is symmetric with respect to said \(x\)-axis.}
-Since \(\sigma\) is then imaginary (its real part is zero), \(\tau := \sigma^2\) is real such that \(\tau = - \Im(\sigma)^2\).
+Since \(\sigma\) is then imaginary (its real part is zero), \(\tau \coloneqq \sigma^2\) is real such that \(\tau = - \Im(\sigma)^2\).
As such, since \(\Im(\sigma)^2\) must be positive, \(\tau\) must then be real-valued and negative.
In practice, we incorporate the \(1 / (\pi i)\) term into the quadrature weights and, since \(\Tilde{w}_q\) also happens to be real-valued and negative, we can redefine the quadrature weights \(w_q = -\Tilde{w}_q\) and the quadrature points \(t_q = -\sigma^2_q\) for simplicity.\marginnote{With this redefinition, both \(w_q > 0\) and \(t_q > 0\). This then means that \(\left(t_q \identity + \bm K\right)\) is positive definite.}
-This yields the quadrature rule result \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1} w_q \bm c_q\), where \(\bm c_q = \left(t_q \identity + \bm K\right)^{-1}\).
+This yields the quadrature rule result \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1} w_q \bm c_q\), where \(\bm c_q = (t_q \identity + \bm K)^{-1}\).
This change of variables from \(\tau\) to \(\sigma\) yields a set of quadrature weights \(w_1, \ldots, w_Q\) given by
\begin{align}\label{eq:quadrature-weights}
- w_q = \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cn(i u_q \mathcal{K}^\prime (k) \given k) \dn (i u_q \mathcal{K}^\prime (k) \given k)
+ w_q = \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cdot \cn(i u_q \mathcal{K}^\prime (k) \given k) \cdot \dn (i u_q \mathcal{K}^\prime (k) \given k)
\end{align}
and a set of quadrature points (shifts) given by
\begin{align}\label{eq:quadrature-shifts}
\sigma^2_q = \lambda_{\mathrm{min}} \cdot \sn(i u_q \mathcal{K}^\prime (k) \given k)^2,
\end{align}
-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\).
+where \(\sn\), \(\cn\), and \(\dn\) are the Jacobi elliptic functions and \(\mathcal{K}^\prime(k)\) is the complete elliptic integral, such that \(\mathcal{K}^\prime(k) = \mathcal{K}(k^\prime) = \mathcal{K}(\sqrt{1 - k^2})\).
+The 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.
+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.
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{marginfigure}
+ \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.}
+ \label{fig:ellipk}
+\end{marginfigure}
+
+\begin{figure}
+ \centering
+ \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. 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}
\begin{marginfigure}
\begin{algorithm}[H]
@@ -185,166 +217,192 @@ We can thereby now compute the quadrature weights and shifts, as seen in \cref{a
\end{algorithm}
\end{marginfigure}
-The last step involves solving the \(Q\) linear systems defined in \cref{eq:ciq-sigma-product} to determine \(\bm c_q\).
-As mentioned above, this is done by using msMINRES (\cref{sec:msminres,alg:msminres}).
-Since the Lanczos algorithm only requires the matrix-vector product \(\bm K \bm v\) with a random vector \(\bm v\), we can use the same random vector for all \(Q\) linear systems.
+The last step involves solving the \(Q\) linear systems defined in \cref{eq:ciq-sigma-product} to determine \(\bm c_q\). % Actually we solve the other linear system.
+As mentioned above, this is done by using msMINRES, an algorithm for solving linear systems of the form \((\bm K + t_q \identity) \bm c_q = \bm b\) where \(\bm K\) is symmetric.
+Like other Krylov methods, msMINRES computes the solutions to the linear systems by constructing a Krylov subspace \(\mathcal{K}_q\) of \(\bm K + t_q \identity\), and it then iteratively minimises the norm of the residual in \(\mathcal{K}_q\).
+While msMINRES only requires that the matrices \(\bm K + t_q \identity\) be symmetric, and not necessarily also positive definite, since the shifts \(t_q\) are positive and \(\bm K\) is positive semidefinite, the sum is also positive definite.
+However, the CIQ implementation from \textcite{pleiss2020fast} uses msMINRES, and there is also evidence suggesting that msMINRES converges faster than, for example, the conjugate gradient method (CG)~\cite{fong2012cg}, a Krylov method that requires the matrices to be positive definite.
+
+Since the Lanczos algorithm only requires the matrix-vector product \(\bm K \bm v\) with a random vector \(\bm v\) to estimate a bound on the eigenvalues, we can use the same bound for all \(Q\) linear systems.
As such, the computational bottleneck of the contour integral quadrature is performing the msMINRES for the shifted systems.
-Additionally, since msMINRES is an iterative Krylov method, it is also sensitive to the conditioning of the linear systems, i.e., the conditioning of \(\bm K\).
+Additionally, since msMINRES is an iterative Krylov method, its convergence bound is sensitive to the conditioning of the linear systems, i.e., the conditioning of \(\bm K\).
To mitigate the effect of ill conditioning, we can use a preconditioner \(\bm P\) such that \(\bm P^{-1} \bm K\) is well-conditioned.
\section{Preconditioning}
\label{sec:preconditioning}
-One drawback of contour integral quadrature if that the accuracy its of approximation of \(\bm \Lambda^{-1/2} \bm \varepsilon_0\) suffers when the precision matrix \(\bm \Lambda\) is poorly conditioned.\sidenote{A matrix is considered poorly conditioned when the ratio of its largest eigenvalue to its lowest eigenvalue, its \emph{conditioning number} (\(\kappa := \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}\)), is very large.}
-
+One drawback of CIQ is that the accuracy its of approximation of \(\bm \Lambda^{-1/2} \bm \varepsilon_0\) suffers when the precision matrix \(\bm \Lambda\) is poorly conditioned.
In the case of the Laplace approximation, as per \cref{eq:ggn-laplace}, we have that the precision matrix is given by
\begin{align}\label{eq:ggn-laplace}
- \bm \Lambda = \sum_i^N \bm J_i\T \bm H \bm J_i + \alpha \identity,
+ \bm \Lambda = \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity,
\\ \bm \theta \sim \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1})
\end{align}
where \(N\) is the number of observations in your dataset.
-Since each \(J_i\T H J_i\) is, for \(i \in 1 \ldots N\), a low-rank positive semi-definite matrix, it has \(o\) non-zero eigenvalues.\sidenote{With rank \(o\), where \(o\) is the number of neural network outputs.}
-Let us define \(\lambda_1, \ldots, \lambda_o\) as the eigenvalues of a matrix sorted in decreasing order.
-We can then
-Assume that the non-zero spectrum of \(\bm J_i\T \bm H \bm J_i\) is relatively flat, i.e., \(\lambda_1 \not\gg \lambda_o\).
+Since each \(J_i\T H_i J_i\) is, for \(i \in 1 \ldots N\), a low-rank positive semi-definite matrix, it has \(O\) non-zero eigenvalues (see \cref{sec:spectrum-ggn}).
+% Let us define \(\lambda_1, \ldots, \lambda_o\) as the eigenvalues of a matrix sorted in decreasing order.
+% Assume that the non-zero spectrum of \(\bm J_i\T \bm H_i \bm J_i\) is relatively flat, i.e., \(\lambda_1 \not\gg \lambda_o\).
-In a hand-wavy way, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are, the more the largest eigenvalue will be to the sum of the eigenvalues.
+% In a hand-wavy way, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are, the more the largest eigenvalue will be to the sum of the eigenvalues.
-For preconditioning, we need to find a preconditioner matrix \(P\) and its inverse \(P^{-1}\) such that
+For preconditioning, we need to find a preconditioner matrix \(P\) and its inverse \(P^{-1}\) such that the matrix \(P^{-1} \bm \Lambda\) is well-conditioned.
+This is done by choosing \(P\) such that the two matrices are ``similar'', i.e., they have close eigenvalues and eigenvectors.
+In this way, the preconditioner acts as an ``initial best guess'' of the value to be computed.
+We are thus seeking a preconditioner \(P\) such that
\begin{align}
\bm P^{-1} \bm A \approx{} & \identity \Rightarrow \bm P \approx \bm A
\\ \bm P^{-1} \bm \Lambda ={}& \bm P^{-1} \left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right)
\end{align}
-where \(\bm \Lambda\) is our precision matrix, approximated by a GGN Hessian.
-
-The eigenvalues of our GGN approximation are given by
+where \(\bm \Lambda\) is the posterior precision matrix given by \cref{eq:ggn-laplace} to improve convergence of the msMINRES and Lanczos algorithms.
+The eigenvalues of our GGN approximation are thus given by
\begin{align}
- \lambda\left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right) = \big\{ \underbrace{s_1, \ldots, s_p}_{\leq N \cdot k}, \underbrace{\alpha, \ldots, \alpha}_{\geq n - N \cdot k} \big\}
+ \lambda\left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right) = \big\{ \underbrace{s_1, \ldots, s_K}_{K \leq N \cdot O}, \underbrace{\alpha, \ldots, \alpha}_{D - K} \big\}
\end{align}
-where \(\{s_1, \ldots, s_p\}, p \leq N \cdot k\) are the non-zero eigenvalues of the rank-\(p\) matrix \(\sum^N_i \bm J_i\T \bm H_i \bm J_i\).
-
+where \(\{s_1, \ldots, s_K\}, K \leq N \cdot O\) are the non-zero eigenvalues of the rank-\(K\) matrix \(\sum^N_i \bm J_i\T \bm H_i \bm J_i\).
If a preconditioner approximates all eigenvalues and eigenvectors of a matrix, then the preconditioner and the matrix are equal and the preconditioner is exact.
However, we often do not have a preconditioner which is exact, but rather an approximation of the matrix.
-In this case, a balance must be struck between the accuracy of the preconditioner and the computational cost of computing the preconditioner (as well as its feasibility).
+In this case, a balance must be struck between the accuracy of the preconditioner and its computational cost.
Furthermore, the inverse of the preconditioner is required for preconditioning, which is often not available.
-To approximate the inverse of the preconditioner, we can use the conjugate gradient method, though this method is then sensitive to the conditioning of the preconditioner.
-Spectrum, shape of spectrum, how steep is it?
-The more values you add, the steeper the spectrum becomes (even the non-zero one).
-
-\begin{enumerate}
- \item Scaling
- \item Eigenvalues and eigenvectors
- \item Conditioning of the preconditioner
- \item Targeting high and low eigenvalues
- \item Rank, summation of low rank matrices
- \item Speed
-\end{enumerate}
-
-Typically, when preconditioning, you need a matrix which is similar to the matrix you are preconditioning.
-In this way, the preconditioner acts as an ``initial best guess'' of the value to be computed.
-As such, for CIQ, it would be expected that the preconditioner should approximate the value of \(\bm \Lambda^{-1/2}\).
-In reality, we only need an approximation of \(\bm \Lambda\).
-WAIT IS THIS TRUE? WE STILL NEED \(\bm \Lambda^{1/2} \bm b\), WE JUST CALCULATE THIS VALUE USING CIQ.
-
-As such, the conditioning number of our preconditioner also becomes an issue, since we have two conflicting objectives: to have a preconditioner which is as similar to \(\bm \Lambda\) as possible and to have a preconditioner which we can invert easily.
-The first point lends itself toward the use of preconditioners with a similar conditioning number as \(\bm \Lambda\).
-However, to invert an arbitrary positive semi-definite matrix, we will often resort to using the conjugate gradient method which is sentitive to the conditioning number of the preconditioner.
-Thus, there is an intrinsic balance to be struck between a preconditioner which is very similar to \(\bm \Lambda\) (but whose conditioning number is itself an issue) and a preconditioner which is less similar to \(\bm \Lambda\) (and will therefore not improve the conditioning of the problem).
+To approximate the inverse of the preconditioner, we can use the conjugate gradient method, though this method is then sensitive to the condition number of the preconditioner; and if the original matrix \(\bm \Lambda\) is ill-conditioned, then a good preconditioner will also tend to be ill-conditioned.
+
+% \begin{enumerate}
+% \item Scaling
+% \item Eigenvalues and eigenvectors
+% \item Conditioning of the preconditioner
+% \item Targeting high and low eigenvalues
+% \item Rank, summation of low rank matrices
+% \item Speed
+% \end{enumerate}
+
+As such, the condition number of our preconditioner also becomes an issue, since we have two conflicting objectives: to have a preconditioner which is as similar to \(\bm \Lambda\) as possible and to have a preconditioner which we can invert easily.
+The first point lends itself toward the use of preconditioners with a similar condition number as \(\bm \Lambda\).
+However, to invert an arbitrary positive semi-definite matrix, we will often resort to using the conjugate gradient method which is sentitive to the condition number of the preconditioner.
+Thus, there is an intrinsic balance to be struck between a preconditioner which is very similar to \(\bm \Lambda\) (but whose condition number is problematic) and a preconditioner which is less similar to \(\bm \Lambda\) (and will therefore not improve the conditioning of the problem).
Let us consider the toy example of a preconditioner which we can guarantee to have the same eigenvectors as \(\bm \Lambda\) and for which we can control the eigenvalues.
-One example of such a preconditioner would be???
-Depending on the rank of the preconditioner, we can then eliminate as many of the largest / lowest eigenvalues of \(\bm \Lambda\) as the rank of the preconditioner.
-If we have a rank-one preconditioner whose only eigenvector is the largest eigenvector of \(\bm \Lambda\), the second-largest eigenvalue becomes the largest.
-In the best-case scenario, you can then eliminate
-If \(\bm \Lambda\) has a very flat spectrum, then the problem cannot be ill conditioned and thus preconditioning is immaterial.
+To do so, one could hypothetically attempt to compute the eigenvectors of \(\bm \Lambda\) and then use these eigenvectors as the columns of the preconditioner.
+% Depending on the rank of the preconditioner, we can then eliminate as many of the largest / smallest eigenvalues of \(\bm \Lambda\) as the rank of the preconditioner.
+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).
% 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?}
-Finally, performance of computing
-In CIQ, we will, in each iteration, perform one matrix-vector product \(\bm \Lambda \bm v\) and one preconditioner-vector product \textbf{check this}, 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 ....
-However, an effective preconditioner will significantly improve the convergence speed of CIQ, and this effect can often outweigh the cost of computing the preconditioner.
+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.
+However, we usually choose preconditioners which are relatively easy to compute relative to the posterior precision itself.
+Additionally, an effective preconditioner can significantly improve the convergence speed of CIQ, and this effect can often outweigh the cost of computing the preconditioner.
+Thus, we will often choose to use a preconditioner which is not exact, but which is still ``close enough'' to \(\bm \Lambda\) and is cheap to compute.
+We will now discuss some potential choices for preconditioners to sample using CIQ.
\subsection{Fully Linear GGN}
\label{sec:adam-preconditioner}
-If we calculate the GGN approximation of the Hessian by linearizing over the whole model and loss (as described in \cref{sec:practical-ggn}), we get
+
+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},
\end{align}
-where \(\nabla\) is the gradient of the loss with regard to the model parameters.
-DO I NEED TO PROVE THIS?
+where \(\nabla\) is the gradient of the loss with regard to the model parameters.\sidenote{DO I NEED TO PROVE THIS?}
The problem with using this approximation is that the rank of \(\bm P\) is 1.
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, it seems that \emph{this is not the case}).
+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 conditioning number of our problem will not be significantly reduced.
-As such, this preconditioner is not worth pursuing.
-
-like adam
+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}
-We could just choose to calculate the Hessian over fewer points and invert it using the conjugate gradient method to find our preconditioner.
-This would have the advantage of
-\begin{align}
- \lambda\left( J\T J + \alpha I \right) ={} & \{ \underbrace{s_1, s_2, \ldots, s_k}_{k}, \underbrace{\alpha, \ldots, \alpha}_{n - k} \}
- \\ \lambda\left( N J\T J + \alpha I \right) ={}& \{ \underbrace{N s_1, N s_2, \ldots, N s_k}_{k}, \underbrace{\alpha, \ldots, \alpha}_{n - k} \}
- \\ \lambda\left( \frac{N}{B} \textstyle{\sum^B_i} J_i\T J_i + \alpha I \right) ={}& \{ \underbrace{?, \ldots, ?}_{\leq B \cdot k}, \underbrace{\alpha, \ldots, \alpha}_{\geq n - B \cdot k} \}
- % \Span_\lambda\left( J\T J + \alpha \identity \right)
-\end{align}
+\label{sec:sub-sampling-preconditioner}
-If the eigenvectors of \(J_i\T J_i\) are all (approximately) the same, then the per-observation Jacobian term \(\textstyle{\sum^N_i} J_i\T J_i\) has \(k\) non-zero eigenvalues (the same as \(J\T J\)).
-If these eigenvectors are all different (orthogonal?), then this per-observation Jacobian term will have \(B \cdot k\) non-zero eigenvalues.
+We could just choose to calculate the true posterior precision over fewer observations and invert it using the conjugate gradient method to use it as a preconditioner.
+This would have the advantage of being guaranteed to converge to the true posterior precision as the number of preconditioner observations increases.
+The spectrum of this preconditioner would thus be similar to that of the true posterior precision.
+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}\)).
+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.
-The issue with a few of these ideas is that, when we scale up the \(J\T J\) term to account for multiple observations, we worsen the conditioning number of our preconditioner, which means our conjugate gradient solution is worse.
-Now, dealing with poor conditioning only in the inversion of our preconditioner might be better than dealing with it everywhere else, but it still feels very weird that the preconditioner suffers from the same problems that it's trying to solve.
-
-\[\textstyle{\sum^N_i} s_1 v_i / \alpha \leq \kappa \leq N s_1 / \alpha\]
-
-TALK ABOUT PRECONDITIONER WITH ALL OBSERVATIONS BUT DIFFERENT ALPHA
-THIS POSES ANOTHER ISSUE: THE PRECONDITIONER BECOMES WAY MORE EXPENSIVE ON ALL OBSERVATIONS
+% \[\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.}
-Since \(N \gg B\), then the cost of computing the preconditioner will be negligible compared to that of the GGN-vector product.
+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.}
+
+\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.\sidenote{EXPLAIN WOODBURY MATRIX IDENTITY MORE MAYBE?}
+This identity can be used to compute the inverse of the GGN matrix with the prior precision, which is given by
+\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
+\end{align}
+where \(\alpha\) is the prior precision.
+
+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.
+This may be a reasonable assumption for many problems, but is not guaranteed to be true.
+
+However, even if the eigenvectors of the GGN matrix are constant across the dataset, the eigenvalues of this preconditioner should be large enough to reduce the condition number of the posterior precision matrix.
+To do so, we can scale the outer product \(\bm J \bm J\T\) term by some factor, e.g., the number of observations, \(N\).
+Since, by scaling up the outer product term we are effectively scaling the eigenvalues of \(\bm J \bm J\T\) by a constant, the condition number of this outer product will remain the same.
+In this way, we can increase the maximum eigenvalues of the preconditioner without increasing the condition number of the outer product \(\bm J \bm J\T\) which we want to invert using the conjugate gradient method.
+
+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 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.
+
+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.
+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\)).
\subsection{Pivoted Cholesky}
\label{sec:pivoted-cholesky}
-\sidecite{harbrecht2012low,bach2013sharp,gardner2018gpytorch}
-
+The pivoted Cholesky factorisation is a method for computing a low-rank approximation to the Cholesky decomposition of a matrix~\cite{harbrecht2012low,bach2013sharp}.
+This decomposition has been successfully applied to preconditioning for kernel matrices in Gaussian processes~\cite{gardner2018gpytorch}, but not to preconditioning for the posterior precision of neural networks using the Laplace approximation.
+Given the posterior precision matrix \(\bm \Lambda\) from the Laplace approximation, we can compute the pivoted Cholesky factorisation as
\begin{align}
\bm \Lambda ={} & \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity
\\ \approx{}& \bm L\T \bm L + \alpha \identity
\end{align}
-where \(\bm L \in \reals^{k \times n}\), with \(k\) being the chosen rank of the pivoted Cholesky factorisation and \(n\) being the number of parameters.
-We can then use this approximation as a preconditioner by computing its inverse via the Woodbury matrix identity (\cref{sec:woodbury}):
+where \(\bm L \in \reals^{K \times D}\), with \(K\) being the chosen rank of the pivoted Cholesky factorisation and \(D\) being the number of parameters.
+We can then use this approximation as a preconditioner by computing its inverse via the Woodbury matrix identity, similarly to \cref{sec:woodbury-preconditioner}, as
\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 then obtain the preconditioner-vector and inverse-preconditioner-vector products as in Equation~\ref{eq:pivoted-cholesky-preconditioner}.
-The ``scariest'' computation in the inverse preconditioner is inverting a \(k \times k \) matrix, where \(k\) is the rank of our pivoted Cholesky factorisation (which we choose).
+The ``scariest'' computation in the inverse preconditioner is inverting a \(K \times K\) matrix, where \(K\) is the rank of our pivoted Cholesky factorisation (which we choose).
\begin{align}\label{eq:pivoted-cholesky-preconditioner}
\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 )
+ \\ \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}
+Here it is possible to precompute the inverse factorised product given by \(\bm L\T ( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{D \times K}\), meaning that we only need to perform two matrix multiplications for both the preconditioner- and inverse preconditioner-vector products.
+This would then require storing the factorisation \(\bm L \in \reals^{K \times D}\) and the inverse factorised product \(\reals^{D \times K}\) matrix.\sidenote{We can also precompute the inverse \(( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{K \times K}\), which means we would only need to store the factorisation \(\bm L \in \reals^{K \times D}\) and this inverse \(K \times K\) matrix, in which case we would just need to perform one more matrix multiplication (multiplying \(\bm L\T\) with this \(K \times K\) matrix) in every preconditioner call.}
-We can precompute the factorisation-inverse product \(\bm L\T ( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{n \times k}\), meaning that we only need to perform two matrix multiplications for both the preconditioner- and inverse preconditioner-vector products.
-This would then require storing the factorisation \(\bm L \in \reals^{k \times n}\) and the factorisation-inverse product \(\reals^{n \times k}\) matrix.\sidenote{We can also precompute the inverse \(( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{k \times k}\), which means we would only need to store the factorisation \(\bm L \in \reals^{k \times n}\) and this inverse \(k \times k\) matrix, in which case we would just need to perform one more matrix multiplication (multiplying \(\bm L\T\) with this \(k \times k\) matrix) in every preconditioner call.}
-
-ADD SOMEWHERE:
-As we sum over multiple low-rank matrices, we will likely find that the minimum non-zero eigenvalue gets progressively smaller compared to the largest.
-When we have a single rank \(o\) GGN matrix, we have
+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:
Regarding the eigenvalues and eigenvectors, we consider a best-case scenario.
-Assume that the pivoted Cholesky factorisation finds the \(k\) largest eigenvectors.
+Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors.
If the eigenspectrum of this factorisation is steep enough, then the largest eigenvalue \(s_1\) will be significantly larger than the smallest eigenvalue \(s_1\) (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 same as the eigenvalues of the individual Hessians.
However, some of the eigenvectors may be close to each other, and so those eigenvalues will be approximately summed in the summed Hessian's spectrum.
@@ -352,84 +410,88 @@ However, the rank of the total Hessian will increase by the number of observatio
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 eigenspectrum and a larger condition number of the total Hessian.
Thus, when the batch size \(B\) is large, the condition number 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).
-REWRITE THIS, 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).\sidenote{I SHOULD REWRITE THIS PARAGRAPH, IT'S KINDA HARD TO EXPLAIN.}
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.
-A common approach to sampling the largest eigenvalues is to sample vectors with probability proportional to the diagonal value~\sidecite{harbrecht2012low}.
+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.
+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.
-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}.
+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.
+Notably, the quality of the preconditioner depends largely on the ease of inverting the \(\identity + \alpha^{-1} \bm L \bm L\T\) matrix, which depends on its condition number.
+In practice, increasing the rank of the pivoted Cholesky factorisation will increase the condition number of this matrix, since we are approximating more of the largest eigenvectors and so if we approximate more eigenvectors, we will capture progressively smaller eigenvalues.
+I've found that it's most effective to keep \(K\) low, such that the \(K \times K\) outer product of the pivoted Cholesky factorisation matrix \(L\) approximates few enough of the largest eigenvectors that the eigenvalue of the largest is close to that of the smallest (meaning this outer product is well-conditioned).
+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.
-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.
-Notably, the quality of the preconditioner depends largely on the ease of inverting the \(\identity + \alpha^{-1} \bm L \bm L\T\) matrix, which depends on its conditioning number.
-In practice, increasing the number
-I've found that it's most effective to keep \(k\) low, such that the \(k \times k\) outer product of the pivoted Cholesky factorisation matrix \(L\) approximates few enough of the largest eigenvectors that the eigenvalue of the largest is close to that of the smallest (meaning this outer product is well-conditioned).
-You thus need a large enough \(k\) to approximate enough large eigenvectors and a small enough \(k\) to ensure the conditioning 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 (why not include all of them, always?).
+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.
+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[Randomly Pivoted Cholesky]{Randomly Pivoted Cholesky}[Random Pivoted Cholesky]
+\subsection{Other Preconditioners}
Similarly to the pivoted Cholesky preconditioner, the randomly pivoted Cholesky preconditioner requires the diagonal of the GGN matrix in order for the decomposition to be close to the GGN.
We can thus approximate this term as we do in the pivoted Cholesky preconditioner (see \cref{sec:pivoted-cholesky}).
+However, the randomly pivoted Cholesky preconditioner also requires the diagonal of the GGN matrix to be computed.
+Early tests of the randomly pivoted Cholesky preconditioner were not successful, and so it was not pursued further.
-\subsection[Other Laplace Approximations]{Other Laplace Approximations}[Other approximations]
-
-
-
-% Last-Layer Laplace
-
-
-% K-FAC Laplace
-
+Additionally, we considered potential preconditioners to be computed as traditional approximations of the posterior precision, like the diagonal approximation or the Kronecker factorisation.
+These methods have been shown to be relatively effective approximations of the posterior precision in the context of Bayesian neural networks~\cite{daxberger2021laplace}, and so we considered them as potential preconditioners, as they may contain enough information about the full posterior precision to be effective preconditioners.
However, due to the speed, effectiveness, and theoretical properties of other preconditioners proposed in this section, these methods were not attempted.
-
-\subsection[A Recap on Preconditioning]{A Recap on Preconditioning}[A Recap]
-
-It is worth noting the issues that arise from scaling the term in the \(\sum \bm J_i\T \bm H_i \bm J_i\) preconditioner to account for the summing over a Hessian of a limited batch size.\sidenote{Or, more precisely, the term which is equivalent to the \(\bm J_i\T \bm H_i \bm J_i\) term in the Hessian and is consequently used to approximate the largest eigenvalue.}
-
-\begin{enumerate}
- \item Never scale your preconditioner up
-\end{enumerate}
+Overall, both the sub-sampled preconditioner and the pivoted Cholesky preconditioners were found to be effective under the correct settings.
+However, the pivoted Cholesky preconditioner is not easy to implement in JAX and was more difficult to tune than the sub-sampled preconditioner.
+The sub-sampled preconditioner is thus the preferred preconditioner in this project.
\section{Sampling Evaluation}
\label{sec:sampling-evaluation}
-Into blah blah blah blah
-Sampling is sensitive to the conditioning number of the posterior precision.
+As we have explained, sampling through CIQ is sensitive to the condition number of the posterior precision.
As such, we need a way to evaluate the quality of the samples drawn from the approximate posterior to determine if sampling occurred correctly.
This means evaluating whether the samples are drawn from a normal distribution with the correct mean and precision.
-What we would like to fundamentally test is whether the CIQ algorithm converged correctly and that the samples were thus correctly generated.
+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 conditioning), we cannot directly evaluate the covariance of the samples.
+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.
-\subsection[Mahalanobis Distance and the Chi-Square Distribution]{Mahalanobis Distance and the Chi-Square Distribution}[Mahalanobis Distance]
+\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-square distribution with \(D\) degrees of freedom, as per
+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
\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 normalised and their sum of square deviations be calculated by computing their squared Mahalanobis distances, which will be Chi-square distributed with \(D\) degrees of freedom, as per
+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
%
\begin{align}
\bm d ={} & \sqrt{(\bm \varepsilon - \bm \mu)\T \bm \Lambda (\bm \varepsilon - \bm \mu)}\label{eq:mahalanobis-distance}
- \\ \Rightarrow{}& \bm d\T \bm d \sim \chi^2(D) \nonumber
+ \\ \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\).
For the Laplace approximation, this can easily be computed as $\bm \Lambda = - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm x)$.
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 \Sigma)\) while only accessing the precision matrix in the evaluation.
+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.
+However, this method does not provide a quantitative measure of the quality of the samples.
+In order to run experiments using the Laplace approximation without manually checking the plots of every set of samples we generate, we also need a method to evaluate the quality of the samples quantitatively.
+Our current methods do not provide this, and so we need to evaluate the quality of the samples in a different way.
+
+The Kolmogorov-Smirnov test is a non-parametric test for comparing two continuous distributions.
+It is a test of whether two samples are drawn from the same distribution.
+Since the chi-squared distribution is continuous, we can use the Kolmogorov-Smirnov test to compare the chi-squared distribution with \(D\) degrees of freedom to the squared Mahalanobis distances of the samples.
+This test is performed by computing the Kolmogorov-Smirnov statistic, which is the maximum difference between the empirical and theoretical CDFs.
+The null hypothesis of the test is that the samples are drawn from the same distribution, and so the test is rejected if the Kolmogorov-Smirnov statistic is greater than the critical value for the test.
+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 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.
-\subsection{Kolmogorov-Smirnov}
+% \section{Limitations}
-\section{Limitations}
+% The CIQ algorithm
-Scaling with data?
-Any comments on the model size and how it scales?
+% Scaling with data?
+% Any comments on the model size and how it scales?
diff --git a/chapters/theory.tex b/chapters/theory.tex
@@ -6,12 +6,12 @@
\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\) parameterised 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\) parameterised in \(\bm\theta\) which maps an input \(\bm x\) 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 \(\mathcal{L}(\bm\theta): \reals^d \rightarrow \reals\) such that the optimal parameters \(\bm\theta^*\) minimise this loss function.
+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[][5cm]{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.
@@ -19,81 +19,98 @@ 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 x \given \bm \theta) ={} & \prod_i^N p(x_i \given \bm \theta) \nonumber
- \\ \mathcal{L}(\bm\theta) \equiv \mathcal{L}(\bm \theta, \bm x) := -\log p(\bm x \given \bm \theta) ={}& - \sum_i^N \log p(x_i \given \bm \theta).
+ p(\bm y \given \bm \theta) ={} & \prod_i^N p(x_i \given \bm \theta) \nonumber
+ \\ \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).
\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}[Gradient Descent]
+\subsection[Stochastic Gradient Descent]{Stochastic Gradient Descent}[SGD]
-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)\).
+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.}
The update rule is given by
%
-\begin{equation*}
+\begin{align}
\bm\theta_{t+1} = \bm\theta_t - \eta_t \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t)
-\end{equation*}
+\end{align}
%
where \(\eta_t\) is the learning rate at iteration \(t\).
The learning rate \(\eta_t\) can be constant or adaptive, and is often chosen to be a decreasing function of \(t\).
The gradient \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t)\) can be computed separately for each single training example \(\bm x_t, y_t\), as seen in \cref{eq:negative-log-likelihood}.
The gradient is computed using the chain rule as
%
-\begin{equation*}
- \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) = \nabla_{\bm\theta} \log \lik(\bm x_t, y_t \given \bm\theta_t) \nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)
-\end{equation*}
-%
-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\), \(\bm\theta_t\), and \(\eta_t\) can be interpreted as the step size of the gradient descent algorithm.
-It can either be constant throughout training (\(\eta_t = \eta\)) or adaptive (if it varies throughout training).
-
-\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.
-
-\sidetextcite{rattray1998natural} introduced the natural gradient as a way to optimize a function \(f\) parameterised 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}
+\begin{align}\label{eq:sgd-gradient}
+ \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t) ={} & \nabla_{\bm\theta} \log p(y_t \given \bm \theta) \nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)
+\end{align}
%
-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.
-\sidecite{kingma2014adam}
-\sidecite{martens2015optimizing}
-\sidecite{botev2017practical}
-
-For a recent, detailed discussion of the natural gradient, see \sidetextcite{martens2020new}.
-\sidecite{wu2019logan}
+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\).
+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\) parameterised 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}
\subsection{Adam}
-
-% How does adam work?
-
-\sidecite{sutskever2013importance}
-\sidecite{kingma2014adam}
-
-% Integrate with the previous section
-Adam approximates the second moment of the gradient of the loss, that this corresponds to the variance of this gradient, and that this is the diagonal of the Fisher information matrix, as we define in \cref{sec:fisher-information} (\cref{eq:gradient-covariance}).
-% This is a common approximation in practice, as the diagonal of the Fisher information matrix is often sufficient to approximate the curvature of the log-likelihood.
-
-However, these optimisation methods all identify a single set of parameters as the optimal solution.
+\label{sec:adam}
+
+Adam~\cite{kingma2014adam} is a popular alternative to SGD for training neural networks.
+It is an approximate second-order optimisation method, which means that it approximates the curvature of the loss function.
+It is a variant of SGD that uses the first and second moments of the gradient of the loss function to scale the learning rate of each parameter.
+In Adam, the first moment of the gradient is approximated by a moving average of the gradient, and the second moment of the gradient is approximated by a moving average of the squared gradient.
+The moving averages are weighted by the parameters \(\beta_1\) and \(\beta_2\) respectively.
+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 m_t \odot \frac{1}{\sqrt{v_t} + \varepsilon},
+ \\ m_t ={}& (1 - \beta_1^t) \left[ \beta_1 m_{t-1} + (1 - \beta_1) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
+ \\ v_t ={}& (1 - \beta_2^t) \left[ \beta_2 v_{t-1} + (1 - \beta_2) \left( \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right)^2 \right],
+\end{align}
+where \(\odot\) denotes element-wise multiplication, \(\varepsilon\) is a small constant to prevent division by zero, and \(m_t\) and \(v_t\) are the first and second moments of the gradient respectively.
+
+The momentum term~\cite{polyak1964some} is inspired by physics, where the momentum of an object is the product of its mass and velocity.
+In Adam, the first moment of the gradient is used to compute the momentum of the parameters, increasing the convergence speed of the optimisation~\cite{sutskever2013importance}.
+The second moment of the gradient is also known as the variance of the gradient.
+Adam approximates the second moment by the diagonal of the Fisher information matrix.
+This is a common approximation in practice, as the diagonal of the Fisher information matrix is easy to compute.
+The Fisher information matrix is a discussed further in \cref{sec:hessian}.
+The use of the second moment is inspired by optimisation methods such as Newton's method and the natural gradient~\cite{rattray1998natural}, which use the second derivative of the loss function to compute the optimal step size by scaling the step by the inverse of the curvature of the loss function.
+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}
@@ -104,36 +121,45 @@ From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and
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}
- p(\bm \theta \given \bm x) = \frac{p(\bm x \given \bm \theta)\,p(\bm \theta)}{p(\bm x)},
+ p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta)\,p(\bm \theta)}{p(\bm y)},
\end{align}
-where \(p(\bm \theta \given \bm x)\) is the posterior distribution over the parameters, \(p(\bm \theta)\) is the prior distribution over parameters, \(p(\bm x \given \bm \theta)\) is the likelihood of the data given the parameters, and \(p(\bm x)\) is the marginal likelihood of the data.
-The marginal likelihood is obtained by marginalising over the parameters, i.e., \(p(\bm x) = \int p(\bm x \given \bm \theta)\,p(\bm \theta)\,d\bm \theta\).
+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.}
-\begin{align}
- p(\bm \theta \given \bm x) \propto p(\bm x \given \bm \theta)\,p(\bm \theta) =: \tilde{p}(\bm \theta \given \bm x) .
+\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}
+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.
Similarly to the negative log-likelihood, we can factorise the negative log-posterior loss as per
\begin{align}\label{eq:negative-log-posterior}
- \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm x) ={} & - \log p(\bm x \given \bm \theta) - \log p(\bm\theta) \nonumber
+ \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta) \nonumber
\\ ={}& \underbrace{-\sum_i^N \log p(x_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}.
\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.
-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.,
-\begin{align}
- p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I).
-\end{align}
+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)\).
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.
-% Instead of learning a single set of parameters, Bayesian deep learning learns a distribution over parameters.
+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.
+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).
+However, they cannot be used to compute the posterior distribution itself directly.
+Others, like variational inference (VI)~\cite{jordan1999introduction,wainwright2008graphical}, are methods which attempt to find a tractable approximation to the posterior distribution by parameterising a family of approximate posteriors and optimising the parameters to minimise the KL divergence between the approximate posterior and the true posterior.
+The performance of these methods depends on the variational family used to approximate the posterior and the choice of the variational parameters.
+One common choice for selecting the variational family is to use a Gaussian distribution with diagonal covariance matrix, which is known as mean-field VI.\sidenote{The mean-field assumption refers to the assumption that the parameters are independent of each other which leads to a diagonal covariance matrix.}
+Many variational methods suffer from the difficulty of their optimisation problem, since they often attempt to train variance parameters, which are difficult to optimise~\cite{skafte2019reliable}.
+In this project, we focus on the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, in which the posterior is approximated by a Gaussian distribution, but the variance is deduced, rather than optimised.
\section[The Laplace Approximation]{The Laplace Approximation}[Laplace Approximation]
\label{sec:laplace}
-In the Laplace approximation (LA)~\sidecite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI\@.
+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}}\).
Since the Taylor expansion is performed around the MAP solution, the first order derivative is zero, and the expansion is given by
\begin{align}
@@ -143,378 +169,438 @@ Since the Taylor expansion is performed around the MAP solution, the first order
\\ \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) \propto p(\bm \theta \given \bm y) \nonumber
\end{align}
where \(\tilde{p}(\bm \theta \given \bm y)\) corresponds to the unnormalised posterior.
-Normalising it by the normalisation constant \(p(\bm x)\) gives the normalised posterior\sidenote{This normalisation constant is computed as the marginal likelihood for the Laplace approximation in \cref{sec:maximising-the-evidence}.}
+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) \nonumber
% \\={}& \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}).
\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.
+This method is known as post-hoc Laplace.
+However, it is also possible to compute the posterior precision during training, which is known as online Laplace.
+This can be done either to sample from the neural network at each iteration of training and evaluating the Hessian at each sample~\cite{miani2022laplacian} or to estimate the log-marginal likelihood~\cite{daxberger2021laplace,immer2021scalable}, enabling training of hyperparameters and improved model selection.
+We approach the latter, as well as post-hoc Laplace, in this work.
-\marginnote{The equality in \cref{eq:laplace-normal} only holds when the precision matrix \(\bm \Lambda\) is positive definite.}
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.
+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.
-In practice, we can determine the Hessian of the loss with respect to the parameters, where the loss is given by the negative log-posterior.
-In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_{\textsc{map}}} = \left. \nabla^2_{\bm \theta} \mathcal{L} \right|_{\bm \theta_{\textsc{map}}}\) and the posterior is approximated by \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\).
-
-\cref{eq:negative-log-likelihood-ggn}
-
-GGN APPROXIMATION FOR POSITIVE DEFINITENESS
-
-EIGENVALUES?
-minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees positive definiteness
-
-As we have seen, the Laplace approximation is obtained from a second-order Taylor expansion of the log-posterior around the MAP solution, which corresponds to approximating the posterior with a Gaussian distribution with a covariance matrix given by the inverse of the Hessian of the negative log-posterior.
-However, this approximation relies on the positive definiteness of the Hessian, which is not guaranteed for non-convex losses.
-We will now show that the GGN approximation can be used as a positive definite Hessian approximate the posterior even when the Hessian itself is not positive definite.
-
-\section{The Hessian}
+\subsection{The Hessian}
+\label{sec:hessian}
The Hessian is a matrix of second-order partial derivatives of a scalar function.
From now on, we will use the notation \(\nabla^2_{\bm\theta} f\) to represent the Hessian matrix of \(f\), where each element is given by \(\left(\nabla^2_{\bm\theta} f\right)_{i, j} = \frac{\partial^2 f}{\partial\theta_i \partial\theta_j}\).
-Suppose we have a function \(\mathcal{L}: \reals^d \rightarrow \reals\) parameterised by \(\bm \theta \in \reals^d\).
-Then the Hessian can be written as \(J_{\bm \theta}(\nabla_{\bm\theta} \mathcal{L})\).
+Suppose we have a function \(\mathcal{L}: \reals^D \rightarrow \reals\) parameterised by \(\bm \theta \in \reals^D\).
+Then the Hessian can be interpreted as the Jacobian matrix of the gradient of the function, as per \(J_{\bm \theta}(\nabla_{\bm\theta} \mathcal{L})\).
We are interested in the Hessian of the loss function \(\mathcal{L}(\bm\theta)\) with respect to the parameters \(\bm\theta\).
-For a neural network with \(d\) parameters, the Hessian is therefore a \(d \times d\) square matrix.
+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 and positive definite.
-In this case, the Hessian is generally dominated by the block-diagonal~\sidecite{martens2015optimizing}.
+In this case, the Hessian is generally dominated by the block-diagonal~\cite{martens2015optimizing}.
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}, we get ... from \cref{eq:negative-log-posterior-hessian}.\sidenote{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\).}
+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}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( \sum_i^N \log p(x_i \given \bm \theta) + \log p(\bm \theta) \right) \nonumber
- \\ ={}& \nabla^2_{\bm\theta} \sum_i^N \log p(x_i \given \bm \theta) + \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
- \\ ={}& \sum_i^N \nabla^2_{\bm\theta} \log p(x_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(y_i \given \bm \theta) - \log p(\bm \theta) \right) \nonumber
+ \\ ={}& -\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).
\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\).
+However, as we discussed previously, the Hessian of the negative log-likelihood is not always positive definite.
+In the next section, we will discuss how to obtain a positive definite Hessian approximation.
-\section[The Generalised Gauss-Newton Approximation]{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
+\subsection[The Generalised Gauss-Newton Approximation]{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
-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 equation of the 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}
={} & J_{\bm \theta}(\nabla_{\bm \theta} \mathcal{L})
= J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot \nabla_{\bm \theta} f\right) \nonumber
\\ ={}& J_{\bm \theta} (\nabla_{\bm \theta} f) \cdot \nabla_f \mathcal{L} + J_{\bm \theta} \left(\nabla_f \mathcal{L}\right) \cdot \nabla_{\bm \theta} f \nonumber
- \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + \nabla_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm \theta} f
+ \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + \nabla_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm \theta} f,
\end{align}
-with \(\nabla_f\) representing the gradient with respect to \(f(\bm x)\).
+where \(\nabla_f\) is the gradient with respect to \(f(\bm x, \bm \theta)\) and \(J_{\bm \theta}\) is the Jacobian with respect to \(\bm \theta\).
+% The first term is the GGN approximation of the Hessian of the loss function.
Assume now we will approximate \(f\) using a first-order Taylor expansion.
+In this way, we are linearising our neural network function.
This approximation is given by
%
\begin{align}\label{eq:nn-taylor}
f_{\bm \theta}(\bm x) \approx{} f_{\bm \theta_0}(\bm x) + \nabla_{\bm \theta} f_{\bm \theta_0}(\bm x) \cdot (\bm \theta - \bm \theta_0) .
- \\ \text{Notice that } \nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0 \nonumber
+ \\ \text{Notice that } \nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0. \nonumber
\end{align}
%
-% In this approximation, we have \(\nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0\).
-In this way, we are linearising our neural network function.
-By combining \cref{eq:ggn-whole-hessian,eq:nn-taylor} we then obtain the Gauss-Newton Hessian approximation,
+By combining \cref{eq:ggn-whole-hessian,eq:nn-taylor} we then obtain the generalised Gauss-Newton approximation of the Hessian,
%
\begin{align*}
\nabla^2_{\bm\theta} \mathcal{L}
- \approx{} & (\nabla_{\bm\theta} f)\T (\nabla^2_f \mathcal{L}) (\nabla_{\bm\theta} f)
+ \approx{} & \nabla_{\bm\theta} f\T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm\theta} f,
\end{align*}
+which corresponds to the first 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}
+ \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}
+\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))\).
-For mean square error (MSE) loss, we have one output variable (\(o = 1\)), so \(f: \reals^n \times \reals^d \rightarrow \reals\).
-From the expression of MSE loss, we get the GGN-approximate Hessian
+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,
%
\begin{align*}%\label{eq:mse-hessian}
- \nabla^2_{f} \mathcal{L} = & {} \frac{1}{N} \sum_{n=1}^N (y_n - f(\bm x_n))^2 = 1 \nonumber
- \\ \Longrightarrow \nabla^2_{\bm\theta} \mathcal{L} =&{} (\nabla_{\bm\theta} f)\T (\nabla_{\bm\theta} f)
+ \bm H_i \approx & {} (y_i - f(\bm x_i, \bm \theta))^2 = 2 \nonumber
+ \\ \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.
+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 alld \(\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.
-For log-likelihood loss, we get
-\begin{align}\label{eq:negative-log-likelihood-ggn}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \sum_i^N \nabla^2_{\bm\theta} \log p(x_i \given \bm \theta)
- \\ \approx{}& \sum_i^N \left[\nabla_{\bm\theta} f (\bm \theta_0, x_i)\right]\T \left[\nabla^2_f \mathcal{L} (f(\bm \theta_0, x_i), t_i)\right] \left[\nabla_{\bm\theta} f(\bm \theta_0, x_i)\right] \nonumber
- \\ ={}& \sum_i^N J_i\T H_i J_i \nonumber
+% 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}
+ \mathcal{F} = \sum_{i=1}^N \E{p(y \given \bm x_i, \theta)}{\nabla_{\bm \theta} \log p(y \given \bm x_i, \bm \theta) \cdot \nabla_{\bm \theta} \log p(y \given \bm x_i, \bm \theta)\T},
\end{align}
+where \(\E{p(y \given \bm x_i, \theta)}{\bullet}\) denotes the expectation with respect to the likelihood function.
+The empirical Fisher matrix is relevant as it motivates second-order methods such as natural gradient descent (and thus the Adam optimiser) and the use of the GGN approximation.
-\subsection{Practical Considerations}
-% \subsection{Practical Considerations for the GGN Approximation}
+\subsection{Practical Considerations for the GGN}[Practical Considerations]
\label{sec:practical-ggn}
The generalised Gauss-Newton approximation requires that the function for which we would like to compute the Hessian, \(h(x)\), can be framed as a composition of two functions \(h(x) = g(f(x))\).
-However, there are many possible choices of functions \(f\) and \(g\) which, when composed, yield \(h\), and these different compositions can yield wildly varying performance for the GGN approximation~\sidecite{kunstner2019limitations}.
-Given that, in our case, \(h\) is the loss function with respect to the model parameters of a neural network, a natural choice of functions (and the choice we have assumed above) are \(g: \reals^o \rightarrow \reals\) and \(f: \reals^d \rightarrow \reals^o\) where \(g\) is the loss function with respect to the model output and \(f\) is a function which maps the model parameters to a model output, for a given model input \(x\).
+% A GGN approximation is any approximation which fulfils \cref{eq:ggn-whole-hessian}.
+However, there are many possible choices of functions \(f\) and \(g\) which, when composed, yield \(h\), and these different compositions can yield wildly varying performance for the GGN approximation~\cite{kunstner2019limitations}.
+Given that, in our case, \(h\) is the loss function with respect to the model parameters of a neural network, a natural choice of functions (and the choice we have assumed above) are \(g: \reals^O \rightarrow \reals\) and \(f: \reals^D \rightarrow \reals^O\) where \(g\) is the loss function with respect to the model output and \(f\) is a function which maps the model parameters to a model output, for a given model input \(x\).
This construction allows us to capture the curvature of the loss with respect to the model output, which can often be easily computed in a simple closed form, without requiring the computation of the curvature of the model function with respect to its parameters; this therefore corresponds to linearising the model function with a Taylor expansion as in \cref{eq:nn-taylor}.
However, this is not the only possible construction.
Another choice is to linearise the whole likelihood function with respect to the model parameters.
-This yields \(g(y) = -\log(y)\), and thus we obtain
+This yields \(g(y) = -\log(y)\), and thus we obtain the GGN approximation of the Hessian with this composition as
\begin{align}\label{eq:fully-linear-ggn}
\nabla^2_{\bm \theta} \mathcal{L}
- ={} & \nabla f \T \cdot \nabla^2 g \cdot \nabla f \nonumber
- \\ ={}& \nabla f \T \cdot \frac{1}{f(\bm\theta)^2} \cdot \nabla f \nonumber
- % \\ ={}& \left(\nabla f / f(\bm\theta)\right)\T \cdot \left(\nabla f / f(\bm\theta)\right) \nonumber
- \\ ={}& \left(\frac{1}{f(\bm\theta)} \nabla f\right)\T \cdot \left(\frac{1}{f(\bm\theta)} \nabla f\right) \nonumber
+ ={} & \nabla_{\bm \theta} f \T \cdot \nabla_{f}^2 g \cdot \nabla_{\bm \theta} f \nonumber
+ \\ ={}& \nabla_{\bm \theta} f \T \cdot f(\bm\theta)^{-2} \cdot \nabla_{\bm \theta} f \nonumber
+ \\ ={}& \left(f(\bm\theta)^{-1} \cdot \nabla_{\bm \theta} f\right)\T \cdot \left(f(\bm\theta)^{-1} \cdot \nabla_{\bm \theta} f\right) \nonumber
\\ ={}& \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right)\T \cdot \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right) \nonumber
- \\ ={}& \left(\nabla_{\bm\theta} g \circ f \right)\T \cdot \left(\nabla_{\bm\theta} g \circ f \right) \nonumber
+ \\ ={}& \left[\nabla_{\bm\theta} (g \circ f) \right]\T \cdot \left[\nabla_{\bm\theta} (g \circ f) \right] \nonumber
\\ ={}& \nabla_{\bm\theta} \mathcal{L}\T \cdot \nabla_{\bm\theta} \mathcal{L}
\end{align}
The diagonal of this matrix is given by the element-wise square of the gradient \( \nabla_{\bm\theta} \mathcal{L} \odot \nabla_{\bm\theta} \mathcal{L} \).
+Notice that this is very similar to 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.
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 its use for optimisation algorithms that require a positive definite Hessian, such as conjugate gradient descent, or the Laplace approximation, which employs the inverse of the Hessian as the covariance matrix of a Gaussian distribution, thereby assuming a positive semi-definite covariance matrix.
-
-The Levenberg-Marquardt heuristic~\sidecite{pascanu2013revisiting}
-
-A GGN approximation is any approximation which fulfils \cref{eq:ggn-whole-hessian}.
-\sidecite{kunstner2019limitations}
-
-\section{Large-Scale Laplace}
-
-
-\section{Limitations}
-
-A sentence or two on multi-modality.
-
-\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~\sidecite{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~\sidecite{martens2020new}
-\sidecite{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?)
-\sidecite{kunstner2019limitations}
-
-\section[The Woodbury matrix identity]{The Woodbury matrix identity}[Woodbury]
-\label{sec:woodbury}
-
-In many methods, such as Natural Gradient Descent~\sidecite{martens2020new} and the Laplace approximation~\sidecite{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~\sidecite{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 ∗.
+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.
+
+% The Levenberg-Marquardt heuristic~\cite{pascanu2013revisiting}
+
+\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 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.
+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.
+While, theoretically speaking, the eigenvalues of the likelihood GGN matrix are greater than or equal to zero, in practice, due to numerical issues, the eigenvalues can be slightly negative.
+As such, \(\alpha\) must be chosen to be sufficiently large to ensure that the posterior precision is positive definite.
+While this may appear to detract from the effectiveness of the prior, it is important to note that, in most cases, the prior precision is large enough to ensure that the posterior precision is positive definite.
+
+Since \(\bm J_i\) is a rectangular matrix of shape \(O \times D\), the rank of each \(\bm J_i\T \bm H_i \bm J_i\) is at most \(O\), where \(O\) is the number of outputs of the model.
+The rank of the sum of these matrices thus depends on the eigenvectors of each \(\bm J_i\T \bm H_i \bm J_i\).
+If we add \(\bm J_n\T \bm H_n \bm J_n\) into the sum \(\sum_{i=1}^{n-1} \bm J_i\T \bm H_i \bm J_i\), then, if any of the eigenvectors of \(\bm J_n\T \bm H_n \bm J_n\) are not contained in the span of the eigenvectors of the sum, then the rank of the sum will increase.
+This is because the rank of a matrix is the number of linearly independent columns.
+Thus, non-rigorously, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are between terms in the sum, the closer the largest eigenvalue will be to the sum of the largest eigenvalues of each term.
+This means that the rank of the sum of the GGN matrices can increase as we add more data points, and will not be greater than \(N \cdot O\).
+Furthermore, as we add more data points, the eigenvalues of the sum of the GGN matrices will increase.
+Specifically, the largest eigenvalue of the sum of the GGN matrices will be the less than or equal to the sum of the largest eigenvalues of the GGN matrices.
+Additionally, the non-zero spectrum of the sum of the GGN matrices will itself get steeper as we add more data points.
+If the \(K\) non-zero eigenvectors are identical across the GGN matrices, then this spectrum of the sum of the GGN matrices will just be the sum of the spectrum of each GGN matrix and should thus not get steeper.
+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.}
+
+\section{Limitations and Large-Scale Laplace}[Limitations]
+
+One limitation of the Laplace approximation is that it can only approximate one mode of the posterior distribution (the MAP).
+However, most approximate posterior distributions, such as MCMC and variational inference, also struggle to approximate multimodal distributions.
+The standard way to solve this problem is to use a multimodal approximation such as a mixture, typically by training multiple models from different initialisations and then combining them in a deep ensemble.
+This is also possible with the Laplace approximation~\cite{eschenhagen2021mixtures}.
+However, this is not a scalable solution, as it requires training multiple models.
+Additionally, it is not clear whether there is a significant benefit to using a multimodal Laplace approximation over a single mode Laplace approximation (i.e., whether generalisation improves).
+
+Another limitation of the Laplace approximation is that it is not scalable to non-trivial problems for the full-rank Laplace approximation without performing crude approximations.
+This is because the full-rank Laplace approximation requires the computation of the posterior precision, which is a \(D \times D\) matrix, where \(D\) is the number of parameters.
+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).
+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.
+
+ADD ANY OF THIS?
\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.
+ \lambda\left( J\T J + \alpha I \right) ={} & \{ \underbrace{s_1, s_2, \ldots, s_k}_{k}, \underbrace{\alpha, \ldots, \alpha}_{n - k} \}
+ \\ \lambda\left( N J\T J + \alpha I \right) ={}& \{ \underbrace{N s_1, N s_2, \ldots, N s_k}_{k}, \underbrace{\alpha, \ldots, \alpha}_{n - k} \}
+ \\ \lambda\left( \frac{N}{B} \textstyle{\sum^B_i} J_i\T J_i + \alpha I \right) ={}& \{ \underbrace{?, \ldots, ?}_{\leq B \cdot k}, \underbrace{\alpha, \ldots, \alpha}_{\geq n - B \cdot k} \}
+ % \Span_\lambda\left( J\T J + \alpha \identity \right)
\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\),
+% \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}
-% 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
+% \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}
-% \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
+% \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}
-% % 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}
+% % \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}
-\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}
+% % % 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}
-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)\).
+% \section[Krylov Subspace Methods]{Krylov Subspace Methods}[Krylov Methods]
-These methods are always exact after \(N\) iterations, however are often almost exact after \(J \ll N\) iterations~\sidecite{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~\sidecite{giordano2018covariances, nocedal1999numerical}.
+% 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}
-\subsection{The Lanczos Algorithm}
-\label{sec:lanczos}
+% 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)\).
-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 := \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\).
+% 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 Conjugate Gradient Method]{The Conjugate Gradient Method}[The Conjugate Gradient]
+% \subsection{The Lanczos Algorithm}
+% \label{sec:lanczos}
-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}
+% 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}
+% 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/training.tex b/chapters/training.tex
@@ -7,50 +7,52 @@
As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model is equal to
\begin{align}
- p(\bm \theta \given \bm x) = \frac{p(\bm x \given \bm \theta) p(\bm \theta)}{p(\bm x)},
+ p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta) p(\bm \theta)}{p(\bm y)},
\end{align}
-where \(p(\bm x)\) 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}, since it's obtained by marginalising out \(\bm \theta\) from the likelihood, as per
\begin{align}
- p(\bm x) ={} & \int_{\bm \theta} p(\bm \theta \given \bm x)\,d\bm \theta.
+ p(\bm y) ={} & \int_{\bm \theta} p(\bm \theta \given \bm y)\,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 it using a variational approximation, such as the Laplace approximation, or by using a Monte Carlo approximation, such as the Markov chain Monte Carlo (MCMC) method. % autogenerated
+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.
-However, since we're integrating over the posterior, maximising the marginal likelihood will seek minima that are flat, since the posterior will be very broad in these regions.
-This is a desirable property, since it means that the model will be less likely to overfit the data, both from a theoretical standpoint~\sidecite{hochreiter1997flat} and in practice~\sidecite{keskar2016large,jiang2019fantastic,maddox2019simple}.
+Since we're integrating over the posterior, maximising the marginal likelihood will seek minima that are flat, since the posterior will be very broad in these regions.
+This is a desirable property, since it means that the model will be less likely to overfit the data, both from a theoretical standpoint~\cite{hochreiter1997flat} and in practice~\cite{keskar2016large,jiang2019fantastic,maddox2019simple}.
-
-Furthermore, \sidetextcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing k-fold cross validation for all values \(k = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(k\) folds and across all values of \(k\)).
-% This is a very strong result, since it shows that the marginal likelihood is a very good proxy for the true cross-validation performance of a model. autogenerated
+Furthermore, \textcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing k-fold cross validation for all values \(k = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(k\) folds and across all values of \(k\)). TODO: FIX THIS TO BE MORE ACCURATE.
+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.
Finally, 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.
-This is useful, since it allows us to use gradient-based optimisation methods to find the hyperparameters that maximise the marginal likelihood.
+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.
-For the Laplace approximation, we approximate the posterior as \(p(\bm \theta \given \bm x) \approx \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\), where the posterior precision \(\bm \Lambda = -\nabla^2_{\bm\theta} \mathcal L(\bm \theta_{\textsc{map}})\) is the Hessian of the negative log-posterior at the maximum a posteriori (MAP) estimate \(\bm \theta_{\textsc{map}}\).
+For the Laplace approximation, we approximate the posterior distribution as \(p(\bm \theta \given \bm y) \approx \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\), where the posterior precision \(\bm \Lambda = -\nabla^2_{\bm\theta} \mathcal L(\bm \theta_{\textsc{map}})\) is the Hessian of the negative log-posterior at the maximum a posteriori (MAP) estimate \(\bm \theta_{\textsc{map}}\).
We thus obtain the approximate Laplace posterior
\begin{align}\label{eq:laplace-posterior}
- \log p(\bm \theta \given \bm x) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm x) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm x) \cdot (\bm \theta - \bm \theta_{\textsc{map}})
- \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm x) (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
- \\={}& \log p(\bm \theta_{\textsc{map}} \given \bm x) - \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
+ \log p(\bm \theta \given \bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) \cdot (\bm \theta - \bm \theta_{\textsc{map}})
+ \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
+ \\={}& \log p(\bm \theta_{\textsc{map}} \given \bm y) - \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}).\nonumber
\end{align}
-The approximate Laplace marginal is then computed as per
+The approximate Laplace marginal is then computed as
\begin{align}\label{eq:laplace-marginal}
- p(\bm x) ={} \stackrel{\textsc{la}}{\approx}{} & \int_{\bm \theta} p(\bm \theta_{\textsc{map}} \given \bm x) \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}) \right)\,d \bm \theta
- \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm x) \int_{\bm \theta} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
- \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm x)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}.\nonumber
+ p(\bm y) ={} \stackrel{\textsc{la}}{\approx}{} & \int_{\bm \theta} p(\bm \theta_{\textsc{map}} \given \bm y) \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}) \right)\,d \bm \theta
+ \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) \int_{\bm \theta} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
+ % \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) \int_{\bm \theta} \frac{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}}{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
+ % \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) (2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2} \int_{\bm \theta} \frac{1}{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
+ \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}.\nonumber
\end{align}
-Taking the log of the approximate Laplace marginal, we obtain
+Taking the log of this approximate Laplace marginal, we obtain
\begin{align}\label{eq:laplace-log-marginal}
- \log p(\bm x) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm x) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda
- \\ ={}& \log p(\bm x \given \bm \theta_{\textsc{map}}) + \log p(\bm \theta_{\textsc{map}}) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda. \nonumber
+ \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda
+ \\ ={}& \log p(\bm y \given \bm \theta_{\textsc{map}}) + \log p(\bm \theta_{\textsc{map}}) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda. \nonumber
\end{align}
For the normal likelihood and normal prior, we have
\begin{align}\label{eq:normal-likelihood}
- p(\bm x \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 x \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity)
- \\ ={}& \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
+ 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)
+ \\ ={}& \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
\end{align}
and
\begin{align}\label{eq:normal-prior}
@@ -58,102 +60,140 @@ and
\\ \log p(\bm \theta) ={}& - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2,
\end{align}
where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of the likelihood.
-
Combining \cref{eq:laplace-log-marginal,eq:normal-likelihood,eq:normal-prior}, we expand the Laplace approximation to the log-marginal likelihood as
\begin{align}\label{eq:expanded-log-marginal}
- \log p(\bm x) \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
+ \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
\\ & - \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\).
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.}
+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.
+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.
+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{\bm y_i - f_{\bm \theta}(\bm x_i)}^2\) & 1. \(N O / 2 \log \rho\)
+% \\\addlinespace 1. \(\rho / 2 \sum_{i=1}^N \norm{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 x) \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
+% \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
% \\ & - \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\).
% However, in doing so, the authors ignore the terms \(D / 2 \log(2 \pi)\)
% While these terms do not affect the optimisation problem, they are useful for comparing models with different numbers of parameters.
-\section[The Determinant Lower Bound]{The Determinant Lower Bound}[The Lower Bound]
+\section{The Determinant Bound}
+\label{sec:determinant-bound}
-We compute an \emph{upper} bound on the log-determinant of the Hessian of the negative log-posterior \(\bm \Lambda \in \reals^{D \times D}\), i.e.,
-\begin{align}
- \log \det \bm \Lambda \leq{} & B_D(\mu_1, \mu_2, \beta)
- \\ :={}& \begin{pmatrix} \log \beta & \log t \end{pmatrix}
+We wish 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)\).
+\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}
\begin{pmatrix} \beta & t \\ \beta^2 & t^2 \end{pmatrix}^{-1}
\begin{pmatrix} \mu_1 \\ \mu_2 \end{pmatrix},
\end{align}
-where
-\begin{align}\label{eq:mu1-mu2}
- \mu_1 ={} & \Tr(\bm \Lambda)
- \\ \mu_2 ={}& \norm{\bm \Lambda}^2_F
- \\ t ={}& \frac{\beta \mu_1 - \mu_2}{\beta D - \mu_1},
-\end{align}
-where \(\beta\) is an upper bound on the largest eigenvalue of \(\bm \Lambda\).
-Unfortunately, well-known methods such as the Lanczos algorithm compute a \emph{lower} bound on the largest eigenvalue.
-
-MMM THIS IS NOT THAT CLEAN WITHOUT RANK O:
-Since \(\bm \Lambda\) is a low-rank matrix given by \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\), it has eigenvalues \(\{\lambda_1 + \alpha, \ldots, \lambda_{O \cdot N} + \alpha, \alpha, \ldots, \alpha\}\).
-\begin{align}
- \log \det \bm \Lambda = \sum_{i=1}^{O} \log(\lambda_i + \alpha) + \sum_{}
-\end{align}
-where \(\lambda_1, \ldots, \lambda_O\) are the eigenvalues of \(\bm \Lambda\).
-
+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.
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.
-WHERE DOES THE BOUND COME FROM? WHICH UNCERTAINTY? THE UPPER BOUND ON THE EIGENVALUE?
-
-We then need to compute the trace of the Hessian matrix, as per \cref{eq:mu1-mu2}.
+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, and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
+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 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}
+ \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 \nonumber
+ \\ ={}& \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
+\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 \nonumber
+ \\ \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 \nonumber
+\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\).
+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]{Hutchinson's Trace Estimator}[The Trace Estimator]
+\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 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 the trace of the precision matrix as
-\begin{align}
- \mu_1 ={} & \Tr(\bm \Lambda) = \Tr\left(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} \bm \Lambda\right)
- \\ ={}& \E{}{\Tr(\bm \varepsilon_0 \bm \varepsilon_0\T \bm \Lambda)} = \E{}{\Tr{\bm \varepsilon_0\T \bm \Lambda \bm \varepsilon_0}}
- \\ ={}& \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \bm \varepsilon_0} + D \alpha.
+Since \(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} = \identity\), we compute \(\mu_1\) as
+\begin{align}\label{eq:hutchinson-trace}
+ \mu_1 ={} & \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) + K \alpha
+ \\ \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) ={}& \Tr\left(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)\right) \nonumber
+ \\ = \E{}{\Tr(\bm \varepsilon_0 \bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right))} \nonumber
+ % \\ = \E{}{\Tr\left(\bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) \bm \varepsilon_0\right)} \nonumber
+ ={} & \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \bm \varepsilon_0}. \nonumber
\end{align}
-Similarly, we compute the trace of the square precision matrix as
-\begin{align}
- \Tr(\bm \Lambda^2) ={} & \Tr\left(\left( {\textstyle {\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) + D \alpha^2
- \\ ={} & \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right)\T \bm \varepsilon_0} + D \alpha.
+Similarly, we compute \(\mu_2\) as
+\begin{align}\label{eq:hutchinson-square-trace}
+ \mu_2 =\Tr\left(\left( {\textstyle {\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
+ \\ \Tr\left(\left( {\textstyle {\textstyle \sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i \right)^2\right) = \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right)\T \bm \varepsilon_0}. \nonumber
\end{align}
-Fortunately, this means that we only need to compute one GGN-vector product
-% \begin{align}
-
-% \end{align}
-
-Furthermore, the trace of
-\begin{align}
- \Tr(\bm \Lambda) ={} & \Tr({\textstyle{\sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i) = \sum_{i=1}^N \Tr(\bm J_i\T \bm H_i \bm J_i)
- \\ ={}& \sum_{i=1}^N \Tr(\bm H_i \bm J_i \bm J_i\T)
+We then apply a Monte Carlo approximation of the expectations by sampling \(S\) samples from \(\bm \varepsilon_0\sim \normal(\bm 0, \identity)\) to compute the traces.
+Fortunately, this means that we only need to compute one GGN-vector product \(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \cdot \bm \varepsilon_0\) to estimate the \(\mu_1\) and \(\mu_2\) terms in \cref{eq:laplace-log-determinant-bound}.
+Furthermore, we have that the trace can be rearranged and simplified to
+\begin{align}\label{eq:trace-rearrangement}
+ \Tr\big({\textstyle{\sum_{i=1}^N}} \underbrace{\bm J_i\T \bm H_i \bm J_i}_{D \times D}\big) ={} & \sum_{i=1}^N \Tr\left(\bm J_i\T \bm H_i \bm J_i\right) = \sum_{i=1}^N \Tr\big(\underbrace{\bm H_i \bm J_i \bm J_i\T}_{O \times O}\big),
\end{align}
+which means we can reduce the dimensionality of \(\bm \varepsilon_0\) to the number of model outputs \(O\) instead of the (significantly larger) number of parameters \(D\), reducing the variance and computational cost of the trace estimator.
+Empirically, we find that, without the trace rearrangement in \cref{eq:trace-rearrangement}, the variance of the trace estimator is too high to be useful.
\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.
+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
+\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)
+ \\ ={} & D \log \frac{N}{B} + \log \det \left( \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \frac{B}{N} \alpha \identity \right). \nonumber
+\end{align}
+The scaling is thus important for determining the relative importance of the prior and likelihood terms in the eigenvalues.
+To scale the posterior precision correctly, we thus need to scale the likelihood term by \(N / B\).
+We do this by scaling the trace and square traces for \cref{eq:laplace-log-determinant-bound} as
+\begin{align}
+ \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.
+% \subsection{Limitations}
-\subsection{Computational Cost}
-
-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.
-
-\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
+% Training hyperparameters
+% Speed for convolutions
+% Scalability to huge networks hasn't been tested
+% Rank estimation
diff --git a/kao.sty b/kao.sty
@@ -879,6 +879,8 @@
}
}
+\RequirePackage[ruled, vlined]{algorithm2e}
+
%----------------------------------------------------------------------------------------
% TOC, LOF & LOT
%----------------------------------------------------------------------------------------
@@ -900,11 +902,13 @@
% Define the style for toc entries
\@ifpackagelater{scrbase}{2019/10/11}{%
\DeclareTOCStyleEntries[indent=0em,numwidth=2em,dynnumwidth=yes,pagenumberwidth=2.1em]{tocline}{figure,table,lstlisting}%
+ \DeclareTOCStyleEntries[beforeskip=-6pt,level:=figure,indent=0em,numwidth=2em,dynnumwidth=yes,pagenumberwidth=2.1em]{tocline}{algocf}%
\DeclareTOCStyleEntries[dynnumwidth=yes]{tocline}{subsubsection,subsection,section,chapter,part}%
\DeclareTOCStyleEntries[pagenumberwidth=2.5em]{tocline}{chapter,part}%
\DeclareTOCStyleEntries[pagenumberwidth=2.1em]{tocline}{subsubsection,subsection,section}%
}{%
\DeclareTOCStyleEntries[indent=0em,numwidth=2em,dynnumwidth=yes]{tocline}{figure,table,lstlisting}%
+ \DeclareTOCStyleEntries[beforeskip=-6pt,level:=figure,indent=0em,numwidth=2em,dynnumwidth=yes,pagenumberwidth=2.1em]{tocline}{algocf}%
\DeclareTOCStyleEntries[dynnumwidth=yes]{tocline}{subsubsection,subsection,section,chapter,part}%
}
@@ -1382,7 +1386,7 @@
%verbatim \preto{\@verbatim}{\topsep=0pt \partopsep=0pt }
% Algorithms
-\RequirePackage[ruled, vlined]{algorithm2e} % Algorithms
+% \RequirePackage[ruled, vlined]{algorithm2e} % Algorithms
% \RequirePackage[linesnumbered, ruled, vlined]{algorithm2e} % Algorithms
% Special gliphs
diff --git a/macros.tex b/macros.tex
@@ -34,6 +34,8 @@
\newcommand{\lik}{p\left(\bm y \given \bm \theta\right)}
+\newcommand{\tabf}[1]{\fontseries{b}\selectfont{#1}}
+
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
diff --git a/main.bib b/main.bib
@@ -5,6 +5,16 @@
year = {2013}
}
+@article{bai1996bounds,
+ title = {Bounds for the Trace of the Inverse and the Determinant of Symmetric Positive Definite Matrices},
+ author = {Bai, Zhaojun and Golub, Gene H},
+ journal = {Annals of Numerical Mathematics},
+ volume = {4},
+ pages = {29--38},
+ year = {1996},
+ publisher = {Citeseer}
+}
+
@book{bishop1995neural,
title = {Neural Networks for Pattern Recognition},
author = {Bishop, Christopher},
@@ -19,6 +29,15 @@
year = {2017}
}
+@inproceedings{damianou2013deep,
+ title = {Deep {G}aussian Processes},
+ author = {Damianou, Andreas and Lawrence, Neil D},
+ booktitle = {Artificial Intelligence and Statistics},
+ pages = {207--215},
+ year = {2013},
+ organization = {PMLR}
+}
+
@inproceedings{daxberger2021laplace,
title = {Laplace Redux-Effortless {B}ayesian Deep Learning},
author = {Daxberger, Erik and Kristiadi, Agustinus and Immer, Alexander and Eschenhagen, Runa and Bauer, Matthias and Hennig, Philipp},
@@ -33,6 +52,23 @@
year = {1990}
}
+@article{eschenhagen2021mixtures,
+ title = {Mixtures of {L}aplace Approximations for Improved Post-Hoc Uncertainty in Deep Learning},
+ author = {Eschenhagen, Runa and Daxberger, Erik and Hennig, Philipp and Kristiadi, Agustinus},
+ journal = {arXiv preprint arXiv:2111.03577},
+ year = {2021}
+}
+
+@article{fong2012cg,
+ title = {{CG} Versus {MINRES}: An Empirical Comparison},
+ author = {Fong, David Chin-Lung and Saunders, Michael},
+ journal = {Sultan Qaboos University Journal for Science [SQUJS]},
+ volume = {17},
+ number = {1},
+ pages = {44--62},
+ year = {2012}
+}
+
@article{fong2020marginal,
title = {On the Marginal Likelihood and Cross-Validation},
author = {Fong, Edwin and Holmes, Chris C},
@@ -51,6 +87,13 @@
year = {2018}
}
+@book{gelman1995bayesian,
+ title = {Bayesian Data Analysis},
+ author = {Gelman, Andrew and Carlin, John B and Stern, Hal S and Rubin, Donald B},
+ year = {1995},
+ publisher = {Chapman and Hall/CRC}
+}
+
@article{giordano2018covariances,
title = {Covariances, Robustness and Variational Bayes},
author = {Giordano, Ryan and Broderick, Tamara and Jordan, Michael I},
@@ -104,6 +147,17 @@
publisher = {MIT Press One Rogers Street, Cambridge, MA 02142-1209, USA journals-info~…}
}
+@article{hutchinson1990stochastic,
+ title = {A Stochastic Estimator of the Trace of the Influence Matrix for {L}aplacian Smoothing Splines},
+ author = {Hutchinson, Michael F},
+ journal = {Communications in Statistics-Simulation and Computation},
+ volume = {19},
+ number = {2},
+ pages = {433--450},
+ year = {1990},
+ publisher = {Taylor \& Francis}
+}
+
@inproceedings{immer2021scalable,
title = {Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning},
author = {Immer, Alexander and Bauer, Matthias and Fortuin, Vincent and R{\"a}tsch, Gunnar and Emtiyaz, Khan Mohammad},
@@ -137,6 +191,16 @@
year = {2020}
}
+@article{jordan1999introduction,
+ title = {An Introduction to Variational Methods for Mraphical Models},
+ author = {Jordan, Michael I and Ghahramani, Zoubin and Jaakkola, Tommi S and Saul, Lawrence K},
+ journal = {Machine learning},
+ volume = {37},
+ pages = {183--233},
+ year = {1999},
+ publisher = {Springer}
+}
+
@article{keskar2016large,
title = {On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima},
author = {Keskar, Nitish Shirish and Mudigere, Dheevatsa and Nocedal, Jorge and Smelyanskiy, Mikhail and Tang, Ping Tak Peter},
@@ -234,6 +298,17 @@
year = {2020}
}
+@article{polyak1964some,
+ title = {Some Methods of Speeding Up the Convergence of Iteration Methods},
+ author = {Polyak, Boris T},
+ journal = {USSR Computational Mathematics and Mathematical Physics},
+ volume = {4},
+ number = {5},
+ pages = {1--17},
+ year = {1964},
+ publisher = {Elsevier}
+}
+
@article{rattray1998natural,
title = {Natural Gradient Descent for On-Line Learning},
author = {Rattray, Magnus and Saad, David and Amari, Shun-ichi},
@@ -267,6 +342,13 @@
year = {2020}
}
+@inproceedings{skafte2019reliable,
+ title = {Reliable Training and Estimation of Variance Networks},
+ author = {Skafte, Nicki and J{\o}rgensen, Martin and Hauberg, S{\o}ren},
+ booktitle = nips,
+ year = {2019}
+}
+
@inproceedings{sutskever2013importance,
title = {On the Importance of Initialization and Momentum in Deep Learning},
author = {Sutskever, Ilya and Martens, James and Dahl, George and Hinton, Geoffrey},
@@ -285,6 +367,26 @@
publisher = {SIAM}
}
+@article{wainwright2008graphical,
+ title = {Graphical Models, Exponential Families, and Variational Inference},
+ author = {Wainwright, Martin J and Jordan, Michael I and others},
+ journal = {Foundations and Trends{\textregistered} in Machine Learning},
+ volume = {1},
+ number = {1--2},
+ pages = {1--305},
+ year = {2008},
+ publisher = {Now Publishers, Inc.}
+}
+
+@book{williams2006gaussian,
+ title = {Gaussian Processes for Machine Learning},
+ author = {Williams, Christopher KI and Rasmussen, Carl Edward},
+ volume = {2},
+ number = {3},
+ year = {2006},
+ publisher = {MIT press Cambridge, MA}
+}
+
@book{woodbury1950inverting,
title = {Inverting Modified Matrices},
author = {Woodbury, Max},
diff --git a/mathjax.tex b/mathjax.tex
@@ -0,0 +1,7 @@
+\newcommand{\bm}[1]{\boldsymbol{#1}}
+% \newcommand{\asfasdfa}{}
+% TeX: {
+% Macros: {
+% bm: ["\\boldsymbol{#1}", 1]
+% }
+% }
diff --git a/tables/mnist_ablation.tex b/tables/mnist_ablation.tex
@@ -1,13 +1,13 @@
\begin{tabular}{rrr}
\toprule
- \(q\) & p-value & Time [m] \\
+ \(q\) & p-value & Time [m] \\
\midrule
- 1 & 0.0000 & 10.6 \\
- 2 & 0.0000 & 18.8 \\
- 3 & 0.0000 & 23.4 \\
- 4 & 0.0001 & 23.4 \\
- 5 & 0.0284 & 23.9 \\
- 10 & 0.0542 & 24.4 \\
- 20 & 0.0561 & 24.9 \\
+ 1 & 0.0000 & 10.6 \\
+ 2 & 0.0000 & 18.8 \\
+ 3 & 0.0000 & 23.4 \\
+ 4 & 0.0001 & 23.4 \\
+ 5 & \tabf{0.0284} & 23.9 \\
+ 10 & \tabf{0.0542} & 24.4 \\
+ 20 & \tabf{0.0561} & 24.9 \\
\bottomrule
\end{tabular}
