The differential geometric perspective on Laplace

Linearised Laplace approximations (Section Linearisation), and specifically their covariance structure, have interpretations in Riemannian and information geometry. They encode the neural network’s parameter space geometry, as informed by the loss.

The GGN Metric¶

In the previous chapter, we discussed in detail the curvature matrix defined in Equation (4), called the Generalised Gauss-Newton matrix (GGN). It is the precision term from the likelihood and, combined with the precision term from the prior and inverting, it yields the covariance matrix of a linearised Laplace approximation.

Recall that the GGN matrix at the parameters $\theta\in\mathbb{R}^{P}$ is given by (Equation (22))

We will now see how the GGN arises as a metric on the parameter space. Recall our training data $\mathcal{D}=\{(x_n,y_n)\}_{n=1,\dots,N}$. We restrict the neural network $f\colon \mathbb{R}^D\times\mathbb{R}^P \to \mathbb{R}^C$ to the training set, to obtain the function

For a single input-label pair $(x_n, y_n)$, assuming $\ell(z,y_n)$ is strictly convex (which is the case, for example, for the squared loss in regression), the positive definite loss Hessian

characterises a Riemannian metric in the output space $\mathbb{R}^{C}$. That is, it defines a symmetric positive definite bilinear form $(v_1,v_2)\mapsto v_1^\top H_n(z)v_2$ at each point $z\in \mathbb{R}^{C}$. Such a metric encodes the curvature of the loss given the target $y_n$. Note that if the likelihood is an exponential family, then $H_n$ is independent of the target $y_n$ (Proposition 2).

In fact, we only require $\ell(z,y_n)$ to be (not necessarily strictly) convex, in which case $H_n$ is only positive semidefinite. This is the case for the cross-entropy loss in classification, where the Hessian is rank-deficient. In this case, $H_n$ defines a degenerate Riemannian metric on $\mathbb{R}^{C}$.

We can then equip the image of $f_X$, $(\mathbb{R}^C)^N$, with the block diagonal (degenerate) Riemannian metric

in standard coordinates, where $z = (z_1,\dots , z_N)$. This metric characterises the curvature of the loss on the training set. It can then be pulled back through $f_X$ to the parameter space:

i.e., the pullback $f_X^*H$ corresponds exactly to the GGN $G$.

Due to symmetries in the neural network architecture (see Part 7 — Identifiability and symmetries) and overparametrisation, the GGN is rank-deficient and only positive semidefinite. It therefore defines again a degenerate Riemannian metric on the parameter space.^[1] Thus, damping is needed to get a proper Riemannian metric on the parameter space. The strictly convex regulariser term $\Omega(\theta)$ is such a dampener: adding its Hessian to the GGN results in a positive definite matrix (cf. Equation (43))

Using this as a preconditioner in second-order optimisation means we perform Riemannian (stochastic) gradient descent in the parameter space.

But how do we interpret the GGN metric? We base our presentation on Roy et al. (2024). Fixing a parameter vector $\theta_0$, consider the set

i.e., the set of parameters $\theta$ for which the neural network defines the same function on the training data as $\theta_0$. This is an equivalence class of parameters, called a fiber. By the preimage theorem Tu, 2011, Theorem 9.9, it is (locally) a manifold of dimension $P-N C$ if $f_X$ is (locally) continuously differentiable and the Jacobians $J_\theta f_X(\theta)$ are surjective for all (local) $\theta\in \mathcal M(\theta_0)$. Intuitively, the parameter space is $P$-dimensional but we constrain the model’s prediction on the training set of $N$ samples, meaning we put $N C$ scalar constraints on the output. This eliminates $N C$ degrees of freedom, yielding a $(P - N C)$-dimensional manifold. The surjectivity condition holds in particular when the eigenvalues of the neural tangent kernel matrix over the entire training set are bounded away from 0, a common assumption for overparametrised models, see Roy et al. (2024, Appendix A) and the references therein.

