masters-thesis

laplace.tex (41491B)

---

 \setchapterpreamble[u]{\margintoc}
 \chapter{The Laplace Approximation}
 \labch{laplace}
 % http://saul.cpsc.ucalgary.ca/pmwiki.php/GradTips/GradTipsLiteratureReview
 
 In this chapter, we discuss training and inference in deep learning from a Bayesian perspective.
 We introduce the Laplace approximation as a method for approximating the posterior distribution of a model's parameters given the data.
 We discuss the construction of the Laplace posterior, its properties, and how it can be computed in practice.
 
 \section{Optimisation in Deep Learning}[Optimisation]
 \label{sec:optimisation-frequentist}
 
 Suppose a neural network is a real-valued function \(f: \reals^N \times \reals^D \rightarrow \reals^O\) parametrised in \(\bm\theta \in \reals^D\) which maps an input \(\bm x \in \reals^N\) to an output \(f(\bm x, \bm\theta) \equiv f_{\bm\theta}(\bm x)\).
 Our goal is to find the parameters \(\bm\theta^*\) which best model the observed data.
 From a frequentist perspective, we can define a loss function on the model \(\mathcal{L}(\bm\theta): \reals^D \rightarrow \reals\) such that the optimal parameters \(\bm\theta^*\) minimise this loss function for the given model.
 Often, we formulate the loss function as the negative log-likelihood of the data under the model, \(\mathcal{L}(\bm\theta) = -\log \lik\).
 In this case, the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the likelihood of the data.\sidenote{Maximising the likelihood is equivalent to minimising the negative log-likelihood, since the log function is strictly monotonically increasing.}
 
 The likelihood function is problem-specific, and is defined as the probability density of the data given the parameters \(\bm\theta\).
 Typically, the likelihood is chosen to be a Gaussian distribution with mean \(f_{\bm\theta}(\bm x)\) and precision \(\rho\) (which is frequently chosen to be equal to one) for regression problems, and a categorical distribution with logits \(f_{\bm\theta}(\bm x)\) for classification problems.
 These two cases correspond to the mean-squared error (MSE) and cross-entropy (CE) negative log-likelihood loss functions, respectively.
 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(\bm y_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(\bm y_i \given \bm \theta).
 \end{align}
 This allows us to compute the loss as the sum of the losses for each individual observation.
 This loss function is defined in an optimisation problem to find the parameters \(\bm\theta^*\) of the model which minimise this loss, thereby maximising the performance of the model on the data.
 
 \subsection{Gradient Descent}
 
 Gradient descent is a first-order optimisation algorithm which iteratively updates the parameters \(\bm\theta\) in the direction of the negative gradient of the loss function \(\mathcal{L}(\bm\theta)\).\sidenote{First-order optimisation methods only use the first derivative of the loss function.}
 The update rule is given by
 %
 \begin{align}
     \bm\theta_{t+1} = \bm\theta_t - \eta_t \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t),
 \end{align}
 %
 where \(\eta_t\) is the learning rate at iteration \(t\).
 The learning rate \(\eta_t\) can be constant or adaptive, and is often chosen to be a non-increasing function of \(t\).
 The gradient \(\nabla_{\bm\theta} \mathcal{L}(\bm\theta_t, \bm x_t, y_t)\) can be computed separately for each individual training example \(\bm x_t, y_t\), as seen in \cref{eq:negative-log-likelihood}.
 Using the chain rule on the composition of the loss into the loss as a function of the model output \(\mathcal{L}\) and the model output as a function of the parameters \(f\) such that \(\mathcal{L}(\bm \theta_t, \bm x_t, y_t) = \mathcal{L}(y_t, f(\bm x_t, \bm \theta_t))\),we compute the gradient as
 %
 \begin{align}\label{eq:sgd-gradient}
     \nabla_{f} \mathcal{L}(\bm\theta_t, \bm x_t, y_t) ={} & \nabla_{\bm\theta} \log p(y_t \given \bm \theta) \nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t),
 \end{align}
 %
 where \(\nabla_{\bm\theta} f_{\bm\theta_t}(\bm x_t)\) is the gradient of the neural network output with respect to its parameters at iteration \(t\) and \(\nabla_{\bm\theta} \log p(y_t \given \bm \theta)\) is the gradient of the loss function with respect to the model output.
 Note that in this case, we perform one gradient update for each data point, which is why we iterate over the data points and update the parameters in the same step \(t\).
 Thus, gradient descent only requires the gradient of the loss function with respect to the parameters of the neural network, which can be computed efficiently using automatic differentiation for a similar cost to a single forward pass through the network.
 However, gradient descent is prone to getting stuck in local minima and can be slow to converge.
 In practice, we cannot compute the gradient of the loss function for the entire dataset at once, since this would require storing the entire dataset in memory and computing the gradient over it.
 Instead, we compute the gradient of the loss for a subset of the data at each iteration (a batch), and then update the parameters based on this unbiased estimate of the gradient.
 This is known as stochastic gradient descent (SGD), and is the most common form of gradient descent used in practice.
 SGD helps avoid getting stuck in local minima, but, like gradient descent, does not take into account the curvature of the loss function, and so can overshoot the minimum.
 
 \subsection{Adam}
 \label{sec:adam}
 
 Adam~\cite{kingma2014adam} is a popular alternative to SGD for training neural networks.
 It is an approximate second-order optimisation method, which means that it approximates the curvature of the loss.
 It is a variant of SGD that uses the first and second moments of the gradient of the loss to scale the learning rate of each parameter.
 In Adam, the first moment of the gradient is approximated by a moving average of the gradient and the second moment of the gradient is approximated by a moving average of the squared gradient.
 The moving averages are weighted by the parameters \(\beta_1\) and \(\beta_2\).
 The learning rate is then scaled by the ratio of the first and second moments.
 The update rule for Adam is thus
 \begin{align}
     \bm\theta_{t+1} ={} & \bm\theta_t - \eta_t \bm m_t \odot \left(\sqrt{\bm v_t} + \varepsilon\right)^{-1},
     \\ \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) \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \odot \nabla_{\bm\theta} \mathcal{L}(\bm\theta_t) \right],
 \end{align}
 where \(\odot\) denotes the Hadamard product, \(\varepsilon\) is a small constant to prevent division by zero, and \(\bm m_t\) and \(\bm v_t\) are the first and second moments of the gradient respectively.
 The momentum term~\cite{polyak1964some} takes inspiration from physics, where the momentum of an object is the product of its mass and velocity.
 In Adam, the first moment of the gradient is used to compute the momentum of the parameters, increasing the convergence speed of the optimisation~\cite{sutskever2013importance}.
 The second moment of the gradient is also known as the variance of the gradient.
 Adam approximates the second moment by the diagonal of the Fisher information matrix where the loss and model function are assumed to be linear.
 This is a common approximation in practice, as this approximation is easy to compute.
 The Fisher information matrix is discussed further in \cref{sec:ggn}.
 The use of the second moment is inspired by optimisation methods such as Newton's method and the natural gradient~\cite{rattray1998natural}, which use the second derivative of the loss to compute the optimal step size by scaling the step by the inverse of the curvature of the loss.
 These optimisation methods identify a single set of parameters as the optimal solution.
 In the next section, we will discuss how to learn a distribution over parameters instead.
 
 \section{Maximum a Posteriori Estimation}[MAP Estimation]
 \label{sec:bayesian-deep-learning}
 
 Bayesian deep learning is a framework for deep learning that allows for modelling the parameters as random variables instead of as single optima in the loss landscape.
 From a Bayesian perspective, we define a prior distribution \(p(\bm\theta)\) and a likelihood function \(\lik\) such that, under maximum a posteriori estimation, the goal of optimisation is to find the parameters \(\bm\theta^*\) which maximise the posterior distribution \(p(\bm\theta \given \bm y)\).
 This way, we can quantify the uncertainty in our model parameters \(\bm\theta\) by computing the posterior distribution \(p(\bm\theta \given \bm y)\).
 It is obtained from Bayes' theorem, which allows for uncertainty quantification by
 \begin{align}\label{eq:bayes-theorem}
     p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta)\,p(\bm \theta)}{p(\bm y)},
 \end{align}
 where \(p(\bm \theta \given \bm y)\) is the posterior distribution over the parameters, \(p(\bm \theta)\) is the prior distribution over parameters, \(p(\bm y \given \bm \theta)\) is the likelihood of the data given the parameters, and \(p(\bm y)\) is the marginal likelihood of the data.
 The marginal likelihood is obtained by marginalising over the parameters, i.e., \(p(\bm y) = \int p(\bm y \given \bm \theta)\,p(\bm \theta)\,d\bm \theta\).
 Due to this integral, the marginal is often intractable.
 However, since the posterior distribution is proportional to the product of the likelihood by the prior, we can obtain the unnormalised posterior from the likelihood and the prior\marginnote{Since the marginal likelihood is a normalisation constant and does not depend on the parameters, it does not affect the optimisation problem.}
 \begin{align}\label{eq:unnormalised posterior}
     p(\bm \theta \given \bm y) \propto p(\bm y \given \bm \theta)\,p(\bm \theta) \eqqcolon \tilde{p}(\bm \theta \given \bm y) .
 \end{align}
 Since we want to determine the parameters which maximise the posterior distribution, we can typically use the same optimisation techniques as for frequentist learning, but with the negative log-posterior loss instead of the negative log-likelihood.
 Furthermore, since the unnormalised posterior from \cref{eq:unnormalised posterior} is proportional to the likelihood and the prior, the exact posterior itself need not be tractable.
 Similarly to the negative log-likelihood, we factorise the negative log-posterior
 \begin{align}\label{eq:negative-log-posterior}
     \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(\bm 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(\bm 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.
 Because of this, the prior is often chosen to be a Gaussian distribution with zero mean and a diagonal covariance matrix with some precision \(\alpha\), i.e., \(p(\bm\theta) = \normal(\bm\theta \given \bm 0, \alpha^{-1}\bm I)\).\sidenote{In the frequentist setting, this type of regularisation is known as weight decay or \(L_2\)-regularisation. Note that \(- 1 / N \log p(\bm y \given \bm \theta) - \gamma / 2 \norm{\bm \theta}^2\) is the typical formulation in this setting, where \(\gamma\) is the regularisation strength and is usually within the range of \(10^{-2}\) to \(10^{-4}\). In our setting, we are instead computing the log-posterior \(-\log p(\bm y \given \bm \theta) - \alpha / 2 \norm{\bm \theta}^2\), and so we must multiply the regularisation strength by \(N\) to obtain \(\alpha\).}
 The likelihood term is defined as in the frequentist setting (see \cref{sec:optimisation-frequentist}).
 However, the posterior distribution itself is typically intractable, so we approximate the posterior instead.
 
 We use the posterior to quantify the uncertainty in our model parameters.
 Often, all that is needed is a set of samples from the posterior.
 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}
     \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.
 There are several methods for approximating the posterior distribution, which we will discuss in the following sections.
 Some, like Markov chain Monte Carlo (MCMC)~\cite{gelman1995bayesian}, are exact methods, which can be used to sample from the posterior distribution directly (though they are often computationally expensive).
 However, they cannot be used to compute the posterior distribution itself directly.
 Others, like variational inference (VI)~\cite{jordan1999introduction,wainwright2008graphical}, are methods which attempt to find a tractable approximation to the posterior distribution by parametrising a family of approximate posteriors and optimising the parameters to minimise the KL divergence between the approximate posterior and the true posterior.
 The performance of these methods depends on the variational family used to approximate the posterior and the choice of the variational parameters.
 One common choice for selecting the variational family is to use a Gaussian distribution with diagonal covariance matrix, which is known as mean-field VI.\sidenote{The mean-field assumption refers to the assumption that the parameters are independent of each other which leads to a diagonal covariance matrix.}
 Many variational methods suffer from the difficulty of their optimisation problem, since they often attempt to train variance parameters, which are difficult to optimise~\cite{skafte2019reliable}.
 In this project, we focus on the Laplace approximation, in which the posterior is approximated by a Gaussian distribution, but where the variance is deduced rather than optimised.
 
 \section{The Laplace Approximation}[Laplace Approximation]
 \label{sec:laplace}
 
 In the Laplace approximation (LA)~\cite{laplace1774memoire,bishop1995neural,mackay2003information,daxberger2021laplace}, the posterior is approximated by a Gaussian, similarly to mean-field VI\@.
 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}
     \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}).
 Dividing 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}
     \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 to the 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 by sampling from the neural network at each iteration of training and evaluating the gradient at each sample~\cite{miani2022laplacian} or by estimating the log-marginal likelihood~\cite{daxberger2021laplace,immer2021scalable}, enabling training of hyperparameters and improved model selection.
 In this work, we approach the latter, as well as post-hoc Laplace.
 
 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}}}\).
 However, this approximation only holds when the posterior precision \(\bm \Lambda\) is positive definite.
 For BNNs, this is not guaranteed as the loss function is not convex with respect to the model parameters, which means that the normal distribution assumed by the Laplace approximation may not be valid.
 Furthermore, the Hessian is only ensured to be positive semi-definite at the MAP. This may not be realistic when training neural networks, which tend to have unreliable convergence~\cite{choromanska2015loss}.
 In the next sections, we will go into the computation of the Hessian for the negative log-posterior loss function from \cref{sec:bayesian-deep-learning} and show how the generalised Gauss-Newton approximation can be used to obtain a positive definite Hessian approximation, even when the Hessian itself is not positive definite.
 
 \subsection{The Hessian}
 \label{sec:hessian}
 
 The Hessian is a matrix of second-order partial derivatives of a scalar function.\sidenote{From now on, we will use the notation \(\nabla^2_{\bm\theta} f\) to represent the Hessian matrix of \(f\), where each element is given by \[\left(\nabla^2_{\bm\theta} f\right)_{i, j} = \frac{\partial^2 f}{\partial\theta_i \partial\theta_j}.\]}
 Suppose we have a function \(\mathcal{L}: \reals^D \rightarrow \reals\) parametrised by \(\bm \theta \in \reals^D\).
 Then the Hessian can be interpreted as the Jacobian matrix of the gradient of the function, as per \(J_{\bm \theta}(\nabla_{\bm\theta} \mathcal{L})\).
 
 We are interested in the Hessian of the loss function \(\mathcal{L}(\bm\theta)\) with respect to the parameters \(\bm\theta\).
 For a neural network with \(D\) parameters, the Hessian is therefore a \(D \times D\) square matrix.
 Furthermore, if all the neural network's second partial derivatives are continuous, then the Hessian is symmetric.
 In this case, the Hessian is generally dominated by the block-diagonal~\cite{martens2015optimizing}.
 However, the Hessian will only be positive definite if the loss function is a convex function in the model parameters \(\bm \theta\).
 
 Typically, the loss function is defined as the negative log-likelihood or negative log-posterior.
 In this framing, if we consider the negative log-posterior loss from \cref{eq:negative-log-posterior}, by the linearity of the derivative, we obtain
 \begin{align}\label{eq:negative-log-posterior-hessian}
     \begin{split}
         \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \nabla^2_{\bm\theta} \left( -\sum_i^N \log p(\bm y_i \given \bm \theta) - \log p(\bm \theta) \right)
         \\  ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta)
         \\  ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(\bm 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(\bm y_i \given \bm \theta) - \log p(\bm \theta) \right)
     % \\  ={}& -\nabla^2_{\bm\theta} \sum_i^N \log p(\bm y_i \given \bm \theta) - \nabla^2_{\bm\theta} \log p(\bm \theta) \nonumber
     % \\  ={}& -\sum_i^N \nabla^2_{\bm\theta} \log p(\bm 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, e.g., the Hessian of the MSE loss.
 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\).
 However, as we discussed previously, the Hessian of the negative log-likelihood is not always positive definite.
 In the next section, we will discuss how to obtain a positive definite Hessian approximation.
 
 \subsection{The Generalised Gauss-Newton Approximation}[The GGN Approximation]
 \label{sec:ggn}
 
 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}
     \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
 \begin{align}
     \nabla^2_{\bm \theta} \mathcal{L}_{ij} = \sum_{n=1}^O \frac{\partial^2}{\partial \theta_i \theta_j} f_n \cdot \frac{\partial}{\partial f_n} \mathcal{L} + \sum_{n,m=1}^O \frac{\partial}{\partial \theta_i} f_n \cdot \frac{\partial^2}{\partial f_n f_m} \mathcal{L} \cdot \frac{\partial}{\partial \theta_j} f_m,
 \end{align}
 where \(f_n\) is the \(n\)\textsuperscript{th} output of the neural network function \(f(\bm x, \bm \theta)\).
 
 Assume now we will approximate \(f\) using a first-order Taylor expansion.
 In this way, we are linearising our neural network function.
 This approximation is given by
 %
 \begin{align}\label{eq:nn-taylor}
     f_{\bm \theta}(\bm x) \approx{} f_{\bm \theta_0}(\bm x) + \nabla_{\bm \theta} f_{\bm \theta_0}(\bm x) \cdot (\bm \theta - \bm \theta_0) .
 \end{align}
 %
 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}
     \nabla^2_{\bm\theta} \mathcal{L}
     \approx{} & \nabla_{\bm\theta} f\T \cdot \nabla^2_f \mathcal{L} \cdot \nabla_{\bm\theta} f,
 \end{align}
 which corresponds to the 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}
     \begin{split}
         \nabla^2_{\bm\theta} \mathcal{L}(\bm \theta) ={} & \sum_i^N \nabla^2_{\bm\theta} \log p(\bm 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}(\bm 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}
     \bm \Lambda \approx \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity,
 \end{align}
 where we have defined \(\bm J_i = J_{\bm \theta} f(\bm x_i, \bm \theta)\) and \(\bm H_i = \nabla^2_f \mathcal{L}(\bm y_i, f(\bm x_i, \bm \theta))\).
 
 For mean square error (MSE) loss (e.g. for regression), we obtain the simple expression for the GGN approximation of the Hessian
 %
 \begin{align}\label{eq:mse-ggn}
     \bm H_i \approx & {}\, (\bm y_i - f(\bm x_i, \bm \theta))^2 = 2 \identity,
     \\ \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.\sidenote{This is equivalent to the loss function being a convex function of the model outputs.}
 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.
 
 It is also interesting to note that, for our definition of the loss, the GGN approximation of the Hessian is equivalent to the Fisher information matrix, which is a measure of the curvature of the log-likelihood function with respect to the model parameters.
 Here, the Fisher information matrix is given by
 \begin{align}
     \mathcal{F} = \sum_{i=1}^N \E{p(y \given \bm x_i, \theta)}{\nabla_{\bm \theta} \log p(y \given \bm x_i, \bm \theta) \cdot \nabla_{\bm \theta} \log p(y \given \bm x_i, \bm \theta)\T},
 \end{align}
 where \(\E{p(y \given \bm x_i, \theta)}{\bullet}\) denotes the expectation with respect to the likelihood function.
 The Fisher matrix (and its empirical equivalent) is relevant as it motivates second-order methods such as natural gradient descent (and thus the Adam optimiser) and the use of the GGN approximation.
 
 \subsection{Practical Considerations for the GGN}[Practical Considerations]
 \label{sec:practical-ggn}
 
 The generalised Gauss-Newton approximation requires that the function for which we would like to compute the Hessian, \(h(x)\), can be framed as a composition of two functions \(h(x) = g(f(x))\).
 % A GGN approximation is any approximation which fulfils \cref{eq:ggn-whole-hessian}.
 However, there are many possible choices of functions \(f\) and \(g\) which, when composed, yield \(h\), and these different compositions can yield wildly varying performance for the GGN approximation~\cite{kunstner2019limitations}.
 Given that, in our case, \(h\) is the loss function with respect to the model parameters of a neural network, a natural choice of functions (and the choice we have assumed above) are \(g: \reals^O \rightarrow \reals\) and \(f: \reals^D \rightarrow \reals^O\) where \(g\) is the loss function with respect to the model output and \(f\) is a function which maps the model parameters to a model output, for a given model input \(x\).
 This construction allows us to capture the curvature of the loss with respect to the model output, which can often be easily computed in a simple closed form, without requiring the computation of the curvature of the model function with respect to its parameters; this therefore corresponds to linearising the model function with a Taylor expansion as in \cref{eq:nn-taylor}.
 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}
     \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}.
 
 The choice of application of the generalised Gauss-Newton approximation is not based on having a reduced computational cost---both the Hessian and the GGN approximation require the same number of forward and backward passes.
 However, the GGN approximation is better behaved than the exact Hessian since it is guaranteed to be positive semidefinite for convex losses, thereby making it more suitable for practical applications.
 This is because the GGN approximation is designed to be positive semi-definite, while the exact Hessian can be vulnerable to negative curvature.
 This feature can limit the use of the exact Hessian for algorithms that require a positive definite Hessian, such as, in our case, the Laplace approximation.
 
 In practice, as we explained in \cref{ch:intro}, we are looking for methods which access the product of the GGN matrix with a vector, rather than the GGN matrix itself.
 This then means we want to compute
 \begin{align}
     \begin{split}
         \bm \Lambda \bm v ={} & \left(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\right) \bm v
         \\  ={}& \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i \bm v + \alpha \bm v,
     \end{split}
 \end{align}
 such that we actually compute the product of the Jacobian of the neural network (as defined in \cref{eq:negative-log-likelihood-ggn}) by a vector \(\bm v\), i.e., \(\bm J_i \bm v\).
 
 \subsection{The Spectrum of the GGN Approximation}[The Spectrum of the GGN]
 \label{sec:spectrum-ggn}
 
 When we construct the posterior precision for the Laplace approximation, we compute the posterior precision as in \cref{eq:laplace-precision-ggn}.
 For methods which we will introduce later, we are interested in the spectrum (i.e., the set of eigenvalues) of the posterior precision.
 The eigenvalues of this precision will thus be the sum of the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) and of \(\alpha \identity\), where \(\alpha > 0\) is the prior precision.
 The first term, being the positive semi-definite GGN matrix, will have non-negative eigenvalues.
 The second term, being a diagonal matrix with entries \(\alpha\), will have eigenvalues that are all equal to \(\alpha\).
 Therefore, the eigenvalues of the posterior precision will be the sum of the eigenvalues of the GGN matrix and \(\alpha\) and will thus all be positive.
 This means that the prior precision \(\alpha\) guarantees the positive definiteness of the posterior precision.
 For some applications, it is desirable to have a weak prior, so that the regularisation effect of the prior is not too strong.
 While, theoretically speaking, the eigenvalues of the likelihood GGN matrix are greater than or equal to zero, in practice, due to numerical issues, the eigenvalues can be slightly negative.
 As such, \(\alpha\) must be chosen to be sufficiently large to ensure that the posterior precision is positive definite.
 While this may appear to detract from the effectiveness of the prior, it is important to note that, in most cases, the prior precision is large enough to ensure that the posterior precision is positive definite.
 
 Since \(\bm J_i\) is a rectangular matrix of shape \(O \times D\), the rank of each \(\bm J_i\T \bm H_i \bm J_i\) is at most \(O\), where \(O\) is the number of outputs of the model.
 The rank of the sum of these matrices thus depends on the eigenvectors of each \(\bm J_i\T \bm H_i \bm J_i\).
 If we add \(\bm J_n\T \bm H_n \bm J_n\) into the sum \(\sum_{i=1}^{n-1} \bm J_i\T \bm H_i \bm J_i\), then, if any of the eigenvectors of \(\bm J_n\T \bm H_n \bm J_n\) are not contained in the span of the eigenvectors of the sum, then the rank of the sum will increase.
 This is because the rank of a matrix is the number of linearly independent columns.
 Thus, non-rigorously, we have that the closer the eigenvectors (assuming that the eigenvectors have the same ordering between the two matrices) are between terms in the sum, the closer the largest eigenvalue will be to the sum of the largest eigenvalues of each term.\sidenote{By ``same ordering'' we mean that they are eigenvectors with respect to the same relative eigenvalue.}
 Practically speaking, this means that the rank of the sum of the GGN matrices can increase as we add more data points (but will not be greater than \(N \cdot O\)).
 Furthermore, as we add more data points, the eigenvalues of the sum of the GGN matrices will increase.
 Specifically, the largest eigenvalue of the sum of the GGN matrices will be the less than or equal to the sum of the largest eigenvalues of the GGN matrices.
 Additionally, the non-zero spectrum of the sum of the GGN matrices will itself get steeper as we add more data points.
 If the \(O\) non-zero eigenvectors are identical across the GGN matrices, then this non-zero spectrum of the sum of the GGN matrices will just be the sum of the spectrum of each GGN matrix and should thus not get steeper.\sidenote{If the GGN matrices \(\bm J\T \bm H \bm J\) have an identical spectrum
     \begin{align*}
         \{ \underbrace{s_1, s_2, \ldots, s_O}_{O}, \underbrace{0, \ldots, 0}_{D - O} \},
     \end{align*}
     then the sum of GGN matrices as per \(\sum^N \bm J\T \bm H \bm J = N \bm J\T \bm H \bm J\) has the spectrum
     \begin{align*}
         \{ \underbrace{N s_1, \ldots, N s_O}_{O}, \underbrace{0, \ldots, 0}_{D - O} \}.
     \end{align*}
 }
 However, if these eigenvectors are not identical, then as the largest eigenvalue of the sum of the GGN increases and the number of non-zero eigenvalues of the sum of the GGN increases, the non-zero spectrum of the sum of the GGN matrices will get steeper.
 For these reasons, as we add more data points, the posterior precision will become more ill-conditioned.\sidenote{A matrix is considered ill-conditioned when the ratio of its largest eigenvalue to its lowest eigenvalue, its \emph{condition number} (\(\kappa \coloneqq \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}\)), is very large.}
 If the sum of the GGN matrices is indeed low-rank, then the prior precision \(\alpha\) will be the smallest eigenvalue of the posterior precision.
 
 Empirically, we find that, for MNIST, the effective rank of the posterior precision is significantly less than \(N \cdot O\) (see \cref{fig:spectrum-mnist}).
 This suggests that the eigenvectors of the posterior precision are relatively similar across different observations, though not identical.
 However, this seems to greatly depend on the structure of the neural network (e.g., the filter size) in ways that are not yet clear.
 
 \begin{marginfigure}
     \centering
     \includegraphics{mnist_spectrum_3vs5.pdf}
     \caption[Spectrum of the posterior precision on MNIST.]{Spectrum of the Laplace posterior precision on 100 MNIST observations with a CNN. Note the log scale and the fact that we are zoomed in on the largest eigenvalues, as can be seen from the eigenvalue indices on the x-axis. For the \(3 \times 3\) filter, there are two distinct clusters in the spectrum, with the first cluster being over an order of magnitude larger than the second. This first cluster consists of nine eigenvalues, and so the precision has an effective rank close to nine. For the \(5 \times 5\) filter, we observe a hyper-exponential distribution of eigenvalues in this spectrum, which does not have any distinct clusters. This suggests that the effective rank of the precision is over 50, which is still significantly lower than the number of observations \(N = 100\) times the number of outputs \(O = 10\).}
     \label{fig:spectrum-mnist}
 \end{marginfigure}
 
 \section{Limitations}
 
 One limitation of the Laplace approximation is that it can only approximate one mode of the posterior distribution (the MAP).
 However, most approximate posterior distributions, such as MCMC and variational inference, also struggle to approximate multimodal distributions.
 The standard way to solve this problem is to use a multimodal approximation such as a mixture, typically by training multiple models from different initialisations and then combining them in a deep ensemble.
 This is also possible with the Laplace approximation~\cite{eschenhagen2021mixtures}.
 However, deep ensembles are quite expensive, as they require training multiple models.
 Additionally, it is not clear whether there is a significant benefit to using a multimodal Laplace approximation over a single mode Laplace approximation (i.e., whether generalisation improves).
 
 Another limitation of the Laplace approximation is that it is not scalable to non-trivial problems for the full Laplace approximation without performing crude approximations.
 This is because the full Laplace approximation requires the computation of the posterior precision, which is a \(D \times D\) matrix, where \(D\) is the number of parameters.
 Thus, in order to scale to large models and large datasets without approximating the posterior precision with, e.g., a diagonal matrix, we need to find a method to perform the Laplace approximation without instantiating the posterior precision.
 During training, we typically only need to train the model in a deterministic way (to find the MAP), and then we can use the Laplace approximation to sample from the approximate posterior post-hoc.
 Thus, for inference, we need to determine a method for sampling from the Gaussian approximate posterior without explicitly computing the posterior precision.
 We have so far only considered the Laplace approximation computed by maximising posterior (i.e., MAP estimation), but it is also possible to perform the Laplace approximation by training a model on its marginal likelihood (i.e., the likelihood of the data under the model).
 This would reduce overfitting and allow us to optimise hyperparameters via backpropagation.
 However, we must find a differentiable way to compute the Laplace approximate log-marginal likelihood without explicitly storing the posterior precision.