Title: The Geometry of Neural Nets' Parameter Spaces Under Reparametrization
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

Neural networks (NNs) are parametrized functions. Since it is often desirable to assign meaning or interpretability to the parameters (weights) of a network, it is interesting to ask whether certain transformations of the parameters leave the network invariant -equivalent in some sense. Various notions of invariance have been studied in NNs, in particular under weight-space symmetry [5, 14, 20, 42, 53] and reparametrization [20, 21, 30, 32, 71]. The former studies the behavior of a function L ( θ ) under some invertible T : Θ ×G → Θ where G is a group; for any θ , the function L is invariant under the symmetry T if and only if L ( θ ) = L ( T ( θ , g )) for all g ∈ G . For example, normalized NNs are symmetric under scaling of the weights, i.e. L ( θ ) = L ( c θ ) for all c > 0 [7]-similar scale-symmetry also presents in ReLU NNs [61]. Meanwhile, invariance under reparametrization studies the behavior of the NN when it is transformed under the change-of-variable rule

Philipp Hennig

L

(

θ

)

L

(

T

(

θ

, g

))

Θ

L

(

θ

)

Θ

˜

L

(

ψ

)

Ψ

ϕ

of calculus: Given a transformed parameter ψ = φ ( θ ) under a bijective differentiable φ : Θ → Ψ that maps the original parameter space onto another parameter space, the function L ( θ ) becomes ˆ L ( ψ ) := L ( φ -1 ( ψ )) . Note their difference (see Fig. 1): In the former, one works on a single parameter space Θ and a single function L -the map T acts as a symmetry of elements of Θ under the group G . In contrast, the latter assumes two parameter spaces Θ and Ψ which are connected by φ and hence two functions L and ˆ L .

/negationslash

/negationslash

Default

(a) Optimization

θ

(

t

)

ϕ

-

1

(

ψ

(

t

)

)

(b) Maximum Density

p

(

θ

)

ˆ

p

(

ψ

)

θ

∗

=

ϕ

-

1

(

ψ

∗

)

Metric-Aware

θ

∗

=

ϕ

-

1

(

ψ

∗

)

(c) Flatness Measure

L

(

θ

)

ˆ

L

(

ψ

)

ˆ

H

(

ψ

∗

)

= det

H

(

θ

∗

)

ˆ

E

(

ψ

∗

) = det

E

(

θ

∗

)

We focus on the latter. While in this scenario, we have ˆ L ( ψ ) = L ( θ ) whenever ψ = φ ( θ ) , it is wellknown that downstream quantities such as optimization trajectories [55, 71], the Hessian [20, Sec. 5], and probability densities over θ [59, Sec. 5.2.1.4] are not invariant under model reparametrization. These non-invariances are detrimental since an arbitrary reparametrization can affect the studied quantities and thus consistency cannot be guaranteed-parametrization muddles the analysis.

For instance, because of this, (i) one cannot relate Hessian-based sharpness measures with generalization 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 modelthis is one reason why invariant methods like natural gradient are desirable [37, 50, 58, 62], and (iii) under a reparametrization, a prior density might put low probability in a region that corresponds to the high-probability 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 a 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 is represented via a coordinate system, and conserved under a change of coordinates [47]. The parameter space of a neural network, which is by default assumed to be Θ = R d with the Euclidean metric and the Cartesian coordinates, is a Riemannian manifold-model reparametrization is such a change in the coordinate system.

Why, then, can reparametrizations bring up the aforementioned inconsistencies? In this work, we 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 reparametrization is guaranteed if the Riemannian metric is not just implicitly assumed, but made explicit, and it 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 to 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 is 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 a guardrail for existing and future methods to avoid pathologies due to reparametrization.

2 Preliminaries

This section provides background and notation on relevant concepts of neural networks and Riemannian 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 a particular parametrization. E.g., a linear map A is represented by a matrix A , a point z is represented by a tuple of numbers θ , an inner product G ( z ) at z is represented by a matrix G ( θ ) . The same goes for the differential ∇ , ∇ ; and the gradient grad , grad .

2.1 Some notation for neural networks

The following concepts are standard, introduced mostly to clarify notation. Let f : R n × R d → R k with ( x , θ ) ↦→ f ( x ; θ ) be a model with input, output, and parameter spaces R n , R k , R d . Let D := { ( x i , y i ) } i m =1 be a dataset and write X := { x i } i m =1 and Y := { y i } i m =1 . We moreover write f ( X ; θ ) := vec( { f ( x i ; θ ) } i m =1 ) ∈ R mk . The standard way of training f is by finding a point θ ∗ ∈ R d s.t. θ ∗ = arg min θ ∈ R d ∑ m i =1 ℓ ( y i , f ( x i ; θ )) =: arg min θ ∈ R d L ( θ ) , for some loss function ℓ . If we add a weight decay γ/ 2 ∥ θ ∥ 2 2 term to L ( θ ) , the minimization problem has a probabilistic interpretation as maximum a posteriori (MAP) estimation of the posterior density p ( θ | D ) under the likelihood function p ( Y | f ( X ; θ )) ∝ exp( -∑ m i =1 ℓ ( y i , f ( x i ; θ ))) and an isotropic Gaussian prior p ( θ ) = N ( θ | 0 , γ -1 I ) . We denote the MAP loss by L MAP. In our present context, this interpretation is relevant because this probabilistic interpretation implies a probability density, and it is widely known that a probability density transforms nontrivially under reparametrization.

Atextbook way of obtaining θ ∗ is gradient descent (GD): At each time step t , obtain the next estimate θ ( t +1) = θ ( t ) -α ∇ L| θ ( t ) , for some α > 0 . This can be considered as the discretization of the gradient flow ODE ˙ θ = -∇ L| θ . Among the many variants of GD is preconditioned GD which considers a positive-definite matrix field R on R d , yielding ˙ θ = -R ( θ ) -1 ∇ L| θ .

2.2 The geometry of the parameter space

Remark 1. Werestrict 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 in the general local-chart formulation. ■

The parameter space R d is the canonical example of a smooth manifold. We can impose upon this manifold a (global) coordinate chart: a homeomorphism that represents a point with an element of R d . The standard choice is the Cartesian coordinate system θ : R d → Θ ∼ = R d , which we have used in the previous section. That is, the image of θ uniquely represents elements of R d -given a point z ∈ R d , we write θ = θ ( z ) for its coordinate representation under θ .

