Title: Abide by the law and follow the flow: conservation laws for gradient flows
Venue: NeurIPS 2023
Accepted: yes
Keywords: Implicit bias

conservation laws

gradient flow

linear neural network

matrix factorization
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.

Abide by the Law and Follow the Flow: Conservation Laws for Gradient Flows

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.

1 Introduction

State-of-the-art approaches in machine learning rely on the conjunction of gradient-based optimization with vastly 'over-parameterized' architectures. A large body of empirical [30] and theoretical [5] works suggest that, despite the ability of these models to almost interpolate the input data, they are still able to generalize well. Analyzing the training dynamics of these models is thus crucial to gain a better understanding of this phenomenon. Of particular interest is to understand what properties of the initialization are preserved during the dynamics, which is often loosely referred to as being an 'implicit bias' of the training algorithm. The goal of this article is to make this statement precise, by properly defining maximal sets of such 'conservation laws', by linking these quantities to algebraic computations (namely a Lie algebra) associated with the model parameterization (in our framework, this parameterization is embodied by a re-parameterization mapping φ ), and finally by exhibiting algorithms to implement these computations in SageMath [29].

Over-parameterized model Modern machine learning practitioners and researchers have found that over-parameterized neural networks (with more parameters than training data points), which are often trained until perfect interpolation, have impressive generalization properties [30, 5]. This performance seemingly contradicts classical learning theory [25], and a large part of the theoretical deep learning literature aims at explaining this puzzle. The choice of the optimization algorithm is crucial to the model generalization performance [10, 21, 14], thus inducing an implicit bias .

Implicit bias The terminology 'implicit bias' informally refers to properties of trained models which are induced by the optimization procedure, typically some form of regularization [22]. For gradient descent, in simple cases such as scalar linear neural networks or two-layer networks with a single neuron, it is actually possible to compute in closed form the implicit bias, which induces some approximate or exact sparsity regularization [10]. Another interesting case is logistic classification on separable data, where the implicit bias selects the max-margin classifier both for linear models [26] and for two-layer neural networks in the mean-field limit [8]. A key hypothesis to explicit the implicit bias is often that the Riemannian metric associated to the over-parameterization is either of Hessian type [10, 3], or can be somehow converted to be of Hessian type [3], which is seemingly always a very strong constraint. For example, even for simple two-layer linear models (i.e., matrix factorization) with more than a single hidden neuron, the Hessian type assumption does not hold, and no closed form is known for the implicit bias [11]. The work of [17] gives conditions on the over-parameterization for this to be possible (for instance certain Lie brackets should vanish). These conditions are (as could be expected) stronger than those required to apply Frobenius theory, as we do in the present work to retrieve conservation laws.

Conservation laws Finding functions conserved during gradient flow optimization of neural networks (a continuous limit of gradient descent often used to model the optimization dynamics) is particularly useful to better understand the flow behavior. One can see conservation laws as a 'weak' form of implicit bias: to explain, among a possibly infinite set of minimizers, which properties (e.g. in terms of sparsity, low-rank, etc.) are being favored by the dynamic. If there are enough conservation laws, one has an exact description of the dynamic (see Section 3.4), and in some cases, one can even determine explicitly the implicit bias. Otherwise, one can still predict what properties of the initialization are retained at convergence, and possibly leverage this knowledge. For example, in the case of linear neural networks, certain balancedness properties are satisfied and provide a class of conserved functions [24, 9, 1, 2, 15, 28, 19]. These conservation laws enable for instance to prove the global convergence of the gradient flow [4] under some assumptions. We detail these laws in Proposition 4.1. A subset of these 'balancedness' laws still holds in the case of a ReLU activation [9], which reflects the rescaling invariance of these networks (see Section 4 for more details). More generally such conservation laws bear connections with the invariances of the model [16]: to each 1-parameter group of transformation preserving the loss, one can associate a conserved quantity, which is in some sense analogous to Noether's theorem [23, 12]. Similar reasoning is used by [31] to show the influence of initialization on convergence and generalization performance of the neural network. Our work is somehow complementary to this line of research: instead of assuming a priori known symmetries, we directly analyze the model and give access to conservation laws using algebraic computations. For matrix factorization as well as for certain ReLU network architectures, this allows us to show that the conservation laws reported in the literature are complete (there are no other independent quantities that would be preserved by all gradient flows).

Contributions

We formalize the notion of a conservation law, a quantity preserved through all gradient flows given a model architecture (e.g. a ReLU neural network with prescribed layers). Our main contributions are:

to show that for several classical losses, characterizing conservation laws for deep linear (resp. shallow ReLU) networks boils down to analyzing a finite dimensional space of vector fields;

to propose an algorithm (coded in SageMath) identifying polynomial conservation laws on linear / ReLU network architectures; it identifies all known laws on selected examples;

to formally define the maximum number of (not necessarily polynomial) independent conservation laws and characterize it a) theoretically via Lie algebra computations; and b) practically via an algorithm (coded in SageMath) computing this number on worked examples;

to illustrate that in certain settings these findings allow to rewrite an over-parameterized flow as an 'intrinsic' low-dimensional flow;

to highlight that the cost function associated to the training of linear and ReLU networks, shallow or deep, with various losses (quadratic and more) fully fits the proposed framework.

A consequence of our results is to show for the first time that conservation laws commonly reported in the literature are maximal: there is no other independent preserved quantity (see Propositions 4.3, 4.2, Corollary 4.4) and Section 4.2).

2 Conservation Laws for Gradient Flows

After some reminders on gradient flows, we formalize the notion of conservation laws.

2.1 Gradient dynamics

We consider learning problems, where we denote x i ∈ R m the features and y i ∈ Y the targets (for regression, typically with Y = R n ) or labels (for classification) in the case of supervised learning, while y i can be considered constant for unsupervised/self-supervised learning. We denote X := ( x i ) i and Y := ( y i ) i . Prediction is performed by a parametric mapping g ( θ , · ) : R m → R n (for instance a neural network) which is trained by empirically minimizing over parameters θ ∈ Θ ⊆ R D a cost

where /lscript is the loss function. In practical examples with linear or ReLU networks, Θ is either R D or an open set of 'non-degenerate' parameters. The goal of this paper is to analyze what functions h ( θ ) are preserved during the gradient flow (the continuous time limit of gradient descent) of E X,Y :

A priori, one can consider different 'levels' of conservation, depending whether h is conserved: during the optimization of E X,Y for a given loss /lscript and a given data set ( x i , y i ) i ; or given a loss /lscript , during the optimization of E X,Y for any data set ( x i , y i ) i . Note that using stochastic optimization methods and discrete gradients would break the exact preservation of the conservation laws, and only approximate conservation would hold, as remarked in [16].

2.2 Conserved functions

As they are based on gradient flows, conserved functions are first defined locally.

Definition 2.1 (Conservation through a flow) . Consider an open subset Ω ⊆ Θ and a vector field χ ∈ C 1 ( Ω , R D ) . By the Cauchy-Lipschitz theorem, for each initial condition θ init, there exists a unique maximal solution t ∈ [0 , T θ init ) ↦→ θ ( t, θ init ) of the ODE ˙ θ ( t ) = χ ( θ ( t )) with θ (0) = θ init. A function h : Ω ⊆ R D → R is conserved on Ω through the vector field χ if h ( θ ( t, θ init )) = h ( θ init ) for each choice of θ init ∈ Ω and every t ∈ [0 , T θ init ) . It is conserved on Ω through a subset W ⊂ C 1 ( Ω , R D ) if h is conserved on Ω during all flows induced by all χ ∈ W .

In particular, one can adapt this definition for the flow induced by the cost (1).

