commit af49311130d1e13c108c05e9f76989879cfb9db5
parent 49b9bbd989f144bcae3e171dc44e5920ffc099a0
Author: Silas Brack <s174433@student.dtu.dk>
Date: Mon, 2 Jan 2023 19:35:42 +0100
Wrote more.
Diffstat:
5 files changed, 136 insertions(+), 86 deletions(-)
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -4,11 +4,41 @@
\section{The Advantage of Working With Hessian-Vector Products}
-How much larger matrices can we work with if we don't need to instantiate the whole Hessian?
-How much faster is it to compute the GVP versus, say, the HVP or the explicit Hessian-vector product?
+In this project we propose performing the Laplace approximation by only storing the computational graph for the implicit Hessian-vector product, thereby avoiding the explicit computation and storage of the quadratically-scaling Hessian matrix.
+How significant of an effect does this have, though?
+To test this, we perform Hessian-vector products \(\bm K \bm v\) and inverse Hessian-vector products \(\bm K^{-1} \bm v\), where \(\bm K \in \reals^{N \times N}\), and compare their performance.
+To compare performance with regard to memory usage, we perform these products for an increasing problem dimensionality \(N \in \{10^1, 10^2, \ldots, 10^9\}\).
+
+It is difficult to benchmark memory consumption in JAX, since the just-in-time compiler optimises computations using XLA.
+This then prevents us from stopping the program and measuring memory usage at a single point in time, as many computations will be performed at once.
+Ideally, it would be possible to measure the \emph{peak} memory usage by JAX during the running of the program.
+This is not trivial, though, as it would require querying the memory usage from the GPU asynchronously at a very high frequency and ``hope'' the peak usage point was captured.
+Instead, we simply measure the value of \(N\) at which the program fails from an out-of-memory error instead, as a proxy for memory usage.
+
+This experiment will then answer two questions:
+\begin{enumerate}
+ \item How much larger problems can we tackle 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}
-describe experimental setup, how were results generated etc.
-Note that computational time includes \emph{computing} the Hessian, and then multiplying it by a vector.
+\begin{marginfigure}
+ \centering
+ \includegraphics{hessian_profile.pdf}
+ \caption{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for Hessian-vector and GGN-vector products.}
+ \label{fig:hessian-profile}
+\end{marginfigure}
+
+We then compose a function \((g \circ f) (\bm x) := \sum^N_{i=1} (2 x_i + 0.5)^2\) in terms of
+\(f(\bm x) = 2 \bm x + 0.5, g(\bm x) := \sum^N_{i=1} x_i^2\).
+This is equivalent to the sum of square errors loss \(g\) for a simple linear model \(f\).
+Since \(f(x)\) is a linear function, the GGN matrix will exactly equal the actual Hessian.
+We then initialise \(\bm x\) to a random vector and compute the matrix-vector product \(\bm K \bm v\) where \(\bm v\) is a random vector using three Hessian-vector products: the JAX efficient HVP and GVP functions and a manual Hessian instantiation and multiplication.
+We then measure the total wall-clock time for computing this Hessian and multiplying it by \(\bm v\).
+
+Furthermore, since in this project we typically compute some variation of the matrix inverse of \(\bm K\), and inverting a high-dimensional matrix can also be very expensive, we also compare the performance of the \emph{inverse} Hessian-vector product \(\bm K^{-1} \bm v\).
+To compute the inverse for the efficient implicit HVP and GVP functions, we use the conjugate gradient method, which only requires the use of matrix-vector multiplications to compute this inverse product.
+
+% RESULTS
for values of 100k parameters or more, the manual computation failed, for both the inverse and the vector product (because it's the explicit \(N \times N\) matrix instantiation that fails)
The largest number of parameters tested is 100M.
@@ -17,36 +47,36 @@ Notably, however, these implicit Hessian-vector product methods allowed for comp
\cref{fig:hessian-profile}
-\begin{marginfigure}
- \centering
- \includegraphics{hessian_profile.pdf}
- \caption{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for Hessian-vector and GGN-vector products.}
- \label{fig:hessian-profile}
-\end{marginfigure}
-
some analysis
comment on performance of manual products
comment on weird dip in inverse, possibly something to do with caching or something like that? like the size of your cache
+\begin{marginfigure}[-4cm]
+ \centering
+ \includegraphics{hessian_profile_ratios.pdf}
+ \caption{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. It can be seen that the two methods have the same time complexity (up to a linear factor).}
+ \label{fig:hessian-profile-ratios}
+\end{marginfigure}
+
however, due to the scale of the figure, it isn't completely clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product.
We can then plot the speedup from going from the HVP to the GVP as a function of the number of model parameters.
This can be calculated as the ratio between the wall-clock time of the HVP to the wall-clock time of the GVP.
Thus, larger values suggest greater performance gains for the GVP, while values closer to one suggest there is no performance gain.
These results can be seen in \cref{fig:hessian-profile-ratios}.
-\begin{marginfigure}
- \centering
- \includegraphics{hessian_profile_ratios.pdf}
- \caption{Comparison of the speedup from computing 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. It can be seen that the two methods have the same time complexity (up to a linear factor).}
- \label{fig:hessian-profile-ratios}
-\end{marginfigure}
-
Notice that the number of parameters on the \(x\)-axis are exponentially increasing.
This is in line with the theory, as noted in \cref{ssec:practical-ggn} (SHOULD IT BE MOVED HERE?)
\section{Sampling Ablation}
+The contour integral quadrature method (CIQ) is a Krylov subspace method which converges to the exact solution for \(\lim_{Q \to \infty} \bm s^{\mathrm{CIQ}}_{Q} = \bm K^{-1/2}\).
+However, CIQ is typically extremely close to the exact solution for low values of \(Q\)\sidenote{\textcite{pleiss2020fast} find that \(Q=20\) is approximately sufficient for most problems.}.
+We then investigated the number of quadrature points \(Q\) required to determine an adequate approximate solution for the Laplace approximation.
+
+Need to define a metric which determines if sampling was successful.
+Root mean square error of my percentiles vs the real chi-sq percentiles? Feels a bit weird.
+
As we vary \(q \in [1, \ldots, 20]\), how does performance vary?
What metrics can we use?
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -36,7 +36,7 @@ architecture.
% \section{My Solution}
\section{Large-Scale Laplace}
-The ideas presented in this article essentially constitute a compilation of advice given by academics from multiple fields and institutions,~\cite{miani2022laplacian}
+The ideas presented in this article essentially constitute a compilation of advice given by academics from multiple fields and institutions,~\sidecite{miani2022laplacian}
% \section{Research Objectives}
@@ -60,4 +60,4 @@ frameworks to efficiently approximate the Hessian.
\section{JAX and XLA}
-\cite{jax2018github,sabne2020xla}
+\sidecite{jax2018github,sabne2020xla}
diff --git a/chapters/literature.tex b/chapters/literature.tex
@@ -7,8 +7,8 @@
% How does adam work?
-\cite{sutskever2013importance}
-\cite{kingma2014adam}
+\sidecite{sutskever2013importance}
+\sidecite{kingma2014adam}
% Integrate with the previous section
Note that Adam approximates the second moment of the gradient of the loss, that this corresponds to the variance of this gradient, and that this is the diagonal of the Fisher information matrix, as we defined in \cref{sec:fisher-information} (\cref{eq:gradient-covariance})!
@@ -17,21 +17,21 @@ Note that Adam approximates the second moment of the gradient of the loss, that
% \section{K-FAC}
% Another alternative to approximating the Fisher information matrix is by calculating the Kronecker-factored Approximate Curvature (K-FAC).
-% \cite{martens2015optimizing}
+% \sidecite{martens2015optimizing}
% \section{Noisy K-FAC}
-% \cite{zhang2018noisy}
+% \sidecite{zhang2018noisy}
% \section{Gauss-Newton Optimization}
-% \cite{botev2017practical}
+% \sidecite{botev2017practical}
% \section{Variational Adaptive-Newton}
-% \cite{khan2017variational}
+% \sidecite{khan2017variational}
% \section{Natural Gradient Variational Inference}
-% Natural gradient VI (VOGN)~\cite{osawa2019practical} approximates the posterior as
+% Natural gradient VI (VOGN)~\sidecite{osawa2019practical} approximates the posterior as
% %
% \begin{equation*}
% p(\bm \theta) \approx \normal(\bm 0, \bm I / \delta)
@@ -44,30 +44,35 @@ Note that Adam approximates the second moment of the gradient of the loss, that
\section{Laplace Approximation}\label{sec:laplace}
-In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI.
+In the Laplace approximation (LA)~\sidecite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI.
However, instead of finding the optimal Gaussian distribution locations and scales by maximising the ELBO, the LA finds the location by computing the MAP solution \(\bm \theta_{\textsc{map}}\) and the scale by approximating the log posterior with a second degree Taylor expansion around this solution (\(\bm \theta_0 = \bm \theta_{\textsc{map}}\)) and determining the curvature via its Hessian matrix \(\bm \Lambda = - \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_\mathrm{MAP}}\).
Since the Taylor expansion is performed around the MAP solution, the first order derivative is zero, and the expansion is simply given by
-\begin{align*}
+\begin{align}
\ln p(\bm \theta \given \bm y)
- \approx{} & \ln p(\bm \theta_{\textsc{map}} \given \bm y) +
- \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \left( \left.\nabla^2_{\bm \theta} \ln p(\bm \theta \given \bm y) \right|_{\bm \theta_\mathrm{MAP}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) %\label{eq:laplace-taylor}
+ \approx{} & \ln p(\bm \theta_{\textsc{map}} \given \bm y) + \nonumber
+ \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \left( \left.\nabla^2_{\bm \theta} \ln p(\bm \theta \given \bm y) \right|_{\bm \theta_\mathrm{MAP}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) \label{eq:laplace-taylor}
\\ \Rightarrow \tilde{p}(\bm \theta \given \bm y) \approx{}& p(\bm \theta_{\textsc{map}} \given \bm y) \exp\left(- \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right) \propto p(\bm \theta \given \bm y) \nonumber
-\end{align*}
+\end{align}
where \(\tilde{p}(\bm \theta \given \bm y)\) corresponds to the unnormalised posterior.
Normalising it yields
-\begin{align*}
+\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
% \\={}& \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_{\textsc{map}}, \bm \Lambda^{-1}).%\label{eq:laplace}
-\end{align*}
+ \\={}& \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}).
+\end{align}
-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_{\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}}}\).
+\marginnote{Note that \cref{eq:laplace-normal} only holds when the precision matrix \(\bm \Lambda\) is positive definite.}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_{\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}}}\).
In practice, we can determine the Hessian of the loss with respect to the parameters, where the loss is given by the negative log-posterior.
In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_{\textsc{map}}} = \left. \nabla^2_{\bm \theta} \mathcal{L} \right|_{\bm \theta_{\textsc{map}}}\) and the posterior is approximated by \(\normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\).
\cref{eq:negative-log-likelihood-ggn}
+GGN APPROXIMATION FOR POSITIVE DEFINITENESS
+
+EIGENVALUES?
+minimum eigenvalue determined by the perturbation \(\alpha\) which guarantees positive definiteness
+
% \begin{align}
% \bm \Lambda = \sum_i^N \bm J_i\T \bm H_i \bm J_i
% \end{align}
@@ -77,14 +82,14 @@ In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \
% % https://martiningram.github.io/vi-with-good-covariances/
% % https://github.com/martiningram/jax_advi/blob/main/jax_advi/lrvb.py#L49
-% \cite{giordano2018covariances}
+% \sidecite{giordano2018covariances}
% INLA: Laplace for nested models (hierarchical models)
-% \cite{rue2017bayesian}
+% \sidecite{rue2017bayesian}
% \section{Online Laplace}
-% \cite{miani2022laplacian}
+% \sidecite{miani2022laplacian}
% \section{Bayes’ Theorem}
@@ -179,8 +184,8 @@ In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \
% \end{align}
% where \(\bm{\theta}^{(s)} \sim \normal(\bm{\mu}, \bm{\sigma})\).
-% The use of the mean-field approximation for BNNs has been found to be unreliable~\cite{wu2018deterministic} and, for increasingly wide neural networks, converges towards the prior~\cite{coker2021wide}.
-% However, for \emph{deep} BNNs, the mean-field assumption may be reasonable~\cite{farquhar2020liberty}.
+% The use of the mean-field approximation for BNNs has been found to be unreliable~\sidecite{wu2018deterministic} and, for increasingly wide neural networks, converges towards the prior~\sidecite{coker2021wide}.
+% However, for \emph{deep} BNNs, the mean-field assumption may be reasonable~\sidecite{farquhar2020liberty}.
% \subsection{Full-Rank Variational Approximation}
@@ -220,7 +225,7 @@ In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \
% \subsection{Deep Ensembles}
-% In deep ensembles~\cite{lakshminarayanan2017simple}, \(M\) neural networks are trained with different initialisations.
+% In deep ensembles~\sidecite{lakshminarayanan2017simple}, \(M\) neural networks are trained with different initialisations.
% In this way, multiple local MAP solutions \(\bm\theta^{(m)}\) are obtained and equally considered, such that \(p(\bm \theta = \bm \theta^{(m)} \given \bm y) = 1 / M\) for every \(\bm \theta^{(m)} \in \{\bm\theta^{(1)}, \ldots, \bm\theta^{(M)}\}\) (and 0 elsewhere).
% To make predictions with a deep ensemble, we use the posterior predictive given by
% \begin{align}
@@ -234,13 +239,13 @@ In this case, we get \(\bm \Lambda = \left. - \nabla^2_{\bm \theta} \log p(\bm \
% Stochastic weight averaging estimates the final point estimate of the weights of a neural network \(\bm \theta_\mathrm{SWA}\) as the average of the model parameters \(\bm \theta\) at the end of each epoch obtained via gradient descent after convergence has been reached.
-% For SWAG~\cite{maddox2019simple}, the standard deviation of the model parameters \(\bm \sigma_\mathrm{SWA}\) is also calculated.
+% For SWAG~\sidecite{maddox2019simple}, the standard deviation of the model parameters \(\bm \sigma_\mathrm{SWA}\) is also calculated.
% From the SWA estimates of the model parameter mean and standard deviations, the posterior is then estimated as
% \begin{align}\label{eq:swag}
% p(\bm \theta \given \bm y) \approx{} & \normal(\bm \theta_\mathrm{SWA}, \bm \sigma_\mathrm{SWA}).
% \end{align}
-% For MultiSWAG~\cite{wilson2020bayesian}, \(M\) SWAG estimates are obtained from different initialisations, as in deep ensembles.
+% For MultiSWAG~\sidecite{wilson2020bayesian}, \(M\) SWAG estimates are obtained from different initialisations, as in deep ensembles.
% Then, a Gaussian mixture model (GMM) is formed as a combination of each Gaussian SWAG posterior estimate \(\normal(\bm \theta_\mathrm{SWA}^{(m)}, \bm \sigma_\mathrm{SWA}^{(m)})\) where each component is equally weighed, as
% \begin{align}\label{eq:multiswag}
% p(\bm \theta \given \bm y) \approx{} & \frac{1}{M} \sum_{m=1}^M \normal(\bm \theta_\mathrm{SWA}^{(m)}, \bm \sigma_\mathrm{SWA}^{(m)}).
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -85,23 +85,6 @@ Since, in our case, the eigenvalues of \(\bm K\) will all be real-valued\marginn
% QUADRATURE
% ------------------------------------------
-We can then apply this to the function \(f(\bm K) = \bm K^{-1/2}\), apply a change of variable, and then approximate the integral using the quadrature rule with \(Q\) quadrature points, as per
-\begin{align}\label{eq:ciq}
- \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 \nonumber
- \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q w_q \left( \sigma^2_q \identity - \bm K \right)^{-1} \nonumber
- \\ \bm K^{-1/2} \bm v \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q w_q \underbrace{\left( \sigma^2_q \identity - \bm K \right)^{-1} \bm v}_{Q\text{ system solves}},
-\end{align}
-where \(\sigma\) is obtained from \(\tau\) as a change of variable and \(w_1, \ldots, w_Q\) are the \(Q\) quadrature weights.
-We can then solve the \(Q\) systems of equations \(\left( \sigma^2_1 \identity - \bm K \right) = \bm v, \ldots, \left( \sigma^2_Q \identity - \bm K \right) = \bm v\) to obtain the solutions \(\bm c_q = \left(\sigma^2_q \identity - \bm K \right)^{-1} \bm v\).
-Note that, since the integrand is symmetric with respect to the real axis\sidenote{This is due to the fact that the circle with its centre on the \(x\)-axis is symmetric with respect to said \(x\)-axis.}, only the imaginary component of \(\Gamma_\sigma\) is required.
-Since \(\sigma\) is then imaginary (its real part is zero), \(\tau := \sigma^2\) is real such that \(\tau = - \Im(\sigma)^2\).
-As such, since \(\Im(\sigma)^2\) must be positive, \(\tau\) must then be real-valued and negative.
-And since, as mentioned before, the integrand is symmetric with respect to the real axis, we can ?????
-In practice, we incorporate the \(1 / (\pi i) \) term into the quadrature weights and, with the quadrature points \(\sigma^2_1, \ldots, \sigma^2_Q\), compute \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1} w_q \bm c_q\), where the solutions \(\bm c_1, \ldots \bm c_Q\) are computed using msMINRES (\cref{ssec:msminres}).
-
-Since we will be sampling a fixed number of quadrature points from the circle, to guarantee 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 of the spectrum of \(\bm K\).
-
\def\ymax{2.0}
\def\ymin{-2.0}
\def\xmax{3.0}
@@ -114,22 +97,43 @@ Since we will be sampling a fixed number of quadrature points from the circle, t
\label{fig:ciq}
\end{marginfigure}
+We can then apply this to the function \(f(\bm K) = \bm K^{-1/2}\), apply a change of variable, and then approximate the integral using the quadrature rule with \(Q\) quadrature points (see \cref{fig:ciq}), as per
+\begin{align}\label{eq:ciq}
+ \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 \nonumber
+ \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1} \nonumber
+ \\ \bm K^{-1/2} \bm v \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \underbrace{\left( \sigma^2_q \identity - \bm K \right)^{-1} \bm v}_{Q\text{ system solves}},
+\end{align}
+where \(\sigma\) is obtained from \(\tau\) as a change of variable and \(\Tilde{w}_1, \ldots, \Tilde{w}_Q\) are the \(Q\) quadrature weights.
+We can then solve the \(Q\) systems of equations \(\left( \sigma^2_1 \identity - \bm K \right) \bm c_1 = \bm v, \ldots, \left( \sigma^2_Q \identity - \bm K \right) \bm c_Q = \bm v\) to obtain the solutions \(\bm c_q = \left(\sigma^2_q \identity - \bm K \right)^{-1} \bm v\).
+Note that, since the integrand is symmetric with respect to the real axis\sidenote{This is due to the fact that the circle with its centre on the \(x\)-axis is symmetric with respect to said \(x\)-axis.}, only the imaginary component of \(\Gamma_\sigma\) is required.
+Since \(\sigma\) is then imaginary (its real part is zero), \(\tau := \sigma^2\) is real such that \(\tau = - \Im(\sigma)^2\).
+As such, since \(\Im(\sigma)^2\) must be positive, \(\tau\) must then be real-valued and negative.\marginnote{For simplicity, since \(w_q\) happens to also be real-valued and negative, from now on we will redefine the variables \(w_q \gets -\Tilde{w}_q\) and \(t_q \gets -\sigma^2_q\).}
+In practice, we incorporate the \(1 / (\pi i) \) term into the quadrature weights and, with the quadrature points \(t_1, \ldots, t_Q\), compute \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1} w_q \bm c_q\).
+
+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 of the spectrum of \(\bm K\).
+
% ------------------------------------------
% JACOBI ELLIPTICAL FUNCTIONS
% ------------------------------------------
-We could sample these quadrature points uniformly, but...
-Instead, we use Jacobi elliptic functions (\cref{sec:jacobi-elliptic})
-With these functions, we can compute the weights for Gaussian quadrature as
-\begin{align}
- w_q = - \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cn(i u_q \mathcal{K}^\prime (k) \given k) \dn (i u_q \mathcal{K}^\prime (k) \given k)
+We could sample these quadrature points from the circle uniformly, which corresponds to the regular trapezoid quadrature rule.
+However, the convergence for this quadrature is linear with regard to the condition number of \(\bm K\) \sidecite{pleiss2020fast,hale2008computing}.
+Since the precision matrix computed from the Laplace approximation as the GGN matrix can often have a very low rank and poor conditioning, uniform sampling would then require a very large number of quadrature points \(Q\) and would therefore be inadequate.
+
+Since our GGN precision matrix \(\bm \Lambda = \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) can be 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}\) before sampling \(\sigma\).
+This change of variables yields a set of quadrature weights \(w_1, \ldots, w_Q\) given by
+\begin{align}\label{eq:quadrature-weights}
+ w_q = - \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cn(i u_q \mathcal{K}^\prime (k) \given k) \dn (i u_q \mathcal{K}^\prime (k) \given k),
\end{align}
-and shifts as
-\begin{align}
- \sigma^2_q = \lambda_{\mathrm{min}} \cdot \sn(i u_q \mathcal{K}^\prime (k) \given k)^2
+and a set of quadrature points (shifts) by
+\begin{align}\label{eq:quadrature-shifts}
+ \sigma^2_q = \lambda_{\mathrm{min}} \cdot \sn(i u_q \mathcal{K}^\prime (k) \given k)^2,
\end{align}
-
-\cref{alg:quad} describes .
+where \(\sn\), \(\cn\), and \(\dn\) are the Jacobi elliptic functions and \(\mathcal{K}^\prime(k) = \mathcal{K}(k^\prime) = \mathcal{K}(\sqrt{1 - k^2})\) (see \cref{sec:jacobi-elliptic}).
+To compute these values, we also use the elliptic modulus \(k = \kappa(\bm K)^{-1/2}\) (where \(\kappa(\bm K)\) denotes the condition number of \(\bm K\)), and \(u_q = (q - 1/2) / Q\).
+We can thereby now compute the quadrature weights and shifts, as seen in \cref{alg:quad}.
\begin{marginfigure}
\begin{algorithm}[H]
@@ -174,6 +178,13 @@ and shifts as
% }
% \end{algorithm}
+The last step involves solving the \(Q\) linear systems defined in \cref{eq:ciq} to determine the solves \(\bm c_q\).
+As mentioned above, this is done by using msMINRES (\cref{ssec:msminres,alg:msminres}) as \texttt{msMINRES(-K, b, t\_q, P, J)}\sidenote{Note that we negate \(\bm K\), since msMINRES computes \(\left(t_q \identity + \bm K \right)^{-1} \bm v\), while we want to compute \(\left(t_q \identity - \bm K \right)^{-1} \bm v\). Since we can only access the matrix-vector product, we simply use \(- \bm K \bm v\) whenever the matrix-vector product is called.}.
+
+This last step is then sensitive to preconditioning
+
+Mention computational bottleneck being performing the msMINRES for the shifted systems.
+
\section{Preconditioning}
One drawback of contour integral quadrature if that the accuracy its of approximation of \(\bm \Lambda^{-1/2} \bm \varepsilon_0\) suffers when the precision matrix \(\bm \Lambda\) is poorly conditioned\sidenote{Recall that a matrix is considered poorly conditioned when the ratio of its largest eigenvalue to its lowest eigenvalue, its \emph{conditioning number} (\(\kappa := \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}\)), is very large.}.
@@ -288,7 +299,7 @@ Since \(N \gg B\), then the cost of computing the preconditioner will be negligi
\subsection{Pivoted Cholesky}
-\cite{harbrecht2012low,bach2013sharp,gardner2018gpytorch}
+\sidecite{harbrecht2012low,bach2013sharp,gardner2018gpytorch}
\begin{align}
\Sigma^{-1} ={} & \sum_i^N J_i\T J_i + \alpha \identity
@@ -321,7 +332,7 @@ If the fall-off in the eigenspectrum of this factorisation is steep enough, then
Unfortunately, if want to use the pivoted Cholesky decomposition to approximate \(\sum^B_i J_i\T J_i\), we are largely trying to approximate the higher eigenvectors of the sum.
In order for this to be the case, we will want to oversample vectors from this matrix which correspond to the largest eigenvalues.
-A common approach to sampling the largest eigenvalues is to sample vectors with probability proportional to the diagonal value~\cite{harbrecht2012low}.
+A common approach to sampling the largest eigenvalues is to sample vectors with probability proportional to the diagonal value~\sidecite{harbrecht2012low}.
However, it is not, in general, possible to efficiently determine the diagonal of a matrix via matrix-vector products.
Often,it is necessary to multiply a vector of zeros with a value of one at the index at which to find the diagonal.
diff --git a/chapters/theory.tex b/chapters/theory.tex
@@ -2,16 +2,16 @@
\chapter{Theory}
\labch{theory}
-\section{Optimization in Deep Learning}
+\section{Optimisation in Deep Learning}
-Suppose a neural network is a real-valued function \(f: \reals^n \times \reals^d \rightarrow \reals^o\) parameterized in \(\bm\theta\) which maps an input \(\bm x\) to an output \(f(\bm x; \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
+Suppose a neural network is a real-valued function \(f: \reals^n \times \reals^d \rightarrow \reals^o\) parameterised in \(\bm\theta\) which maps an input \(\bm x\) to an output \(f(\bm x; \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
Our goal is to find the optimal parameters \(\bm\theta^*\) which best model the observed data.
-From a frequentist perspective, we can define a loss function \(\mathcal{L}(\bm\theta): \reals^d \rightarrow \reals\) such that the optimal parameters \(\bm\theta^*\) minimize this loss function.
+From a frequentist perspective, we can define a loss function \(\mathcal{L}(\bm\theta): \reals^d \rightarrow \reals\) such that the optimal parameters \(\bm\theta^*\) minimise this loss function.
Often the loss function is defined as the negative log-likelihood of the data under the model, \(\mathcal{L}(\bm\theta) = -\log \lik\).
-In this case, the goal of optimization is to find the parameters \(\bm\theta^*\) which maximize the likelihood of the data~\footnote{Maximizing the likelihood is equivalent to minimizing the negative log likelihood, since the log function is strictly monotonically increasing.}.
+In this case, the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the likelihood of the data~\sidenote[][5cm]{Maximising the likelihood is equivalent to minimising the negative log likelihood, since the log function is strictly monotonically increasing.}.
-From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and a likelihood function \(\lik\) such that the goal of optimization is to find the parameters \(\bm\theta^*\) which maximize the posterior distribution \(p(\bm\theta \given \bm y)\).
+From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and a likelihood function \(\lik\) such that the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the posterior distribution \(p(\bm\theta \given \bm y)\).
This way, we can quantify the uncertainty in our model parameters \(\bm\theta\) by computing the posterior distribution \(p(\bm\theta \given \bm y)\).
Due to the assumption of independence between observations, we can factorise the negative log-likelihood loss as
@@ -28,7 +28,7 @@ Similarly, for the negative log-posterior loss, we have
\subsection{Stochastic Gradient Descent}
-Stochastic gradient descent (SGD) is a first-order optimization algorithm which iteratively updates the parameters \(\bm\theta\) in the direction of the negative gradient of the loss function \(\mathcal{L}(\bm\theta)\).
+Stochastic gradient descent (SGD) is a first-order optimisation algorithm which iteratively updates the parameters \(\bm\theta\) in the direction of the negative gradient of the loss function \(\mathcal{L}(\bm\theta)\).
The update rule is given by
%
\begin{equation*}
@@ -53,7 +53,7 @@ It can either be constant throughout training (\(\eta_t = \eta\)) or adaptive (i
The Hessian is a matrix of second-order partial derivatives of a scalar function.
From now on, we will use the notation \(\nabla^2_{\bm\theta} f\) to represent the Hessian matrix of \(f\), where each element is given by \(\left(\nabla^2_{\bm\theta} f\right)_{i, j} = \frac{\partial^2 f}{\partial\theta_i \partial\theta_j}\).
-Suppose we have a function \(\mathcal{L}: \reals^d \rightarrow \reals\) parameterized by \(\bm \theta \in \reals^d\).
+Suppose we have a function \(\mathcal{L}: \reals^d \rightarrow \reals\) parameterised by \(\bm \theta \in \reals^d\).
Then the Hessian can be written as \(J_{\bm \theta}(\nabla_{\bm\theta} \mathcal{L})\).
We are interested in the Hessian of the loss function \(\mathcal{L}(\bm\theta)\) with respect to the parameters \(\bm\theta\).
@@ -69,7 +69,7 @@ In this framing, if considering, e.g., the negative log-posterior loss from \cre
\\ ={}& \sum_i^N \nabla^2_{\bm\theta} \log p(x_i \given \bm \theta) + \nabla^2_{\bm\theta} \log p(\bm \theta).
\end{align}
-\section{The Generalized Gauss-Newton Approximation}
+\section{The Generalised Gauss-Newton Approximation}
From the equation of the Hessian we can apply the chain rule twice and the product rule once to obtain a simpler expression
%
@@ -114,7 +114,9 @@ For log-likelihood loss, we get
\\ ={}& \sum_i^N J_i\T H_i J_i \nonumber
\end{align}
-\subsection{Practical Considerations for the GGN Approximation}\label{ssec:practical-ggn}
+\subsection{Practical Considerations}
+% \subsection{Practical Considerations for the GGN Approximation}
+\label{ssec:practical-ggn}
Note that the generalised Gauss-Newton approximation requires that the function for which we would like to compute the Hessian, \(h(x)\), can be framed as a composition of two functions \(h(x) = g(f(x))\).
However, there are many possible choices of functions \(f\) and \(g\) which, when composed, yield \(h\), and these different compositions can yield wildly varying performance for the GGN approximation~\cite{kunstner2019limitations}.
@@ -239,7 +241,8 @@ Instead, it is often approximated by the diagonal of the Hessian.
For a recent, detailed discussion of the natural gradient, see \cite{martens2020new}.
\cite{wu2019logan}
-\section{Inverting a low-rank matrix}\label{sec:woodbury}
+\section{Inverting a low-rank matrix}
+\label{sec:woodbury}
% Theorem 1.1 ([18, Thm. 1.35]). Let U, V ∈ Cn×k with k ≤ n and assume that V ∗U is nonsingular. Let f be defined on the spectrum of A = αIn + U V ∗, and if k = n let f be defined at α. Then (1.2) f (A) = f (α)In + U (V ∗U )−1(f (αIk + V ∗U ) − f (α)Ik )V ∗.
\begin{align}
@@ -290,7 +293,7 @@ The power iteration method involves iteratively computing the normalised product
\bm b_{j+1} = \frac{\bm K \bm b_{j}}{\norm{\bm K \bm b_{j}}},
\end{align}
for some number of iterations \(J\), such that \(\lim_{j \to \infty} \bm b_j\) equals the eigenvector corresponding to the largest eigenvalue of \(\bm K\).
-This thereby yields a set of vectors which is proportional to\sidenote{This is because the norm \(\zeta_j := \norm{\bm K \bm b_{j}}\) is a scalar, which means that \(\bm b_j = \bm K \bm b_{j-1} \zeta_{j-1}^{-1} = \zeta_{j-1}^{-1} \bm K \bm b_{j-1}\), so \(\bm b_j = Z \bm K \bm b_{j-1}\) where \(Z = \prod_{i=0}^{j-1} \zeta_{i}^{-1} \).} \(\{ \bm b_0, \bm K \bm b_0, \bm K^2 \bm b_0, \ldots, \bm K^{j-1} \bm b_0 \}\) which is clearly a spanning set for a Krylov subspace.
+This thereby yields a set of vectors which is proportional to\sidenote{This is because the norm of the Lanczos vectors \(\zeta_j := \norm{\bm K \bm b_{j}}\) is a scalar, which means that \(\bm b_j = \bm K \bm b_{j-1} \zeta_{j-1}^{-1} = \zeta_{j-1}^{-1} \bm K \bm b_{j-1}\), so \(\bm b_j = Z \bm K \bm b_{j-1}\) where \(Z = \prod_{i=0}^{j-1} \zeta_{i}^{-1} \).} \(\{ \bm b_0, \bm K \bm b_0, \bm K^2 \bm b_0, \ldots, \bm K^{j-1} \bm b_0 \}\) which is clearly a spanning set for a Krylov subspace.
However, the power iteration method only uses the last vector \(\bm b_J\), thereby effectively ``wasting'' the other vectors \(\bm b_1, \ldots, \bm b_{J-1}\).
To make use of the entire Krylov subspace \(\mathcal{K}_j(\bm K, \bm b_0)\), the Lanczos algorithm combines the power iteration method with the Gram–Schmidt process, along with the assumption that \(\bm K\) is Hermitian, to produce an orthonormal basis from the Krylov subspace.
This method then converts an eigendecomposition problem for \(\bm K\) into an eigendecomposition problem for a tridiagonal matrix of rank \(J\) whose maximum and minimum eigenvalues will be close to those of \(\bm K\).
@@ -349,7 +352,8 @@ The algorithm itself can be seen in \cref{alg:cg}.
% }
% \end{algorithm}
-\subsection{The Minimum Residual Method}\label{ssec:msminres}
+\subsection{The Minimum Residual Method}
+\label{ssec:msminres}
Another method which is derived from the Lanczos algorithm is the minimum residual method (MINRES).
Given a certain number of iterations \(J\), the MINRES algorithm constructs a Krylov subspace by exploiting the three-vector recurrence (like conjugate gradient does) and solves a linear system \(\bm K \bm c = \bm b\), where \(\bm K\) is a square matrix (though not necessarily a symmetric one), by finding the least-squares solution which lies within this Krylov subspace of dimension \(J\), i.e., by solving the constrained optimisation problem characterised by
@@ -397,7 +401,7 @@ The MINRES algorithm can then be trivially extended to solve for multiple shifts
Complete elliptic integral of the first kind is given by
\begin{align}
- \mathcal{K}p) = \int_0^{\pi/2} \left( 1 - m \sin(t)^2 \right)^{-1/2}\,dt, \quad p = 1 - m,
+ \mathcal{K}(p) = \int_0^{\pi/2} \left( 1 - m \sin(t)^2 \right)^{-1/2}\,dt, \quad p = 1 - m,
\end{align}
and is defined in the domain\sidenote{Note that if \(p > 1\), then what should be computed is instead the identity \(\mathcal{K}(p) = \mathcal{K}(1/p) / \sqrt{p}\).} \(0 < p \leq 1\).
This integral is approximated by