The choice of a coordinate system is not unique. Any homeomorphism ψ : R d → Ψ ∼ = R d can be used as an alternative coordinate chart. For instance, the polar coordinates can be used to represent points in R d instead. Crucially, the images of any pair of coordinates must be con-

θ

ψ

R

d

z

θ

=

θ

(

z

)

ψ

=

ψ

(

z

)

T

z

R

d

v

w

〈

v, w

〉

=

v

/latticetop

Gw

Θ

∼

=

R

d

ϕ

Ψ

∼

=

R

d

nected through a diffeomorphism (a smooth function with a smooth inverse) φ : Θ → Ψ . Such a map is called a change of coordinates or reparametrization since it acts on the parameter space.

Example 2 (Reparametrization). Reparametrization is ubiquitous in machine learning:

(a) Mean-Field Variational Inference. Let q ( x ; θ ) be a variational approximation, which is often chosen to be N ( x | µ , diag( σ 2 )) , i.e. θ = { µ ∈ R d , σ 2 ∈ R d > 0 } . Common choices of reparametrization of σ 2 include the log-space [52] or softplus [12] parametrization.

(b) Weight normalization. Given a NN f θ , WeightNorm [69] applies the reparametrization ψ = r η where r ∈ R > 0 and η := θ / ∥ θ ∥ ∈ S d -1 . This is akin to the polar coordinates. (Note that we assume θ ∈ R d \ { 0 } since otherwise φ is not a global diffeomorphism.) ■

At each point z ∈ R d , there exists a vector space T z R d called the tangent space at z , consisting of the so-called tangent vectors . An important example of a tangent vector is the gradient vector grad L| z

of L : R d → R at z . The dual space T ∗ z R d of T z R d is referred to as the cotangent space and consists of linear functionals of T z R d , called tangent covectors . An example of a tangent covector is the differential ∇L| z of L 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.

One can take an inner product of tangent vectors by equipping the manifold with a Riemannian metric G , which, at each point in R d is represented by a positive-definite d × d matrix G whose coefficients depend on the choice of coordinates. With the help of the metric, there is an isomorphism between T z R d and T ∗ z R d . In coordinates, it is given by the map T z R d → T ∗ z R d : v ↦→ Gv and its inverse T ∗ z R d → T z R d : ω ↦→ G -1 ω . One important instance of this isomorphism is the fact that grad L| z is represented by G -1 ∇ L in coordinates. The natural gradient is a direct consequence of this fact. In practice, it is standard to assume the Cartesian coordinates on R d , implying G ≡ I . We have thus grad L ≡ ∇ L , which only seems trivial at 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 ; θ )) : θ ∈ R d } and assume that p ( Y | f ( X ; θ )) ↦→ θ is the coordinate chart [19, 55, etc.]. The problem with this construction is that there is no bijection R d → M in general. Indeed, the Jacobian of the map f ( X ; · ) : R d → M is in practice surjective everywhere due to overparametrization [34, 78]. Therefore, one cannot define a proper metric on the parameter space R d 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 a 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, dampingwhich breaks the interpretation of the Fisher on R d as the pullback metric from M -is 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 to 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-independent 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 φ : Θ → Ψ is a reparametrization, with ψ = φ ( θ ) . Coordinate independence is encoded in the following transformation rules, which involve the Jacobian J ( θ ) = ( ∂ ψ i /∂ θ j ) of φ , and its inverse J -1 ( ψ ) = ( J ( θ )) -1 = ( ∂ θ i /∂ ψ j ) -the Jacobian of φ -1 . (Color codes are used for clarity when different objects are present in a single expression later on.)

(a) A function h : Θ → R in θ -coordinates transforms into ˆ h = h ◦ φ -1 in ψ -coordinates.

(b) A tangent vector v in θ -coordinates transforms into J ( θ ) v . In particular, a gradient vector grad L| θ transforms into J ( θ ) grad L| θ .

(c) A tangent covector ω in θ -coordinates transforms into J -1 ( ψ ) ⊤ ω . In particular, given a transformed function L◦ φ -1 : Ψ → R , we have ∇ ( L◦ φ -1 ) | ψ = J -1 ( ψ ) ⊤ ∇ L| φ -1 ( ψ ) , which we recognize as the standard chain rule.

(d) Ametric G ( θ ) becomes J ( θ ) -⊤ G ( θ ) J ( θ ) -1 = J -1 ( ψ ) ⊤ G ( ψ ) J -1 ( ψ ) . In general, this rule applies to any tensor that takes tangent vectors as its arguments, e.g. a bilinear map.

The following examples show how these rules ensure the invariance of geometric objects.

Example 4. Let h , v , ω , G respectively be a function, vector, covector, and metric in θ -coordinates at a point θ = θ ( z ) , and let φ : Θ → Ψ be a reparametrization.

(a) Let ψ := φ ( θ ) and ˆ h := h ◦ φ -1 be θ ∈ Θ and h expressed in ψ -coordinates, respectively. Then, ˆ h ( ψ ) = h ( φ -1 ( ψ )) = h ( θ ) . That is, the actions of h and ˆ h agree in both coordinates and thus they represent the same abstract function on R d .

(b) The action of ω on v in θ -coordinates is given by the product ω ⊤ v . Let ˆ v := J ( θ ) v and ˆ ω := J ( θ ) -⊤ ω be the transformed vector and covectors. Then,

That is, both ω and ˆ ω are the representations of the same linear functional; v and ˆ v are the representations of the same tangent vector.

(c) Let ˆ G := J -⊤ GJ -1 be the transformed metric in the ψ -coordinates. Then,

and thus the transformation rules ensure that inner products are also invariant.

■

Finally, we say that an ODE's dynamics is invariant if the trajectory in parametrization corresponds to the trajectory in another parametrization. Concretely, a trajectory ( θ t ) t in Θ is invariant if under the reparametrization φ : Θ → Ψ , the transformed trajectory ( ψ t ) t is related to ( θ t ) t by θ t = φ -1 ( ψ t ) 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. 2. 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

