commit f731fc765fc8d1175bd07f9e738d3a1c368b2aab
parent 2c88600f7a817437dd08b1125c5d9c25ee6c9394
Author: Silas Brack <s174433@student.dtu.dk>
Date: Mon, 6 Feb 2023 18:29:19 +0100
Close to final.
Diffstat:
15 files changed, 245 insertions(+), 174 deletions(-)
diff --git a/chapters/ablation_experiments.tex b/chapters/ablation_experiments.tex
@@ -0,0 +1,121 @@
+\setchapterpreamble[u]{\margintoc}
+\chapter{Ablation Experiments}
+\label{ch:ablation-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.
+These experiments were run on a single NVIDIA A100 GPU.
+
+% 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 depend on the \(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}.
+These experiments were run on a single NVIDIA A100 GPU.
+
+\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\,k 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 Laplace methods that 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.
+
+\marginnote{Results for the experiments in this chapter run on CPU are shown in the appendix, in \cref{fig:hessian-profile-cpu}.}
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -23,6 +23,9 @@ By analysing quantile--quantile plots of our posterior samples and visualising o
We then benchmark the performance of Hessian-vector products and estimate that our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
This is the first method for performing the Laplace approximation while only accessing the Hessian implicitly.
+BLEH, STILL NOT TOO HAPPY WITH MY ABSTRACT.
+I FEEL LIKE I SHOULD MAKE IT MORE CONCISE AND FOCUS ON THE RESULTS.
+
% Our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
% We compute the full Laplace approximation on both a sine function and MNIST.
% To do this, we demonstrate that the samples generated by our approximate sampling procedure are correctly distributed and that training the marginal likelihood using our method is feasible and produces meaningful results.
diff --git a/chapters/boring_experiments.tex b/chapters/boring_experiments.tex
@@ -1,121 +0,0 @@
-\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.
-These experiments were run on a single NVIDIA A100 GPU.
-
-% 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 depend on the \(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}.
-These experiments were run on a single NVIDIA A100 GPU.
-
-\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\,k 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 Laplace methods that 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.
-
-\marginnote{Results for the experiments in this chapter run on CPU are shown in the appendix, in \cref{fig:hessian-profile-cpu}.}
diff --git a/chapters/conclusion.tex b/chapters/conclusion.tex
@@ -5,7 +5,7 @@
\section{Summary and Key Results}
In this project, we aimed to develop a scalable full Laplace approximation for Bayesian neural networks.
-In doing so, we developed a novel method for approximately sampling from the full Laplace posterior and methods for preconditioning the Hessian of the negative log posterior, which involved attempting to better understand the spectrum of the Generalised Gauss-Newton Laplace posterior precision.
+In doing so, we developed a novel method for approximately sampling from the full Laplace posterior and methods for preconditioning the Hessian of the negative log posterior, which involved attempting to better understand the spectrum of the generalised Gauss-Newton Laplace posterior precision.
To test the quality of our sampling procedure, we showed how to evaluate the quality of the samples it produces while only accessing the samples themselves and the desired posterior precision, enabling us and to choose a suitable preconditioner for it.
This evaluation method showed that, with an appropriate preconditioner, we can reliably produce samples which are Gaussian-distributed with the correct mean and covariance, as desired.
As such, we have shown that our approximate sampling method is a valid (albeit slow) method for sampling from the full Laplace posterior.
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -1,48 +1,95 @@
\setchapterpreamble[u]{\margintoc}
-\chapter{Cool Experiments}
-\label{ch:cool-experiments}
+\chapter{Experiments}
+\label{ch:experiments}
+
+In \cref{ch:ablation-experiments} we tested whether the approximate sampling scheme is able to produce correct posterior samples on an arbitrary network.
+In this chapter, we use our custom bound on the marginal likelihood to train a Bayesian neural network on the marginal via online Laplace.
+We then use our approximate sampling scheme to sample from the posterior and propagate the samples through the network to produce a predictive distribution.
+In this way, we test the entire Bayesian neural network pipeline, from training to inference.
\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.\
+We first use a simple sine function as a toy example to demonstrate the Bayesian neural network pipeline for regression.
+First, we must train the network on the marginal likelihood to demonstrate that the model can learn a good posterior distribution.
+We then sample from the posterior and visualise the posterior predictive samples to showcase the quality of the predictive distribution.
\begin{marginfigure}
\centering
- \includegraphics{sin_posterior_predictive_joint.pdf}
- \caption[Laplace predictions for a sine curve.]{Visualisation of the predicted output of multilayer perceptron (MLP) modelling a sine curve.}
+ \includegraphics{sin_posterior_predictive_bothalphas_clear.pdf}
+ \caption[Laplace predictive posterior for a sine curve.]{Laplace predictive posterior (with a confidence interval given by one standard deviation) for a multilayer perceptron (MLP) modelling \(f(x) = \sin(5 x + 1)\) with parameters trained on the marginal likelihood for a fixed prior precision \(\alpha = 70\) (top) and \(\alpha = 1\) (bottom). We compute the marginal using our approximate lower bound and sample from the posterior using our approximate sampling scheme. The choice of prior precision has a large effect on the posterior uncertainty, but, when it is chosen appropriately, our approximate methods are able to learn a reasonable predictive distribution (i.e., more uncertainty in the regions with no data) with a similar mean prediction to the MSE-optimised network.}
\label{fig:sin}
\end{marginfigure}
+In \cref{fig:sin} we show the posterior predictive distribution for a fully connected neural network modelling a sine curve and compare it to a neural network trained on the MSE loss, i.e., maximum likelihood for regression.
+We vary the prior precision \(\alpha\) to show how the posterior predictive distribution changes.
+We have the same prior precision for all weights and biases, with values \(\alpha = 70\) for the top figure and \(\alpha = 1\) for the bottom figure.
+In \cref{fig:chisq-sin} we show the chi-squared sample plot for the posterior samples on this sine regression task to demonstrate that the posterior samples are multivariate normal with the correct covariance, as we explain in \cref{sec:sampling-evaluation}.
+
+We can see from \cref{fig:sin} that for \(\alpha = 70\) the posterior predictive distribution has a low uncertainty in the regions with data, but a high uncertainty in the regions without data.
+Furthermore, for this value of \(\alpha\) the mean of the posterior predictive distribution is relatively close to the mean of the MSE-optimised network.
+
+For a very low posterior precision \(\alpha = 1\), the posterior predictive distribution has a high uncertainty throughout the whole function domain.
+This is because the prior uncertainty is very high, so the posterior uncertainty is also high.
+Furthermore, because of the noise in the posterior samples, the mean of the posterior predictive distribution is not close to that of the MSE-optimised network.
+
+We can thus see that maximising the marginal likelihood is able to learn a reasonable predictive distribution (i.e., more uncertainty in the regions with no data) with a similar mean prediction to the MSE-optimised network, when the prior precision is chosen appropriately.
+However, if the prior precision is too low, the posterior predictive distribution will have a high uncertainty throughout the whole function domain, and the mean of the posterior predictive distribution will not be close to that of the MSE-optimised network.
+In these experiments, we have chosen the prior precision to be fixed, but in practice, we would choose the prior precision by maximising the marginal likelihood or by cross-validation.
+However, training the hyperparameters on the marginal likelihood is a sensitive procedure, and so we have not yet managed to successfully optimise the prior precision.
+
+Furthermore, we can see from \cref{fig:chisq-sin} that the sample square Mahalanobis distances appear chi-squared distributed, so the samples themselves are multivariate normal with the Laplace posterior precision.
+We deliberately choose an extremely large number of posterior samples (1000) to demonstrate the accuracy of the posterior samples.
+However, it would not be practical to compute this many posterior samples for larger problems, since we would have to either store them all in memory or compute them in batches.
+
\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)\), using CIQ. 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(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}
\section{MNIST}
-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.
+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.}
+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).
+Because of this, it is difficult to choose an appropriate prior precision \(\alpha\) for the problem, since we would have to either learn \(\alpha\) by maximising the marginal likelihood or choose \(\alpha\) via cross-validation.
+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.}
+We can see these results with \(\alpha = 0.1\) in \cref{fig:chisq-mnist}.
+
+\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}
+
+\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 the MNIST problem.
+
+We also show the spectrum of the posterior precision matrix for the MNIST problem in \cref{fig:spectrum-mnist} for a CNN with a \(3 \times 3\) filter and a \(5 \times 5\) filter.
+The spectrum of the posterior precision matrix for the \(3 \times 3\) filter has two distinct clusters, with the first cluster being over an order of magnitude larger than the second.
+The first cluster (which dominates the spectrum) consists of nine eigenvalues, and so the precision has an effective rank close to 9.
+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 distribution of eigenvalues suggests that the effective rank of the precision is over 50, rather than 25 which would be plausible when considering that a \(3 \times 3\) filter has rank 9.
+We do not yet understand this phenomenon and do not have a good explanation for why the posterior precision matrix for the \(5 \times 5\) filter has such a large effective rank and a smooth spectrum, unlike for the \(3 \times 3\) filter.
\begin{marginfigure}
\centering
\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.}
+ \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, which does not have any distinct clusters. This suggests that the effective rank of the precision is over 50.}
\label{fig:spectrum-mnist}
\end{marginfigure}
\section{Discussion}
The choice of prior precision \(\alpha\) is non-trivial.
-Furthermore, since the prior precision is the lowest eigenvalue of the posterior precision matrix, we
+Furthermore, since the prior precision is the lowest eigenvalue of the posterior precision matrix, we must choose \(\alpha\) to be large enough to ensure that the posterior precision matrix is positive definite (i.e., the prior precision must counteract small numerical issues which may lead to some negative eigenvalues).
+Additionally, the condition number of the posterior precision matrix is proportional to \(\alpha^{-1}\), so \(\alpha\) should be large enough to ensure that the posterior precision matrix is well-conditioned.
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.
+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.
-\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
@@ -3,6 +3,10 @@
\label{ch:laplace}
% http://saul.cpsc.ucalgary.ca/pmwiki.php/GradTips/GradTipsLiteratureReview
+In this chapter, we discuss training and inference in deep learning from a Bayesian perspective.
+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]
\label{sec:optimisation-frequentist}
@@ -132,17 +136,24 @@ However, since the posterior distribution is proportional to the product of the
\begin{align}\label{eq:unnormalised posterior}
p(\bm \theta \given \bm y) \propto p(\bm y \given \bm \theta)\,p(\bm \theta) =: \tilde{p}(\bm \theta \given \bm y) .
\end{align}
+\marginnote{SHOULD I TALK ABOUT HOW WE MAKE PREDICTIONS? I.E., HOW WE COMPUTE THE POSTERIOR PREDICTIVE BY MC SAMPLING FROM THE POSTERIOR?}
+% From here, we can make predictions in a Bayesian setting by computing the posterior predictive distribution
+% \begin{align}\label{eq:posterior predictive}
+% p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
+% \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y), \nonumber
+% \end{align}
+% where \(y^*\) is the new data point, \(S\) is the number of Monte Carlo posterior samples \(\bm \theta_s\).\sidenote{DOES THIS HOLD IN REGRESSION?}
Since we want to determine the parameters which maximise the posterior distribution, we can typically use the same optimisation techniques as for frequentist learning, but with the negative log-posterior loss instead of the negative log-likelihood.
Furthermore, since the unnormalised posterior from \cref{eq:unnormalised posterior} is proportional to the likelihood and the prior, the exact posterior itself need not be tractable.
Similarly to the negative log-likelihood, we can factorise the negative log-posterior loss as per
\begin{align}\label{eq:negative-log-posterior}
- \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta) \nonumber
- \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}.
+ \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
+ \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}. \nonumber
\end{align}
Here the regularisation term is given by the prior distribution \(p(\bm\theta)\).
Due to the black-box nature of neural networks, it is often difficult to define a prior distribution over the parameters using previous knowledge.
-Because of this, the prior is often chosen to be a Gaussian distribution with zero mean and a diagonal covariance matrix with some precision \(\alpha\), i.e., \(p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I)\).\sidenote{In the frequentist setting, this type of regularisation is known as weight decay or \(L_2\)-regularisation.}
+Because of this, the prior is often chosen to be a Gaussian distribution with zero mean and a diagonal covariance matrix with some precision \(\alpha\), i.e., \(p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I)\).\sidenote{In the frequentist setting, this type of regularisation is known as weight decay or \(L_2\)-regularisation. Note that \(- 1 / N \log p(\bm y \given \bm \theta) - \gamma / 2 \norm{\bm \theta}^2\) is the typical formulation in this setting, where \(\gamma\) is the regularisation strength and is usually within the range of \(10^{-2}\) to \(10^{-4}\). In our setting, we are instead computing the log posterior \(-\log p(\bm y \given \bm \theta) - \alpha / 2 \norm{\bm \theta}^2\), and so we must multiply the regularisation strength by \(N\) to obtain \(\alpha\).}
The likelihood term is defined as in the frequentist setting (see \cref{sec:optimisation-frequentist}).
However, the posterior distribution itself is typically intractable, so approximations to the posterior are used instead.
@@ -173,9 +184,9 @@ Since the Taylor expansion is performed around the MAP solution, the first order
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}.}
\begin{align}\label{eq:laplace-normal}
- p(\bm\theta \given \bm y) \approx{} & \sqrt{\frac{\det (\bm \Lambda)}{(2 \pi)^D}} \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right) \nonumber
+ 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)
% \\={}& \sqrt{\frac{\det (- \bm H)}{(2 \pi)^D}} \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_0)\T (- \bm H) (\bm \theta - \bm \theta_0)\right) \nonumber
- \\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}).
+ \\={}& \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.
This method is known as post-hoc Laplace.
@@ -186,7 +197,7 @@ We approach the latter, as well as post-hoc Laplace, in this work.
We have now shown that approximating the log-posterior with a second degree Taylor expansion around the \(\bm\theta_{\textsc{map}}\) corresponds to approximating the posterior with a Gaussian distribution given by \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\) where \(\bm \Lambda = - \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_{\textsc{map}}}\).
However, this approximation only holds when the posterior precision \(\bm \Lambda\) is positive definite.
For BNNs, this is not guaranteed, which means that the normal distribution assumed by the Laplace approximation may not be valid.
-In the next sections, we will go into the computation of the Hessian for the negative log-posterior loss function from \cref{sec:bayesian-deep-learning} and show how the Generalised Gauss-Newton approximation can be used to obtain a positive definite Hessian approximation, even when the actual Hessian itself is not positive definite.
+In 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.
\subsection{The Hessian}
\label{sec:hessian}
@@ -204,9 +215,9 @@ In this case, the Hessian is generally dominated by the block-diagonal~\cite{mar
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
\begin{align}\label{eq:negative-log-posterior-hessian}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right) \nonumber
+ \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
\\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
- \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta).
+ \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
\end{align}
In this construction, the first term is the Hessian of the negative log-likelihood, i.e., the Hessian of the loss function we are interested in.
For the common prior given by a zero-mean Gaussian \(\bm \theta \sim \normal(\bm 0, \alpha^{-1} \identity)\), the log-prior regularisation term is given by \(\nabla^2_{\bm\theta} \log p(\bm \theta) = - \alpha \identity\).
@@ -221,9 +232,9 @@ From the equation of the Hessian we can apply the chain rule twice and the produ
\begin{align}\label{eq:ggn-whole-hessian}
\nabla^2_{\bm \theta} \mathcal{L}
={} & J_{\bm \theta}(\nabla_{\bm \theta} \mathcal{L})
- = J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot J_{\bm \theta} f\right) \nonumber
+ = J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot J_{\bm \theta} f\right)
\\ ={}& J_{\bm \theta} (J_{\bm \theta} f) \cdot \nabla_f \mathcal{L} + J_{\bm \theta} \left(\nabla_f \mathcal{L}\right) \cdot J_{\bm \theta} f \nonumber
- \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + J_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot J_{\bm \theta} f,
+ \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + J_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot J_{\bm \theta} f, \nonumber
\end{align}
where \(\nabla_f\) is the gradient with respect to \(f(\bm x, \bm \theta)\) and \(J_{\bm \theta}\) is the Jacobian with respect to \(\bm \theta\).
More precisely, we obtain
@@ -255,7 +266,7 @@ This means that the GGN approximation is an approximation to the Hessian of the
For a given negative log-likelihood loss \(\mathcal{L}\), we thus have the GGN approximation of the Hessian
\begin{align}\label{eq:negative-log-likelihood-ggn}
\nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta)
- \\ \approx{}& \sum_{i=1}^N {\underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i\T}}\T \underbrace{\left(\nabla^2_f \mathcal{L}(y_i, f(\bm x_i, \bm \theta))\right)}_{\bm H_i} \underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i}.
+ \\ \approx{}& \sum_{i=1}^N {\underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i\T}}\T \underbrace{\left(\nabla^2_f \mathcal{L}(y_i, f(\bm x_i, \bm \theta))\right)}_{\bm H_i} \underbrace{J_{\bm \theta} f(\bm x_i, \bm \theta)}_{\bm J_i}. \nonumber
\end{align}
For a normal prior with precision \(\alpha\) on \(\bm \theta\), we obtain the GGN approximation of the posterior precision
\begin{align}\label{eq:laplace-precision-ggn}
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -2,7 +2,12 @@
\chapter{Sampling}
\label{ch:sampling}
-\section{Motivation}
+We will now discuss how to perform inference by sampling from the Laplace posterior distribution.
+We first give an overview of how we typically sample from the Laplace posterior and discuss the memory constraints of this sampling method when using the full posterior precision matrix.
+In \cref{sec:ciq}, we propose an algorithm for sampling from the Laplace posterior that doesn't require storing the full precision matrix, allowing us to scale sampling to larger models.
+Finally, we discuss how to mitigate some of the issues of our method (\cref{sec:preconditioning}) and propose a practical technique for determining whether sampling was successful (\cref{sec:sampling-evaluation}).
+
+\section{Sampling from the Laplace Posterior}[Posterior Sampling]
In many methods, such as the Laplace approximation (\cref{sec:laplace}), we sample from a multivariate normal distribution parametrised by a mean vector \(\bm \mu\) and a covariance matrix \(\bm \Sigma\).
To sample from a Gaussian distribution, you can apply the reparametrisation trick to compute a vector of samples as
@@ -20,10 +25,10 @@ Because of this, the covariance matrix is given by the inverse of this Hessian,
where \(\bm \Lambda = \nabla^2_{\bm\theta} \log p(\bm \theta \given \bm y)\).
In \cref{eq:sample-reparam}, the square root of a matrix yields another matrix which, when multiplied by itself, is equal to the original matrix.
However, since a Gaussian distribution is invariant to rotation, in this special case, we just need to obtain a matrix which returns the original matrix \emph{up to a rotation}.\sidenote{That is, if we have a vector of standard normal samples \(\bm \varepsilon \sim \normal(\bm 0, \identity)\), then, for an orthogonal matrix \(\bm R\) (i.e., a rotation matrix), \(\bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\). We then have that we can sample \emph{non}-standard normal samples as \(\bm \Lambda^{-1/2} \bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\).}
-We thus need to find a method which computes the inverse square root of a high-dimensional matrix up to a rotation.
+We thus need to find a method which computes the inverse square root of the full \(D \times D\) posterior precision matrix up to a rotation, where \(D\) is the number of model parameters.
Some methods exist which compute the inverse square root of a matrix up to a rotation, such as the Cholesky decomposition.
-However, these methods are not suitable for large matrices, as they require the computation of the full matrix, which is not feasible for large matrices.
+However, these methods are not suitable for large matrices, as they require the instantiation of the full precision matrix, which is not feasible for large matrices.
There are also methods which can compute the inverse square root of a low-rank matrix that is constructed as the inner product of two rectangular matrices by performing the inverse square root on the outer product (which is low-dimensional for low-rank problems), such as the Woodbury matrix identity.
However, the precision matrix we are interested in is not itself low-rank but a sum of low-rank matrices, and thus this method is not suitable for our problem.
This is discussed in more detail in \cref{sec:laplace,sec:preconditioning}.
@@ -119,7 +124,7 @@ To estimate these values, we can use the Lanczos algorithm, a Krylov subspace me
We can then apply Cauchy's integral formula to the function \(f(\bm K) = \bm K^{-1/2}\), apply a change of variable, and then approximate this contour integral using the quadrature rule with \(Q\) quadrature points (see \cref{fig:ciq}), as per
\begin{align}
\bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau \label{eq:ciq-integral-tau}
- \\ \approx{}& \frac{1}{2 \pi i} \sum_{q=1}^Q \Tilde{w}_q \tau_q^{-1/2} \left( \tau_q \identity - \bm K \right)^{-1}, \label{eq:ciq-quadrature-tau}
+ \\ \approx{}& \frac{1}{2 \pi i} \sum_{q=1}^Q \Tilde{w}_q \tau_q^{-1/2} \left( \tau_q \identity - \bm K \right)^{-1}, \nonumber %\label{eq:ciq-quadrature-tau}
\end{align}
where \(\tau_1, \ldots, \tau_Q\) are the sampled quadrature points and \(\Tilde{w}_1, \ldots, \Tilde{w}_Q\) are the quadrature weights.
Since we will be sampling a fixed number of quadrature points from the circle, to improve the closeness of the approximation of the integral, it would be beneficial to choose a closed circle which intersects the real axis \emph{at} the minimum and maximum eigenvalues, since this will avoid sampling quadrature points which lie outside the spectrum of \(\bm K\).
@@ -134,11 +139,11 @@ Since the precision matrix (which will be computed from the Laplace approximatio
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,eq:ciq-quadrature-tau} as
-\begin{align}
- \bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau \nonumber
- \\ ={}& \frac{1}{\pi i} \oint_{\Gamma_\sigma} \left( \sigma^2 \identity - \bm K \right)^{-1}\,d\sigma \label{eq:ciq-integral-sigma}
- \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1}. \label{eq:ciq-quadrature-sigma}
+To then 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}
+ \bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau
+ \\ ={}& \frac{1}{\pi i} \oint_{\Gamma_\sigma} \left( \sigma^2 \identity - \bm K \right)^{-1}\,d\sigma \nonumber
+ \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1}. \nonumber % \label{eq:ciq-quadrature-sigma}
\end{align}
We can now adapt the equation to allow for computation using only matrix-vector products.
From here, the inverse square root vector product is given by
@@ -378,8 +383,8 @@ The pivoted Cholesky factorisation is a method for computing a low-rank approxim
This decomposition has been successfully applied to preconditioning for kernel matrices in Gaussian processes~\cite{gardner2018gpytorch}, but not to preconditioning for the posterior precision of neural networks using the Laplace approximation.
Given the posterior precision matrix \(\bm \Lambda\) from the Laplace approximation, we can compute the pivoted Cholesky factorisation as
\begin{align}
- \bm \Lambda ={} & \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity \nonumber
- \\ \approx{}& \bm L\T \bm L + \alpha \identity,
+ \bm \Lambda ={} & \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity
+ \\ \approx{}& \bm L\T \bm L + \alpha \identity, \nonumber
\end{align}
where \(\bm L \in \reals^{K \times D}\), with \(K\) being the chosen rank of the pivoted Cholesky factorisation and \(D\) being the number of parameters.
We can then use this approximation as a preconditioner by computing its inverse via the Woodbury matrix identity, similarly to \cref{sec:woodbury-preconditioner}, as
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -2,11 +2,14 @@
\chapter{Training}
\label{ch:training}
+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.
+Finally, we discuss the implementation of online Laplace with mini-batching using our upper bound.
+
\section{Maximising the Evidence}
\label{sec:maximising-the-evidence}
-% In Online Laplace, we seek to maximise the evidence, which is the marginal likelihood of the data given the model parameters.
-
As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model is equal to
\begin{align}
p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta)\,p(\bm \theta)}{p(\bm y)},
@@ -21,16 +24,18 @@ For this reason, we typically approximate the posterior by some tractable distri
Since we are integrating over the posterior, maximising the marginal likelihood will seek minima that are flat, since the posterior will be very broad in these regions.
This is a desirable property, since it means that the model will be less likely to overfit the data, both from a theoretical standpoint~\cite{hochreiter1997flat} and in practice~\cite{keskar2016large,jiang2019fantastic,maddox2019simple}.
-\textcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing leave-p-out cross validation for all values \(p = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(p\) folds and across all values of \(p\)).
+To further motivate training on the evidence term, \textcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing leave-p-out cross validation for all values \(p = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(p\) folds and across all values of \(p\)).
This result shows that the marginal likelihood is a good proxy for the true cross-validation performance of a model.
Because of this, it can be used for model selection, since it will choose the model that performs best on average across all permutations of the data.
-Finally, the marginal likelihood can be used to optimise hyperparameters, since it can be used to compute the gradient of the log-marginal with respect to the hyperparameters.
+Additionally, the marginal likelihood can be used to optimise hyperparameters, since it can be used to compute the gradient of the log-marginal with respect to the hyperparameters.
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}}\).
-We thus obtain the approximate Laplace posterior
+As in \cref{eq:laplace-taylor}, we have the approximate Laplace log-posterior
\begin{align}\label{eq:laplace-posterior}
\log p(\bm \theta \given \bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) \cdot (\bm \theta - \bm \theta_{\textsc{map}})
\\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
@@ -95,7 +100,7 @@ However, we can compute an upper bound on the log-determinant of the posterior p
\section{The Determinant Bound}
\label{sec:determinant-bound}
-We wish to compute an upper bound on the log-determinant of the Hessian of the negative log-posterior \(\bm \Lambda \in \reals^{D \times D}\) such that we have \(\log \det \bm \Lambda \leq B_D(\mu_1, \mu_2, \beta)\).
+We want to compute an upper bound on the log-determinant of the Hessian of the negative log-posterior \(\bm \Lambda \in \reals^{D \times D}\) such that we have \(\log \det \bm \Lambda \leq B_D(\mu_1, \mu_2, \beta)\).
\textcite{bai1996bounds} provide one such bound,
\begin{align}\label{eq:bai-golub-bound}
B_D(\mu_1, \mu_2, \beta) \coloneqq{} & \begin{pmatrix} \log \beta & \log t \end{pmatrix}
@@ -120,7 +125,7 @@ The former can then be bounded by \cref{eq:bai-golub-bound} with an upper bound
We then compute the trace of this term as
\begin{align}\label{eq:laplace-trace-bound}
\Tr{\bm \Lambda} ={} & \sum_{k=1}^{K} (\lambda_k + \alpha) + (D - K) \alpha
- \\ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha \nonumber
+ \\ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha
\\ ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha, \nonumber
\end{align}
and for the square trace we have
diff --git a/figures/overview.tex b/figures/overview.tex
@@ -1,14 +1,14 @@
\begin{tikzpicture}
[
font=\sffamily,
- box/.style = {draw, rounded corners, minimum width=22mm, minimum height=5mm, align=center},
+ box/.style = {rounded corners, minimum width=22mm, minimum height=5mm, align=center},
]
% \definecolor{color1}{rgb}{0.86, 0.82, 1.0}
% \definecolor{color2}{rgb}{0.97, 0.96, 1.0}
\definecolor{color1}{rgb}{0.74, 0.83, 0.9}
\definecolor{color2}{rgb}{0.94, 0.97, 1.0}
\definecolor{kaogreen}{HTML}{91C889}
- \node (nn) [box,fill=color1] {Neural Network};
+ \node (nn) [box,fill=color1,minimum height=1.0cm,minimum width=3cm] {Neural Network};
\node (max) [below=of nn,yshift=0.5cm] {\(\argmax_{\bm \theta}\)};
\node (posterior) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=-2cm,yshift=0.5cm] {\(\log p(\bm \theta \given \bm y)\)};
\node (marginal) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=2cm,yshift=0.5cm] {\(\overbrace{\log p(\bm \theta \given \bm y) \quad \text{\colorbox{kaogreen}{$\log \det \bm \Lambda$}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
@@ -16,7 +16,7 @@
\draw[dashed,->] (max) to [bend right] (posterior);
\draw[dashed,->] (max) to [bend left] (marginal);
\begin{scope}[on background layer]
- \node (training) [box,fill=color1,minimum width=8.05cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
+ \node (training) [box,fill=color1,minimum width=8.05cm,inner sep=0.25cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
\end{scope}
\node [above=0cm and 0cm of training.north west,anchor=south west] {Training};
\node [above=0cm and 0cm of posterior.north west,anchor=south west] {{Post-Hoc}};
@@ -26,9 +26,9 @@
\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] {\(\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=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}\)};
\begin{scope}[on background layer]
- \node (inference) [box,fill=color1,minimum width=5.05cm,fit={(sample)(samplingeq)}] {}; % minimum width=8.05cm,minimum height=3.2cm,
+ \node (inference) [box,fill=color1,inner sep=0.25cm,fit={(sample)(samplingeq)}] {}; % minimum width=8.05cm,minimum height=3.2cm,
\end{scope}
\node [above=0cm and 0cm of inference.north west,anchor=south west] {Inference};
\draw[dashed,->] (map) to (inference);
diff --git a/figures/sin_posterior_predictive_bad.pdf b/figures/sin_posterior_predictive_bad.pdf
Binary files differ.
diff --git a/figures/sin_posterior_predictive_bothalphas.pdf b/figures/sin_posterior_predictive_bothalphas.pdf
Binary files differ.
diff --git a/figures/sin_posterior_predictive_bothalphas_clear.pdf b/figures/sin_posterior_predictive_bothalphas_clear.pdf
Binary files differ.
diff --git a/figures/sin_posterior_predictive_joint_bad.pdf b/figures/sin_posterior_predictive_joint_bad.pdf
Binary files differ.
diff --git a/figures/sin_posterior_predictive_samples_bad.pdf b/figures/sin_posterior_predictive_samples_bad.pdf
Binary files differ.
diff --git a/main.tex b/main.tex
@@ -249,7 +249,7 @@
\input{chapters/laplace.tex}
\input{chapters/training.tex}
\input{chapters/sampling.tex}
-\input{chapters/boring_experiments.tex}
+\input{chapters/ablation_experiments.tex}
\input{chapters/experiments.tex}
\input{chapters/conclusion.tex}