Definition 2.2 (Conservation during the flow (2) with a given dataset) . Consider an open subset Ω ⊆ Θ and a dataset ( X,Y ) such that E X,Y ∈ C 2 ( Ω , R ) . A function h : Ω ⊆ R D → R is conserved on Ω during the flow (2) if it is conserved through the vector field χ ( · ) := ∇ E X,Y ( · ) .

Our goal is to study which functions are conserved during 'all' flows defined by the ODE (2). This in turn leads to the following definition.

Definition 2.3 (Conservation during the flow (2) with 'any' dataset) . Consider an open subset Ω ⊂ Θ and a loss /lscript ( z, y ) such that /lscript ( · , y ) is C 2 -differentiable for all y ∈ Y . A function h : Ω ⊆ R D → R is conserved on Ω for any data set if, for each data set ( X,Y ) such that g ( · , x i ) ∈ C 2 ( Ω , R ) for each i , the function h is conserved on Ω during the flow (2). This leads us to introduce the family of vector fields:

so that being conserved on Ω for any dataset is the same as being conserved on Ω through W g Ω .

The above definitions are local and conditioned on a choice of open set of parameters Ω ⊂ Θ . We are rather interested in functions defined on the whole parameter space Θ , hence the following definition.

Definition 2.4. A function h : Θ ↦→ R is locally conserved on Θ for any data set if for each open subset Ω ⊆ Θ , h is conserved on Ω for any data set.

A basic property of C 1 conserved functions (which proof can be found in Appendix A) corresponds to an 'orthogonality' between their gradient and the considered vector fields.

Proposition 2.5. Given a subset W ⊂ C 1 ( Ω , R D ) , its trace at θ ∈ Ω is defined as the linear space

A function h ∈ C 1 ( Ω , R ) is conserved on Ω through W if, and only if ∇ h ( θ ) ⊥ W ( θ ) , ∀ θ ∈ Ω .

Therefore, combining Proposition 2.5 and Definition 2.4, the object of interest to study locally conserved functions is the union of the traces

Corollary 2.6. A function h : Θ ↦→ R is locally conserved on Θ for any data set if and only if ∇ h ( θ ) ⊥ W g θ for all θ ∈ Θ .

It will soon be shown (cf Theorem 2.14) that W g θ can be rewritten as the trace W ( θ ) of a simple finitedimensional functional space W . Meanwhile, we keep the specific notation. For the moment, this set is explicitly characterized via the following proposition (which proof can be found in Appendix B).

Proposition 2.7. Assume that for each y ∈ Y the loss /lscript ( z, y ) is C 2 -differentiable with respect to z ∈ R n . For each θ ∈ Θ we have:

where X θ is the set of data points x such that g ( · , x ) is C 2 -differentiable in the neighborhood of θ .

Example 2.8 . As a first simple example, consider a two-layer linear neural network in dimension 1 (both for the input and output), with a single neuron. For such - admittedly trivial - architecture, the parameter is θ = ( u, v ) ⊆ R 2 and the model writes g ( θ , x ) = uvx . One can directly check that the function: h ( u, v ) = u 2 -v 2 is locally conserved on R 2 for any data set . Indeed in that case ∇ h ( u, v ) = (2 u, -2 v ) /latticetop ⊥ W g θ = span ( x,y ) ∈ R × Y { ( vx, ux ) /latticetop ∇ z /lscript ( g ( θ , x ) , y ) } = R × ( v, u ) /latticetop given that the gradient ∇ z /lscript ( g ( θ , x ) , y ) is an arbitrary scalar.

In this example we obtain a simple expression of W g θ , however in general cases it is not possible to obtain such a simple expression from Proposition 2.7. We will show that in some cases, it is possible to express W g θ as the trace W ( θ ) of a simple finite-dimensional space W (cf. Theorem 2.14).

2.3 Reparametrization

To make the mathematical analysis tractable and provide an algorithmic procedure to determine these functions, our fundamental hypothesis is that the model g ( θ , x ) can be (locally) factored via a reparametrization φ as f ( φ ( θ ) , x ) . We require that the model g ( θ , x ) satisfies the following central assumption.

Assumption 2.9 (Local reparameterization) . There exists d and φ ∈ C 2 ( Θ , R d ) such that: for each parameter θ 0 in the open set Θ ⊆ R D , for each x ∈ X such that θ ↦→ g ( θ , x ) is C 2 in a neighborhood of θ 0 1 , there is a neighborhood Ω of θ 0 and f ( · , x ) ∈ C 2 ( φ ( Ω ) , R n ) such that

Note that if the model g ( · , x ) is smooth on Ω then (6) is always satisfied with φ := id and f ( · , x ) := g ( · , x ) , yet this trivial factorization fails to capture the existence and number of conservation laws as studied in this paper. This suggests that, among all factorizations shaped as (6), there may be a notion of an optimal one.

Example 2.10 . (Factorization for linear neural networks) In the two-layer case, with r neurons, denoting θ = ( U, V ) ∈ R n × r × R m × r (so that D = ( n + m ) r ), we can factorize g ( θ , x ) := UV /latticetop x by the reparametrization φ ( θ ) := UV /latticetop ∈ R n × m using f ( φ , x ) = φ · x . More generally for q layers, with θ = ( U 1 , · · · , U q ) , we can still factorize g ( θ , x ) := U 1 · · · U q x using φ ( θ ) := U 1 · · · U q and the same f . This factorization is globally valid on Ω = Θ = R D since f ( · , x ) does not depend on θ 0 .

The notion of locality of the factorization (6) is illustrated by the next example.

Example 2.11 (Factorization for two-layer ReLU networks) . Consider g ( θ , x ) = ( ∑ r j =1 u k,j σ ( 〈 v j , x 〉 + b j ) + c k ) n k =1 , with σ ( t ) := max( t, 0) the ReLU activation function and v j ∈ R m , u k,j ∈ R , b j , c k ∈ R . Then, denoting θ = ( U, V, b, c ) with U = ( u k,j ) k,j =: ( u 1 , · · · , u r ) ∈ R n × r , V = ( v 1 , · · · , v r ) ∈ R m × r , b = ( b 1 , · · · , b r ) /latticetop ∈ R r and c = ( c 1 , · · · , c n ) ∈ R n (so that D = ( n + m + 1) r + n ), we rewrite g ( θ , x ) = ∑ r j =1 u j ε j,x ( v /latticetop j x + b j ) + c where, given x , ε j,x = 1 ( v /latticetop j x + b j > 0) is piecewise constant with respect to θ . Consider θ 0 = ( U 0 , V 0 , b 0 , c 0 ) ∈ R D where V 0 = ( v 0 1 , · · · , v 0 r ) and b 0 = ( b 0 1 , · · · , b 0 r ) /latticetop . Then the set X θ 0 introduced in Proposition 2.7 is X θ 0 = R m -∪ j { v 0 j /latticetop x + b 0 j = 0 } . Let x ∈ X θ 0 . Then on any domain Ω ⊂ R D such that θ 0 ∈ Ω and ε j,x ( θ ) := 1 ( v /latticetop j x + b j > 0) is constant over θ ∈ Ω , the model g θ ( x ) can be factorized by the reparametrization φ ( θ ) = (( u j v /latticetop j , u j b j ) r j =1 , c ) . In particular, in the case without bias ( ( b, c ) = (0 , 0) ), the reparametrization is defined by φ ( θ ) = ( φ j ) r j =1 where φ j = φ j ( θ ) := u j v /latticetop j ∈ R n × m (here d = rmn ) using f ( φ , x ) = ∑ j ε j,x φ j x : the reparametrization φ ( θ ) contains r matrices of size m × n (each of rank at most one) associated to a 'local' f ( · , x ) valid in a neighborhood of θ . A similar factorization is possible for deeper ReLU networks [27], as further discussed in the proof of Theorem 2.14 in Appendix C.

