Linearised Laplace approximations

The previous chapter showed a computationally lightweight Gaussian approximation $\mathcal{N}(\theta;\theta_*, \precision(\theta_*)^{-1})$ to the true posterior $p(\theta \mid \mathcal{D})$ on the parameter space $\mathbb{R}^{P}$. Our final goal is to use this Gaussian measure to make decisions based on the model’s output: its predictions feed into the risks and utilities that downstream tasks optimise (Chapter Analytic, low-cost functionality from Laplace). To this end, we first quantify the predictive uncertainty of our network on the output space $\mathbb{R}^{C}$ by turning it into a Gaussian process (GP). Figure 1 previews the pipeline, from the maximum a posteriori (MAP) estimate $\theta_*$ to the predictive distribution.

Linearisation¶

To turn an arbitrary (differentiable) neural network parameterised by $\theta \in \mathbb{R}^P$ into a Gaussian process, we use two ingredients: linearisation in the parameter space and the previously introduced Laplace approximation.

We first introduce the concept of the linearised network $f^{\theta_0}\colon \mathbb{R}^{D} \times \mathbb{R}^{P} \to \mathbb{R}^{C}$ as the first-order Taylor approximation of network $f$ around parameters $\theta_0 \in \mathbb{R}^{P}$:

where $\jaclogit(x, \theta_0) \in \mathbb{R}^{C \times P}$ is the network Jacobian evaluated at $(x, \theta_0)$.^[1] Analogously to Equation (1), we define the regularised empirical risk of the linearised network $f^{\theta_0}$ with parameters $\theta$ as

Unless explicitly stated otherwise (e.g., in the prior-only discussion in Section The Prior), we fix the linearisation point at a minimiser of the original regularised empirical risk $\regrisk(\theta)$:

that is, we set $\theta_0 \equiv \theta_*$. Linearisation itself only requires differentiability of $f$ at $\theta_0$. The Laplace approximation needs more: $\theta_*$ must be a local minimiser of $\regrisk$ with zero gradient, and the resulting precision must be positive definite after adding the prior curvature. We note that $f^{\theta_*}$ shares its value and its parameter Jacobian with $f$ at $\theta_*$, so $\nabla_\theta \regrisk^{\theta_*}(\theta_*)=\nabla_\theta \regrisk(\theta_*)=0$ and $\theta_*$ is a stationary point of $\regrisk^{\theta_*}$, too. Because $f^{\theta_*}$ is affine in $\theta$, convexity of $\ell(\cdot,y)$ in $z$ makes each loss term convex in $\theta$; with a strictly convex $\Omega$, the risk $\regrisk^{\theta_*}$ is then strictly convex in $\theta$, hence $\theta_*$ is its unique minimiser.

Following Section Laplace approximations are local and differential, the Laplace approximation using the regularised empirical risk of the linearised network yields the parameter-space Gaussian measure $\mathcal{N}(\theta; \theta_*, \precision(\theta_*)^{-1})$ with

The curvature $\precision$ of the linearised network’s regularised empirical risk differs from the Hessian in Equation (7) for networks that are not piecewise affine in their parameters (except at the non-generic parameter values where all training residuals vanish). The difference is that $\precision(\theta_*)$, as defined in Equation (4), discards the network’s parameter curvature (and higher-order derivatives) and only retains the curvature of the loss function over the model outputs. The advantages of this modelling choice are summarised in Section Advantages of Using the Linearised Network.

An important property of Gaussian measures is their closure under affine transformations, stated formally below.

As such, the linearised network $f^{\theta_*}$ pushes the parameter posterior forward, via the affine map of Equation (1), to a Gaussian in the output space (the logit space in classification). In the network-output variable $z\in\mathbb{R}^{C}$ (the logits, which the loss $\ell(z,y)$ takes as its first argument), this measure is

This result already equips the network $f$ with a probability measure in its output space (or logit space in classification) for every input $x_*$. A Gaussian process is an even more powerful construction, capable of modelling a Gaussian measure over any finite set of inputs, not just a single one. The linearised network $f^{\theta_*}$ equipped with its parameter-space Laplace approximation defines a Gaussian process $\GP(f(\cdot, \theta_*), k)$ with the following (Mercer) kernel:

