Abide by the Law and Follow the Flow: Conservation Laws for Gradient Flows
Sibylle Marcotte, Rémi Gribonval, Gabriel Peyré
TL;DR
This work develops a rigorous, constructive framework to identify conservation laws in gradient-flow dynamics of over-parameterized neural networks. By introducing a local reparameterization φ and analyzing the generated Lie algebra Lie(W_φ), the authors reduce the problem to finite-dimensional linear-algebraic computations and Frobenius-type criteria to count independent conserved quantities. They provide algorithms (implemented in SageMath) to compute polynomial conservation laws and, under mild loss assumptions, the maximal number of independent (not necessarily polynomial) laws; they validate completeness for linear and shallow ReLU networks and precisely characterize the matrix-factorization case (q=2). The results offer a principled way to understand initialization-preserved properties (implicit bias) and to recast certain high-dimensional gradient flows as low-dimensional Riemannian dynamics, with potential to guide initialization choices in large models and to inform theoretical analyses of generalization.
Abstract
Understanding the geometric properties of gradient descent dynamics is a key ingredient in deciphering the recent success of very large machine learning models. A striking observation is that trained over-parameterized models retain some properties of the optimization initialization. This "implicit bias" is believed to be responsible for some favorable properties of the trained models and could explain their good generalization properties. The purpose of this article is threefold. First, we rigorously expose the definition and basic properties of "conservation laws", that define quantities conserved during gradient flows of a given model (e.g. of a ReLU network with a given architecture) with any training data and any loss. Then we explain how to find the maximal number of independent conservation laws by performing finite-dimensional algebraic manipulations on the Lie algebra generated by the Jacobian of the model. Finally, we provide algorithms to: a) compute a family of polynomial laws; b) compute the maximal number of (not necessarily polynomial) independent conservation laws. We provide showcase examples that we fully work out theoretically. Besides, applying the two algorithms confirms for a number of ReLU network architectures that all known laws are recovered by the algorithm, and that there are no other independent laws. Such computational tools pave the way to understanding desirable properties of optimization initialization in large machine learning models.