Combining Proposition 2.7 and using chain rules, we get a new characterization of W g θ :

Proposition 2.12. Assume that the loss /lscript ( z, y ) is C 2 -differentiable with respect to z . We recall (cf (5) ) that W g θ := ∪ Ω ⊆ Θ : Ω is open and Ω /owner θ W g Ω ( θ ) . Under Assumption 2.9, for all θ ∈ Θ :

with ∂φ ( θ ) ∈ R d × D the Jacobian of φ and W f φ ( θ ) := span ( x,y ) ∈ X θ × Y { ∂ f x ( φ ( θ )) /latticetop ∇ z /lscript ( g ( θ , x ) , y ) } , where f x ( · ) := f ( · , x ) .

We show in Section 2.4 that, under mild assumptions on the loss /lscript , W f φ ( θ ) = R d , so that Proposition 2.12 yields W g θ = range( ∂φ ( θ ) /latticetop ) . Then by Corollary 2.6, a function h that is locally conserved on Θ for any data set is entirely characterized via the kernel of ∂φ ( θ ) /latticetop : ∂φ ( θ ) /latticetop ∇ h ( θ ) = 0 for all θ ∈ Θ . The core of our analysis is then to analyze the (Lie algebraic) structure of range( ∂φ ( · ) /latticetop ) .

2.4 From conserved functions to conservation laws

For linear and ReLU networks we show in Theorem 2.14 and Proposition 2.16 that under (mild) assumptions on the loss /lscript ( · , · ) , being locally conserved on Θ for any data set (according to Definition 2.4) is the same as being conserved (according to Definition 2.1) on Θ through the finite-dimensional subspace

where we write ∂φ ( θ ) /latticetop = ( ∇ φ 1 ( θ ) , · · · , ∇ φ d ( θ )) ∈ R D × d , with ∇ φ i ∈ C 1 ( Θ , R D ) .

The following results (which proofs can be found in Appendix C) establish that in some cases, the functions locally conserved for any data set are exactly the functions conserved through W φ .

Lemma 2.13. Assume that the loss ( z, y ) ↦→ /lscript ( z, y ) is C 2 -differentiable with respect to z ∈ R n and satisfies the condition:

Then for linear neural networks (resp. for two-layer ReLU networks) and all θ ∈ Θ we have W f φ ( θ ) = R d , with the reparametrization φ from Example 2.10 and Θ := R D (resp. with φ from Example 2.11 and Θ consisting of all parameter θ of the network such that hidden neurons are associated to pairwise distinct 'hyperplanes', cf Appendix C for details).

Condition (9) holds for classical losses /lscript (e.g. quadratic/logistic losses), as shown in Lemma C.2 in Appendix C. Note that the additional hypothesis of pairwise distinct hyperplanes for the two-layer ReLU case is a generic hypothesis and is usual (see e.g. the notion of twin neurons in [27]). The tools from Appendix C extend Theorem 2.14 beyond (deep) linear and shallow ReLU networks. An open problem is whether the conclusions of Lemma 2.13 still hold for deep ReLU networks.

Theorem 2.14. Under the same assumptions as in Lemma 2.13, we have that for linear neural networks, for all θ ∈ Θ := R D :

The same result holds for two-layer ReLU networks with φ from Example 2.11 and Θ the (open) set of all parameters θ such that hidden neurons are associated to pairwise distinct 'hyperplanes'.

This means as claimed that for linear and two-layer ReLU networks, being locally conserved on Θ for any data set exactly means being conserved on Θ through the finite-dimensional functional space W φ ⊆ C 1 ( Θ , R D ) . This motivates the following definition

Definition 2.15. A real-valued function h is a conservation law of φ if it is conserved through W φ .

Proposition 2.5 yields the following intermediate result.

Proposition 2.16. h ∈ C 1 ( Ω , R ) is a conservation law for φ if and only if

Thanks to Theorem 2.14, the space W φ defined in (8) introduces a much simpler proxy to express W g θ as a trace of a subset of C 1 ( Θ , R D ) . Moreover, when φ is C ∞ , W φ is a finite-dimensional space of infinitely smooth functions on Θ , and this will be crucial in Section 4.1 to provide a tractable scheme (i.e. operating in finite dimension) to compute the maximum number of independent conservation laws, using the Lie algebra computations that will be described in Section 3.

Example 2.17 . Revisiting Example 2.8, the function to minimize is factorized by the reparametrization φ : ( u ∈ R , v ∈ R ) ↦→ uv ∈ R with θ := ( u, v ) . We saw that h (( u, v )) := u 2 -v 2 is conserved: and indeed 〈∇ h ( u, v ) , ∇ φ ( u, v ) 〉 = 2 uv -2 vu = 0 , ∀ ( u, v ) .

In this simple example, the characterization of Proposition 2.16 gives a constructive way to find such a conserved function: we only need to find a function h such that 〈∇ h ( u, v ) , ∇ φ ( u, v ) 〉 = 〈∇ h ( u, v ) , ( v, u ) /latticetop 〉 = 0 . The situation becomes more complex in higher dimensions, since one needs to understand the interplay between the different vector fields in W φ .

2.5 Constructibility of some conservation laws

Observe that in Example 2.17 both the reparametrization φ and the conservation law h are polynomials, a property that surprisingly systematically holds in all examples of interest in the paper, making it possible to algorithmically construct some conservation laws as detailed now.

By Proposition 2.16, a function h is a conservation law if it is in the kernel of the linear operator h ∈ C 1 ( Ω , R ) ↦→ ( θ ∈ Ω ↦→ ( 〈∇ h ( θ ) , ∇ φ i ( θ ) 〉 ) i =1 , ··· ,d ) . Thus, one could look for conservation laws in a prescribed finite-dimensional space by projecting these equations in a basis (as in finite-element methods for PDEs). Choosing the finite-dimensional subspace could be generally tricky, but for the linear and ReLU cases all known conservation laws are actually polynomial 'balancedness-type conditions' [1, 2, 9], see Section 4. In these cases, the vector fields in W φ are also polynomials (because φ is polynomial, see Theorem C.4 and Theorem C.5 in Appendix C), hence θ ↦→〈∇ h ( θ ) , ∇ φ i ( θ ) 〉 is a polynomial too. This allows us to compute a basis of independent polynomial conservation laws of a given degree (to be freely chosen) for these cases, by simply focusing on the corresponding subspace of polynomials. We coded the resulting equations in SageMath, and we found back on selected examples (see Appendix J) all existing known conservation laws both for ReLU and linear networks. Open-source code is available at [18].

2.6 Independent conserved functions

Having an algorithm to build conservation laws is nice, yet how can we know if we have built 'all' laws? This requires first defining a notion of a 'maximal' set of functions, which would in some sense be independent. This does not correspond to linear independence of the functions themselves (for instance, if h is a conservation law, then so is h k for each k ∈ N but this does not add any other constraint), but rather to pointwise linear independence of their gradients. This notion of independence is closely related to the notion of 'functional independence' studied in [7, 20]. For instance, it is shown in [20] that smooth functionally dependent functions are characterized by having dependent gradients everywhere. This motivates the following definition.

Definition 2.18. A family of N functions ( h 1 , · · · , h N ) conserved through W ⊂ C 1 ( Ω , R D ) is said to be independent if the vectors ( ∇ h 1 ( θ ) , · · · , ∇ h N ( θ )) are linearly independent for all θ ∈ Ω .