Advantages of Using the Linearised Network¶

We perform the Laplace approximation w.r.t. $\regrisk^{\theta_*}$ (with $\theta_* = \operatorname*{arg\,min}_{\theta\in\mathbb{R}^{P}} \regrisk(\theta)$) instead of the regularised empirical risk of $f$ for multiple reasons.

First, the Hessian of the regularised empirical risk of $f$ (Equation (7)) is not guaranteed to be positive definite, preventing us from interpreting it as the inverse covariance matrix of our approximate Gaussian posterior over the parameters $\theta$.^[2] In second-order optimisation, the lack of positive definiteness also prohibits using the inverse Hessian as a preconditioner. However, the Hessian of the regularised empirical risk of $f^{\theta_*}$, with a strongly convex regulariser, is guaranteed to be positive definite.

Second, the Hessian of the regularised empirical risk of $f^{\theta_*}$ is considerably cheaper to use than that of $f$, allowing Laplace approximations to scale to deep learning architectures (see Section The Fisher).

There are thus two uses of the linearised model. First, it defines the curvature $\precision$ and hence the parameter-space Gaussian posterior approximation. Second, it defines the predictive pushforward in Equation (13). Earlier work on Laplace approximations for neural networks MacKay, 1992Ritter et al., 2018 sampled $M$ parameter configurations $\{\theta_s\}_{s=1}^{M}$ from the parameter-space Gaussian and calculated predictive quantities from the nonlinear predictions $\{f(x, \theta_s)\}_{s=1}^{M}$. This loses the exact output-space Gaussian distribution available from linearisation and yields no posterior GP. Lawrence (2000) shows empirically that this approach is prone to underfitting: the predictive becomes too uncertain even on in-distribution data points. Immer et al. (2021) identify the core reason for performance loss as a mismatch between the model in which we perform approximate Bayesian inference (the linearised network $f^{\theta_*}$) and the model we use for prediction (the nonlinear network $f$). Figure 2 visualises this underfitting and the low-density samples drawn under the true posterior. See also Immer et al. (2021, Fig. 3). By using the linearised network to make predictions and not only to calculate the curvature, the approximation remains internally consistent.

The Prior¶

The Gaussian process determined by the posterior parameter-space covariance matrix $\precision^{-1}$ is tied to the dataset $\mathcal{D}$, as $\precision$ is the Hessian of the linearised model’s regularised empirical risk over $\mathcal{D}$. It is worth spending some time on the case when the neural network $f\colon \mathbb{R}^{D}\times\mathbb{R}^{P}\to\mathbb{R}^{C}$ has not been exposed to any dataset $\mathcal{D}$ yet. In this case, the model’s behaviour is determined by its architecture and by the prior $p(\theta)$ over the parameters.

For a non-Gaussian parameter prior $p(\theta)$, a Laplace approximation to the prior is centred at a local maximum $\theta_0$ of $p(\theta)$, analogously to Equation (6). Its precision is the negative log-prior Hessian,

For the construction that follows, we specialise to a Gaussian prior centred at the linearisation point,

where $\theta_0$ can be any chosen initialisation. The prior’s mode is $\theta_0$, so the linearisation and the Laplace approximation of the prior share the same centre, and the precision above evaluates to $\precision = I$. Pushing this prior through the linearised network gives the kernel $k\colon \mathbb{R}^{D} \times \mathbb{R}^{D} \to \mathbb{R}^{C \times C}$,

The kernel comes from the same pushforward as the posterior GP kernel in Equation (14). Compared with it, the data-dependent pair $(\theta_*,\precision(\theta_*))$ is replaced by the dataset-independent pair $(\theta_0,I)$. This kernel is known as the Neural Tangent Kernel (NTK) Jacot et al., 2018.^[3] It differs from the neural network Gaussian process (NNGP) kernel: the NNGP kernel is the covariance of random network outputs $f(x,\theta)$ and $f(x',\theta)$ under a random initialisation distribution, while the NTK is the covariance of the linearised network outputs $f^{\theta_0}(x,\theta)$ and $f^{\theta_0}(x',\theta)$ under the prior $\mathcal{N}(\theta;\theta_0,I)$. As the network width is taken to infinity in a specific way, both kernels are convergent. Where these limits exist, they are different deterministic kernels.

