commit 1e8f4f61931768f795eb54d71d844eb495ab3e05
parent 1f230ab33424712515fca44f99ddf6693a750abe
Author: Silas Brack <s174433@student.dtu.dk>
Date: Sat, 4 Feb 2023 19:20:07 +0100
Some fixes from feedback.
Diffstat:
7 files changed, 81 insertions(+), 49 deletions(-)
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -3,7 +3,7 @@
% https://www.student.unsw.edu.au/writing-abstracts-honours-theses
-Bayesian methods have been used for decades, but have recently seen an increased interest in the domain of deep learning, since they promise to provide a principled way to quantify uncertainty in neural networks.
+Bayesian methods promise to provide a principled way to quantify uncertainty in neural networks.
These methods rely on Bayesian inference to learn the posterior distribution of the model parameters given the data.
In doing so, they allow for uncertainty estimates, which is important for many applications in machine learning, such as those involving safety-critical decisions which rely on risk assessment and interpretability.
However, Bayesian inference is often computationally intractable, and as such, approximate methods which compromise performance for feasibility are used.
@@ -11,8 +11,9 @@ The Laplace approximation is one such method which enables post-hoc estimation o
Additionally, since the Laplace approximation deduces uncertainty rather than learning it, it has many advantages over other Bayesian methods that attempt to learn variance.
Due to its construction, the Laplace approximation allows for training by maximising the marginal likelihood, which facilitates model selection and hyperparameter optimisation.
However, the Laplace approximation typically relies on crude approximations to the posterior precision to make it computationally feasible.
-% Often, strong assumptions are made about the Hessian's structure, but these assumptions are not always accurate.
-Conventional full Laplace uses the full posterior precision, which cannot be instantiated for all but toy problems, since it scales quadratically with the number of parameters.
+When we do not factor this precision matrix into more practical approximations, the full posterior precision must be instantiated, which is intractable for all but toy problems, since this matrix scales quadratically with the number of parameters.
+When we perform the Laplace approximation using this full posterior precision, we call this the full Laplace approximation.
+% Conventional full Laplace uses the full posterior precision matrix, which cannot be instantiated for all but toy problems, since it scales quadratically with the number of parameters.
In this work, we propose a method for performing the Laplace approximation using only Hessian-vector products, which allows us to perform marginal training and inference using the full Laplace approximation without storing the entire Hessian.
To accomplish this, we estimate the Hessian's log-determinant and inverse square root using only Hessian-vector products.
We also suggest a technique for evaluating the quality of posterior samples without instantiating nor inverting the Hessian based on results from traditional statistics.
@@ -20,7 +21,6 @@ To implement this method, we use JAX, a library for automatic differentiation wh
We show that our method is able to perform the full Laplace approximation on both a sine function and MNIST and that training the marginal likelihood using our method is feasible and produces meaningful results.
This is the first method for performing the Laplace approximation while only accessing the Hessian implicitly.
Future work will focus on carrying out hyperparameter optimisation and large-scale experiments using this method.
-MEH, I THINK I SPEND TOO MUCH TIME ON BACKGROUND AND INTRO.
% \chapter*{Resum\'e}
% \addcontentsline{toc}{chapter}{Resum\'e}
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -11,8 +11,8 @@
% Why this problem is important.
Bayesian methods allow for estimates of uncertainty which enable more efficient usage of data (e.g., via active learning) and avoid overfitting.
-They accomplish this by modelling a distribution on the model parameters, which are then used to make predictions.
-These predictions will then have uncertainty estimates, thus improving model interpretability and enabling the assessment of model predictive confidence.
+They accomplish this by modelling a distribution on the model parameters, which is then used to make predictions.
+From here we can sample from this distribution to make predictions, thereby yielding uncertainty estimates, improving model interpretability and enabling the assessment of model predictive confidence.
These methods have been successfully applied to a wide range of classical problems in statistics, though attempts to apply them to deep learning have had limited success.
For deep learning, the Bayesian approach is often limited by the computational cost of Bayesian inference.
@@ -30,7 +30,7 @@ However, this is computationally infeasible, and so approximate Bayesian methods
\section{Current Methods}
-Currently, approximate Bayesian methods are either expensive to compute (Markov Chain Monte Carlo), are significantly more difficult to implement (such as variational inference), or simply perform poorly and are limited in their Bayesian interpretation (MC dropout).
+Currently, approximate Bayesian methods are either expensive to compute (Markov Chain Monte Carlo), are significantly more difficult to implement (such as variational inference), or simply perform poorly and are limited in their Bayesian interpretation (Monte Carlo dropout, deep ensembles, SWAG, etc.).
Of these, MC dropout and Markov Chain Monte Carlo are the most widely used, though both have significant drawbacks.
As such, there is demand for a method which exhibits the same computational cost as the optimization algorithms for deterministic neural networks while providing accurate posterior approximations and working out-of-the-box for any given architecture.
@@ -48,15 +48,15 @@ To overcome this, we propose a method which performs the full Laplace approximat
\section{Large-Scale Laplace}
-We are seeking an implementation of the full Laplace approximation which doesn't require instantiation of the full posterior precision matrix.
+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.
-Sampling thus requires computing the covariance matrix as the inverse of the Hessian of the loss with respect to the model parameters, and then taking the square root of this covariance matrix.
+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.
However, for large models, which can contain millions or billions of parameters, we cannot instantiate the Hessian matrix, for reasons mentioned above.
-Instead, we would like to compute the inverse square root of the Hessian without explicitly storing it.
-% , and instead only storing the output of Hessian-vector products.
+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.
@@ -76,8 +76,16 @@ We would thus like to develop a method which performs the full Laplace approxima
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.
+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
+ \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{figure}
+
% This should be discussed in a bit more detail in a thesis, since there
% are certain objectives discussed in the project plan.
% In an article, a quick sentence which summarises the rest of the
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -41,10 +41,10 @@ The gradient \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t)\) can b
Using the chain rule on the composition of the loss into the loss as a function of the model output \(\mathcal{L}\) and the model output as a function of the parameters \(f\) such that \(\mathcal{L}(\bm \theta_t, \bm x_t, y_t) = \mathcal{L}(y_t, f(\bm x_t, \bm \theta_t))\),we compute the gradient as
%
\begin{align}\label{eq:sgd-gradient}
- \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t) ={} & \nabla_{\bm\theta} \log p(y_t \given \bm \theta) \nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t),
+ \nabla_{f} \mathcal{L}(\bm\theta_t, \bm x_t, y_t) ={} & \nabla_{\bm\theta} \log p(y_t \given \bm \theta) \nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t),
\end{align}
%
-where \(\nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)\) is the gradient of the neural network output with respect to its parameters at iteration \(t\).
+where \(\nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)\) is the gradient of the neural network output with respect to its parameters at iteration \(t\) and \(\nabla_{\bm\theta} \log p(y_t \given \bm \theta)\) is the gradient of the loss function with respect to the model output.
Note that in this case, we perform one gradient update for each data point, which is why we iterate over the data points and update the parameters in the same step \(t\).
Thus, SGD only requires the gradient of the loss function with respect to the parameters of the neural network, which can be computed efficiently using automatic differentiation for a similar cost to a single forward pass through the network.
This makes SGD a popular choice for training neural networks.
@@ -100,16 +100,16 @@ The update rule for Adam is thus given by
\begin{align}
\bm\theta_{t+1} ={} & \bm\theta_t - \eta_t \bm m_t \odot \frac{1}{\sqrt{\bm v_t} + \varepsilon},
\\ \bm m_t ={}& (1 - \beta_1^t) \left[ \beta_1 \bm m_{t-1} + (1 - \beta_1) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
- \\ \bm v_t ={}& (1 - \beta_2^t) \left[ \beta_2 \bm v_{t-1} + (1 - \beta_2) \left( \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right)^2 \right],
+ \\ \bm v_t ={}& (1 - \beta_2^t) \left[ \beta_2 \bm v_{t-1} + (1 - \beta_2) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \odot \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
\end{align}
where \(\odot\) denotes the Hadamard product, \(\varepsilon\) is a small constant to prevent division by zero, and \(\bm m_t\) and \(\bm v_t\) are the first and second moments of the gradient respectively.
The momentum term~\cite{polyak1964some} is inspired by physics, where the momentum of an object is the product of its mass and velocity.
In Adam, the first moment of the gradient is used to compute the momentum of the parameters, increasing the convergence speed of the optimisation~\cite{sutskever2013importance}.
The second moment of the gradient is also known as the variance of the gradient.
-Adam approximates the second moment by the diagonal of the Fisher information matrix.
-This is a common approximation in practice, as the diagonal of the Fisher information matrix is easy to compute.
-The Fisher information matrix is a discussed further in \cref{sec:hessian}.
+Adam approximates the second moment by the diagonal of the Fisher information matrix where the loss and model function are assumed to be linear.
+This is a common approximation in practice, as this approximation is easy to compute.
+The Fisher information matrix is a discussed further in \cref{sec:ggn}.
The use of the second moment is inspired by optimisation methods such as Newton's method and the natural gradient~\cite{rattray1998natural}, which use the second derivative of the loss function to compute the optimal step size by scaling the step by the inverse of the curvature of the loss function.
These optimisation methods all identify a single set of parameters as the optimal solution.
In the next section, we will discuss how to learn a distribution over parameters instead.
@@ -141,7 +141,7 @@ Similarly to the negative log-likelihood, we can factorise the negative log-post
\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)\).
+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.}
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.
@@ -199,7 +199,6 @@ We are interested in the Hessian of the loss function \(\mathcal{L}(\bm\theta)\)
For a neural network with \(D\) parameters, the Hessian is therefore a \(D \times D\) square matrix.
Furthermore, if all the neural network's second partial derivatives are continuous, then the Hessian is symmetric.
In this case, the Hessian is generally dominated by the block-diagonal~\cite{martens2015optimizing}.
-However, the Hessian will only be positive definite if the loss function is a convex function of the model parameters \(\bm \theta\).
Typically, the loss function is considered as the negative log-likelihood or negative log-posterior.
In this framing, if considering, e.g., the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we get
@@ -213,7 +212,8 @@ For the common prior given by a zero-mean Gaussian \(\bm \theta \sim \normal(\bm
However, as we discussed previously, the Hessian of the negative log-likelihood is not always positive definite.
In the next section, we will discuss how to obtain a positive definite Hessian approximation.
-\subsection[The Generalised Gauss-Newton Approximation]{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
+\subsection{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
+\label{sec:ggn}
From the equation of the Hessian we can apply the chain rule twice and the product rule once to obtain a simpler expression,
%
@@ -248,7 +248,7 @@ By combining \cref{eq:ggn-whole-hessian,eq:nn-taylor} we then obtain the general
\nabla^2_{\bm\theta} \mathcal{L}
\approx{} & \nabla_{\bm\theta} f\T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm\theta} f,
\end{align}
-which corresponds to the first term in \cref{eq:ggn-whole-hessian}.
+which corresponds to the second term in \cref{eq:ggn-whole-hessian}.
% In this approximation, we have \(\nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0\).
This means that the GGN approximation is an approximation to the Hessian of the composition of two functions, \(\mathcal{L}(\bm \theta) \equiv (\mathcal{L} \circ f) (\bm \theta)\) where we linearise the inner function \(f\) around a point \(\bm \theta_0\).
For a given negative log-likelihood loss \(\mathcal{L}\), we thus have the GGN approximation of the Hessian
@@ -269,7 +269,7 @@ We then get the simple expression for the GGN approximation of the Hessian for M
\bm H_i \approx & {}\, (y_i - f(\bm x_i, \bm \theta))^2 = 2,
\\ \bm \Lambda \approx&{}\, 2 \sum_{i=1}^N \bm J_i\T \bm J_i + \alpha \identity.
\end{align}
-If the Hessian of the loss function with respect to the model outputs \(\bm H_i\) is positive semidefinite, the GGN approximation is positive semidefinite, since each term in the sum \(\bm J_i\T \bm H_i \bm J_i\) is positive semidefinite.
+If the Hessian of the loss function with respect to the model outputs \(\bm H_i\) is positive semidefinite, the GGN approximation is positive semidefinite, since each term in the sum \(\bm J_i\T \bm H_i \bm J_i\) is positive semidefinite.\sidenote{This is equivalent to the loss function being a convex function of the model parameters \(\bm \theta\).}
Furthermore, this is the case for most common loss functions, such as MSE loss and cross-entropy loss.
For the case of the normal prior, the GGN approximation of the posterior precision will be positive definite, since the eigenvalues of \(\alpha \identity\) are all \(\alpha > 0\) and the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) are all greater than or equal to zero.
@@ -304,7 +304,7 @@ This yields \(g(y) = -\log(y)\), and thus we obtain the GGN approximation of the
={} \nabla_{\bm\theta} \mathcal{L}\T \cdot \nabla_{\bm\theta} \mathcal{L}.
\end{align}
The diagonal of this matrix is given by the element-wise square of the gradient \( \nabla_{\bm\theta} \mathcal{L} \odot \nabla_{\bm\theta} \mathcal{L} \).
-Notice that this is very similar to the second-order term used in the Adam optimiser, which we discussed in \cref{sec:adam}.
+Notice that this is the second-order term used in the Adam optimiser, which we discussed in \cref{sec:adam}.
The choice of application of the generalised Gauss-Newton approximation is not based on having a reduced computational cost---both the Hessian and the GGN approximation require the same number of forward and backward passes.
However, the GGN approximation is better behaved than the exact Hessian, thereby making it more suitable for practical applications.
@@ -330,7 +330,7 @@ Since \(\bm J_i\) is a rectangular matrix of shape \(O \times D\), the rank of e
The rank of the sum of these matrices thus depends on the eigenvectors of each \(\bm J_i\T \bm H_i \bm J_i\).
If we add \(\bm J_n\T \bm H_n \bm J_n\) into the sum \(\sum_{i=1}^{n-1} \bm J_i\T \bm H_i \bm J_i\), then, if any of the eigenvectors of \(\bm J_n\T \bm H_n \bm J_n\) are not contained in the span of the eigenvectors of the sum, then the rank of the sum will increase.
This is because the rank of a matrix is the number of linearly independent columns.
-Thus, non-rigorously, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are between terms in the sum, the closer the largest eigenvalue will be to the sum of the largest eigenvalues of each term.
+Thus, non-rigorously, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are between terms in the sum, the closer the largest eigenvalue will be to the sum of the largest eigenvalues of each term.\sidenote{By ``same ordering'' we mean that they are eigenvectors with respect to the same relative eigenvalue.}
This means that the rank of the sum of the GGN matrices can increase as we add more data points, and will not be greater than \(N \cdot O\).
Furthermore, as we add more data points, the eigenvalues of the sum of the GGN matrices will increase.
Specifically, the largest eigenvalue of the sum of the GGN matrices will be the less than or equal to the sum of the largest eigenvalues of the GGN matrices.
@@ -353,7 +353,7 @@ One limitation of the Laplace approximation is that it can only approximate one
However, most approximate posterior distributions, such as MCMC and variational inference, also struggle to approximate multimodal distributions.
The standard way to solve this problem is to use a multimodal approximation such as a mixture, typically by training multiple models from different initialisations and then combining them in a deep ensemble.
This is also possible with the Laplace approximation~\cite{eschenhagen2021mixtures}.
-However, this is not a scalable solution, as it requires training multiple models.
+However, deep ensembles are quite expensive, as they require training multiple models.
Additionally, it is not clear whether there is a significant benefit to using a multimodal Laplace approximation over a single mode Laplace approximation (i.e., whether generalisation improves).
Another limitation of the Laplace approximation is that it is not scalable to non-trivial problems for the full Laplace approximation without performing crude approximations.
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -2,7 +2,7 @@
\chapter{Sampling}
\label{ch:sampling}
-\section{Introduction}
+\section{Motivation}
In many methods, such as the Laplace approximation (\cref{sec:laplace}), we sample from a multivariate normal distribution parameterised by a mean vector \(\bm \mu\) and a covariance matrix \(\bm \Sigma\).
To sample from a Gaussian distribution, you can apply the reparametrisation trick to compute a vector of samples as
@@ -12,9 +12,10 @@ To sample from a Gaussian distribution, you can apply the reparametrisation tric
where \(\bm \varepsilon_0\) are samples drawn from a standard normal distribution, i.e., \(\bm \varepsilon_0 \sim \normal(\bm 0, \identity)\).
Using the Laplace approximation, we estimate the precision matrix as the Hessian matrix of the loss with respect to the parameters, as discussed in \cref{sec:laplace}.
-Because of this, the covariance matrix is given by the inverse of this Hessian, and we thus sample as per
+Because of this, the covariance matrix is given by the inverse of this Hessian, and we thus sample from the Laplace-approximate posterior distribution as
\begin{align}\label{eq:sample-laplace}
- \bm \varepsilon = \bm \mu + \bm \Lambda^{-1/2} \bm \varepsilon_0,
+ \bm \varepsilon = \bm \theta_{\textsc{map}} + \bm \Lambda^{-1/2} \bm \varepsilon_0,
+ % \bm \varepsilon = \bm \mu + \bm \Lambda^{-1/2} \bm \varepsilon_0,
\end{align}
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.
@@ -92,8 +93,9 @@ it is then also trivial to compute the square root product by multiplying \(\bm
The first step of CIQ invokes Cauchy's integral formula.
Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed circular contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as\marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
\begin{align}\label{eq:cauchy-integral-formula}
- f(\bm K) ={} & \frac{1}{2 \pi i} \oint_\Gamma f(\tau) \left(\tau \identity - \bm K\right)^{-1}\,d\tau.
+ f(\bm K) ={} & \frac{1}{2 \pi i} \oint_\Gamma f(\tau) \left(\tau \identity - \bm K\right)^{-1}\,d\tau,
\end{align}
+where \(\tau\) is a complex variable and \(\oint_\Gamma\) is the contour integral along the contour \(\Gamma\).
Since, in our case, the eigenvalues of \(\bm K\) will all be real-valued, a circle centred on the real axis which encloses the minimum and maximum eigenvalues of \(\bm K\) will suffice.\marginnote{If a real-valued matrix is symmetric, its eigenvalues are real. In our case, \(\bm K\) will be the precision matrix obtained from the Laplace approximation.}
To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).\sidenote{Krylov subspace methods are a class of iterative methods which estimate \(f(\bm K) \bm v\) for some matrix \(\bm K\) and vector \(\bm v\) by using a Krylov subspace. The Krylov subspace is a subspace spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), given by \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\left\{ \bm v, \bm K \bm v, \bm K^2 \bm v, \ldots, \bm K^{j-1} \bm v \right\}\).}
@@ -264,15 +266,6 @@ In this case, a balance must be struck between the accuracy of the preconditione
Furthermore, the inverse of the preconditioner is required for preconditioning, which is often not available.
To approximate the inverse of the preconditioner, we can use the conjugate gradient method, though this method is then sensitive to the condition number of the preconditioner; and if the original matrix \(\bm \Lambda\) is ill-conditioned, then a good preconditioner will also tend to be ill-conditioned.
-% \begin{enumerate}
-% \item Scaling
-% \item Eigenvalues and eigenvectors
-% \item Conditioning of the preconditioner
-% \item Targeting high and low eigenvalues
-% \item Rank, summation of low rank matrices
-% \item Speed
-% \end{enumerate}
-
As such, the condition number of our preconditioner also becomes an issue, since we have two conflicting objectives: to have a preconditioner which is as similar to \(\bm \Lambda\) as possible and to have a preconditioner which we can invert easily.
The first point lends itself toward the use of preconditioners with a similar condition number as \(\bm \Lambda\).
However, to invert an arbitrary positive semi-definite matrix, we will often resort to using the conjugate gradient method which is sentitive to the condition number of the preconditioner.
@@ -305,7 +298,7 @@ If we calculate the GGN approximation of the Hessian by linearizing over the who
\\ \bm P^{-1} ={}& \beta \cdot \nabla \nabla\T + \alpha^{-1} \identity
\\ \beta ={}& \left(\alpha \norm{\nabla}^2 + \norm{\nabla}^4\right)^{-1} - \left(\alpha \norm{\nabla}^2\right)^{-1},
\end{align}
-where \(\nabla\) is the gradient of the loss with regard to the model parameters.\sidenote{DO I NEED TO PROVE THIS? BLAH BLAH BLAH EIGENDECOMPOSITION}
+where \(\nabla\) is the gradient of the loss with regard to the model parameters.
The problem with using this approximation is that the rank of \(\bm P\) is 1.
This means that, best-case scenario, we can reduce the highest eigenvalue of the Hessian.
This occurs when the gradient corresponds to the direction of the highest precision (empirically, this may not be the case).
@@ -343,8 +336,13 @@ However, this preconditioner will be very expensive to compute, since it effecti
\subsection{Using the Woodbury Matrix Identity}[Woodbury]
\label{sec:woodbury-preconditioner}
-The Woodbury matrix identity is a useful identity for computing the inverse of a matrix which is the sum of two matrices of which one is low-rank and the other can easily be inverted.\sidenote{SHOULD I EXPLAIN WOODBURY MATRIX IDENTITY MORE MAYBE?}
-This identity can be used to compute the inverse of the GGN matrix with the prior precision, which is given by
+The Woodbury matrix identity is a useful identity for computing the inverse of a matrix which is the sum of two matrices of which one is low-rank and the other can easily be inverted.
+The identity is given by
+\begin{align}
+ \left( \bm A + \bm B \bm C \bm D \right)^{-1} = \bm A^{-1} - \bm A^{-1} \bm B \left( \bm C^{-1} + \bm D \bm A^{-1} \bm B \right)^{-1} \bm D \bm A^{-1},
+\end{align}
+where \(\bm A\) is a \(n \times n\) matrix and can easily be inverted, \(\bm B\) is a \(n \times m\) matrix, \(\bm C\) is a \(m \times m\) matrix, and \(\bm D\) is a \(m \times n\) matrix.
+This identity can be used to compute the inverse of the GGN matrix with the prior precision as
\begin{align}
\left( \alpha \identity + \bm J\T \bm H \bm J \right)^{-1} \bm v = \alpha^{-1} \identity - \alpha^{-2} \bm J\T \left( \bm H^{-1} + \alpha^{-1} \bm J \bm J\T \right)^{-1} \bm J,
\end{align}
@@ -389,7 +387,7 @@ We can then use this approximation as a preconditioner by computing its inverse
\\ \bm P^{-1} ={}& \alpha^{-1} \identity - \alpha^{-2} \bm L\T \left( \identity + \alpha^{-1} \bm L \bm L\T \right)^{-1} \bm L.
\end{align}
We then obtain the preconditioner-vector and inverse-preconditioner-vector products as in Equation~\ref{eq:pivoted-cholesky-preconditioner}.
-The ``scariest'' computation in the inverse preconditioner is inverting a \(K \times K\) matrix, where \(K\) is the rank of our pivoted Cholesky factorisation (which we choose).
+The most expensive computation in the inverse preconditioner is inverting a \(K \times K\) matrix, where \(K\) is the rank of our pivoted Cholesky factorisation (which we choose).
\begin{align}\label{eq:pivoted-cholesky-preconditioner}
\bm P \bm v ={} & \bm L\T ( \bm L \bm v ) + \alpha \bm v,
\\ \bm P^{-1} \bm v ={}& \alpha^{-1} \bm v - \alpha^{-2} \bm L\T ( \underbrace{\identity + \alpha^{-1} \bm L \bm L\T}_{K \times K} )^{-1} ( \bm L \bm v ).
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -34,7 +34,7 @@ We thus obtain the approximate Laplace posterior
\\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
\\={}& \log p(\bm \theta_{\textsc{map}} \given \bm y) - \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}).\nonumber
\end{align}
-The approximate Laplace marginal is then computed as\marginnote{MEH, this page feels a bit too equation-heavy.}
+The approximate Laplace marginal is then computed as
\begin{align}\label{eq:laplace-marginal}
p(\bm y) ={} \stackrel{\textsc{la}}{\approx}{} & \int_{\bm \theta} p(\bm \theta_{\textsc{map}} \given \bm y) \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}) \right)\,d \bm \theta
\\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) \int_{\bm \theta} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
@@ -50,17 +50,17 @@ Taking the log of this approximate Laplace marginal, we obtain
For the normal likelihood and normal prior, we have
\begin{align}\label{eq:normal-likelihood}
p(\bm y \given \bm \theta) ={} & \normal(\bm y \given f_{\bm\theta}(\bm x), \rho^{-1} \identity) = \prod_{i=1}^N \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity),
- \\ \log p(\bm y \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity)
+ \\ \log p(\bm y \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity) \label{eq:normal-log-likelihood}
\\ ={}& \sum_{i=1}^N \left( - \frac{O}{2} \log(2 \pi) - \frac{1}{2} \log \det (\rho^{-1} \identity) - \frac{1}{2} \rho \norm{y_i - f_{\bm\theta}(\bm x_i)}^2\right)\nonumber
\\={}& - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2\nonumber
\end{align}
and
\begin{align}\label{eq:normal-prior}
p(\bm \theta) ={} & \normal(\bm \theta \given \bm 0, \alpha^{-1} \identity),
- \\ \log p(\bm \theta) ={}& - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2,
+ \\ \log p(\bm \theta) ={}& - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2, \label{eq:normal-log-prior}
\end{align}
where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of the likelihood.
-Combining \cref{eq:laplace-log-marginal,eq:normal-likelihood,eq:normal-prior}, we expand the Laplace approximation to the log-marginal likelihood as
+Combining \cref{eq:laplace-log-marginal,eq:normal-log-likelihood,eq:normal-log-prior}, we expand the Laplace approximation to the log-marginal likelihood as
\begin{align}\label{eq:expanded-log-marginal}
\log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{y_i - f_{\bm \theta}(\bm x_i)}^2
\\ & - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2 \nonumber
@@ -106,7 +106,7 @@ Since \(B_D(\mu_1, \mu_2, \beta)\) is an upper bound on the log-determinant of \
This is effective because we want to maximise the log-marginal.
Assume \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) has rank \(K \leq O \cdot N\).
-Then, \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues, and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
+Then, \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues greater than \(\alpha\), and the remaining \(D - K\) eigenvalues are equal to \(\alpha\).
We then decompose the log determinant as
\begin{align}\label{eq:log-determinant-decomposition}
\log \det \bm \Lambda ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + \sum_{k=K+1}^D \log \alpha
@@ -131,7 +131,7 @@ We can bound the log determinant of \(\bm \Lambda\) using \cref{eq:log-determina
\\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,
\\ \mu_2 ={}& \Tr\left(\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)^2\right) + 2 \alpha \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha^2.
\end{align}
-For the upper bound on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\).
+For the upper bound on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\) as an upper bound.\sidenote{Since the matrix is positive semidefinite, all eigenvalues are non-negative and thus the trace is bigger than the largest eigenvalue.}
From \cref{eq:laplace-log-determinant-bound}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log determinant of the posterior precision.
% WHERE DOES THE BOUND COME FROM? WHICH UNCERTAINTY? THE UPPER BOUND ON THE EIGENVALUE?
diff --git a/figures/overview.tex b/figures/overview.tex
@@ -0,0 +1,26 @@
+\begin{tikzpicture}
+ [
+ % font=\sffamily,
+ box/.style = {draw, rounded corners, minimum width=22mm, minimum height=5mm, align=center},
+ ]
+ \definecolor{kaogreen}{HTML}{91C889}
+ \node (nn) [box] {Neural Network};
+ \node (max) [below=of nn,yshift=0.5cm] {\(\argmax_{\bm \theta}\)};
+ \node (posterior) [box,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,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)\)}}\)};
+ \draw[dashed] (nn) to (max);
+ \draw[dashed,->] (max) to [bend right] (posterior);
+ \draw[dashed,->] (max) to [bend left] (marginal);
+ \node (training) [box,minimum width=8.05cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
+ \node [above=0cm and 0cm of training.north west,anchor=south west] {\textsc{Training}};
+ \node [above=0cm and 0cm of posterior.north west,anchor=south west] {\textsc{post-hoc}};
+ \node [above=0cm and 0cm of marginal.north east,anchor=south east] {\textsc{online}};
+ %
+ \node (map) [below=of training,yshift=0.7cm] {\(\bm \theta_{\textsc{map}}\)};
+ \draw[dashed] (training) to (map);
+ \node (sample) [below=of map] {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 (inference) [box,minimum width=8.05cm,minimum height=3.2cm,fit={(sample)(samplingeq)}] {};
+ \node [above=0cm and 0cm of inference.north west,anchor=south west] {\textsc{Inference}};
+ \draw[dashed,->] (map) to (inference);
+\end{tikzpicture}
diff --git a/main.tex b/main.tex
@@ -56,7 +56,7 @@
\usepackage{censor}
\usepackage{tikz}
-\usetikzlibrary{patterns,shapes.arrows,plotmarks,angles,quotes}
+\usetikzlibrary{fit,patterns,shapes.arrows,quotes,positioning}
\usepackage{pgfplots}
\pgfplotsset{compat=newest}
\DeclareUnicodeCharacter{2212}{−}
