commit b7f4c02547354368267ecbc0b8eb47d8b03ad215
parent f731fc765fc8d1175bd07f9e738d3a1c368b2aab
Author: Silas Brack <s174433@student.dtu.dk>
Date: Mon, 6 Feb 2023 19:13:35 +0100
Getting close to a version I can send to Soren.
Diffstat:
3 files changed, 19 insertions(+), 51 deletions(-)
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -13,8 +13,9 @@ For all but toy problems, however, the full Laplace approximation is intractable
In this work, we propose a method for computing the Laplace approximation using only Hessian-vector products, which allows us to perform marginal training and posterior sampling using the full Laplace approximation without storing the entire Hessian.
To accomplish this, we show that it is possible to estimate the Hessian's log-determinant and inverse square root using only Hessian-vector products.
+To overcome conditioning issues in the Hessian matrix, we propose several potential preconditioners which can be used to improve convergence during sampling.
Next, we suggest a technique for evaluating the quality of posterior samples without instantiating nor inverting the Hessian based on results from traditional statistics.
-To implement this method, we use JAX, a library for automatic differentiation which enables us to efficiently perform Hessian-vector products without explicitly storing the entire Hessian.
+To implement these methods, we use JAX, a library for automatic differentiation which enables us to efficiently perform Hessian-vector products without explicitly storing the entire Hessian.
We perform the full Laplace approximation on both a sine function and MNIST using our method.
This consists two steps, (a) training a neural network by maximising either its posterior probability or marginal likelihood and (b) sampling from this posterior.
@@ -24,7 +25,7 @@ We then benchmark the performance of Hessian-vector products and estimate that o
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.
+I FEEL LIKE I SHOULD MAKE IT MORE CONCISE AND FOCUS MORE 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.
diff --git a/chapters/appendix.tex b/chapters/appendix.tex
@@ -1,5 +1,5 @@
\chapter{Algorithms}
-% \setchapterstyle{lines}
+\setchapterstyle{lines}
% \setchapterpreamble[u]{\margintoc}
\begin{algorithm}[h!]
@@ -53,7 +53,7 @@
\end{algorithm}
\chapter{Additional Results}
-% \setchapterstyle{lines}
+\setchapterstyle{lines}
\begin{table*}[h]
\centering
@@ -66,7 +66,7 @@
\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.}
+ \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for manual matrix-vector, Hessian-vector, and GGN-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.}
\end{subfigure}
\hspace{0.5cm}
\begin{subfigure}[h]{0.45\textwidth}
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -51,42 +51,32 @@ This last point is important, since it also means that the Laplace approximation
In the Laplace approximation, to compute the posterior distribution of the model parameters, we need to compute the Hessian of the loss with respect to the model parameters.
Typically, this is done by storing the Hessian matrix explicitly and computing the inverse of this matrix.
Since the Hessian matrix is a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters, this is intractable for large models, which can contain millions or billions of parameters, thus limiting the applicability of full Laplace outside of toy problems of a few thousand parameters.
-To overcome this, we can use a crude approximation, such as only storing the Hessian matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.%~\cite{}.
+To overcome this, we can use a crude approximation, such as only storing the Hessian matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.
Many of these Hessian approximations have successfully been applied to the Laplace approximation~\cite{ritter2018scalable,ritter2018online}, though are limited by the quality of the approximation.
To overcome this, we propose a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
\section{Large-Scale Laplace}
We propose an implementation of the full Laplace approximation which doesn't require instantiation of the full posterior precision matrix.
-To perform inference using the Laplace approximation we need to sample from the posterior distribution, which is approximated by a normal distribution with a precision given as the Hessian of the loss with respect to the model parameters.
+First, we would like to learn a posterior distribution over the model parameters by maximising either the marginal likelihood or the posterior distribution.
+Then, to perform inference using the Laplace approximation, we need to sample from the posterior distribution, which is approximated by a normal distribution with a precision given as the Hessian of the loss with respect to the model parameters.
The standard procedure for sampling thus requires computing the covariance matrix as the inverse of this precision matrix, and then taking the square root of the covariance matrix.
+Then, we can sample from this posterior distribution to make predictions.
+To train using the negative log marginal likelihood loss for the Laplace approximation, we must compute the log determinant of the Hessian.
However, for large models, which can contain millions or billions of parameters, we cannot instantiate the Hessian matrix, for reasons mentioned above.
Instead, we would like to access the inverse square root of the Hessian without explicitly storing it.
-Furthermore, we would like to perform training using the negative log marginal likelihood loss, which, for the Laplace approximation, requires computing the log determinant of the Hessian.
-To perform these computations we need to approximate these quantities without explicitly storing the Hessian matrix.
-One way to do this is to instead compute Hessian-vector products, which are matrix-vector products with the Hessian matrix.
-This would allow us to only store Hessian-vector products instead of the Hessian matrix itself, thus using significantly less memory.
-Fortunately, it is possible to compute Hessian-vector products while only storing the implicit Hessian computational graph.
-% \textcite{pleiss2020fast} propose a contour integral quadrature (CIQ), a method which approximates the product of a square root of a matrix with a vector without requiring this matrix to be stored explicitly.
-% Instead of instantiating the matrix, they approximate this square-root matrix-vector product by performing repeated matrix-vector products with this matrix, and then approximating the result using a quadrature rule.
-% While they use it for computing the square root of the kernel matrix in Gaussian processes, we use it for computing the inverse square root of the posterior precision in Bayesian neural networks using the Laplace approximation.
-% Since, to sample from the Laplace posterior distribution we need to compute the square root of the covariance matrix, we need be able to compute Hessian-vector products (i.e., matrix-vector products with the Hessian matrix) without explicitly storing the Hessian matrix.
-% Fortunately, it is possible to compute Hessian-vector products while only storing the implicit Hessian computational graph. (MAYBE SHOULD REMOVE THIS PARAGRAPH, TOO MUCH INTO METHODS)
+Our method allows us approximate the Hessian's inverse square root and log determinant using only implicit Hessian-vector products (HVPs).
+These HVPs compute products of the Hessian matrix with a vector without explicitly storing the Hessian matrix itself.
+Since the HVP outputs a vector, the memory requirements are linear in the number of parameters, rather than quadratic, allowing us to perform the full Laplace approximation with significantly larger models.
+We use JAX~\cite{jax2018github}, a Python library for GPU-accelerated automatic differentiation and computational graph compilation, to efficiently compute Hessian-vector while only storing the implicit Hessian computational graph.
+JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the Hessian-vector product computation into efficient machine code.
-% \section{JAX and XLA}
-
-JAX~\cite{jax2018github} is a Python library for GPU-accelerated automatic differentiation and compilation of computational graphs into efficient machine code using XLA~\cite{sabne2020xla}, a compiler for linear algebra.
-JAX enables us to efficiently compute Hessian-vector products without explicitly storing the Hessian matrix by compiling the Hessian-vector product computation with XLA.
-
-% \section{Research Objectives}
-
-We would thus like to develop a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
+We have thus developed a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
To do this, we will use JAX and XLA to compute Hessian-vector products without explicitly storing the Hessian matrix.
-We will then use these Hessian-vector products to compute the log determinant and inverse square root of the Hessian.
-These quantities will be used to perform training and inference using the Laplace approximation.
+We will then exploit these Hessian-vector products to compute the log determinant and inverse square root of the Hessian.
+These quantities enable us to perform training and inference using the Laplace approximation.
An overview of the Laplace approximation and our contributions is shown in \cref{fig:overview}.
-% SHOULD THIS BE LONGER? TALK ABOUT EVALUATION SOMEWHERE?
\begin{figure}
\centering
@@ -94,26 +84,3 @@ An overview of the Laplace approximation and our contributions is shown in \cref
\caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Hessian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the Hessian. Our contributions allow us to forgo the typical crude approximations of the Hessian by performing the full Laplace approximation without explicitly storing the Hessian matrix.}
\label{fig:overview}
\end{figure}
-% \begin{marginfigure}
-% \centering
-% \resizebox{\linewidth}{!}{\input{figures/overview.tex}}
-% \caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight our contributions in green, i.e., the computations which we will approximate using only Hessian-vector products.}
-% \label{fig:overview}
-% \end{marginfigure}
-
-% This should be discussed in a bit more detail in a thesis, since there
-% are certain objectives discussed in the project plan.
-% In an article, a quick sentence which summarises the rest of the
-% introduction and states succinctly exactly which problem will be approached
-% and which approach will be made should suffice.
-
-% In this project, we strive to develop an effortless Bayesian method.
-% More specifically, this project investigates the connections between the
-% Laplace approximation and the approximate second-order derivatives used
-% in modern optimizers, such as Adam. We use a sampling-based training
-% procedure, where a sample is first drawn from a Gaussian
-% weight-posterior, a gradient step is performed on this sampled neural
-% network, and the variance of the weight-posterior is updated with an
-% approximate Hessian. Approximating the Hessian is the most time-consuming and painful-to-engineer step of this training procedure. In
-% this project, we tap into the potential of modern machine learning
-% frameworks to efficiently approximate the Hessian.