The NTK determines the function-space trajectory of $f$ during full-batch gradient flow training on the average empirical loss via the Ordinary Differential Equation (ODE)

where $t$ is the time variable of the gradient flow and $k_t$ is the NTK of $f_t$. Jacot et al. (2018) show that in the infinite-width limit of fully-connected neural networks with the so-called NTK initialisation, the NTK stays constant during training, and the convergence of optimisation is related to the positive definiteness of the NTK. However, the NTK of regular, finite-width neural networks is not constant and must be recomputed at every evaluation point.

From a modelling perspective, putting an isotropic Gaussian prior over the parameters is seldom well-motivated. However, the framework of Laplace approximations is compatible with interpretable function-space priors. Cinquin et al. (2024) define Laplace priors using a prior Gaussian process. In particular, the negative log-prior is, up to an additive constant, equal to the squared Reproducing Kernel Hilbert Space (RKHS) norm of the difference between the linearised network $f^{\theta_0}$ and the mean of the prior Gaussian process $\mu$. This negative log-prior is also used as a regulariser during training. As specifying the Gaussian prior in the function space only requires knowledge about the function’s behaviour and not about the parameter space’s structure, the framework of Cinquin et al. (2024) allows domain experts to directly encode functional behaviour in the Bayesian inference problem.

The Fisher¶

In Section Linearisation, we glossed over the curvature matrix defined in Equation (4). This matrix and its components give an information-geometric interpretation to Laplace approximations, warranting its own chapter (Chapter The differential geometric perspective on Laplace).

Consider the regularised empirical risk of the linearised network $f^{\theta_0}$ (Equations (1) and (2)). Under the assumption of Equation (2) applied to the linearised network, the term

is equal, up to an additive constant, to the negative log-likelihood of the dataset under the model, $-\log p(Y \mid X, \theta)$, whereas the term $\Omega(\theta)$ equals the negative log-prior, $-\log p(\theta)$, again up to an additive constant.

Writing out the Hessian of $\mathcal{L}^{\theta_0}$ evaluated at $\theta_0$, we obtain

called the Generalised Gauss-Newton matrix (GGN). In this section, we show that this matrix is equivalent to the Fisher Information Matrix for a family of loss functions, which is a natural (usually degenerate) Riemannian metric for statistical manifolds parameterised by the model (Section The Fisher Metric).

Consider the statistical manifold

with conditional independence structure $p(x, y \mid \theta) = p(y \mid x, \theta)\cdot p(x)$, whose in-depth study we defer to Chapter The differential geometric perspective on Laplace. The $p(y \mid x, \theta)$ distribution represents the likelihood parameterised by a neural network and $p(x)$ is the ground-truth input distribution. A natural notion of separation between probability distributions on this manifold is the Kullback-Leibler (KL) divergence:

The KL divergence is not a valid metric, as it is non-symmetric. However, for parameters $\theta_1$ and $\theta_2$ and $\Delta\theta = \theta_1 - \theta_2$, the KL divergence’s second-order Taylor approximation around $\theta_1$ is given by

Both the zeroth- and first-order terms vanish (the former because $D_{\text{KL}}(p\ \Vert\ p)=0$, the latter because the expected score is zero, Equation (27)), leaving the quadratic Fisher term as the leading contribution. At each parameter $\theta$, the bilinear form $g_\theta(u, v) := u^{\top} \fisher(\theta) v$ on tangent vectors $u, v$ is a valid Riemannian metric over manifold $\mathcal{M}$ (up to the identifiability of parameters; see Chapter The differential geometric perspective on Laplace). The matrix

is called the Fisher Information Matrix (FIM) of the random variable $(\mathrm X, \mathrm Y)$ distributed according to $p(x, y \mid \theta)$.

Differentiating the normalisation $\int p(x, y \mid \theta) \,\mathrm{d} x \,\mathrm{d} y = 1$ under the integral sign and using $\nabla_\theta p = p \nabla_\theta \log p$ gives the zero-mean score,