A growing body of literature connects the flatness of minima found by optimizers to generalization performance [8, 13, 23, 27, 36, 38, 48, 53]. However, as Dinh et al. [20] observed, this association does not have a solid foundation since standard sharpness measures derived from the Hessian of the loss function are not invariant under reparametrization.

From the Riemannian-geometric perspective, the Hessian of a function L (or L MAP or any other twice-differentiable function) on θ -coordinates is represented by a d × d matrix with coefficients H ij ( θ ) := ∂ 2 L ∂ θ i ∂ θ j ( θ ) -∑ d k =1 Γ k ij ( θ ) ∂ L ∂ θ k ( θ ) , for any θ ∈ Θ and i, j = 1 , . . . , d , where Γ k ij is a particular three-dimensional array. While it might seem daunting, when θ is a local minimum of L , the second term is zero since the partial derivatives of L are all zero at θ . Thus, H ( θ ) equals the standard Hessian matrix ( ∂ 2 L / ∂ θ i ∂ θ j ) at θ ∈ arg min L .

Considering the Hessian as a bilinear function that takes two vectors and produces a number, it follows the covariant transformation rule, just like the metric (see Appendix B.1 for a derivation): Let φ : Θ → Ψ and ψ = φ ( θ ) . The Hessian matrix H ( θ ) at a local minimum in θ -coordinates transforms into ˆ H ( ψ ) = J -1 ( ψ ) ⊤ H ( φ -1 ( ψ )) J -1 ( ψ ) in ψ -coordinates-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

and so in general, we do not have the relation (det ˆ H ) = (det H ) ◦ φ -1 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 a linear map/operator on the tangent space. This is done by simply multiplying H with the inverse of the metric G , i.e. E := G -1 H . The determinant of E is thus an invariant quantity. 1 To show this, under φ , the linear map E transforms into

due to the transformation of both G and H . Hence, det E transforms into

in ψ -coordinates. Thus, we have the desired invariant transformation (det ˆ E ) = (det E ) ◦ φ -1 . Note that G is an arbitrary metric-this invariance thus also holds for the Euclidean case where G ≡ I . Note further that the trace and eigenvalues of E are also invariant; see Appendices B.2 and B.3. Finally, see Appendix A for a simple working example.

To obtain invariance in the Hessian-determinant, -trace, and -eigenvalues at a minimum of L , we explicitly write it as E = G -1 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 ˙ θ = -G ( θ ) -1 ∇ L| θ . But NGD is widely presented as invariant under reparametrization, while GD is not [51, 71, etc.]. Is the choice of the metric G the cause? Here we will show from the framework laid out in Section 2.2 that any metric is invariant if its transformation rule is faithfully followed. And thus, the invariance of NGD is not unique. Rather, the Fisher metric used in NGD is part of the family of metrics that transform correctly under autodiff, and thus it is 'automatically' invariant under standard deep learning libraries like PyTorch, TensorFlow, and JAX [1, 24, 66]-see Appendix C.

The underlying assumption of GD is that one works in the Cartesian coordinates and that G ≡ I . For this reason, one can ignore the metric G in ˙ θ = -G ( θ ) -1 ∇ L| θ , and simply write ˙ θ = -∇ L| θ . This is correct but this simplification is exactly what makes GD not invariant under a reparametrization φ : Θ → Ψ where ψ = φ ( θ ) . To see this, notice that while GD transforms ∇ L correctly via the chain rule ˙ ψ = -J -1 ( ψ ) ⊤ ∇ L| φ -1 ( ψ ) , by ignoring the metric I , one would miss the important fact that it must also be transformed into ˆ G ( ψ ) = J -1 ( ψ ) ⊤ J -1 ( ψ ) . It is clear that ˆ G ( ψ ) does not equal I in general. Thus, we cannot ignore this term in the transformed dynamics since it would imply that one uses a different metric-the dynamics are thus different in the θ - and ψ -coordinates. When one explicitly considers the above transformation, one obtains

This dynamics in ψ -coordinates preserves the assumption that the metric is I in θ -coordinates and thus invariant. In contrast, the dynamics ˙ ψ under just the chain rule implicitly changes the metric in θ -coordinates, i.e. from I into J ( θ ) ⊤ J ( θ ) , and thus the trajectories are not invariant.

This discussion can be extended to any metric R : Simply use the transformed metric ˆ R ( ψ ) = J -1 ( ψ ) ⊤ R ( φ -1 ( ψ )) J -1 ( ψ ) , and we obtain the invariant dynamics of any preconditioned GD with preconditioner R given by ˙ ψ = -J ( φ -1 ( ψ )) R ( φ -1 ( ψ )) -1 ∇ L| φ -1 ( ψ ) .

To obtain invariance in optimizers with any metric/preconditioner, explicitly write down the metric even if it is trivial, and perform the proper transformation under reparametrization.

Remark 5. For the discretized dynamics, the larger the step size, the less exact the invariance. For instance, discrete natural gradient update rule is only invariant up to the first-order [50, 71]. This is orthogonal to our work since it is about improving ODE solvers [71]. ■

The consequence is that we need a 'geometric-aware' autodiff library such that the invariant dynamics above is always satisfied. In this case, any preconditioner R yields invariance under any reparametrization, even nonlinear ones. This is in contrast to the current literature, e.g. under standard autodiff, Newton's method is only affine-invariant, and structured approximate NGD methods such as K-FAC are only invariant under a smaller class of reparametrizations [56].

3.3 Invariance in probability densities

Let q Θ ( θ ) be a probability density function (pdf) under the Lebesgue measure d θ on Θ . Under a reparametrization φ : Θ → Ψ with ψ = φ ( θ ) , it transforms into q Ψ ( ψ ) = q Θ ( φ -1 ( ψ )) | det J -1 ( ψ ) | . This transformation rule ensures q Ψ to be a valid pdf under the Lebesgue measure d ψ on Ψ , i.e. ∫ Ψ q Ψ ( ψ ) d ψ = 1 . Notice, in general q Ψ = q Θ ◦ φ -1 due to the changeof-random-variable rule. Hence, pdfs transform differently than standard functions (Example 4a)

̸