Further, when the Jacobians $J_\theta f_X(\theta)$ are surjective locally around $\theta_0$, the preimage theorem further tells us that $\ker J_{\theta}f_X(\theta_0) = T_{\theta_0}\mathcal M(\theta_0)$, the tangent space of $\mathcal M(\theta_0)$ at $\theta_0$ (see Figure 1). If the Hessians $\nabla_{z}\nabla_{z}^{\top} \ell(z, y_n)$ are full rank (for example, when $\ell(z,y_n) = \|z-y_n\|^2$ the squared loss), we get $\ker G(\theta_0) = T_{\theta_0}\mathcal M(\theta_0)$. These are directions in which the training data gives us no information about the parameters. Thus, in such directions, linearised Laplace approximations solely rely on the prior. The orthogonal complement of this null space consists of the directions in which the training data gives us information.

Now having defined the fibers in parameter space, we can quotient them out to obtain the effective parameter space: $\mathbb{R}^P/\sim$ where $\sim$ is the equivalence relation induced by $f_X$, $\theta_0 \sim \theta_1 \Leftrightarrow f_X(\theta_0) = f_X(\theta_1)$. The effective parameters are parameters for which the neural network $f$ takes distinct values on the training data. Formally, the effective parameter space is a fiber bundle over the parameter space. It can be thought of as a manifold perpendicular to the fibers $\mathcal M(\theta)$; Roy et al. (2024) construct such a manifold explicitly. We note, however, that this is an arbitrary choice, as there is no canonical way to choose representatives for the equivalence classes and hence such a manifold. As a consequence, the non-linear effective parameter space is difficult to leverage computationally. However, the GGN’s degenerate metric on the original parameter space $\mathbb{R}^P$ captures the similarity of training outputs between different parameters. One therefore does not have to construct the effective parameter space explicitly, but rather use the GGN’s geometry in the parameter space.

The Fisher Metric¶

In the previous chapter we saw that the GGN metric can be seen as the pullback of the loss curvature metric on the training outputs. We now ask whether there is a more intrinsic, parametrisation-independent justification for this metric. We offer a such a perspective on the GGN metric via information geometry.

Recall from Chapter Linearised Laplace approximations, Equation (26) the Fisher information metric

which we now view as a degenerate Riemannian metric on the parameter space $\mathbb{R}^P$. If we take the empirical distribution on the $x$, $\frac{1}{N}\sum_{n=1}^N \delta_{x_n}$, we get (Equations (35) and (38))

We saw in Chapter Linearised Laplace approximations that $\tilde{F}(\theta) = \frac{1}{N}G(\theta)$, a renormalisation of the GGN metric, under the mild condition that the Hessian of the negative log-likelihood loss does not depend on the target $y$ (see Equation (40)).

This gives an important connection between the covariance used in linearised approximation and the fundamental object from information geometry, which is the Fisher information metric. The goal of this chapter is to explain the importance of the Fisher information metric, why it can be seen as the natural metric for any probabilistic model, and its interpretation in deep learning.

A statistical manifold is a manifold $\mathcal M$ consisting of probability measures. Any parametric statistical model defines a parametrisation of a statistical manifold $\mathbb{R}^P \to \mathcal M$. What is $\mathcal M$ for a neural network? So far, we have viewed a neural network as a map $f\colon \mathbb{R}^D\times \mathbb{R}^P \to \mathbb{R}^C$. However, recall that we assume it induces probability distributions $p(y\mid x,\theta)$, since the loss function $\ell(f(x,\theta), y)$ is assumed to be the negative log-likelihood $-\log p(y \mid x,\theta)$. Therefore we implicitly viewed $f$ as a probabilistic model; for each $x\in \mathbb{R}^D$, and $\theta\in \mathbb{R}^P$, we have a probability measure $\mathrm f(x,\theta) \in \mathcal P(\mathcal{Y})$ defined by the neural network, where $\mathcal{Y}$ is the set of labels ($\mathbb{R}^C$ for regression, $\{1,\dots,C\}$ for classification) and $\mathcal P(\mathcal{Y})$ is the space of probability measures on the set of labels. $\mathrm f$ is given by the composition of $f$ with some appropriate $\mathrm g\colon \mathbb{R}^C \to \mathcal P(\mathcal{Y})$. In regression, we most often take $\mathrm g(z) = \mathcal N(z, \epsilon^2 I) \in \mathcal P(\mathbb{R}^C)$ for some fixed $\epsilon>0$, while in classification, most often $\mathrm g(z) = \operatorname{softmax}(z) \in \Delta^{C-1} \cong \mathcal P(\{1,\dots,C\})$. In Chapter Linearised Laplace approximations, we saw that both of these choices guarantee the equivalence between the Fisher and the GGN (Equation (40)).