Differentiating once more yields the two equivalent forms of the Fisher information,

The FIM of independent random variables is additive. In particular, for $N$ i.i.d. random variables $\{(\mathrm X_n, \mathrm Y_n)\}_{n=1}^N$, the FIM is $N$ times that of a single random variable $(\mathrm X, \mathrm Y)$:

In the following, we use the notation $\fisher(\theta)$ for the FIM of a single random variable $(\mathrm X, \mathrm Y)$ distributed according to $p(x, y \mid \theta)$.

Using Equation (25), the utility of the FIM becomes clear in second-order optimisation. In particular, consider the objective

for some constant $c \in \mathbb{R}_{> 0}$. The intuition behind this objective is that we want to minimise our model’s non-regularised empirical risk while limiting how much the likelihood induced by the model changes, irrespective of the curvature of the loss surface. To find an approximate minimiser of the objective in Equation (30), we consider the quadratic approximation to our constraint in Equation (25) and linearise the empirical risk around the current parameter $\theta$:

Substituting in the linearised empirical risk and quadratic constraint, our objective becomes

The above is a convex quadratically constrained linear program. When $\fisher(\theta)$ is positive definite and $\nabla_\theta\mathcal{L}(\theta)\neq 0$, the optimum lies on the boundary of the constraint ellipsoid and can be found in closed form with Lagrange’s method, yielding

In particular, the optimal update direction $\Delta\theta^*$ is given by the so-called natural gradient $\fisher(\theta)^{-1}\nabla_\theta \mathcal{L}(\theta)$ with a magnitude that is a function of the constant $c$. Natural-gradient methods are developed in detail in Chapter Natural-gradient-based variational inference.

In practice, we replace the unknown distribution of inputs $p(x)$ with the empirical distribution w.r.t. training set $\mathcal{D}$:

The corresponding FIM approximation becomes

Using the chain rule w.r.t. the model’s output $z = f(x_n, \theta)$, we obtain

Here the network’s own second-order term vanishes under the expectation by the zero-mean score in $z$ (Equation (27) applied to $p(y \mid z)$), which is why the linearised and nonlinear networks share the same $\tilde{\fisher}(\theta_0)$. The expectation over model outputs for each input can be expensive to compute. One approach to mitigate the cost is the Monte Carlo approximation:

with $y_s \sim p(y \mid x_n, \theta)$. For small $M$, we obtain a rather crude but cheap approximation. An alternative approach is to first rewrite $\tilde{\fisher}(\theta)$ as

(Equation (28), applied to $p(y \mid x_n, \theta)$), then substitute the ground-truth labels instead of approximating the expectation:

called the Empirical Fisher Information Matrix (EFIM). While this substitution may seem natural, Kunstner et al. (2019) show that it has a detrimental effect for optimisation, where the predicted and ground-truth labels can be significantly different, especially at the beginning of training. However, for the Laplace approximation, we use a MAP estimate, which, if it generalises well, approximates the ground-truth label distribution well.

For a wide family of likelihoods, one does not have to approximate the Fisher $\tilde{\fisher}(\theta)$ by Monte Carlo. This allows us to connect the FIM to the GGN under some assumptions. Identifying $\ell(z, y_n) = -\log p(y_n \mid z)$ and $f(x_n, \theta_0) = f^{\theta_0}(x_n, \theta_0)$ using Equation (1), in the case when the Hessian of the negative log-likelihood loss $\ell$ does not depend on the target $y$, the GGN is equal to the FIM (up to a constant scaling factor):

In fact, the difference in the scaling factor has a natural interpretation. Previously, we saw that $\fisher(\theta)$ is the FIM of the random variable $(\mathrm X, \mathrm Y)$ distributed according to $p(x, y \mid \theta)$ and that the FIM is additive for independent random variables. Therefore, under this assumption, the GGN is equal to the FIM w.r.t. $N$ i.i.d. observations $\{(x_n, y_n)\}_{n=1}^N$ distributed as $p(x, y \mid \theta)$ (and not as $p(x, y)$).

The requirement of the loss Hessian not depending on the target might seem arbitrary. However, this property is satisfied by a wide family of likelihoods, called exponential families.