and can thus have non-invariant modes [cf. 59, Sec. 5.2.1.4]: Density maximization, such as the optimization of L 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 p J ( θ ) ∝ | det F ( θ ) | 1 2 with normalization constant ∫ Θ | det F ( θ ) | 1 2 d θ . Is the Fisher actually necessary to obtain invariance in pdf maximization?

Here, we show that the same principle of 'being aware of the implicit metric and following proper transformation rules' can be applied. A pdf q Θ ( θ ) can be written as q Θ ( θ ) = q Θ ( θ ) d θ d θ to explicitly show the dependency of the base measure d θ -this is the Lebesgue measure on Θ , which is the natural unit-volume 2 measure in Euclidean space. Given a metric G , the natural volume-measurement device on Θ is the Riemannian volume form dV G which has θ -coordinate representation dV G := | det G ( θ ) | 1 2 d θ . This object takes the role of the Lebesgue measure on a Riemannian manifold: Intuitively, it behaves like the Lebesgue measure but takes the (local) distortion due to the metric G into account. Indeed, when G ≡ I we recover d θ .

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

This is a valid probability density under dV G on Θ since ∫ Θ q G Θ ( θ ) dV G = 1 . This formulation generalizes the standard Lebesgue density. And, it becomes clear that the Jeffreys prior is simply the uniform distribution under dV F ; its density is q F Θ ≡ 1 under dV F .

We can now address the invariance question. Under φ , considering the transformation rules for both q Θ ( θ ) and G , the density (3) thus becomes

This means, q G Θ transforms into q G Θ ◦ φ -1 and is thus invariant since it transforms as standard function-notice the lack of the Jacobian-determinant term here, compared to the standard change-ofdensity rule. In particular, just like standard unconstrained functions, the modes of q G Θ are invariant under reparametrization. Since G is arbitrary, this also holds for G ≡ I , and thus the modes of Lebesgue-densities are invariant, as long as I is transformed correctly (Fig. 2b).

Note that, even if the transformation rule is now different from the one in standard probability theory, q G Ψ is a valid density under the transformed volume form. This is because due to the transformation of the metric G ↦→ ˆ G , we have dV ˆ G = | det G ( φ -1 ( ψ )) | 1 2 | det J -1 ( ψ ) | d ψ . Thus, together with (4), we have ∫ Ψ q G Ψ ( ψ ) dV ˆ G = 1 . This also shows that the | det J -1 ( ψ ) | term in the standard change-of-density formula (i.e. when G ≡ I ) is actually part of the transformation of d θ .

Put another way, the standard change-of-density formula ignores the metric. The resulting density thus must integrate w.r.t. d ψ -essentially assuming a change of geometry, not just a simple reparametrization-and thus q Ψ can have different modes than the original q Θ . While this noninvariance and change of geometry are useful, e.g. for normalizing flows [68], they cause issues when invariance is desirable, such as in Bayesian inference.

To obtain invariance in density maximization, transform the density function under the Riemannian volume form as an unconstrained function. In particular when G ≡ I in Θ , this gives the invariant transformation of a Lebesgue density q Θ , i.e. q Ψ = q Θ ◦ φ -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

at odds with them. Instead, it gives further insights into the inner working of those methods: They are formulated by not following the geometric rules laid out in Section 2.2, and thus in this case, reparametrization implies a change of metric and hence a change of geometry of the parameter space. They can thus be seen as methods for metric learning [35, 77], i.e. finding the 'best' G for the problem at hand, and are compatible with our work since we do not assume a particular metric.

Hessian-based sharpness measures have been extensively used to measure the generalization of neural nets [13, 27, 53]. However, as Dinh et al. [20] pointed out, they are not invariant under reparametrization. Previous work has proposed the Fisher metric to obtain invariant flatness measures [30, 38, 49]. In this work, we have argued that while the Fisher metric is a good choice due to its automatic-invariance property among other statistical benefits [49], any metric is invariant if one follows the proper transformation rules faithfully. That is, the Fisher metric is not necessary if invariance is the only criterion. Similar reasoning about the Fisher metric has been argued in optimization [3, 37, 51, 55, 65, 71] and MAP estimation [21, 32]: the Fisher metric is often used due to its invariance. However, we have discussed that it is not even the unique automatically invariant metric (Appendix C), so the invariance of the Fisher should not be the decisive factor when selecting a metric. By removing invariance as a factor, our work gives practitioners more freedom in selecting a more suitable metric for the problem at hand, beyond the usual Fisher metric.

Finally, the present work is not limited to just the parameter space. For instance, it is desirable for the latent spaces of variational autoencoders [40] to be invariant under reparametrization [2, 25]. The insights of our work can be implemented directly to latent spaces since they are also manifolds [9].

5 Some Applications

We present applications in infinite-width Bayesian NNs, model selection, and optimization, to show some directions where the results presented above, and invariance theory in general, can be useful. They are not exhaustive, but we hope they can be an inspiration and foundation for future research.

5.1 Infinite-width neural networks

Bayesian NNs tend to Gaussian processes (GPs) as their widths go to infinity [57, 60]. It is widely believed that different parametrizations of an NN yield different limiting kernels [75]. For instance, the NTK parametrization (NTP) yields the NTK [29], and the standard parametrization (SP) yields the NNGP kernel [45].

Due to their nomenclatures and the choice of priors, it is tempting to treat the NTP and SP as reparametrization of each other. That is, at each layer, the SP parameter θ transforms into √

-

2

0

2

Θ

-

2

0

2

Ψ

ϕ

-

1

ψ = φ ( θ ) = σ/ h θ where h is the previous layer's width. This induces the same prior on both θ and ψ , i.e. N ( 0 , σ 2 /h I ) , and thus one might guess that invariance should be immediate. However, if they are indeed a reparametrization of the other, the NN f θ must transform as f ψ := f θ ◦ φ -1 ( ψ ) and any downstream quantities, including the NTK, must also be invariant. This is a contradiction since the functional form of the NTP shown in Jacot et al. [29] does not equal f ψ , and the NTK diverges in the SP but not in the NTP [70]. The SP and NTP are thus not just reparametrizations.

Instead, we argue that the SP and NTP are two completely different choices of the NN's architectures, hyperparameters (e.g. learning rate), and priors-see Fig. 4 for intuition and Appendix E.1 for the details. Seen this way, it is thus not surprising that different 'parametrization' yields different limiting behavior. Note that this argument applies to other 'parametrizations' [e.g. 15, 70, 75, 76].

