The Geometry of Neural Nets' Parameter Spaces Under Reparametrization

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.

Paper

Similar papers

Peer review

Reviewer togv7/10 · confidence 4/52023-06-29

Summary

The paper discusses reparametrizations of parameter spaces and the implied transformation rules for quantities like gradients, Hessians or probability densities. Parameter 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. The transformation laws of such objects are well known in differential geometry, however, as the authors argue, they are in deep learning applications often disregarded. In particular, the metric tensor on $M=\mathbb{R}^d$ is often taken to be $G=\mathbb{I}$ and therefore dropped from the equations. Its 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. The contribution of the paper is to point out these disregarded transformation rules and to discuss how the quantities should actually transform. Section 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. The third section discusses more specific quantities of interest in deep learning. Firstly, it considers the Hessian matrix determinant as measure of flatness of parameter landscapes. The authors argue that this measure is not well-suited since it depends on the choice of coordinates. They propose to use the determinant of $G^{-1}H$ instead, which is invariant under reparametrizations. Secondly, it considers loss gradients $\mathrm{grad}\mathcal{L} := G^{-1}\nabla\mathcal{L}$. On $M=\mathbb{R}^d$, the trivial metric $G=\mathbb{I}$ is usually dropped, which leads again to inconsistent transformations of the gradient. Lastly, Section 3.3 discusses probability density functions (pdfs). When being expressed relative to the Lebesgue measure on $M=\mathbb{R}^d$, the pdfs transform with the well known Jacobian determinant factor. However, as the density may be stretched out or condensed in this procedure, the densities' mode may not be preserved. It 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. After discussing related work in section 4, the fifth section considers some applications, arguing in particular that NTK and standard parametrizations or neural networks are not just reparametrizations, but geometrically truly different models (section 5.1) and that the Laplace marginal likelihood is invariant under reparametrization (section 5.2). It investigates furthermore the effect of using the Hessian or $G^{-1}H$ in (preconditioned) optimizers like ADAM.

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. Its main contribution lies in pointing out this issue and discussing how the equations can be fixed. While 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. It 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. I 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. These weaknesses should be easy to fix - more details follow in the next paragraphs. Firstly, the considered coordinate charts and hence reparametrizations are _global_ homeomorphisms. This is in principle possible, but excludes practically relevant choices like polar or spherical coordinates, which are not global homeomorphisms. Polar coordinates are, in fact, used as an example right after saying that charts should be global homeomorphisms. The 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. Specifically, there should be charts $\theta: U^\theta \to \theta(U^\theta) \subseteq \mathbb{R}^d$ and $\psi: U^\psi \to \psi(U^\psi) \subseteq \mathbb{R}^d$ on domains $U^\theta\subseteq M=\mathbb{R}^d$ and $U^\psi\subseteq M=\mathbb{R}^d$ with transition maps $\varphi: \theta(U^\theta\cap U^\psi) \to \psi(U^\theta\cap U^\psi)$. Note that the main results of the paper will still hold in this more general setting. Secondly, 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. It is introduced at the end of the second section, assuming parameter space reparametrizations $\varphi:\Theta\to\Psi$ and functions which simultaneously seem to satisfy $F: \Theta\to\Theta$ and $F: \Psi\to\Psi$. This would in principle require $\Theta=\Psi$, while actually only $\Theta\cong\mathbb{R}^d\cong\Psi$ is demanded initially. Furthermore, this does not work in the clean formulation with local charts suggested above, since then $\varphi: \theta(U^\theta\cap U^\psi) \to \psi(U^\theta\cap U^\psi)$ and the upper and lower arrows would be $F^\theta: \theta(U^\theta\cap U^\psi) \to \theta(U^\theta\cap U^\psi)$ and $F^\psi: \psi(U^\theta\cap U^\psi) \to \psi(U^\theta\cap U^\psi)$, respectively. The 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. It 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. The concept of equivariance is subsequently used in a rather unspecific way in the paper. Section 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. Furthermore, equivariance is a property of _functions_, and it is not entirely clear to me how one should interpret the gradient descent algorithm as such. Section 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))$. However, 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). The issue of this dubious notion of "equivariance" is easily fixed by removing it from the paper and referring to it as usual as coordinate independence or covariance.

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. I 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? Further minor suggestions and corrections beyond what I wrote above follow. These should be very easy to fix. Line 50 states that _"Intrinsic ... means that objects ... must be independent of ... coordinate system"_. This concept is called coordinate independence or covariance. "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). Charts are defined as homeomorphisms, but the paper considers smooth manifolds. The smooth structure is actually only respected (and defined via) smooth charts, i.e. diffeomorphisms. It 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}$. That the choice of coordinate system is not unique is somewhat tricky, as there exists the _canonical_ global chart $\mathrm{id}_{\mathbb{R}^d}$ mentioned by the authors. If the manifold was just taken as Euclidean space, i.e. with metric but without the canonical coordinates of $\mathbb{R}^d$, global 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)$. This raises the question why we are considering general diffeomorphisms in the first place? Is the metric just an arbitrary choice (which should of course still be respected, requiring coordinate independence)? It would be great if the authors could discuss this point. Line 120 mentions that "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", however, also abstract coordinate free vectors are part of linear algebra. I would just write that vectors and covectors are in coordinates represented by numerical coefficient vectors, i.e. tuples of numbers. Line 123: The metric is not only positive definite, but also symmetric. I am confused about the "surjective everywhere" in line 134. Is the problem with bijectivity not that the mapping from parameters to models is in general non-injective? It 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