First, we state some facts about simple parameterisations.

Parameterisations via neural networks are more delicate. The following proposition gives a necessary and sufficient condition for Equation (40).

The above statement on neural network parameterisations covers the most popular loss functions in classification and regression. In classification, the cross-entropy loss is the negative log-likelihood of a categorical distribution parameterised by our neural network. The logits are the categorical’s natural parameters, so with the identity link $g = \operatorname{Id}$ the cross-entropy loss satisfies Proposition 3 in its affine special case. Similarly, the $L_2$ loss for regression satisfies Proposition 3 with $g = \operatorname{Id}$.

A word of caution is in order: the GGN/FIM alone is usually degenerate; thus, it does not serve as a valid precision matrix for the Laplace approximation. Considering a single input-label pair $(x_n, y_n)$, the term

is a positive semidefinite (PSD) matrix, as the loss Hessian $\nabla_{z}\nabla_{z}^{\top} \ell(z, y_n)\big|_{z = f^{\theta_0}(x_n, \theta_0)}$ is positive semidefinite. For softmax cross-entropy, this Hessian is $\operatorname{diag}(p)-pp^{\top}$, which is rank deficient as the softmax is shift-invariant. Usually, the GGN, that is, the sum of these matrices over the training dataset, is still rank deficient. However, adding the Hessian of the strongly convex regulariser term $\Omega(\theta_0)$ to the GGN results in the positive definite matrix

which is the Hessian of the regularised empirical risk of the linearised model, $\regrisk^{\theta_0}(\theta_0)$. This justifies using this matrix as a parameter-space precision in the Laplace approximation (Equation (6)).

With the curvature $\precision$ in hand, we can place the NTK of Section The Prior and the posterior GP kernel of Equation (14) in a single frame. Both are Jacobian Gram matrices $\jaclogit M^{-1}\jaclogit^{\top}$ in a parameter-space metric $M$; they differ in the expansion point ($\theta_0$ for the NTK, $\theta_*$ for the posterior kernel) and in the metric $M$. Preconditioning the gradient flow of Equation (19) by $M$ multiplies the parameter velocity by $M^{-1}$, inserting $M^{-1}$ between the Jacobians and replacing its kernel $k_t$ with $\jaclogit(x,\theta_t)M^{-1}\jaclogit(x_n,\theta_t)^{\top}$, holding the expansion point fixed and changing only the metric. Vanilla gradient descent uses the Euclidean metric $M=I$ and recovers the NTK, the isotropic-prior pushforward. Regularised natural-gradient descent Amari, 1998 uses $M=\precision(\theta_t)$, the GGN/Fisher curvature plus the prior precision (Equation (43)) at the current parameters. We call the resulting kernel $\jaclogit(\cdot,\theta_t)\precision(\theta_t)^{-1}\jaclogit(\cdot,\theta_t)^{\top}$ the natural-gradient NTK. Because $\precision(\theta_t)$ is positive definite at every $\theta_t$, the natural-gradient NTK governs the dynamics all along training, even where the gradient is nonzero and no Laplace posterior exists. It coincides with the posterior GP kernel $\jaclogit\precision^{-1}\jaclogit^{\top}$ Immer et al., 2021 only at a posterior mode $\theta_*$, as the Laplace approximation requires. The parameter-space metric thus determines both the predictive uncertainty, through the Gaussian we push forward, and the optimisation geometry, through the descent the function follows.

Scalable Curvature Approximations¶

The GGN of the linearised network (cf. Equation (22))

is dense in general and far too large to form for modern neural networks. We therefore seek structured surrogates that (i) are cheap to store and invert, yet (ii) retain meaningful parameter correlations. Table 1 compares the three curvature proxies we develop: the diagonal, low-rank, and Kronecker-Factored Approximate Curvature (KFAC) approximations. Scalable variants of these curvature approximations are developed further in Chapter Scalable Laplace approximations, and last-layer inference is treated in Chapter Last-layer inference.

A simple baseline is the diagonal $\operatorname{diag} \ggn(\theta)$, which is trivial to store and invert, though calculating it efficiently requires care. However, it encodes an independence assumption between all parameters, destroying the often-crucial within-layer correlations carried by $\ggn(\theta)$. This tends to hurt the utility of the downstream uncertainty estimates Daxberger et al., 2021.