Altogether, our work complements previous work and opens up the avenue for constructing non-trivial infinite-width NNs in a 'Bayesian' way, in the sense that we argue to achieve a desired limiting behavior by varying the model (i.e. architecture, functional form) and the prior (including over hyperparameters) instead of varying the parametrization. This way, one may leverage Bayesian analysis, e.g. model selection via the marginal likelihood [28, 54], for studying infinite-width NNs.

5.2 The Laplace marginal likelihood

The Laplace log-marginal likelihood (LML) of a model M with parameter in R d -useful for Bayesian model comparison [11, Sec. 4.4.1]-is defined by [54]

Let φ : θ MAP ↦→ ψ MAP be a reparametrization of R d . It is of interest to answer whether log Z is invariant under φ , because if it is not, then φ 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 ( M 1 ) ≥ log ˜ Z ( M 2 ) even though originally log Z ( M 1 ) < log Z ( M 2 ) . The question of 'which model is better' thus cannot be answered definitively.

As we have established in Section 3.3, the first and second terms, i.e. L MAP, are invariant under the Riemannian volume measure. At a glance, the remaining terms do not seem to be invariant due to the transformation of the bilinear-Hessian as discussed in Sec. 3.1. Here, we argue that these terms are invariant when one considers the derivation of the Laplace LML, not just the individual terms in (5).

In the Laplace approximation, the last two terms of log Z are the normalization constant of the unnormalized density [18] h ( θ ) := exp ( -1 2 d ⊤ H ( θ MAP ) d ) , where d := ( θ -θ MAP ) . Notice that d is a tangent vector at θ MAP and H ( θ MAP ) is a bilinear form acting on d . Using the transformation rules in Example 4, it is straightforward to show that h ↦→ h ◦ φ -1 -see

Appendix E.3. Since Z is fully determined by h , this suggests that log Z ( M ; θ MAP ) also transforms into log ˜ Z ( M ; φ -1 ( ψ MAP )) , where ψ MAP = φ ( θ MAP ) . This is nothing but the transformation of a standard, unconstrained function on R d , thus log Z is invariant.

We show this numerically in Table 1. We train a network in the Cartesian parametrization and obtain its log Z . Then we reparametrize the net with WeightNorm and naïvely compute log Z again. These log Z 's are different because the WeightNorm introduces more parameters than the Cartesian one, even though the degrees of freedom are the same. Moreover, the Hessian-determinant is not invariant under autodiff. However, when transformed as argued above, log Z is trivially invariant.

5.3 Biases of preconditioned optimizers

The loss landscape depends on the metric assumed in the parameter space through the Hessian operator E = G -1 H . The assumption on G in practice depends on the choice of the optimizer. For instance, under the default assumption of the Cartesian coordinates, using GD implies the Euclidean metric G ≡ I , and ADAM uses the gradient-2nd-moment metric [39].

We argue that explicitly including the metric is not just theoretically the correct thing to do (since it induces invariance, and by Riemannian geometry), but also practically beneficial. Fig. 5 compares measurements of sharpness of minima (Hessian-trace).

MNIST

FMNIST

CIFAR10

CIFAR100

0

10

20

30

Dataset

Log-Trace Hessian

SGD

H

=

E

ADAM

H

ADAM

E

Using the metric-aware Hessian operator E , one can show definitively (i.e. independent of the choice of parametrization) that ADAM tends to obtain much sharper minima than SGD. The benefits of using the Hessian operator have also been confirmed in previous work. Cohen et al. [16] argue that when analyzing the optimization dynamic of an adaptive/preconditioned GD algorithm, one should take the preconditioner into account when measuring Hessian-based sharpness measures. From this, they demonstrate a sharpness-evolution behavior known in vanilla GD, allowing a comparison between vanilla and adaptive GD.

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 a Riemannian manifold, albeit often a trivial one. This raises the question of why one should observe non-invariance in neural nets, whereas, by definition, Riemannian manifolds are invariant under a change of coordinates. As we showed, this discrepancy only arises if the transformations used in the construction of a neural net along with an optimizer ignore the implicitly assumed metric. By acknowledging the metric and using the transformation rules associated with geometric objects, invariance and invariance under reparametrization are then guaranteed. Our results provide a geometric solution towards full invariance of neural nets-it is compatible with and complementary to other works that focus on invariance under group-action symmetries, both in the weight and input spaces of neural networks.

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: 01IS18039A); the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) in the frame of the priority programme SPP 2298 'Theoretical Foundations of Deep Learning'-FKZ HE 7114/5-1; and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. AK & FD are grateful to the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. AK, FD, and PH are also grateful to Frederik Künstner and Runa Eschenhagen for fruitful discussions.

References

Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. TensorFlow: a system for large-scale machine learning. In OSDI , 2016.

Francisco E Acosta, Sophia Sanborn, Khanh Dao Duc, Manu Madhav, and Nina Miolane. Quantifying local extrinsic curvature in neural manifolds. arXiv preprint arXiv:2212.10414 , 2022.

Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018 , 2020.

Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Hernández-Lobato. Adapting the linearised Laplace model evidence for modern deep learning. In ICML , 2022.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In NeurIPS , 2019.

Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In ICLR , 2019.

Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09745 , 2022.

Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: on the curvature of deep generative models. In ICLR , 2018.

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.

Christopher M. Bishop. Pattern Recognition and Machine Learning . Springer, 2006.

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In ICML , 2015.

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.

An Mei Chen, Haw-minn Lu, and Robert Hecht-Nielsen. On the geometry of feedforward neural network error surfaces. Neural computation , 5(6), 1993.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In NeurIPS , 2018.

Jeremy M 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 arXiv:2207.14484 , 2022.

Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral CNN. In ICML , 2019.

Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless Bayesian deep learning. In NeurIPS , 2021.

Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, and Koray Kavukcuoglu. Natural neural networks. In NIPS , 2015.

Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In ICML , 2017.

Pierre Druilhet and Jean-Michel Marin. Invariant HPD credible sets and MAP estimators. Bayesian Analysis , 2(4), 2007.

Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. NeurIPS , 2018.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In ICLR , 2021.

Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. SysML , 4(9), 2018.

