commit d926ca597a95ff63857ecc405b16b83acb5ee430
parent b6ef1c9c9a8da2729f18917d1ba916723b0d7994
Author: Silas Brack <s174433@student.dtu.dk>
Date: Thu, 2 Feb 2023 19:33:01 +0100
Finished ablation experiments section.
Diffstat:
19 files changed, 188 insertions(+), 143 deletions(-)
diff --git a/chapters/acknowledgments.tex b/chapters/acknowledgments.tex
@@ -3,14 +3,14 @@
\blackout
{
- To S{\o}ren's group --- especially Marco --- for always making me feel like a part of the group and that my contribution was valued.
+ To S{\o}ren's group---especially Marco---for always making me feel like a part of the group and that my contribution was valued.
To Frederik and Marco for giving me the chance to dip my toes in ``real'' research in the most undiluted way possible, with what I perceive to be all the highs and the lows which constitute it.
To S{\o}ren, for always being up for a great chat, for giving me the opportunity to work on difficult and novel problems (like preconditioning), and for approaching everything with kindness.
I don't think I could have chosen a better supervisor.
- To J\'ulia, Brandon and Julia for the constant laughs and endless opportunities to relax on the couch when working was too much (memes surrounding me working on the couch aside) and Adam, with who I shared the experience of writing a thesis --- highs and lows both.
+ To J\'ulia, Brandon and Julia for the constant laughs and endless opportunities to relax on the couch when working was too much (memes surrounding me working on the couch aside) and Adam, with who I shared the experience of writing a thesis---highs and lows both.
To Jason and Kristi for always cultivating/nurturing in me a sense of curiosity and independence.
And for never forgetting me in the supermarket.
diff --git a/chapters/appendix.tex b/chapters/appendix.tex
@@ -1,9 +1,8 @@
\chapter{Algorithms}
% \setchapterstyle{lines}
-\label{ch:appendix-algorithms}
% \setchapterpreamble[u]{\margintoc}
-\begin{algorithm}
+\begin{algorithm}[h!]
\caption{Multi-Shift Minimum Residual (msMINRES)}\label{alg:msminres}
\Input{$\bm K \succ 0$, $\bm b$, $\bm P \succ 0$, $t_1, \ldots, t_Q$, $J > 0$}
\Output{$\bm c_1 = (\bm K + t_1)^{-1} \bm b, \ldots, \bm c_Q = (\bm K + t_Q)^{-1} \bm b$}
@@ -27,7 +26,7 @@
$\alpha_{j} \gets \bm p \cdot \bm q_{j-1}$\;
$\bm z_j \gets \bm p - \alpha_j \bm z_{j-1} - \beta_{j-1} \bm z_{j-2}$\;
$\bm q_j \gets \bm P \bm z_j$\;
- $\beta_j \gets \left({\bm z_j \bm q_j}\right)^{1/2}$\;
+ $\beta_j \gets ({\bm z_j \bm q_j})^{1/2}$\;
% $\beta_j \gets \sqrt{\bm z_j \bm q_j}$\;
$\bm z_j \gets \bm z_j / \beta_{j}$\;
$\bm q_j \gets \bm q_j / \beta_{j}$\;
@@ -37,14 +36,13 @@
$\alpha^{(q)}_{j} \gets \alpha_{j} + t_q$\;%\Comment{Compute shifted alpha}
$\eta^{(q)}_j \gets \cos^{(q)}_{j-1} \alpha^{(q)}_{j} - \sin^{(q)}_{j-1} \zeta^{(q)}_{j}$\;%\Comment{Givens rotation from 1 step ago}
$\zeta^{(q)}_{j} \gets \cos^{(q)}_{j-1} \zeta^{(q)}_{j} + \sin^{(q)}_{j-1} \alpha^{(q)}_{j}$\;
- $r^{(q)}_{j} \gets \left({\eta^{(q)2}_{j} + \beta^2_{j}}\right)^{1/2}$\;%\Comment{Compute next Givens terms}
+ $r^{(q)}_{j} \gets ({\eta^{(q)2}_{j} + \beta^2_{j}})^{1/2}$\;%\Comment{Compute next Givens terms}
% $r^{(q)}_{j} \gets \sqrt{\eta^{(q)2}_{j} + \beta^2_{j}}$\;%\Comment{Compute next Givens terms}
$\cos^{(q)}_{j} \gets \eta^{(q)}_{j} / r^{(q)}_{j}$\;
$\sin^{(q)}_{j} \gets \beta^{(q)}_{j} / r^{(q)}_{j}$\;
$\eta^{(q)}_{j} \gets \eta^{(q)}_{j} \cos^{(q)}_{j} + \sin^{(q)}_{j} \beta_{j}$\;%\Comment{Apply current Givens rotation}
- $\varphi^{(q)}_{j} \gets - \varphi^{(q)}_{j-1} \sin^{(q)}_{j}$\;%\Comment{Apply the latest Givens rotation to the Lanczos-rhs, computing the scale terms for the search vectors}
- $\varphi^{(q)}_{j-1} \gets \varphi^{(q)}_{j-1} \cos^{(q)}_{j}$\;
- $\bm d^{(q)}_j \gets \left( \bm q_{j-1} - \zeta^{(q)}_{j} \bm d^{(q)}_{j-1} - \varepsilon^{(q)}_{j} \bm d^{(q)}_{j-2} \right) / \eta^{(q)}_{j}$\;%\Comment{Get the new search vector}
+ $\varphi^{(q)}_{j} \gets - \varphi^{(q)}_{j-1} \sin^{(q)}_{j}$; $\varphi^{(q)}_{j-1} \gets \varphi^{(q)}_{j-1} \cos^{(q)}_{j}$\;%\Comment{Apply the latest Givens rotation to the Lanczos-rhs, computing the scale terms for the search vectors}
+ $\bm d^{(q)}_j \gets ( \bm q_{j-1} - \zeta^{(q)}_{j} \bm d^{(q)}_{j-1} - \varepsilon^{(q)}_{j} \bm d^{(q)}_{j-2} ) / \eta^{(q)}_{j}$\;%\Comment{Get the new search vector}
$\Delta \bm c^{(q)}_{j} \gets \bm d^{(q)}_{j} \varphi^{(q)}_{j-1}$\;
$\bm c^{(q)}_{j} \gets \bm c^{(q)}_{j-1} + \Delta \bm c^{(q)}_{j}$\;%\Comment{Update the solution}
}
@@ -56,7 +54,30 @@
\chapter{Additional Results}
% \setchapterstyle{lines}
-\label{ch:appendix-results}
+
+\begin{table*}[h]
+ \centering
+ \caption[Comparison of Hessian-vector product performance.]{Comparison of the wall-clock time for Hessian-vector, GGN-vector, and manual matrix-vector products, as well as the equivalent inverse-vector products, on a toy problem. Time measured in milliseconds, measured on an A100 GPU across 5 runs that are reported as a mean and standard deviation. Implicit inverse performed via conjugate gradient for HVP and GVP. We highlight in bold the extreme slowdown of computing the manual Hessian-vector product for a 1\,M-dimensional problem. We can see that manual computations fail above small problems, while the implicit GVP and HVP functions and their inverses scale well. We also see that the GVP is slightly faster than the HVP, both both are significantly faster than the manual Hessian-vector product.}
+ \label{tab:hessian-profile}
+ \include{tables/hessian-profile.tex}
+\end{table*}
+
+\begin{figure*}[h!]
+ \begin{subfigure}[h]{0.45\textwidth}
+ \centering
+ \includegraphics{hessian_profile_cpu.pdf}
+ \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for HVP and GVP. We can see that the explicit Hessian-vector product fails for problems with more than 10\,000 dimensions and that the GVP is faster than the HVP.}
+ \label{fig:hessian-profile-cpu}
+ \end{subfigure}
+ \hspace{0.5cm}
+ \begin{subfigure}[h]{0.45\textwidth}
+ \centering
+ \includegraphics{hessian_profile_ratios_cpu.pdf}
+ \caption[Comparison of speedup from HVP to GVP.]{Comparison of the speedup from computation of the Hessian-vector product (HVP) versus the GGN-vector product (GVP) (top) and inverse HVP versus inverse GVP using the conjugate gradient method (bottom). Speedup is calculated as the ratio of the wall-clock time of the HVP to the GVP. We can see that the two methods have the same time complexity (up to a linear factor).}
+ \label{fig:hessian-profile-ratios-cpu}
+ \end{subfigure}
+ \caption{Performance comparison of different Hessian-vector products on CPU.}
+\end{figure*}
\begin{figure*}
\centering
diff --git a/chapters/boring_experiments.tex b/chapters/boring_experiments.tex
@@ -0,0 +1,117 @@
+\setchapterpreamble[u]{\margintoc}
+\chapter{Boring Experiments}
+\label{ch:boring-experiments}
+
+\section{Hessian-Vector Products}[Hessian-Vector Products]
+
+In this project, we propose performing the Laplace approximation by only storing the computational graph for the implicit Hessian-vector product, thereby avoiding the explicit computation and storage of the quadratically-scaling Hessian matrix.
+How significant of an effect does this have, though?
+To test this, we perform Hessian-vector products \(\bm K \bm v\) and inverse Hessian-vector products \(\bm K^{-1} \bm v\), where \(\bm K \in \reals^{N \times N}\), and compare their performance.
+To compare performance with regard to memory usage and runtime, we perform these products for an increasing problem dimensionality \(N \in \{10^1, 10^2, \ldots, 10^9\}\).
+
+It is difficult to benchmark memory consumption in JAX, since the just-in-time compiler optimises computations using XLA.
+This then prevents us from stopping the program and measuring memory usage at a single point in time, as many computations will be performed at once.
+Ideally, it would be possible to measure the \emph{peak} memory usage by JAX during the running of the program.
+This is not trivial, though, as it would require querying the memory usage from the GPU asynchronously at a very high frequency and ``hope'' the peak usage point was captured.
+Instead, we simply measure the value of \(N\) at which the program fails from an out-of-memory error instead, as a proxy for memory usage.
+This experiment should then answer two questions:
+\begin{enumerate}
+ \item How much larger problems can we solve if we don't need to instantiate the whole Hessian?
+ \item How much faster is it to compute the GVP versus, say, the HVP or the explicit Hessian-vector product?
+\end{enumerate}
+
+To perform this experiment, we define the the loss function composition \(\mathcal{L}(\bm x) \equiv (g \circ f) (\bm x) \coloneqq \sum^N_{i=1} (2 x_i + 0.5)^2\) as
+\(f(\bm x) \coloneqq 2 \bm x + 0.5, g(\bm x) \coloneqq \sum^N_{i=1} x_i^2\).
+This is equivalent to the sum-of-square-error loss \(g\) for a simple linear model \(f\).
+Since \(f(x)\) is a linear function, the GGN matrix will exactly equal the actual Hessian.\sidenote{This allows us to compare our implementation of the GGN-vector product to the Hessian-vector product. While this test is not perfect (since it only holds for linear functions), it is a good sanity check. We find that the two are indeed the same in these experiments.}
+We then initialise \(\bm x\) to a random vector and compute the matrix-vector product \(\bm K \bm v\) where \(\bm v\) is a random vector using three Hessian-vector products: the JAX efficient HVP and GVP functions and a manual Hessian instantiation and multiplication.
+We then measure the total wall-clock time for computing this Hessian product by \(\bm v\).
+
+\begin{marginfigure}
+ \centering
+ \includegraphics{hessian_profile_broken.pdf}
+ \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for manual Hessian-vector products, HVPs, and GVPs for increasing numbers of parameters \(D\). Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse is performed via conjugate gradient for HVP and GVP. We can see that the explicit inversion fails for \(D > 10 k\) and that explicit Hessian-vector products fail for \(D > 1 M\) (and, for \(D = 1 M\), the manual product is extremely slow). Implicit products, however, can be computed so long as the output vector fits in memory, while implicit inversion fails for \(D = 1 B\). Notice the broken axis and thus, in particular, the performance difference between the explicit and implicit inversion methods.}
+ \label{fig:hessian-profile}
+\end{marginfigure}
+
+Furthermore, since we are looking to compute some variation of the matrix inverse of \(\bm K\), and inverting a high-dimensional matrix can also be very expensive, we also compare the performance of the \emph{inverse} Hessian-vector product \(\bm K^{-1} \bm v\).
+To compute the inverse for the efficient implicit HVP and GVP functions, we use the conjugate gradient method, which only requires the use of matrix-vector multiplications to compute this inverse product.
+To compute the explicit inverse, we use the \(\texttt{scipy.linalg}\) package, which uses the LAPACK library to compute the inverse.
+
+% RESULTS
+The results for this experiment have been visualised in \cref{fig:hessian-profile}.
+Additionally, the raw results for this experiment have been summarised in the appendix, in \cref{tab:hessian-profile}.
+We can see that the explicit Hessian-vector product fails for \(D > 1 M\) parameters, while the explicit inversion fails for \(D > 10 k\) parameters.
+This is not surprising, as the explicit Hessian is a \(D \times D\) matrix, and so the memory requirements for this matrix are \(D^2\) parameters.
+Additionally, while the wall clock time for computing the Hessian-vector product is sublinear with \(D\), explicit Hessian-vector product is over an order of magnitude slower for \(D = 1 \mathrm{M}\) than for \(D = 100 \mathrm{k}\).
+This may be due to the explicit Hessian, at this size, being too large to fit in the cache of the GPU.
+
+The largest number of parameters tested is 1\,B.
+Above 1\,B parameters, even implicit methods fail, although this failure occurs in the instantiation of a single random vector of that size (so it would fail on this hardware for any method, since we can't even store the \(\bm v\) nor the product \(\bm K \bm v\)).
+Notably, however, these implicit Hessian-vector product methods allowed for computation of matrix products for problems three orders of magnitude larger, and computation of solves for problems four orders of magnitude larger, than the explicit methods.
+This is a significant improvement in scalability, and is a key reason why we use these methods in our experiments.
+Additionally the implicit inverse products are significantly faster than explicit inversion, even for small problems.
+Performance also appears identical regardless of whether the explicit inverse is computed by actual inversion or by solving a system of linear equations.
+
+\begin{marginfigure}
+ \centering
+ \includegraphics{hessian_profile_ratios.pdf}
+ \caption[Comparison of speedup from HVP to GVP.]{Comparison of the speedup from computation of the Hessian-vector product (HVP) versus the GGN-vector product (GVP) (top) and inverse HVP versus inverse GVP using the conjugate gradient method (bottom). Speedup is calculated as the ratio of the wall-clock time of the HVP to the GVP. We can see that the two methods have the same time complexity (up to a linear factor).}
+ \label{fig:hessian-profile-ratios}
+\end{marginfigure}
+
+However, due to the scale of the figure, it isn't clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product.
+We can then plot the speedup from going from the HVP to the GVP as a function of the number of model parameters.
+This can be calculated as the ratio of the wall-clock time of the HVP to the wall-clock time of the GVP.
+Thus, larger values suggest greater performance gains for the GVP, while values closer to one suggest there is no performance gain.
+These results can be seen in \cref{fig:hessian-profile-ratios}.
+The number of parameters on the \(x\)-axis are exponentially increasing.
+We can see that the computation of the GVP is always faster than the computation of the HVP, and that the speedup is approximately linear in the number of parameters.
+This is in line with what we expected (see \cref{sec:practical-ggn}).
+We can also see a slight decrease in the speedup as the number of parameters increases for the Hessian-vector products.
+However, it is not clear why this is the case.
+
+\section{Sampling Ablation}
+
+\begin{marginfigure}
+ \centering
+ \includegraphics{ablation/mnist_ablation.pdf}
+ \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST. We can see that sampling fails categorically for \(Q < 5\) and succeeds categorically for \(Q \geq 10\). SHOULD I CHANGE THRESHOLD TO 0.01 OR SO?}
+ \label{fig:mnist-ablation}
+\end{marginfigure}
+
+Contour integral quadrature (CIQ) converges to the exact solution for \(\lim_{Q \to \infty} \bm s^{\mathrm{CIQ}}_{Q} = \bm K^{-1/2} \bm v\).
+However, CIQ is typically close to the exact solution for low values of \(Q\).\sidenote{\textcite{pleiss2020fast} find that \(Q=20\) is approximately sufficient for most problems.}
+We now attempt to determine the number of quadrature points \(Q\) required to compute an adequate posterior samples.
+
+We vary \(Q \in \{1, \ldots, 20\}\) to determine the number of quadrature points required to effectively sample from the Laplace approximation.
+We then compute the Mahalanobis distance between the samples and the true posterior mean, and compare this to the chi-squared distribution with \(D\) degrees of freedom.
+By comparing the empirical CDF of the Mahalanobis distances to the theoretical \(\chi^2_D\) CDF, we can perform the Kolmogorov-Smirnov test for each value of \(Q\).
+This p-value is then used to determine whether the null hypothesis that the samples are multivariate normal can be rejected based on some significance threshold.
+Sample evaluation is explained in more detail in \cref{sec:sampling-evaluation}.
+The results of this experiment are shown in \cref{tab:mnist-ablation,fig:mnist-ablation}.
+Quantile--quantile plots for each value of \(Q\) are shown in the appendix, in \cref{fig:ablation-chisq}.
+
+\begin{margintable}
+ \centering
+ \caption[Ablation over quadrature points for MNIST.]{Sampling time and p-values for Kolmogorov-Smirnov with varying quadrature points \(Q\) on MNIST with prior precision \(\alpha = 0.1\), for 200 posterior samples. Lower p-value indicates a higher probability of rejecting normality. Bold indicates sampling was successful based on visual inspection of quantile--quantile plots. Since, for values of \(Q \geq 5\), samples appear visually to be multivariate normal, we can conclude that we should choose a significance threshold of 0.01 or lower.}
+ \label{tab:mnist-ablation}
+ \include{tables/mnist_ablation.tex}
+\end{margintable}
+
+% RESULTS, DISCUSSION
+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\).
+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.
+Inspection of \cref{fig:ablation-chisq} in the appendix suggests that there is not a significant difference between using 5, 10, or 20 quadrature points.
+Using less than 4 quadrature points, however, results in sampling failure, while using exactly 4 quadrature points results in a borderline success.
+For 200 posterior samples, a more appropriate significance threshold of \(0.01\) or lower would reject sampling for \(Q \leq 4\), since the cases where sampling fails completely have p-values that are below single floating point precision (i.e., they are effectively zero).
+
+In taking as much as 20 to 30 minutes to compute 200 posterior samples for an MNIST model with 15\,000 parameters, the time required to compute the samples is not negligible.
+In practice, it is likely that only a small number of samples will be obtained (possibly around 5--15), since for each sample it is necessary to store a parameter vector of size \(D\).
+However, the sampling procedure still takes over 20 minutes to compute 10 samples.
+As such, the time required to compute the samples may be a significant factor in the overall time required to compute the Laplace approximation.
+In particular, for models with a very large number of parameters, the time required to compute the samples may be prohibitive when compared when the near-instant time required to sample from the diagonal Laplace approximation.
+Even if the sampling time scales sublinearly with the number of parameters, it would still prevent the application in methods such as online Laplace, which require the Laplace approximation to be computed at each iteration, from being used in practice.
+However, this sampling time constraint is also significant for MCMC methods, which are the most common method for sampling from the posterior distribution of a BNN.
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -16,7 +16,7 @@ This method allows us to compute a differentiable log-marginal likelihood bound,
We showed that this method is a valid method for training the full Laplace approximation, and that it can be used to train the full Laplace approximation on a convolutional neural network.
However, we also found that the marginal training procedure to optimise the hyperparameters of the full Laplace approximation is not yet feasible, since the hyperparameters converge to unreasonable values.
-Finally, we showed that it is possible to perform the full Laplace approximation on a convolutional neural network, train its weights with marginal likelihood training, and sample from its posterior.
+Finally, we showed that it is possible to perform the full Laplace approximation with marginal likelihood training, and sample from its posterior.
We demonstrated on a simple sine wave regression problem and an MNIST classification problem that our method can be used to train the full Laplace approximation and yield predictive performance which rivals deterministic training.
However, we found that the sampling procedure is too slow for online Laplace, and may become a significant bottleneck for inference using post-hoc Laplace on large-scale problems.
@@ -39,7 +39,7 @@ As such, there is still much work to be done to understand the methods we have d
% Also diagonal raw values maybe, compare the two?
% Diagonal lets you know exactly which parameters contribute to the largest eigenvalues
We also have not shown that the full Laplace approximation shows improved predictive performance or uncertainty calibration when compared to the diagonal or Kronecker-factored Laplace approximations.
-Furthermore, the full Laplace approximation should be compared to other methods for Bayesian inference, such as variational approximations, deep ensembles, Monte Carlo dropout, and Markov chain Monte Carlo --- both in terms of performance and runtime.
+Furthermore, the full Laplace approximation should be compared to other methods for Bayesian inference, such as variational approximations, deep ensembles, Monte Carlo dropout, and Markov chain Monte Carlo---both in terms of performance and runtime.
% RELATED TO ACTUAL USEFULNESS OF THE METHOD
Furthermore, the feasibility of the full Laplace approximation must be demonstrated on a large-scale problem.
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -1,130 +1,13 @@
\setchapterpreamble[u]{\margintoc}
-\chapter{Experiments}
-\labch{experiments}
+\chapter{Cool Experiments}
+\label{ch:cool-experiments}
-WARNING: THIS CHAPTER IS BY NO MEANS FINISHED
+WARNING: THIS CHAPTER IS NOT FINISHED
-\section{Hessian-Vector Products}[Hessian-Vector Products]
-
-THINKING OF MOVING THIS SECTION TO \cref{ch:intro} OR \cref{ch:laplace}
-
-In this project, we propose performing the Laplace approximation by only storing the computational graph for the implicit Hessian-vector product, thereby avoiding the explicit computation and storage of the quadratically-scaling Hessian matrix.
-How significant of an effect does this have, though?
-To test this, we perform Hessian-vector products \(\bm K \bm v\) and inverse Hessian-vector products \(\bm K^{-1} \bm v\), where \(\bm K \in \reals^{N \times N}\), and compare their performance.
-To compare performance with regard to memory usage, we perform these products for an increasing problem dimensionality \(N \in \{10^1, 10^2, \ldots, 10^9\}\).
-
-It is difficult to benchmark memory consumption in JAX, since the just-in-time compiler optimises computations using XLA.
-This then prevents us from stopping the program and measuring memory usage at a single point in time, as many computations will be performed at once.
-Ideally, it would be possible to measure the \emph{peak} memory usage by JAX during the running of the program.
-This is not trivial, though, as it would require querying the memory usage from the GPU asynchronously at a very high frequency and ``hope'' the peak usage point was captured.
-Instead, we simply measure the value of \(N\) at which the program fails from an out-of-memory error instead, as a proxy for memory usage.
-
-This experiment will then answer two questions:
-\begin{enumerate}
- \item How much larger problems can we solve if we don't need to instantiate the whole Hessian?
- \item How much faster is it to compute the GVP versus, say, the HVP or the explicit Hessian-vector product?
-\end{enumerate}
-
-\begin{marginfigure}
- \centering
- \includegraphics{hessian_profile.pdf}
- \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for Hessian-vector and GGN-vector products. We can see that the explicit Hessian-vector product fails for problems with more than 10\,000 dimensions and that the GVP is faster than the HVP.}
- \label{fig:hessian-profile}
-\end{marginfigure}
-
-We then compose the loss function \((g \circ f) (\bm x) \coloneqq \sum^N_{i=1} (2 x_i + 0.5)^2\) as
-\(f(\bm x) \coloneqq 2 \bm x + 0.5, g(\bm x) \coloneqq \sum^N_{i=1} x_i^2\).
-This is equivalent to the sum-of-square-error loss \(g\) for a simple linear model \(f\).
-Since \(f(x)\) is a linear function, the GGN matrix will exactly equal the actual Hessian.
-We then initialise \(\bm x\) to a random vector and compute the matrix-vector product \(\bm K \bm v\) where \(\bm v\) is a random vector using three Hessian-vector products: the JAX efficient HVP and GVP functions and a manual Hessian instantiation and multiplication.
-We then measure the total wall-clock time for computing this Hessian and multiplying it by \(\bm v\).
-
-Furthermore, since in this project we typically compute some variation of the matrix inverse of \(\bm K\), and inverting a high-dimensional matrix can also be very expensive, we also compare the performance of the \emph{inverse} Hessian-vector product \(\bm K^{-1} \bm v\).
-To compute the inverse for the efficient implicit HVP and GVP functions, we use the conjugate gradient method, which only requires the use of matrix-vector multiplications to compute this inverse product.
-
-% RESULTS
-
-for values of 100k parameters or more, the manual computation failed, for both the inverse and the vector product (because the explicit \(N \times N\) matrix instantiation fails)
-The largest number of parameters tested is 100M.
-Even the implicit methods failed for 1B parameters, although this failure occurred in the instantiation of a single random vector of that size (and so would fail on this hardware for any method).
-Notably, however, these implicit Hessian-vector product methods allowed for computation of matrix products and solves for problems up to 1000 times greater than the explicit method.
-
-\cref{fig:hessian-profile}
-
-some analysis
-
-comment on performance of manual products
-comment on weird dip in inverse, possibly something to do with caching or something like that? like related to the size of the L2 cache
-
-\begin{marginfigure}%[-4cm]
- \centering
- \includegraphics{hessian_profile_ratios.pdf}
- \caption[Comparison of speedup from HVP to GVP.]{Comparison of the speedup from computation of the Hessian-vector product (HVP) versus the GGN-vector product (GVP) (top) and inverse HVP versus inverse GVP using the conjugate gradient method (bottom). Speedup is calculated as the ratio of the wall-clock time of the HVP to the GVP. We can see that the two methods have the same time complexity (up to a linear factor).}
- \label{fig:hessian-profile-ratios}
-\end{marginfigure}
-
-however, due to the scale of the figure, it isn't completely clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product.
-We can then plot the speedup from going from the HVP to the GVP as a function of the number of model parameters.
-This can be calculated as the ratio between the wall-clock time of the HVP to the wall-clock time of the GVP.
-Thus, larger values suggest greater performance gains for the GVP, while values closer to one suggest there is no performance gain.
-These results can be seen in \cref{fig:hessian-profile-ratios}.
-
-The number of parameters on the \(x\)-axis are exponentially increasing.
-This is in line with the theory, as noted in \cref{sec:practical-ggn} (SHOULD IT BE MOVED HERE?)
-
-\section{Sampling Ablation}
-
-THINKING OF MOVING THIS SECTION TO \cref{ch:sampling}
-
-\begin{marginfigure}
- \centering
- \includegraphics{ablation/mnist_ablation.pdf}
- \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST. We can see that sampling fails categorically for \(Q < 5\) and succeeds categorically for \(Q \geq 10\). FIX THIS PLOT.}
- \label{fig:mnist-ablation}
-\end{marginfigure}
-
-Contour integral quadrature (CIQ) converges to the exact solution for \(\lim_{Q \to \infty} \bm s^{\mathrm{CIQ}}_{Q} = \bm K^{-1/2} \bm v\).
-However, CIQ is typically close to the exact solution for low values of \(Q\).\sidenote{\textcite{pleiss2020fast} find that \(Q=20\) is approximately sufficient for most problems.}
-We then investigated the number of quadrature points \(Q\) required to determine an adequate approximate solution for the Laplace approximation.
-
-We vary \(Q \in \{1, \ldots, 20\}\) to determine the number of quadrature points required for the Laplace approximation.
-We then compute the Mahalanobis distance between the samples and the true posterior mean, and compare this to the chi-squared distribution with \(D\) degrees of freedom.
-We then compute the p-value of the Kolmogorov-Smirnov test for each value of \(Q\) by comparing the empirical CDF of the Mahalanobis distances to the theoretical \(\chi^2_D\) CDF.
-This p-value is then used to determine whether the null hypothesis that the samples are multivariate normal can be rejected based on some significance threshold.
-Sample evaluation is explained in more detail in \cref{sec:sampling-evaluation}.
-The results of this experiment are shown in \cref{tab:mnist-ablation,fig:mnist-ablation}.
-
-\begin{margintable}
- \centering
- \caption[Ablation over quadrature points for MNIST.]{Sampling time and p-values for Kolmogorov-Smirnov with varying quadrature points \(Q\) on MNIST with prior precision \(\alpha = 0.1\), for 200 posterior samples. Lower p-value indicates a higher probability of rejecting normality. Bold indicates sampling was successful based on visual inspection of quantile--quantile plots (see \cref{ch:appendix-results}, \cref{fig:ablation-chisq}). Since, for values of \(Q \geq 5\), samples appear visually to be multivariate normal, we can conclude that we should choose a significance threshold of 0.01 or lower.}
- \label{tab:mnist-ablation}
- \include{tables/mnist_ablation.tex}
-\end{margintable}
-
-% RESULTS
-
-Based on \cref{fig:mnist-ablation}, sampling appears to be successful (based on a significance threshold of \(0.05\)) for \(Q \in \{10, \ldots, 20\}\).
-
-% DISCUSSION
-However, inspection of CHISQ PLOTS suggests that there is not a significant difference between these values and \(Q = 5\).
-In general, significance testing can be sensitive to the number of samples used, and, for large numbers of samples, these tests can excessively reject the null hypothesis.
-However, the number of samples is generally fixed as a function of the number of parameters in the model and the amount of memory available.
-As such, it is important to select an appropriate significance threshold for the test based on the number of posterior samples that are required.
-For 200 posterior samples, a more appropriate significance threshold of \(0.01\) or lower would reject sampling for \(Q \leq 3\), since the cases where sampling completely fails have p-values that are below single floating point precision.
-Furthermore, the plots in CHISQ PLOTS suggest that the performance of the sampling method is effectively equal for \(Q \in \{5, \ldots, 20\}\).
-
-In terms of the time required to compute the samples, the time required to compute the samples is XXXX in \(Q\).
-However, in taking as much as 20 to 30 minutes to compute 200 samples of an MNIST model with 15\,000 parameters, the time required to compute the samples is not negligible.
-In practice, it is likely that only a small number of samples will be obtained (around 5-15), since for each sample it is necessary to store a parameter vector of size \(D\).
-However, for 10 samples, it still takes XXXX minutes.
-As such, the time required to compute the samples may be a significant factor in the overall time required to compute the Laplace approximation.
-In particular, for models with a very large number of parameters, the time required to compute the samples may be prohibitive when compared when the near-instant time required to sample from the diagonal Laplace approximation.
-Even if the sampling time scales sublinearly with the number of parameters, it would still prevent methods such as online Laplace, which require the Laplace approximation to be computed at each iteration, from being used in practice.
-However, this sampling time constraint is also significant for MCMC methods, which are the most common method for sampling from the posterior distribution of a BNN.
-
-\section{Sampling for the Sine Function}[The Sine Function]
+\section{The Sine Function}
The choice of \(\alpha\) is a non-trivial one
+Show plot of predictions for alpha 100 and alpha 1, to show that alpha 1 looks stoopid.\
\begin{marginfigure}
\centering
@@ -136,18 +19,19 @@ The choice of \(\alpha\) is a non-trivial one
\begin{figure}
\centering
\includegraphics{chisq_sin.pdf}
- \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for visualising normality of 1000 posterior samples from the full Laplace approximation for a fully connected neural network modelling \(f(x) = \sin(x)\). Since the sample square Mahalanobis distances appear chi-squared distributed, the samples themselves are multivariate normal with the correct (known) covariance.}
+ \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(x)\), 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}
-\section[Sampling for MNIST]{Sampling for MNIST}[MNIST]
+\section{MNIST}
-Sampling was successful for values of alpha as low as 0.1,
+Sampling was successful for values of alpha as low as 0.1, as can be seen in \cref{fig:chisq-mnist}.
+Thus, we sample 200 posterior samples instead of the 1000 samples used in the previous experiment to demonstrate what these plots look like for a more realistic number of samples.
\begin{marginfigure}
\centering
- \includegraphics{eigenvalues_mnist_55k_0.1_old.pdf}
- \caption[Spectrum of the posterior precision on MNIST.]{Hyperexponential distribution of eigenvalues of Laplace posterior precision on MNIST with a convolutional NN with 17k parameters and a 3 by 3 filter. There are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so we can conclude that the posterior covariance has an effective rank close to nine. FIX THIS PLOT.}
+ \includegraphics{mnist_spectrum_3vs5.pdf}
+ \caption[Spectrum of the posterior precision on MNIST.]{Spectrum of the Laplace posterior precision on MNIST with a CNN. For the \(3 \times 3\) filter, there are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so the precision has an effective rank close to nine. For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, with the spectrum appearing smooth (i.e., not having any distinct clusters). This suggests that the effective rank of the precision is over 50.}
\label{fig:spectrum-mnist}
\end{marginfigure}
@@ -157,3 +41,10 @@ The choice of prior precision \(\alpha\) is non-trivial.
Furthermore, since the prior precision is the lowest eigenvalue of the posterior precision matrix, we
While we can ``choose'' \(\alpha\) to be a value large enough to reduce conditioning issues, this somewhat invalidates the Bayesian interpretation of the prior, which should theoretically be chosen based on prior knowledge of the problem, can also be learned from the data.
Choosing the prior in order to reduce ill-conditioning of our posterior precision means we sacrifice some of the benefits of the prior.
+
+\begin{figure}
+ \centering
+ \includegraphics{figures/ablation/chisq_55k_0.1_20q_100precondpoints_cg.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}
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -311,8 +311,6 @@ However, the GGN approximation is better behaved than the exact Hessian, thereby
This is because the GGN approximation is designed to be positive semi-definite, while the exact Hessian can be vulnerable to negative curvature.
This feature can limit the use of the exact Hessian for optimisation algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
-% The Levenberg-Marquardt heuristic~\cite{pascanu2013revisiting}
-
\subsection{The Spectrum of the GGN Approximation}[The Spectrum of the GGN]
\label{sec:spectrum-ggn}
diff --git a/chapters/preface.tex b/chapters/preface.tex
@@ -11,7 +11,7 @@
I had a slow start to the thesis, since for the first month or so I just explored the software and the literature, and tried to get a feel for the problem.
As I tried to generate some preliminary results, I encountered a few roadblocks.
- As it turned out, these problems were quite fundamental to the problem --- but S{\o}ren had recently read a paper which addressed them and invited me to a Hackathon with his research group, where we would try to implement the methods from the paper.
+ As it turned out, these problems were quite fundamental to the problem---but S{\o}ren had recently read a paper which addressed them and invited me to a Hackathon with his research group, where we would try to implement the methods from the paper.
Initially, I was a bit hesitant to go, since I had never been to a Hackathon before, and I was worried that I would be out of my depth and that I would not be able to contribute.
However, Lorenza helped me get over my fear of marginalisation and I decided to go, and I am glad that I did.
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -424,7 +424,7 @@ Notably, the quality of the preconditioner depends largely on the ease of invert
In practice, increasing the rank of the pivoted Cholesky factorisation will increase the condition number of this matrix, since we are approximating more of the largest eigenvectors and so if we approximate more eigenvectors, we will capture progressively smaller eigenvalues.
I've found that it's most effective to keep \(K\) low, such that the \(K \times K\) outer product of the pivoted Cholesky factorisation matrix \(L\) approximates few enough of the largest eigenvectors that the eigenvalue of the largest is close to that of the smallest (meaning this outer product is well-conditioned).
We thus need to choose \(K\) such that it is large enough to approximate enough large eigenvectors and small enough to ensure the condition number of the outer product is low.
-Furthermore, you want enough observations \(B\) to be included in the preconditioner such that it's possible to accurately approximate the eigenvectors of the problem --- it may be optimal to use the whole dataset, though this may worsen the conditioning of the outer product.
+Furthermore, you want enough observations \(B\) to be included in the preconditioner such that it's possible to accurately approximate the eigenvectors of the problem---it may be optimal to use the whole dataset, though this may worsen the conditioning of the outer product.
Furthermore, the pivoted Cholesky preconditioner is not easy to implement in JAX.
The pivoted Cholesky decomposition is not implemented in JAX, and so we would need to implement it manually.
diff --git a/figures/ablation/mnist_ablation.pdf b/figures/ablation/mnist_ablation.pdf
Binary files differ.
diff --git a/figures/hessian_profile_all.pdf b/figures/hessian_profile_all.pdf
Binary files differ.
diff --git a/figures/hessian_profile_broken.pdf b/figures/hessian_profile_broken.pdf
Binary files differ.
diff --git a/figures/hessian_profile.pdf b/figures/hessian_profile_cpu.pdf
Binary files differ.
diff --git a/figures/hessian_profile_hvp.pdf b/figures/hessian_profile_hvp.pdf
Binary files differ.
diff --git a/figures/hessian_profile_implicit.pdf b/figures/hessian_profile_implicit.pdf
Binary files differ.
diff --git a/figures/hessian_profile_ratios.pdf b/figures/hessian_profile_ratios.pdf
Binary files differ.
diff --git a/figures/hessian_profile_ratios.pdf b/figures/hessian_profile_ratios_cpu.pdf
Binary files differ.
diff --git a/figures/mnist_spectrum_3vs5.pdf b/figures/mnist_spectrum_3vs5.pdf
Binary files differ.
diff --git a/main.tex b/main.tex
@@ -71,6 +71,7 @@
\tikzstyle{mydashed}=[dash pattern=on 2pt off 2pt]
\newcommand{\tick}[2]{\draw[thick] (#1) ++ (#2:0.1) --++ (#2-180:0.2)} %0.03*\xmax
+\usepackage{float}
\usepackage{caption}
\usepackage{subcaption}
@@ -248,6 +249,7 @@
\input{chapters/laplace.tex}
\input{chapters/training.tex}
\input{chapters/sampling.tex}
+\input{chapters/boring_experiments.tex}
\input{chapters/experiments.tex}
\input{chapters/conclusion.tex}
diff --git a/tables/hessian-profile.tex b/tables/hessian-profile.tex
@@ -0,0 +1,16 @@
+\begin{tabular}{rrrrrrrrrrrrr}
+ \toprule
+ & \multicolumn{2}{c}{HVP} & \multicolumn{2}{c}{GVP} & \multicolumn{2}{c}{Manual} & \multicolumn{2}{c}{HVP Inverse} & \multicolumn{2}{c}{GVP Inverse} & \multicolumn{2}{c}{Manual Inverse} \\
+ {\(D\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} & \multicolumn{1}{c}{\(\mu\)} & \multicolumn{1}{c}{\(\sigma\)} \\
+ \midrule
+ 10\, & 91.2 & 3.6 & 35.7 & 1.9 & 105.6 & 5.2 & 277.8 & 7.9 & 178.6 & 9.0 & 1703.6 & 646.7 \\
+ 100\, & 91.6 & 2.7 & 36.5 & 2.3 & 111.6 & 2.1 & 288.3 & 8.7 & 200.8 & 26.9 & 1732.7 & 672.1 \\
+ 1\,k & 89.5 & 3.5 & 36.8 & 3.3 & 119.9 & 2.3 & 308.6 & 7.5 & 256.5 & 34.3 & 1767.3 & 653.4 \\
+ 10\,k & 91.4 & 3.8 & 36.7 & 3.4 & 118.0 & 2.4 & 312.4 & 11.1 & 228.9 & 38.1 & 2165.0 & 720.8 \\
+ 100\,k & 89.7 & 1.9 & 36.7 & 3.3 & 179.9 & 13.5 & 325.8 & 5.3 & 206.2 & 8.5 & --- & --- \\
+ 1\,M & 91.3 & 1.5 & 38.5 & 3.0 & \tabf{2961.8} & \tabf{1060.9} & 351.7 & 6.6 & 252.0 & 33.0 & --- & --- \\
+ 10\,M & 101.2 & 5.1 & 43.8 & 2.5 & --- & --- & 343.4 & 13.5 & 216.5 & 12.6 & --- & --- \\
+ 100\,M & 100.0 & 2.4 & 45.5 & 2.1 & --- & --- & 372.3 & 13.9 & 271.5 & 34.4 & --- & --- \\
+ 1\,B & 105.9 & 4.8 & 52.4 & 7.5 & --- & --- & --- & --- & --- & --- & --- & --- \\
+ \bottomrule
+\end{tabular}