Low-rank GGN Approximations¶

A scalable yet expressive alternative is a low-rank approximation of $\ggn(\theta)$. As $\ggn(\theta)$ is symmetric for any $\theta \in \mathbb{R}^P$, it possesses an eigendecomposition with real eigenvalues of the form

where $Q(\theta) \in \mathbb{R}^{P \times P}$ is an orthogonal matrix (i.e., a matrix with orthonormal columns/rows) and $\Lambda(\theta) \in \mathbb{R}^{P \times P}$ is a diagonal matrix.

For $q_{:,i}$, the $i$th column of $Q(\theta)$, and $\lambda_i$, the $i$th diagonal entry of $\Lambda(\theta)$, it holds that $\ggn(\theta) q_{:, i} = \lambda_i q_{:, i}$. We follow the convention $\lambda_1 \ge \lambda_2 \ge \cdots \ge \lambda_{P}$. In the following, we drop the dependence of $Q$ and its factors on $\theta$ for brevity.

We now seek a good low-rank approximation $\tilde{\ggn}_k$ to $\ggn$, as this reduces the sizes of the objects we need to handle. Since $\ggn\succeq0$, its eigenvalues are nonnegative, so the $k$ algebraically largest coincide with the $k$ largest in magnitude. The Eckart–Young–Mirsky theorem Eckart & Young, 1936 states that the objective

is minimised by $\tilde{\ggn}_k := Q_{1:k} \Lambda_{1:k} Q_{1:k}^{\top}$ where $Q_{1:k} \in \mathbb{R}^{P \times k}$ is the first $k$ columns of $Q$ and $\Lambda_{1:k} \in \mathbb{R}^{k \times k}$ holds the first $k$ eigenvalues of $\ggn$. This result holds for both the spectral norm $\Vert~\cdot~\Vert_2$ and the Frobenius norm $\Vert~\cdot~\Vert_F$. Therefore, one obtains an optimal low-rank approximation $\tilde{\ggn}_k$ to $\ggn$ in the above sense by computing the first $k$ eigenvalues and eigenvectors of $\ggn$. Storing this low-rank approximation requires storing $(P + 1) k$ values instead of $P^2$.

The most naive way to compute $\tilde{\ggn}_k$ is first calculating the eigendecomposition $\ggn = Q \Lambda Q^{\top}$ and then truncating it to the first $k$ eigenpairs. However, this strategy does not resolve the infeasibility of storing the full eigendecomposition with $P(P + 1)$ values. Instead, one uses efficient matrix-free schemes such as the Locally Optimal Block Preconditioned Conjugate Gradient (LOBPCG) method Knyazev, 2001 or Lanczos iterations Lanczos, 1950 to obtain the low-rank decomposition. Notably, these algorithms only require matrix-vector products $v \mapsto \ggn v$ with the GGN to obtain the top eigenpairs, which does not require representing the entire $\ggn$ matrix in memory.

The low-rank representation is particularly handy, as for simple priors, such as isotropic Gaussians, calculating the approximate covariance matrix is just as efficient as working with the low-rank GGN. Let $\Gamma\succ0$ be the Hessian of the negative log-prior $\Omega$, that is, the prior-curvature term $\nabla_\theta\nabla_\theta^{\top}\Omega$ of the precision $\precision=\ggn+\nabla_\theta\nabla_\theta^{\top}\Omega$ (Equation (43)). Then the Woodbury matrix identity yields

In particular, for isotropic Gaussians, $\Gamma=\tau I$,

In words: a diagonal-plus-low-rank matrix has a diagonal-plus-low-rank inverse. Multiplying by such a matrix is efficient and does not require forming a $P \times P$ matrix. For a vector $v \in \mathbb{R}^{P}$, we first calculate $u=Q_{1:k}^{\top} v \in \mathbb{R}^k$, then elementwise multiply $u$ by the diagonal entries in $\operatorname{diag}\left( -\frac{\lambda_1}{\tau(\lambda_1+\tau)},\ldots,-\frac{\lambda_k}{\tau(\lambda_k+\tau)} \right)$, then map the result back to $\mathbb{R}^{P}$ with $Q_{1:k}$ and add the diagonal term $\tau^{-1}v$.