Søren Hauberg. Only Bayes should learn a manifold (on the estimation of differential geometric structure from data). arXiv preprint arXiv:1806.04994 , 2018.

Dan Hendrycks, Mantas Mazeika, and Thomas Dietterich. Deep anomaly detection with outlier exposure. In ICLR , 2019.

Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In NIPS , 1994.

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.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In NIPS , 2018.

Cheongjae Jang, Sungyoon Lee, Frank C. Park, and Yung-Kyun Noh. A reparametrization-invariant sharpness measure based on information geometry. In NeurIPS , 2022.

Harold Jeffreys. An invariant form for the prior probability in estimation problems. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences , 186(1007), 1946.

Ian H Jermyn. Invariant Bayesian estimation on manifolds. The Annals of Statistics , 33(2), 2005.

Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In ICLR , 2020.

Ryo Karakida and Kazuki Osawa. Understanding approximate Fisher information for fast convergence of natural gradient descent in wide neural networks. In NeurIPS , 2020.

Mahmut Kaya and Hasan ¸ Sakir Bilge. Deep metric learning: A survey. Symmetry , 11(9), 2019.

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In ICLR , 2017.

Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In ICML , 2022.

SungYub Kim, Sihwan Park, Kyung-Su Kim, and Eunho Yang. Scale-invariant Bayesian neural networks with connectivity tangent kernel. In ICLR , 2023.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In ICLR , 2014.

Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. An infinite-feature extension for Bayesian ReLU nets that fixes their asymptotic overconfidence. In NeurIPS , 2021.

Vˇ era K˚ urková and Paul C Kainen. Functionally equivalent feedforward neural networks. Neural Computation , 6(3), 1994.

Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In ICLR , 2021.

Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical Fisher approximation for natural gradient descent. In NeurIPS , 2019.

Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In ICLR , 2018.

John M Lee. Introduction to smooth manifolds . Springer, 2013.

John M Lee. Introduction to Riemannian manifolds . Springer, 2018.

Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? a mathematical framework. In NeurIPS , 2021.

Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-Rao metric, geometry, and complexity of neural networks. In AISTATS , 2019.

Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Structured second-order methods via natural gradient descent. In ICML Workshop on beyond first-order methods in machine learning systems , 2021.

Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Tractable structured natural gradient descent using local parameterizations. In ICML , 2021.

Christos Louizos and Max Welling. Structured and efficient variational deep learning with matrix Gaussian posteriors. In ICML , 2016.

Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.07085 , 2022.

David JC MacKay. The evidence framework applied to classification networks. Neural Computation , 4(5), 1992.

James Martens. New insights and perspectives on the natural gradient method. JMLR , 21(146), 2020.

James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417, 2015.

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.

Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning. In ICLR , 2021.

Kevin P. Murphy. Machine Learning: A Probabilistic Perspective . The MIT Press, 2012.

Radford M Neal. Bayesian Learning for Neural Networks . PhD thesis, University of Toronto, 1995.

Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-SGD: Path-normalized optimization in deep neural networks. In NIPS , 2015.

Kazuki Osawa, Shigang Li, and Torsten Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv preprint arXiv:2211.14133 , 2022.

George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR , 22(57), 2021.

Omiros Papaspiliopoulos, Gareth O Roberts, and Martin Sköld. A general framework for the parametrization of hierarchical models. Statistical Science , 2007.

Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. In ICLR , 2014.

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, high-performance deep learning library. In NeurIPS , 2019.

Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurIPS , 2021.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML , 2015.

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In NIPS , 2016.

Jascha Sohl-Dickstein, Roman Novak, Samuel S Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with a standard parameterization. arXiv preprint arXiv:2001.07301 , 2020.

Yang Song, Jiaming Song, and Stefano Ermon. Accelerating natural gradient with higher-order invariance. In ICML , 2018.

Michalis K Titsias. Learning model reparametrizations: Implicit variational inference by fitting MCMC distributions. arXiv preprint arXiv:1708.01529 , 2017.

Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC-Bayesian analysis. In ICML , 2020.

Jesse van Oostrum, Johannes Müller, and Nihat Ay. Invariance properties of the natural gradient in overparametrised systems. Information Geometry , 2022.

Greg Yang and Edward J Hu. Tensor programs IV: Feature learning in infinite-width neural networks. In ICML , 2021.

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. In NeurIPS , 2022.

Liu Yang and Rong Jin. Distance metric learning: A comprehensive survey. Michigan State Universiy , 2 (2), 2006.

Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for over-parameterized neural networks. In NeurIPS , 2019.

Appendix A A Simple Working Example

Let Θ := R d and let L : R d → R defined by L ( θ ) := ∑ d i =1 ( θ i ) 2 be a twice-differentiable function on Θ . Suppose φ : Θ → Ψ defined by φ ( θ ) = 3 θ =: ψ be 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 J ( θ ) = diag(3 , . . . , 3) ∈ R d × d and its inverse is J -1 ( ψ ) = diag( 1 / 3 , . . . , 1 / 3 ) . The autodiff Hessian matrix is given by H ( θ ) ≡ 2 I ∈ R d × d . When we transform θ → ψ , the function L ( θ ) becomes ˆ L ( ψ ) := ∑ d i =1 ( 1 / 3 ψ i ) 2 . The autodiff will take into account this change and thus will compute the Hessian ˆ H of ˆ L as

Notice that, det H ( θ ) ≡ 2 d , meanwhile, we have det ˆ H ( ψ ) ≡ ( 2 / 9 ) d . Therefore, the Hessian determinant computed by an autodiff system is not invariant under reparametrization.

But if instead we 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 ( θ ) = G ( θ ) -1 H ( θ ) = I -1 (2 I ) , and perform the correct transformation on both identity matrices, we obtain

Thus, the determinant of the transformed 'metric-conditioned Hessian' equals

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 as a lower index, we sum them over. For example: z = x i y i means z = ∑ i x i y i and B k j = A k ij x i means B k j = ∑ k A k ij x i , etc. Specifically for partial derivatives, the index of the denominator is always treated as a lower index. ■

B.1 The Riemannian Hessian Under Reparametrization

Let L : M → R be a function on a Riemannian manifold M with metric G . The Riemannian Hessian Hess L of L is defined in coordinates θ by

where Γ k ij is the connection coefficient.