As the paper does not propose a new method but comments on the mathematical formulation of theories, limitations do not really apply.

Reviewer LwSr8/10 · confidence 3/52023-07-05

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.

Strengths

- Viewing reparameterizations from the perspective of Riemannian geometry brings much needed clarity to the discussion. - The paper is excellently written and insightful, a pleasure to read. - 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).

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. - The relation between flatness and generalization in light of reparameterizations has been established theoretically in [1]. - Is it possible to interpret relative flatness [1] as an invariant transformation? That is, is it possible that $G^-1(\theta) = ||\theta||^2_2$? [1] Petzka, Henning, et al. "Relative flatness and generalization." Advances in neural information processing systems 34 (2021): 18420-18432.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The paper is clear about the assumptions and limitations.

Reviewer Poha6/10 · confidence 1/52023-07-07

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Discussion of limitations on Page 2.

Reviewer ecDa6/10 · confidence 3/52023-07-07

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.

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. I 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. Some 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. The 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? [1] Kunin, Daniel, et al. "Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics." International Conference on Learning Representations. 2020. (Appendix A)

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? - 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors state limitations at the end of the introduction section. There are no potential negative societal impacts of the work.

Reviewer togv2023-08-12

The authors addressed the issues raised in "weaknesses" by promising to rewrite the paper accordingly. I don't have any remaining questions. Please make sure to include a discussion of the relation to Cohen et al. and Weiler et al. in the paper. Note that they also do not require symmetries of the manifold itself, but just consider "gauge symmetries" in the parametrization/coordinates of tangent spaces. Equivariance under symmetries of the manifold may be induced by this coordinate independence.

Reviewer LwSr2023-08-14

Answer to authors

I thank the authors for their reply. Both the reply and the other reviews keep me convinced that this is a good paper. I keep my score.

Reviewer ecDa2023-08-15

Thank you for the response. I now have a better understanding of the significance of invariance under reparametrization. I also appreciate the clarification on the difference between invariance under reparametrization and symmetry. I have increased my score accordingly.

Reviewer Poha2023-08-17

We thank the author for their rebuttal. I think certain readers would appreciate the addition of discussion for these two papers. I have no further questions.

Program Chairsdecision2023-09-21

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC