commit c8a9b994df54f16ea7e8edf90558dfe019b488b5
parent 67c36c33f35fc14721b35c2bcd5c3cde326b40bb
Author: Silas Brack <s174433@student.dtu.dk>
Date: Thu, 9 Feb 2023 13:06:21 +0100
Just missing LSL in intro.
Diffstat:
9 files changed, 212 insertions(+), 114 deletions(-)
diff --git a/chapters/abstract.tex b/chapters/abstract.tex
@@ -9,43 +9,41 @@ However, Bayesian inference is often computationally intractable, so approximate
The Laplace approximation (LA) is one such method, though it itself typically relies on crude approximations to the posterior precision matrix to make it computationally feasible.
When we do not factor the posterior precision into more practical approximations, the full posterior precision must be instantiated.
We call this the full Laplace approximation.
-For all but toy problems, however, the full Laplace approximation is intractable, since the posterior precision matrix scales quadratically with the number of parameters.
+For all but toy problems, since the full precision matrix scales quadratically with the number of parameters, the full posterior precision cannot be stored in memory.
-In this work, we propose a method for computing the Laplace approximation using only Hessian-vector products, which allows us to perform marginal training and posterior sampling using the full Laplace approximation without storing the entire Hessian.
-To accomplish this, we show that it is possible to estimate the Hessian's log-determinant and inverse square root using only Hessian-vector products.
-To overcome conditioning issues in the Hessian matrix, we propose several potential preconditioners which can be used to improve convergence during sampling.
-Next, we suggest a technique for evaluating the quality of posterior samples without instantiating nor inverting the Hessian based on results from traditional statistics.
-To implement these methods, we use JAX, a library for automatic differentiation which enables us to efficiently perform Hessian-vector products without explicitly storing the entire Hessian.
+In this work, we propose a method for computing the Laplace approximation using only Jacobian-vector products, which allows us to perform marginal training and posterior sampling using the full Laplace approximation without storing the entire posterior precision matrix.
+To accomplish this, we show that we can estimate the posterior precision's log-determinant and inverse square root using only Jacobian-vector products.
+To overcome the conditioning issues that arise during sampling, we propose several potential preconditioners which can be used to improve convergence of the inverse square root approximation.
+Next, we suggest a technique for evaluating the quality of the posterior samples without instantiating nor inverting the posterior precision based on results from traditional statistics.
+To implement these methods, we use JAX, a library for automatic differentiation which enables us to efficiently perform Jacobian-vector products without explicitly storing the entire posterior precision.
We perform the full Laplace approximation on both a sine function and MNIST using our method.
This consists two steps, (a) training a neural network by maximising either its posterior probability or marginal likelihood and (b) sampling from this posterior.
-We find that our approximate maximum marginal training procedure is able to learn a set of parameters which yields performance comparable to that of the maximum posterior and maximum likelihood training procedures.
-By analysing quantile--quantile plots of our posterior samples and visualising our posterior predictive samples, we find that our approximate sampling method produces samples which are correctly distributed.
-We then benchmark the performance of Hessian-vector products and estimate that our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
-This is the first method for performing the Laplace approximation while only accessing the Hessian implicitly.
-
-% Our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full Hessian.
-% We compute the full Laplace approximation on both a sine function and MNIST.
-% To do this, we demonstrate that the samples generated by our approximate sampling procedure are correctly distributed and that training the marginal likelihood using our method is feasible and produces meaningful results.
-% This is the first method for performing the Laplace approximation while only accessing the Hessian implicitly.
-
-% \chapter*{Resum\'e}
-% \addcontentsline{toc}{chapter}{Resum\'e}
-
-% 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.
+Our approximate maximum marginal training procedure is able to learn a set of parameters which yields performance comparable to that of the maximum posterior and maximum likelihood training procedures.
+By analysing quantile--quantile plots of our posterior samples and visualising these samples, we find that our approximate sampling method produces samples which are correctly distributed.
+Our benchmarks of the performance of Jacobian-vector products estimate that our method yields a 10\,000x memory reduction over the conventional full Laplace approximation, since it does not store the quadratically-scaling full posterior precision matrix.
+This is the first method for performing the Laplace approximation while only accessing the posterior precision implicitly.
+
+\chapter*{Resum\'e}
+\addcontentsline{toc}{chapter}{Resum\'e}
+
+Bayesianske metoder lover at give en principiel m{\aa}de at kvantificere usikkerheden i neurale netv{\ae}rk p{\aa}
+Dette er vigtigt for mange anvendelser inden for maskinl{\ae}ring, f.eks. i forbindelse med sikkerhedskritiske beslutninger, som er afh{\ae}ngige af risikovurdering og fortolkningsmuligheder.
+Bayesiansk inferens er imidlertid ofte beregningsm{\ae}ssigt uoverkommelig, s{\aa} der anvendes tiln{\ae}rmede metoder, som g{\aa}r p{\aa} kompromis med ydeevnen for at opn{\aa} praktiske fordele.
+Laplace-approximationen (LA) er en s{\aa}dan metode, selv om den typisk er afh{\ae}ngig af grove tiln{\ae}rmelser af den efterf{\o}lgende pr{\ae}cisionsmatrix for at g{\o}re den beregningsm{\ae}ssigt gennemf{\o}rlig.
+N{\aa}r vi ikke faktoriserer den efterf{\o}lgende pr{\ae}cision i mere praktiske tiln{\ae}rmelser, skal den fulde efterf{\o}lgende pr{\ae}cision instantieres.
+Vi kalder dette den fulde Laplace-approximation.
+Da den fulde pr{\ae}cisionsmatrix skalerer kvadratisk med antallet af parametre, kan den fulde posteriorpr{\ae}cision ikke lagres i hukommelsen for alle problemer undtagen leget{\o}jsproblemer, da den fulde pr{\ae}cisionsmatrix skalerer kvadratisk med antallet af parametre.
+
+I dette arbejde foresl{\aa}r vi en metode til beregning af Laplace-approksimationen udelukkende ved hj{\ae}lp af jacobian-vektorprodukter, som g{\o}r det muligt at udf{\o}re marginal tr{\ae}ning og posterior sampling ved hj{\ae}lp af den fulde Laplace-approksimation uden at lagre hele matrixen for posterior pr{\ae}cision.
+For at opn{\aa} dette viser vi, at vi kan estimere den posteriore pr{\ae}cisionens log-determinant og den inverse kvadratrod ved hj{\ae}lp af kun jacobianiske vektorprodukter.
+For at overvinde de konditioneringsproblemer, der opst{\aa}r under pr{\o}veudtagning, foresl{\aa}r vi flere potentielle pr{\ae}konditioneringsv{\ae}rkt{\o}jer, som kan anvendes til at forbedre konvergensen af den inverse kvadratrodsapproximation.
+Dern{\ae}st foresl{\aa}r vi en teknik til at evaluere kvaliteten af de efterf{\o}lgende stikpr{\o}ver uden at instantiere eller invertere den efterf{\o}lgende pr{\ae}cision baseret p{\aa} resultater fra traditionel statistik.
+Til at gennemf{\o}re disse metoder anvender vi JAX, et bibliotek til automatisk differentiering, som g{\o}r det muligt at udf{\o}re jacobian-vektorprodukter effektivt uden eksplicit lagring af hele den efterf{\o}lgende pr{\ae}cision.
+
+Vi udf{\o}rer den fulde Laplace-approximation p{\aa} b{\aa}de en sinusfunktion og MNIST ved hj{\ae}lp af vores metode.
+Denne best{\aa}r af to trin: (a) tr{\ae}ning af et neuralt netv{\ae}rk ved at maksimere enten dets efterf{\o}lgende sandsynlighed eller marginale sandsynlighed og (b) pr{\o}veudtagning fra denne efterf{\o}lgende sandsynlighed.
+Vi finder, at vores tiln{\ae}rmede maksimale marginale tr{\ae}ningsprocedure er i stand til at l{\ae}re et s{\ae}t parametre, som giver en ydeevne, der kan sammenlignes med den maksimale posterior- og maksimale sandsynlighedstr{\ae}ningsprocedure.
+Ved at analysere quantile--quantile-plots af vores posteriorpr{\o}ver og visualisere disse pr{\o}ver finder vi, at vores tiln{\ae}rmede pr{\o}veudtagningsmetode giver pr{\o}ver, der er korrekt fordelt.
+Vores benchmarks af Jacobian-vektorprodukternes ydeevne ansl{\aa}r, at vores metode giver en 10\ 000x hukommelsesreduktion i forhold til den konventionelle fulde Laplace-approximation, da den ikke lagrer den kvadratisk skalerende matrix med fuld pr{\ae}cision for den efterf{\o}lgende periode.
+Dette er den f{\o}rste metode til at udf{\o}re Laplace-approximationen, hvor der kun er implicit adgang til den efterf{\o}lgende pr{\ae}cision.
diff --git a/chapters/introduction.tex b/chapters/introduction.tex
@@ -19,10 +19,10 @@
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 using Bayes' theorem to model 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.
-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 understanding risk and interpretability, where traditional deterministic neural networks are overconfident or fail to provide uncertainty estimates at all.
+From here we can sample from this distribution to make predictions, thereby yielding uncertainty estimates, thus enabling the assessment of model predictive confidence.
+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 understanding risk, where traditional deterministic neural networks are overconfident or fail to provide uncertainty estimates at all.
The estimation of model confidence is particularly important in domains where silent failures can have catastrophic consequences~\cite{jospin2022hands}, such as autonomous driving~\cite{rao2018deep}, medical diagnosis~\cite{ker2017deep}, and financial markets~\cite{cavalcante2016computational}.
-The improved interpretability, however, is very valuable in the context of deep learning, where models are black-boxes which are difficult to interpret.
+% The improved interpretability, however, is very valuable in the context of deep learning, where models are black-boxes which are difficult to interpret.
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.
This is because Bayesian inference requires computing the posterior distribution of the model parameters given the data, which, due to the nonlinear complexity of the neural network architecture, is often intractable.
@@ -48,7 +48,16 @@ The Laplace approximation has been applied in, for example, the prediction of ea
The primary advantages of the Laplace approximation are that it is simple to implement, effective at out-of-distribution detection~\cite{miani2022laplacian} and can be performed post-hoc on a pre-trained model.
This last point is important, since it also means that the Laplace approximation can be as cheap to train as the optimization algorithms for deterministic neural networks, while only requiring a small amount of additional computation to compute the posterior samples for inference.
-In the Laplace approximation, to compute the posterior distribution of the model parameters, we need to compute the Hessian of the loss with respect to the model parameters.
+In the Laplace approximation corresponds to performing a second-order Taylor expansion of the log-posterior around the \emph{maximum a posteriori} (MAP) estimate of the model parameters as
+\begin{align}
+ \log p(\bm \theta\,|\,\bm y)
+ \approx{} & \log p(\bm \theta_{\textsc{map}}\,|\,\bm y) + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y) (\bm \theta - \bm \theta_{\textsc{map}}). \nonumber
+\end{align}
+This is equivalent to approximating the log-posterior as a Gaussian distribution, with mean and covariance given by the MAP estimate and the Hessian of the log-posterior respectively, i.e.,
+\begin{align}
+ p(\bm \theta\,|\,\bm y) \approx{} & \normal\left(\bm \theta\,\Big|\,\bm \theta_{\textsc{map}}, \left(- \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y)\right)^{-1}\right). \nonumber
+\end{align}
+To compute the posterior distribution of the model parameters, we thus need to compute the Hessian of the negative log-posterior loss with respect to the model parameters.
Typically, this is done by storing the Hessian matrix explicitly and computing the inverse of this matrix.
Since the Hessian matrix is a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters, this is intractable for large models, which can contain millions or billions of parameters, thus limiting the applicability of full Laplace outside of toy problems of a few thousand parameters.
To overcome this, we can use a crude approximation, such as only storing the Hessian matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.
@@ -66,21 +75,28 @@ To train using the negative log-marginal likelihood loss for the Laplace approxi
However, for large models, which can contain millions or billions of parameters, we cannot instantiate the Hessian matrix, for reasons mentioned above.
Instead, we would like to access the inverse square root of the Hessian without explicitly storing it.
-Our method allows us approximate the Hessian's inverse square root and log-determinant using only implicit Hessian-vector products (HVPs).
+Our method will allow us to approximate the Hessian's inverse square root and log-determinant using only implicit Hessian-vector products (HVPs).
These HVPs compute products of the Hessian matrix with a vector without explicitly storing the Hessian matrix itself.
Since the HVP outputs a vector, the memory requirements are linear in the number of parameters, rather than quadratic, allowing us to perform the full Laplace approximation with significantly larger models.
We use JAX~\cite{jax2018github}, a Python library for GPU-accelerated automatic differentiation and computational graph compilation, to efficiently compute Hessian-vector while only storing the implicit Hessian computational graph.
JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the Hessian-vector product computation into efficient machine code.
-We have thus developed a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
+We will thus develop a method which performs the full Laplace approximation without requiring the Hessian to be stored explicitly.
To do this, we will use JAX and XLA to compute Hessian-vector products without explicitly storing the Hessian matrix.
We will then exploit these Hessian-vector products to compute the log-determinant and inverse square root of the Hessian.
These quantities enable us to perform training and inference using the Laplace approximation.
An overview of the Laplace approximation and our contributions is shown in \cref{fig:overview}.
-\begin{figure}
+\begin{marginfigure}[-15cm]
\centering
\input{figures/overview.tex}
\caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Hessian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the Hessian. Our contributions allow us to forgo the typical crude approximations of the Hessian by performing the full Laplace approximation without explicitly storing the Hessian matrix.}
\label{fig:overview}
-\end{figure}
+\end{marginfigure}
+
+% \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 in green the computations which we approximate in this project using only Hessian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the Hessian. Our contributions allow us to forgo the typical crude approximations of the Hessian by performing the full Laplace approximation without explicitly storing the Hessian matrix.}
+% \label{fig:overview}
+% \end{figure}
diff --git a/chapters/laplace.tex b/chapters/laplace.tex
@@ -105,8 +105,12 @@ Furthermore, since the unnormalised posterior from \cref{eq:unnormalised posteri
Similarly to the negative log-likelihood, we can factorise the negative log-posterior loss as per
\begin{align}\label{eq:negative-log-posterior}
- \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
- \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}. \nonumber
+ \begin{split}
+ \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
+ \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}.
+ \end{split}
+ % \mathcal{L}(\bm\theta) = -\log p(\bm \theta \given \bm y) ={} & - \log p(\bm y \given \bm \theta) - \log p(\bm\theta)
+ % \\ ={}& \underbrace{-\sum_i^N \log p(y_i \given \bm \theta)}_{\text{Negative log-likelihood}} - \underbrace{\log p(\bm\theta)}_{\text{Regularisation}}. \nonumber
\end{align}
Here the regularisation term is given by the prior distribution \(p(\bm\theta)\).
Due to the black-box nature of neural networks, it is often difficult to define a prior distribution over the parameters using previous knowledge.
@@ -118,8 +122,12 @@ We use the posterior distribution to quantify the uncertainty in our model param
Often, all that is needed is a set of samples from the posterior distribution.
From here, we can make predictions in a Bayesian setting by computing the posterior predictive distribution and using Monte Carlo sampling as
\begin{align}\label{eq:posterior predictive}
- p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
- \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y), \nonumber
+ \begin{split}
+ p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
+ \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y),
+ \end{split}
+ % p(y^* \given \bm y) ={} & \int p(y^* \given \bm x^*, \bm \theta)\,p(\bm \theta \given \bm y)\,d\bm \theta,
+ % \\ \approx{}& \frac{1}{S} \sum_{s=1}^S p(y^* \given \bm x^*, \bm \theta_s), \quad \bm \theta_s \sim p(\bm \theta \given \bm y), \nonumber
\end{align}
where \(y^*\) is the new data point and \(S\) is the number of Monte Carlo posterior samples \(\bm \theta_s\).
However, in some cases, it is useful to have a closed-form expression for the posterior distribution.
@@ -139,21 +147,27 @@ In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mack
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_\textsc{map}}\).
Since the Taylor expansion is performed around the MAP solution, the first order derivative is zero, and the expansion is given by
\begin{align}
- \log p(\bm \theta \given \bm y)
- \approx{} & \log 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} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_\textsc{map}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) \label{eq:laplace-taylor}
+ \begin{split}
+ \log p(\bm \theta \given \bm y)
+ \approx{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) +
+ \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \left( \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_\textsc{map}} \right) (\bm \theta - \bm \theta_{\textsc{map}}) \label{eq:laplace-taylor}
+ \end{split}
\\ \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 (\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)
- \\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}). \nonumber
+ \begin{split}
+ 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)
+ \\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}).
+ \end{split}
+ % 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)
+ % \\={}& \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1}). \nonumber
\end{align}
The Laplace approximation can simply be trained as an MAP solution, where the Hessian is computed to obtain the posterior precision at inference time.
This method is known as post-hoc Laplace.
However, it is also possible to compute the posterior precision during training.
-This can be done either to sample from the neural network at each iteration of training and evaluating the Hessian at each sample~\cite{miani2022laplacian} or to estimate the log-marginal likelihood~\cite{daxberger2021laplace,immer2021scalable}, enabling training of hyperparameters and improved model selection.
+This can be done either by sampling from the neural network at each iteration of training and evaluating the Hessian at each sample~\cite{miani2022laplacian} or by estimating the log-marginal likelihood~\cite{daxberger2021laplace,immer2021scalable}, enabling training of hyperparameters and improved model selection.
We approach the latter, as well as post-hoc Laplace, in this work.
We have now shown that approximating the log-posterior with a second degree Taylor expansion around the \(\bm\theta_{\textsc{map}}\) corresponds to approximating the posterior with a Gaussian distribution given by \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\) where \(\bm \Lambda = - \left.\nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y) \right|_{\bm \theta_{\textsc{map}}}\).
@@ -178,9 +192,14 @@ However, the Hessian will only be positive definite if the loss function is a co
Typically, the loss function is considered as the negative log-likelihood or negative log-posterior.
In this framing, if considering, e.g., the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we get
\begin{align}\label{eq:negative-log-posterior-hessian}
- \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
- \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
- \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
+ \begin{split}
+ \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
+ \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta)
+ \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta).
+ \end{split}
+ % \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(y_i \given \bm \theta) - \log p(\bm \theta) \right)
+ % \\ ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
+ % \\ ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta). \nonumber
\end{align}
In this construction, the first term is the Hessian of the negative log-likelihood, i.e., the Hessian of the loss function we are interested in.
For the common prior given by a zero-mean Gaussian \(\bm \theta \sim \normal(\bm 0, \alpha^{-1} \identity)\), the log-prior regularisation term is given by \(\nabla^2_{\bm\theta} \log p(\bm \theta) = - \alpha \identity\).
@@ -193,11 +212,18 @@ In the next section, we will discuss how to obtain a positive definite Hessian a
From the expression of the posterior precision as the Hessian of the negative log-posterior in \cref{eq:negative-log-posterior-hessian} we can apply the chain rule twice and the product rule once to obtain a simpler expression,
%
\begin{align}\label{eq:ggn-whole-hessian}
- \nabla^2_{\bm \theta} \mathcal{L}
- ={} & J_{\bm \theta}(\nabla_{\bm \theta} \mathcal{L})
- = J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot J_{\bm \theta} f\right)
- \\ ={}& J_{\bm \theta} (J_{\bm \theta} f) \cdot \nabla_f \mathcal{L} + J_{\bm \theta} \left(\nabla_f \mathcal{L}\right) \cdot J_{\bm \theta} f \nonumber
- \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + J_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot J_{\bm \theta} f, \nonumber
+ \begin{split}
+ \nabla^2_{\bm \theta} \mathcal{L}
+ ={} & J_{\bm \theta}(\nabla_{\bm \theta} \mathcal{L})
+ = J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot J_{\bm \theta} f\right)
+ \\ ={}& J_{\bm \theta} (J_{\bm \theta} f) \cdot \nabla_f \mathcal{L} + J_{\bm \theta} \left(\nabla_f \mathcal{L}\right) \cdot J_{\bm \theta} f
+ \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + J_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot J_{\bm \theta} f,
+ \end{split}
+ % \nabla^2_{\bm \theta} \mathcal{L}
+ % ={} & J_{\bm \theta}(\nabla_{\bm \theta} \mathcal{L})
+ % = J_{\bm \theta}\left(\nabla_f \mathcal{L} \cdot J_{\bm \theta} f\right)
+ % \\ ={}& J_{\bm \theta} (J_{\bm \theta} f) \cdot \nabla_f \mathcal{L} + J_{\bm \theta} \left(\nabla_f \mathcal{L}\right) \cdot J_{\bm \theta} f \nonumber
+ % \\ ={}& \nabla^2_{\bm \theta} f \cdot \nabla_f \mathcal{L} + J_{\bm \theta} f \T \cdot \nabla^2_f \mathcal{L} \cdot J_{\bm \theta} f, \nonumber
\end{align}
where \(\nabla_f\) is the gradient with respect to \(f(\bm x, \bm \theta)\) and \(J_{\bm \theta}\) is the Jacobian with respect to \(\bm \theta\).
More precisely, we obtain
@@ -225,8 +251,10 @@ which corresponds to the second term in \cref{eq:ggn-whole-hessian}.
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}. \nonumber
+ \begin{split}
+ \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}.
+ \end{split}
\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,13 +293,15 @@ However, this is not the only possible construction.
Another choice is to linearise the whole likelihood function with respect to the model parameters.
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
- ={} \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}.
+ \begin{split}
+ \nabla^2_{\bm \theta} \mathcal{L}
+ ={} & \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)
+ \\ ={}& \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right)\T \cdot \left(-\nabla_f g \cdot \nabla_{\bm\theta} f\right)
+ \\ ={}& \left[\nabla_{\bm\theta} (g \circ f) \right]\T \cdot \left[\nabla_{\bm\theta} (g \circ f) \right]
+ ={} \nabla_{\bm\theta} \mathcal{L}\T \cdot \nabla_{\bm\theta} \mathcal{L}.
+ \end{split}
\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 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
@@ -123,10 +123,12 @@ 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}, \nonumber %\label{eq:ciq-quadrature-tau}
+ \begin{split}
+ \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}
+ \end{split}
\end{align}
-where \(\tau_1, \ldots, \tau_Q \in \Gamma\) are the sampled quadrature points and \(\Tilde{w}_1, \ldots, \Tilde{w}_Q \in \reals\) are the quadrature weights.
+where we have the sampled quadrature points \(\tau_1, \ldots, \tau_Q \in \Gamma\) and the quadrature weights \(\Tilde{w}_1, \ldots, \Tilde{w}_Q \in \reals\).
% 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\).
Since we will be sampling a fixed number of quadrature points from the circle, the closer the contour is to the singular points (i.e., the eigenvalues), the better the approximation to the integral will be.
@@ -142,9 +144,11 @@ Since our GGN precision matrix \(\bm \Lambda = \sum^N_{i=1} \bm J_i\T \bm H_i \b
This is accomplished by applying a change of variable from \(\tau\) to \(\sigma = \tau^{1/2}\) and sampling the quadrature points from \(\sigma\).
To then account for this change of variable, we reformulate \cref{eq:ciq-integral-tau} as % eq:ciq-quadrature-tau
\begin{align}\label{eq:ciq-integral-sigma}
- \bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau
- \\ ={}& \frac{1}{\pi i} \oint_{\Gamma_\sigma} \left( \sigma^2 \identity - \bm K \right)^{-1}\,d\sigma \nonumber
- \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1}. \nonumber % \label{eq:ciq-quadrature-sigma}
+ \begin{split}
+ \bm K^{-1/2} ={} & \frac{1}{2 \pi i} \oint_\Gamma \tau^{-1/2} \left(\tau \identity - \bm K\right)^{-1}\,d\tau
+ \\ ={}& \frac{1}{\pi i} \oint_{\Gamma_\sigma} \left( \sigma^2 \identity - \bm K \right)^{-1}\,d\sigma
+ \\ \approx{}& \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \left( \sigma^2_q \identity - \bm K \right)^{-1}.
+ \end{split}
\end{align}
We can now adapt the equation to allow for computation using only matrix-vector products.
From here, the inverse square root vector product is given by
@@ -161,7 +165,7 @@ This yields the quadrature rule result \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1}
\begin{marginfigure}
\includegraphics[scale=1]{ellipk.pdf}
- \caption[Comparison of SciPy and JAX \texttt{ellipk} implementations.]{Comparison of SciPy and JAX \texttt{ellipk} implementations. For values for which the complete elliptic integral \(K(p)\) is well-defined, both implementations yield the same results. Furthermore, due to its static nature, JAX implementation is compatible with \emph{just-in-time} compilation.}
+ \caption[Comparison of SciPy and JAX \texttt{ellipk} implementations.]{Comparison of SciPy and JAX \texttt{ellipk} implementations. For values for which the complete elliptic integral \(\mathcal{K}(p)\) is well-defined, both implementations yield the same results. Furthermore, due to its static nature, JAX implementation is compatible with \emph{just-in-time} compilation.}
\label{fig:ellipk}
\end{marginfigure}
@@ -386,8 +390,10 @@ 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, \nonumber
+ \begin{split}
+ \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,
+ \end{split}
\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
@@ -474,8 +480,10 @@ As per the definition of the chi-squared distribution, for \(D\)-dimensional sta
Samples \(\bm \varepsilon\) which are normally distributed but do not have a mean of zero and a standard deviation of one can be standardised and their sum of square deviations be calculated by computing their squared Mahalanobis distances, which will be chi-squared distributed with \(D\) degrees of freedom, as per
%
\begin{align}
- \bm d ={} & \sqrt{(\bm \varepsilon - \bm \mu)\T \bm \Lambda (\bm \varepsilon - \bm \mu)}\label{eq:mahalanobis-distance}
- \\ \Rightarrow{}& \bm d\T \bm d \sim \chi^2(D). \nonumber
+ \begin{split}
+ \bm d ={} & \sqrt{(\bm \varepsilon - \bm \mu)\T \bm \Lambda (\bm \varepsilon - \bm \mu)}\label{eq:mahalanobis-distance}
+ \\ \Rightarrow{}& \bm d\T \bm d \sim \chi^2(D).
+ \end{split}
\end{align}
%
This has the advantage of only requiring computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it doesn't require instantiation of the matrix.
diff --git a/chapters/training.tex b/chapters/training.tex
@@ -37,17 +37,19 @@ This is known as \emph{online} Laplace, and is an alternative to the \emph{post-
For the Laplace approximation, we approximate the posterior distribution as \(p(\bm \theta \given \bm y) \approx \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\), where the posterior precision \(\bm \Lambda = -\nabla^2_{\bm\theta} \mathcal L(\bm \theta_{\textsc{map}})\) is the Hessian of the negative log-posterior at the maximum a posteriori (MAP) estimate \(\bm \theta_{\textsc{map}}\).
As in \cref{eq:laplace-taylor}, we have the approximate Laplace log-posterior
\begin{align}\label{eq:laplace-posterior}
- \log p(\bm \theta \given \bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) \cdot (\bm \theta - \bm \theta_{\textsc{map}})
- \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}})\nonumber
- \\={}& \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
+ \begin{split}
+ \log p(\bm \theta \given \bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) \cdot (\bm \theta - \bm \theta_{\textsc{map}})
+ \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}})
+ \\={}& \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}}).
+ \end{split}
\end{align}
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
- % \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) \int_{\bm \theta} \frac{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}}{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
- % \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) (2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2} \int_{\bm \theta} \frac{1}{(2\pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta\nonumber
- \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}.\nonumber
+ \begin{split}
+ 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
+ \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}.
+ \end{split}
\end{align}
Taking the log of this approximate Laplace marginal, we obtain
\begin{align}\label{eq:laplace-log-marginal}
@@ -58,7 +60,7 @@ 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) \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
+ \\ ={}& \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
@@ -69,9 +71,11 @@ and
where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of the likelihood.
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
- \\ & + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda, \nonumber
+ \begin{split}
+ \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
+ \\ & + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda,
+ \end{split}
\end{align}
where it can be seen that there are constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\).
While these terms are not necessary for the optimisation problem, we choose to include them.
@@ -118,8 +122,10 @@ 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 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
- \\ ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + (D - K) \log \alpha, \nonumber
+ \begin{split}
+ \log \det \bm \Lambda ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + \sum_{k=K+1}^D \log \alpha
+ \\ ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + (D - K) \log \alpha,
+ \end{split}
\end{align}
where \(\lambda_1, \ldots, \lambda_K\) are the non-zero eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\).
The first term is the log-determinant of a positive definite matrix with eigenvalues \(\{\lambda_1 + \alpha, \ldots, \lambda_K + \alpha\}\), and the second term is a constant.
@@ -127,8 +133,10 @@ The former can then be bounded by \cref{eq:bai-golub-bound} with an upper bound
We then compute the trace of this term as
\begin{align}\label{eq:laplace-trace-bound}
\Tr{\bm \Lambda} ={} & \sum_{k=1}^{K} (\lambda_k + \alpha) + (D - K) \alpha
- \\ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha \label{eq:laplace-trace-bound-2}
- \\ ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha, \nonumber
+ \\ \begin{split}
+ \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha \label{eq:laplace-trace-bound-2}
+ \\ ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,
+ \end{split}
\end{align}
and for the square trace we have
\begin{align}\label{eq:laplace-square-trace-bound}
@@ -152,10 +160,9 @@ 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)
- \\ ={}& \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
+ \\ \begin{split}\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))}
+ \\ ={}& \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}.\end{split}
\end{align}
Similarly, we compute \(\mu_2\) as
\begin{align}\label{eq:hutchinson-square-trace}
diff --git a/figures/overview.tex b/figures/overview.tex
@@ -3,20 +3,22 @@
font=\sffamily,
box/.style = {rounded corners, minimum width=22mm, minimum height=5mm, align=center},
]
+ \newcommand{\myrel}[2]{\mathrel{\ooalign{\raisebox{2.2\height}{#1}\cr\hidewidth#2\hidewidth\cr}}}
% \definecolor{color1}{rgb}{0.86, 0.82, 1.0}
% \definecolor{color2}{rgb}{0.97, 0.96, 1.0}
\definecolor{color1}{rgb}{0.74, 0.83, 0.9}
\definecolor{color2}{rgb}{0.94, 0.97, 1.0}
\definecolor{kaogreen}{HTML}{91C889}
\node (nn) [box,fill=color1,minimum height=1.0cm,minimum width=3cm] {Neural Network};
- \node (max) [below=of nn,yshift=0.5cm] {\(\argmax_{\bm \theta}\)};
- \node (posterior) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=-2cm,yshift=0.5cm] {\(\log p(\bm \theta \given \bm y)\)};
- \node (marginal) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=2cm,yshift=0.5cm] {\(\overbrace{\log p(\bm \theta \given \bm y) \quad \text{\colorbox{kaogreen}{$\log \det \bm \Lambda$}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
+ \node (max) [below=of nn,yshift=0.1cm] {\(\argmax_{\bm \theta}\)};
+ \node (posterior) [box,fill=color2,minimum width=1.8cm,minimum height=2.1cm,below=of max,xshift=-1.1cm,yshift=0.5cm] {\(\log p(\bm \theta \given \bm y)\)};
+ \node (marginal) [box,fill=color2,minimum width=1.8cm,minimum height=2.1cm,below=of max,xshift=1.1cm,yshift=0.5cm] {\(\overbrace{\myrel{\text{\normalsize \(\log p(\bm \theta \given \bm y)\)}}{\text{\normalsize \colorbox{kaogreen}{$\log \det \bm \Lambda$}}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
+ % \node (marginal) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=2cm,yshift=0.5cm] {\(\overbrace{\log p(\bm \theta \given \bm y) \quad \text{\colorbox{kaogreen}{$\log \det \bm \Lambda$}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
\draw[dashed] (nn) to (max);
- \draw[dashed,->] (max) to [bend right] (posterior);
- \draw[dashed,->] (max) to [bend left] (marginal);
+ \draw[dashed,->] (max) to [bend left] (posterior.60);
+ \draw[dashed,->] (max) to [bend right] (marginal.120);
\begin{scope}[on background layer]
- \node (training) [box,fill=color1,minimum width=8.05cm,inner sep=0.25cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
+ \node (training) [box,fill=color1,minimum width=3.8cm,inner sep=0.25cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
\end{scope}
\node [above=0cm and 0cm of training.north west,anchor=south west] {Training};
\node [above=0cm and 0cm of posterior.north west,anchor=south west] {{Post-Hoc}};
@@ -34,3 +36,41 @@
\draw[dashed,->] (map) to (inference);
% \node [above=0cm and 0cm of inference.north east,anchor=south east] {\cref{ch:sampling}};
\end{tikzpicture}
+
+
+% \begin{tikzpicture}
+% [
+% font=\sffamily,
+% box/.style = {rounded corners, minimum width=22mm, minimum height=5mm, align=center},
+% ]
+% % \definecolor{color1}{rgb}{0.86, 0.82, 1.0}
+% % \definecolor{color2}{rgb}{0.97, 0.96, 1.0}
+% \definecolor{color1}{rgb}{0.74, 0.83, 0.9}
+% \definecolor{color2}{rgb}{0.94, 0.97, 1.0}
+% \definecolor{kaogreen}{HTML}{91C889}
+% \node (nn) [box,fill=color1,minimum height=1.0cm,minimum width=3cm] {Neural Network};
+% \node (max) [below=of nn,yshift=0.5cm] {\(\argmax_{\bm \theta}\)};
+% \node (posterior) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=-2cm,yshift=0.5cm] {\(\log p(\bm \theta \given \bm y)\)};
+% \node (marginal) [box,fill=color2,minimum width=3.8cm,minimum height=2cm,below=of max,xshift=2cm,yshift=0.5cm] {\(\overbrace{\log p(\bm \theta \given \bm y) \quad \text{\colorbox{kaogreen}{$\log \det \bm \Lambda$}}}^{\text{\normalsize \(\log p(\bm y)\)}}\)};
+% \draw[dashed] (nn) to (max);
+% \draw[dashed,->] (max) to [bend right] (posterior);
+% \draw[dashed,->] (max) to [bend left] (marginal);
+% \begin{scope}[on background layer]
+% \node (training) [box,fill=color1,minimum width=8.05cm,inner sep=0.25cm,minimum height=3.2cm,fit={(max)(posterior)(marginal)}] {};
+% \end{scope}
+% \node [above=0cm and 0cm of training.north west,anchor=south west] {Training};
+% \node [above=0cm and 0cm of posterior.north west,anchor=south west] {{Post-Hoc}};
+% \node [above=0cm and 0cm of marginal.north east,anchor=south east] {Online};
+% % \node [above=0cm and 0cm of training.north east,anchor=south east] {\cref{ch:training}};
+% %
+% \node (map) [below=of training,yshift=0.7cm] {\(\bm \theta_{\textsc{map}}\)};
+% \draw[dashed] (training) to (map);
+% \node (sample) [below=of map,yshift=0.5cm] {Sample};
+% \node (samplingeq) [below=0cm and 0cm of sample,box,fill=color2,minimum width=3.8cm,minimum height=2cm] {\(\begin{aligned}\bm \varepsilon \sim{}& \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}) \\ ={}& \bm \theta_{\textsc{map}} + \text{\colorbox{kaogreen}{$\bm \Lambda^{-1/2} \bm \varepsilon_0$}}\end{aligned}\)};
+% \begin{scope}[on background layer]
+% \node (inference) [box,fill=color1,inner sep=0.25cm,fit={(sample)(samplingeq)}] {}; % minimum width=8.05cm,minimum height=3.2cm,
+% \end{scope}
+% \node [above=0cm and 0cm of inference.north west,anchor=south west] {Inference};
+% \draw[dashed,->] (map) to (inference);
+% % \node [above=0cm and 0cm of inference.north east,anchor=south east] {\cref{ch:sampling}};
+% \end{tikzpicture}
diff --git a/kaorefs.sty b/kaorefs.sty
@@ -196,9 +196,9 @@
\newcommand{\reftab}[1]{\hyperref[tab:#1]{\tablename}\xspace\ref{tab:#1}}
\newcommand{\vreftab}[1]{\hyperref[tab:#1]{\tablename\xspace\vref{tab:#1}}}
-\newcommand{\refeqshort}[1]{\hyperref[eq:#1]\eqnameshort\xspace(\ref{eq:#1})}
-\newcommand{\refeq}[1]{\hyperref[eq:#1]\eqname\xspace\ref{eq:#1}}
-\newcommand{\vrefeq}[1]{\hyperref[eq:#1]\eqname\xspace\vref{eq:#1}}
+% \newcommand{\refeqshort}[1]{\hyperref[eq:#1]\eqnameshort\xspace(\ref{eq:#1})}
+% \newcommand{\refeq}[1]{\hyperref[eq:#1]\eqname\xspace\ref{eq:#1}}
+% \newcommand{\vrefeq}[1]{\hyperref[eq:#1]\eqname\xspace\vref{eq:#1}}
\newcommand{\refdef}[1]{\hyperref[def:#1]\defname\xspace\ref{def:#1}}
\newcommand{\vrefdef}[1]{\hyperref[def:#1]\defname\xspace\vref{def:#1}}
diff --git a/macros.tex b/macros.tex
@@ -22,8 +22,7 @@
\newcommand{\entropy}[1]{\mathbb{H}\left[#1\right]}
\newcommand{\crossentropy}[2]{\mathbb{H}\left[#1, #2\right]}
-\newcommand{\given}{\, | \,}
-% \newcommand{\given}{\mid}
+\newcommand{\given}{\;|\;}
\newcommand{\norm}[1]{\left\lVert#1\right\rVert}
\newcommand{\abs}[1]{\left\lvert#1\right\rvert}
diff --git a/main.tex b/main.tex
@@ -52,7 +52,7 @@
\usepackage{kaorefs}
\usepackage[utf8]{inputenc}
-\usepackage{amsmath,amssymb,amsthm,bm}
+\usepackage{amsmath,amssymb,amsthm,bm,mathtools}
\usepackage{censor}
\usepackage{tikz}