So we have a map $\mathrm f\colon \mathbb{R}^D\times \mathbb{R}^P \to \mathcal P(\mathcal{Y})$. Write $p(y\mid x,\theta)$ for the density of $\mathrm f(x,\theta)$, so $p(y\mid x,\theta) = \frac{1}{\sqrt{2\pi}\epsilon}\exp\left(-\frac{1}{2\epsilon^2}\|f(x,\theta)-y\|^2\right)$ in regression, $p(y\mid x,\theta) = \operatorname{softmax}(f(x,\theta))_y$ in classification. Then the statistical model $\mathrm m\colon \mathbb{R}^P \to \mathcal P(\mathbb{R}^D\times \mathcal{Y})$ is defined as the joint distribution over $x$ and $y$. Namely, if the distribution over the inputs $x$ has a density $p(x)$, the density of $\mathrm m(\theta)$ is $p(x,y\mid \theta) := p(x)p(y\mid x,\theta)$. As discussed above, in practice we take the empirical distribution over $x$. The statistical manifold $\mathcal M$ defined by the neural network $f$ is the image of $\mathrm m$ in $\mathcal P(\mathbb{R}^D\times \mathcal{Y})$.

This is where the Fisher information metric comes into play. The Fisher metric is an object intrinsic to the statistical manifold $\mathcal M$, and is independent of how the manifold is parametrised by $\mathbb{R}^P$. Formally, suppose we have a different parametrisation $\tilde {\mathrm m}\colon \mathbb{R}^P \to \mathcal M$. Then $\mathrm m$ and $\tilde {\mathrm m}$ differ by some diffeomorphism^[2] of the parameter space, i.e. there is a diffeomorphism $\varphi\colon \mathbb{R}^P\to \mathbb{R}^P$ such that $\mathrm m = \tilde {\mathrm m}\circ \varphi$. Then we see that for $\tilde \theta = \varphi(\theta)$, $\mathrm m(\theta) = \tilde {\mathrm m}(\tilde \theta)$ or, abusing notation, $p(x,y\mid \theta) = p(x,y\mid \tilde \theta)$. Therefore, from the definition of the Fisher information metric (Equation (9) or Equation (10)), the Fisher information metric on the two distinct parametrisations coincide $F_{\mathrm m}(\theta) = F_{\tilde{\mathrm m}}(\tilde \theta)$.

The fact that the Fisher metric does not depend on the parametrisation implies that it is intrinsic to $\mathcal M$. Therefore, there should exist a way to define it directly on $\mathcal M$ in an abstract, measure-theoretic way. Doing this rigorously requires some mathematical groundwork, which the keen reader can find in Ay et al. (2017). Briefly, for $\mu\in \mathcal M$ and tangent vectors $\nu_1,\nu_2\in T_\mu\mathcal M$, which are signed measures integrating to 0, the intrinsic Fisher information metric is defined by

where $\frac{\,\mathrm{d}\nu_i}{\,\mathrm{d}\mu}$ is the Radon-Nikodym derivative of $\nu_i$ with respect to $\mu$, i.e. the densities of $\nu_i$ with respect to $\mu$. Then the degenerate metric on the parameter space $\mathbb{R}^P$ is given by the pullback of this intrinsic metric onto $\mathbb{R}^P$ (see Ay et al. (2017)).