Under a change of coordinates φ : θ ↦→ ψ , we have ˜ L = L◦ φ -1 and

where m , n , o are just dummy indices-present to express summations. Note that the transformation rule for Γ k ij implies that it is not a tensor-to be a tensor, there must not be the second term in the previous formula.

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 :

In the matrix form, we can write the above as ˜ H = J -⊤ HJ -1 , where J is the Jacobian of φ . Thus, the Riemannian Hessian at any θ (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 Γ k ij is explicitly considered. In particular, the Euclidean Hessian does not follow this tensorial transformation under autodiff due to the fact that (i) Γ k ij = 0 for any i , j , k and thus dropped from the equation, and (ii) autodiff is not designed to handle advanced geometric objects like Γ k ij .

B.2 Hessian-Trace Under Reparametrization

Let L : R d → R be a function of R d under the Cartesian coordinates and G a Riemannian metric. The Riemannian trace of the Hessian matrix H of L is defined by [47]:

That is, it is defined as the standard trace of the Hessian operator E .

Let φ : θ ↦→ ψ be a reparametrization on R d . Then, using (2) and the property tr( AB ) = tr( BA ) twice, the Riemannian trace of the Hessian transforms into

Since ψ = φ ( θ ) , we have that (tr ˜ G ˜ H )( ψ ) = (tr G H )( θ ) for any given θ . 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 an eigenvalue of the linear map E ( θ ) = G ( θ ) -1 H ( θ ) on the tangent space at z ∈ R d that is represented θ if E ( θ ) v = λ v for an eigenvector v ∈ T z R d . We shall show that ˜ λ , the eigenvalue under under the reparametrization φ : θ ↦→ ψ , equals the original eigenvalue λ .

Using the transformation rule of E ( θ ) in (2) and the transformation rule of tangent vectors in (4), along with the relation ( J -1 ( ψ )) -1 = J ( φ -1 ( ψ )) , we get

Since φ -1 ( ψ ) = θ , we arrive back at the original equation before the reparametrization and thus we conclude that ˜ λ = λ .

Appendix C The Invariance of the Fisher Metric

We have seen that any metric on the parameter space yields 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. It 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 inv(equi)variances can thus be obtained effortlessly.

Martens [55, Sec. 12] mentions the following family of curvature matrices satisfying this transformation behavior

where J ( θ ; · ) is the network's Jacobian ∂f ( · ; θ ) /∂ θ , with an arbitrary data distribution D and an invertible matrix A that transforms like A ( φ -1 ( ψ )) under φ : θ ↦→ ψ . Under a reparametrization φ , an autodiff library will compute

given L ◦ φ -1 , due to the elementary transformation rule of J ( θ ; · ) , 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].

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 to 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) as an example. Its ψ -representation computed by the autodiff library will only match the correct transformation for element-wise reparametrizations, whose Jacobian J ( θ ) 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 is theoretically in(equi)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 as its statistical efficiency [3] and guarantees in optimization [34, 78]. Automatic transformation is but a cherry on top.

It is interesting for future work to extend autodiff libraries to take into account the invariant transformation rules of geometric objects. By doing so, any metric-not just the Fisher metric-will yield invariance automatically .

Appendix D General Manifold: Local Coordinate Charts

In the main text, we have exclusively used global coordinate charts ( R d , θ ) and ( R d , ψ ) -these charts cover the entire R d and θ , ψ are homeomorphisms on R d . However, there exist other coordinate systems that are not global, i.e. they are constructed using multiple charts { ( U i ⊆ R d , θ i : U i → Θ i ⊆ R d ) } i s.t. R d = ∪ i U i and each θ i is a diffeomorphism on U i .

̸

If { ( U i , θ i ) } i and { ( V j , ψ j ) } j be two local coordinate systems of R d . Let ( U i , θ i ) and ( V j , ψ j ) be an arbitrary pair of charts from the above collections where U i ∩ V j = ∅ . In this setting, reparametrization between these two charts amounts to the condition that the transition map φ between θ ( U i ∩ V j ) and ψ ( U i ∩ V j ) 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.

Note that this definition is a generalization of global coordinate charts. In this case, there is only a single chart for each coordinate system, i.e. ( U, θ ) and V, ψ , and they trivially overlap since U = R d and V = R d . Moreover, there is only a single diffeomorphism of concern, as shown in Fig. 3.

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 a pair of two charts, instead of the whole R d .

Appendix E Details on Applications

E.1 Infinite-Width Neural Networks

R

d

U

V

θ

ψ

ϕ

Note We use a 2-layer NN without bias for clarity. The extension to deep networks follows directly by induction-see e.g. [6, 45]. We also use the same prior variance σ 2 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: z ∼ N ( µ , Σ ) = ⇒ Az ∼ N ( Aµ , A Σ A ⊤ ) .

Neural-network Gaussian process (NNGP)

Let f ( x ) : R d → R defined by f ( x ) := w ⊤ ϕ ( Wx ) be a real-valued, 2-layer NN with parameter θ := { W ∈ R h × n , w ∈ R h } and component-wise nonlinearity ϕ . Note that we assume x is i.i.d. Let vec( W ) ∼ N ( 0 , σ 2 I ) and w ∼ N ( 0 , σ 2 h I ) be priors over the weights. This parametrization of the weights and the priors is called the standard parametrization (SP) [45, 60].

Step (a) Given a particular preactivation value z , we have a linear model f ( x ) = w ⊤ ϕ ( z ) . We can view this as a Gaussian process with mean and covariance

Taking the limit as h →∞ , the mean stays trivially zero and we have by the law of large numbers:

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 we will obtain now.

Step (b) Notice that each z i = W ⊤ i x where W i is the i -th row of W . This is a linear model and thus, z is distributed as a GP with mean and covariance

So, z i ∼ GP (0 , K z ) Since the prior over W is i.i.d., this holds for all i = 1 , . . . , h . We can thus now compute the expectation in K ( x , x ′ ) : it is done w.r.t. this GP over z i .

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 is that as 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 the prior, and not on particular values (e.g. learned) of the previous hidden units. In this sense, an infinite-width L -layer NN is simply a linear model with a constant feature extractor induced by the network's first L -1 layers that are fixed at initialization. Note that the kernel K over the function output is called the NNGP kernel [45].