An immediate upper bound holds on the largest possible number N of functionally independent functions h 1 , . . . , h N conserved through W : for θ ∈ Ω ⊆ R D , the space span { ∇ h 1 ( θ ) , . . . , ∇ h N ( θ ) } ⊆ R D is of dimension N (by independence) and (by Proposition 2.5) orthogonal to W ( θ ) . Thus, it is necessary to have N ≤ D -dim W ( θ ) . As we will now see, this bound can be tight under additional assumptions on W related to Lie brackets (corresponding to the so-called Frobenius theorem). This will in turn lead to a characterization of the maximum possible N .

3 Conservation Laws using Lie Algebra

The study of hyper-surfaces trapping the solution of ODEs is a recurring theme in control theory, since the existence of such surfaces is the basic obstruction of controllability of such systems [6]. The basic result to study these surfaces is the so-called Frobenius theorem from differential calculus (See Section 1.4 of [13] for a good reference for this theorem). It relates the existence of such surfaces, and their dimensions, to some differential condition involving so-called 'Lie brackets' [ u, v ] between pairs of vector fields (see Section 3.1 below for a more detailed exposition of this operation). However, in most cases of practical interest (such as for instance matrix factorization), the Frobenius theorem is not suitable for a direct application to the space W φ because its Lie bracket condition is not satisfied. To identify the number of independent conservation laws, one needs to consider the algebraic closure of W φ under Lie brackets. The fundamental object of interest is thus the Lie algebra generated by the Jacobian vector fields, that we recall next. While this is only defined for vector fields with stronger smoothness assumption, the only consequence is that φ is required to be infinitely smooth, unlike the loss /lscript ( · , y ) and the model g ( · , x ) that can be less smooth. All concretes examples of φ in this paper are polynomial hence indeed infinitely smooth.

Notations Given a vector subspace of infinitely smooth vector fields W ⊆ X ( Θ ) := C ∞ ( Θ , R D ) , where Θ is an open subset of R D , we recall (cf Proposition 2.5) that its trace at some θ is the subspace

For each open subset Ω ⊆ Θ , we introduce the subspace of X ( Ω ) : W | Ω := { χ | Ω : χ ∈ W } .

3.1 Background on Lie algebra

A Lie algebra A is a vector space endowed with a bilinear map [ · , · ] , called a Lie bracket, that verifies for all X,Y,Z ∈ A : [ X,X ] = 0 and the Jacobi identity: [ X, [ Y, Z ]] + [ Y, [ Z, X ]] + [ Z, [ X,Y ]] = 0 .

For the purpose of this article, the Lie algebra of interest is the set of infinitely smooth vector fields X ( Θ ) , endowed with the Lie bracket [ · , · ] defined by

with ∂χ ( θ ) ∈ R D × D the jacobian of χ at θ . The space R n × n of matrices is also a Lie algebra endowed with the Lie bracket [ A,B ] := AB -BA. This can be seen as a special case of (12) in the case of linear vector fields, i.e. χ ( θ ) = A θ .

Generated Lie algebra Let A be a Lie algebra and let W ⊂ A be a vector subspace of A . There exists a smallest Lie algebra that contains W . It is denoted Lie( W ) and called the generated Lie algebra of W . The following proposition [6, Definition 20] constructively characterizes Lie( W ) , where for vector subspaces [ W,W ′ ] := { [ χ 1 , χ 2 ] : χ 1 ∈ W, χ 2 ∈ W ′ } , and W + W ′ = { χ 1 + χ 2 : χ 1 ∈ W, χ 2 ∈ W ′ } .

Proposition 3.1. Given any vector subspace W ⊆ A we have Lie( W ) = ⋃ k W k where:

We will see in Section 3.2 that the number of conservation laws is characterized by the dimension of the trace Lie( W φ )( θ ) defined in (11). The following lemma (proved in Appendix D) gives a stopping criterion to algorithmically determine this dimension (see Section 3.3 for the algorithm).

Lemma3.2. Given θ ∈ Θ , if for a given i , dim W i +1 ( θ ′ ) = dim W i ( θ ) for every θ ′ in a neighborhood of θ , then there exists a neighborhood Ω of θ such that W k ( θ ′ ) = W i ( θ ′ ) for all θ ′ ∈ Ω and k ≥ i , where the V i are defined by Proposition 3.1. Thus Lie( W )( θ ′ ) = W i ( θ ′ ) for all θ ′ ∈ Ω . In particular, the dimension of the trace of Lie( W ) is locally constant and equal to the dimension of W i ( θ ) .

3.2 Number of conservation laws

The following theorem uses the Lie algebra generated by W φ to characterize the number of conservation laws. The proof of this result is based on two successive uses of the Frobenius theorem and can be found in Appendix E (where we also recall Frobenius theorem for the sake of completeness).

Theorem 3.3. If dim(Lie( W φ )( θ )) is locally constant then each θ ∈ Ω ⊆ R D admits a neighborhood Ω ′ such that there are D -dim(Lie( W φ )( θ )) (and no more) independent conserved functions through W φ | Ω ′ , i.e., there are D -dim(Lie( W φ )( θ )) independent conservation laws of φ on Ω ′ .

Remark 3.4 . The proof of the Frobenius theorem (and therefore of our generalization Theorem 3.3) is actually constructive. From a given φ , conservation laws are obtained in the proof by integrating in time ( i.e. solving an advection equation) the vector fields belonging to W φ . Unfortunately, this cannot be achieved in closed form in general, but in small dimensions, this could be carried out numerically (to compute approximate discretized laws on a grid or approximate them using parametric functions such as Fourier expansions or neural networks).

A fundamental aspect of Theorem 3.3 is to rely only on the dimension of the trace of the Lie algebra associated with the finite-dimensional vector space W φ . Yet, even if W φ is finite-dimensional, it might be the case that Lie( W φ ) itself remains infinite-dimensional. Nevertheless, what matters is not the dimension of Lie( W φ ) , but that of its trace Lie( W φ )( θ ) , which is always finite (and potentially much smaller that dimLie( W φ ) even when the latter is finite) and computationally tractable thanks to Lemma 3.2 as detailed in Section 3.3. In section 4.1 we work out the example of matrix factorization, a non-trivial case where the full Lie algebra Lie( W φ ) itself remains finite-dimensional.

Theorem 3.3 requires that the dimension of the trace at θ of the Lie algebra is locally constant. This is a technical assumption, which typically holds outside a set of pathological points. A good example is once again matrix factorization, where we show in Section 4.1 that this condition holds generically.

3.3 Method and algorithm, with examples

Given a reparametrization φ for the architectures to train, to determine the number of independent conservation laws of φ , we leverage the characterization 3.1 to algorithmically compute dim(Lie( W φ )( θ )) using an iterative construction of bases for the subspaces W k starting from W 0 := W φ , and stopping as soon as the dimension stagnates thanks to Lemma 3.2. Our open-sourced code is available at [18] and uses SageMath. As we now show, this algorithmic principle allows to fully work out certain settings where the stopping criterion of Lemma 3.2 is reached at the first step ( i = 0 ) or the second one ( i = 1 ). Section 4.2 also discusses its numerical use for an empirical investigation of broader settings.

Example where the iterations of Lemma 3.2 stop at the first step. This corresponds to the case where Lie W φ ( θ ) = W 1 ( θ ) = W 0 ( θ ) := W φ ( θ ) on Ω . This is the case if and only if W φ satisfies that

i.e., when Frobenius Theorem (see Theorem E.1 in Appendix E) applies directly. The first example is a follow-up to Example 2.11.

/negationslash