There is an infinite number of ways to define an intrinsic Riemannian metric on $\mathcal M$. So what makes the Fisher information metric special? It is not special for its invariance under transformations of the parameter space, but instead for its invariance under transformations of the data space.

The factorisation condition means that the likelihood depends on $\theta$ only through $T(x)$, since the factor $h(x)$ has no $\theta$-dependence. Thus, all information about $\theta$ contained in the full data point $x$ is also contained in the statistic $T(x)$, justifying the term ‘sufficient’.

In our deep learning setting above, the data space is $\mathcal D = \mathbb{R}^D\times\mathcal Y$, the product of the input space with the label space.

Note that, for a statistical manifold $\mathcal M$ over $\mathcal D$, a statistic $T\colon \mathcal D \to \mathcal D'$ induces a map $T_*\colon \mathcal M \to \mathcal M'$ acting by pushforwards: $T_*\mu = \mu\circ T^{-1}$. So $\mathcal M'$ is a statistical manifold over $\mathcal D'$. Then, under mild assumptions on $\mathcal M$, we obtain the following theorem:

Up to a multiplicative constant, the Fisher information metric is the unique metric satisfying this. This is Chentsov’s theorem. We leave as an exercise to verify this theorem for the special case of the deep learning setting above.

As explained above, sufficient statistics should be thought of as maps of the data space which do not induce a loss of information on the parameters with respect toa given statistical model. The Fisher information metric, being invariant under such transformations, should be thought of as quantifying ‘information’, hence its name. Note in particular that it is invariant under diffeomorphisms of the data space (which are trivially sufficient statistics owing to their bijective nature), and so does not depend on how the data is ‘presented’. This makes it the natural metric for statistical manifold $\mathcal M$.

Geometric Laplace Approximations¶

Laplace approximations, being Gaussian approximations to the neural network’s posterior, are limited in their expressivity. As a consequence, there have been attempts at leveraging the geometry of the parameter space to define more expressive variants of Laplace approximations.

A key observation is that the Laplace approximations are not diffeomorphism-invariant: given a diffeomorphism of the parameter space $\varphi\colon \mathbb{R}^P\to \mathbb{R}^P$, if such diffeomorphism is non-linear, then performing Laplace with respect to the parameters $\varphi$, i.e., a reparametrised Laplace approximation, yields a distinct posterior approximation than performing Laplace with respect to the parameters $\theta$. In the extreme case, one can show there exists a diffeomorphism $\varphi$ for which the Laplace approximation with respect to the corresponding parameters is exact Costa et al., 2025. In practice, computing such $\varphi$ is as difficult as computing the exact posterior itself. However, this result gives hope that one may be able to construct computationally feasible diffeomorphisms $\varphi$ yielding high-quality Laplace approximations, see for instance MacKay (1998). An important note is that, while Laplace approximations are not invariant under diffeomorphisms of the parameter space, the predictives of linearised Laplace approximations are, up to changing the prior. We leave this as an exercise to the reader.

Another attempt at leveraging geometry for Laplace is to perform the Laplace approximation in the tangent space of the parameter space and then project it through an exponential map under some Riemannian metric Bergamin et al., 2023Yu et al., 2024, coined Riemannian Laplace approximations. The latter work uses the Fisher as the Riemannian metric. Work remains to be done to make the computation of the exponential map efficient and scalable. See Costa et al. (2025) for the relationship between such Riemannian Laplace approximations and the aforementioned reparametrised Laplace approximations.

Finally, diffusion Laplace approximations Roy et al., 2024 sample the approximate posterior through random walks with respect to a Riemannian metric in the weight space. Once again, while there are promising preliminary results, further work is needed to scale such approaches computationally. Moreover, the relationship between diffusion Laplace approximations and Riemannian Laplace approximations is yet to be described.