Neural tangent kernel (NTK)

Let us transform w into v := σ √ h w and W into V := σ √ h W and define the prior to be w ∼ N ( 0 , I ) and vec( W ) ∼ N ( 0 , I ) . Then, we define the transformed network as ˆ f ( x ) := v ⊤ ϕ ( V x ) = σ √ h w ⊤ ϕ ( σ √ h Wx ) with parameter ψ := { V ∈ R h × n , v ∈ R h } . This is called the NTK parametrization (NTP) [29]. We will see below that even though v , V have the same prior distributions as w , W in the SP, they have different behavior in terms of the NTK.

As before, let us assume a particular preactivation value z . The empirical NTK (i.e. finite-width NTK) on the last layer is defined by:

The (asymptotic) NTK is obtained by taking the limit of h →∞ :

which coincides the NNGP kernel K . Crucially, this is obtained via a backward propagation from the output of the network and thus the linear-Gaussian property we have used to derive the NNGP via forward propagation does not apply. 3 This is why the scaling of σ h is required in the NTP. That is, using the SP, the empirical NTK is not scaled by σ 2 h 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 a 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 are reparametrization of each other, they should be invariant, as we have discussed in the main text. Instead, we 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 a foundation for future work in this field: To obtain a desired limiting behavior, study the network architecture and its prior, instead of the parametrization.

Suppose ψ in the NTP is a reparametrization of θ in the SP. Then the function φ : θ ↦→ ψ defined by θ ↦→ σ √ h θ is obviously the smooth reparametrization with an invertible (diagonal) Jacobian J ( θ ) = σ √ h I . In this case, the network in the NTP must be defined by ˜ f = f ◦ φ -1 , where f is the SP-network, by Example 4. That is, with some abuse of notation,

This is different from the definition of the NTP-network ˆ f ( x ) = v ⊤ ϕ ( V x ) . So, obviously, the NTP is not the reparametrization of the SP. Therefore, a clearer way of thinking about the NTP and SP is to treat them as two separate network functions (i.e. two separate architectures)-the scaling factor σ √ h should be thought of as part of the layer's functional form instead of as part of the parameter. In particular, they are not two representations of a 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 on Section 3.2. First, notice that the differential ∇ w f ( x ) transforms into J -1 ( v ) ⊤ ∇ ˜ f ( x ) | φ -1 ( v ) for any x ∈ R n . Next, notice that the Euclidean metric transforms into ˜ G ( v ) := J -1 ( v ) ⊤ J -1 ( v ) . So the gradient transforms into J ( φ -1 ( v )) ∇ f ( x ) | φ -1 ( v ) . 4 Therefore,

the empirical NTK kernel ˆ K Ψ for ˜ f is given by

Thus, the empirical NTK is invariant and the asymptotic NTK also is. Therefore, we still have a problem with the NTK blow-up in this parametrization. This reinforces the fact that the difference between the SP and NTP is not because of parametrization.

Additionally, let us now inspect the priors in the SP and NTP. In the SP, the prior is N ( θ | 0 , σ 2 / h I ) . Therefore, so that we have the same prior in both Θ and Ψ , the prior of ψ = φ ( θ ) must be N ( ψ | 0 , I ) . This is obviously not the case since we have N ( ψ | 0 , σ 2 / h I ) because the NTP explicitly defines N ( θ | 0 , I ) as the prior of θ . Thus, not only that the SP and NTP assume two different architectures, but they also assume two different prior altogether. It is thus not surprising that the distribution over their network outputs f ( x ) , ˆ f ( x ) are different, both in the finite- and infinite-width regimes.

Implication In his seminal work, Neal [60] concluded that the fact that infinite-width NNs are 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 et 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.

E.2 Biases of Preconditioned Optimizers

For MNIST and FMNIST, the network is LeNet. Meanwhile, we use the WiderResNet-16-4 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 are in Table 2. Additionally, in Table 3, we discuss the effect of reparametrization to sharpness on ADAM and SGD.

E.3 Laplace Marginal Likelihood

Let θ MAP be a MAP estimate in an arbitrary θ -coordinates of R d , obtained by minimizing the MAP loss L MAP. Let log h = -L MAP-note that L MAP itself is a log-density function. The Laplace marginal likelihood [18, 28, 54] is obtained by performing a second-order Taylor's expansion:

where H ( θ MAP ) is the Hessian matrix of L MAP at θ MAP. Then, by exponentiating and taking the integral over θ , we have

Since the integral is the normalization constant of the Gaussian N ( θ | θ MAP , H ( θ MAP )) , we obtain the Laplace log-marginal likelihood (LML):

Notice that H ( θ ) is a bilinear form, acting on the tangent vector d ( θ MAP ) := ( θ -θ MAP ) . Under a reparametrization φ : θ ↦→ ψ with ψ MAP = φ ( θ MAP ) , the term inside the exponent in (15) transforms into

due to the transformations of the tangent vector and the bilinear-Hessian. This simplifies into

which always equals the original integrand in (15). Thus, the integral evaluates to the same value. Hence, the last two terms of log Z ( θ MAP ) transform into -d 2 log(2 π ) + log det H ( φ -1 ( ψ MAP )) in ψ -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 x + ϵ , where ϵ ∼ N (0 , 0 . 3 2 ) . The network is a 1-hidden layer TanH network trained for 1000 epochs.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nA 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nAs the paper does not propose a new method but comments on the mathematical formulation of theories, limitations do not really apply.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- The discussion on flatness-based generalization measures does not address existing generalization bounds for reparameterization-invariant flatness measures (see questions).\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe paper is clear about the assumptions and limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: 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.

STRENGTHS:
1. Covers applications of this type of thinking in several areas of machine learning.
2. Good, careful exposition of geometric concepts and calculations.
3. The note on the utility of non-invariant reparameterization for normalizing flows and optimization is interesting.
4. 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).

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].

[Du et al. 2018] Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced. NeurIPS 2018  
[Ding et al. 2023] Flat minima generalize for low-rank matrix recovery. arXiv 2023.

QUESTIONS:
n/a

LIMITATIONS:
Discussion of limitations on Page 2.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUnder 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\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors state limitations at the end of the introduction section. There are no potential negative societal impacts of the work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}