The Geometry of Neural Nets' Parameter Spaces Under Reparametrization

Published: 21 Sept 2023, Last Modified: 19 Dec 2023NeurIPS 2023 spotlightEveryoneRevisionsBibTeX
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.
Supplementary Material: pdf
Submission Number: 940