commit cb24c7dcc378f15be0735100084c55ea7665adce
parent 35cb369b3befee5f85202f97ab58ed19b1ff4ab7
Author: Silas Brack <s174433@student.dtu.dk>
Date: Sat, 11 Feb 2023 15:46:45 +0100
Almost final.
Diffstat:
16 files changed, 191 insertions(+), 151 deletions(-)
diff --git a/chapters/ablation_experiments.tex b/chapters/ablation_experiments.tex
@@ -94,7 +94,7 @@ This p-value is then used to determine whether the null hypothesis that the samp
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}.
Quantile--quantile plots for each value of \(Q\) are shown in the appendix, in \cref{fig:ablation-chisq}.
-These experiments were run on a single NVIDIA A100 GPU.
+These experiments were, again, run on a single NVIDIA A100 GPU.
\begin{margintable}
\centering
@@ -103,7 +103,7 @@ These experiments were run on a single NVIDIA A100 GPU.
\include{tables/mnist_ablation.tex}
\end{margintable}
-Analysis of the p-values in \cref{tab:mnist-ablation} suggests that sampling is successful for \(Q \geq 10\) with a significance threshold of \(0.05\).
+The p-values in \cref{tab:mnist-ablation} suggest that sampling is successful for \(Q \geq 10\) with a significance threshold of \(0.05\).
However, 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.
In our case, 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.
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -27,23 +27,23 @@ This is the first method for performing the Laplace approximation while only acc
\addcontentsline{toc}{chapter}{Resum\'e}
Bayesianske metoder lover at give en principiel m{\aa}de at kvantificere usikkerheden i neurale netv{\ae}rk p{\aa}.
-Dette er vigtigt for mange anvendelser inden for maskinl{\ae}ring, f.eks. i forbindelse med sikkerhedskritiske beslutninger, der er afh{\ae}ngige af risikovurdering og fortolkningsmuligheder.
+Dette er vigtigt for mange anvendelser inden for machine learning, f.eks. i forbindelse med sikkerhedskritiske beslutninger, der er afh{\ae}ngige af risikovurdering og fortolkningsmuligheder.
Bayesiansk inferens er imidlertid ofte beregningsm{\ae}ssigt uh{\aa}ndterbar, s{\aa} der anvendes tiln{\ae}rmede metoder, som g{\aa}r p{\aa} kompromis med ydeevnen for at opn{\aa} praktisk anvendelighed.
-Laplace-approximationen (LA) er en s{\aa}dan metode, selv om den typisk er afh{\ae}ngig af grove tiln{\ae}rmelser af den efterfd{\o}lgende pr{\ae}cisionsmatrix for at gd{\o}re den beregningsm{\ae}ssigt gennemfd{\o}rlig.
-N{\aa}r vi ikke faktoriserer den efterfd{\o}lgende pr{\ae}cision i mere praktiske tiln{\ae}rmelser, kr{\ae}ver de eksisterende metoder, at den fulde efterfd{\o}lgende pr{\ae}cision skal instantieres.
+Laplace-approximationen (LA) er en s{\aa}dan metode, selv om den typisk er afh{\ae}ngig af grove tiln{\ae}rmelser af den efterf{\o}lgende pr{\ae}cisionsmatrix for at g{\o}re den beregningsm{\ae}ssigt gennemf{\o}rlig.
+N{\aa}r vi ikke faktoriserer den efterf{\o}lgende pr{\ae}cision i mere praktiske tiln{\ae}rmelser, kr{\ae}ver de eksisterende metoder, at den fulde efterf{\o}lgende pr{\ae}cision skal instantieres.
Vi kalder dette den fulde Laplace-approximation.
-Da den fulde pr{\ae}cisionsmatrix imidlertid for alle problemer undtagen legetd{\o}jsproblemer skalerer kvadratisk med antallet af parametre, er denne matrix for stor til at blive lagret i hukommelsen.
+Da den fulde pr{\ae}cisionsmatrix imidlertid for alle problemer undtagen leget{\o}jsproblemer skalerer kvadratisk med antallet af parametre, er denne matrix for stor til at blive lagret i hukommelsen.
I dette arbejde foresl{\aa}r vi en metode til beregning af Laplace-approksimationen ved hj{\ae}lp af Jacobian-vektorprodukter alene.
-Dette giver os mulighed for at udfd{\o}re marginal tr{\ae}ning og efterfd{\o}lgende stikprd{\o}veudtagning ved hj{\ae}lp af den fulde Laplace-approximation uden at lagre hele den efterfd{\o}lgende pr{\ae}cisionsmatrix.
+Dette giver os mulighed for at udf{\o}re marginal tr{\ae}ning og efterf{\o}lgende stikpr{\o}veudtagning ved hj{\ae}lp af den fulde Laplace-approximation uden at lagre hele den efterf{\o}lgende pr{\ae}cisionsmatrix.
For at opn{\aa} dette viser vi, at vi kan estimere den posteriore pr{\ae}cisionens log-determinant og den inverse kvadratrod ved hj{\ae}lp af kun jacobianiske vektorprodukter.
-For at overvinde de konditioneringsproblemer, der opst{\aa}r under sampling, tilbyder vi flere potentielle pr{\ae}konditioneringsv{\ae}rktd{\o}jer, som kan anvendes til at forbedre konvergensen af den inverse kvadratrodsapproximation.
-Dern{\ae}st giver vi en teknik til at evaluere kvaliteten af de efterfd{\o}lgende stikprd{\o}ver uden at instantiere eller invertere den efterfd{\o}lgende pr{\ae}cision baseret p{\aa} resultater fra traditionel statistik.
-Til at gennemfd{\o}re disse metoder anvender vi JAX, et bibliotek til automatisk differentiering, som gd{\o}r det muligt at udfd{\o}re jacobian-vektorprodukter effektivt uden eksplicit lagring af hele den efterfd{\o}lgende pr{\ae}cision.
+For at overvinde de konditioneringsproblemer, der opst{\aa}r under sampling, tilbyder vi flere potentielle pr{\ae}konditioneringsv{\ae}rkt{\o}jer, som kan anvendes til at forbedre konvergensen af den inverse kvadratrodsapproximation.
+Dern{\ae}st giver vi en teknik til at evaluere kvaliteten af de efterf{\o}lgende stikpr{\o}ver uden at instantiere eller invertere den efterf{\o}lgende pr{\ae}cision baseret p{\aa} resultater fra traditionel statistik.
+Til at gennemf{\o}re disse metoder anvender vi JAX, et bibliotek til automatisk differentiering, som g{\o}r det muligt at udf{\o}re jacobian-vektorprodukter effektivt uden eksplicit lagring af hele den efterf{\o}lgende pr{\ae}cision.
-Vi udfd{\o}rer den fulde Laplace-approximation p{\aa} b{\aa}de en sinusfunktion og MNIST ved hj{\ae}lp af vores metode.
-Denne best{\aa}r af to trin: (a) tr{\ae}ning af et neuralt netv{\ae}rk ved at maksimere enten dets posterior sandsynlighed eller marginale sandsynlighed og (b) prd{\o}veudtagning fra denne posterior.
+Vi udf{\o}rer den fulde Laplace-approximation p{\aa} b{\aa}de en sinusfunktion og MNIST ved hj{\ae}lp af vores metode.
+Denne best{\aa}r af to trin: (a) tr{\ae}ning af et neuralt netv{\ae}rk ved at maksimere enten dets posterior sandsynlighed eller marginale sandsynlighed og (b) pr{\o}veudtagning fra denne posterior.
Vores tiln{\ae}rmede maksimale marginale tr{\ae}ningsprocedure er i stand til at l{\ae}re et s{\ae}t parametre, som giver en ydeevne, der er sammenlignelig med den maksimale posterior- og maksimale sandsynlighedstr{\ae}ningsprocedure.
-Ved at analysere quantile--quantile-plots af vores posteriorprd{\o}ver og visualisere disse prd{\o}ver finder vi, at vores tiln{\ae}rmede prd{\o}veudtagningsmetode giver prd{\o}ver, der er korrekt fordelt.
-Vores benchmarks af Jacobian-vektorprodukternes ydeevne ansl{\aa}r, at vores metode giver en 10\ 000x hukommelsesreduktion i forhold til den konventionelle fulde Laplace-approximation, da den ikke lagrer den kvadratisk skalerende matrix med fuld pr{\ae}cision for den efterfd{\o}lgende periode.
-Dette er den fd{\o}rste metode til at udfd{\o}re Laplace-approximationen, hvor der kun er implicit adgang til den efterfd{\o}lgende pr{\ae}cision.
+Ved at analysere quantile--quantile-plots af vores posteriorpr{\o}ver og visualisere disse pr{\o}ver finder vi, at vores tiln{\ae}rmede pr{\o}veudtagningsmetode giver pr{\o}ver, der er korrekt fordelt.
+Vores benchmarks af Jacobian-vektorprodukternes ydeevne ansl{\aa}r, at vores metode giver en 10\,000x hukommelsesreduktion i forhold til den konventionelle fulde Laplace-approximation, da den ikke lagrer den kvadratisk skalerende matrix med fuld pr{\ae}cision for den efterf{\o}lgende periode.
+Dette er den f{\o}rste metode til at udf{\o}re Laplace-approximationen, hvor der kun er implicit adgang til den efterf{\o}lgende pr{\ae}cision.
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -22,7 +22,7 @@ However, we found that the sampling procedure is too slow to sample outside of a
Overall, while we have developed novel methods for training and inference in the full Laplace approximation, we have yet to demonstrate its improved performance over the diagonal Laplace approximation nor its practical application to large-scale problems.
-\section[Outlook and Future Developments]{Outlook and Future Developments}[Future Work]
+\section{Outlook and Future Developments}[Outlook and Future Work]
As mentioned in the previous section, we have not yet demonstrated the practicality of our method for full Laplace approximation, particularly for large-scale problems.
Hyperparameter selection typically requires performing a grid search over a range of hyperparameters, which is computationally expensive.
@@ -39,7 +39,7 @@ Furthermore, the full Laplace approximation should be compared to other methods
Furthermore, the feasibility of the full Laplace approximation must be demonstrated on a large-scale problem.
It is not clear how our method for the full Laplace approximation will scale with the number of parameters, data points, and model outputs.
-This regards not only the actual predictive performance and uncertainty estimation of the method, but also the post-hoc inference time for the approximate posterior.
+This concerns not only the actual predictive performance and uncertainty estimation of the method, but also the post-hoc inference time for the approximate posterior.
As such, important experiments must be performed to determine the scalability of the full 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 in the case of zero-shot learning.
These experiments would thus also require models with many parameters, such as a ResNet.
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -67,7 +67,7 @@ However, it would not be practical to compute this many posterior samples for la
\begin{figure}[b]
\centering
- \includegraphics{chisq_sin.pdf}
+ \includegraphics{chisq_sin_lower.pdf}
\caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for visualising normality of 1000 posterior samples from the full Laplace approximation for a fully connected neural network modelling \(f(x) = \sin(5 x + 1)\), using CIQ. Since the sample square Mahalanobis distances appear chi-squared distributed, the samples themselves are multivariate normal with the correct (known) covariance.}
\label{fig:chisq-sin}
\end{figure}
@@ -75,23 +75,31 @@ However, it would not be practical to compute this many posterior samples for la
\section{MNIST}
We now attempt to use our approximate sampling method on a more realistic problem, namely, the classification of handwritten digits from the MNIST dataset.
-To demonstrate that our approximate sampling method can be used on a more complex network, we train a convolutional neural network on the posterior distribution to compute the post-hoc Laplace approximation.\sidenote{Currently, the inclusion of convolutions in the network causes a major slowdown in the training procedure, so we maximise the posterior during training (i.e., post-hoc Laplace) instead of the marginal likelihood. We are working on a fix.}
+To demonstrate that our approximate sampling method can be used on a more complex network, we train a convolutional neural network on MNIST to compute the Laplace approximation.
+Currently, the inclusion of convolutions in the network causes a major slowdown in our approximate log-marginal training procedure, so we maximise the posterior during training (i.e., performing post-hoc Laplace) instead of the marginal likelihood.
+Because of this, these experiments have been performed with a fixed posterior precision \(\alpha\) which may not be the best choice for the posterior precision.
+
Unfortunately, evaluating the quality of the posterior predictive distribution for MNIST is not as straightforward as for the sine curve.
-The reason for this is that we cannot easily visually determine the appropriate magnitude of the uncertainty in the posterior predictive distribution for MNIST (since it is a 10-dimensional vector and the input images are \(28 \times 28\) pixels).
-We instead only evaluate the quality of the posterior distribution by comparing the posterior samples to the theoretical posterior distribution, as we did for the sine function in \cref{fig:chisq-sin}, where we sampled 1000 posterior samples.
-Since, in practice, we would not be able to store this many posterior samples for large-scale problems, we instead sample 200 parameter vectors from the posterior for MNIST, to demonstrate what these plots look like for a more realistic number of samples.\sidenote{In reality, for large-scale problems, we would not be able to store even 200 posterior samples, but between 5 and 20.}
+The reason for this is that we cannot easily visually determine the appropriate magnitude of the uncertainty in the posterior predictive distribution for MNIST (since the output is a 10-dimensional vector).
+We instead only evaluate the quality of the posterior distribution by comparing the posterior samples to the theoretical posterior distribution, as we did for the sine function in \cref{fig:chisq-sin}.
+For the sine function, we sampled 1000 posterior samples.
+Since, in practice, we would not be able to store this many posterior samples for large-scale problems, we instead sample 200 parameter vectors from the posterior for MNIST, to demonstrate what these plots look like for a more realistic number of samples.
+In reality, for large-scale problems, we would not be able to store even 200 posterior samples, but between 5 and 20.
We can see these results with \(\alpha = 0.1\) in \cref{fig:chisq-mnist}.
+\cref{fig:chisq-mnist} shows that sampling succeeds for values of \(\alpha\) as low as 0.1.
+This value of the prior precision is much lower than the value of \(\alpha\) we would likely use in practice.
+Additionally, as we noted in \cref{sec:spectrum-ggn}, since \(\alpha\) is the lowest eigenvalue of the posterior precision, the condition number is inversely proportional to \(\alpha\).
+Because of this, by sampling with a low value of \(\alpha\), we are sampling from the most ill-conditioned posterior precision we can expect to encounter.
+As such, this experiment allows us to conclude that our approximate sampling method would successfully be able to sample from the posterior distribution for a wide range of values of \(\alpha\) on MNIST without being too ill-conditioned for our method.
+
\begin{figure}
\centering
- \includegraphics{figures/ablation/chisq_55k_0.1_20q_100precondpoints_cg.pdf}
+ \includegraphics{figures/chisq_mnist_lower.pdf}
\caption[Chi-squared sample plot for MNIST.]{Chi-squared sample plot for visualising normality of 200 posterior samples from the full Laplace approximation for a convolutional neural network trained on MNIST with prior precision \(\alpha = 0.1\), using CIQ. Though the sample square Mahalanobis distances do not appear \emph{exactly} chi-squared distributed, this is because of the relatively low sample size used to compute the empirical histogram and quantiles. Thus, the plots themselves strongly suggest that the samples are multivariate normal with the correct (known) covariance.}
\label{fig:chisq-mnist}
\end{figure}
-\cref{fig:chisq-mnist} shows that sampling succeeds for values of \(\alpha\) as low as 0.1, a value which is much smaller (and more ill-conditioned) than the value of \(\alpha\) we would likely use in practice.
-As such, we can conclude that our approximate sampling method would be able to sample from the posterior distribution for a wide range of values of \(\alpha\) on MNIST.
-
\section{Discussion}
The choice of prior precision \(\alpha\) is non-trivial.
@@ -101,4 +109,3 @@ While we can ``choose'' \(\alpha\) to be a value large enough to reduce conditio
Choosing the prior in order to reduce ill-conditioning of our posterior precision means we sacrifice some of the benefits of the prior.
As such, we would like to learn \(\alpha\) from the data by maximising the marginal likelihood.
Current issues with training prevent us from reliably learning \(\alpha\) from the data across models of different sizes.
-
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -74,7 +74,7 @@ The first two terms are easy to compute (and differentiate), but computing the l
To overcome this, we thus need to estimate the log-determinant of the posterior precision matrix without instantiating it.
The second step of the Laplace approximation is to perform inference and make predictions from the posterior predictive distribution.
-We estimate the posterior predictive via Monte Carlo sampling by sampling from the posterior distribution and averaging the model predictions over the posterior samples.
+We estimate the posterior predictive via Monte Carlo by sampling from the posterior distribution and averaging the model predictions over the posterior samples.
Since the posterior distribution is approximated by a Gaussian, we sample from the posterior as
\begin{align*}
\bm \varepsilon = \bm \theta_{\textsc{map}} + \bm \Lambda^{-1/2} \bm \varepsilon_0,
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -7,7 +7,7 @@ In this chapter, we discuss training and inference in deep learning from a Bayes
We introduce the Laplace approximation as a method for approximating the posterior distribution of a model's parameters given the data.
We analyse the construction of the Laplace approximate posterior and discuss its properties and how it can be computed in practice.
-\section[Optimisation in Deep Learning]{Optimisation in Deep Learning}[Optimisation]
+\section{Optimisation in Deep Learning}[Optimisation]
\label{sec:optimisation-frequentist}
Suppose a neural network is a real-valued function \(f: \reals^N \times \reals^D \rightarrow \reals^O\) parametrised in \(\bm\theta \in \reals^D\) which maps an input \(\bm x \in \reals^N\) to an output \(f(\bm x, \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
@@ -40,7 +40,7 @@ The update rule is given by
%
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 non-increasing 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 \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t)\) can be computed separately for each individual training example \(\bm x_t, y_t\), as seen in \cref{eq:negative-log-likelihood}.
Using the chain rule on the composition of the loss into the loss as a function of the model output \(\mathcal{L}\) and the model output as a function of the parameters \(f\) such that \(\mathcal{L}(\bm \theta_t, \bm x_t, y_t) = \mathcal{L}(y_t, f(\bm x_t, \bm \theta_t))\),we compute the gradient as
%
\begin{align}\label{eq:sgd-gradient}
@@ -50,9 +50,9 @@ Using the chain rule on the composition of the loss into the loss as a function
where \(\nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)\) is the gradient of the neural network output with respect to its parameters at iteration \(t\) and \(\nabla_{\bm\theta} \log p(y_t \given \bm \theta)\) is the gradient of the loss function with respect to the model output.
Note that in this case, we perform one gradient update for each data point, which is why we iterate over the data points and update the parameters in the same step \(t\).
Thus, gradient descent only requires the gradient of the loss function with respect to the parameters of the neural network, which can be computed efficiently using automatic differentiation for a similar cost to a single forward pass through the network.
-However, gradient descent is prone to getting stuck in local minima, and can be slow to converge.
-In practice, we cannot compute the gradient of the loss function with respect to the parameters of the neural network for the entire dataset at once, since this would require storing the entire dataset in memory and computing the gradient over it.
-Instead, we compute the gradient of the loss function with respect to the parameters of the neural network for a subset of the data at each iteration (a batch), and then update the parameters based on this unbiased estimate of the gradient.
+However, gradient descent is prone to getting stuck in local minima and can be slow to converge.
+In practice, we cannot compute the gradient of the loss function for the entire dataset at once, since this would require storing the entire dataset in memory and computing the gradient over it.
+Instead, we compute the gradient of the loss for a subset of the data at each iteration (a batch), and then update the parameters based on this unbiased estimate of the gradient.
This is known as stochastic gradient descent (SGD), and is the most common form of gradient descent used in practice.
SGD helps avoid getting stuck in local minima, but, like gradient descent, does not take into account the curvature of the loss function, and so can overshoot the minimum.
@@ -62,7 +62,7 @@ SGD helps avoid getting stuck in local minima, but, like gradient descent, does
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.
+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
@@ -73,12 +73,12 @@ The update rule for Adam is thus given by
\end{align}
where \(\odot\) denotes the Hadamard product, \(\varepsilon\) is a small constant to prevent division by zero, and \(\bm m_t\) and \(\bm 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.
+The momentum term~\cite{polyak1964some} takes inspiration from 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 where the loss and model function are assumed to be linear.
This is a common approximation in practice, as this approximation is easy to compute.
-The Fisher information matrix is a discussed further in \cref{sec:ggn}.
+The Fisher information matrix is discussed further in \cref{sec:ggn}.
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.
@@ -98,7 +98,7 @@ The marginal likelihood is obtained by marginalising over the parameters, i.e.,
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 does not depend on the parameters, it does not affect the optimisation problem.}
\begin{align}\label{eq:unnormalised posterior}
- p(\bm \theta \given \bm y) \propto p(\bm y \given \bm \theta)\,p(\bm \theta) =: \tilde{p}(\bm \theta \given \bm y) .
+ p(\bm \theta \given \bm y) \propto p(\bm y \given \bm \theta)\,p(\bm \theta) \eqqcolon \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.
@@ -138,9 +138,9 @@ Others, like variational inference (VI)~\cite{jordan1999introduction,wainwright2
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.
+In this project, we focus on the Laplace approximation, in which the posterior is approximated by a Gaussian distribution, but where the variance is deduced rather than optimised.
-\section[The Laplace Approximation]{The Laplace Approximation}[Laplace Approximation]
+\section{The Laplace Approximation}[Laplace Approximation]
\label{sec:laplace}
In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI\@.
@@ -155,7 +155,7 @@ 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),
\end{align}
where \(\tilde{p}(\bm \theta \given \bm y)\) corresponds to the unnormalised posterior (\cref{eq:unnormalised posterior}).
-Normalising it by the normalisation constant \(p(\bm y)\) gives the normalised posterior\marginnote{This normalisation constant is computed as the marginal likelihood for the Laplace approximation in \cref{sec:maximising-the-evidence}.}
+Dividing 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}
\begin{split}
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)
@@ -164,22 +164,22 @@ Normalising it by the normalisation constant \(p(\bm y)\) gives the normalised p
% 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)
% \\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}). \nonumber
\end{align}
-The Laplace approximation can simply be trained as an MAP solution, where the Hessian is computed to obtain the posterior precision at inference time.
+The Laplace approximation can simply be trained to the 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.
-This can be done either by sampling from the neural network at each iteration of training and evaluating the Hessian at each sample~\cite{miani2022laplacian} or by estimating 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.
+This can be done either by sampling from the neural network at each iteration of training and evaluating the gradient at each sample~\cite{miani2022laplacian} or by estimating the log-marginal likelihood~\cite{daxberger2021laplace,immer2021scalable}, enabling training of hyperparameters and improved model selection.
+In this work, we approach the latter, as well as post-hoc Laplace.
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 as the loss function may not be convex with respect to the model parameters, which means that the normal distribution assumed by the Laplace approximation may not be valid.
-In the next sections, we will go into the computation of the Hessian for the negative log-posterior loss function from \cref{sec:bayesian-deep-learning} and show how the generalised Gauss-Newton approximation can be used to obtain a positive definite Hessian approximation, even when the actual Hessian itself is not positive definite.
+For BNNs, this is not guaranteed as the loss function is not convex with respect to the model parameters, which means that the normal distribution assumed by the Laplace approximation may not be valid.
+Furthermore, the Hessian is only ensured to be positive semi-definite at the MAP. This may not be realistic when training neural networks, which tend to have unreliable convergence~\cite{choromanska2015loss}.
+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 Hessian itself is not positive definite.
\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}\).
+The Hessian is a matrix of second-order partial derivatives of a scalar function.\sidenote{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\) parametrised 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})\).
@@ -189,8 +189,8 @@ Furthermore, if all the neural network's second partial derivatives are continuo
In this case, the Hessian is generally dominated by the block-diagonal~\cite{martens2015optimizing}.
However, the Hessian will only be positive definite if the loss function is a convex function in the model parameters \(\bm \theta\).
-Typically, the loss function is considered as the negative log-likelihood or negative log-posterior.
-In this framing, if considering, e.g., the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we get
+Typically, the loss function is defined as the negative log-likelihood or negative log-posterior.
+In this framing, if we consider the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we obtain
\begin{align}\label{eq:negative-log-posterior-hessian}
\begin{split}
\nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(\bm y_i \given \bm \theta) - \log p(\bm \theta) \right)
@@ -201,7 +201,7 @@ In this framing, if considering, e.g., the negative log-posterior loss from \cre
% \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
% \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
\end{align}
-In this construction, the first term is the Hessian of the negative log-likelihood, i.e., the Hessian of the loss function we are interested in.
+In this construction, the first term is the Hessian of the negative log-likelihood, e.g., the Hessian of the MSE loss.
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.
@@ -262,9 +262,9 @@ For a normal prior with precision \(\alpha\) on \(\bm \theta\), we obtain the GG
\end{align}
where we have defined \(\bm J_i = J_{\bm \theta} f(\bm x_i, \bm \theta)\) and \(\bm H_i = \nabla^2_f \mathcal{L}(\bm y_i, f(\bm x_i, \bm \theta))\).
-For mean square error (MSE) loss (e.g. for regression), we obtain the simple expression the GGN approximation of the Hessian
+For mean square error (MSE) loss (e.g. for regression), we obtain the simple expression for the GGN approximation of the Hessian
%
-\begin{align}%\label{eq:mse-hessian}
+\begin{align}\label{eq:mse-ggn}
\bm H_i \approx & {}\, (\bm y_i - f(\bm x_i, \bm \theta))^2 = 2 \identity,
\\ \bm \Lambda \approx&{}\, 2 \sum_{i=1}^N \bm J_i\T \bm J_i + \alpha \identity.
\end{align}
@@ -278,7 +278,7 @@ Here, the Fisher information matrix is given by
\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.
+The Fisher matrix (and its empirical equivalent) 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 for the GGN}[Practical Considerations]
\label{sec:practical-ggn}
@@ -307,7 +307,7 @@ The diagonal of this matrix is given by the element-wise square of the gradient
Notice that this is the second-order term used in the Adam optimiser, which we discussed in \cref{sec:adam}.
The choice of application of the generalised Gauss-Newton approximation is not based on having a reduced computational cost---both the Hessian and the GGN approximation require the same number of forward and backward passes.
-However, the GGN approximation is better behaved than the exact Hessian since it is guaranteed to be positive semidefinite for positive definite loss functions, thereby making it more suitable for practical applications.
+However, the GGN approximation is better behaved than the exact Hessian since it is guaranteed to be positive semidefinite for convex losses, thereby making it more suitable for practical applications.
This is because the GGN approximation is designed to be positive semi-definite, while the exact Hessian can be vulnerable to negative curvature.
This feature can limit the use of the exact Hessian for algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
@@ -319,7 +319,7 @@ This then means we want to compute
\\ ={}& \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i \bm v + \alpha \bm v,
\end{split}
\end{align}
-such that we actually compute the product of the Jacobian of the neural network (as defined in \cref{eq:negative-log-likelihood-ggn}) by a vector \(\bm v\).
+such that we actually compute the product of the Jacobian of the neural network (as defined in \cref{eq:negative-log-likelihood-ggn}) by a vector \(\bm v\), i.e., \(\bm J_i \bm v\).
\subsection{The Spectrum of the GGN Approximation}[The Spectrum of the GGN]
\label{sec:spectrum-ggn}
@@ -330,7 +330,7 @@ The eigenvalues of this precision will thus be the sum of the eigenvalues of \(\
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\).
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.
+This means that 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.
@@ -340,8 +340,8 @@ Since \(\bm J_i\) is a rectangular matrix of shape \(O \times D\), the rank of e
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.\sidenote{By ``same ordering'' we mean that they are eigenvectors with respect to the same relative eigenvalue.}
-Practically speaking, 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\).
+Thus, non-rigorously, we have that the closer the eigenvectors (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.\sidenote{By ``same ordering'' we mean that they are eigenvectors with respect to the same relative eigenvalue.}
+Practically speaking, this means that the rank of the sum of the GGN matrices can increase as we add more data points (but 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.
@@ -356,6 +356,7 @@ If the \(O\) non-zero eigenvectors are identical across the GGN matrices, then t
}
However, if these eigenvectors are not identical, then as the largest eigenvalue of the sum of the GGN increases and the number of non-zero eigenvalues of the sum of the GGN increases, the non-zero spectrum of the sum of the GGN matrices will get steeper.
For these reasons, as we add more data points, the posterior precision will become more ill-conditioned.\sidenote{A matrix is considered ill-conditioned when the ratio of its largest eigenvalue to its lowest eigenvalue, its \emph{condition number} (\(\kappa \coloneqq \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}\)), is very large.}
+If the sum of the GGN matrices is indeed low-rank, then the prior precision \(\alpha\) will be the smallest eigenvalue of the posterior precision.
Empirically, we find that, for MNIST, the effective rank of the posterior precision is significantly less than \(N \cdot O\) (see \cref{fig:spectrum-mnist}).
This suggests that the eigenvectors of the posterior precision are relatively similar across different observations, though not identical.
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -59,10 +59,10 @@ There are then essentially three steps to CIQ, namely:
\Input{\(\bm K \succ 0\), \(\bm b\), \(\bm P \succ 0\), \(J > 0\), \(Q > 0\)}
\Output{\(\bm s = \bm K^{-1/2} b\)}
\(\_, \bm\alpha, \bm\beta \gets \mathrm{CG}(\bm K, \bm b, \bm P, J)\)\;
- \(M_{\mathrm{Lanczos}} \gets \mathrm{TriDiag}(\bm\alpha, \bm\beta)\)\;
- \(\lambda_M \gets \lambda(M_{\mathrm{Lanczos}})\)\;
- \(\lambda_{\min} \gets \min{\lambda_M}\)\;
- \(\lambda_{\max} \gets \max{\lambda_M}\)\;
+ \(\bm M_{\mathrm{Lanczos}} \gets \mathrm{TriDiag}(\bm\alpha, \bm\beta)\)\;
+ \(\lambda_{\bm M} \gets \lambda(\bm M_{\mathrm{Lanczos}})\)\;
+ \(\lambda_{\min} \gets \min{\lambda_{\bm M}}\)\;
+ \(\lambda_{\max} \gets \max{\lambda_{\bm M}}\)\;
\(w_1, \ldots, w_Q, t_1, \ldots, t_Q \gets \mathrm{Quadrature}(\lambda_{\min}, \lambda_{\max}, Q)\)\;
\(\bm c_1, \ldots, \bm c_Q \gets \mathrm{msMINRES}(\bm K, \bm b, t_q, \bm P, J)\) \Comment*[r]{\(\left( t_q \identity + \bm K \right)^{-1} \bm b\).}
% \(s_q \gets \left( t_q \identity + H \right)^{-1} b\) \Comment*[r]{Computed via msMINRES(K, t, J).}
@@ -78,6 +78,7 @@ it is then also trivial to compute the square root product by multiplying \(\bm
\begin{align}
\bm K^{1/2} \bm \varepsilon_0 = \bm K \cdot \underbrace{\bm K^{-1/2} \bm \varepsilon_0}_{\text{\cref{eq:ciq-asdf}}}.
\end{align}
+\marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
\def\ymax{2.0}
\def\ymin{-2.0}
@@ -96,14 +97,13 @@ it is then also trivial to compute the square root product by multiplying \(\bm
% ------------------------------------------
The first step of CIQ invokes Cauchy's integral formula.
-Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as\marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
+Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as
\begin{align}\label{eq:cauchy-integral-formula}
f(\bm K) ={} & \frac{1}{2 \pi i} \oint_\Gamma f(\tau) \left(\tau \identity - \bm K\right)^{-1}\,d\tau,
\end{align}
where \(\oint_\Gamma\) is the contour integral along the contour \(\Gamma\) parametrised in the complex variable \(\tau\).
Since, in our case, the eigenvalues of \(\bm K\) will all be real-valued, a circle centred on the real axis which encloses the minimum and maximum eigenvalues of \(\bm K\) will suffice.\marginnote{If a real-valued matrix is symmetric, its eigenvalues are real. In our case, \(\bm K\) will be the precision matrix obtained from the Laplace approximation.}
-To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).
-\sidenote{Krylov subspace methods are a class of iterative methods which estimate certain properties (e.g., the rank, eigenvalues, span) of some matrix \(\bm K\) by using a Krylov subspace for some vector \(\bm v\). The Krylov subspace is a subspace spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), given by \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\left\{ \bm v, \bm K \bm v, \bm K^2 \bm v, \ldots, \bm K^{j-1} \bm v \right\}\).}
+To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).\sidenote{Krylov subspace methods are a class of iterative methods which estimate certain properties (e.g., the rank, eigenvalues) of a matrix \(\bm K\) with a Krylov subspace \(\mathcal{K}_j(\bm K, \bm v)\) for some vector \(\bm v\). This subspace is spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\{ \bm v, \bm K \bm v, \ldots, \bm K^{j-1} \bm v \}\).}
% ------------------------------------------
% QUADRATURE
@@ -121,7 +121,7 @@ To estimate these values, we can use the Lanczos algorithm, a Krylov subspace me
\label{fig:ciq}
\end{marginfigure}
-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
+We then apply Cauchy's integral formula to \(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}
\begin{split}
\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}
@@ -138,11 +138,11 @@ 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\)~\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 the precision matrix (which will be computed from the Laplace approximation as the GGN matrix) can often have a 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\).
-To then account for this change of variable, we reformulate \cref{eq:ciq-integral-tau} as % eq:ciq-quadrature-tau
+To account for this change of variable, we reformulate \cref{eq:ciq-integral-tau} as % eq:ciq-quadrature-tau
\begin{align}\label{eq:ciq-integral-sigma}
\begin{split}
\bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau
@@ -150,7 +150,7 @@ To then account for this change of variable, we reformulate \cref{eq:ciq-integra
\\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1}.
\end{split}
\end{align}
-We can now adapt the equation to allow for computation using only matrix-vector products.
+We can now adapt this equation to allow for computation using only matrix-vector products.
From here, the inverse square root vector product is given by
\begin{align}\label{eq:ciq-sigma-product}
\bm K^{-1/2} \bm v \approx{} & \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \underbrace{\left( \sigma^2_q \identity - \bm K \right)^{-1} \bm v}_{Q\text{ system solves}}.
@@ -188,12 +188,12 @@ This integral is approximated by
\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)\)
-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\).
+Additionally, we need to compute the Jacobian elliptic functions \(\sn(u \given m)\), \(\cn(u \given m)\), and \(\dn(u \given m)\).
+To compute these values, we first calculate the elliptic modulus \(k = \kappa(\bm K)^{-1/2}\) (where \(\kappa(\bm K)\) denotes the condition number of \(\bm K\)) and \(u_q = (q - 1/2) / Q\).
These functions are not implemented in JAX, and SciPy uses the Cephes numerical library in C to compute them.
Because of this, we ported the Cephes Jacobi elliptical functions to JAX, and \cref{fig:ellipk,fig:ellipj} show the functions evaluated for the two implementations.
Our static JAX implementation allows for efficient \emph{just-in-time} compilation, which is not possible with the SciPy implementation.
-We can thereby now compute the quadrature weights and shifts, as seen in \cref{alg:quad}.
+We can thereby now compute the quadrature weights and shifts as described in \cref{alg:quad}.
\begin{figure}
\centering
@@ -227,35 +227,36 @@ We can thereby now compute the quadrature weights and shifts, as seen in \cref{a
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.
+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 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 we follow their implementation.
+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, 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.
+To mitigate the effect of ill conditioning, we use a preconditioner \(\bm P\) such that \(\bm P^{-1} \bm K\) is well-conditioned.
\section{Preconditioning}
\label{sec:preconditioning}
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
+In the case of the Laplace approximation, the precision matrix is given by
\begin{align}\label{eq:ggn-laplace}
\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_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}).
+% Since each \(J_i\T H_i J_i\) is, for \(i \in 1 \ldots N\), a low-rank positive semi-definite matrix of rank \(O\), 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.
-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
+For preconditioning, we need to find a preconditioner matrix \(\bm P\) and its inverse \(\bm P^{-1}\) such that the matrix \(\bm P^{-1} \bm \Lambda\) is well-conditioned.
+This is done by choosing \(\bm 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 \(\bm 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),
@@ -268,7 +269,7 @@ The eigenvalues of our GGN approximation are thus given by
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 its computational cost.
+In this case, there is a tradeoff 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 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.
@@ -277,22 +278,22 @@ The first point lends itself toward the use of preconditioners with a similar co
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.
+Let us consider the toy example of a preconditioner which we can guarantee to have the same eigenvectors as \(\bm \Lambda\) and whose eigenvalues we can freely choose.
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.
+Thus, if we have a rank-\(K\) preconditioner, then we can eliminate at most the \(K\) challenging (largest or smallest) eigenvalues of \(\bm \Lambda\) by using a rank-\(K\) preconditioner.
The steeper the spectrum of \(\bm \Lambda\), the more effective this preconditioning will be (however, in this case, \(\bm \Lambda\) will probably also tend to be worse conditioned).
% This analogy can be extended in a hand-wavy way to an imperfect preconditioner which does not exactly approximate the eigenvalue
-For a preconditioner with a condition number of \(\kappa^\prime\), we can at best obtain a preconditioned system with conditioning \(\kappa / \kappa^\prime\).
+For a preconditioner with a condition number of \(\kappa^\prime\), we can at best obtain a preconditioned system with conditioning \(\kappa / \kappa^\prime\), where \(\kappa\) is the condition number of \(\bm \Lambda\).
This is because we want the highest eigenvalue of the preconditioner to neutralise (i.e., align with) the highest eigenvalue of \(\bm \Lambda\), and the same for the lowest eigenvalues.
-In CIQ, we will, in each iteration, perform one matrix-vector product \(\bm \Lambda \bm v\) and one preconditioner-vector product, as seen in \cref{alg:ciq}.%\sidenote{Because of this, to avoid the preconditioner becoming a performance bottleneck, the speed of the preconditioner-vector products should be significantly greater than that of the matrix-vector products.}
+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.
+Additionally, an effective preconditioner can significantly improve the convergence speed of CIQ, and this effect can 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 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}
@@ -306,10 +307,10 @@ If we calculate the GGN approximation of the Hessian by linearizing over the who
\\ \beta ={}& \left(\alpha \norm{\nabla}^2 + \norm{\nabla}^4\right)^{-1} - \left(\alpha \norm{\nabla}^2\right)^{-1}, \label{eq:adam-preconditioner-beta}
\end{align}
where \(\nabla\) is the gradient of the loss with regard to the model parameters.
-Since \cref{eq:adam-preconditioner} is in the form of a diagonal perturbation to an outer product, we have the closed-form inverse given by \cref{eq:adam-preconditioner-inverse,eq:adam-preconditioner-beta}.
-The problem with using this approximation is that the rank of the outer product is one, where this term determines the largest eigenvalue of the preconditioner.
-This means that, best-case scenario, we can reduce the highest eigenvalue of the Hessian.
-This occurs when the gradient corresponds to the direction of the highest precision (empirically, this may not be the case).
+Since \cref{eq:adam-preconditioner} is in the form of a diagonal perturbation to a vector outer product, we have the closed-form inverse given by \cref{eq:adam-preconditioner-inverse,eq:adam-preconditioner-beta}.
+The problem with using this approximation is that the rank of the outer product is one and this term determines the largest eigenvalue of the preconditioner.
+This means that, best-case scenario, we can only reduce the highest eigenvalue of the Hessian.
+This occurs when the gradient corresponds to the direction of the highest precision (this may not be the case).
This would then leave the second-highest eigenvalue unaffected.
% If the rank of our precision matrix is greater than one, then this value seems likely to be similarly large to originally the largest eigenvalue.
% Therefore, the condition number of our problem will not be significantly reduced.
@@ -321,9 +322,9 @@ We therefore need to find a higher-rank approximation of the posterior precision
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.
+Since the lowest eigenvalue of the posterior precision is the prior precision \(\alpha\), we can (even with few preconditioner observations) 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\).
+Alternatively, we can estimate the eigenvectors and eigenvalues of \(\textstyle{N / B \sum^B_i} J_i\T J_i\), where \(B\) is the number of preconditioner observations.
This would \emph{hopefully} have an eigenbasis that is close to that of \(\textstyle{\sum^N_i} J_i\T J_i\) without needing to calculate the precision matrix over the whole dataset (we may need to guarantee all classes are represented in \(x_i, i \sim \mathcal{B}\)).\sidenote{Empirically, it seems likely that the eigenvectors of the posterior precision are relatively similar across different observations---see \cref{fig:spectrum-mnist} and \cref{sec:spectrum-ggn}.}
However, scaling the precondititioner up to the size of the entire dataset will lead to a preconditioner which is ill-conditioned, making inversion via conjugate gradient difficult.
@@ -347,15 +348,15 @@ The identity is given by
\begin{align}
\left( \bm A + \bm B \bm C \bm D \right)^{-1} = \bm A^{-1} - \bm A^{-1} \bm B \left( \bm C^{-1} + \bm D \bm A^{-1} \bm B \right)^{-1} \bm D \bm A^{-1},
\end{align}
-where \(\bm A\) is a \(n \times n\) matrix and can easily be inverted, \(\bm B\) is a \(n \times m\) matrix, \(\bm C\) is a \(m \times m\) matrix, and \(\bm D\) is a \(m \times n\) matrix.
+where \(\bm A\) is a \(D \times D\) matrix and can easily be inverted, \(\bm B\) is a \(D \times O\) matrix, \(\bm C\) is a \(O \times O\) matrix, and \(\bm D\) is a \(O \times D\) matrix.
This identity can be used to compute the inverse of the GGN matrix with the prior precision as
\begin{align}
\left( \alpha \identity + \bm J\T \bm H \bm J \right)^{-1} = \alpha^{-1} \identity - \alpha^{-2} \bm J\T \left( \bm H^{-1} + \alpha^{-1} \bm J \bm J\T \right)^{-1} \bm J,
\end{align}
-where \(\alpha\) is the prior precision and \(\bm H\) and \(\bm J\) are the Hessian and the Jacobian as per \cref{eq:negative-log-likelihood-ggn} evaluated at a single observation.
+where \(\alpha\) is the prior precision and \(\bm H \in \reals^{O \times O}\) and \(\bm J \in \reals^{O \times D}\) are the Hessian and Jacobian from \cref{eq:negative-log-likelihood-ggn} evaluated at a single observation.
-Here, to exploit the low-rank structure of the preconditioner, we only compute the GGN matrix over a single observation.
-As such, for this preconditioner to be effect, we assume that the eigenvectors of the GGN matrix are constant across the dataset.
+Here, to exploit the low-rank structure of the preconditioner, we only compute the GGN matrix over one observation.
+For this preconditioner to be effective, we assume that the eigenvectors of the GGN matrix are approximately 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.
@@ -363,24 +364,24 @@ To do so, we can scale the outer product \(\bm J \bm J\T\) term by some factor,
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.
+Before the Woodbury matrix identity can be used, the Hessian of the loss with respect to the model outputs \(\bm H\) must first be inverted.
+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, depending on the number of outputs.
+For regression tasks where the Hessian is diagonal (\cref{eq:mse-ggn}), this is not a problem.
For classification tasks using the cross-entropy loss, we can compute the inverse of the Hessian as a diagonal and outer product matrix, so this is also not a problem.
For models with a very large number of outputs, such as autoencoders, the matrix \(\bm H\) will be fairly large, and so computing its inverse manually can become expensive.
-However, in the case of autoencoders, the loss can be the MSE loss, which has a diagonal Hessian, and so can be inverted easily.
+However, in the case of autoencoders, the loss is usually the MSE loss, and so has a diagonal Hessian which 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 did not end up using this preconditioner.
+Thus, if \(\bm H\) is positive definite (which is necessary for the Laplace approximation to be valid) 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 when the eigenvectors of the GGN matrix are not constant across the dataset.
+However, due to the effectiveness of the sub-sampled GGN preconditioner and the assumption of stationary eigenvectors, we did not end up using this preconditioner.
Additionally, this preconditioner cannot assume a higher rank than the number of outputs \(O\), since adding more observations to the preconditioner will increase its rank to some unknown value and prevent the application of the Woodbury matrix identity.
-Ideally, we would like to be able to find a preconditioner for which we can choose the rank of the low-rank approximation \(\bm A \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) without simply computing the GGN on a single observation (\(\bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) for some observation \(i\)) as in \cref{sec:woodbury-preconditioner}, while still allowing us to apply the Woodbury matrix identity, unlike in the case of the sub-sampling preconditioner in \cref{sec:sub-sampling-preconditioner} (\(\sum^B_{i=1} \bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\)).
+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\)), 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}
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.
+This decomposition has been successfully applied to preconditioning for kernel matrices in Gaussian processes~\cite{gardner2018gpytorch}, but not yet 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}
\begin{split}
@@ -394,7 +395,7 @@ We can then use this approximation as a preconditioner by computing its inverse
\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}.
+We then obtain the preconditioner-vector and inverse preconditioner-vector products as in Equation~\ref{eq:pivoted-cholesky-preconditioner}.
The most expensive 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,
@@ -405,7 +406,7 @@ This would then require storing the factorisation \(\bm L \in \reals^{K \times D
As explained in \cref{sec:spectrum-ggn}, as we sum over multiple low-rank GGN matrices, we will likely find that the minimum non-zero eigenvalue gets progressively smaller compared to the largest.
The condition number will therefore get progressively worse.
-The preconditioning of \(\sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) (which should be low-rank) will progressively become both \emph{less} low-rank and its condition number will progressively become higher.% MERGE WITH:
+The preconditioning of \(\sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) (which should be low-rank) will progressively become \emph{less} low-rank and its condition number will progressively become higher.
Regarding the non-zero eigenvalues and eigenvectors, we consider a best-case scenario.
Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors of the precision matrix.
@@ -418,7 +419,7 @@ This will then lead to a steeper spectrum and a larger condition number of the t
Thus, when the batch size \(B\) is large, the condition number of the sum of GGN Hessians will increase and the preconditioner will become less effective.
How much the condition number increases with the batch size depends on the eigenspectrum of the per-observation Hessians (i.e., how close the eigenvectors are to each other).
-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.
+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~\cite{harbrecht2012low}.
However, it is not possible to generally and efficiently determine the diagonal of a matrix via matrix-vector products.
@@ -428,35 +429,42 @@ Alternatively, we can approximate the diagonal of the GGN matrix as the element-
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.
+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 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.
Furthermore, the pivoted Cholesky decomposition is not implemented in JAX.
Since the pivoted Cholesky decomposition involves incrementally adding columns to the factorisation, and JAX is not friendly towards mutable variables, this preconditioner is not straightforward to implement.
-Additionally, the pivoted Cholesky decomposition traditionally uses the diagonal of the matrix to determine which columns to add to the factorisation, which is difficult without instantiating the matrix,we would need to approximate the diagonal of the GGN matrix.
+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, so we would need to approximate the diagonal of the GGN matrix.
\subsection{Other Preconditioners}
+The randomly pivoted Cholesky~\cite{chen2022randomly} preconditioner is simple and inexpensive variant of the partial Cholesky decomposition family of algorithms.
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.
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.
+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 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 used.
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.
+We provide a high-level comparison of the various preconditioners we have presented in this chapter in \cref{tab:preconditioner-comparison}.
+
+\begin{table*}
+ \caption[Comparison of various preconditioners.]{Comparison of various approximations of the likelihood contribution \(\sum_i^N \bm J_i\T \bm H_i \bm J_i\) to the Laplace posterior precision for use in preconditioning. We compare the rank of the approximation, the speed of computing preconditioner-vector products, the closeness of the approximation, and the method used for inverting the preconditioner.}
+ \label{tab:preconditioner-comparison}
+ \include{tables/preconditioner-comparison.tex}
+\end{table*}
\section{Sampling Evaluation}
\label{sec:sampling-evaluation}
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.
+As such, we need a technique for evaluating 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 thus whether the samples were generated successfully.
@@ -466,7 +474,7 @@ Instead, we require a method to evaluate the quality of the samples without acce
\subsection{The Chi-Squared Distribution}[Chi-Squared Distribution]
-As per the definition of the chi-squared distribution, for \(D\)-dimensional standard normal samples \(\bm \varepsilon_0\), the sum of square deviations from the mean is distributed as a chi-squared distribution with \(D\) degrees of freedom, as per
+As per the definition of the chi-squared distribution, for \(D\)-dimensional standard normal samples \(\bm \varepsilon_0\), the sum of square deviations from the mean is chi-squared distributed with \(D\) degrees of freedom, as per
\begin{align}
{\bm \varepsilon_0}\T \bm \varepsilon_0 \sim \chi^2(D).
\end{align}
@@ -479,7 +487,7 @@ Samples \(\bm \varepsilon\) which are normally distributed but do not have a mea
\end{split}
\end{align}
%
-This has the advantage of only requiring computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it does not require instantiation of the matrix.
+This has the advantage of only requiring the computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it does not require instantiation of the matrix.
For the Laplace approximation, this can easily be computed as $\bm \Lambda = - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y)$.
Since this precision matrix can be computed exactly, this calculation can be performed to evaluate the whether a set of samples \(\bm \varepsilon\) are drawn from a distribution \(\normal(\bm \mu, \bm \Lambda^{-1})\) while only implicitly accessing this precision matrix in the evaluation via matrix-vector products.
To do so, we can compute the squared Mahalanobis distance of each sample \(\bm \varepsilon_i\) from the mean \(\bm \mu\) and compare it to the chi-squared distribution with \(D\) degrees of freedom.
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -4,7 +4,7 @@
In this chapter, we will discuss training neural networks by maximising the evidence under the Laplace approximation.
We first motivate maximising the evidence as opposed to maximising the posterior and discuss how the Laplace approximation guarantees the feasibility of the evidence, as well as the issues that arise when computing the marginal likelihood with the full posterior precision.
-Next, we propose a novel method for computing an upper bound on the log-determinant of the full posterior precision matrix without the issues that typically arise when using the full posterior precision.
+Next, we propose a novel method for computing an upper bound on the log-determinant of the full posterior precision matrix which avoids the aforementioned issues.
Finally, we discuss the implementation of online Laplace with mini-batching using our upper bound.
\section{Maximising the Evidence}
@@ -29,12 +29,12 @@ This result shows that the marginal likelihood is a good proxy for the true cros
Because of this, it can be used for model selection, since it will choose the model that performs best on average across all permutations of the data.
Additionally, unlike the posterior, the marginal likelihood can be used to optimise both model parameters and hyperparameters, since we can meaningfully minimise the log-marginal with respect to the hyperparameters (and compute the respective gradients).
-This is useful, since it allows us to use gradient-based optimisation methods to find the hyperparameters that maximise the marginal likelihood, as has been commonly done in Gaussian Processes~\cite{williams2006gaussian,damianou2013deep}.
+This 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.
The Laplace approximation provides an approximation to the posterior which allows us to compute and maximise the marginal likelihood.
-This is known as \emph{online} Laplace, and is an alternative to the \emph{post-hoc Laplace} approach, which is the standard approach to training under the Laplace approximation~\cite{daxberger2021laplace}.
-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}}\).
+This is known as \emph{online} Laplace, and is an alternative to the \emph{post-hoc} Laplace approach, which is the standard approach to training under the Laplace approximation~\cite{daxberger2021laplace}.
+Under this 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}}\).
As in \cref{eq:laplace-taylor}, we have the approximate Laplace log-posterior
\begin{align}\label{eq:laplace-posterior}
\begin{split}
@@ -43,6 +43,7 @@ As in \cref{eq:laplace-taylor}, we have the approximate Laplace log-posterior
\\={}& \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}}).
\end{split}
\end{align}
+\marginnote{In reality, we will not be at the MAP solution when optimising the marginal, neither during training nor after convergence (the marginal and the posterior have different minima). In this case, we cannot neglect the first derivative term in \cref{eq:laplace-posterior}. However, when this term is not zero, the Laplace posterior is \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}^*, \bm \Lambda^{-1})\), where \(\bm \theta_{\textsc{map}}^* \coloneqq \bm \theta_{\textsc{map}} + \bm \Lambda^{-1} \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y)\). This corresponds to shifting the parameters to the MAP with a Newton step \(\bm \theta_{\textsc{map}}^* \coloneqq \bm \theta_{\textsc{map}} - \nabla^2_{\bm \theta} \mathcal{L}(\bm \theta_{\textsc{map}})^{-1} \nabla_{\bm \theta} \mathcal{L}(\bm \theta_{\textsc{map}})\), since Newton's method finds the minimum in one step for quadratic functions (like the second order Taylor expansion to our loss).}
The approximate Laplace marginal is then computed as
\begin{align}\label{eq:laplace-marginal}
\begin{split}
@@ -51,7 +52,7 @@ The approximate Laplace marginal is then computed as
\\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}.
\end{split}
\end{align}
-Taking the log of this approximate Laplace marginal, we obtain
+Taking the log of this approximate marginal, we obtain
\begin{align}\label{eq:laplace-log-marginal}
\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
@@ -78,11 +79,7 @@ Combining \cref{eq:laplace-log-marginal,eq:normal-log-likelihood,eq:normal-log-p
\end{split}
\end{align}
where it can be seen that there are constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\).
-While these terms are not necessary for the optimisation problem, we choose to include them.
-% In particular, the terms \(D /\,2 \log(2 \pi)\) are important for comparing models with different numbers of parameters.
-% \sidenote{We note that these constant terms are not included in the computation of the \texttt{log\_marginal\_likelihood} function from the \texttt{laplace} Python library~\cite{daxberger2021laplace}. In spite of this, \textcite{immer2021scalable} use this library in their experiments, where they compare models with different numbers of parameters.}
-% Furthermore, the other terms are important for interpreting the log-marginal as a likelihood.
-These terms allow us to interpret the log-marginal as a likelihood and do not significantly affect performance.
+While these terms are not necessary for the optimisation problem, we choose to include them, as these terms allow us to interpret the log-marginal as a likelihood and do not significantly affect performance.
From \cref{eq:expanded-log-marginal}, we can see that the Laplace log-marginal likelihood requires computation of the log-determinant of the posterior precision matrix.
This log-determinant term gets minimised in our optimisation, which favours models with small eigenvalues in the Hessian of the loss, i.e., points of low curvature in the loss landscape, as we argued earlier in this section.
This is a challenging computation, as there is no \emph{cheap} closed-form expression for the log-determinant of a matrix without instantiating the matrix.
@@ -115,12 +112,11 @@ We want to compute an upper bound on the log-determinant of the posterior preci
\begin{pmatrix} \mu_1 \\ \mu_2 \end{pmatrix},
\end{align}
with \(\mu_1 = \Tr(\bm \Lambda)\), \(\mu_2 = \norm{\bm \Lambda}^2_F = \Tr\left(\bm \Lambda^2\right)\), \(t = \frac{\beta \mu_1 - \mu_2}{\beta D - \mu_1}\), where \(\beta\) is an upper bound on the largest eigenvalue of \(\bm \Lambda\).
-This bound involves approximating the log eigenvalues of \(\bm \Lambda\) with a third-order polynomial.
-Since \(B_D(\mu_1, \mu_2, \beta)\) is an upper bound on the log-determinant of \(\bm \Lambda\), we can use this to compute a lower bound on the log-marginal likelihood, as per \cref{eq:expanded-log-marginal}.
-This is effective because we want to maximise the log-marginal.
+This bound amounts to 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 are computing a lower bound on the log-marginal likelihood (which we maximise), as per \cref{eq:expanded-log-marginal}.
Assume \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) has rank \(K \leq O \cdot N\).
-Then, \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues greater than \(\alpha\), and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
+Then \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues greater than \(\alpha\) and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
We then decompose the log-determinant as
\begin{align}\label{eq:log-determinant-decomposition}
\begin{split}
@@ -130,8 +126,8 @@ We then decompose the log-determinant as
\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
+The former can then be bounded using \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 in \cref{eq:laplace-trace-bound-2} via
\begin{align}\label{eq:laplace-trace-bound}
\Tr{\bm \Lambda} ={} & \sum_{k=1}^{K} (\lambda_k + \alpha) + (D - K) \alpha
\\ \begin{split}
@@ -143,14 +139,14 @@ 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-2,eq:laplace-square-trace-bound} as
+We can thus bound \(\log \det \bm \Lambda\) with \cref{eq:log-determinant-decomposition,eq:laplace-trace-bound-2,eq:laplace-square-trace-bound} as
\begin{align}\label{eq:laplace-log-determinant-bound}
\log \det \bm \Lambda \leq{} & B_K\left(\mu_1, \mu_2, \beta\right) + (D - K) \log \alpha,
- \\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,
- \\ \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.
+ \\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,\label{eq:mu1}
+ \\ \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.\label{eq:mu2}
\end{align}
-For the upper bound on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\) as an upper bound.\sidenote{Since the matrix is positive semidefinite, all eigenvalues are non-negative and thus the trace is bigger than the largest eigenvalue.}
-From \cref{eq:laplace-log-determinant-bound}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log-determinant of the posterior precision.
+For the upper bound \(\beta\) on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\) as an upper bound.\sidenote{Since the matrix is positive semidefinite, all eigenvalues are non-negative and thus the trace is bigger than the largest eigenvalue.}
+From \cref{eq:mu1,eq:mu2}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log-determinant of the posterior precision.
% WHERE DOES THE BOUND COME FROM? WHICH UNCERTAINTY? THE UPPER BOUND ON THE EIGENVALUE?
\subsection{Hutchinson's Trace Estimator}[The Trace Estimator]
@@ -194,10 +190,9 @@ The log-determinant of this mini-batched precision is equivalent to
\\ ={} & 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
+To correctly compute the log-determinant of the posterior precision, we scale the trace and square traces for \cref{eq:mu1,eq:mu2} 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 Jacobian-vector products without the need to instantiate the posterior precision.
+By combining \cref{eq:laplace-log-determinant-bound,eq:bai-golub-bound,eq:hutchinson-trace,eq:hutchinson-square-trace,eq:trace-rearrangement,eq:laplace-trace-scaling,eq:laplace-square-trace-scaling}, we can therefore compute a differentiable upper bound on the log-marginal likelihood under the Laplace approximation using only Jacobian-vector products without needing to instantiate the posterior precision.
diff --git a/figures/chisq_mnist_lower.pdf b/figures/chisq_mnist_lower.pdf
Binary files differ.
diff --git a/figures/chisq_mnist_upper.pdf b/figures/chisq_mnist_upper.pdf
Binary files differ.
diff --git a/figures/chisq_sin_lower.pdf b/figures/chisq_sin_lower.pdf
Binary files differ.
diff --git a/figures/chisq_sin_upper.pdf b/figures/chisq_sin_upper.pdf
Binary files differ.
diff --git a/figures/overview.tex b/figures/overview.tex
@@ -28,7 +28,7 @@
\node (map) [below=of training,yshift=0.7cm] {\(\bm \theta_{\textsc{map}}\)};
\draw[dashed] (training) to (map);
\node (sample) [below=of map,yshift=0.5cm] {Sample};
- \node (samplingeq) [below=0cm and 0cm of sample,box,fill=color2,minimum width=3.8cm,minimum height=2cm] {\(\begin{aligned}\bm \varepsilon \sim{}& \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}) \\ ={}& \bm \theta_{\textsc{map}} + \text{\colorbox{kaogreen}{$\bm \Lambda^{-1/2} \bm \varepsilon_0$}}\end{aligned}\)};
+ \node (samplingeq) [below=0cm and 0cm of sample,box,fill=color2,minimum width=4.17cm,minimum height=2.1cm] {\(\begin{aligned}\bm \varepsilon \sim{}& \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}) \\ ={}& \bm \theta_{\textsc{map}} + \text{\colorbox{kaogreen}{$\bm \Lambda^{-1/2} \bm \varepsilon_0$}}\end{aligned}\)};
\begin{scope}[on background layer]
\node (inference) [box,fill=color1,inner sep=0.25cm,fit={(sample)(samplingeq)}] {}; % minimum width=8.05cm,minimum height=3.2cm,
\end{scope}
diff --git a/main.bib b/main.bib
@@ -39,6 +39,15 @@
publisher = {Elsevier}
}
+@inproceedings{choromanska2015loss,
+ title = {The Loss Surfaces of Multilayer Networks},
+ author = {Choromanska, Anna and Henaff, Mikael and Mathieu, Michael and Arous, G{\'e}rard Ben and LeCun, Yann},
+ booktitle = {Artificial Intelligence and Statistics},
+ pages = {192--204},
+ year = {2015},
+ organization = {PMLR}
+}
+
@inproceedings{damianou2013deep,
title = {Deep {G}aussian Processes},
author = {Damianou, Andreas and Lawrence, Neil D},
@@ -98,7 +107,7 @@
}
@inproceedings{gardner2018gpytorch,
- title = {{GP}yTorch: Blackbox Matrix-Matrix Gaussian Process Inference With {GPU} Acceleration},
+ title = {{GP}yTorch: Blackbox Matrix-Matrix {G}aussian Process Inference With {GPU} Acceleration},
author = {Gardner, Jacob and Pleiss, Geoff and Weinberger, Kilian Q and Bindel, David and Wilson, Andrew G},
booktitle = nips,
year = {2018}
@@ -134,7 +143,7 @@
@article{harbrecht2012low,
title = {On the Low-Rank Approximation by the Pivoted {C}holesky Decomposition},
author = {Harbrecht, Helmut and Peters, Michael and Schneider, Reinhold},
- journal = {Applied numerical mathematics},
+ journal = {Applied Numerical Mathematics},
volume = {62},
number = {4},
pages = {428--440},
@@ -150,7 +159,7 @@
number = {4},
pages = {881--901},
year = {2000},
- publisher = {MIT Press}
+ publisher = {MIT Press Cambridge}
}
@article{hochreiter1997flat,
@@ -161,7 +170,7 @@
number = {1},
pages = {1--42},
year = {1997},
- publisher = {MIT Press One Rogers Street, Cambridge, MA 02142-1209, USA journals-info~…}
+ publisher = {MIT Press Cambridge}
}
@article{hutchinson1990stochastic,
@@ -209,9 +218,9 @@
}
@article{jordan1999introduction,
- title = {An Introduction to Variational Methods for Mraphical Models},
+ title = {An Introduction to Variational Methods for Graphical Models},
author = {Jordan, Michael I and Ghahramani, Zoubin and Jaakkola, Tommi S and Saul, Lawrence K},
- journal = {Machine learning},
+ journal = {Machine Learning},
volume = {37},
pages = {183--233},
year = {1999},
@@ -419,6 +428,13 @@
publisher = {SIAM}
}
+@article{chen2022randomly,
+ title = {Randomly Pivoted {C}holesky: Practical Approximation of a Kernel Matrix with Few Entry Evaluations},
+ author = {Chen, Yifan and Epperly, Ethan N and Tropp, Joel A and Webber, Robert J},
+ journal = {arXiv preprint arXiv:2207.06503},
+ year = {2022}
+}
+
@article{wainwright2008graphical,
title = {Graphical Models, Exponential Families, and Variational Inference},
author = {Wainwright, Martin J and Jordan, Michael I and others},
@@ -443,7 +459,7 @@
volume = {2},
number = {3},
year = {2006},
- publisher = {MIT press Cambridge, MA}
+ publisher = {MIT Press Cambridge}
}
@book{woodbury1950inverting,
diff --git a/main.tex b/main.tex
@@ -52,7 +52,7 @@
\usepackage{kaorefs}
\usepackage[utf8]{inputenc}
-\usepackage{amsmath,amssymb,amsthm,bm,mathtools}
+\usepackage{amsmath,amssymb,amsthm,bm,mathtools,xfrac}
\usepackage{censor}
\usepackage{tikz}
@@ -77,6 +77,8 @@
\input{macros.tex}
+\hyphenation{Hes-s-ian}
+
\graphicspath{{figures/}{./}} % Paths in which to look for images
\makeindex[columns=3, title=Alphabetical Index, intoc] % Make LaTeX produce the files required to compile the index
diff --git a/tables/preconditioner-comparison.tex b/tables/preconditioner-comparison.tex
@@ -0,0 +1,10 @@
+\begin{tabular}{lllll}
+ \toprule
+ & Rank & Speed & Closeness & Inversion \\
+ \midrule
+ \(\bm \nabla\T \bm \nabla\) & 1 & Very fast & Bad & Closed-form \\
+ \(\sum_i \bm J_i\T \bm H_i \bm J_i\) & Flexible, implicit & Slow & Arbitrarily close & Conjugate Gradient \\
+ \(\bm J\T \bm H \bm J\) & \(O\) & Fast & Okay & Woodbury \\
+ \(\bm L\T \bm L\) & Flexible, explicit & Fast (precomputable) & Good & Woodbury \\
+ \bottomrule
+\end{tabular}
+\ No newline at end of file
