The Geometry of Neural Nets' Parameter Spaces Under Reparametrization
Authors: Agustinus Kristiadi
Felix Dangel
Philipp Hennig
Venue: NeurIPS 2023 spotlight (2023)
Keywords: neural network
invariance
equivariance
reparametrization
riemannian geometry
parameter space

Abstract: Model reparametrization, which follows the change-of-variable rule of calculus, is a popular way to improve the training of neural nets. But it can also be problematic since it can induce inconsistencies in, e.g., Hessian-based flatness measures, optimization trajectories, and modes of probability densities. This complicates downstream analyses: e.g. one cannot definitively relate flatness with generalization since arbitrary reparametrization changes their relationship. In this work, we study the invariance of neural nets under reparametrization from the perspective of Riemannian geometry. From this point of view, invariance is an inherent property of any neural net if one explicitly represents the metric and uses the correct associated transformation rules. This is important since although the metric is always present, it is often implicitly assumed as identity, and thus dropped from the notation, then lost under reparametrization. We discuss implications for measuring the flatness of minima, optimization, and for probability-density maximization. Finally, we explore some interesting directions where invariance is useful.

The Geometry of Neural Nets' Parameter Spaces Under Reparametrization
Agustinus Kristiadi Vector Institute, University of Tübingen akristiadi,
Felix Dangel
Philipp Hennig
University of Tübingen, Tübingen Al Center philipp hewnigeuni-tuebbee.. -tubbingen.dd
far.ssttadd..faaaeeeeeeettrrrrrt Feaageelleecttrrsstiiuuu ai
Abstract
Model reparametrization, which follows the change-of-variable rule of calculus, is popular way to improve the training of neural nets. But it can also be problem- atic since it can induce inconsistencies in, e.g.. Hessian-based flatness measures, optimization trajectories, and modes of probability densities. This complicates downstream analyses: e.g. one cannot definitively relate flatness with generaliza- tion since arbitrary reparametrization changes their relationship. In this work, we study the invariance neural nets under reparametrization from the perspective Riemannian geometry. From this point of view, invariance any neural net if one explicitly represents the metric and uses the correct associated transformation rules. This is important since although the metric is always present, it is often implicitly assumed as identity, and thus dropped from the notation, then lost under reparametrization. We discuss implications for measuring the flatness minima, optimization, and for biopabilitt-denniii maximization. Finally, we explore some interesting directions where invariance is useful.
inherent property of
Introduction
Neural networks (NNs) are parametrized functions. Since it is often desirable to assign meaning interpretability to the parameters (weights) of a network, interesting to ask whether certain transformations of the parameters leave the network invariant invariant-ccquivalent in some sense. Various notions of invariance have been studied in NNs, in particular under weight-space symmetry [5, 14, 20, 42, 42, and reparametrization [20, 21, 30. 32, 71]. The former studies the behavior of a function L(0) under some invert- ible T θ G - ④ where g is group; for any O, the function L is invariant under the symmetry T if and only if (()) C(T(O, .g)) for all g G. For example, normal- ized NNs are symmetric under scaling of the weights, i.e. [7m-sim scale-symmetry also presents in ReLU NNs [61]. Mcanwhile, invariance under reparametrization studies the behavior of the NN when it is transformed under the of-igeab--aaiibll... rule of calculus: Given a transformed parameter &b- = ((0) under bijective differentiable y that maps the original parameter space onto another parameter space, the function L(θ) becomes [(Φ) = C(p (+)). Note their difference (see Fig. 1): In the former, one works on single parameter space θ and single function L-the map 7 acts as symmetry of elements of θ under the group G. contrast, the latter assumes two parameter spaces and which are connected by y and hence two functions L and L.
(())
C(T(0,9) .1))
((
((w)
Figure 1: The difference between the symmetry (top) and reparametrization problems (bottom). θ and v are two different parameter spaces.
(()) L(cl) for all
v
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a) Optimiization
(b) Maximum Density
(c) Flatness Mcasure
o(o)
pio)
Default
1p (1p(
((
(14
Metrie- Metric-Award Aware
5p (1p")
C(0)
((P)
H()") H(()) det B(b)) = des (OO")
Figure 2: Invariance is retained (bottom) if the Riemannian metric is explicitly tracked and one transforms geometric objects such as vectors, covectors, and the metric itself properly under a reparametrization o 0 + v. (a) Gradient descent's trajectories are therefore invariant under change of parametrization. (b) Invariance of the modes of probability density function is an inherent property under natural base measure induced by the metric.
(c) When the Hessian at a critical point is seen as linear map E with the help of the metric, its determinant is invariant.
We focus on the latter. While this scenario, we have ((() = C()) whenever &b = (0)). it well- known that downstream quantities such optimization trajectories [55, 71], the Hessian [20, Sec. 5], and probability densities over 0 [59, Sec. 5.2.1.4) are not invariant under model reparamctrization. These non-in non-invariances are detrimental since an arbitrary reparametrization can affect the studied quantities and thus consistency cannot be guaranteed paranerrraaaaaatiioioon muddles the analysis.
For instance, because of this, (i) one cannot relate Hessian- based sharpness measures with gener- alization and the correlation analyses between them [27, 33, 73] become meaningless, (ii) a good preconditioner cannot be guaranteed to be good anymore for optimizing the reparametrized model- this is one reason why invariant methods like natural gradient are desirable [37, 50, 58, 62], and (iii) under the reparametrization, prior density might put low probability in region that corresponds to region in the original space, making posterior inference pathological [25].
To analyze this issue, in this work we adopt the framework of Riemannian geometry, which is generalization of calculus studying the intrinsic properties of manifolds. "Intrinsic" in this context means that objects such as functions, vectors, and tensors defined on the manifold must be independent of how the manifold coordinates [47]. The parameter space of neural network, which is by default assumed to be
represented via a coordinate system, and conserved under a change of
= Rd with the Euclidean metric and the Cartesian coordinates, a Riemannian manifold-mode reparametrization is such change in the coordinate system.
Why, then, can reparametrizations bring up the aforementioned inconsistencies? In this work, discuss this discrepancy. We observe that this issue often arises because the Riemannian metric is left implicit, and dropped when computing downstream quantities such as gradients, Hessians, and volumes on the parameter space. This directly suggests the solution: Invariance under reparametriza- tion is guaranteed if the Riemannian metric is not just implicitly assumed, but made explicit, and is made sure that the associated transformation rules of objects such as vectors, covectors, and tensors are performed throughout. We show how these insights apply common cases (Fig. 2).
Limitation Our work focuses on the reparametrization consistency of prior and future methods that leverage geometric objects such as gradients and the Hessian. Thus, we leave the geometric analysis for invariance under symmetry as future work. In any case, our work complimentary to other works that analyze invariance under symmetry [17, 22, 43, 67, etc]. Moreover, this work's focus is not to propose "better" methods e.g. for better preconditioners or better generalization metrics. Rather, we provide guardrail for existing and future methods to avoid pathologies due reparametrization.
2 Preliminaries
This section provides background and notation relevant concepts of neural networks and Rieman- nian geometry. For the latter, we frame the discussion in terms of linear algebra as close as possible to the notation of the ML community: We use regular faces to denote abstract objects and bold faces for

---

their concrete representations in particular parametrization. E.g., a linear map A is represented by matrix A, point is represented by a tuple of numbers 0, an inner product G(z) is represented by matrix G( The same goes for the differential V. V; and the gradient grad, grad.
2.1 Some notation for neural networks
The following concepts standard, introduced mostly to clarify notation. Let f R Rd R with (z.0) 0) - f(x: 0) be a model with input, output, and parameter spaces R", R R Let J and Y == ) We moreover write D (( y: Th be dataset and write X == {a:}
Rmk The standard way of training f is by finding point s.t. 0" arg mi >x eyy, fæz 0)) arg millgeRd [(0), for some loss function L. If we add a weight decay 2/21 1 term to [()), the minimization problem has probabilistic interpretation as maximum a posteriori (MAP) estimation of the posterior density p(O D) under the likelihood function p(Y [Xx; 0)) & exp(- - -==1 / 0))) and an isotropic Gaussian prior (0) ( 0, 11 We denote the MAP loss by LMAP- our present context, this interpretation is relevant because this probabilistic interpretation implies probability density, and it is widely known that a probability density transforms nontrivially under reparametrization.
f(X:0)
vec([f(x }m
1==1 E.
O" E R
A textbook way of obtaining 0 is gradient descent (GD): At each time step t, obtain the next estimate ((++1) 0( a Clore, for some > O. This can be considered the discretization of the gradient flow ODE o = -V Cle. Among the many variants of GD is preconditioned GD which considers a positive-definite matrix field R on Rd, yielding = - - VLoo-
2.2 The geometry of the parameter space
Remark 1. We restrict our discussion to global coordinate charts since they are the default assumption in neural networks, and in a bid to make our discussion clear for people outside of the differential geometry community. We refer the curious reader to Appendix D for a discussion on local coordinate charts. Note that, our results hold the general local-chart formulation.
The parameter space Rd is the canonical ex- ample of a smooth manifold. We can impose upon this manifold a (global) coordinate chart: a homeomorphism that represents point with an element of Rd. The standard choice is the Cartesian coordinate system 0 : RRd which we have used in the previous section. That is, the image of 9 uniquely represents elements of IRd -given a point E Rd. we write O = #(2) for its coordinate representation under
14
T R
(u,w) 0 Go
W
= Rdd
+ = (+))
0(=) @ Tr gd
The choice of a coordinate system is not unique. Any homeomorphism &b Rd + V = Fd can be used as an alternative coordinate chart. For instance, the polar coordinates can be used to represent points in IR instead. Crucially, the images of any pair of coordinates must con- nected through diffeomorphism (a smooth function with smooth inverse) 40 θ - v. Such map is called a change of coordinates or reparametrization since it acts on the parameter space.
M
Figure 3: The implicit geometric assumption on the parameter space R... 6, z are two different (global) coordinates on Rdd.
Example 2 (Reparametrizationn. Reparametriization is ubiquitous in machine learning:
(a) Mean-Field Variational Inference. Let q( 0) be a variational approximation, which is often chosen to be N(x μ,, diag(o²)) i.e. 0 = {u E Rd o² E RR Common choices reparametrization o² include the log-space [52] or softplus [12] parametrization.
(b) Weight normalization. Given a NN fe. WeightNorm [69] applies the reparametrization 2/ = rm where r E R>0 and 17 := 0/10|l E 3d-1, This is akin to the polar coordinates. (Note that we assume E Rd {0) since otherwise & is not global diffeomorphism.)
At each point = RRd. there exists vector space T.Rd called the tangent space at z, consisting of the so-called tangent vectors. An important example of tangent vector is the gradient vector grad C)2
3

---

of C: of linear functionals of T- R., called tangent covectors. An example of a tangent covector is the differential VCl2 of C at z. Under a coordinate system, one can think of both tangent vectors and covectors as vectors in the sense of linear algebra, i.e., tuples of numbers.
IR :.L:R
R z. z. The dual space TŞRº of T,R" is referred to as the cotangent space and consists
One can take inner product of tangent vectors by equipping the manifold with a Riemannian metric G, which, at each point in IR is represented by positive-definit d x d matrix G whose coefficients depend on the choice of coordinates. With the help of the metric, there is isomorphism between T,R² and T.Rd. In coordinates, is given by the map T. - T.R" v + Gv and its inverse T - T,Re w - G-tw. One important instance of this isomorphism is the fact that grad C|2 is represented by G C in coordinates. The natural gradient is direct consequence of this fact. In practice, it is standard to assume the Cartesian coordinates on R. implying G = I. We have thus grad C V L. which only seems trivial first sight, but reveals the relationship between the tangent vector on the l.h.s. and the cotangent vector on the r.h.s.
Remark 3. Another common way to define a NN's parameter space is to define a manifold of probabilistic models M == {p(Y f(x; 0)) 0 E [Rd] and assume that p(Y 0)) + 0 is the coordnintee chart [19. 55, etc.]. The problem with this construction is that there is no bijection F - MM in general. Indeed, the Jacobian of the map f(X; E M is in practice surjective everywhere due to overparametrization [34, 78]. Therefore. one cannot define proper metric on the parameter space IRd that corresponds to a metric on the distribution space M [46, Prop. 13.9]. For instance, using this construction, the Fisher metric is singular for overparametrized NNs [10, 34, 74, 78] and thus not valid Riemannian metric. The pseudoinverse has been used to handle this but is mostly limited to theoretical analyses [10, 74]: As far as we know, in practice, damping- which breaks the interpretation of the Fisher on Rd as the pullback metric from M de facto for handling this due to its numerical stability [4].
By detaching from the distribution space, the definition used in this work does not have this issue. It enables more general constructions of metrics in the parameter space since any positive-definite matrix is admissible. E.g. one is free add damping or use any approximation to the Fisher- our results still apply to this case. Thus, our construction is closer to practice.
2.2.1 Transformation rules
Differential geometry is the study of coordinate-independee objects: Geometric objects must be invariant under change-of-coordinates in the sense that any pair of representations must refer to the same (abstract) object. Suppose o : Q - Vp reparametrization, with &b = (0)). Coordinate independence is encoded in the following transformation rules, which involve the Jacobian (() (drb/00) /00,) of ip, and its inverse J-1 J (vb) = (J())) -1 = (00,/ (00,/ (oob,,)tth the Jacobian of a (Color codes are used for clarity when different objects are present in a single expression later on.)
(a) A function h θ + R in 0-coordinates transforms into h h 9-1 in vi-coordinates.
(b) A tangent vector e in f-coordinates transforms into J(0)) v. In particular, gradient vector grad C/o transforms into (@) grad Clo.
(c) A tangent covector w in f-coordinates transforms into (1b) T w. In particular, given a transformed function L 9-11, v + I. we have ( o p
I) -- (1b) V ((pp-))))
which we recognize as the standard chain rule.
(d) A metric G()) becomes (())
(nb). In general, this rule applies to any tensor that takes tangent vectors as its arguments, e.g. a bilinear map.
-00_((()11=-((()))) ((
The following examples show how these rules ensure the invariance of geometric objects. Example 4. Let h, v, w, G respectively be a function, vector, covector, and metric in d-coordinates at point 0= = 0(2), and let yp: ④ + & be reparametrization.
(a) Let 1/ = (0) and h == h yp be O E ② and h expressed in wi-coordinates, respectively. Then, h {)) = h (Φ)) = hO) That is, the actions of h and h agree in both coordinates and thus they represent the same abstract function on TRd
(b) The action of w on in f-coordinates is given by the product w v. Let ú (()v and ت = J(0)-Ta w be the transformed vector and covectors. Then,
w o (1(0) Ww) () w
(I)
4

---

That is, both w and w are the representations of the same linear functional: v and e are the representations of the same tangent vector.
(c) Let G JJ GJJ
be the transformed metric in the vi-coordinates. Then,
Go Go (Jv) GJJJJO) v Gv,
and thus the transformation rules ensure that inner products are also invariant.
Finally, we say that an ODE's dynamics invariant if the trajectory in parametrization corresponds to the trajectory in another parametrization. Concretely, a trajectory (0t); in θ is invariant if under the reparametrization yp: θ - V, the transformed trajectory (1/e), is related (0t); by 0. = 4 (1bb) for each t. See Fig. 2a for an illustration.
3 Neural Networks and Reparametrization
We discuss three aspects of the parameter space under reparametrization, as illustrated in Fig. First, we address the non-invariance of Hessian-based flatness measures [e.g.. 20], and show how taking the metric into account provides invariance. Second, we show that the invariance property often cited in favor of natural gradient is not unique, but an inherent property of any gradient descent algorithm when considering its ODE. Finally, we show that modes of probability density functions on the parameter space are invariant when the Lebesgue measure is generalized using the metric.
3.1 Invariance of the Hessian
growing body of literature connects the flatness of minima found by optimizers to generalization performance [8. 13, 23, 27, 36, 38, 48, 53J. However, as Dinh et al. [20] observed, this association does not have solid foundation since standard sharpness measures derived from the Hessian of the loss function not invariant under reparametrization.
From the Riemannian geomeetri perspective, the Hessian of function C (or CAP or any other wice-differentiable function) on fl-corddinates is represented by a d d matrix with coefficients ()) == 80,00, (0) Σx= T (0) 0C (0), for any 0 θ and i, 1,... d, where 1 is 30, particular three- frre--dimeniioall array. While it might seem daunting, when 0 is a local minimum of C. the second term is zero since the partial derivatives of are all zero at Thus, (()) equals the standard Hessian matrix (/00/000)) at O arg min L.
8²
Considering the Hessian as bilinear function that takes vectors and produces number, follows the covariant transformation rule, just like the metric (see Appendix B. for derivation): Let yp θ - and +/ (0)) The Hessian matrix H (0) at a local minimum f-coordinates transforms into H (b)) = J (xb) (*b)) (ab) in gi-coordinates-tie this is correctly computed by automatic differentiation, i.e. by chain and product rules. But, while this gives invariance of H as a bilinear map (Example 4c), the determinant of H-a popular flatness measure-is not invariant because
(det H (xb) = det (J- (xb)" H [ (1)) = (det J ((de) H (x)-)))))))
and so in general, we do not have the relation (det H) (det H) sp- that would make this function invariant under reparametrization (Example 4a).
The key to obtaining invariance is to employ the metric G to transform the bilinear Hessian into linear map/operator on the langent space. This is done by simply multiplying H with the inverse of the metric G, i.e. E = G H. The determinant of E is thus an invariant quantity." To show this, under ye, the linear map E transforms into
E(b) G HL
(2)
= (J- G( (+)) (())J-((()))
due to the transformation of both G and H. Hence, det E transforms into
(det (+b) (det (o) (Gl (C)
(det ()
'This is because one can view the loss landscape a hypersurface: H is connected to the second fundamental form and E to the shape operator. det E is thus related to the invariant Gaussian curvature [47, Ch. 8).
5

---

in cb-coordinates. Thus, we have the desired invariant transformation (det E) = (det o Note that G is an arbitrary metric- this invariance thus also holds for the Euclidean case where Note further that the trace and eigenvalues of E are also invariant; see Appendices B.2 and B.3. Finally, see Appendix for a simple working example.
To obtain invariance in the Hessian-determinan -trace, and -eigenvalues at minimum of C, we explicitly write it as E G- "H, even when G = 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 0 - - Llo. But NGD is widely presented as invariant under reparametrization, while GD not [51, 71, etc. Is the choice the metric G the cause? Here will show from the framework laid 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 is "automatically" invariant under standard deep learning libraries like PyTorch, TensorFlow, and JAX [1, 24, 66] see Appendix C.
Section 2.2 that any metric is invariant its
The underlying assumption of GD that one works in the Cartesian coordinates and that G = For this reason, one ignore the metric G in 0 (G))-lly Lle. and simply write O -V This is correct but this simplification is exactly what makes GD not invariant under a reparametrization 9 θ + T where v/ (0)). . To see this, notice that while GD transforms V correctly via the chain rule v/ - (lp-)( by ignoring the metric I, one would miss the important fact that it must also be transformed into G(ab) = JJ1 (xb) ( It is clear that ((tb) does not equal in general. Thus, we cannot ignore this term in the transformed dynamics since would imply that one uses a different metric tte dynamics are thus different in the 6- and ci-coordinates.
When one explicitly considers the above transformation, one obtains
xb -
V ((0)
V -+())
=- ([)-)
This dynamics co-oorddssates preserves the assumption that the metric is I in f-coordinates and thus invariant. In contrast, the dynamics vp under just the chain rule implicitly changes the metric in 0-coordinates, i.e. from into (0) (0), and thus the trajectories are not invariant.
This discussion can be extended to any metric R: Simply use the transformed metric R(bb) ( T R( (1b) and we obtain the invariant dynamics of any preconditioned GD with preconditioner R given by xb- [ (xb) R(p-~(()))) V ((()--
To obtain invariance in optimizers with any metric/preconnitionee, 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, 71J. This is orthogonal work since it is about improving ODE solvers [71].
The consequence is that we need a "geometric-aware" autodiff library such that the invariant dy- namics above is always satisfied. In this case, any preconditioner R yields invariance under any reparametrization, even nonlincar ones. This is in contrast to the current literature, e.g. under standard autodiff, Newton's 's method is only affine-invariane invariant, and structured approximate NGD methods such as K-FAC are only invariant under a smaller class of reparametrivations [56].
3.3 Invariance probability densities
Let (0) be a probability density function (pdf) under the Lebesgue measure d0 on O. Under a reparametrization &
- with &b (0)). transforms into g (b)) = (xb)) J- This transformation rule ensures go to be valid pdf under the I Lbeegue measure dap on V. i.e. go (xb) dub 1. Notice, general gs * de o p due to the change- of-radoom-variabl rule. Hence, pdfs transform differently than standard functions (Example 4a)

---

and can thus have non non-invariant modes [cf. 59, Sec. 5.2.1.4]: Density maximization, such as the optimization of 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 F [21]. When applied to a prior, this gives rise to the famous Jeffreys [31] prior pO( & det F(O))² with normalization constant Sel det F(O) 2 d0. Is the Fisher actually necessary 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 qe (6) can be written as qe (0) de (0) de
do to explicitly show the dependency of the base measure d0 this is the Lebesgue measure on e, which is the natural unit- volt-oolmme measure in Euclidean space. Given metric G, the natural volume- measurement device on θ is the Riemannian volume form dVG which has 0-coordinate representation dVG := det G d0. This object takes the role of the Lehesgue measure on Riemannian manifold: Intuitively, it behaves like the Lebesgue measure but takes the (local) distortion due to the metric G into account. Indeed, when G we recover d0.
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
ge v0)))d de
(e det =
det G de
This is valid probability density under dVG on since So 98 ()) dVG 1. This formulation generalizes the standard Lebesgue density. And, it becomes clear that the Jeffreys prior is simply the uniform distribution under F; its density is o under dVF.
We can now address the invariance question. Under yp considering the transformation rules for both qe (e(O) and G, the density (3) thus becomes
d (vb) go (xb)) det J (e))l (v)) 110"0//060000000000000"000""1010
(4)
(p det (())) = - (g
This means, C transforms into g o - and is thus invariant since it transforms as standard term here, compared to the standard change-of-
function-notice the lack of the density rule. In particular, just like standard unconstrained functions, the modes of g are invariant under reparametrization. Since G is arbitrary, this also holds for G I, and thus the modes of debesgue-densitee are invariant, as long as I is transformed correctly (Fig. 2b).
Note that, even if the transformation rule is different from the one in standard probability theory, G ty is a valid density under the transformed volume form. This is because due to the transformation of the metric G - G, have dVG det (<b)1 det (xb) dep. Thus, together with (4), we have fy dh (xb) dVG = 1. This also shows that the det J [xb) term in the standard change-of-density formula (i.e. when G I) is actually part of the transformation of dO.
G
Put another way, the standard change-of-density formula ignores the metric. The resulting den- sity thus must integrate w.r.t. ees-nnsitt assuming change of geometry, not just a simple reparametrization- and thus qu can have different modes than the original qe- While this non- invariance and change of geometry are useful, e.g. for normalizing flows [68], they cause issues when invariance desirable, such as in Bayesian inference.
To obtain invariance density maximization, transform the density function under the Riemannian volume form as an unconstrained function. In particular when G = I in O this gives the invariant transformation of Lebesgue density ge, i.e. Qv = qe o -1
4
Related Work
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.
7

---

at odds with them. Instead, it gives further insights into the inner working of those methods: They are formulated by nor following the geometric rules laid out in Section 2.2, and thus in this case. reparametrization implies 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 particular metric.
Hessian-baid based sharpness measures have been extensively used 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 invariant one follows the proper transformation rules faithfully. That is, the Fisher metric is not necessary 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 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 more suitable metric for the problem at hand, beyond the usual Fisher metric.
Finally, the present work not limited to just the parameter space. For instance, 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 latent spaces since they are also manifolds [9].
5 Some Applications
We present applications in infinite-width Bayesian NNs, model selection, and optimization, 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). 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].
2
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 transforms into sb (0)) = a/ ho where h is the previous layer's width. This induces the same prior both 0 and s i.e. (/ o2/h. //II. and thus one might guess that invariance should be immediate. However, they are indeed a reparametrization of the other, the NN fo must transform f == fo o ((b) 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 al. [29] does equal fap and the NTK diverges in the SP but not in the NTP [70]. The SP and NTP are thus not just reparametrizations.
Figure 4: Even though A/O 0, /) in the SP and Wntb 0, a² /2/TT] in the NTP, they correspond to different priors. Here, we use o = 1. = 4.
Instead, we argue that the SP and NTP are two completely different choices of the NN's architectures, hyperparameters (c.g. learning rate), and priors-see see Fig. 4 for intuition and Appendix E. I 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- non-trivial infinite-width NNs "Bayesian" way, 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-widtt NNs.
8

---

5.2 The Laplace marginal likelihood
The Laplace marginal likelihood (LML) of model .M with parameter in R useful for Bayesian model comparison [11, Sec. 4.4.1] is defined by [54]
log Z (M; MMAP) == - [ ((MMA)) + log p(0MA)) - 5 l2g( + log det H (0MAP).
(5)
Let +p MAP + MMAP be a reparametrization of IR. It is of interest to answer whether log Z is invariant under ip, because if it is not. then yp introduces an additional confounder in the model selection and therefore might yield spurious/inconsistent results. For instance, there might exists a reparametrization s.t. log Z (M1) 2 log (M22) even though originally log (M1) < log ((M2). The question of "which model is better" thus cannot be answered definitively.
As have established in Section 3.3, the first and second terms, i.e. LAPP invariant under the Riemannian volume measure. At 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 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(0) == exp {-2 H (fMA))) dd), where = (0 MAP).. Notice that d is tan- gent vector at OMAP and H (0MAP) is a bi- linear form acting on d. Using the trans- formation rules in Example 4, it is straight- forward to show that h + h o y -1-s see Appendix E.3. Since Z is fully determined by h, this suggests that log Z(M; MMPP) also transforms into log ((M: 9 ((MMA))) where dỘMAP = (o ((0A)). This nothing but the transformation of a standard, unconstrained function on R4, thus log Z is invariant.
Table I: The Laplace LML log Z and its decomposition NN under different parametrizations. L and "Rest" stand for the first and the remaining terms in (5).
Param.
log Z
C
Rest
Cartesian WeightNorm
211.77334
-143.4 -1.3.4
-6,.3+3.4
-227.1=3.6
-8..743.6
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. Morcover, the dessiann-deerminnat under autodiff. However, when transformed as argued above, log Z is trivially invariant.
not invariant
5.3 Biases of preconditioned optimizers
The loss landscape depends on the metric assumed in the parameter space through the Hessian operator
G-1 H. The assumption on G in practice depends on the choice of the optimizer. For instance, under the default assumption of the Carte- sian coordinates, using GD implies the Eu- clidean metric G = I, and ADAM uses the gradient-2nm mommomentt metric [39].
Log-Truc Taice Hessian
IH ADAM H
SGD
10
ADAM
MNIST' FMNIST CIFARIO CIFARIOD
Dataset
We argue that explicitly including the met- ric 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 (ressian-tracee., Using the metric-aware Hessian operator E, one can show definitively (i.e. independent of the choice of parametrization) that ADAM tends 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/rreoondtioned GD algorithm, one should take the preconditioner into account when measuring flessian-based sharpness measures. From this, they demonstrate charpness-evolution behavior known in vanilla GD, allowing a comparison between vanilla and adaptive GD.
Figure 5: The effect of viewing the Hessian as a lin- ear map E to measure sharpness at minima. ADAM finds much sharper minima when the geometry of the parameter space is taken into account.

---

6 Conclusion
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 Riemannian manifold, albeit often a trivial one. This raises the question of why one should observe non-invariance neural nets, whereas, by definition, Riemannian manifolds are invariant under change of coordinates. As we showed, this discrepancy only arises 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 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.
then guaranteed. Our results provide geometric solution towards full
Acknowledgments and Disclosure of Funding
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: 01 80099);; the Deutsche Forschungsgemminnchaaff (DFG, German Research Foundation) in the frame of the priority programme SPP 2298 "Theoretical Foundations of Deep Learning" Searmingg-FZZ -KKZ HE 71 14/55-1; and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemin 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 the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. AK, FD, and PH are also grateful Frederik Künstner and Runa Eschenhagen for fruitful discussions.
Resources used
References
[1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemaaat, Geoffrey Irving, Michael Isard, et al. TensorFlow: system for large-scale machine learning. In OSDI, 2016.
[2] Francisco E Acosta, Sophia Sanborn, Khanh Dao Duc, Manu Madhav, and Nina Miolane, Quantifying local extrinsic curvature in neural manifolds. arXiv preprint arXiv:2212.10411 2022.
[3] Shun- Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-277 1998.
[4] Rohan Anil, Vineet Gupta. Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv preprint arXiv:2002.0901. 2020.
[5] Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Hernández-Lobate Adapting the linearised Laplace model evidence for modern deep learning. In ICML, 2022.
[6] Sanjeev Arora, Simon S Du. Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. NeurlPS, 2019.
[7] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto ratc-tuning by batch normaliza- tion. In ICLR, 2019.
[8] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent edge of stability in deep learning. arXiv preprint 2205::2205009455 2022.
Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: on the curvature of deep generative models. In ICLR, 2018.
[10] Alberto Bernacchia, Máté Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In NIPS, 2018
[11] Christopher M. Bishop. Pattern Recognition and Machine Learning Springer, 2006.
10

---

[12] Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In ICML, 2015.
[13] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In ICLR, 2017.
[14] An Mei Chen, Haw-minn Lu, and Robert Hecht-Niclsen. On the geometry of feedforward neural network error surfaces. Neural computation, 5(6), 1993.
[15] Lenaic Chizat and Francis Bach. On global convergence of gradient descent for over- wer-paramettrized models using optimal transport. In NewrIPS, 2018.
[16] Jeremy Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint vrxiv:2207.14484 2022.
[17] Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral CNN. In ICML, 2019.
[18] Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless Bayesian deep learning. In NeurlPS, 2021.
[19] Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, and Koray Kavukcuoglu. Natural neural networks. In NIPS, 2015.
[20] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. ICMI, 2017.
[21] Pierre Druilhet and Jean Mech-Michel Marin. Invariant HPD credible sets and MAP estimators. Bayesion Analysis, 2(4), 2007.
[22] Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. NeurlPS, 2018
[23] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. ICLR, 2021.
[24] Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine lcarning programs via high-level tracing. SysML, 4(9), 2018.
[25] Søren Hauberg. Only Bayes should learn manifold (on estimation of differential geometric structure from data). arXiv preprint 18KK::8800009994 2018.
[26] Dan Hendrycks, Mantas Mazeika, and Thomas Dietterich. Deep anomaly detection with outlier exposure. In ICLR, 2019.
[27] Sepp Hochreiter and Jürgen Schmidhuher. Simplifying neural nets by discovering flat minima. In NIPS, 1994.
[28] Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection in deep learning. In ICML, 2021.
[29] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In NIPS, 2018
[30] Cheongjae Jang, Sungyoon Lee, Frank C. Park, and Yung-Kyun Noh. A reparametiization-nvaaiat sharpness measure based on information geometry. In NeurlPS, 2022.
[31] Harold Jeffreys. An invariant form for the prior probability in estimation problems. Proceedings the Royal Society of London. Series A. Mathematical and Physical Sciences, 186(1007), 1946.
[32] lan H Jermyn. Invariant Bayesian estimation on manifolds. The Annals of Statistics, 33(2). 2005.
[33] Yiding Jiang, Behnam Neyshabur, Hosscin Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic general- ization measures and where to find them. In ICLR, 2020.
[34] Ryo Karakida and Kazuki Osawa. Understanding approximate Fisher information for fast convergence of natural gradient descent in wide neural networks. In NeurlPS, 2020.
11

---

[35] Mahmut Kaya and Hasan Şakir Bilge. Deep metric learning: A survey. Symmetry, 11(9), 2019.
[36] Nitish Shirish Keskar, Dheevatsa Mudigere, Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.
On large-batch training for deep learning: Generalization gap and sharp minima. In ICLR, 2017.
[37] Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In JCML, 2022.
[38] SungYun Yuh Kim, Sihwan Park, Kyung-Su Kim, and Eunho Yang. Scalc-invariant Bayesian neural networks with connectivity tangent kemee.. In ICIR, 2023.
[39] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[40] Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In JCLR, 2014.
[41] Agustinus Kristiadi, Matthias Hein, Philipp Hennig. An infinite-feature feature extension for Bayesian ReLU nets that fixes their asymptotic overconfidence. In NewrlPS, 2021.
[42] Vêra Kürková and Paul Kainen. Functionally equivalent feedforward neural networks. Neurul Computa- tion, 6(3), 1994.
[43] Daniel Kunin, Javier Sagastuy-Bren Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In ICLR, 2021.
[44] Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations the empirical Fisher approximation for natural gradient descent. In NewrlPS, 2019.
[45] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In ICLR, 2018.
[46] John Lee. Introduction to smooth manifolds. Springer, 2013.
[47] John M Lee. Introduction to Riemannian manifolds. Springer, 2018.
[48] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? mathematical framework. In NeurlPS, 2021.
[49] Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-Rao metric, geometry, and complexity of neural networks. In AISTATS, 2019.
[50] Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Structured second-order methods natural gradient descent. ICML Workshop on beyond first-order methods machine learning systems. 2021.
[51] Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Tractable structured natural gradient descent using local parameterizations. In ICML, 2021.
[52] Christos Louizos and Max Welling. Structured and efficient variational deep learning with matrix Gaussian posteriors, In ICML, 2016.
[53] Kaifeng Lyu. Zhiyuan Li, Sanjeev Arora. Understanding the generalization benefit normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.00 07085, 2022.
[54] David JC MacKay. The evidence framework applied to classification networks. Neural Computation, 4(5), 1992.
[55] James Martens. New insights and perspectives on the natural gradient method. JMLR, 21(146), 2020.
[56] James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature, International conference on machine learning, pages 2408-2417, 2015.
[57] Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In ICLR, 2018
[58] Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning. ICLR, 2021.
[59] Kevin Murphy, Machine Learning: Probabilistic Perspective. The MIT Press, 2012.
[60] Radford M Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, 1995.
12

---

[61] Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-SGD: Path-normalized optimization in deep neural networks. In NIPS, 2015.
[62] Kazuki Osawa, Shigang Li, and Torsten Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv preprint arXiv:2211.J1133 2022.
[63] George Papamakarios, Eric Nalisnick, Danilo Jimenez. Rezende, Shakir Mohamed, and Balaji Lakshmi- narayanan. Normalizing flows for probabilistic modeling and inference. JMLR, 22(57), 2021.
[64] Omiros Papaspiliopoulos, Gareth O Roberts, and Martin Sköld. A general framework for the parametriza- tion of hierarchical models. Statistical Science, 2007.
[65]
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. In ICLR, 2014.
[66] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, nigh-performane deep learning library. In NeurIPS, 2019.
[67] Henning Petzka, Michael Kamp. Linara Adilova, Cristian Sminchisescu, and Mario Boley, Relative flatness and generalization. In NeurlPS, 2021.
[68] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015,
[69] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterizatioo to accelcrate training of deep neural networks. In NIPS, 2016.
[70] Jascha Sohl-Dicke Dickstein, Roman Novak, Samuel S Schoenholz, and Jachoon Lee. On the infinite width limit of neural networks with standard parameterization. arXiv preprint arXiv:2001.03301, 2020.
[71] Yang Song, Jiaming Song, and Stefano Ermon. Accelcrating natural gradient with higher-order invariance. In ICML, 2018.
[72] Michalis K Titsias. Learning model reparametrizations: Implicit variational inference by fitting MCMC distributions. arXiv preprint arXiv:1708.01522 2017.
[73] Yusuke Tsuzuku, Issci Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC- MAC-Baycsian analysis. In ICML, 2020.
[74] Jesse van Oostrum, Johannes Müller, and Nihat Ay. Invariance properties of the natural gradient in overparametrrsed systems. Information Geometry, 2022.
[75] Greg Yang and Edward J Hu. Tensor programs IV: Feature learning in infinite-width neural networks. In ICML, 2021.
[76] Greg Yang. Edward J Hu, Igor Babuschkin. Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer. NeurIPS, 2022.
[77] Liu Yang and Rong Jin. Distance metric learning: A comprehensive survey. Michigan State Universiy.
(2), 2006.
[78] Guodong Zhang, James Martens, and Roger Grosse. Fast convergence of natural gradient descent for ever-parameterized neural networks. NeurIPS, 2019.
13

---

Appendix A A Simple Working Example
Let θ := R2 and let 2 R² - R defined by :1 ==0) (0:)² be a wiice-differentiabl function on O. Suppose yp @
defined by -00) 30 =: 2i the reparametrization of choice. In this example, we shall see how this often-used reparametrization leads to pathologies in the values derived from the Hessian of L computed by an autodiff system.
Notice that the Jacobian of & is given by (() = diag(3,... ...,. E Rdxd and its inverse is (w) diag('/s, ................. 1/s). The autodiff Hessian matrix is given by H()) = 21 E Rdxd When we transform
......3................
0 + 1 the function (()) becomes L( ( := E/sub/// The autodiff will take into account this
change and thus will compute the Hessian H of L as
Hb)(}) = diag(//................... 1/3) (2T) diag(!/s..... 1/3) = diag(3/0 (/o,
Notice that, det (0) 24, meanwhile, we have det ( (2/9) Therefore, the Hessian determinant computed by an autodiff system is not invariant under reparametrization.
But if instead explicitly take into account the metric into the Hessian, which by default is (implicitly) chosen to be G())) = I, resulting in the seemingly redundant expression E(O) = I (2I), and perform the correct transformation both identity matrices,
G(Ø) (() I
obtain
((bb) (diag(//3,... ..., 1/3) diag(//a, 1/a) 1 diag(//3,...... 1/3) (2T) diag(t/a..... 1/a) diag(3, ... 3) diag(3,.:;
1/3) (2I) diag(t/a, ...,, .... 1/a)
diag(3..... ..... 3) (2f) diag(t/a,... ...,, 1/a)).
Thus, the determinant of the transformed "metrrc-conditioned Hessian" equals
det É( ( (det diag(3,.... .... 3)) (det (det diag(!/s.... 1/s)) 2 =2 which coincides with the determinant of the original E.
Appendix B
Derivations
Note. Throughout this section, we use the Einstein summation convention: If the same index appears twice, once as an upper index and once lower index, we sum them over. For example: = x²y means = Σ, x'y; y; and B = Alf=² means B = Eg Abjxr, r", etc. Specifically for partial derivatives, the index of the denominator is always treated as lower index.
B.1 The Riemannian Hessian Under Reparametrization
Let L: M -> TR be function on Riemannian manifold M with metric G. The Riemannian Hessian Hess L of C is defined in coordinates 0 by
8
r 0C
H jj 00:000
(6)
00k*
where n is the connection coefficient. 1)
Under a change of coordinates yp 0 + th have E = L yp" and
no O 3m a0m
8 Oip
80 Oup Oupi
Onp 00
where are just dummy indices-presedd present to express summations. Note that the transformation rule for P implies that is not a tensor-to be a tensor, there must not be the second term in the previous formula.
14

---

Using (7) and the standard chain & product rules to transform the partial derivatives in (6), we obtain the coordinate representation of the transformed Hessian Hess L:
(Lop
O(CopT)
H kg OupiOupi
Oaph
G d 00"
ac 0200
Ōrph 00mh d0"
0
Oup"
OC 00
To
00000" Oupi Orps)
00
mn 80 Orpi Orpi
Orbi
000
00 Ouph
G 00 00
dC 0200
Onky 00 00
OC 00%
0
Orly OL 00%
00m000 Onii 2ub
d 00
mn 00 Orp' dubi
00 Bibh
irpi Trbi A0 00 Bub
8 00mh 00
0C
0
00 00m OC
0 ag-
no
00 Jepi
Jeby
00
Oups
mn
Ozbi Orpi
8rb) 00
00 00m
L
OL
no
aapi Örbi
30
mn 00"
00 00m
H H
Orpi Örpi
(8) In the matrix form, we can write the above as =- where J is the Jacobian of y. Thus, the Riemannian Hessian at any I (not just at critical points) transforms just like the metric and thus invariant as discussed in Example 4. Note: this only holds when the term containing the connection coefficients r
ij is explicitly considered. In particular, the Euclidean Hessian does not follow this tensorial transformation under autodiff due the fact that (i) r 0 for any J. k and thus dropped from the equation, and (ii) autodiff is not designed to handle advanced geometric objects like '
sj
B.2 Hessian- Trace Under Reparametrization
Let C: Rd + R be function of Rd under the Cartesian coordinates and G Riemannian metric. The Riemannian trace of the Hessian matrix H of C is defined by [47]:
C:R
(trg H)) = (r(G(())''HV)
(9)
That is, it is defined as the standard trace of the Hessian operator E.
Let y : + be reparametrization on Rd. Then, using (2) and the property tr(AB) tr( twice, the Riemannian trace of the Hessian transforms into
(tr& (v}) tr(E(v))
= tr(((J (xb) (Q)
(10)
(4)
= -tre H) (b)).
Since 1bb (0). we have that (tr& (b) (tre (0) for any given 0. Therefore, the trace of the Hessian operator (or the Riemannian trace of the Hessian) is invariant.
B.3 Hessian-Eigenvalues Under Reparametrization
We use the setting from the preceding section. Recall that À is eigenvalue of the linear map (() = - (0) on the tangent space at E Rd that is represented 0 if E(O) = Xv for an eigenvector v E T, R.dd We shall show that X, the eigenvalue under under the reparametrization xp 0 + wb, equals the original eigenvalue A.
Using the transformation rule of E(0) in (2) and the transformation rule of tangent vectors in (4), along with the relation (J-1(b) (+)) -'= [(p (2b)), we get
E Āo
ÁJ (xb))
(11)
Ely ( Āv.
15

---

Since -0- ([b) 0, we arrive back at the original equation before the reparametrization and thus we conclude that = X.
Appendix C The Invariance of the Fisher Metric
We have seen that any metric on the parameter space yiclds inv(equi)variance when all transformation rules of geometric objects are followed. However, in practical settings, one uses autodiff libraries [1. 66]. whose job is to compute only the elementary rules of derivatives such as the chain and product rules. Notice that derivatives and Hessians (at least at critical points) are transformed correctly under autodiff, while the metric is not in general. is thus practically interesting to find a family of metrics that transform correctly and automatically under reparametrization given only an autodiff library. Under those metrics, the aforementioned inveeui)yarriane can thus be obtained effortlessly.
Martens [55, Sec. 12] mentions the following family of curvature matrices satisfying this transforma- tion behavior
B(0)
x E a(): A(0, æ. (z..)) ((: EE/D
(12)
where J(O:-) is the network's Jacobian 0s( 0)/0, with an arbitrary data distribution D) and an invertible matrix A that transforms like A (xpb((()) under yp sb. Under reparametrization yp, an autodiff library will compute
(13) given Lo --_ due to the elementary transformation rule of ((; just like Example 4(c). This family includes the Fisher and the generalized Gauss-Newton matrices, as well as the empirical Fisher matrix [44, 55].
B(tb) J- B ( [
However, note that any B as above is sufficient. Thus, this family is much larger than the Fisher metric, indicating that automatic invariance is not unique that metric and its denominations, Moreover, in practice, these metrics are often substituted by structural approximations, such as their diagonal and Kronecker factorization [56]. This restricts the kinds of reparametrizations under which these approximate metrics transform automatically: Take the diagonal of (12) an example. Its d-representation computed by the autodiff library will only match the correct transformation for element- wise reparametrizations. whose Jacobian ((0) is diagonal. On top of that, in practical algorithms, such approximations are usually combined with additional techniques such as damping or momentum, which further break their automaticness.
Due to the above and because any metric theoretically variant if manual intervention is performed, the inv(equi)variance property of the Fisher metric should not be the determining factor of using the Fisher metric. Instead, one should look into its more unique properties such its statistical efficiency [3] and guarantees in optimization [34, 78J. Automatic transformation is but cherry on top.
is interesting for future work to extend autodiff libraries to take into account the invariant transfor- mation rules of geometric objects. By doing so, any metric-not just the Fisher metric-wii wwlll yield invariance automatcally.
Appendix D General Manifold: Local Coordinate Charts
In the main text, we have exclusively used global coordinate charts (Rd 0) and (IRd, these charts cover the entire Rd and 0, to are homeomorphisms on Rd. However, there exist other coordinate systems that are not global, i.e. they are constructed using multiple charts {(U, Rd. 0 U. O S R s.t. Rd U,U; and each 0, is diffeomorphism on U,.
If {(U,, 0,)); and {VV,. dy)}; be two local coordinate systems of TR.d. Let (U,, 0,) and (V,, w/y)) be arbitrary pair of charts from the above collections where U, n V, * Ø. In this setting, reparametrization between these two charts amounts to the condition that the transition map between O(U; n V2) and vUUU, n V2) is a diffeomorphism, see Fig. 6. Reparametrization on the whole manifold can then be defined if for all pairs of two charts from two coordinate systems, their transition maps are all diffeomorphism whenever they overlap with each other.
16

---

Note that this definition is a generalization of global coordinate charts. In this case, there is only single chart for each coordinate system, i.e. (U, and V, 4/, and they trivially overlap since U I and is only single diffeomorphism of concern, as shown in Fig. 3.
Rd, Moreover, there
8
Our results in the main text hold in this general case by simply replacing the domain for the discussion to be the overlapping regions of pair of two charts, instead of the whole Rd
Appendix E Details on Applications
Figure 6: The diffeomorphism defined from O(U V) (U V)..
E.1
Infinite-Width Neural Networks
Note We use 2-layer NN without bias for clarity. The extension to deep networks follows directly by induction-sed e.g. [6, 45]. We also use the same prior variance o² without loss of generalization for simplicity. See also Kristiadi et al. [41, Appendix A] for further intuition. Recall the property of Gaussians under linear transformation:
(1z, Σ) = ~ N((N A
Neural-network Gaussian process (NNGP)
Let f(x) Rdd - I defined by f(x) == T di Wx) be a real-valued, 2-layer NN with parameter = ( E ARxx w E } and component-wise nonlinearity ø. Note that we assume æ is i.i.d.
Let vec(W) (N (²)) and ~ NO I be priors the weights. This parametrization of the weights and the priors called the standard parametrization (SP) [45, 60].
Step (a) Given a particular preactivation value 2, we have a linear model f(x) w (()). We can view this Gaussian process with mean and covariance
[[f(2)] 0 o(z) 0,
a² xz))
d² h
for(f()), h
Σ d(z)) (z2)). h i=1
Taking the limit as h > oo, the mean stays trivially zero and we have by the law of large numbers:
x² h
K (xx = lim CovlS(x),, 1++0
lim 11000 h
Σ ((2)) i=1
o² E (2z)) 2.
In particular, both the mean and covariance over the output does not depend on the particular realization of the hidden units ((z). That is, they only depend on the distribution of z induced by the prior, which will obtain now.
Step (b) Notice that each 2, = W. x where W, is the i-th row of W. This a linear model and thus, is distributed as GP with mean and covariance
E 2 = 0 00 0,
Cov 2,, 2!] 0² Ta' ''' K,(K,, 2()-
So, z, GO(O, K2) Since the prior over W is i..d., this holds for all i I,..,,h. h. We can thus now compute the expectation in K (x, æ²): it done w.r.t. this GP over zi.
To obtain the GP over the function output of a deep network, simply apply steps (a) and (b) above recursively. The crucial message from this derivation that the width of each layer of a deep net goes to infinity, the network loses representation power- the output of each layer only depends on prior, and not on particular values (e.g. learned) of previous hidden units. this sense, an infinite-width I.-layer NN is simply linear model with constant feature extractor induced by the network's first J. layers that are fixed at initialization. Note that the kernel K over the function output called the NNGP kernel [45].
17

---

Neural tangent kernel (NTK)
Let us transform u into v -h w and W into : a W and define the prior to be w ~ NO. I) and vec( W) ~ (O, Then, we define the transformed network as f(xx) == v (Vv)) =
vi
Vh w Φ Wæ) with parameter +b == {V E R E IR*). This is called the NTK parametrization (NTP) [29]. We will below that even though v, V have the same prior distribu- tions as w, W in the SP, they have different behavior in terms of the NTK.
As before, let us assume particular preactivation value z. The empirical NTK (i.e. finite-width NTK) on the last layer is defined by:
o² h Σ ((x,) o(21)
k(x, x') = (.)l =
h i=1
The (asymptotic) NTK is obtained by taking the limit of - xo:
K(x,, x') = lim K( x²) 1-+00
((z)) 0(1) = E [d(z))(())) =
(14)
which coincides the NNGP kernel K. Crucially, this obtained via backward propagation from the output of the network and thus the Gaaessan property we have used to derive the NNGP via forward propagation does not apply. This is why the scaling of Å required in the NTP. That is, o² using the SP, the empirical NTK is not scaled by hn and thus when taking the limit to obtain K, the sum diverges and the limit does not exist.
Is the NTP a reparametrization of the SP?
It is tempting to treat the NTP as reparametrization of the SP-in fact, it is standard in the literature to treat them as two different parametrizations of the same network. However, we show that geometrically, this is inaccurate. Indeed from the geometric perspective, if two functions reparametrization of each other, they should be invariant, as we have discussed in the main text. Instead, show that the different limiting behaviors are present because the NTP and SP assume two different functions and two different priors- they are not connected by a reparametrization. This clears up confusion and provides foundation for future work in this field: To obtain desired limiting behavior, study the network architecture and its prior, instead of the parametrization.
Suppose va in the NTP is reparametrization of 0 in the SP. Then the function yp 0 + vb defined by + a 0 is obviously the smooth reparametrization with an invertible (diagonal) Jacobian /h (() - I. In this case, the network in the NTP must be defined by f f= fo -1, where f is the sP-newwork, by Example That is, with some abuse of notation,
f(æ) = (v) o(p (V)x) w ((Wz)=f(z))
This is different from the definition of the NTP-network f(x) TT ((Vx.. So, obviously, the NTP is not the reparametrization of the SP. Therefore, clearer way of thinking about the NTP and SP is treat them as two separate network functions (i.e. two separate architectures) the scaling factor VE should be thought of as part of the layer's functional form instead of as part of the parameter. In particular, they
two representations of single abstract function.
To verify this, let us compute the NTK of f(x) (i.e. treating the scaling as a reparametrization) at its last layer. The derivation is based Section 3.2. First, notice that the differential / (z: transforms into J (v)) V f(xz)) [1 (w)) for any x E R". Next, notice that the Euclidean metric transforms into
J (v). So the gradient transforms into J ( (w)) ( f i(w) Therefore,
'It still applies for obtaining the distribution of 2.. The NTK can thus be thought of as a kernel that arises from performing forward and backward propagations once at initialization [6. 75J. This can be seen in the expression of the NTKs on lower layers which decompose into the NNGP and an expression similar to (14), but involving the derivative of [29].
"We use the gradient to get the NTK since otherwise it does not make sense to take the inner product of differentials w.r.t the metric.
18

---

Table 2: Test accuracies, averaged over random seeds.
Methods
MNIST
FMNIST
CIFARIO
CIFARI00
SGD
99.3
92.9
94,9
76.8
ADAM
99.2
92.6
92.4
71.9
the empirical NTK kernel Kv for f is given by
(x( (J( f(x)(())))))))) (m), (o))a
(J(z) (2))
G(o) ( (
=VV((=)l (()) (B)) (O) ( (a)) ((0)
(V)(=)) 1()) (= ((0))
Thus, the empirical NTK is invariant and the asymptotic NTK also is. Therefore, still have problem with the NTK blow-up in this parametrization, This reinforces the fact that the difference between the SP and NTP not because of parametrization.
Additionally, us now inspect the priors in the and NTP. In the SP, the prior N/O 0. /bI). Therefore, so that we have the same prior in both @ and V, the prior of a/ (0) must be N( 0, I). This is obviously not the case since we have N/ú 0, o//d) because the NTP explicitly defines NO 0, I) as the prior of 0. Thus, not only that the SP and NTP assume two different architectures, but they also assume two different prior altogether. is thus not surprising that the distribution over their network outputs f(x). f(x) are different, both in the finite- and infinite-widihe regimes.
Implication In his seminal work, Neal [60] concluded that the fact that infinite-wied NNs Gaussian processes disappointing. However, as we have seen in the discussion above, different functional forms, architectures, and priors of NNs yield different limiting behaviors. Therefore, this gives us hope that meaningful, non-GP infinite-width NNs can be obtained, Indeed, Yang and Hu [75], Yang al. [76] have recently shown us a way to do so. However, they argue that their feature-learning limiting behavior is due to a different parametrization, contrary to the present work. Our work thus complements theirs and opens up the avenue for constructing non-trivial infinite-width NNs in a "Bayesian" way, in the sense that we achieve the desired limiting behaviors by varying the model and the prior.
F.2 Biases of Preconditioned Optimizers
For MNIST and FMNIST, the network is LeNet. Meanwhile, we use the WiderResNett16-44 model for CIFAR- 10 and -100. For ADAM, we use the default setting suggested by Kingma and Ba [39]. For SGD, we use the commonly-used learning rate of 0.1 with Nesterov momentum 0.9 [26]. The cosine annealing method is used to schedule the learning rate for 100 epochs. The test accuracies in Table 2. Additionally, in Table 3, we discuss the effect of reparametrization sharpness ADAM and SGD.
E.3
Laplace Marginal Likelihood
Let OMAP be MAP estimate in arbitrary 0-coordinates of Rd obtained by minimizing the MAP -note loss CMAP. Let log h - -LMAP- note that LMAP itself a log-density function. The Laplace marginal likelihood [18, 28, 54] is obtained by performing second-order Taylor's expansion:
1 - -MAPP) H (0MAP)) - 0MAP).
log (() R log le ((0AP))
2
where H (0MAP) is the Hessian matrix of LMAP at OMAP. Then, by exponentiating and taking the integral over 0, we have
2(0MAP) h(OMAP) / exp (- (MAP) TH ((OMAP)) (MAP) d0.
(15)
19

---

Table Hessian-bass based sharpness measures can change under reparametrization without affecting the model's generalization (results on CIFAR-10). The generalization gap is the test accuracy, subtracted from the train accuracy-lower is better. Under the default parametrization, SGD achieves lower sharpness and generalizes better than ADAM which achieves higher sharpness. However, one can reparametrize SGD's minimum s.t. it achieves much higher (or lower) sharpness than ADAM while retaining the same generalization performance. Hence, is hard to study the correlation between sharpness and generalization. This highlights the need for invariance.
Optimizer
Reparametrization dDMAP = (((0AA))
Generalization gap [%]
Sharpness tr
ADAM
wMMA OMAP
7.2 + 0.2
1929.8 +: 61.2
AAMA = AMAP
1531.8 + 14.2
SGD
UMMA 7 +0MAP
5.2 + 0.2
6143.7 + 60.8
UMMAP 20MAP
383.6 1: 3.3
Since the integral is the normalization constant of the Gaussian NN OMAP, ((MAAP)), we obtain the Laplace log marginal likelihood (LML):
log (0MAP) = -LMMP(OMMAT 2
log(2x) + log det H (0MAP).
Notice that (0) is bilinear form, acting on the tangent vector d(OMAP) (0- 10--00 0MAP). Under reparametrization 0 ab with 4PMAP (((MA))) the term inside the exponent in (15) transforms into
(J (bMa)))() ((MMA))) ((bxx)))) 2
TJ H (W(P)((aa)) (z/MP))
(
J [ (((MA)) d(<p (WMAP)),
due to the transformations of the tangent vector and the -eessaa. This simplifies into
exp -- (HMAA)) ((
H( ((
which always equals the original integrand in (15). Thus, the integral evaluates the same value. Hence, the last two terms of log (0MAP) transform into - & log(2 + log det ((pMa))) in vi-coordinates. This quantity is thus invariant under reparametrization since it behaves like standard functions.
E.3.1 Experiment Setup
We use the toy regression dataset of size 150. Training inputs are sampled uniformly from [0, 8). while training targets are obtained via y sin - e, where e (N(0,, 0.32). The network is 1 -hidden layer TanH network trained for 1000 epochs.
20


Reviews:
[{"review_id": "YIVQs4ufdV", "reviewer": "Reviewer_ecDa", "summary": "Under model reparametrization, Hessian-based flatness measures, optimization trajectories, and probability densities are not invariant. Motivate by these inconsistencies, this paper studies the invariance associated with the reparametrization of neural networks. By viewing the parameter space as a Riemannian manifold, the authors show that the invariance and equivariance under reparametrization is preserved by explicitly including the metric when computing geometric objects such as the Hessian. The authors point out that acknowledging the metric helps in measuring the flatness of minima, optimization, and probability-density maximization. \n", "questions": "- By comparing the definitions on page 1, it seems that invariance under symmetry is a special case of invariance under reparametrization. Can the reparametrization studied in this paper be viewed as more general than previous works on weight-space symmetries?\n- Could the authors elaborate on the intuition on why we should explicitly include the metric when comparing the sharpness of the solution found by different optimizers, as suggested in section 5.3?", "limitations": "The authors state limitations at the end of the introduction section. There are no potential negative societal impacts of the work.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper draws attention to the nature of reparametrization through Riemannian geometry. By introducing a framework that transforms representations of geometric objects to keep them invariant under reparametrizations, the paper provides a useful tool in comparing properties of neural networks after reparametrization.\n\nI appreciate the authors’ effort to make the derivations both mathematically rigorous and accessible. In particular, since the parameter space is usually Euclidean, it makes sense to present most of the material in linear algebra terminologies instead of the more general Riemannian geometry. ", "weaknesses": "The novelty of this paper seems limited. As the authors also mention, the lack of invariance under reparametrization has been observed before. The transformation of various quantities in reparametrization has also been discussed (see below). The discoveries on the application side are also not well-presented. As a result, it is not clear what the main contributions are.\n\nSome parts of the paper could be explained in more details. (a) The goal of section 5.1 is not clear. Is the goal to show that SP and NPT are different because they cannot be seen as reparametrizations of each other? (b) The significance of section 5.2 might be clearer if the authors could add a sentence to give a general intuition for Equation 5. \n\nThe transformation of the Hessian and gradient under reparametrization has been discussed in a previous paper that is not cited [1]. Could the authors comment on how their approach in section 3 is different?\n\n[1] Kunin, Daniel, et al. \"Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics.\" International Conference on Learning Representations. 2020. (Appendix A)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your feedback! Here we address your major comments/questions. We incorporated your suggestions into the text directly. Please see also our [\"global\" response](https://openreview.net/forum?id=vtLNwa6uX0&noteId=CvzNeGzNp7) for a general discussion.\n\n**Novelty** The main contribution of our present work is to show that invariance under reparametrization of many quantities relevant to neural nets is natural when considering the correct transformation of geometric objects. While indeed previous work has discussed non-invariance under reparametrization, they either accept it as a fact (e.g. Sec. 5 of [1]) or come up with a special metric/method to overcome this issue (e.g. [2]). Our argument, meanwhile, is broadly applicable since no assumption about the metric is imposed. Our work thus provides a guardrail against past & future confusions regarding invariance. We hope to invoke more interest in this topic by providing some example applications in Sec. 5. Finally, please note that other reviewers mentioned that our work provides a new perspective (**_R.Poha_**), brings much-needed clarity (**_R.LwSr_**), and is useful since making ML researchers aware of the discussed topic is of utmost significance (**_R. togv_**).\n\n**Sec. 5.1** Our goal is to clear up confusion about the term “parametrization” in SP and NTP. In the preceding sections we have argued that if two parametrizations are connected by a diffeomorphism, they represent a single function. But since SP and NTP, despite their monikers, are not reparametrization of each other, then they represent two different functions. This clears up the confusion why they have different limiting behavior (in terms of the NTK and NNGP kernels) and calls for a more suitable way of analyzing them.\n\n**Reparametrization vs symmetry** When one fixes the bijective transformation $T$, and it happens to be smooth, then indeed they can be seen as the same, e.g. in the normalized NN example, the scaling $c$ is fixed. However, in general $T$ is not fixed, i.e. we consider instead a _family_ of invertible $\\Theta \\to \\Theta$ and so they are inherently different. Invariance under symmetries is better expressed in terms of group actions from (Lie) group theory, i.e. the map $T$ should instead be defined as $T: \\mathcal{G} \\times \\Theta \\to \\Theta$ where $\\mathcal{G}$ is a group. For normalized NNs, $\\mathcal{G} = \\mathbb{R}_{>0}$, i.e. we take into account _all_ scaling factors $c > 0$ in $\\mathcal{G}$ instead of fixing it. The definition of symmetry on page 1 is indeed incomplete (missing the group term) and we have updated it. Please note that this missing term does not affect the discussion since we focus on reparametrization, not symmetry. Note also that the group-action definition is standard in the literature, including in Kunin et al.\n\n**Kunin et al.** They analyze the gradient and the Hessian under symmetry, not reparametrization (see above). Their work therefore tackles a different problem than ours. In particular they study invariance under group actions, while we study invariance under diffeomorphism. Both are important (as written in our paper and in e.g. [1] Sec. 3 and Sec. 5) and our work is complementary to theirs—our geometric insights further enhance their symmetry-invariant gradient and Hessian formulations with reparametrization invariance. We added their work to the citation list. Thanks for pointing out their work! \n\n**Why include the metric in sharpness** There are at least two reasons why. _First_, as we have shown in Sec. 2.2.1 and Sec. 3.1, this is the geometrically principled way to compute Hessian-based sharpness (trace, det, eigenvalues), and this naturally yields invariance under reparametrization and solves the problem shown by [1, Sec. 5]. _Second_, when using a preconditioned gradient descent (PGD), the metric-infused sharpness yields similar behaviors as standard gradient descent [3], allowing for direct comparisons. Indeed, PGD is essentially just a GD acting on a space with different geometry induced by the metric. Moreover, by taking into account the parameter-space metric, this reveals that the loss landscape geometry under ADAM’s metric is actually much sharper than that assumed in GD (our Fig. 5)---this information might be useful for future work.\n\n[Our answer to **_R.LwSr_**](https://openreview.net/forum?id=vtLNwa6uX0&noteId=OHnkKsN22C) regarding symmetry vs reparametrization might also interest you.\n\nPlease let us know if you have further comments/questions/suggestions!\n\n**References**\n\n[1] Dinh, Laurent, et al. \"Sharp minima can generalize for deep nets.\" ICML, 2017.\n\n[2] Jang, Cheongjae, et al. \"A reparametrization-invariant sharpness measure based on information geometry.\" NeurIPS, 2022.\n\n[3] Cohen, Jeremy M., et al. \"Adaptive gradient methods at the edge of stability.\" arXiv preprint arXiv:2207.14484 (2022)."}, {"review_id": "uub67ACYqy", "reviewer": "Reviewer_Poha", "summary": "This work analyzes the invariances and non-invariances of model reparameterization in machine learning. The authors show that, if we account for Riemannian metrics in parameter spaces, then many quantities thought to be not invariant are in fact invariant to reparameterization. Thus, by properly applying transformation rules on geometric quantities, we can obtain equivariant or invariant functions on parameter space.\n", "questions": "n/a", "limitations": "Discussion of limitations on Page 2.\n", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Covers applications of this type of thinking in several areas of machine learning.\n2. Good, careful exposition of geometric concepts and calculations.\n3. The note on the utility of non-invariant reparameterization for normalizing flows and optimization is interesting.\n4. Overall, this work gives a useful perspective that helps analyze ML models (e.g. Section 5.2), and will hopefully give actionable insights to improve them (e.g. other metrics instead of Fisher).\n", "weaknesses": "1. I am not sure about the utility of the suggested method of measuring sharpness, and I would appreciate if the authors could comment on this. Indeed, the sharpness of ReLU networks depend on the scale of the weights chosen. However, [Du et al. 2018] shows that there is some implicit bias so that arbitrary scales of the weights are not converged to by GD in practice. Also, there are generalization results in terms of (I believe) non-invariant Hessian trace [Ding et al. 2023].\n\n[Du et al. 2018] Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced. NeurIPS 2018  \n[Ding et al. 2023] Flat minima generalize for low-rank matrix recovery. arXiv 2023.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review! Please see our [\"global\" response](https://openreview.net/forum?id=vtLNwa6uX0&noteId=CvzNeGzNp7) for a general discussion. Here, we address your specific comments.\n\n**Du et al.** They focus on the invariance of ReLU networks under the scaling symmetry, while we are tackling the problem of invariance under reparametrization. This difference is also made clear by Dinh et al. 2015, Sec. 3 (symmetry) vs. Sec. 5 (reparametrization). See also [our answers to **_R.ecDA_**](https://openreview.net/forum?id=vtLNwa6uX0&noteId=1K4jej38ao) and our \"global\" response.\n\n**Ding et al.** Indeed they showed generalization results with the non-invariant Hessian trace. However, their results will have the pathologies that we discussed in our paper and in Sec. 5 of Dinh et al. Our work is compatible with them in that it provides a guardrail for their work: we give the necessary steps to make their results resistant to pathologies under reparametrization.\n\nIn any case, we added both works to the related work section of our paper. Please let us know if you have further comments or questions!\n\n**References**\n\n[1] Dinh, Laurent, et al. \"Sharp minima can generalize for deep nets.\" ICML, 2017.\n"}, {"review_id": "1tKC1qbCT2", "reviewer": "Reviewer_LwSr", "summary": "The paper shows that reparameterizations of neural networks can be understood uisng Riemannian geometry. They first show how a reparameterization of a neural network's parameters can be expressed via a Riemannian metric which then yields transformation rules that can be applied to any function on the parameters. The paper uses this as basis to show why the determinant, trace, and Eigenvalues of the loss Hessian are not invariant under reparameterization, and how applying the correct transformation rules yields reparameterization invariance. The same is shown for gradient descent and probability densities. Lastly, the paper applies this to infinite-width neural networks (showing that the NTK is not a reparameterization of a standard infinite-width Bayesian NN), to the Laplace marginal likelihood, and to preconditioned optimizers. ", "questions": "- It might be interesting to see how a very simple reparameterization (e.g., multiplying one layer with a constant c and the next with 1/c for ReLU NNs) can be interpreted in Riemannian geometry. I.e., what is G in that case, how would the transformed Hessian look like? While this might not fit into the main text, it would make for a great practical example in the appendix for readers (like myself) not too familiar with Riemannian geometry.\n- The relation between flatness and generalization in light of reparameterizations has been established theoretically in [1].\n- Is it possible to interpret relative flatness [1] as an invariant transformation? That is, is it possible that $G^-1(\\theta) = ||\\theta||^2_2$? \n\n[1] Petzka, Henning, et al. \"Relative flatness and generalization.\" Advances in neural information processing systems 34 (2021): 18420-18432.", "limitations": "The paper is clear about the assumptions and limitations.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Viewing reparameterizations from the perspective of Riemannian geometry brings much needed clarity to the discussion.\n- The paper is excellently written and insightful, a pleasure to read.\n- The paper discusses a wide range of implications relevant to machine learning.", "weaknesses": "- The discussion on flatness-based generalization measures does not address existing generalization bounds for reparameterization-invariant flatness measures (see questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your positive review! Please see also our [\"global\" response](https://openreview.net/forum?id=vtLNwa6uX0&noteId=CvzNeGzNp7) for a general discussion. Here, we address your major comments. All other comments and suggestions are implemented directly in the paper.\n\nYou are completely right that our work is positioned in such a way as to bring much-needed clarity to the behavior of neural networks’ parameter spaces under reparametrization. We are glad to hear that our work can be followed easily and gives insights to you, as a researcher outside of Riemannian geometry.\n\n**Relative flatness of Petzka et al.** We note that Petzka et al., while using the term “reparametrization” in their paper, actually tackle the “symmetry” problem discussed in Sec. 1 in our paper. Reparametrization is a specific transformation of the parameter space under a smooth invertible map (diffeomorphism)—in the language of calculus, it is essentially the change-of-variable formula, e.g. in integration by substitution. Symmetry, meanwhile is defined through group actions [1], e.g. under rescaling of the weights where not just one scaling factor $c > 0$ is considered, but all of $c \\in \\mathbb{R}\\_{>0}$ (the space $\\mathbb{R}_{>0}$ here is the group). \n\nIn any case, both kinds of invariance (under symmetry and reparametrization) are important as we mentioned in our paper (see also [2], Sec. 3 and Sec. 5) since in general one does not imply the other. (See our \"global\" response.) They thus complement each other and our work complements Petzka et al.’s relative flatness which provides an invariant measure for generalization under rescaling group actions, but is not invariant under reparametrization. (Notice that relative flatness is defined through the non-invariant Hessian trace, and we show in our paper how to make it invariant.) In other words, our work makes relative flatness invariant to _both_ rescaling group actions and reparametrization.\n\n**Interpretation of as an invariant transformation** The symmetry $f(c \\theta) = f(\\theta)$ can be written as a group action of the group $\\mathcal{G} := \\mathbb{R}_{>0}$ on the manifold $\\Theta := \\mathbb{R}^n$ by multiplication. Your intuition is spot on that $\\mathcal{G}$-invariant quantities such as relative flatness can be seen as quantities in the \"symmetry-free space\". In the above case, one can think of relative flatness as a generalization metric on the quotient space $\\Theta / \\mathcal{G}$, which happens to be the sphere $\\mathbb{S}^{n-1} = \\\\{ \\theta \\in \\Theta : \\Vert \\theta \\Vert^2_2 = 1 \\\\}$. \n\n**Example with simple reparametrization** As we have discussed before, the rescaling \"reparametrization\" is better described as \"symmetry\" and thus is not suited as an example in this paper. But, we are happy to give a simple step-by-step example of reparametrization and its effect on the metric, Hessian, etc. in the appendix. We will do so by expanding Example 1a, i.e., using the transformation $\\theta = \\log \\psi$.\n\n[Our answer to **_R.ecDa_**](https://openreview.net/forum?id=vtLNwa6uX0&noteId=1K4jej38ao) might also be of interest to you. \n\nThanks again and please let us know if you have further comments!\n\n**References**\n\n[1] Kunin, Daniel, et al. \"Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics.\" ICLR, 2020\n\n[2] Dinh, Laurent, et al. \"Sharp minima can generalize for deep nets.\" ICML, 2017."}, {"review_id": "U2IVMRBZaD", "reviewer": "Reviewer_togv", "summary": "The paper discusses reparametrizations of parameter spaces and the implied transformation rules for quantities like gradients, Hessians or probability densities. \nParameter spaces are interpreted as Riemannian manifolds $M=\\mathbb{R}^d$ and the quantities of interest are coordinate independent geometric objects like tangent vectors, covectors, tensors, or volume forms.\nThe transformation laws of such objects are well known in differential geometry,\nhowever, as the authors argue, they are in deep learning applications often disregarded.\nIn particular, the metric tensor on $M=\\mathbb{R}^d$ is often taken to be $G=\\mathbb{I}$ and therefore dropped from the equations.\nIts transformation when changing coordinates is then naively forgotten, that is, one uses again metric coefficients $\\hat{G}=\\mathbb{I}$ in the new coordinates, which correspond geometrically to a _different Riemannian geometry_, such that the algorithms become coordinate dependent.\nThe contribution of the paper is to point out these disregarded transformation rules and to discuss how the quantities should actually transform.\n\nSection 2 introduces the mathematical setting and explains transformation rules of geometric quantities, giving in particular examples of how their transformation laws are consistent by canceling out in different tensor contractions.\nThe third section discusses more specific quantities of interest in deep learning.\nFirstly, it considers the Hessian matrix determinant as measure of flatness of parameter landscapes.\nThe authors argue that this measure is not well-suited since it depends on the choice of coordinates.\nThey propose to use the determinant of $G^{-1}H$ instead, which is invariant under reparametrizations.\nSecondly, it considers loss gradients $\\mathrm{grad}\\mathcal{L} := G^{-1}\\nabla\\mathcal{L}$.\nOn $M=\\mathbb{R}^d$, the trivial metric $G=\\mathbb{I}$ is usually dropped, which leads again to inconsistent transformations of the gradient.\nLastly, Section 3.3 discusses probability density functions (pdfs).\nWhen being expressed relative to the Lebesgue measure on $M=\\mathbb{R}^d$, the pdfs transform with the well known Jacobian determinant factor.\nHowever, as the density may be stretched out or condensed in this procedure, the densities' mode may not be preserved.\nIt is hence more suitable to express the density relative to the Riemannian volume form: as this form transforms itself already with the Jacobian determinant factor, the density relative to it remains invariant, which preserves in particular its modes.\nAfter discussing related work in section 4, the fifth section considers some applications, arguing in particular\nthat NTK and standard parametrizations or neural networks are not just reparametrizations, but geometrically truly different models (section 5.1) and\nthat the Laplace marginal likelihood is invariant under reparametrization (section 5.2).\nIt investigates furthermore the effect of using the Hessian or $G^{-1}H$ in (preconditioned) optimizers like ADAM.", "questions": "The issue of a coordinate independent formulation of deep learning algorithms was previously studied in the publication \"Equivariant and coordinate independent convolutional networks\" by Weiler et al. (2021), which should be mentioned.\nI would furthermore be very interested in how their coordinate independence of feature vectors and neural network operations relates to the coordinate independence of parameter spaces in the current submission?\n\nFurther minor suggestions and corrections beyond what I wrote above follow. These should be very easy to fix.\n\nLine 50 states that\n_\"Intrinsic ... means that objects ... must be independent of ... coordinate system\"_.\nThis concept is called coordinate independence or covariance.\n\"Intrinsic\" refers instead to \"not extrinsic\", where extrinsic properties are geometric properties depending on an embedding of the manifold in some ambient space (e.g. sectional curvatures).\n\nCharts are defined as homeomorphisms, but the paper considers smooth manifolds.\nThe smooth structure is actually only respected (and defined via) smooth charts, i.e. diffeomorphisms.\n\nIt would probably be helpful to note that the \"standard choice\" of global chart in line 97 is the canonical identity map $\\mathrm{id}_{\\mathbb{R}^d}$.\n\nThat the choice of coordinate system is not unique is somewhat tricky,\nas there exists the _canonical_ global chart $\\mathrm{id}_{\\mathbb{R}^d}$ mentioned by the authors.\nIf the manifold was just taken as Euclidean space, i.e. with metric but without the canonical coordinates of $\\mathbb{R}^d$,\nglobal charts would still not be arbitrary, but one could restrict to isometries, which are defined up to transition maps in the Euclidean group $\\mathrm{E}(d)$.\nThis raises the question why we are considering general diffeomorphisms in the first place?\nIs the metric just an arbitrary choice (which should of course still be respected, requiring coordinate independence)?\nIt would be great if the authors could discuss this point.\n\nLine 120 mentions that\n\"Under a coordinate system, one can think of both tangent vectors and covectors as vectors in the sense of linear algebra, i.e., tuples of numbers\",\nhowever, also abstract coordinate free vectors are part of linear algebra.\nI would just write that vectors and covectors are in coordinates represented by numerical coefficient vectors, i.e. tuples of numbers.\n\nLine 123: The metric is not only positive definite, but also symmetric.\n\nI am confused about the \"surjective everywhere\" in line 134.\nIs the problem with bijectivity not that the mapping from parameters to models is in general non-injective?\n\nIt might be helpful to mention somewhere around line 193 that the $\\Gamma_{ij}^k$ are called Christoffel symbols. This way the reader unfamiliar with these concepts can read up on it.", "limitations": "As the paper does not propose a new method but comments on the mathematical formulation of theories, limitations do not really apply.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors observe that many of the mathematical formulations used in deep learning are from a geometric perspective not coordinate independent, which is a fundamental feature that any consistent mathematical theory should satisfy.\nIts main contribution lies in pointing out this issue and discussing how the equations can be fixed.\nWhile this contribution is not \"original\" in the sense that it would be a novel idea, making researchers aware of it is of utmost significance since a coordinate independent formulation of algorithms is a fundamental requirement.\nIt is hard to judge how clear the paper is for someone without background in differential geometry, but it is kept simple and is certainly easy to reed for someone knowing about differential geometry.", "weaknesses": "A main weakness of the paper is that the mathematical formulation could be more precise at some points.\nI would usually not be too strict in deep learning, however, as the paper sets out to fix the inconsistent use of mathematics, it should be more precise.\nThese weaknesses should be easy to fix - more details follow in the next paragraphs.\n\nFirstly, the considered coordinate charts and hence reparametrizations are _global_ homeomorphisms.\nThis is in principle possible, but excludes practically relevant choices like polar or spherical coordinates, which are not global homeomorphisms.\nPolar coordinates are, in fact, used as an example right after saying that charts should be global homeomorphisms.\nThe only clean way around this issue (and to include polar coordinates) would be to admit the usual atlases of _local_ charts and study reparametrizations as usual on the intersections of charts.\nSpecifically, there should be charts\n$\\theta: U^\\theta \\to \\theta(U^\\theta) \\subseteq \\mathbb{R}^d$ and $\\psi: U^\\psi \\to \\psi(U^\\psi) \\subseteq \\mathbb{R}^d$\non domains $U^\\theta\\subseteq M=\\mathbb{R}^d$ and $U^\\psi\\subseteq M=\\mathbb{R}^d$\nwith transition maps $\\varphi: \\theta(U^\\theta\\cap U^\\psi) \\to \\psi(U^\\theta\\cap U^\\psi)$.\nNote that the main results of the paper will still hold in this more general setting.\n\nSecondly, the paper keeps talking about a dubious concept of \"_equivariance under reparametrization_\", which is non-standard and, while looking somewhat similar to the usual concept of group equivariance, is different from it and more confusing than it is enlightening. \nIt is introduced at the end of the second section,\nassuming parameter space reparametrizations $\\varphi:\\Theta\\to\\Psi$\nand functions which simultaneously seem to satisfy $F: \\Theta\\to\\Theta$ and $F: \\Psi\\to\\Psi$.\nThis would in principle require $\\Theta=\\Psi$, while actually only $\\Theta\\cong\\mathbb{R}^d\\cong\\Psi$ is demanded initially.\nFurthermore, this does not work in the clean formulation with local charts suggested above, since then \n$\\varphi: \\theta(U^\\theta\\cap U^\\psi) \\to \\psi(U^\\theta\\cap U^\\psi)$\nand the upper and lower arrows would be\n$F^\\theta: \\theta(U^\\theta\\cap U^\\psi) \\to \\theta(U^\\theta\\cap U^\\psi)$ and\n$F^\\psi: \\psi(U^\\theta\\cap U^\\psi) \\to \\psi(U^\\theta\\cap U^\\psi)$, respectively.\nThe corresponding diagram would be the usual coordinate independence transformation rule $F^\\psi = \\varphi\\circ F^\\theta\\circ \\varphi^{-1}$, which is _not_ an equivariance condition.\nIt can in general also not be made to one by setting $F^\\psi = F^\\theta$, as this would require the equality of $\\theta(U^\\theta\\cap U^\\psi) = \\psi(U^\\theta\\cap U^\\psi)$ in the first place.\n\nThe concept of equivariance is subsequently used in a rather unspecific way in the paper.\nSection 3.2 talks about and is titled \"equivariance of gradient descent\", however, this is not made precise, i.e. no equation is mentioned which follows the commutative diagram in the author's definition of equivariance.\nFurthermore, equivariance is a property of _functions_, and it is not entirely clear to me how one should interpret the gradient descent algorithm as such.\nSection 3.3 is titled \"equivariance of probability densities\", which seems to refer to equation 4, i.e. $q_\\Psi^G(\\psi) = q_\\Theta^G(\\varphi^{-1}(\\psi))$.\nHowever, this is just a coordinate independence equation, since the functions $q_\\Psi^G$ and $q_\\Theta^G$ on the left and right hand side differ from each other (they are not the same $F$ as in the diagram).\n\nThe issue of this dubious notion of \"equivariance\" is easily fixed by removing it from the paper\nand referring to it as usual as coordinate independence or covariance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your extensive review! Your write-up on the summary and strengths of our paper is completely spot on!\n\nHere we will address your major comments and questions. Minor comments and suggestions will be implemented directly in the text. See also our [\"global\" response](https://openreview.net/forum?id=vtLNwa6uX0&noteId=CvzNeGzNp7) for general discussion.\n\n**Precise mathematical formulation** You are right that it can indeed be more precise. Your suggestions on the usage of local-chart formulation and removing the current notion of “equivariance” & changing it into “coordinate-independence” are spot on and we will implement them in the paper. On the other hand, we are trying to make the text accessible to a broader audience in deep learning—this is our target community. Note that other reviewers, who seem to fall into this category, mentioned that our paper is accessible (**_R.ecDa_**, **_R.Poha_**) and insightful (**_R.LwSR_**). So, we will follow your suggestion by showing predominantly intuition and figures in the main text and deferring the extra mathematical details to the appendix. \n\n**Weiler et al. (2021)** They tackle the problem of “geometric deep learning”, where the manifold of interest is the input (and feature) space. So, at a high level, our work—which focuses on the parameter space—is complementary to theirs. E.g., one benefits from our work when measuring the sharpness of the loss landscape of a gauge-equivariant NN. While indeed their theory tells us how to transform parameters (i.e. convolution kernels in their case) under a gauge transformation (a group action), they are compatible with our work in the same way that the symmetry of a manifold is compatible with coordinate independence of the same manifold.\n\n**Choice of coordinate systems, is the metric arbitrary?** $\\mathbb{R}^n$ with the canonical global coordinates is often the default choice for the parameter space of deep networks and only the metric is varied (e.g. in natural gradient methods and normalizing flows—the latter can be seen as a metric-learning mechanism through a non-invariant coordinate transformation acting on the canonical coordinates). Even for geometric-focused NNs like gauge equivariant nets [1], where strong manifold assumptions are imposed in the input/feature space, to our knowledge no further manifold assumption is applied on the parameter space, other than possibly the metric (e.g. using SGD vs. ADAM during the optimization). In this sense, considering general diffeomorphisms is useful since our work then provides the coordinate-independence guarantee and preservation of the metric in broad deep-learning applications.\n\n**Surjective everywhere** The problem is that the map $\\theta \\mapsto f(X; \\theta)$ is almost always a submersion for an overparametrized network [2, 3], and certainly not a diffeomorphism. Meanwhile, the requirement for pulling back a metric is that the map must be an immersion. (Note that (local) diffeomorphism implies immersion.) We will rephrase “surjective everywhere” into “non-injective everywhere” to make this point clearer.\n\nIf you have further suggestions, we always welcome them! Thanks again for your great comments and suggestions!\n\n**References**\n\n[1] Cohen, Taco, et al. \"Gauge equivariant convolutional networks and the icosahedral CNN.\" ICML 2019.\n\n[2] Zhang, Guodong, James Martens, and Roger B. Grosse. \"Fast convergence of natural gradient descent for over-parameterized neural networks.\" NeurIPS 2019.\n\n[3] Karakida, Ryo, and Kazuki Osawa. \"Understanding approximate Fisher information for fast convergence of natural gradient descent in wide neural networks.\" NeurIPS 2020.\n"}]

Author rebuttal:
**To all reviewers:** Thank you very much for your input! To supplement the responses to your individual reviews, here we would like to address the common questions and comments.

Our work focuses on addressing _invariance under reparametrization_, i.e. under change of variable from the point of view of calculus or the concept of coordinate independence from the point of view of differential geometry. In particular, it studies the invariance of quantities like gradients and Hessians when the coordinates of the parameter space are mapped into _new_ coordinates. This is different from _invariance under symmetry_, where the invariance is studied under a _group_ acting on a _fixed_ choice of coordinates of the parameter space.

Crucially, these two concepts are compatible with each other. Indeed, from the differential geometry point of view, coordinate independence is an _inherent_ property of a manifold as **_R.togv_** also pointed out; symmetry is a property that _can be_ studied further on that manifold—every manifold is coordinate-independent, but not every manifold has symmetry. Thus, our work complements previous works and provides a guardrail for future works that focus on studying invariance under symmetry.  For example, our work enhances the works of Kunin et al. [1], Petzka et al. [2], and Cohen et al. [3, 4], where they addressed invariances under various symmetries but are still susceptible to pathologies under reparametrization. (See our individual responses for more details.) Please note also that even though these previous works often used the term "reparametrization", their problem is better termed as "symmetry", as per the differential-geometric definitions we use.

**References**

[1] Kunin, Daniel, et al. "Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics." ICLR, 2020

[2] Petzka, Henning, et al. "Relative flatness and generalization." NeurIPS, 2021.

[3] Cohen, Taco, et al. "Gauge equivariant convolutional networks and the icosahedral CNN." ICML 2019.

[4] Weiler, Maurice, et al. "Coordinate Independent Convolutional Networks--Isometry and Gauge Equivariant Convolutions on Riemannian Manifolds." arXiv preprint arXiv:2106.06020 (2021).


Meta-review: The reviewers reached unanimous agreement that the paper merits acceptance. They noted that the paper gives a rigorous yet accessible treatment of reparameterization of neural networks from a differential geometry perspective, providing much needed clarity to the topic. Reviewers also appreciated the fact that the paper discusses a wide range of implications for ML. As such I recommend accepting the paper as a spotlight presentation.
Decision: Accept (spotlight)