commit 45831861696087adb4c89c4e447cad87549cd062
parent 4aec31e55061f3a0495d9d6afcc9cd4cbfa9a385
Author: Silas Brack <s174433@student.dtu.dk>
Date: Wed, 1 Feb 2023 23:35:04 +0100
Fixies.
Diffstat:
6 files changed, 78 insertions(+), 78 deletions(-)
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -3,42 +3,42 @@
% https://www.student.unsw.edu.au/writing-abstracts-honours-theses
-Bayesian methods have been used in machine learning 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 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.
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.
-The Laplace approximation is one such method which allows for post-hoc estimation of the posterior distribution by assuming a normal distribution for the posterior and using the Hessian to approximate the posterior precision.
-Additionally, since the Laplace approximation deduces uncertainty rather than learning it, it has many advantages over other Bayesian methods which attempt to learn variance.
-Due to its construction, the Laplace approximation enables training by maximising the marginal likelihood, which facilitates model selection and hyperparameter optimisation.
+The Laplace approximation is one such method which enables post-hoc estimation of the posterior distribution by assuming a normal distribution for the posterior and using the Hessian to approximate the posterior precision.
+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.
-In this work, we propose a method for performing the Laplace approximation using only Hessian-vector products, which enables us to perform marginal training and inference using the full Laplace approximation without storing the entire Hessian.
-To accomplish this, we compute the Hessian's log-determinant and inverse square root using only Hessian-vector products.
-We suggest a technique for evaluating the quality of posterior samples without instantiating the Hessian based on results from traditional statistics.
-To implement this, we use JAX, a library for automatic differentiation which enables us to efficiently perform Hessian-vector products without explicitly storing the entire Hessian.
+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.
+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.
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}
+% \chapter*{Resum\'e}
+% \addcontentsline{toc}{chapter}{Resum\'e}
-Bayesianske metoder har v{\ae}ret anvendt inden for maskinl{\ae}ring i {\aa}rtier, men har for nylig oplevet en {\o}get interesse inden for dyb l{\ae}ring, da de lover at give en principiel m{\aa}de at kvantificere usikkerheden i neurale netv{\ae}rk p{\aa}.
-Disse metoder er baseret p{\aa} bayesiansk inferens til at l{\ae}re den efterf{\o}lgende fordeling af modelparametrene i betragtning af dataene.
-Derved giver de mulighed for usikkerhedsestimater, hvilket er vigtigt for mange anvendelser inden for maskinl{\ae}ring, f.eks. anvendelser, der vedr{\o}rer sikkerhedskritiske beslutninger, som er afh{\ae}ngige af risikovurdering og fortolkningsmuligheder.
-Bayesiansk inferens er imidlertid ofte beregningsm{\ae}ssigt uoverkommelig, og derfor anvendes tiln{\ae}rmede metoder, som g{\aa}r p{\aa} kompromis med ydeevnen for at opn{\aa} gennemf{\o}rlighed.
-Laplace-approximationen er en af disse metoder, som giver mulighed for post-hoc-estimation af den efterf{\o}lgende fordeling ved at antage en normalfordeling for den efterf{\o}lgende fordeling og anvende Hessian til at tiln{\ae}rme den efterf{\o}lgende pr{\ae}cision.
-Da Laplace-approksimationen desuden udleder usikkerheden i stedet for at l{\ae}re den, har den mange fordele i forhold til andre bayesianske metoder, som fors{\o}ger at l{\ae}re variansen.
-P{\aa} grund af sin konstruktion g{\o}r Laplace-approksimationen det muligt at tr{\ae}ne ved at maksimere den marginale sandsynlighed, hvilket letter modelvalg og optimering af hyperparametre.
-Laplace-approksimationen er imidlertid typisk afh{\ae}ngig af grove approksimationer af den efterf{\o}lgende pr{\ae}cision for at g{\o}re den computerm{\ae}ssigt gennemf{\o}rlig.
-% Ofte g{\o}res der st{\ae}rke antagelser om Hessians struktur, men disse antagelser er ikke altid n{\o}jagtige.
-Konventionel fuld Laplace anvender den fulde posteriore pr{\ae}cision, som ikke kan instantieres for alle undtagen leget{\o}jsproblemer, da den skalerer kvadratisk med antallet af parametre.
-I dette arbejde foresl{\aa}r vi en metode til at udf{\o}re Laplace-approximationen ved kun at bruge Hessian-vektorprodukter, hvilket g{\o}r det muligt at udf{\o}re marginal tr{\ae}ning og inferens ved hj{\ae}lp af den fulde Laplace-approximation uden at lagre hele Hessian.
-For at opn{\aa} dette beregner vi Hessians log-determinant og den omvendte kvadratrod ved kun at bruge Hessian-vektorprodukter.
-Vi foresl{\aa}r en teknik til at evaluere kvaliteten af efterf{\o}lgende pr{\o}ver uden at instantiere Hessian'en baseret p{\aa} resultater fra traditionel statistik.
-Til at gennemf{\o}re dette bruger vi JAX, et bibliotek til automatisk differentiering, som g{\o}r det muligt for os at udf{\o}re Hessian-vektorprodukter effektivt uden eksplicit at lagre hele Hessian.
-Vi viser, at vores metode er i stand til at udf{\o}re den fulde Laplace-approximation p{\aa} b{\aa}de en sinusfunktion og MNIST, og at tr{\ae}ning af den marginale sandsynlighed ved hj{\ae}lp af vores metode er gennemf{\o}rlig og giver meningsfulde resultater.
-Dette er den f{\o}rste metode til at udf{\o}re Laplace-approximationen, samtidig med at der kun er implicit adgang til Hessian.
-Det fremtidige arbejde vil fokusere p{\aa} optimering af hyperparametre og eksperimenter i stor skala ved hj{\ae}lp af denne metode.
+% Bayesianske metoder har v{\ae}ret anvendt i {\aa}rtier, men har for nylig oplevet en {\o}get interesse inden for dyb l{\ae}ring, da de lover at give en principiel m{\aa}de at kvantificere usikkerheden i neurale netv{\ae}rk p{\aa}.
+% Disse metoder er baseret p{\aa} bayesiansk inferens til at l{\ae}re den efterf{\o}lgende fordeling af modelparametrene i betragtning af dataene.
+% Derved giver de mulighed for usikkerhedsestimater, hvilket er vigtigt for mange anvendelser inden for maskinl{\ae}ring, f.eks. anvendelser, der vedr{\o}rer sikkerhedskritiske beslutninger, som er afh{\ae}ngige af risikovurdering og fortolkningsmuligheder.
+% Bayesiansk inferens er imidlertid ofte beregningsm{\ae}ssigt uoverkommelig, og derfor anvendes tiln{\ae}rmede metoder, som g{\aa}r p{\aa} kompromis med ydeevnen for at opn{\aa} gennemf{\o}rlighed.
+% Laplace-approximationen er en af disse metoder, som giver mulighed for post-hoc-estimation af den efterf{\o}lgende fordeling ved at antage en normalfordeling for den efterf{\o}lgende fordeling og anvende Hessian til at tiln{\ae}rme den efterf{\o}lgende pr{\ae}cision.
+% Da Laplace-approksimationen desuden udleder usikkerheden i stedet for at l{\ae}re den, har den mange fordele i forhold til andre bayesianske metoder, som fors{\o}ger at l{\ae}re variansen.
+% P{\aa} grund af sin konstruktion g{\o}r Laplace-approksimationen det muligt at tr{\ae}ne ved at maksimere den marginale sandsynlighed, hvilket letter modelvalg og optimering af hyperparametre.
+% Laplace-approksimationen er imidlertid typisk afh{\ae}ngig af grove approksimationer af den efterf{\o}lgende pr{\ae}cision for at g{\o}re den computerm{\ae}ssigt gennemf{\o}rlig.
+% % Ofte g{\o}res der st{\ae}rke antagelser om Hessians struktur, men disse antagelser er ikke altid n{\o}jagtige.
+% Konventionel fuld Laplace anvender den fulde posteriore pr{\ae}cision, som ikke kan instantieres for alle undtagen leget{\o}jsproblemer, da den skalerer kvadratisk med antallet af parametre.
+% I dette arbejde foresl{\aa}r vi en metode til at udf{\o}re Laplace-approximationen ved kun at bruge Hessian-vektorprodukter, hvilket g{\o}r det muligt at udf{\o}re marginal tr{\ae}ning og inferens ved hj{\ae}lp af den fulde Laplace-approximation uden at lagre hele Hessian.
+% For at opn{\aa} dette beregner vi Hessians log-determinant og den omvendte kvadratrod ved kun at bruge Hessian-vektorprodukter.
+% Vi foresl{\aa}r en teknik til at evaluere kvaliteten af efterf{\o}lgende pr{\o}ver uden at instantiere Hessian'en baseret p{\aa} resultater fra traditionel statistik.
+% Til at gennemf{\o}re dette bruger vi JAX, et bibliotek til automatisk differentiering, som g{\o}r det muligt for os at udf{\o}re Hessian-vektorprodukter effektivt uden eksplicit at lagre hele Hessian.
+% Vi viser, at vores metode er i stand til at udf{\o}re den fulde Laplace-approximation p{\aa} b{\aa}de en sinusfunktion og MNIST, og at tr{\ae}ning af den marginale sandsynlighed ved hj{\ae}lp af vores metode er gennemf{\o}rlig og giver meningsfulde resultater.
+% Dette er den f{\o}rste metode til at udf{\o}re Laplace-approximationen, samtidig med at der kun er implicit adgang til Hessian.
+% Det fremtidige arbejde vil fokusere p{\aa} optimering af hyperparametre og eksperimenter i stor skala ved hj{\ae}lp af denne metode.
diff --git a/chapters/experiments.tex b/chapters/experiments.tex
@@ -24,7 +24,7 @@ This experiment will then answer two questions:
\begin{marginfigure}
\centering
\includegraphics{hessian_profile.pdf}
- \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for Hessian-vector and GGN-vector products.}
+ \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for Hessian-vector, GGN-vector, and manual matrix-vector products. Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse performed via conjugate gradient for Hessian-vector and GGN-vector products. We can see that the explicit Hessian-vector product fails for problems with more than 10\,000 dimensions and that the GVP is faster than the HVP.}
\label{fig:hessian-profile}
\end{marginfigure}
@@ -55,7 +55,7 @@ comment on weird dip in inverse, possibly something to do with caching or someth
\begin{marginfigure}%[-4cm]
\centering
\includegraphics{hessian_profile_ratios.pdf}
- \caption[Comparison of speedup from HVP to GVP.]{Comparison of the speedup from computation of the Hessian-vector product (HVP) versus the GGN-vector product (GVP) (top) and inverse HVP versus inverse GVP using the conjugate gradient method (bottom). Speedup is calculated as the ratio of the wall-clock time of the HVP to the GVP. It can be seen that the two methods have the same time complexity (up to a linear factor).}
+ \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}
@@ -73,7 +73,7 @@ This is in line with the theory, as noted in \cref{sec:practical-ggn} (SHOULD IT
\begin{marginfigure}
\centering
\includegraphics{mnist_ablation.pdf}
- \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST.}
+ \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\).}
\label{fig:mnist-ablation}
\end{marginfigure}
@@ -92,7 +92,7 @@ The results of this experiment are shown in \cref{tab:mnist-ablation,fig:mnist-a
\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\). Lower p-value indicates a higher probability of rejecting normality. Bold indicates sampling was successful based on visual inspection of Q--Q plots.}
+ \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}
@@ -132,7 +132,7 @@ The choice of \(\alpha\) is a non-trivial one
\begin{figure}
\centering
\includegraphics{chisq_sin.pdf}
- \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for visualising normality of 1000 posterior samples from the full Laplace approximation for a fully connected neural network modelling \(f(x) = \sin(x)\). Since the sample square Mahalanobis distances appear chi-squared distributed, the samples themselves are multivariate normal with correct (fixed) covariance. As can be seen by looking at the histogram and the Q--Q plot, samples are multivariate normal.}
+ \caption[Chi-squared sample plot for sine curve.]{Chi-squared sample plot for visualising normality of 1000 posterior samples from the full Laplace approximation for a fully connected neural network modelling \(f(x) = \sin(x)\). Since the sample square Mahalanobis distances appear chi-squared distributed, the samples themselves are multivariate normal with correct (fixed) covariance. As can be seen by looking at the histogram and the quantile--quantile plot, samples are multivariate normal.}
\label{fig:chisq-sin}
\end{figure}
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -19,7 +19,7 @@ These two cases correspond to the mean-squared error (MSE) and cross-entropy (CE
Due to the assumption of independence between observations, we can factorise the likelihood and the negative log-likelihood loss as
\begin{align}\label{eq:negative-log-likelihood}
- p(\bm y \given \bm \theta) ={} & \prod_i^N p(x_i \given \bm \theta), \nonumber
+ p(\bm y \given \bm \theta) ={} & \prod_i^N p(x_i \given \bm \theta),
\\ \mathcal{L}(\bm\theta) \equiv \mathcal{L}(\bm \theta, \bm y) \coloneqq -\log p(\bm y \given \bm \theta) ={}& - \sum_i^N \log p(y_i \given \bm \theta).
\end{align}
This allows us to compute the loss as the sum of the losses for each individual observation.
@@ -98,11 +98,11 @@ The moving averages are weighted by the parameters \(\beta_1\) and \(\beta_2\) r
The learning rate is then scaled by the ratio of the first and second moments.
The update rule for Adam is thus given by
\begin{align}
- \bm\theta_{t+1} ={} & \bm\theta_t - \eta_t m_t \odot \frac{1}{\sqrt{v_t} + \varepsilon},
- \\ m_t ={}& (1 - \beta_1^t) \left[ \beta_1 m_{t-1} + (1 - \beta_1) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
- \\ v_t ={}& (1 - \beta_2^t) \left[ \beta_2 v_{t-1} + (1 - \beta_2) \left( \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right)^2 \right],
+ \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],
\end{align}
-where \(\odot\) denotes element-wise multiplication, \(\varepsilon\) is a small constant to prevent division by zero, and \(m_t\) and \(v_t\) are the first and second moments of the gradient respectively.
+where \(\odot\) denotes element-wise multiplication, \(\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}.
@@ -167,9 +167,9 @@ Since the Taylor expansion is performed around the MAP solution, the first order
\ln p(\bm \theta \given \bm y)
\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_\textsc{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
+ \\ \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),
\end{align}
-where \(\tilde{p}(\bm \theta \given \bm y)\) corresponds to the unnormalised posterior.
+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
@@ -244,17 +244,17 @@ This approximation is given by
Notice that \(\nabla^2_{\bm \theta} f_{\bm \theta_0}(\bm x) = 0\) under the Taylor expansion assumption of neural network linearity.
By combining \cref{eq:ggn-whole-hessian,eq:nn-taylor} we then obtain the generalised Gauss-Newton approximation of the Hessian,
%
-\begin{align*}
+\begin{align}
\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*}
+\end{align}
which corresponds to the first 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
\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}.
\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}
@@ -265,10 +265,10 @@ where we have defined \(\bm J_i = J_{\bm \theta} f(\bm x_i, \bm \theta)\) and \(
For mean square error (MSE) loss, we have one output variable (\(O = 1\)), so \(f: \reals^N \times \reals^D \rightarrow \reals\).
We then get the simple expression for the GGN approximation of the Hessian for MSE loss,
%
-\begin{align*}%\label{eq:mse-hessian}
- \bm H_i \approx & {} (y_i - f(\bm x_i, \bm \theta))^2 = 2 \nonumber
- \\ \bm \Lambda \approx&{} 2 \sum_{i=1}^N \bm J_i\T \bm J_i + \alpha \identity.
-\end{align*}
+\begin{align}%\label{eq:mse-hessian}
+ \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.
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.
@@ -296,12 +296,12 @@ Another choice is to linearise the whole likelihood function with respect to the
This yields \(g(y) = -\log(y)\), and thus we obtain the GGN approximation of the Hessian with this composition as
\begin{align}\label{eq:fully-linear-ggn}
\nabla^2_{\bm \theta} \mathcal{L}
- ={} & \nabla_{\bm \theta} f \T \cdot \nabla_{f}^2 g \cdot \nabla_{\bm \theta} f \nonumber
- \\ ={}& \nabla_{\bm \theta} f \T \cdot f(\bm\theta)^{-2} \cdot \nabla_{\bm \theta} f \nonumber
+ ={} & \nabla_{\bm \theta} f \T \cdot \nabla_{f}^2 g \cdot \nabla_{\bm \theta} f
+ ={} \nabla_{\bm \theta} f \T \cdot f(\bm\theta)^{-2} \cdot \nabla_{\bm \theta} f
\\ ={}& \left(f(\bm\theta)^{-1} \cdot \nabla_{\bm \theta} f\right)\T \cdot \left(f(\bm\theta)^{-1} \cdot \nabla_{\bm \theta} f\right) \nonumber
\\ ={}& \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right)\T \cdot \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right) \nonumber
\\ ={}& \left[\nabla_{\bm\theta} (g \circ f) \right]\T \cdot \left[\nabla_{\bm\theta} (g \circ f) \right] \nonumber
- \\ ={}& \nabla_{\bm\theta} \mathcal{L}\T \cdot \nabla_{\bm\theta} \mathcal{L}
+ ={} \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}.
diff --git a/chapters/sampling.tex b/chapters/sampling.tex
@@ -116,7 +116,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}, \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\).
@@ -152,7 +152,7 @@ This yields the quadrature rule result \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1}
This change of variables from \(\tau\) to \(\sigma\) 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) \cdot \cn(i u_q \mathcal{K}^\prime (k) \given k) \cdot \dn (i u_q \mathcal{K}^\prime (k) \given k)
+ w_q = \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cdot \cn(i u_q \mathcal{K}^\prime (k) \given k) \cdot \dn (i u_q \mathcal{K}^\prime (k) \given k),
\end{align}
and a set of quadrature points (shifts) given by
\begin{align}\label{eq:quadrature-shifts}
@@ -166,7 +166,7 @@ The complete elliptic integral of the first kind is given by
where \(p = 1 - m\), and is defined in the domain \(0 < p \leq 1\).\sidenote{If \(p > 1\), then what should be computed is instead the identity \(\mathcal{K}(p) = \mathcal{K}(1/p) / \sqrt{p}\).}
This integral is approximated by
\begin{align}
- \mathcal{K}(p) \approx P(p) - \log(p) Q(p)
+ \mathcal{K}(p) \approx P(p) - \log(p) Q(p),
\end{align}
where \(P\) and \(Q\) are tenth-order polynomials.
Jacobian elliptic functions \(\sn(u \given m)\), \(\cn(u \given m)\), and \(\dn(u \given m)\)
@@ -235,7 +235,7 @@ One drawback of CIQ is that the accuracy its of approximation of \(\bm \Lambda^{
In the case of the Laplace approximation, as per \cref{eq:ggn-laplace}, we have that the precision matrix is given by
\begin{align}\label{eq:ggn-laplace}
\bm \Lambda = \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity,
- \\ \bm \theta \sim \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1})
+ \\ \bm \theta \sim \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}),
\end{align}
where \(N\) is the number of observations in your dataset.
Since each \(J_i\T H_i J_i\) is, for \(i \in 1 \ldots N\), a low-rank positive semi-definite matrix, it has \(O\) non-zero eigenvalues (see \cref{sec:spectrum-ggn}).
@@ -249,13 +249,13 @@ This is done by choosing \(P\) such that the two matrices are ``similar'', i.e.,
In this way, the preconditioner acts as an ``initial best guess'' of the value to be computed.
We are thus seeking a preconditioner \(P\) such that
\begin{align}
- \bm P^{-1} \bm A \approx{} & \identity \Rightarrow \bm P \approx \bm A
- \\ \bm P^{-1} \bm \Lambda ={}& \bm P^{-1} \left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right)
+ \bm P^{-1} \bm A \approx{} & \identity \Rightarrow \bm P \approx \bm A,
+ \\ \bm P^{-1} \bm \Lambda ={}& \bm P^{-1} \left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right),
\end{align}
where \(\bm \Lambda\) is the posterior precision matrix given by \cref{eq:ggn-laplace} to improve convergence of the msMINRES and Lanczos algorithms.
The eigenvalues of our GGN approximation are thus given by
\begin{align}
- \lambda\left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right) = \big\{ \underbrace{s_1, \ldots, s_K}_{K \leq N \cdot O}, \underbrace{\alpha, \ldots, \alpha}_{D - K} \big\}
+ \lambda\left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right) = \big\{ \underbrace{s_1, \ldots, s_K}_{K \leq N \cdot O}, \underbrace{\alpha, \ldots, \alpha}_{D - K} \big\},
\end{align}
where \(\{s_1, \ldots, s_K\}, K \leq N \cdot O\) are the non-zero eigenvalues of the rank-\(K\) matrix \(\sum^N_i \bm J_i\T \bm H_i \bm J_i\).
If a preconditioner approximates all eigenvalues and eigenvectors of a matrix, then the preconditioner and the matrix are equal and the preconditioner is exact.
@@ -346,7 +346,7 @@ However, this preconditioner will be very expensive to compute, since it effecti
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
\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
+ \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}
where \(\alpha\) is the prior precision.
@@ -379,20 +379,20 @@ 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
- \\ \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 \nonumber
+ \\ \approx{}& \bm L\T \bm L + \alpha \identity,
\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
\begin{align}
- \bm P ={} & \bm L\T \bm L + \alpha \identity
- \\ \bm P^{-1} ={}& \alpha^{-1} \identity - \alpha^{-2} \bm L\T \left( \identity + \alpha^{-1} \bm L \bm L\T \right)^{-1} \bm L
+ \bm P ={} & \bm L\T \bm L + \alpha \identity,
+ \\ \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}.
+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).
\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 )
+ \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 ).
\end{align}
Here it is possible to precompute the inverse factorised product given by \(\bm L\T ( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{D \times K}\), meaning that we only need to perform two matrix multiplications for both the preconditioner- and inverse preconditioner-vector products.
This would then require storing the factorisation \(\bm L \in \reals^{K \times D}\) and the inverse factorised product \(\reals^{D \times K}\) matrix.\sidenote{We can also precompute the inverse \(( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{K \times K}\), which means we would only need to store the factorisation \(\bm L \in \reals^{K \times D}\) and this inverse \(K \times K\) matrix, in which case we would just need to perform one more matrix multiplication (multiplying \(\bm L\T\) with this \(K \times K\) matrix) in every preconditioner call.}
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -16,7 +16,7 @@ where \(p(\bm y)\) is known as the evidence. It's also known as the \emph{margin
This term is typically intractable, since it involves integrating over the entire parameter space.
For this reason, we typically approximate the posterior by some tractable distribution, such as a Gaussian, and then use this to approximate the marginal likelihood.
-Since we're integrating over the posterior, maximising the marginal likelihood will seek minima that are flat, since the posterior will be very broad in these regions.
+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}.
Furthermore, \textcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing k-fold cross validation for all values \(k = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(k\) folds and across all values of \(k\)). TODO: FIX THIS TO BE MORE ACCURATE.
@@ -49,14 +49,14 @@ Taking the log of this approximate Laplace marginal, we obtain
\end{align}
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)
+ 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)
\\ ={}& \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)
+ 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,
\end{align}
where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of the likelihood.
@@ -127,9 +127,9 @@ and for the square trace we have
\end{align}
We can bound the log determinant of \(\bm \Lambda\) using \cref{eq:log-determinant-decomposition,eq:laplace-trace-bound,eq:laplace-square-trace-bound} as
\begin{align}\label{eq:laplace-log-determinant-bound}
- \log \det \bm \Lambda \leq{} & B_K\left(\mu_1, \mu_2, \beta\right) + (D - K) \log \alpha
- \\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha \nonumber
- \\ \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 \nonumber
+ \log \det \bm \Lambda \leq{} & B_K\left(\mu_1, \mu_2, \beta\right) + (D - K) \log \alpha,
+ \\ \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\).
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.
@@ -142,16 +142,16 @@ To compute these, we can use Hutchinson's trace estimator~\cite{hutchinson1990st
Let \(\bm \varepsilon_0 \sim \normal(\bm 0, \identity)\).
Since \(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} = \identity\), we compute \(\mu_1\) as
\begin{align}\label{eq:hutchinson-trace}
- \mu_1 ={} & \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) + K \alpha
- \\ \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) ={}& \Tr\left(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)\right) \nonumber
- \\ = \E{}{\Tr(\bm \varepsilon_0 \bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right))} \nonumber
+ \mu_1 ={} & \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) + K \alpha,
+ \\ \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) ={}& \Tr\left(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)\right)
+ \\ ={}& \E{}{\Tr(\bm \varepsilon_0 \bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right))} \nonumber
% \\ = \E{}{\Tr\left(\bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) \bm \varepsilon_0\right)} \nonumber
- ={} & \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \bm \varepsilon_0}. \nonumber
+ \\ ={}& \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \bm \varepsilon_0}. \nonumber
\end{align}
Similarly, we compute \(\mu_2\) as
\begin{align}\label{eq:hutchinson-square-trace}
- \mu_2 =\Tr\left(\left( {\textstyle {\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
- \\ \Tr\left(\left( {\textstyle {\textstyle \sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i \right)^2\right) = \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right)\T \bm \varepsilon_0}. \nonumber
+ \mu_2 =\Tr\left(\left( {\textstyle {\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,
+ \\ \Tr\left(\left( {\textstyle {\textstyle \sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i \right)^2\right) = \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right)\T \bm \varepsilon_0}.
\end{align}
We then apply a Monte Carlo approximation of the expectation by sampling \(S\) samples from \(\bm \varepsilon_0\sim \normal(\bm 0, \identity)\) to compute the traces.
Fortunately, this means that we only need to compute one GGN-vector product \(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \cdot \bm \varepsilon_0\) to estimate the \(\mu_1\) and \(\mu_2\) terms in \cref{eq:laplace-log-determinant-bound}.
@@ -184,7 +184,7 @@ The scaling is thus important for determining the relative importance of the pri
To scale the posterior precision correctly, we thus need to scale the likelihood term by \(N / B\).
We do this by scaling the trace and square traces for \cref{eq:laplace-log-determinant-bound} as
\begin{align}
- \Tr\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) ={} & \frac{N}{B} \Tr\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) \label{eq:laplace-trace-scaling}
+ \Tr\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) ={} & \frac{N}{B} \Tr\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) \label{eq:laplace-trace-scaling},
\\ \Tr\left(\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right) ={}& \frac{N^2}{B^2} \Tr\left(\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right). \label{eq:laplace-square-trace-scaling}
\end{align}
By combining \cref{eq:laplace-log-determinant-bound,eq:bai-golub-bound,eq:hutchinson-trace,eq:hutchinson-square-trace,eq:trace-rearrangement,eq:laplace-trace-scaling,eq:laplace-square-trace-scaling}, we can therefore compute a differentiable upper bound on the log-marginal likelihood under the Laplace approximation using only Hessian-vector products without the need to instantiate the posterior precision.
diff --git a/tables/mnist_ablation.tex b/tables/mnist_ablation.tex
@@ -1,6 +1,6 @@
\begin{tabular}{rrr}
\toprule
- \(q\) & p-value & Time [m] \\
+ \(Q\) & p-value & Time [m] \\
\midrule
1 & 0.0000 & 10.6 \\
2 & 0.0000 & 18.8 \\