Kronecker-Factored Approximate Curvature¶

A complementary path towards scalable approximate curvature representations assumes independence between layers only, while retaining rich within-layer parameter correlations via a Kronecker factorisation. This idea underlies KFAC, which has recently seen wide use as a scalable curvature proxy Martens & Grosse, 2015. In this section, we give an introduction to KFAC for dense linear layers and refer to Dangel et al. (2025) for a more detailed and general exposition of KFAC.

Consider layer $i$ with weights $W^{(i)}\in\mathbb{R}^{d_i\times d_{i-1}}$, input activation $a_n^{(i-1)}\in\mathbb{R}^{d_{i-1}}$, and pre-activation $s_n^{(i)}=W^{(i)}a_n^{(i-1)}\in\mathbb{R}^{d_i}$.^[4] Let $z_n=f(x_n,\theta)\in\mathbb{R}^{C}$ denote the network output. Let $J_{W^{(i)}} s_n^{(i)} \in \mathbb{R}^{d_i \times (d_i d_{i-1})}$ denote the Jacobian of $s_n^{(i)}$ w.r.t. $\operatorname{vec}(W^{(i)})$, where $\operatorname{vec}$ flattens its input by stacking its columns. For datum $n$, define the output-space curvature $H_n:=\nabla_{z}\nabla_{z}^{\top}\ell(z,y_n)\big|_{z=z_n}\succeq0$, so $H_n\in\mathbb{R}^{C \times C}$. Writing $H_n=S_nS_n^{\top}$ with $S_n\in\mathbb{R}^{C \times C}$, the corresponding layer block of the GGN can be expressed as

with $g_{n,c}^{(i)}:=(J_{s_n^{(i)}} z_n)^{\top}[S_n]_{:,c}\in\mathbb{R}^{d_i}$, where $J_{s_n^{(i)}} z_n$ is the Jacobian of the network output $z_n$ w.r.t. this layer’s pre-activation $s_n^{(i)}$.^[5] Thus, the sum over $c$ is over the network output curvature factor, while $g_{n,c}^{(i)}$ lives in the layer pre-activation space.

For a linear layer, the output-parameter Jacobian has the form $J_{W^{(i)}} s_n^{(i)} = (a_n^{(i-1)})^{\top} \otimes I_{d_i}$ where $\otimes$ denotes the Kronecker product. Noting the Kronecker properties $(A \otimes B)^{\top} = A^{\top} \otimes B^{\top}$ and $(A \otimes B) (C \otimes D) = (AC) \otimes (BD)$ for compatible matrices $A, B, C, D$, we get

The KFAC approximation then pulls the Kronecker product outside the sum:

Thus, the layer block is approximated by

Stacking layers and neglecting inter-layer couplings yields a block-diagonal curvature with Kronecker blocks $A^{(i)}\otimes B^{(i)}$.

The Kronecker form is efficient to store. While the dense GGN $\ggn^{(i)}$ for a single linear layer has $(d_{i-1}d_i)^2$ entries, the Kronecker factorisation only requires storing $A^{(i)}$ and $B^{(i)}$, that is, a total of $d_{i-1}^2+d_i^2$ values. The block-diagonal layer approximation then avoids storing cross-layer blocks. One can also calculate efficient covariance-vector products without forming a dense $P \times P$ matrix by leveraging the cached eigendecompositions of the Kronecker factors and standard Kronecker identities. This property is more involved to show than in the case of low-rank approximations; refer to Dangel et al. (2025) for more details.

The Predictive¶

Equation (13) gives a Gaussian distribution in the output space for input $x_*$. For regression, this distribution is called the predictive distribution: its mean encodes the average prediction under the parameter-space posterior and the linearised model, and its covariance structure quantifies the uncertainty in the prediction, including correlations of the outputs. Figure 3 shows a one-dimensional regression example of the linearised Laplace approximation.

