training.tex (21216B)
1 \setchapterpreamble[u]{\margintoc} 2 \chapter{Training} 3 \labch{training} 4 5 In this chapter, we will discuss training neural networks by maximising the evidence under the Laplace approximation. 6 We first motivate maximising the evidence as opposed to maximising the posterior and discuss how the Laplace approximation guarantees the feasibility of the evidence, as well as the issues that arise when computing the marginal likelihood with the full posterior precision. 7 Next, we propose a novel method for computing an upper bound on the log-determinant of the full posterior precision matrix which avoids the aforementioned issues. 8 Finally, we discuss the implementation of online Laplace with mini-batching using our upper bound. 9 10 \section{Maximising the Evidence} 11 \label{sec:maximising-the-evidence} 12 13 As seen in \cref{eq:bayes-theorem}, the posterior probability for a given model is 14 \begin{align} 15 p(\bm \theta \given \bm y) = \frac{p(\bm y \given \bm \theta)\,p(\bm \theta)}{p(\bm y)}, 16 \end{align} 17 where \(p(\bm y)\) is known as the evidence. It is also known as the \emph{marginal likelihood}, as it is obtained by marginalising \(\bm \theta\) from the likelihood as 18 \begin{align} 19 p(\bm y) ={} & \int_{\bm \theta} p(\bm y \given \bm \theta)\,p(\bm \theta)\,d\bm \theta. 20 \end{align} 21 This term is typically intractable, since it involves integrating over the entire parameter space. 22 For this reason, we typically approximate the posterior by some tractable distribution, such as a Gaussian, and then use this to approximate the marginal likelihood. 23 24 Since we are integrating over the posterior, maximising the marginal likelihood will seek minima that are flat, since the posterior will be very broad in these regions. 25 This is a desirable property, since it means that the model will be less likely to overfit the data, both from a theoretical standpoint~\cite{hochreiter1997flat} and in practice~\cite{keskar2016large,jiang2019fantastic,maddox2019simple}. 26 27 To further motivate training on the evidence term, \textcite{fong2020marginal} show that maximising the marginal (and consequently the log-marginal) is equivalent to performing leave-p-out cross validation for all values \(p = 1, \ldots, \infty\) and choosing the model with the highest average posterior probability (across each of \(p\) folds and across all values of \(p\)). 28 This result shows that the marginal likelihood is a good proxy for the true cross-validation performance of a model. 29 Because of this, it can be used for model selection, since it will choose the model that performs best on average across all permutations of the data. 30 31 Additionally, unlike the posterior, the marginal likelihood can be used to optimise both model parameters and hyperparameters, since we can meaningfully minimise the log-marginal with respect to the hyperparameters (and compute the respective gradients). 32 This is useful, since it allows us to use gradient-based optimisation methods to find the hyperparameters that maximise the marginal likelihood, as has been commonly done in Gaussian processes~\cite{williams2006gaussian,damianou2013deep}. 33 This provides an alternative to finding the hyperparameters by grid search, which is computationally expensive. 34 35 The Laplace approximation provides an approximation to the posterior which allows us to compute and maximise the marginal likelihood. 36 This is known as \emph{online} Laplace, and is an alternative to the \emph{post-hoc} Laplace approach, which is the standard approach to training under the Laplace approximation~\cite{daxberger2021laplace}. 37 Under this approximation, we approximate the posterior distribution as \(p(\bm \theta \given \bm y) \approx \normal(\bm \theta \given \bm \theta_{\textsc{map}}, \bm \Lambda^{-1})\), where the posterior precision \(\bm \Lambda = \nabla^2_{\bm\theta} \mathcal L(\bm \theta_{\textsc{map}})\) is the Hessian of the negative log-posterior at the maximum a posteriori (MAP) estimate \(\bm \theta_{\textsc{map}}\). 38 As in \cref{eq:laplace-taylor}, we have the approximate Laplace log-posterior 39 \begin{align}\label{eq:laplace-posterior} 40 \begin{split} 41 \log p(\bm \theta \given \bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) \cdot (\bm \theta - \bm \theta_{\textsc{map}}) 42 \\ & + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \nabla^2_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y) (\bm \theta - \bm \theta_{\textsc{map}}) 43 \\={}& \log p(\bm \theta_{\textsc{map}} \given \bm y) - \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}). 44 \end{split} 45 \end{align} 46 \marginnote{In reality, we will not be at the MAP solution when optimising the marginal, neither during training nor after convergence (the marginal and the posterior have different minima). In this case, we cannot neglect the first derivative term in \cref{eq:laplace-posterior}. However, when this term is not zero, the Laplace posterior is \(\normal(\bm \theta \given \bm \theta_{\textsc{map}}^*, \bm \Lambda^{-1})\), where \(\bm \theta_{\textsc{map}}^* \coloneqq \bm \theta_{\textsc{map}} + \bm \Lambda^{-1} \nabla_{\bm \theta} \log p(\bm \theta_{\textsc{map}} \given \bm y)\). This corresponds to shifting the parameters to the MAP with a Newton step \(\bm \theta_{\textsc{map}}^* \coloneqq \bm \theta_{\textsc{map}} - \nabla^2_{\bm \theta} \mathcal{L}(\bm \theta_{\textsc{map}})^{-1} \nabla_{\bm \theta} \mathcal{L}(\bm \theta_{\textsc{map}})\), since Newton's method finds the minimum in one step for quadratic functions (like the second order Taylor expansion to our loss).} 47 The approximate Laplace marginal is then computed as 48 \begin{align}\label{eq:laplace-marginal} 49 \begin{split} 50 p(\bm y) \stackrel{\textsc{la}}{\approx}{} & \int_{\bm \theta} p(\bm \theta_{\textsc{map}} \given \bm y) \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}}) \right)\,d \bm \theta 51 \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y) \int_{\bm \theta} \exp\left( -\frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}}) \bm \Lambda (\bm \theta - \bm \theta_{\textsc{map}})\right)\,d \bm \theta 52 \\ ={}& p(\bm \theta_{\textsc{map}} \given \bm y)\,(2 \pi)^{D / 2}\,(\det \bm \Lambda)^{-1/2}. 53 \end{split} 54 \end{align} 55 Taking the log of this approximate marginal, we obtain 56 \begin{align}\label{eq:laplace-log-marginal} 57 \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & \log p(\bm \theta_{\textsc{map}} \given \bm y) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda 58 \\ ={}& \log p(\bm y \given \bm \theta_{\textsc{map}}) + \log p(\bm \theta_{\textsc{map}}) + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda. \nonumber 59 \end{align} 60 For the normal likelihood and normal prior, we have 61 \begin{align}\label{eq:normal-likelihood} 62 p(\bm y \given \bm \theta) ={} & \normal(\bm y \given f_{\bm\theta}(\bm x), \rho^{-1} \identity) = \prod_{i=1}^N \normal(\bm y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity), 63 \\ \log p(\bm y \given \bm \theta) ={}& \sum_{i=1}^N \log \normal(\bm y_i \given f_{\bm\theta}(\bm x_i), \rho^{-1} \identity) \label{eq:normal-log-likelihood} 64 \\ ={}& \sum_{i=1}^N \left( - \frac{O}{2} \log(2 \pi) - \frac{1}{2} \log \det (\rho^{-1} \identity) - \frac{1}{2} \rho \norm{\bm y_i - f_{\bm\theta}(\bm x_i)}^2\right) \nonumber 65 \\={}& - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2, \nonumber 66 \end{align} 67 and 68 \begin{align}\label{eq:normal-prior} 69 p(\bm \theta) ={} & \normal(\bm \theta \given \bm 0, \alpha^{-1} \identity), 70 \\ \log p(\bm \theta) ={}& - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2, \label{eq:normal-log-prior} 71 \end{align} 72 where \(\alpha\) is the precision of the prior and \(\rho\) is the precision of the likelihood. 73 Combining \cref{eq:laplace-log-marginal,eq:normal-log-likelihood,eq:normal-log-prior}, we expand the Laplace approximation to the log-marginal likelihood as 74 \begin{align}\label{eq:expanded-log-marginal} 75 \begin{split} 76 \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{N O}{2} \log(2 \pi) + \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2 77 \\ & - \frac{D}{2} \log(2 \pi) + \frac{D}{2} \log \alpha - \frac{1}{2} \alpha \norm{\bm \theta}^2 78 \\ & + \frac{D}{2} \log (2 \pi) - \frac{1}{2} \log \det \bm \Lambda, 79 \end{split} 80 \end{align} 81 where it can be seen that there are constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\). 82 While these terms are not necessary for the optimisation problem, we choose to include them, as these terms allow us to interpret the log-marginal as a likelihood and do not significantly affect performance. 83 From \cref{eq:expanded-log-marginal}, we can see that the Laplace log-marginal likelihood requires computation of the log-determinant of the posterior precision matrix. 84 This log-determinant term gets minimised in our optimisation, which favours models with small eigenvalues in the Hessian of the loss, i.e., points of low curvature in the loss landscape, as we argued earlier in this section. 85 This is a challenging computation, as there is no \emph{cheap} closed-form expression for the log-determinant of a matrix without instantiating the matrix. 86 However, we can compute an upper bound on the log-determinant of the posterior precision matrix, which is useful for our optimisation problem which minimises the log-determinant term. 87 88 % \begin{tabular}{ll} 89 % Minimise & Maximise 90 % \\\addlinespace 1. \(\rho / 2 \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2\) & 1. \(N O / 2 \log \rho\) 91 % \\ 2. \(1 / 2 \log \det \bm \Lambda\) & 2. \(D \log \alpha\) 92 % \\ 3. \(\alpha / 2 \norm{\bm \theta}^2\) 93 % \end{tabular} 94 95 % The \texttt{laplace} Python library provides a function to compute the log-marginal likelihood of the Laplace approximation. The function is called \texttt{log\_marginal\_likelihood} and is given by 96 % \begin{align} 97 % \log p(\bm y) \mathrel{\stackrel{\textsc{la}}{\approx}}{} & \frac{N O}{2} \log \rho - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2 98 % \\ & - \frac{1}{2} \left( \log \det \bm \Lambda - D \log \alpha + \alpha \norm{\bm \theta}^2 \right). \nonumber 99 % \end{align} 100 % As compared to the exact log-marginal from \cref{eq:expanded-log-marginal}, this implementation is obtained by omitting the constant terms which do not depend on \(\bm \theta\), \(\alpha\), or \(\rho\). 101 % However, in doing so, the authors ignore the terms \(D / 2 \log(2 \pi)\) 102 % While these terms do not affect the optimisation problem, they are useful for comparing models with different numbers of parameters. 103 104 \section{The Determinant Bound} 105 \label{sec:determinant-bound} 106 107 We want to compute an upper bound on the log-determinant of the posterior precision \(\bm \Lambda \in \reals^{D \times D}\) such that we have \(\log \det \bm \Lambda \leq B_D(\mu_1, \mu_2, \beta)\). 108 \textcite{bai1996bounds} provide one such bound, 109 \begin{align}\label{eq:bai-golub-bound} 110 B_D(\mu_1, \mu_2, \beta) \coloneqq{} & \begin{pmatrix} \log \beta & \log t \end{pmatrix} 111 \begin{pmatrix} \beta & t \\ \beta^2 & t^2 \end{pmatrix}^{-1} 112 \begin{pmatrix} \mu_1 \\ \mu_2 \end{pmatrix}, 113 \end{align} 114 with \(\mu_1 = \Tr(\bm \Lambda)\), \(\mu_2 = \norm{\bm \Lambda}^2_F = \Tr\left(\bm \Lambda^2\right)\), \(t = \frac{\beta \mu_1 - \mu_2}{\beta D - \mu_1}\), where \(\beta\) is an upper bound on the largest eigenvalue of \(\bm \Lambda\). 115 This bound amounts to approximating the log eigenvalues of \(\bm \Lambda\) with a second-order polynomial. 116 Since \(B_D(\mu_1, \mu_2, \beta)\) is an upper bound on the log-determinant of \(\bm \Lambda\), we are computing a lower bound on the log-marginal likelihood (which we maximise), as per \cref{eq:expanded-log-marginal}. 117 118 Assume \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\) has rank \(K \leq O \cdot N\). 119 Then \(\bm \Lambda = \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i + \alpha \identity\) has \(K\) distinct eigenvalues greater than \(\alpha\) and the remaining \(D - K\) eigenvalues are equal to \(\alpha\). 120 We then decompose the log-determinant as 121 \begin{align}\label{eq:log-determinant-decomposition} 122 \begin{split} 123 \log \det \bm \Lambda ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + \sum_{k=K+1}^D \log \alpha 124 \\ ={} & \sum_{k=1}^{K} \log(\lambda_k + \alpha) + (D - K) \log \alpha, 125 \end{split} 126 \end{align} 127 where \(\lambda_1, \ldots, \lambda_K\) are the non-zero eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\). 128 The first term is the log-determinant of a positive definite matrix with eigenvalues \(\{\lambda_1 + \alpha, \ldots, \lambda_K + \alpha\}\), and the second term is a constant. 129 The former can then be bounded using \cref{eq:bai-golub-bound} with an upper bound on its eigenvalues \(\beta \geq \lambda_1 + \alpha\) as \(B_D\left(\sum_{k=1}^{K} (\lambda_k + \alpha), \sum_{k=1}^{K} (\lambda_k + \alpha)^2, \beta\right)\). 130 We then compute the trace of this term as in \cref{eq:laplace-trace-bound-2} via 131 \begin{align}\label{eq:laplace-trace-bound} 132 \Tr{\bm \Lambda} ={} & \sum_{k=1}^{K} (\lambda_k + \alpha) + (D - K) \alpha 133 \\ \begin{split} 134 \Rightarrow \sum_{k=1}^{K} (\lambda_k + \alpha) ={}& \Tr{\bm \Lambda} - (D - K) \alpha \label{eq:laplace-trace-bound-2} 135 \\ ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha, 136 \end{split} 137 \end{align} 138 and for the square trace we have 139 \begin{align}\label{eq:laplace-square-trace-bound} 140 \sum_{k=1}^{K} (\lambda_k + \alpha)^2 ={} & \Tr\left(\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)^2\right) + 2 \alpha \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha^2. 141 \end{align} 142 We can thus bound \(\log \det \bm \Lambda\) with \cref{eq:log-determinant-decomposition,eq:laplace-trace-bound-2,eq:laplace-square-trace-bound} as 143 \begin{align}\label{eq:laplace-log-determinant-bound} 144 \log \det \bm \Lambda \leq{} & B_K\left(\mu_1, \mu_2, \beta\right) + (D - K) \log \alpha, 145 \\ \mu_1 ={}& \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha,\label{eq:mu1} 146 \\ \mu_2 ={}& \Tr\left(\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)^2\right) + 2 \alpha \Tr\left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right) + K \alpha^2.\label{eq:mu2} 147 \end{align} 148 For the upper bound \(\beta\) on the eigenvalues of \(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i\), we can exploit that this matrix is positive semidefinite and use \(\beta = \Tr( \sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i ) + \alpha\) as an upper bound.\sidenote{Since the matrix is positive semidefinite, all eigenvalues are non-negative and thus the trace is bigger than the largest eigenvalue.} 149 From \cref{eq:mu1,eq:mu2}, we can see that we need to calculate the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\) to compute the upper bound on the log-determinant of the posterior precision. 150 % WHERE DOES THE BOUND COME FROM? WHICH UNCERTAINTY? THE UPPER BOUND ON THE EIGENVALUE? 151 152 \subsection{Hutchinson's Trace Estimator}[The Trace Estimator] 153 154 To compute the upper bound on the log-determinant of the posterior precision, we need to compute \(\mu_1\) and \(\mu_2\), which depend on the traces \(\Tr(\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)\) and \(\Tr((\sum_{i=1}^N \bm J_i\T \bm H_i \bm J_i)^2)\). 155 To compute these, we can use Hutchinson's trace estimator~\cite{hutchinson1990stochastic}. 156 Let \(\bm \varepsilon_0 \sim \normal(\bm 0, \identity)\). 157 Since \(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} = \identity\), we compute \(\mu_1\) as 158 \begin{align}\label{eq:hutchinson-trace} 159 \mu_1 ={} & \Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) + K \alpha, 160 \\ \begin{split}\Tr(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i) ={}& \Tr\left(\E{}{\bm \varepsilon_0 \bm \varepsilon_0\T} \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right)\right) 161 \\ ={}& \E{}{\Tr(\bm \varepsilon_0 \bm \varepsilon_0\T \left(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i\right))} 162 \\ ={}& \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \bm \varepsilon_0}.\end{split} 163 \end{align} 164 Similarly, we compute \(\mu_2\) as 165 \begin{align}\label{eq:hutchinson-square-trace} 166 \mu_2 =\Tr\left(\left( {\textstyle {\textstyle \sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i \right)^2\right) + 2 \alpha \Tr\left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) + K \alpha^2, 167 \\ \Tr\left(\left( {\textstyle {\textstyle \sum_{i=1}^N}} \bm J_i\T \bm H_i \bm J_i \right)^2\right) = \E{}{\bm \varepsilon_0\T \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right) \left( {\textstyle \sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \right)\T \bm \varepsilon_0}. 168 \end{align} 169 We then apply a Monte Carlo approximation of the expectation by sampling \(S\) samples from \(\bm \varepsilon_0\sim \normal(\bm 0, \identity)\) to compute the traces. 170 Fortunately, this means that we only need to compute one GGN-vector product \(\textstyle{\sum_{i=1}^N} \bm J_i\T \bm H_i \bm J_i \cdot \bm \varepsilon_0\) to estimate the \(\mu_1\) and \(\mu_2\) terms in \cref{eq:laplace-log-determinant-bound}. 171 Furthermore, we have that the trace can be rearranged and simplified to 172 \begin{align}\label{eq:trace-rearrangement} 173 \Tr\big({\textstyle{\sum_{i=1}^N}} \underbrace{\bm J_i\T \bm H_i \bm J_i}_{D \times D}\big) ={} & \sum_{i=1}^N \Tr\left(\bm J_i\T \bm H_i \bm J_i\right) = \sum_{i=1}^N \Tr\big(\underbrace{\bm H_i \bm J_i \bm J_i\T}_{O \times O}\big), 174 \end{align} 175 which means we can reduce the dimensionality of \(\bm \varepsilon_0\) to the number of model outputs \(O\) instead of the (significantly larger) number of parameters \(D\), reducing the variance and computational cost of the trace estimator. 176 Empirically, we find that, without the trace rearrangement in \cref{eq:trace-rearrangement}, the variance of the trace estimator is too high to be useful. 177 178 \subsection{Scaling} 179 180 In standard negative log-likelihood training, the only term in the loss is \(-\log p(\bm y \given \bm \theta)\). 181 For regression, the loss function which has a Bayesian interpretation as the negative log-likelihood is the sum of squared errors (SSE), though often the mean squared error is used instead, since computing the correct gradient just involves scaling the step size \(\eta\) up by the number of observations \(N\). 182 For negative log-posterior training, the same scaling can be applied with a Gaussian prior by dividing the prior precision by \(N\). 183 However, for the log-marginal likelihood, scaling the terms is no longer possible, since scaling the loss by a constant will not affect the gradients of the log-determinant term. 184 This means that it is important to compute the negative log-likelihood correctly as the SSE, since the gradients of the log-marginal likelihood are not invariant to scaling the log-marginal by a constant. 185 To compute the posterior precision using a mini-batch, we thus have to scale the likelihood term by \(N / B\) to account for the likelihood being computed on a mini-batch of size \(B\) instead of the full dataset of size \(N\), yielding \(N / B \sum_{i=1}^B \bm J_i \bm H_i \bm J_i + \alpha \identity\). 186 The log-determinant of this mini-batched precision is equivalent to 187 \begin{align} 188 \log\det \left( \frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \alpha \identity \right) 189 ={} & \log\det \left( \frac{N}{B} \left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \frac{B}{N} \alpha \identity \right) \right)\hspace{-0.05cm} 190 \\ ={} & D \log \frac{N}{B} + \log \det \left( \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i + \frac{B}{N} \alpha \identity \right). \nonumber 191 \end{align} 192 The scaling is thus important for determining the relative importance of the prior and likelihood terms in the eigenvalues. 193 To correctly compute the log-determinant of the posterior precision, we scale the trace and square traces for \cref{eq:mu1,eq:mu2} as 194 \begin{align} 195 \Tr\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) ={} & \frac{N}{B} \Tr\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right) \label{eq:laplace-trace-scaling}, 196 \\ \Tr\left(\left(\frac{N}{B} \textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right) ={}& \frac{N^2}{B^2} \Tr\left(\left(\textstyle{\sum_{i=1}^B} \bm J_i \bm H_i \bm J_i\right)^2\right). \label{eq:laplace-square-trace-scaling} 197 \end{align} 198 By combining \cref{eq:laplace-log-determinant-bound,eq:bai-golub-bound,eq:hutchinson-trace,eq:hutchinson-square-trace,eq:trace-rearrangement,eq:laplace-trace-scaling,eq:laplace-square-trace-scaling}, we can therefore compute a differentiable upper bound on the log-marginal likelihood under the Laplace approximation using only Jacobian-vector products without needing to instantiate the posterior precision.