masters-thesis

sampling.tex (53643B)

---

 \setchapterpreamble[u]{\margintoc}
 \chapter{Sampling}
 \labch{sampling}
 
 We will now discuss how to perform inference by sampling from the Laplace posterior distribution.
 We first give an overview of how we typically sample from the Laplace posterior and discuss the memory constraints of this sampling method when using the full posterior precision matrix.
 In \cref{sec:ciq}, we propose an algorithm for sampling from the Laplace posterior that does not require storing the full precision matrix, allowing us to scale sampling to larger models.
 Finally, we discuss how to mitigate some of the issues of our method (\cref{sec:preconditioning}) and propose a practical technique for determining whether sampling was successful (\cref{sec:sampling-evaluation}).
 
 \section{Sampling from the Laplace Posterior}[Posterior Sampling]
 
 In many methods, such as the Laplace approximation (\cref{sec:laplace}), we sample from a multivariate normal distribution parametrised by a mean vector \(\bm \mu\) and a covariance matrix \(\bm \Sigma\).
 To sample from a Gaussian distribution, you can apply the reparametrisation trick to compute a vector of samples as
 \begin{align}\label{eq:sample-reparam}
     \bm \varepsilon ={} & \bm \mu + \bm \Sigma^{1/2} \bm \varepsilon_0,
 \end{align}
 where \(\bm \varepsilon_0\) are samples drawn from a standard normal distribution, i.e., \(\bm \varepsilon_0 \sim \normal(\bm 0, \identity)\).
 
 Using the Laplace approximation, we estimate the precision matrix as the Hessian matrix of the loss with respect to the parameters, as discussed in \cref{sec:laplace}.
 Because of this, the covariance matrix is given by the inverse of this Hessian, and we thus sample from the Laplace-approximate posterior distribution as
 \begin{align}\label{eq:sample-laplace}
     \bm \varepsilon = \bm \theta_{\textsc{map}} + \bm \Lambda^{-1/2} \bm \varepsilon_0,
 \end{align}
 where \(\bm \Lambda = \nabla^2_{\bm\theta} \log p(\bm \theta \given \bm y)\).
 In \cref{eq:sample-reparam}, the square root of a matrix yields another matrix which, when multiplied by itself, is equal to the original matrix.
 However, since an isotropic Gaussian distribution is invariant to rotation, in this special case, we just need to obtain a matrix which returns the original matrix \emph{up to a rotation}.\sidenote{That is, if we have a vector of standard normal samples \(\bm \varepsilon \sim \normal(\bm 0, \identity)\), then, for an orthogonal matrix \(\bm R\) (i.e., a rotation matrix), \(\bm R\T \bm \varepsilon \sim \normal(\bm 0, \identity)\). We then have that we can sample \emph{non}-standard normal samples as \(\bm \Lambda^{-1/2} \bm R\T \bm \varepsilon \sim \normal(\bm 0, \bm \Lambda^{-1})\).}
 We thus need to find a method which computes the inverse square root of the full \(D \times D\) posterior precision matrix up to a rotation, where \(D\) is the number of model parameters.
 
 Some methods exist which compute the inverse square root of a matrix up to a rotation, such as the Cholesky decomposition.
 However, these methods are not suitable for large matrices, as they require the instantiation of the full precision matrix, which is not feasible for large matrices.
 There are also methods which can compute the inverse square root of a low-rank matrix that is constructed as the inner product of two rectangular matrices by performing the inverse square root on the outer product (which is low-dimensional for low-rank problems), such as the Woodbury matrix identity.
 However, the precision matrix we are interested in is not itself low-rank but a sum of low-rank matrices, and thus this method is not suitable for our problem.
 This is discussed in more detail in \cref{sec:laplace,sec:preconditioning}.
 % As such, these are not adequate methods for computing the inverse square root using only matrix-vector products.
 
 Furthermore, methods which do not require computation of the square root of the covariance matrix typically attempt to exploit properties such as a block structure or sparsity of the covariance matrix~\cite{vono2022high}.
 However, we do not have a good understanding of the structure of the Laplace covariance matrix, and thus these methods are not suitable for our problem.
 % Talk about what we do know about the structure? Wrt the sum of low-rank GGNs
 When sampling from high-dimensional Gaussian distributions, it is common to use Markov chain Monte Carlo (MCMC) sampling methods such as Gibbs sampling~\cite{vono2022high}.
 However, these methods are tricky to tune and tend to converge very slowly for high-dimensional problems.
 As such, we do not use these methods for sampling from the Laplace approximation, and instead sample using the reparametrisation trick.
 For this, we need a method to compute the inverse square root using only matrix-vector products.
 
 \section{Contour Integral Quadrature}[Contour Integral Quad]
 \label{sec:ciq}
 
 \textcite{pleiss2020fast} propose one such method.
 Contour integral quadrature (CIQ) attempts to solve the problem of computing the inverse square root of a matrix, up to a rotation, by approximating the expression \(f(\bm K) \bm \varepsilon_0 = \bm K^{-1/2} \bm \varepsilon_0\) where \(\bm K\) is a positive semi-definite matrix using Cauchy's integral formula.
 To approximate this integral, it exploits the Lanczos algorithm, Gaussian quadrature, and the msMINRES algorithm, methods which only require the computation of matrix-vector products \(\bm K \bm v\).
 There are then essentially three steps to CIQ, namely:
 \begin{enumerate}
     \item Perform the Lanczos algorithm to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).
     \item Sample \(Q\) quadrature points from a complex-plane circle to approximate the contour integral from Cauchy's integral formula, using these quadrature points to create \(Q\)  shifted linear systems.
     \item Solve these \(Q\) shifted systems using the msMINRES algorithm.
 \end{enumerate}
 
 \begin{algorithm}
     \caption{Contour Integral Quad (CIQ)}\label{alg:ciq}
     \Input{\(\bm K \succ 0\), \(\bm b\), \(\bm P \succ 0\), \(J > 0\), \(Q > 0\)}
     \Output{\(\bm s = \bm K^{-1/2} b\)}
     \(\_, \bm\alpha, \bm\beta \gets \mathrm{CG}(\bm K, \bm b, \bm P, J)\)\;
     \(\bm M_{\mathrm{Lanczos}} \gets \mathrm{TriDiag}(\bm\alpha, \bm\beta)\)\;
     \(\lambda_{\bm M} \gets \lambda(\bm M_{\mathrm{Lanczos}})\)\;
     \(\lambda_{\min} \gets \min{\lambda_{\bm M}}\)\;
     \(\lambda_{\max} \gets \max{\lambda_{\bm M}}\)\;
     \(w_1, \ldots, w_Q, t_1, \ldots, t_Q \gets \mathrm{Quadrature}(\lambda_{\min}, \lambda_{\max}, Q)\)\;
     \(\bm c_1, \ldots, \bm c_Q \gets \mathrm{msMINRES}(\bm K, \bm b, t_q, \bm P, J)\) \Comment*[r]{\(\left( t_q \identity + \bm K \right)^{-1} \bm b\).}
     % \(s_q \gets \left( t_q \identity + H \right)^{-1} b\) \Comment*[r]{Computed via msMINRES(K, t, J).}
     \(\bm s \gets \sum^Q_{q=1} w_q \bm c_q\)\;
     % \(\mathrm{samples} \gets \sum^Q_{q=1} w_q \left( t_q \identity + H \right)^{-1} b\)
 \end{algorithm}
 
 Since this procedure (see \cref{alg:ciq}) computes the approximate inverse square root product \(\bm K\) by \(\bm \varepsilon_0\) up to a rotation, i.e.,
 \begin{align}\label{eq:ciq-asdf}
     \bm K^{-1/2} \bm \varepsilon_0,
 \end{align}
 it is then also trivial to compute the square root product by multiplying \(\bm K\) onto this vector, as per
 \begin{align}
     \bm K^{1/2} \bm \varepsilon_0 = \bm K \cdot \underbrace{\bm K^{-1/2} \bm \varepsilon_0}_{\text{\cref{eq:ciq-asdf}}}.
 \end{align}
 \marginnote{The inverse square root function is holomorphic for all positive definite real matrices.}
 
 \def\ymax{2.0}
 \def\ymin{-2.0}
 \def\xmax{3.0}
 \def\xmin{-0.5}
 \def\R{1.3}
 \begin{marginfigure}%[-7cm]
     \centering
     \input{figures/cauchy_integral.tex}
     \caption[Integration circle for the Cauchy integral formula.]{Integration circle \(\Gamma\) contained within the domain of \(f\) as defined according to the Cauchy integral formula.}
     \label{fig:cauchy-integral-formula}
 \end{marginfigure}
 
 % ------------------------------------------
 % CAUCHY INTEGRAL FORMULA
 % ------------------------------------------
 
 The first step of CIQ invokes Cauchy's integral formula.
 Cauchy's integral formula is a central theorem in complex analysis which states that, for a holomorphic function \(f\), there exists a closed contour \(\Gamma\) in the complex plane (see \cref{fig:cauchy-integral-formula}) which encloses the eigenvalues of \(\bm K\) such that \(f(\bm K)\) can be approximated as
 \begin{align}\label{eq:cauchy-integral-formula}
     f(\bm K) ={} & \frac{1}{2 \pi i} \oint_\Gamma f(\tau) \left(\tau \identity - \bm K\right)^{-1}\,d\tau,
 \end{align}
 where \(\oint_\Gamma\) is the contour integral along the contour \(\Gamma\) parametrised in the complex variable \(\tau\).
 Since, in our case, the eigenvalues of \(\bm K\) will all be real-valued, a circle centred on the real axis which encloses the minimum and maximum eigenvalues of \(\bm K\) will suffice.\marginnote{If a real-valued matrix is symmetric, its eigenvalues are real. In our case, \(\bm K\) will be the precision matrix obtained from the Laplace approximation.}
 To estimate these values, we can use the Lanczos algorithm, a Krylov subspace method which adapts the power iteration method, to compute a lower bound on the maximum eigenvalue and an upper bound on the minimum eigenvalue of \(\bm K\).\sidenote{Krylov subspace methods are a class of iterative methods which estimate certain properties (e.g., the rank, eigenvalues) of a matrix \(\bm K\) with a Krylov subspace \(\mathcal{K}_j(\bm K, \bm v)\) for some vector \(\bm v\). This subspace is spanned by the set of images of \(\bm v\) under the first \(j\) powers of \(\bm K\), \(\mathcal{K}_j( \bm K, \bm v) = \mathrm{span}\{ \bm v, \bm K \bm v, \ldots, \bm K^{j-1} \bm v \}\).}
 
 % ------------------------------------------
 % QUADRATURE
 % ------------------------------------------
 
 \def\ymax{2.0}
 \def\ymin{-2.0}
 \def\xmax{3.0}
 \def\xmin{-0.5}
 \def\R{1.3}
 \begin{marginfigure}
     \centering
     \input{figures/cauchy_quadrature.tex}
     \caption[Sampling of quadrature points for CIQ.]{The shifts \(\tau_1, \ldots, \tau_Q\) are sampled from a circle with its centre on the real axis which intersects the real axis at the points \(\lambda_{\mathrm{min}}\) and \(\lambda_{\mathrm{max}}\). Since we use Jacobi elliptic functions, we over-sample quadrature points that are closer to the minimum eigenvalue \(\lambda_{\mathrm{min}}\).}
     \label{fig:ciq}
 \end{marginfigure}
 
 We then apply Cauchy's integral formula to \(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}
     \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 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.
 
 % ------------------------------------------
 % JACOBI ELLIPTICAL FUNCTIONS
 % ------------------------------------------
 
 We could sample these quadrature points from the circle uniformly, which corresponds to the regular trapezoid quadrature rule.
 However, the convergence for this quadrature is linear with regard to the condition number of \(\bm K\)~\cite{pleiss2020fast,hale2008computing}.
 Since the precision matrix (which will be computed from the Laplace approximation as the GGN matrix) can often have a low rank and poor conditioning, uniform sampling would then require a very large number of quadrature points \(Q\) and would therefore be inadequate~\cite{pleiss2020fast}.
 
 Since our GGN precision matrix \(\bm \Lambda = \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) is often low-rank, we would instead prefer to oversample quadrature points which lie close to the minimum eigenvalue.
 This is accomplished by applying a change of variable from \(\tau\) to \(\sigma = \tau^{1/2}\) and sampling the quadrature points from \(\sigma\).
 To 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}
     \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 this equation to allow for computation using only matrix-vector products.
 From here, the inverse square root vector product is given by
 \begin{align}\label{eq:ciq-sigma-product}
     \bm K^{-1/2} \bm v \approx{} & \frac{1}{\pi i} \sum_{q=1}^Q \Tilde{w}_q \underbrace{\left( \sigma^2_q \identity - \bm K \right)^{-1} \bm v}_{Q\text{ system solves}}.
 \end{align}
 We can then solve the \(Q\) systems of equations to obtain the solutions \(\bm c_q = \left(\tau_q \identity - \bm K \right)^{-1} \bm v\) for \(q \in 1, \ldots, Q\).
 
 Since the integrand is symmetric with respect to the real axis, only the imaginary component of \(\Gamma_\sigma\) is required.\sidenote{This is because the circle with its centre on the \(x\)-axis is symmetric with respect to said \(x\)-axis.}
 Since \(\sigma\) is then imaginary (its real part is zero), \(\tau \coloneqq \sigma^2\) is real such that \(\tau = - \Im(\sigma)^2\).
 As such, since \(\Im(\sigma)^2\) must be positive, \(\tau\) must then be real-valued and negative.
 In practice, we incorporate the \(1 / (\pi i)\) term into the quadrature weights and, since \(\Tilde{w}_q\) also happens to be real-valued and negative, we can redefine the quadrature weights \(w_q = -\Tilde{w}_q\) and the quadrature points \(t_q = -\sigma^2_q\) for simplicity.\marginnote{With this redefinition, both \(w_q > 0\) and \(t_q > 0\). This then means that \(\left(t_q \identity + \bm K\right)\) is positive definite.}
 This yields the quadrature rule result \(\bm K^{-1/2} \bm v \approx \sum^Q_{q=1} w_q \bm c_q\), where \(\bm c_q = (t_q \identity + \bm K)^{-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 \(\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}
 
 This change of variables from \(\tau\) to \(\sigma\) yields a set of quadrature weights \(w_1, \ldots, w_Q\) given by
 \begin{align}\label{eq:quadrature-weights}
     w_q = \frac{2 \sqrt{\lambda_\mathrm{min}}}{\pi Q} \mathcal{K}^\prime (k) \cdot \cn(i u_q \mathcal{K}^\prime (k) \given k) \cdot \dn (i u_q \mathcal{K}^\prime (k) \given k),
 \end{align}
 and a set of quadrature points (shifts) given by
 \begin{align}\label{eq:quadrature-shifts}
     \sigma^2_q = \lambda_{\mathrm{min}} \cdot \sn(i u_q \mathcal{K}^\prime (k) \given k)^2,
 \end{align}
 where \(\sn\), \(\cn\), and \(\dn\) are the Jacobi elliptic functions and \(\mathcal{K}^\prime(k)\) is the complete elliptic integral, such that \(\mathcal{K}^\prime(k) = \mathcal{K}(k^\prime) = \mathcal{K}(\sqrt{1 - k^2})\).
 The complete elliptic integral of the first kind is given by
 \begin{align}
     \mathcal{K}(p) = \int_0^{\pi/2} \left( 1 - m \sin(t)^2 \right)^{-1/2}\,dt,
 \end{align}
 where \(p = 1 - m\), and is defined in the domain \(0 < p \leq 1\).\sidenote{If \(p > 1\), then what we compute is instead the identity \(\mathcal{K}(p) = \mathcal{K}(1/p) / \sqrt{p}\).}
 This integral is approximated by
 \begin{align}
     \mathcal{K}(p) \approx P(p) - \log(p) Q(p),
 \end{align}
 where \(P\) and \(Q\) are tenth-order polynomials.
 Additionally, we need to compute the Jacobian elliptic functions \(\sn(u \given m)\), \(\cn(u \given m)\), and \(\dn(u \given m)\).
 To compute these values, we first calculate the elliptic modulus \(k = \kappa(\bm K)^{-1/2}\) (where \(\kappa(\bm K)\) denotes the condition number of \(\bm K\)) and \(u_q = (q - 1/2) / Q\).
 These functions are not implemented in JAX, and SciPy uses the Cephes numerical library in C to compute them.
 Because of this, we ported the Cephes Jacobi elliptical functions to JAX, and \cref{fig:ellipk,fig:ellipj} show the functions evaluated for the two implementations.
 Our static JAX implementation allows for efficient \emph{just-in-time} compilation, which is not possible with the SciPy implementation.
 We can thereby now compute the quadrature weights and shifts as described in \cref{alg:quad}.
 
 \begin{figure}
     \centering
     \includegraphics[scale=1]{ellipj.pdf}
     \caption[Comparison of SciPy and JAX \texttt{ellipj} implementations.]{Comparison of SciPy and JAX \texttt{ellipj} implementations. For values for which the functions are well-defined, both implementations yield the same results. Plotted are $\sn(u \given m = 0.7)$, $\cn(u \given m = 0.7)$, and $\dn(u \given m = 0.7)$. Furthermore, due to its static nature, JAX implementation is compatible with \emph{just-in-time} compilation.}
     \label{fig:ellipj}
 \end{figure}
 
 \begin{marginfigure}
     \begin{algorithm}[H]
         \caption{Quadrature}\label{alg:quad}
         \Input{\(\lambda_{\min}\), \(\lambda_{\max}\), \(Q > 0\)}
         \Output{\(w_q, t_q \given q \in 1, \ldots, Q\)}
         \(k^2 \gets \lambda_{\min} / \lambda_{\max}\)\;
         \({k^{\prime}}^2 \gets \sqrt{1 - k^2}\)\;
         \(K^\prime \gets \mathcal{K}({k^{\prime}}^2)\)\;
         \For{\(q \gets 1\) \KwTo \(Q\)}{
             \(u_q \gets (q - 1/2) / Q\)\;
             \(\overline{\sn}_q \gets \sn(u_q K^\prime \given {k^{\prime}}^2)\)\;
             \(\overline{\cn}_q \gets \cn(u_q K^\prime \given {k^{\prime}}^2)\)\;
             \(\overline{\dn}_q \gets \dn(u_q K^\prime \given {k^{\prime}}^2)\)\;
             \(\sn_q \gets i ( \overline{\sn}_q / \overline{\cn}_q )\)\;
             \(\dn_q \gets ( \overline{\dn}_q / \overline{\cn}_q )\)\;
             \(\cn_q \gets ( 1 / \overline{\cn}_q )\)\;
             \(w_q \gets \lambda_{\min}^{1/2} K^\prime \cn_q \dn_q\)\;
             \(w_q \gets 2 w_q / \left(\pi Q\right) \)\;
             \(t_q \gets -\lambda_{\min} \sn_q^2\)\;
         }
     \end{algorithm}
 \end{marginfigure}
 
 The last step involves solving the \(Q\) linear systems defined in \cref{eq:ciq-sigma-product} to determine \(\bm c_q\).  % Actually we solve the other linear system.
 As mentioned above, this is done by using msMINRES, an algorithm for solving linear systems of the form \((\bm K + t_q \identity) \bm c_q = \bm b\) where \(\bm K\) is symmetric.
 Like other Krylov methods, msMINRES computes the solutions to the linear systems by constructing a Krylov subspace \(\mathcal{K}_q\) of \(\bm K + t_q \identity\) and then iteratively minimises the norm of the residual in \(\mathcal{K}_q\).
 While msMINRES only requires that the matrices \(\bm K + t_q \identity\) be symmetric, and not necessarily also positive definite (since the shifts \(t_q\) are positive and \(\bm K\) is positive semidefinite), the sum is also positive definite.
 However, the CIQ implementation from \textcite{pleiss2020fast} uses msMINRES, and we follow their implementation.
 There is also evidence suggesting that msMINRES converges faster than, for example, the conjugate gradient method (CG)~\cite{fong2012cg}, a Krylov method that requires the matrices to be positive definite.
 
 Since the Lanczos algorithm only requires the matrix-vector product \(\bm K \bm v\) with a random vector \(\bm v\) to estimate a bound on the eigenvalues, we can use the same bound for all \(Q\) linear systems.
 As such, the computational bottleneck of the contour integral quadrature is performing the msMINRES for the shifted systems.
 Additionally, since msMINRES is an iterative Krylov method, its convergence bound is sensitive to the conditioning of the linear systems, i.e., the conditioning of \(\bm K\).
 To mitigate the effect of ill conditioning, we use a preconditioner \(\bm P\) such that \(\bm P^{-1} \bm K\) is well-conditioned.
 
 \section{Preconditioning}
 \label{sec:preconditioning}
 
 One drawback of CIQ is that the accuracy its of approximation of \(\bm \Lambda^{-1/2} \bm \varepsilon_0\) suffers when the precision matrix \(\bm \Lambda\) is poorly conditioned.
 In the case of the Laplace approximation, the precision matrix is given by
 \begin{align}\label{eq:ggn-laplace}
     \bm \Lambda = \sum_i^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity,
     \\  \bm \theta \sim \normal(\bm \theta_{\textsc{map}}, \bm \Lambda^{-1}),
 \end{align}
 where \(N\) is the number of observations in your dataset.
 % Since each \(J_i\T H_i J_i\) is, for \(i \in 1 \ldots N\), a low-rank positive semi-definite matrix of rank \(O\), it has \(O\) non-zero eigenvalues (see \cref{sec:spectrum-ggn}).
 % Let us define \(\lambda_1, \ldots, \lambda_o\) as the eigenvalues of a matrix sorted in decreasing order.
 % Assume that the non-zero spectrum of \(\bm J_i\T \bm H_i \bm J_i\) is relatively flat, i.e., \(\lambda_1 \not\gg \lambda_o\).
 
 % In a hand-wavy way, we have that the ``more similar'' the eigenvectors (and assuming that the eigenvectors have the same ordering between the two matrices) are, the more the largest eigenvalue will be to the sum of the eigenvalues.
 
 For preconditioning, we need to find a preconditioner matrix \(\bm P\) and its inverse \(\bm P^{-1}\) such that the matrix \(\bm P^{-1} \bm \Lambda\) is well-conditioned.
 This is done by choosing \(\bm P\) such that the two matrices are ``similar'', i.e., they have close eigenvalues and eigenvectors.
 In this way, the preconditioner acts as an initial best guess of the value to be computed.
 We are thus seeking a preconditioner \(\bm P\) such that
 \begin{align}
     \bm P^{-1} \bm A \approx{} & \identity \Rightarrow \bm P \approx \bm A,
     \\  \bm P^{-1} \bm \Lambda ={}& \bm P^{-1} \left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right),
 \end{align}
 where \(\bm \Lambda\) is the posterior precision matrix given by \cref{eq:ggn-laplace} to improve convergence of the msMINRES and Lanczos algorithms.
 The eigenvalues of our GGN approximation are thus given by
 \begin{align}
     \lambda\left( \textstyle{\sum^N_i} \bm J_i\T \bm H \bm J_i + \alpha \identity \right) = \big\{ \underbrace{s_1, \ldots, s_K}_{K \leq N \cdot O}, \underbrace{\alpha, \ldots, \alpha}_{D - K} \big\},
 \end{align}
 where \(\{s_1, \ldots, s_K\}, K \leq N \cdot O\) are the non-zero eigenvalues of the rank-\(K\) matrix \(\sum^N_i \bm J_i\T \bm H_i \bm J_i\).
 If a preconditioner approximates all eigenvalues and eigenvectors of a matrix, then the preconditioner and the matrix are equal and the preconditioner is exact.
 However, we often do not have a preconditioner which is exact, but rather an approximation of the matrix.
 In this case, there is a tradeoff between the accuracy of the preconditioner and its computational cost.
 Furthermore, the inverse of the preconditioner is required for preconditioning, which is often not available.
 To approximate the inverse of the preconditioner, we can use the conjugate gradient method, though this method is then sensitive to the condition number of the preconditioner; and if the original matrix \(\bm \Lambda\) is ill-conditioned, then a good preconditioner will also tend to be ill-conditioned.
 
 As such, the condition number of our preconditioner also becomes an issue, since we have two conflicting objectives: to have a preconditioner which is as similar to \(\bm \Lambda\) as possible and to have a preconditioner which we can invert easily.
 The first point lends itself toward the use of preconditioners with a similar condition number as \(\bm \Lambda\).
 However, to invert an arbitrary positive semi-definite matrix, we will often resort to using the conjugate gradient method which is sentitive to the condition number of the preconditioner.
 Thus, there is an intrinsic balance to be struck between a preconditioner which is very similar to \(\bm \Lambda\) (but whose condition number is problematic) and a preconditioner which is less similar to \(\bm \Lambda\) (and will therefore not improve the conditioning of the problem).
 
 Let us consider the toy example of a preconditioner which we can guarantee to have the same eigenvectors as \(\bm \Lambda\) and whose eigenvalues we can freely choose.
 To do so, one could hypothetically attempt to compute the eigenvectors of \(\bm \Lambda\) and then use these eigenvectors as the columns of the preconditioner.
 % Depending on the rank of the preconditioner, we can then eliminate as many of the largest / smallest eigenvalues of \(\bm \Lambda\) as the rank of the preconditioner.
 If we have a rank-one preconditioner whose only eigenvector is the largest eigenvector of \(\bm \Lambda\), then the eigenvalue corresponding to this eigenvector becomes one and the second-largest eigenvalue becomes the largest (assuming that the second-largest eigenvalue is greater than one).
 Thus, in the best-case scenario, you can then eliminate the largest eigenvalue of \(\bm \Lambda\) by using a rank-one preconditioner.
 Thus, if we have a rank-\(K\) preconditioner, then we can eliminate at most the \(K\) challenging (largest or smallest) eigenvalues of \(\bm \Lambda\) by using a rank-\(K\) preconditioner.
 The steeper the spectrum of \(\bm \Lambda\), the more effective this preconditioning will be (however, in this case, \(\bm \Lambda\) will probably also tend to be worse conditioned).
 % This analogy can be extended in a hand-wavy way to an imperfect preconditioner which does not exactly approximate the eigenvalue
 For a preconditioner with a condition number of \(\kappa^\prime\), we can at best obtain a preconditioned system with conditioning \(\kappa / \kappa^\prime\), where \(\kappa\) is the condition number of \(\bm \Lambda\).
 This is because we want the highest eigenvalue of the preconditioner to neutralise (i.e., align with) the highest eigenvalue of \(\bm \Lambda\), and the same for the lowest eigenvalues.
 
 In CIQ we will, in each iteration, perform one matrix-vector product \(\bm \Lambda \bm v\) and one preconditioner-vector product, as seen in \cref{alg:ciq}.%\sidenote{Because of this, to avoid the preconditioner becoming a performance bottleneck, the speed of the preconditioner-vector products should be significantly greater than that of the matrix-vector products.}
 As such, if the preconditioner-vector product is similarly or more expensive to compute compared to the matrix-vector product, then preconditioning will significantly slow down the algorithm.
 However, we usually choose preconditioners which are relatively easy to compute relative to the posterior precision itself.
 Additionally, an effective preconditioner can significantly improve the convergence speed of CIQ, and this effect can outweigh the cost of computing the preconditioner.
 Thus, we will often choose to use a preconditioner which is not exact, but which is still close to \(\bm \Lambda\) and is cheap to compute.
 We will now discuss some potential choices for preconditioners to sample using CIQ.
 
 \subsection{Fully Linear GGN}
 \label{sec:adam-preconditioner}
 
 One of the simplest preconditioners is similar to the second moment estimation used in Adam (\cref{sec:adam}).
 If we calculate the GGN approximation of the Hessian by linearizing over the whole likelihood (as described in \cref{sec:practical-ggn}), we get
 \begin{align}
     \bm P ={} & \nabla \nabla\T + \alpha \identity \label{eq:adam-preconditioner}
     \\  \bm P^{-1} ={}& \beta \cdot \nabla \nabla\T + \alpha^{-1} \identity \label{eq:adam-preconditioner-inverse}
     \\  \beta ={}& \left(\alpha \norm{\nabla}^2 + \norm{\nabla}^4\right)^{-1} - \left(\alpha \norm{\nabla}^2\right)^{-1}, \label{eq:adam-preconditioner-beta}
 \end{align}
 where \(\nabla\) is the gradient of the loss with regard to the model parameters.
 Since \cref{eq:adam-preconditioner} is in the form of a diagonal perturbation to a vector outer product, we have the closed-form inverse given by \cref{eq:adam-preconditioner-inverse,eq:adam-preconditioner-beta}.
 The problem with using this approximation is that the rank of the outer product is one and this term determines the largest eigenvalue of the preconditioner.
 This means that, best-case scenario, we can only reduce the highest eigenvalue of the Hessian.
 This occurs when the gradient corresponds to the direction of the highest precision (this may not be the case).
 This would then leave the second-highest eigenvalue unaffected.
 % If the rank of our precision matrix is greater than one, then this value seems likely to be similarly large to originally the largest eigenvalue.
 % Therefore, the condition number of our problem will not be significantly reduced.
 We therefore need to find a higher-rank approximation of the posterior precision.
 
 \subsection{Sub-Sampling the Data}
 \label{sec:sub-sampling-preconditioner}
 
 We could just choose to calculate the true posterior precision over fewer observations and invert it using the conjugate gradient method to use it as a preconditioner.
 This would have the advantage of being guaranteed to converge to the true posterior precision as the number of preconditioner observations increases.
 The spectrum of this preconditioner would thus be similar to that of the true posterior precision.
 Since the lowest eigenvalue of the posterior precision is the prior precision \(\alpha\), we can (even with few preconditioner observations) increase the lowest eigenvalue of the preconditioned system to be equal to one.
 
 Alternatively, we can estimate the eigenvectors and eigenvalues of \(\textstyle{N / B \sum^B_i} J_i\T J_i\), where \(B\) is the number of preconditioner observations.
 This would \emph{hopefully} have an eigenbasis that is close to that of \(\textstyle{\sum^N_i} J_i\T J_i\) without needing to calculate the precision matrix over the whole dataset (we may need to guarantee all classes are represented in \(x_i, i \sim \mathcal{B}\)).\sidenote{Empirically, it seems likely that the eigenvectors of the posterior precision are relatively similar across different observations---see \cref{fig:spectrum-mnist} and \cref{sec:spectrum-ggn}.}
 However, scaling the precondititioner up to the size of the entire dataset will lead to a preconditioner which is ill-conditioned, making inversion via conjugate gradient difficult.
 
 As mentioned before, each msMINRES iteration performs one matrix-vector product and one preconditioner-vector product.
 However, even though the preconditioner is constructed in nearly the same way as the sum of GGN matrices in the posterior precision, the preconditioner is significantly cheaper to compute.
 This is because the preconditioner-vector product will cost \(B / N \cdot t_{\mathrm{GGN}}\), where \(t_{\mathrm{GGN}}\) is the computational cost of performing the GGN-vector product.\sidenote{This analysis assumes that the computation of the GGN-vector product scales linearly in the number of observations used in the GGN matrix. Mathematically speaking, this is a valid assumption, due to the number of elementary operations that are performed in the algorithm. However, the efficient JAX XLA-compiled implementation of Jacobian-vector products can scale differently to na\"ive implementations, so the validity of the assumption is less clear in our case.}
 Since we will have \(N \gg B\), then the cost of computing the preconditioner will be negligible compared to that of the GGN-vector product.
 
 % This preconditioner can also be adjusted to cases where the smallest eigenvalue \(\alpha\) is not very small, but for which the largest eigenvalue is very large.
 % In this case, we can compute this preconditioner with a larger value of \(\alpha\) and a large number of preconditioner observations.
 % This will lead to a preconditioner which has a larger smallest eigenvalue than the true posterior precision, but whose largest eigenvalue is closer to that of the true posterior precision.
 % This will lead to a preconditioner which can reduce the largest eigenvalues of the posterior precision.
 % This may be useful, for instance, in problems for which the eigenvectors of the posterior precision are not constant across the dataset, but are instead highly variable.
 % However, this preconditioner will be very expensive to compute, since it effectively requires computing the entire posterior precision.\sidenote{MEH Empirically, I found that this is useful for the sine problem, though, which seems to have a very constant posterior precision.}
 
 \subsection{Using the Woodbury Matrix Identity}[Woodbury]
 \label{sec:woodbury-preconditioner}
 
 The Woodbury matrix identity~\cite{woodbury1950inverting} is a useful identity for computing the inverse of a matrix which is the sum of two matrices of which one is low-rank and the other can easily be inverted.
 The identity is given by
 \begin{align}
     \left( \bm A + \bm B \bm C \bm D \right)^{-1} = \bm A^{-1} - \bm A^{-1} \bm B \left( \bm C^{-1} + \bm D \bm A^{-1} \bm B \right)^{-1} \bm D \bm A^{-1},
 \end{align}
 where \(\bm A\) is a \(D \times D\) matrix and can easily be inverted, \(\bm B\) is a \(D \times O\) matrix, \(\bm C\) is a \(O \times O\) matrix, and \(\bm D\) is a \(O \times D\) matrix.
 This identity can be used to compute the inverse of the GGN matrix with the prior precision as
 \begin{align}
     \left( \alpha \identity + \bm J\T \bm H \bm J \right)^{-1} = \alpha^{-1} \identity - \alpha^{-2} \bm J\T \left( \bm H^{-1} + \alpha^{-1} \bm J \bm J\T \right)^{-1} \bm J,
 \end{align}
 where \(\alpha\) is the prior precision and \(\bm H \in \reals^{O \times O}\) and \(\bm J \in \reals^{O \times D}\) are the Hessian and Jacobian from \cref{eq:negative-log-likelihood-ggn} evaluated at a single observation.
 
 Here, to exploit the low-rank structure of the preconditioner, we only compute the GGN matrix over one observation.
 For this preconditioner to be effective, we assume that the eigenvectors of the GGN matrix are approximately constant across the dataset.
 This may be a reasonable assumption for many problems, but is not guaranteed to be true.
 
 However, even if the eigenvectors of the GGN matrix are constant across the dataset, the eigenvalues of this preconditioner should be large enough to reduce the condition number of the posterior precision matrix.
 To do so, we can scale the outer product \(\bm J \bm J\T\) term by some factor, e.g., the number of observations, \(N\).
 Since, by scaling up the outer product term we are effectively scaling the eigenvalues of \(\bm J \bm J\T\) by a constant, the condition number of this outer product will remain the same.
 In this way, we can increase the maximum eigenvalues of the preconditioner without increasing the condition number of the outer product \(\bm J \bm J\T\) which we want to invert using the conjugate gradient method.
 
 Before the Woodbury matrix identity can be used, the Hessian of the loss with respect to the model outputs \(\bm H\) must first be inverted.
 To use the conjugate gradient method to invert this Hessian, we would have to perform nested conjugate gradient to also compute the inverse of the outer product term \(\bm H^{-1} + \alpha^{-1} \bm J \bm J\T\), which could be very expensive, depending on the number of outputs.
 For regression tasks where the Hessian is diagonal (\cref{eq:mse-ggn}), this is not a problem.
 For classification tasks using the cross-entropy loss, we can compute the inverse of the Hessian as a diagonal and outer product matrix, so this is also not a problem.
 For models with a very large number of outputs, such as autoencoders, the matrix \(\bm H\) will be fairly large, and so computing its inverse manually can become expensive.
 However, in the case of autoencoders, the loss is usually the MSE loss, and so has a diagonal Hessian which can be inverted easily.
 
 Thus, if \(\bm H\) is positive definite (which is necessary for the Laplace approximation to be valid) and easy to invert, then it should be possible to use the Woodbury matrix identity to compute the inverse of this preconditioner.
 This preconditioner is fast to compute, but may not be effective when the eigenvectors of the GGN matrix are not constant across the dataset.
 However, due to the effectiveness of the sub-sampled GGN preconditioner and the assumption of stationary eigenvectors, we did not end up using this preconditioner.
 Additionally, this preconditioner cannot assume a higher rank than the number of outputs \(O\), since adding more observations to the preconditioner will increase its rank to some unknown value and prevent the application of the Woodbury matrix identity.
 Ideally, we would like to be able to find a preconditioner for which we can choose the rank of the low-rank approximation \(\bm A \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) without simply computing the GGN on a single observation (\(\bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\)), while still allowing us to apply the Woodbury matrix identity, unlike in the case of the sub-sampling preconditioner in \cref{sec:sub-sampling-preconditioner} (\(\sum^B_{i=1} \bm J_i\T \bm H_i \bm J_i \approx \sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\)).
 
 \subsection{Pivoted Cholesky}
 \label{sec:pivoted-cholesky}
 
 The pivoted Cholesky factorisation is a method for computing a low-rank approximation to the Cholesky decomposition of a matrix~\cite{harbrecht2012low,bach2013sharp}.
 This decomposition has been successfully applied to preconditioning for kernel matrices in Gaussian processes~\cite{gardner2018gpytorch}, but not yet 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}
     \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
 \begin{align}
     \bm P ={} & \bm L\T \bm L + \alpha \identity,
     \\  \bm P^{-1} ={}& \alpha^{-1} \identity - \alpha^{-2} \bm L\T \left( \identity + \alpha^{-1} \bm L \bm L\T \right)^{-1} \bm L.
 \end{align}
 We then obtain the preconditioner-vector and inverse preconditioner-vector products as in Equation~\ref{eq:pivoted-cholesky-preconditioner}.
 The most expensive computation in the inverse preconditioner is inverting a \(K \times K\) matrix, where \(K\) is the rank of our pivoted Cholesky factorisation (which we choose).
 \begin{align}\label{eq:pivoted-cholesky-preconditioner}
     \bm P \bm v ={} & \bm L\T ( \bm L \bm v ) + \alpha \bm v,
     \\  \bm P^{-1} \bm v ={}& \alpha^{-1} \bm v - \alpha^{-2} \bm L\T ( \underbrace{\identity + \alpha^{-1} \bm L \bm L\T}_{K \times K} )^{-1} ( \bm L \bm v ).
 \end{align}
 Here it is possible to precompute the inverse factorised product given by \(\bm L\T ( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{D \times K}\), meaning that we only need to perform two matrix multiplications for both the preconditioner- and inverse preconditioner-vector products.
 This would then require storing the factorisation \(\bm L \in \reals^{K \times D}\) and the inverse factorised product \(\reals^{D \times K}\) matrix.\sidenote{We can also precompute the inverse \(( \identity + \alpha^{-1} \bm L \bm L\T )^{-1} \in \reals^{K \times K}\), which means we would only need to store the factorisation \(\bm L \in \reals^{K \times D}\) and this inverse \(K \times K\) matrix, in which case we would just need to perform one more matrix multiplication (multiplying \(\bm L\T\) with this \(K \times K\) matrix) in every preconditioner call.}
 
 As explained in \cref{sec:spectrum-ggn}, as we sum over multiple low-rank GGN matrices, we will likely find that the minimum non-zero eigenvalue gets progressively smaller compared to the largest.
 The condition number will therefore get progressively worse.
 The preconditioning of \(\sum^N_{i=1} \bm J_i\T \bm H_i \bm J_i\) (which should be low-rank) will progressively become \emph{less} low-rank and its condition number will progressively become higher.
 
 Regarding the non-zero eigenvalues and eigenvectors, we consider a best-case scenario.
 Assume that the pivoted Cholesky factorisation finds the \(K\) largest eigenvectors of the precision matrix.
 If the eigenspectrum of this factorisation is steep enough, then the largest eigenvalue \(s_1\) will be significantly larger than the smallest eigenvalue \(s_K\) (i.e., there will be a steep drop-off from the largest to the smallest non-zero eigenvalue).
 Since the eigenvectors of each per-observation term \(\bm J_i\T \bm H_i \bm J_i\) will not be the same, when we sum over them, the eigenvalues of the sum will not be the sum of the eigenvalues of the individual Hessians.
 However, some of the eigenvectors may be close to each other, and so the spectrum of the GGN sum may be approximately the sum of the GGN eigenvalues.
 Nevertheless, the rank of the total Hessian will increase with the number of observations and so the number of non-zero eigenvalues consequently increase as well.
 These eigenvalues will be smaller than the summed eigenvalues (because the eigenvectors may not be contained in the same subspace), but they will be larger than the smallest eigenvalue of the per-observation Hessian.
 This will then lead to a steeper spectrum and a larger condition number of the total Hessian.
 Thus, when the batch size \(B\) is large, the condition number of the sum of GGN Hessians will increase and the preconditioner will become less effective.
 How much the condition number increases with the batch size depends on the eigenspectrum of the per-observation Hessians (i.e., how close the eigenvectors are to each other).
 
 If we want to use the pivoted Cholesky decomposition to approximate \(\sum^B_i \bm J_i\T \bm H_i \bm J_i\), we are largely trying to approximate the higher eigenvectors of the sum.
 In order for this to be the case, we will want to oversample vectors from this matrix which correspond to the largest eigenvalues.
 A common approach to sampling the largest eigenvalues is to sample vectors with probability proportional to the diagonal value~\cite{harbrecht2012low}.
 However, it is not possible to generally and efficiently determine the diagonal of a matrix via matrix-vector products.
 Often, it is necessary to multiply a vector of zeros with a value of one at the index at which to find the diagonal.
 This operation needs to be repeated as many times as there are elements in the diagonal with no shared computation between each element, which is very expensive.
 Alternatively, we can approximate the diagonal of the GGN matrix as the element-wise product of the gradient with itself, similarly to the preconditioner described in \cref{sec:adam-preconditioner} and the GGN approximation computed in \cref{eq:fully-linear-ggn}, though this may not be the most accurate approximation.
 
 In conclusion, to find the optimal parameters for the pivoted Cholesky preconditioner (the number of observations \(B\) and the rank of the pivoted Cholesky factorisation \(K\)), there are some rules of thumb to consider.
 Notably, the quality of the preconditioner depends largely on the ease of inverting the \(\identity + \alpha^{-1} \bm L \bm L\T\) matrix, which depends on its condition number.
 In practice, increasing the rank of the pivoted Cholesky factorisation will increase the condition number of this matrix, since we are approximating more of the largest eigenvectors and so we will capture progressively smaller eigenvalues.
 
 We have found that it is most effective to keep \(K\) low, such that the \(K \times K\) outer product of the pivoted Cholesky factorisation matrix \(L\) approximates few enough of the largest eigenvectors that the eigenvalue of the largest is close to that of the smallest (meaning this outer product is well-conditioned).
 We thus need to choose \(K\) such that it is large enough to approximate enough large eigenvectors and small enough to ensure the condition number of the outer product is low.
 Furthermore, you want enough observations \(B\) to be included in the preconditioner such that it is possible to accurately approximate the eigenvectors of the problem---it may be optimal to use the whole dataset, though this may worsen the conditioning of the outer product.
 
 Furthermore, the pivoted Cholesky decomposition is not implemented in JAX.
 Since the pivoted Cholesky decomposition involves incrementally adding columns to the factorisation, and JAX is not friendly towards mutable variables, this preconditioner is not straightforward to implement.
 Additionally, the pivoted Cholesky decomposition traditionally uses the diagonal of the matrix to determine which columns to add to the factorisation, which is difficult without instantiating the matrix, so we would need to approximate the diagonal of the GGN matrix.
 
 \subsection{Other Preconditioners}
 
 The randomly pivoted Cholesky~\cite{chen2022randomly} preconditioner is simple and inexpensive variant of the partial Cholesky decomposition family of algorithms.
 Similarly to the pivoted Cholesky preconditioner, the randomly pivoted Cholesky preconditioner requires the diagonal of the GGN matrix in order for the decomposition to be close to the GGN.
 We can thus approximate this term as we do in the pivoted Cholesky preconditioner (see \cref{sec:pivoted-cholesky}).
 Early tests of the randomly pivoted Cholesky preconditioner were not successful, and so it was not pursued further.
 
 Additionally, we considered potential preconditioners to be computed as traditional approximations of the posterior precision, like the diagonal approximation or the Kronecker factorisation.
 These methods have been shown to be relatively effective approximations of the posterior precision in the context of Bayesian neural networks~\cite{daxberger2021laplace}, and so they may contain enough information about the full posterior precision to be effective preconditioners.
 
 However, due to the speed, effectiveness, and theoretical properties of other preconditioners proposed in this section, these methods were not used.
 Overall, both the sub-sampled preconditioner and the pivoted Cholesky preconditioners were found to be effective under the correct settings.
 However, the pivoted Cholesky preconditioner is not easy to implement in JAX and was more difficult to tune than the sub-sampled preconditioner.
 The sub-sampled preconditioner is thus the preferred preconditioner in this project.
 We provide a high-level comparison of the various preconditioners we have presented in this chapter in \cref{tab:preconditioner-comparison}.
 
 \begin{table*}
     \caption[Comparison of various preconditioners.]{Comparison of various approximations of the likelihood contribution \(\sum_i^N \bm J_i\T \bm H_i \bm J_i\) to the Laplace posterior precision for use in preconditioning. We compare the rank of the approximation, the speed of computing preconditioner-vector products, the closeness of the approximation, and the method used for inverting the preconditioner.}
     \label{tab:preconditioner-comparison}
     \include{tables/preconditioner-comparison.tex}
 \end{table*}
 
 \section{Sampling Evaluation}
 \label{sec:sampling-evaluation}
 
 As we have explained, sampling through CIQ is sensitive to the condition number of the posterior precision.
 As such, we need a technique for evaluating the quality of the samples drawn from the approximate posterior to determine if sampling occurred correctly.
 This means evaluating whether the samples are drawn from a normal distribution with the correct mean and precision.
 
 What we would like to fundamentally test is whether the CIQ algorithm converged correctly and thus whether the samples were generated successfully.
 However, for cases when sampling fails to converge, this is typically caused by ill-conditioning of the posterior precision matrix.
 Since we only have access to the precision matrix and computing the covariance involves inverting the precision matrix using, for example, the conjugate gradient method (which is itself sensitive to ill-conditioning), we cannot directly evaluate the covariance of the samples.
 Instead, we require a method to evaluate the quality of the samples without access to the covariance matrix, which is what we will discuss in the next section.
 
 \subsection{The Chi-Squared Distribution}[Chi-Squared Distribution]
 
 As per the definition of the chi-squared distribution, for \(D\)-dimensional standard normal samples \(\bm \varepsilon_0\), the sum of square deviations from the mean is chi-squared distributed with \(D\) degrees of freedom, as per
 \begin{align}
     {\bm \varepsilon_0}\T \bm \varepsilon_0 \sim \chi^2(D).
 \end{align}
 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}
     \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 the computation of the product of the precision matrix \(\bm \Lambda\) with a vector, which we can compute efficiently, as it does not require instantiation of the matrix.
 For the Laplace approximation, this can easily be computed as $\bm \Lambda = - \nabla^2_{\bm \theta} \log p(\bm \theta \given \bm y)$.
 Since this precision matrix can be computed exactly, this calculation can be performed to evaluate the whether a set of samples \(\bm \varepsilon\) are drawn from a distribution \(\normal(\bm \mu, \bm \Lambda^{-1})\) while only implicitly accessing this precision matrix in the evaluation via matrix-vector products.
 To do so, we can compute the squared Mahalanobis distance of each sample \(\bm \varepsilon_i\) from the mean \(\bm \mu\) and compare it to the chi-squared distribution with \(D\) degrees of freedom.
 We do this by comparing the histogram of these squared Mahalanobis distances to the appropriate PDF and by visualising the sample empirical percentiles against the theoretical percentiles in a quantile-quantile plot.
 As such, this method can be used to evaluate whether the approximate normal samples obtained from CIQ are correctly distributed.
 However, this method does not provide a quantitative measure of the quality of the samples.
 In order to run experiments using the Laplace approximation without manually checking the plots of every set of samples we generate, we also need a method to evaluate the quality of the samples quantitatively.
 Our current methods do not provide this, and so we need to evaluate the quality of the samples in a different way.
 
 The Kolmogorov-Smirnov test is a non-parametric test for comparing two continuous distributions.
 It is a test of whether two samples are drawn from the same distribution.
 Since the chi-squared distribution is continuous, we can use the Kolmogorov-Smirnov test to compare the chi-squared distribution with \(D\) degrees of freedom to the squared Mahalanobis distances of the samples.
 This test is performed by computing the Kolmogorov-Smirnov statistic, which is the maximum difference between the empirical and theoretical CDFs.
 The null hypothesis of the test is that the samples are drawn from the same distribution, and so the test is rejected if the Kolmogorov-Smirnov statistic is greater than the critical value for the test.
 This means that sample normality is rejected if the p-value of the test is less than the significance level of the test.
 However, this test is very sensitive to the number of samples used to compute the statistic, and so a large number of samples will lead to a very low p-value even if the samples are drawn from the same distribution.
 Since we can sample an arbitrarily large number of samples from the approximate posterior, this may lead to low p-values even when sampling succeeds.
 
 % \section{Limitations}
 
 % The CIQ algorithm 
 
 % Scaling with data?
 % Any comments on the model size and how it scales?