In classification, however, we are interested in obtaining a second-order distribution (a distribution over distributions) over the categorical distributions of a discrete set of labels. To this end, one has to map the logit-space Gaussian measure in Equation (13) through an activation function onto the ($C - 1$)-dimensional probability simplex:

where $\mathrm f$ is the mapping from an input $x_*$ to the logit-space Gaussian distribution given by the Laplace-approximated network, $\mathcal G(\mathbb{R}^{C})$ is the set of Gaussian measures on the logit space $\mathbb{R}^{C}$, $\mathcal P\left(\Delta^{C-1}\right)$ is the set of probability measures on the simplex, and $\varphi_*$ pushes forward the Gaussian probability measures through the activation function. The roman $\mathrm{f}$ is the network with its parameter posterior pushed through it: the affine pushforward of the Laplace posterior $\mathcal{N}(\theta;\theta_*,\precision(\theta_*)^{-1})$ along the linearised network, giving $\mathrm{f}(x_*)=\mathcal{N}(f(x_*,\theta_*), \jaclogit(x_*,\theta_*)\precision(\theta_*)^{-1}\jaclogit(x_*,\theta_*)^{\top})\in\mathcal G(\mathbb{R}^{C})$ as in Equation (13). The italic $f(x_*,\theta_*)\in\mathbb{R}^{C}$ plugs in the point estimate $\theta_*$ and returns a single logit vector; the roman $\mathrm{f}(x_*)$ returns the full logit-space Gaussian, whose mean is exactly that italic vector. All parameters in the support of the Laplace posterior are taken into account rather than a single one being held fixed, so $\mathrm{f}$ takes only the input. In words, $\mathrm h$ is the composition of the Laplace-approximated neural network outputting logit-space Gaussian distributions, $\mathrm{f}$, and the pushforward map through the activation function that gives us a distribution over probability vectors.

In most cases, this activation function is the softmax, defined as $\operatorname{softmax}\colon \mathbb{R}^{C} \to \Delta^{C - 1}$,

In this case, there exists no closed-form expression for $\mathrm h(x_*) \in \mathcal{P}\left(\Delta^{C - 1}\right)$ given input $x_*$. Further, $\mathbb{E}_{P \sim \mathrm h(x_*)}\left[P\right]$, called the posterior predictive, is also intractable. However, sampling from $\mathrm h(x_*)$ is efficient, as it only requires sampling from the logit-space Gaussian $\mathrm{f}(x_*)$ and pushing these logit samples forward through the activation function. Thus, in practice, one resorts to Monte Carlo integration, that is, drawing $M$ samples from $\mathrm h(x_*)$ and calculating their empirical average. Similarly, one can approximate other derived quantities, such as the expected entropy $\mathbb{E}_{P \sim \mathrm h(x_*)}\left[\ShanEntSymb(P)\right]$, through Monte Carlo sampling. Such derived quantities are often used in the field of uncertainty disentanglement, where different quantities capture different sources of uncertainty Mucsányi et al., 2024. These uncertainty estimates are applied to selective classification in Chapter Confidence-based selective classification.

The main issue with the above schemes is the requirement of sampling: the error of Monte Carlo integration decreases as the inverse square root of the number of samples $M$, meaning that to reduce the error by three orders of magnitude, one needs a million samples. Standard softmax-cross-entropy integration cannot avoid this sampling cost without incurring significant approximation error. However, some output activation functions admit a closed-form posterior predictive. For a single output, integrating a normCDF (probit) link against a Gaussian is exact, and the logistic sigmoid admits a classical probit approximation MacKay, 1992. Mucsányi et al. (2025) show that the normalised elementwise sigmoid and normCDF functions yield closed-form predictives that have strong approximation guarantees and are naturally compatible with more performant training objectives than the standard softmax-cross-entropy objective. Further, derived quantities such as the expected entropy can also be obtained in closed form by matching the moments of the Gaussian pushforward distributions with Dirichlets. The change in the training objective comes at no cost: models trained this way become better calibrated while providing sample-free posterior predictives and other derived quantities. Therefore, Laplace approximations in classification can be formulated in a purely sample-free fashion, with theoretical guarantees. Figure 4 compares the MAP, Monte Carlo softmax, and sample-free probit predictives on a two-dimensional example, together with the predictive entropy.