Example 3.5 (two-layer ReLU networks without bias) . Consider θ = ( U, V ) with U ∈ R n × r , V ∈ R m × r , n, m, r ≥ 1 (so that D = ( n + m ) r ), and the reparametrization φ ( θ ) := ( u i v /latticetop i ) i =1 , ··· ,r ∈ R n × m × r , where U = ( u 1 ; · · · ; u r ) and V = ( v 1 ; · · · ; v r ) . As detailed in Appendix F.1, since φ ( θ ) is a collection of r rank-one n × m matrices, dim( W φ ( θ )) = rank ∂φ ( θ ) = ( n + m -1) r is constant on the domain Ω such that u i , v j = 0 , and W φ satisfies (13), hence by Theorem 3.3 each θ has a neighborhood Ω ′ such that there exists r (and no more) independent conserved function through W φ | Ω ′ . The r known conserved functions [9] given by h i : ( U, V ) ↦→‖ u i ‖ 2 -‖ v i ‖ 2 , i = 1 , · · · , r , are independent, hence they are complete.

Example where the iterations of Lemma 3.2 stop at the second step (but not the first one). Our primary example is matrix factorization, as a follow-up to Example 2.10.

Example 3.6 (two-layer linear neural networks) . With θ = ( U, V ) , where ( U ∈ R n × r , V ∈ R m × r ) the reparameterization φ ( θ ) := UV /latticetop ∈ R n × m (here d = nm ) factorizes the functions minimized during the training of linear two-layer neural networks (see Example 2.10). As shown in Appendix I, condition (13) is not satisfied when r > 1 and max ( n, m ) > 1 . Thus, the stopping criterion of Lemma 3.2 is not satisfied at the first step. However, as detailed in Proposition H.3 in Appendix H, ( W φ ) 1 = ( W φ ) 2 = Lie( W φ ) , hence the iterations of Lemma 3.2 stop at the second step.

We complete this example in the next section by showing (Corollary 4.4) that known conservation laws are indeed complete. Whether known conservation laws remain valid and/or complete in this settings and extended ones is further studied in Section 4 and Appendix F.

3.4 Application: recasting over-parameterized flows as low-dimensional Riemannian flows

