The Geometry of Neural Nets' Parameter Spaces Under Reparametrization

TL;DR

This paper investigates how neural network parameter spaces behave under reparametrization using Riemannian geometry, emphasizing the importance of explicit metrics for consistent analysis of flatness, optimization, and probability densities.

Contribution

It introduces a geometric framework that accounts for reparametrization invariance by explicitly representing the metric, clarifying the relationship between flatness and generalization.

Findings

Reparametrization invariance is an inherent property when using explicit metrics.

Implicit assumptions of identity metrics lead to inconsistencies under reparametrization.

The framework impacts measurements of flatness, optimization trajectories, and density estimation.

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.