SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
(\det  \hat {\boldsymbol {E}}) (\boldsymbol {\psi}) = (\det  \boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})) ^ {- 1} (\det  \boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})) \det  (\boldsymbol {G} (\varphi^ {- 1} (\boldsymbol {\psi})) \boldsymbol {H} (\varphi^ {- 1} (\boldsymbol {\psi}))) = (\det  \boldsymbol {E}) (\varphi^ {- 1} (\boldsymbol {\psi}))
$$
This is because one can view the loss landscape as a hypersurface:  $\pmb{H}$  is connected to the second fundamental form and  $\pmb{E}$  to the shape operator. det  $\pmb{E}$  is thus related to the invariant Gaussian curvature [47, Ch. 8].
in  $\psi$ -coordinates. Thus, we have the desired invariant transformation  $(\det \hat{\pmb{E}}) = (\det \pmb{E}) \circ \varphi^{-1}$ . Note that  $\pmb{G}$  is an arbitrary metric—this invariance thus also holds for the Euclidean case where  $\pmb{G} \equiv \pmb{I}$ . Note further that the trace and eigenvalues of  $\pmb{E}$  are also invariant; see Appendices B.2 and B.3. Finally, see Appendix A for a simple working example.
To obtain invariance in the Hessian-determinant, -trace, and -eigenvalues at a minimum of  $\mathcal{L}$ , we explicitly write it as  $E = G^{-1}H$ , even when  $G \equiv I$ , and transform it according to Section 2.2.1.
3.2 Invariance of gradient descent
Viewed from the geometric perspective, both gradient descent (GD) and natural gradient descent (NGD) come from the same ODE framework  $\dot{\pmb{\theta}} = -G(\pmb{\theta})^{-1}\nabla \mathcal{L}|_{\pmb{\theta}}$ . But NGD is widely presented as invariant under reparametrization, while GD is not [51, 71, etc.]. Is the choice of the metric  $\pmb{G}$  the cause? Here we will show from the framework laid out in Section 2.2 that any metric is invariant if its transformation rule is faithfully followed. And thus, the invariance of NGD is not unique. Rather, the Fisher metric used in NGD is part of the family of metrics that transform correctly under autodiff, and thus it is "automatically" invariant under standard deep learning libraries like PyTorch, TensorFlow, and JAX [1, 24, 66]—see Appendix C.
The underlying assumption of GD is that one works in the Cartesian coordinates and that  $G \equiv I$ . For this reason, one can ignore the metric  $G$  in  $\dot{\theta} = -G(\theta)^{-1}\nabla \mathcal{L}|_{\theta}$ , and simply write  $\dot{\theta} = -\nabla \mathcal{L}|_{\theta}$ . This is correct but this simplification is exactly what makes GD not invariant under a reparametrization  $\varphi : \Theta \to \Psi$  where  $\psi = \varphi(\theta)$ . To see this, notice that while GD transforms  $\nabla \mathcal{L}$  correctly via the chain rule  $\dot{\psi} = -J^{-1}(\psi)^{\top}\nabla \mathcal{L}|_{\varphi^{-1}(\psi)}$ , by ignoring the metric  $I$ , one would miss the important fact that it must also be transformed into  $\hat{G}(\psi) = J^{-1}(\psi)^{\top}J^{-1}(\psi)$ . It is clear that  $\hat{G}(\psi)$  does not equal  $I$  in general. Thus, we cannot ignore this term in the transformed dynamics since it would imply that one uses a different metric—the dynamics are thus different in the  $\theta$ - and  $\psi$ -coordinates. When one explicitly considers the above transformation, one obtains
$$
\begin{array}{l} \dot {\psi} = - \hat {G} (\psi) ^ {- 1} J ^ {- 1} (\psi) ^ {\top} \nabla \mathcal {L} | _ {\varphi^ {- 1} (\psi)} = - \left(J ^ {- 1} (\psi) ^ {\top} J ^ {- 1} (\psi)\right) ^ {- 1} J ^ {- 1} (\psi) ^ {\top} \nabla \mathcal {L} | _ {\varphi^ {- 1} (\psi)} \\ = - \boldsymbol {J} (\varphi^ {- 1} (\psi)) \nabla \mathcal {L} | _ {\varphi^ {- 1} (\psi)}. \\ \end{array}
$$
This dynamics in  $\psi$ -coordinates preserves the assumption that the metric is  $\pmb{I}$  in  $\theta$ -coordinates and thus invariant. In contrast, the dynamics  $\psi$  under just the chain rule implicitly changes the metric in  $\theta$ -coordinates, i.e. from  $\pmb{I}$  into  $\pmb{J}(\pmb{\theta})^{\top}\pmb{J}(\pmb{\theta})$ , and thus the trajectories are not invariant.
This discussion can be extended to any metric  $\pmb{R}$ : Simply use the transformed metric  $\hat{\pmb{R}}(\psi) = J^{-1}(\psi)^{\top}\pmb{R}(\varphi^{-1}(\psi))\pmb{J}^{-1}(\psi)$ , and we obtain the invariant dynamics of any preconditioned GD with preconditioner  $\pmb{R}$  given by  $\psi = -J(\varphi^{-1}(\psi))\pmb{R}(\varphi^{-1}(\psi))^{-1}\nabla\mathcal{L}|_{\varphi^{-1}(\psi)}$ .
To obtain invariance in optimizers with any metric/preconditioner, explicitly write down the metric even if it is trivial, and perform the proper transformation under reparametrization.
Remark 5. For the discretized dynamics, the larger the step size, the less exact the invariance. For instance, discrete natural gradient update rule is only invariant up to the first-order [50, 71]. This is orthogonal to our work since it is about improving ODE solvers [71].
The consequence is that we need a "geometric-aware" autodiff library such that the invariant dynamics above is always satisfied. In this case, any preconditioner  $R$  yields invariance under any reparametrization, even nonlinear ones. This is in contrast to the current literature, e.g. under standard autodiff, Newton's method is only affine-invariant, and structured approximate NGD methods such as K-FAC are only invariant under a smaller class of reparametrizations [56].
3.3 Invariance in probability densities
Let  $q_{\Theta}(\pmb{\theta})$  be a probability density function (pdf) under the Lebesgue measure  $d\pmb{\theta}$  on  $\Theta$ . Under a reparametrization  $\varphi: \Theta \to \Psi$  with  $\psi = \varphi(\pmb{\theta})$ , it transforms into  $q_{\Psi}(\psi) = q_{\Theta}(\varphi^{-1}(\psi))|\det J^{-1}(\psi)|$ . This transformation rule ensures  $q_{\Psi}$  to be a valid pdf under the Lebesgue measure  $d\psi$  on  $\Psi$ , i.e.  $\int_{\Psi} q_{\Psi}(\psi) d\psi = 1$ . Notice, in general  $q_{\Psi} \neq q_{\Theta} \circ \varphi^{-1}$  due to the change-of-random-variable rule. Hence, pdfs transform differently than standard functions (Example 4a)
and can thus have non-invariant modes [cf. 59, Sec. 5.2.1.4]: Density maximization, such as the optimization of  $\mathcal{L}_{\mathrm{MAP}}$ , is not invariant even if an invariant optimizer is employed.
Just like the discussion in the previous section, it is frequently suggested that to obtain invariance here, one must again employ the help of the Fisher matrix  $\pmb{F}$  [21]. When applied to a prior, this gives rise to the famous Jeffreys [31] prior  $p_J(\pmb{\theta}) \propto |\det \pmb{F}(\pmb{\theta})|^{\frac{1}{2}}$  with normalization constant  $\int_{\Theta} |\det \pmb{F}(\pmb{\theta})|^{\frac{1}{2}} d\pmb{\theta}$ . Is the Fisher actually necessary to obtain invariance in pdf maximization?
Here, we show that the same principle of "being aware of the implicit metric and following proper transformation rules" can be applied. A pdf  $q_{\Theta}(\theta)$  can be written as  $q_{\Theta}(\theta) = \frac{q_{\Theta}(\theta)d\theta}{d\theta}$  to explicitly show the dependency of the base measure  $d\theta$  — this is the Lebesgue measure on  $\Theta$ , which is the natural unit-volume<sup>2</sup> measure in Euclidean space. Given a metric  $G$ , the natural volume-measurement device on  $\Theta$  is the Riemannian volume form  $dV_G$  which has  $\theta$ -coordinate representation  $dV_G \coloneqq |\det G(\theta)|^{\frac{1}{2}}d\theta$ . This object takes the role of the Lebesgue measure on a Riemannian manifold: Intuitively, it behaves like the Lebesgue measure but takes the (local) distortion due to the metric  $G$  into account. Indeed, when  $G \equiv I$  we recover  $d\theta$ .
Here, we instead argue that invariance is an inherent property of the modes of probability densities, as long as we remain aware of the metric and transform it properly under reparametrization. Explicitly acknowledging the presence of the metric, we obtain
$$
\frac {q _ {\Theta} (\theta) d \theta}{| \det  G (\theta) | ^ {\frac {1}{2}} d \theta} = q _ {\Theta} (\theta) | \det  G (\theta) | ^ {- \frac {1}{2}} =: q _ {\Theta} ^ {G} (\theta). \tag {3}
$$
This is a valid probability density under  $dV_{G}$  on  $\Theta$  since  $\int_{\Theta}q_{\Theta}^{G}(\theta)dV_{G} = 1$ . This formulation generalizes the standard Lebesgue density. And, it becomes clear that the Jeffreys prior is simply the uniform distribution under  $dV_{F}$ ; its density is  $q_{\Theta}^{F}\equiv 1$  under  $dV_{F}$ .
We can now address the invariance question. Under  $\varphi$ , considering the transformation rules for both  $q_{\Theta}(\theta)$  and  $G$ , the density (3) thus becomes
$$
\begin{array}{l} q _ {\Psi} ^ {\boldsymbol {G}} (\boldsymbol {\psi}) = q _ {\Theta} \left(\varphi^ {- 1} (\boldsymbol {\psi})\right) | \det  \boldsymbol {J} ^ {- 1} (\boldsymbol {\psi}) | | \det  (\boldsymbol {J} ^ {- 1} (\boldsymbol {\psi}) ^ {\top} \boldsymbol {G} \left(\varphi^ {- 1} (\boldsymbol {\psi})\right) \boldsymbol {J} ^ {- 1} (\boldsymbol {\psi})) | ^ {- \frac {1}{2}} \\ = q _ {\Theta} \left(\varphi^ {- 1} (\psi)\right) | \det  G \left(\varphi^ {- 1} (\psi)\right) | ^ {- \frac {1}{2}} = q _ {\Theta} ^ {G} \left(\varphi^ {- 1} (\psi)\right). \tag {4} \\ \end{array}
$$
This means,  $q_{\Theta}^{G}$  transforms into  $q_{\Theta}^{G} \circ \varphi^{-1}$  and is thus invariant since it transforms as standard function—notice the lack of the Jacobian-determinant term here, compared to the standard change-of-density rule. In particular, just like standard unconstrained functions, the modes of  $q_{\Theta}^{G}$  are invariant under reparametrization. Since  $G$  is arbitrary, this also holds for  $G \equiv I$ , and thus the modes of Lebesgue-densities are invariant, as long as  $I$  is transformed correctly (Fig. 2b).
Note that, even if the transformation rule is now different from the one in standard probability theory,  $q_{\Psi}^{G}$  is a valid density under the transformed volume form. This is because due to the transformation of the metric  $G \mapsto \hat{G}$ , we have  $dV_{\hat{G}} = |\det G(\varphi^{-1}(\psi))|^{\frac{1}{2}}|\det J^{-1}(\psi)|d\psi$ . Thus, together with (4), we have  $\int_{\Psi}q_{\Psi}^{G}(\psi)dV_{\hat{G}} = 1$ . This also shows that the  $|\det J^{-1}(\psi)|$  term in the standard change-of-density formula (i.e. when  $G \equiv I$ ) is actually part of the transformation of  $d\theta$ .
Put another way, the standard change-of-density formula ignores the metric. The resulting density thus must integrate w.r.t.  $d\psi$  — essentially assuming a change of geometry, not just a simple reparametrization—and thus  $q_{\Psi}$  can have different modes than the original  $q_{\Theta}$ . While this non-invariance and change of geometry are useful, e.g. for normalizing flows [68], they cause issues when invariance is desirable, such as in Bayesian inference.
To obtain invariance in density maximization, transform the density function under the Riemannian volume form as an unconstrained function. In particular when  $G \equiv I$  in  $\Theta$ , this gives the invariant transformation of a Lebesgue density  $q_{\Theta}$ , i.e.  $q_{\Psi} = q_{\Theta} \circ \varphi^{-1}$ .
While reparametrization has been extensively used specifically due to its "non-invariance", e.g. in normalizing flows [63, 68], optimization [19, 69], and Bayesian inference [64, 72], our work is not
In the sense that the parallelepiped spanned by an orthonormal basis has volume one.
at odds with them. Instead, it gives further insights into the inner working of those methods: They are formulated by not following the geometric rules laid out in Section 2.2, and thus in this case, reparametrization implies a change of metric and hence a change of geometry of the parameter space. They can thus be seen as methods for metric learning [35, 77], i.e. finding the "best"  $G$  for the problem at hand, and are compatible with our work since we do not assume a particular metric.
Hessian-based sharpness measures have been extensively used to measure the generalization of neural nets [13, 27, 53]. However, as Dinh et al. [20] pointed out, they are not invariant under reparametrization. Previous work has proposed the Fisher metric to obtain invariant flatness measures [30, 38, 49]. In this work, we have argued that while the Fisher metric is a good choice due to its automatic-invariance property among other statistical benefits [49], any metric is invariant if one follows the proper transformation rules faithfully. That is, the Fisher metric is not necessary if invariance is the only criterion. Similar reasoning about the Fisher metric has been argued in optimization [3, 37, 51, 55, 65, 71] and MAP estimation [21, 32]: the Fisher metric is often used due to its invariance. However, we have discussed that it is not even the unique automatically invariant metric (Appendix C), so the invariance of the Fisher should not be the decisive factor when selecting a metric. By removing invariance as a factor, our work gives practitioners more freedom in selecting a more suitable metric for the problem at hand, beyond the usual Fisher metric.
Finally, the present work is not limited to just the parameter space. For instance, it is desirable for the latent spaces of variational autoencoders [40] to be invariant under reparametrization [2, 25]. The insights of our work can be implemented directly to latent spaces since they are also manifolds [9].
5 Some Applications
We present applications in infinite-width Bayesian NNs, model selection, and optimization, to show some directions where the results presented above, and invariance theory in general, can be useful. They are not exhaustive, but we hope they can be an inspiration and foundation for future research.
5.1 Infinite-width neural networks
Bayesian NNs tend to Gaussian processes (GPs) as their widths go to infinity [57, 60]. It is widely believed that different parametrizations of an NN yield different limiting kernels [75]. For instance, the NTK parametrization (NTP) yields the NTK [29], and the standard parametrization (SP) yields the NNGP kernel [45].
Due to their nomenclatures and the choice of priors, it is tempting to treat the NTP and SP as reparametrization of each other. That is, at each layer, the SP parameter  $\theta$  transforms into
Figure 4: Even though  $\mathcal{N}(\boldsymbol{\theta} \mid \mathbf{0}, \sigma^2 / h\mathbf{I})$  in the SP and  $\mathcal{N}(\psi \mid \mathbf{0}, \sigma^2 / h\mathbf{I})$  in the NTP, they correspond to different priors. Here, we use  $\sigma = 1$ ,  $h = 4$ .
$\psi = \varphi (\pmb {\theta}) = \sigma /\sqrt{h}\pmb{\theta}$  where  $h$  is the previous layer's width. This induces the same prior on both  $\pmb{\theta}$  and  $\psi$ , i.e.  $\mathcal{N}(\mathbf{0},\sigma^2 /h\mathbf{I})$ , and thus one might guess that invariance should be immediate. However, if they are indeed a reparametrization of the other, the NN  $f_{\theta}$  must transform as  $f_{\psi}\coloneqq f_{\theta}\circ \varphi^{-1}(\psi)$  and any downstream quantities, including the NTK, must also be invariant. This is a contradiction since the functional form of the NTP shown in Jacot et al. [29] does not equal  $f_{\psi}$ , and the NTK diverges in the SP but not in the NTP [70]. The SP and NTP are thus not just reparametrizations.
Instead, we argue that the SP and NTP are two completely different choices of the NN's architectures, hyperparameters (e.g. learning rate), and priors—see Fig. 4 for intuition and Appendix E.1 for the details. Seen this way, it is thus not surprising that different "parametrization" yields different limiting behavior. Note that this argument applies to other "parametrizations" [e.g. 15, 70, 75, 76].
Altogether, our work complements previous work and opens up the avenue for constructing non-trivial infinite-width NNs in a "Bayesian" way, in the sense that we argue to achieve a desired limiting behavior by varying the model (i.e. architecture, functional form) and the prior (including over hyperparameters) instead of varying the parametrization. This way, one may leverage Bayesian analysis, e.g. model selection via the marginal likelihood [28, 54], for studying infinite-width NNs.
5.2 The Laplace marginal likelihood
The Laplace log-marginal likelihood (LML) of a model  $\mathcal{M}$  with parameter in  $\mathbb{R}^d$ —useful for Bayesian model comparison [11, Sec. 4.4.1]—is defined by [54]
$$
\log Z (\mathcal {M}; \boldsymbol {\theta} _ {\mathrm {M A P}}) := - \mathcal {L} (\boldsymbol {\theta} _ {\mathrm {M A P}}) + \log p (\boldsymbol {\theta} _ {\mathrm {M A P}}) - \frac {d}{2} \log (2 \pi) + \frac {1}{2} \log \det  \boldsymbol {H} (\boldsymbol {\theta} _ {\mathrm {M A P}}). \tag {5}
$$
Let  $\varphi : \theta_{\mathrm{MAP}} \mapsto \psi_{\mathrm{MAP}}$  be a reparametrization of  $\mathbb{R}^d$ . It is of interest to answer whether  $\log Z$  is invariant under  $\varphi$ , because if it is not, then  $\varphi$  introduces an additional confounder in the model selection and therefore might yield spurious/inconsistent results. For instance, there might exist a reparametrization s.t.  $\log \tilde{Z}(\mathcal{M}_1) \geq \log \tilde{Z}(\mathcal{M}_2)$  even though originally  $\log Z(\mathcal{M}_1) < \log Z(\mathcal{M}_2)$ . The question of "which model is better" thus cannot be answered definitively.
As we have established in Section 3.3, the first and second terms, i.e.  $\mathcal{L}_{\mathrm{MAP}}$ , are invariant under the Riemannian volume measure. At a glance, the remaining terms do not seem to be invariant due to the transformation of the bilinear-Hessian as discussed in Sec. 3.1. Here, we argue that these terms are invariant when one considers the derivation of the Laplace LML, not just the individual terms in (5).
In the Laplace approximation, the last two terms of  $\log Z$  are the normalization constant of the unnormalized density [18]  $h(\theta)\coloneqq \exp \left(-\frac{1}{2} d^{\top}H(\theta_{\mathrm{MAP}})d\right)$ , where  $\pmb {d}\coloneqq (\pmb {\theta} - \pmb{\theta}_{\mathrm{MAP}})$ . Notice that  $\pmb{d}$  is a tangent vector at  $\pmb{\theta}_{\mathrm{MAP}}$  and  $H(\theta_{\mathrm{MAP}})$  is a bilinear form acting on  $\pmb{d}$ . Using the transformation rules in Example 4, it is straightforward to show that  $h\mapsto h\circ \varphi^{-1}$  see
Table 1: The Laplace LML log  $Z$  and its decomposition on a NN under different parametrizations.  $\mathcal{L}$  and "Rest" stand for the first and the remaining terms in (5).
Param.	log Z	- L	Rest
Cartesian	-212.7±3.4	-143.4±0.0	-69.3±3.4
WeightNorm	-227.1±3.6	-143.4±0.0	-83.7±3.6
Appendix E.3. Since  $Z$  is fully determined by  $h$ , this suggests that  $\log Z(\mathcal{M};\pmb{\theta}_{\mathrm{MAP}})$  also transforms into  $\log \tilde{Z} (\mathcal{M};\varphi^{-1}(\psi_{\mathrm{MAP}}))$ , where  $\psi_{\mathrm{MAP}} = \varphi (\pmb{\theta}_{\mathrm{MAP}})$ . This is nothing but the transformation of a standard, unconstrained function on  $\mathbb{R}^d$ , thus  $\log Z$  is invariant.
We show this numerically in Table 1. We train a network in the Cartesian parametrization and obtain its  $\log Z$ . Then we reparametrize the net with WeightNorm and naively compute  $\log Z$  again. These  $\log Z$ 's are different because the WeightNorm introduces more parameters than the Cartesian one, even though the degrees of freedom are the same. Moreover, the Hessian-determinant is not invariant under autodiff. However, when transformed as argued above,  $\log Z$  is trivially invariant.
5.3 Biases of preconditioned optimizers
The loss landscape depends on the metric assumed in the parameter space through the Hessian operator  $\pmb{E} = \pmb{G}^{-1}\pmb{H}$ . The assumption on  $\pmb{G}$  in practice depends on the choice of the optimizer. For instance, under the default assumption of the Cartesian coordinates, using GD implies the Euclidean metric  $\pmb{G} \equiv \pmb{I}$ , and ADAM uses the gradient-2nd-moment metric [39].
We argue that explicitly including the metric is not just theoretically the correct thing to do (since it induces invariance, and by Riemannian geometry), but also practically beneficial. Fig. 5 compares measurements of sharpness of minima (Hessian-trace).
Using the metric-aware Hessian operator  $\pmb{E}$ , one can show definitively (i.e. independent of the choice of parametrization) that ADAM tends to obtain much sharper minima than SGD. The benefits of using the Hessian operator have also been confirmed in previous work. Cohen et al. [16] argue that when analyzing the optimization dynamic of an adaptive/preconditioned GD algorithm, one should take the preconditioner into account when measuring Hessian-based sharpness measures. From this, they demonstrate a sharpness-evolution behavior known in vanilla GD, allowing a comparison between vanilla and adaptive GD.
Figure 5: The effect of viewing the Hessian as a linear map  $\pmb{E}$  to measure sharpness at minima. ADAM finds much sharper minima when the geometry of the parameter space is taken into account.
In this work, we addressed the invariance and invariance associated with the reparametrization of neural nets. We started with the observation that the parameter space is a Riemannian manifold, albeit often a trivial one. This raises the question of why one should observe non-invariance in neural nets, whereas, by definition, Riemannian manifolds are invariant under a change of coordinates. As we showed, this discrepancy only arises if the transformations used in the construction of a neural net along with an optimizer ignore the implicitly assumed metric. By acknowledging the metric and using the transformation rules associated with geometric objects, invariance and invariance under reparametrization are then guaranteed. Our results provide a geometric solution towards full invariance of neural nets—it is compatible with and complementary to other works that focus on invariance under group-action symmetries, both in the weight and input spaces of neural networks.
The authors gratefully acknowledge financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence "Machine Learning - New Perspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) in the frame of the priority programme SPP 2298 "Theoretical Foundations of Deep Learning"—FKZ HE 7114/5-1; and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. AK & FD are grateful to the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. AK, FD, and PH are also grateful to Frederik Künstner and Runa Eschenhagen for fruitful discussions.