As we now show, one striking application of Theorem 3.3 (in simple cases where dim( W φ ( θ )) = dim(Lie W φ ( θ )) is constant on Ω , i.e., rank ( ∂φ ( θ )) is constant on Ω and W φ satisfies (13)) is to fully rewrite the high-dimensional flow θ ( t ) ∈ R D as a low-dimensional flow on z ( t ) := φ ( θ ( t )) ∈ R d , where this flow is associated with a Riemannian metric tensor M that is induced by φ and depends on the initialization θ init . We insist on the fact that this is only possible in very specific cases, but this phenomenon is underlying many existing works that aim at writing in closed form the implicit bias associated with some training dynamics (see Section 1 for some relevant literature. Our analysis sheds some light on cases where this is possible, as shown in the next proposition.

Proposition 3.7. Assume that rank ( ∂φ ( θ )) is constant on Ω and that W φ satisfies (13) . If θ ( t ) ∈ R D satisfies the ODE (2) where θ init ∈ Ω , then there is 0 < T /star θ init ≤ T θ init such that z ( t ) := φ ( θ ( t )) ∈ R d satisfies the ODE

where M ( z ( t ) , θ init ) ∈ R d × d is a symmetric positive semi-definite matrix.

See Appendix G for a proof. Revisiting Example 3.5 leads to the following analytic example. Example 3.8 . Given the reparametrization φ : ( u ∈ R ∗ , v ∈ R d ) ↦→ uv ∈ R d , the variable z := uv satisfies (14) with M ( z, θ init ) = ‖ z ‖ δ I d + ‖ z ‖ -1 δ zz /latticetop , with ‖ z ‖ δ := δ + √ δ 2 + ‖ z ‖ 2 , δ := 1 / 2( u 2 init -‖ v init ‖ 2 ) .

Another analytic example is discussed in Appendix G. In light of these results, an interesting perspective is to better understand the dependance of the Riemannian metric with respect to initialization, to possibly guide the choice of initialization for better convergence dynamics.

Note that the metric M ( z, θ init ) can have a kernel. Indeed, in practice, while φ is a function from R D to R d , the dimensions often satisfy rank ∂φ ( θ ) < min( d, D ) , i.e., φ ( θ ) lives in a manifold of lower dimension. The evolution (14) should then be understood as a flow on this manifold. The kernel of M ( z, θ init ) is orthogonal to the tangent space at z of this manifold.

4 Conservation Laws for Linear and ReLU Neural Networks

To showcase the impact of our results, we show how they can be used to determine whether known conservation laws for linear (resp. ReLU) neural networks are complete, and to recover these laws algorithmically using reparametrizations φ adapted to these two settings. Concretely, we study the conservation laws for neural networks with q layers, and either a linear or ReLU activation, with an emphasis on q = 2 . We write θ = ( U 1 , · · · , U q ) with U i ∈ R n i -1 × n i the weight matrices and we assume that θ satisfies the gradient flow (2). In the linear case the reparametrization is φ Lin ( θ ) := U 1 · · · U q . For ReLU networks, we use the (polynomial) reparametrization φ ReLu of [27, Definition 6], which is defined for any (deep) feedforward ReLU network, with or without bias. In the simplified setting of networks without biases it reads explicitly as:

with U [ i, j ] the ( i, j ) -th entry of U . This covers φ ( θ ) := ( u j v /latticetop j ) r j =1 ∈ R n × m × r from Example 2.11. Some conservation laws are known for the linear case φ Lin [1, 2] and for the ReLu case φ ReLu [9].

Proposition 4.1 ( [1, 2, 9] ) . If θ := ( U 1 , · · · , U q ) satisfies the gradient flow (2) , then for each i = 1 , · · · , q -1 the function θ ↦→ U /latticetop i U i -U i +1 U /latticetop i +1 (resp. the function θ ↦→ diag ( U /latticetop i U i -U i +1 U /latticetop i +1 ) ) defines n i × ( n i +1) / 2 conservation laws for φ Lin (resp. n i conservation laws for φ ReLu ).

Proposition 4.1 defines ∑ q -1 i =1 n i × ( n i +1) / 2 conserved functions for the linear case. In general they are not independent, and we give below in Proposition 4.2, for the case of q = 2 , the exact

number of independent conservation laws among these particular laws. Establishing whether there are other (previously unknown) conservation laws is an open problem for q > 2 . We already answered negatively to this question in the two-layer ReLu case without bias (See Example 3.5). In the following Section (Corollary 4.4), we show the same result in the linear case q = 2 . Numerical computations suggest this is still the case for deeper linear and ReLU networks as detailed in Section 4.2.

4.1 The matrix factorization case ( q = 2 )

To simplify the analysis when q = 2 , we rewrite θ = ( U, V ) as a vertical matrix concatenation denoted ( U ; V ) ∈ R ( n + m ) × r , and φ ( θ ) = φ Lin ( θ ) = UV /latticetop ∈ R n × m .

How many independent conserved functions are already known? The following proposition refines Proposition 4.1 for q = 2 by detailing how many independent conservation laws are already known. See Appendix H.1 for a proof.

Proposition 4.2. Consider Ψ : θ = ( U ; V ) ↦→ U /latticetop U -V /latticetop V ∈ R r × r and assume that ( U ; V ) has full rank noted rk . Then the function Ψ gives rk · (2 r +1 -rk ) / 2 independent conserved functions.

There exist no more independent conserved functions. We now come to the core of the analysis, which consists in actually computing Lie( W φ ) as well as its traces Lie( W φ )( θ ) in the matrix factorization case. The crux of the analysis, which enables us to fully work out theoretically the case q = 2 , is that W φ is composed of linear vector fields (that are explicitly characterized in Proposition H.2 in Appendix H), the Lie bracket between two linear fields being itself linear and explicitly characterized with skew matrices, see Proposition H.3 in Appendix H. Eventually, what we need to compute is the dimension of the trace Lie( W φ )( U, V ) for any ( U, V ) . We prove the following in Appendix H.

Proposition 4.3. If ( U ; V ) ∈ R ( n + m ) × r has full rank noted rk , then: dim(Lie( W φ )( U ; V )) = ( n + m ) r -(2 r +1 -rk ) / 2 .

With this explicit characterization of the trace of the generated Lie algebra and Proposition 4.2, we conclude that Proposition 4.1 has indeed exhausted the list of independent conservation laws.

Corollary 4.4. If ( U ; V ) has full rank, then all conserved functions are given by Ψ : ( U, V ) ↦→ U /latticetop U -V /latticetop V . In particular, there exist no more independent conserved functions.

4.2 Numerical guarantees in the general case

The expressions derived in the previous section are specific to the linear case q = 2 . For deeper linear networks and for ReLU networks, the vector fields in W φ are non-linear polynomials, and computing Lie brackets of such fields can increase the degree, which could potentially make the generated Lie algebra infinite-dimensional. One can however use Lemma 3.2 and stop as soon as dim(( W φ ) k ( θ )) stagnates. Numerically comparing this dimension with the number N of independent conserved functions known in the literature (predicted by Proposition 4.1) on a sample of depths/widths of small size, we empirically confirmed that there are no more conservation laws than the ones already known for deeper linear networks and for ReLU networks too (see Appendix J for details). Our code is open-sourced and is available at [18]. It is worth mentioning again that in all tested cases φ is polynomial, and there is a maximum set of conservation laws that are also polynomial, which are found algorithmically (as detailed in Section 2.5).

Conclusion

In this article, we proposed a constructive program for determining the number of conservation laws. An important avenue for future work is the consideration of more general classes of architectures, such as deep convolutional networks, normalization, and attention layers. Note that while we focus in this article on gradient flows, our theory can be applied to any space of displacements in place of W φ . This could be used to study conservation laws for flows with higher order time derivatives, for instance gradient descent with momentum, by lifting the flow to a higher dimensional phase space. A limitation that warrants further study is that our theory is restricted to continuous time gradient flow. Gradient descent with finite step size, as opposed to continuous flows, disrupts exact conservation. The study of approximate conservation presents an interesting avenue for future work.

References

S. ARORA, N. COHEN, N. GOLOWICH, AND W. HU, A convergence analysis of gradient descent for deep linear neural networks , arXiv preprint arXiv:1810.02281, (2018).

S. ARORA, N. COHEN, AND E. HAZAN, On the optimization of deep networks: Implicit acceleration by overparameterization , in Int. Conf. on Machine Learning, PMLR, 2018, pp. 244253.

S. AZULAY, E. MOROSHKO, M. S. NACSON, B. E. WOODWORTH, N. SREBRO, A. GLOBERSON, AND D. SOUDRY, On the implicit bias of initialization shape: Beyond infinitesimal mirror descent , in Proceedings of the 38th International Conference on Machine Learning, M. Meila and T. Zhang, eds., vol. 139 of Proceedings of Machine Learning Research, PMLR, 18-24 Jul 2021, pp. 468-477.

B. BAH, H. RAUHUT, U. TERSTIEGE, AND M. WESTDICKENBERG, Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers , Information and Inference: A Journal of the IMA, 11 (2022), pp. 307-353.

M. BELKIN, D. HSU, S. MA, AND S. MANDAL, Reconciling modern machine-learning practice and the classical bias-variance trade-off , Proc. of the National Academy of Sciences, 116 (2019), pp. 15849-15854.

B. BONNARD, M. CHYBA, AND J. ROUOT, Geometric and Numerical Optimal Control - Application to Swimming at Low Reynolds Number and Magnetic Resonance Imaging , SpringerBriefs in Mathematics, Springer Int. Publishing, 2018.

A. B. BROWN, Functional dependence , Transactions of the American Mathematical Society, 38 (1935), pp. 379-394.

L. CHIZAT AND F. BACH, Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss , in Conf. on Learning Theory, PMLR, 2020, pp. 1305-1338.

S. S. DU, W. HU, AND J. D. LEE, Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced , Adv. in Neural Inf. Proc. Systems, 31 (2018).

S. GUNASEKAR, J. LEE, D. SOUDRY, AND N. SREBRO, Characterizing implicit bias in terms of optimization geometry , in Int. Conf. on Machine Learning, PMLR, 2018, pp. 1832-1841.

S. GUNASEKAR, B. E. WOODWORTH, S. BHOJANAPALLI, B. NEYSHABUR, AND N. SREBRO, Implicit regularization in matrix factorization , Adv. in Neural Inf. Proc. Systems, 30 (2017).

G. G Ł UCH AND R. URBANKE, Noether: The more things change, the more stay the same , 2021.

A. ISIDORI, Nonlinear system control , New York: Springer Verlag, 61 (1995), pp. 225-236.

Z. JI, M. DUDÍK, R. E. SCHAPIRE, AND M. TELGARSKY, Gradient descent follows the regularization path for general losses , in Conf. on Learning Theory, PMLR, 2020, pp. 21092136.

Z. JI AND M. TELGARSKY, Gradient descent aligns the layers of deep linear networks , arXiv preprint arXiv:1810.02032, (2018).

D. KUNIN, J. SAGASTUY-BRENA, S. GANGULI, D. L. YAMINS, AND H. TANAKA, Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics , arXiv preprint arXiv:2012.04728, (2020).

Z. LI, T. WANG, J. D. LEE, AND S. ARORA, Implicit bias of gradient descent on reparametrized models: On equivalence to mirror descent , in Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, eds., vol. 35, Curran Associates, Inc., 2022, pp. 34626-34640.

S. MARCOTTE, R. GRIBONVAL, AND G. PEYRÉ, Code for reproducible research. Abide by the Law and Follow the Flow: Conservation Laws for Gradient Flows , Oct. 2023.

H. MIN, S. TARMOUN, R. VIDAL, AND E. MALLADA, On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks , in Int. Conf. on Machine Learning, PMLR, 2021, pp. 7760-7768.

W. F. NEWNS, Functional dependence , The American Mathematical Monthly, 74 (1967), pp. 911-920.

B. NEYSHABUR, Implicit regularization in deep learning , arXiv preprint arXiv:1709.01953, (2017).

B. NEYSHABUR, R. TOMIOKA, AND N. SREBRO, In search of the real inductive bias: On the role of implicit regularization in deep learning , arXiv preprint arXiv:1412.6614, (2014).

E. NOETHER, Invariante variationsprobleme , Nachrichten von der Gesellschaft der Wissenschaften zu Göttingen, Mathematisch-Physikalische Klasse, 1918 (1918), pp. 235-257.

A. M. SAXE, J. L. MCCLELLAND, AND S. GANGULI, Exact solutions to the nonlinear dynamics of learning in deep linear neural networks , arXiv preprint arXiv:1312.6120, (2013).

S. SHALEV-SHWARTZ AND S. BEN-DAVID, Understanding machine learning: From theory to algorithms , Cambridge university press, 2014.

D. SOUDRY, E. HOFFER, M. S. NACSON, S. GUNASEKAR, AND N. SREBRO, The implicit bias of gradient descent on separable data , The Journal of Machine Learning Research, 19 (2018), pp. 2822-2878.

P. STOCK AND R. GRIBONVAL, An Embedding of ReLU Networks and an Analysis of their Identifiability , Constructive Approximation, (2022). Publisher: Springer Verlag.

S. TARMOUN, G. FRANCA, B. D. HAEFFELE, AND R. VIDAL, Understanding the dynamics of gradient flow in overparameterized linear models , in Int. Conf. on Machine Learning, PMLR, 2021, pp. 10153-10161.

THE SAGE DEVELOPERS, SageMath, the Sage Mathematics Software System (Version 9.7) , 2022. https://www.sagemath.org .

C. ZHANG, S. BENGIO, M. HARDT, B. RECHT, AND O. VINYALS, Understanding deep learning requires rethinking generalization , in Int. Conf. on Learning Representations, 2017.

B. ZHAO, I. GANEV, R. WALTERS, R. YU, AND N. DEHMAMY, Symmetries, flat minima, and the conserved quantities of gradient flow , arXiv preprint arXiv:2210.17216, (2022).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study conservation laws in the gradient flow dynamics of neural networks.\nThey introduce a notion of local factorisation of the loss, stating that the loss in the neighbourhood of a given wieght vector can be decomposed in a composition of functions, a data-independent term $\\phi$ followed by a data dependent one $f$.\n\nUnder this decomposition, the authors set out to characterise conserved quantities for fixed $\\phi$ and for all $f$ giving rise to a proper ERM loss.\nThey claim that under some assumption, and for ReLU networks, this is equivalent of characterising conserved quantities for fixed $\\phi$ and for all smooth $f$.\nMore strongly, they claim that this is also equivalent to characterising conserved quantities for fixed $\\phi$ and for a special finite dimensional subspace of function, living in the linear span of the $d$ rows of the Jacobian of $\\phi$, where $d$ is the dimension of the codomain of $\\phi$.\n\nThe authors then notice that the conservation laws of linear and ReLU networks will be polynomial in the weights, and consider the question of characterising maximal sets of conservation laws.\nThey link such number to the dimension of the Lie algebra generated locally by the $d$-dimensional vector field associated to the Jacobian of $\\phi$, and work-out explicitly some examples.\n\nFinally, the authors consider the question whether the known conservation laws for two-layers linear and ReLU networks form a maximal set, and find an affirmative answer.\nThey conjecture, and briefly discuss this point in Appendix 9, that the same results hold for larger depths.\n\nSTRENGTHS:\n- The article is well written, and guides the reader nicely through quite technical results without requiring too much previous knowledge.\n- The results seem novel to the extent of my knowledge of the literature, which is not comprehensive on this subject. They seem to fit nicely into a pre-existing line of works (see ref [13] and [26] for example).\n- The main text provides reasonable justification for most of the formal results, which are proven in the Appendix. While I did not check the correctness of the proofs, the results seem reasonable.\n\nWEAKNESSES:\n- It is not clear to what extent the factorisation property the authors introduce is necessary to their analysis. Said differently, it is not clear wehther the analysis could be performed similarly on a concrete example (say fixed dataset) without any mention of a factorisation. The authors could add additional high-level explanations/justifications of this concept.\n\nQUESTIONS:\n## Major questions\n\n- It is not clear to me whether the requirement in Eq. 2 is vacuous, meaning that any architecture automatially satisfies it. Can't I just take $\\phi$ as the identity map, and $f = \\mathcal{E}$? If this trivial factorisation is possible, are the results still non-trivial? Maybe I am losing some very simple nuance. In any case, adding a counter-example, or clarifying better the requirements on $\\phi$ and $f$ could avoid doubts.\n\n- Example 2.2 provides only a local factorisation for ReLU networks. Of course locality is ok as we are considering gradient flow dynamics, which is local. But I wonder whether something special can happen at the boundaries of the set $\\Omega$ defined in Example 2.2, i.e. if there is some gluing condition/gluing phenomenon that may affect the results presented by the authors.\n\n- Is there a commonly used loss for which Eq. 7 is not satisfied?\n\n- The authors stress that their results allow for explicit construction of maximal sets of conserved quantities, yet in the manuscript they provide only an a posteriori verification that known conserved quantities in previously studied architectures indeed form maximal sets. Is there an architecture where new conservation laws can be found through the presented techniques?\n\n## Suggestions for manuscript improvement\n\nline 50: missing closing bracket\n\nline 106: specifying that $D$ is the number of weigths, and $d$ is the dimension of the \"internal representation\" of the decomposition $\\mathcal{E} = f \\circ \\phi$ would be helpful for the reader.\n\nline 119: in the inlined equation, one has a tensor $\\phi$ with three free indices equal to an expression without the same three free indices. I suggest to clarify this writing somehow.\n\nLIMITATIONS:\nThe authors briefly discuss limitations in the main text.\nI would add that the analysis seems to be limited to linear and ReLU architectures 2 layer architectures. It is not clear whether a local factorisation of the form Eq. 2 can be found for other architectures. The authors could maybe add a discussion on this point.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper studies the conservation law of gradient flow dynamics for training neural networks. The authors propose a method to determine the number of conservation laws in given gradient flow dynamics using Lie algebra generated by the Jacobian vector fields. It is shown, either theoretically or empirically, that the known conservation laws in training linear networks and ReLU networks are also maximal.

STRENGTHS:
1. The paper is well-written and clear. The conversation laws in gradient flow dynamics for training neural networks have facilitated the analysis of convergence and implicit bias, thus deserving formal analysis on finding these conversation laws given any network architecture.
2. This paper has an in-depth discussion of the conservation law under the gradient flow on a class of loss functions.
3. Analysis via Lie algebra that determines the maximum number of conservation laws.
4. Showing existing conservation laws studied for linear and ReLU networks are maximal

WEAKNESSES:
none

QUESTIONS:
Do the results stated for ReLU networks hold for any network with homogeneous activation function?

LIMITATIONS:
authors discussed the limitation.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper discusses the geometric properties of gradient descent dynamics in ML models. The authors aim to understand the properties of the optimization initialization that are preserved during the dynamics, which is often referred to as being an \"implicit bias\" of the training algorithm. They also focus on the maximal sets of independent quantities conserved during gradient flows. They have an interesting approach to find the exact number of these conserved quantities by performing algebraic manipulations on the Lie algebra generated by the Jacobian of the model.\n\nThe paper's contributions include formalizing the notion of a conservation law in the setting of training neural networks, proposing an algorithm to identify simply expressed (e.g., polynomial) conservation laws on ReLU NNs, and illustrating how these findings can rewrite an over-parameterized flow as an \"intrinsic\" low-dimensional flow. I find it very intriguing that the commonly reported conservation laws in the literature happen to be maximal (at least empirically).\n\nSTRENGTHS:\nThe manuscript has many strengths and overall I think this is a welcomed contribution to the literature. The manuscript covers various aspects of gradient dynamics, conserved functions, conservation laws, Lie algebra, and their applications in neural networks with illustrative examples. There is also potentially interesting practical consequences of this work. It seems to be a relatively practical approach for determining the number of conservation laws using Lie Group computations, at least on NNs with piecewise linear activation functions. It is interesting that their algorithms confirm (at least empirically) that the conservation laws for ReLU NNs match the laws already known.\n\nWEAKNESSES:\nA few weaknesses are: \n\n- Mostly restricted to deep shallow NNs, continuous-time gradient descent, and simple NN architectures. This limits the applications of the theory to practical situations. The continuous-time restriction on the gradient descent training algorithm is perhaps the more \n\n- In most situations, the generated Lie algebra is going to be infinite-dimensional.  In fact, the two examples in the manuscript are contrived so that the Lie algebra ends up being finite-dimensional. The discussion on the case when the Lie algebra is infinite-dimensional, is only briefly discussed. I would suggest that the author discussed this more. In particular, the stopping criteria are based on the trace of Lie group algebra. \n\n- The above discussion is particularly important for hoping to apply these techniques on NNs with activation functions that are not piecewise linear.\n\nQUESTIONS:\n- Can you add more discussion on the situation where the Lie algebra is infinite-dimensional? \n\n- Can you add more discussion on what happens when the activation functions of the NN are not piecewise linear?\n\nLIMITATIONS:\nThere are no potential negative societal impacts of this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the conservation laws, which are quantities that remain constant, in over-parametrized gradient flows. The authors provide a formal definition for independent conserved functions, which are required to have linearly independent gradients. By applying Frobenius theorem, the authors show that the number of independent conservation laws is linked to the dimension of the trace of the Lie algebra generated by the vector fields spanned by the Jacobian. When this vector field is a Lie algebra, Frobenius theorem can be applied directly to obtain the number of independent conservation laws. When the vector field is not a Lie algebra, the generated Lie algebra need to be computed before applying Frobenius theorem. The authors explicitly compute the Lie algebra and its trace for two layer linear networks and certain ReLU networks, obtained the number of independent conservation laws, and prove that the conservation laws discovered in previous literature are complete. They implement their algorithm in SageMath that constructs a basis of polynomial conservation laws for the above examples, and successfully verify the number of independent conservation laws.\n\nSTRENGTHS:\n- This paper is the first to formally define and study the number of independent conservation laws in gradient flow. Previous works mostly focus on finding conserved quantities in different architectures or using them in convergence proofs. This paper provides a new perspective and contributes to a unified framework of conservation laws in gradient flows.  \n- Using Frobenius theorem to characterize the number of independent conservation laws is novel. By linking to the dimension of the trace of the Lie algebra generated by the Jacobian, the authors present the first known method to determine the number of conservation laws. This method yields the interesting result that known conservation laws are complete in 2-layer linear networks.\n- The idea that conserved functions define invariant hyper-surfaces which trap the gradient flow is interesting and useful. Studying the dimension of these surfaces directly leads to the proposed definition of independent conserved functions.\n- The authors provide a condition under which gradient flows can be recast as low-dimensional Riemannian flows (proposition 3.8), which has potential applications on choosing initializations for better convergence.\n\nWEAKNESSES:\n- The paper’s contribution is overall limited in the aspect of applications. Explicit conservations laws are only given for two-layer linear networks and certain two-layer ReLU networks. The analysis is applied to continuous gradient flow only, and there is no discussion or experiment that verify how well the conservation laws hold in gradient descent. Many neural networks today have more complicated architectures, such as residual connections and various activations other than ReLU, and are often trained with different optimization algorithms, such as Adam. Therefore, while this paper is a promising start to understand implicit bias, more work is needed to obtain insights useful for common machine learning tasks.\n- The abstract vaguely mentions “understanding desirable properties of optimization initialization in large machine learning models”, but the paper provides little supporting arguments. It is not clear what the desirable properties are, and whether it is possible to extend the conservation laws to more realistic settings in large models.\n- The requirement to factor the cost in equation 2 seems strict - $f$ cannot depend on $\\theta$ and $\\phi$ cannot depend on the data and the loss $l$. Factorization for two-layer ReLU network (Example 2.2) is a good example that extends beyond linear networks. However, it is not clear whether this is possible with other activation functions, where the pre-activation is not piecewise linear.\n- There are a few cases where definitions and theorems are mentioned well before the formal statement, for example, in line 134-135, 156-157, 168, 270, etc. Perhaps the organization could be improved to reduce the complexity of the logic flow.\n\nQUESTIONS:\n- The factorization for two-layer ReLU network requires that $\\epsilon_{j,x_i} = \\mathbb{1} (v_j^T x_i) > 0$ is constant. How likely does this condition hold throughout the gradient flow? \n- Would it be possible to include a brief summary of what the Frobenius theorem is about? This theorem appears to be an important foundation, but the form used in this paper (theorem A.12) appears different from the theorem in the given reference [10] (“Theorem 1.4.1. A nonsingular distribution is completely integrable if and only if it is involutive.”)\n- Conservation laws is also an important concept in physics. Is the algorithm that constructs conservation laws related to methods in physics, such as the the Noether’s theorem or conserved quantities from the Killing vector field? Has there been similar analysis on the number of conservation laws for physical systems?\n\nLIMITATIONS:\nThe authors included limitations by clearly stating the assumptions. There are no potential negative societal impacts of the work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies conservation laws for deep neural networks when trained under gradient flow. Here, conservation laws refer to functions of network parameters that are invariant under gradient flow and such laws can potentially help us understand the training dynamics but constraining the manifold of parameters to a low-dimensional space and define robust symmetries of the underlying flow. This line of research had proven to be extremely fruitful for other areas of science, such as physics. The paper studies the number of such conservation laws for generic loss functions and datasets by factorizing the network function. They derive an analytical recipe to derive these numbers and provide explicit examples for simple network architectures. Finally, they provide an algorithm to compute these numbers.\n\nSTRENGTHS:\nThe paper is mathematically engaging, well-written and the content is presented clearly.\n\nWhile previous work extensively studied conservation laws, these work often restricted to specific architectures such as deep linear networks and shallow ReLU networks. This study, as far as I am aware, is the first to generically study symmetries in training dynamics for arbitrary loss functions and datasets, and can be considered as a first step towards understanding the implicit constraints coming from symmetries of various network architectures.\n\nI believe this line of research may be impactful in many problems in DNN community such as pruning deep neural networks, principled approaches to building DNNs and more efficient training strategies.\n\nWEAKNESSES:\n1. The paper is highly technical and only provides generic mathematical tools without brining additional insights over previous findings. Space permitting, it would be helpful to show more examples that go beyond what is already known.\n\n2. The main contribution of the paper is to derive the number of conserved quantities of a given neural network. There is no comment on the explicit constructions of such quantities in generic cases, and no application cases where these numbers can be helpful.\n\n3. The main text lacks experiments, and the ones discussed in supplementary material are not sufficient. Explicit demonstration of conservation laws in simple neural network training might strengthen the paper.\n\n4. The definition of the main algorithm is obscure. A step-by-step implementation might help for clarity.\n\nQUESTIONS:\n1. At line 120, if I am not mistaken, the fidelity function should be $f(\\phi) = \\sum_i \\ell\\left(\\sum_{j,l} \\varepsilon_{j, x_i} \\phi_{j,k,l}  (x_i)_l, y_i\\right)$, i.e. there is no summation over $k$ and the input is forgotten. Depending on how the authors feel, a subscript $f_\\Omega$ can be added to emphasize locality.\n\n2. In Eq. 3, a more generic expression should include the explicit derivative of the data fidelity function, since gradient flow may take us out the domain $\\Omega$ for which $df/d\\phi = \\partial_\\phi f$. Or are you implicitly assuming that infinitesimal gradient flow guarantees such deviations (for example in lazy learning regime)?\n\n3. The statement at line 168 is ambiguous to me; do you mean $\\nabla h \\perp \\chi, \\forall \\chi \\in V$?\n\n4. It seems to me that the effect of loss function decouples from the analysis, since the arguments only depend on $\\phi$. Is it because the number of conserved quantities do not depend on the loss landscape but only the explicit form?\n\nLIMITATIONS:\nThe applicability of the theory to practical neural networks is very restricted at the moment.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-11"}