Abide by the law and follow the flow: conservation laws for gradient flows
Authors: Sibylle Marcotte
Rémi Gribonval
Gabriel Peyré
Venue: NeurIPS 2023 oral (2023)
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
Sibylle Marcotte ENS PSL Univ.
Rémi Gribonval Univ Lyon, EnsL, UCBL, CNRS, Inria, LIP.
Gabriel Peyré CNRS, ENS PSL Univ. gabriel.peyre@ens
sibylle.mare wow.Llee.marctteeens...
remi wew....rrooaal@iiia...m
Abstract
Understanding the geometric properties gradient descent dynamics key ingredient in deciphering the recent success of very large machine learning mod- els. A striking observation is that trained over- wer-parameteriiee 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 given model (e.g. of a ReLU network with a given architecture) with any training data and any loss. Then we explain how find the maximal number of independent conservation laws by performing finiteddimensional algebraic manipulations on the Lie algebra generated by the Jacobian of the model. Finally, we provide algorithms 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 number of ReLU network architectures that all known laws are recovered by the algorithm, and that there are other independent laws. Such computational tools pave the way to understanding desirable properties of optimization initialization large machine learning models.
1 Introduction
State-of-the-art approaches machine learning rely on the conjunction of gradient-based optimization with vastly "over- parametttetteii architectures. large body of empirical 300 and theoretical [5] works suggest that, despite the ability of these models to almost interpolate the input data, they are still able generalize well. Analyzing the training dynamics of these models is thus crucial to gain better understanding of this phenomenon. Of particular interest to understand what properties of the initialization are preserved during the dynamics, which is often loosely referred as being "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 algebraic computations (namely Lie algebra) associated with the model parameterization (in our framework, this parameterization is embodied by re-parameterization mapping Φ)) and finally by exhibiting algorithms to implement these computations in SageMath 29
Over- peeetparamttei model that over-parann peee-paramtterieed neural networks (with more parameters than training data points), which often trained until perfect interpolation, have impressive generalization properties 60 5 This performance seemingly contradicts classical learning theory [25]. and a large part of the theoretical deep learning literature aims explaining this puzzle. The choice of the optimization algorithm is crucial to the model generalization performance [T0 [4], thus inducing an implicit bias.
Modern machine learning practitioners and researchers have found
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 single neuron, it is actually possible to compute in closed form the implicit bias, which induces some approximate or exact sparsity regularization [TO] 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 g A key hypothesis to explicit the implicit bias often that the Riemannian metric associated to the over- wwer-parameterization is either of Hessian type (10 [), or can be somehow converted be of Hessian type B] which is seemingly always very strong constraint. For example, even for simple two-layer linear models (i.e., matrix factorization) with more than single hidden neuron, the Hessian type assumption does not hold, and no closed form is known for the implicit bias [II]. The work of [17] gives conditions on the over-parameterizatioon for this to be possible (for instance certain Lie brackets should vanish). These conditions (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 dynamies) 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 (.44, 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 to prove the global convergence of the gradient flow under some assumptions. We detail these laws in Proposition 4.1 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 the model [To]: 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 's theorem p3 12 Similar reasoning used by 3I to show the influence of initialization on convergence and generalization performance of the neural network. Our work 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 to show that the conservation laws reported in the literature are complete (there no other independent quantities that would be preserved by all gradient flows).
? S
19 These conservation laws enable for instance
Contributions
We formalize the notion of a conservation law, quantity preserved through all gradient flows given a model architecture (e.g. ReL.U 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 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 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 Corollary 4.4 and Section

---

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; E RR the features and y4 E 3 the targets (for regression, typically with y = IR") labels (for classification) in the case of supervised learning, while y/ can be considered constant for seteervvieedssellsupevviiissssss and Y == (v Prediction is performed by parametric mapping g(0,-) Rm - R (for instance neural network) which is trained by empirically minimizing over parameters E θ S CR. a cost
learning. We denote == (xi);
Ex =
Σ (gg(0, si).
where é is the loss function. In practical examples with linear or ReLU networks, θ is either IRO or an open set of "non-degenerate" parameters. The goal of this paper is to analyze what functions h()) preserved during the gradient flow (the continuous time limit of gradient descent) of Ex,y:
0(t) = -VExxx((((t)) with ((0) O
(2)
A priori, one can consider different "levels" of conservation, depending whether h is conserved: during the optimization of Ex.y y for a given loss E and given data set (24, yi)s or given a loss f, during the optimization of Ex.y for any data set (I,, y) 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 [6]
2.2 Conserved functions
As they are based on gradient flows, conserved functions are first defined locally.
Definition 2.1 (Conservation through flow). Consider an open subset Q S O and a vector field E CS, RD). By the Cauchy-Lipschizz theorem, for each initial condition Oinit. there exists unique maximal solution t E [0, To... + ft, Omit) of the ODE 0(t) = x(((t)) with 8(0) = Omn A function h Ω C RD + IR is conserved on Q through the vector field if h(O(t, Omi)) = h(Oinit) for each choice of Inii 52 and every E [0, Te..). ). is conserved on Q through subset W C² (2, RD) if h is conserved on Ω during all flows induced by all X E W.
In particular, one can adapt this definition for the flow induced by the cost
Definition 2.2 (Conservation during the flow 2 with given dataset). Consider an open subset Ω S θ and dataset (X, Y) such that Ex.y E C (22, R). A function h : C
RD - R is conserved
on Ω during the flow 2 it is conserved through the vector field x(-) = (-).v().
Our goal is to study which functions are conserved during "all" flows defined by the ODE This in turn leads to the following definition.
Definition 2.3 (Conservation during the flow 2 with "any" dataset). Consider an open subset C + and loss ((2, y) such that e(., y) is -i-differentiable for all y E V. A function h Ω C RD + R is conserved on Ω for any data set if, for each data set Y) such that g(.,2;) E C² R) for each i, the function h is conserved on Ω2 during the flow This leads to introduce the family of vector fields:
W = {x(+) 3X, Vi 2., E c² (2, x = vεx,} S C(,
so that being conserved on Ω for any dataset is the
being conserved on Ω through W
The above definitions are local and conditioned on choice of open set of parameters Ω C O. 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 θ for any data set if for each open subset Ω C θ, h is conserved on Ω for any data set.
basic property of CD conserved functions (which proof can be found in Appendix to an "orthogonality" between their gradient and the considered vector fields.
corresponds
3

---

Proposition 2.5. Given a subset W C C¹ ($2, RP), its trace ar 0 E Q is defined as the linear space
(0) = span (x(0) x W) C CR.
A ffunction h E (22, R) is conserved on Ω through W if. and only if Vh( 1 ( v0 E
Therefore, combining Proposition 2.5 and Definition 2.4 the object of interest to study locally
conserved functions is the union of the traces
W U WHO
a with Ω neighborhood of 0
(5)
Corollary 2.6. A function h O + is locally conserved on for any data set if and only if Vh(() 1 W for all 0 θ.
that W can be rewritten as the trace (()) of a simple finite- dimensional functional space W. Meanwhile, we keep the specific notation. For the moment, this set explicitly characterized via the following proposition (which proof can be found in Appendix B
will soon be shown (cf Theorem 2. 14
Proposition 2.7. Assume that for each y E y the loss {(2,y) is Cs( Ci-differentiaabe with respect to E RD. For each 0 Q we have:
W
span {(0og(0,)) V (.((g(0,x),))
(x.9)
where Xg is the set of data points x such that ((. x) is C² di-ifferrnntiblee in the neighborhood of 8. Example 2.8. As 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 0 (u, u)
C R² and the model writes g(0, x.) uux. One can directly check that the function: h(u, v) u² n2 is locally conserved on IR2 for any data set. Indeed in that case
Vh(u, v) = (2u, 2v) T W
span {(Vz, ux) V (,l(g(0,z)) x). = R uu, given that
(r.
the gradient V. Vel(g(0,z)) x), y) is an arbitrary scalar.
In this example we obtain simple expression of W. however in general cases it is not possible to obtain such a simple expression from Proposition 2.77 We will show that in some cases, it is possible to express W as the trace W (0) of simple finite-dimensonnal 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(0, x) can be (locally) factored via reparametrization á as ((((), .). We require that the model g(0, x) satisfies the following central assumption.
Assumption 2.9 (Local reparameterization) There exists d and & E C² Rd) such that: for each parameter Og in the open set θ RD for each X such that 0 + g(0, x)) C2 in neighborhood of 6 there is neighborhood Q of 0g and f(,,x) E C² (O()), such that
ve Ω x(0,z)) f f(((0),
(6)
Note that if the model g(-, 2:) is smooth on then 4 is always satisfied with = id and f(;xx) g(-,z). 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 6 there may be notion of an optimal one.
Example 2010. (Factorization for linear neural networks) In the two-layer case, with neurons, denoting 0 = (U, V) E Raxr, x Rmxr (so that = (n + m)r), we can factorize g(6, x) == UVT x by the reparametrization (0) == UVT & Raxm using [(), x) = Φ x. More generally for g layers, with 0 = (UI,... Ug), we can still factorize g0,,z) = U1 ... Uę using s(0) == Un ... U,, and the same f. This factorization is globally valid on f2 = θ = RD since (() does not depend on fn.
The notion of locality of the factorization 6 is illustrated by the next example.
lie., belongs to the set No, defined in Proposition 2.7

---

Example 2.11 (Factorization for two-layer ReLU networks). Consider 9(0,x)
(Ej=1 ( + b3) + ok) ct) with ot)) = max(t. the ReLU activation func- tion and bj E RRT uks E R. bj. Ck E R. Then, denoting 0 (U,V,b,c) with U = (uag)es =: (ui, .u..., u,) E R V (V1, ,.r E Rm b (b1,... b,)) T R and c = (cI.-., Gn) E R" (so that D = (n + m + 1)r + n), we rewrite g(0, x) = Σj= uyEz,z (o] x b,) to where, given z, EJ,r 1(1) x+ bj= > 0) is piecewise constant with re- spect to 0. Consider 00 = (U", vo, 6", c0) E Rp where V = (ef. ,2") and 6 = (66, ... 6,)T. Then the set Xyo introduced in Proposition 2.7 Xgo -u, z+b] 0}. Let E Xgo. Then on any domain S2 C RD such that 0 E Z and (0) = 1(v > 0) constant over 0 E Q2 the model go (z) can be factorized by the reparamettrizatoon 6(0) ( ,u,b,]ī-1.c) In particular, in the case without bias ((b (0, 0)), the reparamettrizaiion is defined by 6(0) = (6:)ī=1 where rmm) using l( x) Σj UZm0 the reparametrization (0) contains r matrices of size m x n (each of rank at most one) associated to a "local" f(,
= (0) == "zyj E RR (here
valid in neighborhood of 0. 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
Proposition 2.12. Assume that the loss f(z. y) is C² -2-differentiable with respect to z. We recall (cf
5 that W == Ungei is open and Rae W (0). Under Assumption 2.9 for all A E A:
W a5(()) WT TW
(7)
((0)
with 0+( RaxD the Jacobian of & and w! 0(0) (x.))e where f=(()= = z(,
= span
We show in Section 2.4 that, under mild assumptions on the loss E, W 6(0) = Rd, so that Proposi- tion 2.12 yields Wy = range(8e(()) T). Then by Corollary 2.6 function h that locally conserved on θ for any data set is entirely characterized via the kernel of (o( 06(0) T Vh(0) = 0 for all
0 E θ The core of our analysis is then to analyze the (Lie algebraic) structure of range( (+)()) ).
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) as- sumptions the loss ((..), being locally conserved on θ for any data set (according Definition 2.4 is the same as being conserved (according to Definition 2.T on through the finite- dimensional subspace
Vol(-)) = (0- + Σ (o) (a1 (aI.....ad) ad) Rd}
We = /pann//nn
(8)
where we write 00(0) = (O), Vor(O)) E RDxd, with Vol (O,
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, v) + f(z, v) is C2 d'-differentiable with respect to z E R" and satisfies the condition:
span Span(V V=((2. = E2,
(9)
yEY
Then for linear neural networks (resp. for two-layer ReLU networks) and all 0 E θ we have 5 =(0) = Rd. with the reparametrization o from Example 2.10 and O = RD (resp. with o from Example 2.17 and θ consisting of all parameter 0 of the network such that hidden neurons are associated to pairwise distinct "hyperplanes" ", cf Appendix for details).
Condition holds for classical losses E (e.g. quadratic/logistic losses), as shown in Lemma C in
Appendix Note that the additional hypothesis of pairwise distinct hyperplanes for the two-layer
ReL.U case is a generic hypothesis and is usual (see e.g. the notion of twin neurons in (27). The tools from Appendix Meadixxexxxo extend Theorem 2.14 beyond (deep) linear and shallow ReL.U networks. An
open problem is whether the conclusions of Lemma 2.13 still hold for deep ReLU networks.
5

---

Theorem 2.14. Under the same assumptions as in Lemma 2.13 we have that for linear neural networks, for all 0 E θ RD:
W = (-)
(10)
The same result holds for two-layer ReLU networks with & from Example 2.17 and O the (open) set of all purameters 0 such that hidden neurons are associated to puirwise distinct "hyperplanes".
This means as claimed that for linear and two-layer ReLU networks, being locally conserved on Q for any data set exactly means being conserved on @ through the finite-dimensions functional space We S C* (O, RP). This motivates the following definition
Definition 2.15. A real-valued function h is conservation law of & if it conserved through We.
Proposition 2.5 yields the following intermediate result.
Proposition 2.16. h E C12, R) conservation law for & and only if
(h(()) (0) 4 E Q, Vj E (1.....d).
Thanks to Theorem 2.14 the space Wa defined in e introduces a much simpler proxy to express W as a trace of a subset of C* (e, (RR). Moreover, when & is C We is a finite-dimensional space of infinitely smooth functions
θ, and this will be crucial in Section 4.1 to provide 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 B
Example 2.17 Revisiting Example 2.8 the function minimize is factorized by the reparametrization
6 (u E IR, U E R) - we E R with 0 = (u, v). We saw that (uu,, b)) == u² y² is conserved: and indeed (Vh(u, v), Vo(u, v)) 2uv 2um 0, V(u,v)
In this simple example, the characterization of Proposition 2.16 gives constructive way to find such a conserved function: we only need to find a function such that (Vh(u, v), Vo(u, v)) = (Vh(u, v). (v, u) ) 0. The situation becomes more complex in higher dimensions, since one needs to understand the interplay between the different vector fields in Wa.
2.5 Constructibility of some conservation laws
Observe that in Example2.11 both the reparametrization o and the conservation law are polynomials, property that surprisingly systematically holds in all examples of interest in the paper, making possible to algorithmically construct some conservation laws as detailed now.
By Proposition 2.16 E 1- ((V)(0), (o)) Thus, one could look for conservation laws in prescribed umite-dimensiona space by projecting these equations in a basis (as in finite-elemen methods for PDEs). Choosing the dinis-admensional subspace could be generally tricky, but for the linear and ReLU cases all known conservation laws are actually polynomial "balancedness-type con- see Section 4 In these cases, the vector fields in We are also polynomials (because is polynomial. see Theorem C.4 and Theorem C.5 in Appendix C hence 0 + (Vh(0), (o) 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 D all existing known conservation laws both for ReLU and linear networks. Open-source code is available at
function h is conservation law if it is in the kernel of the linear operator
E C² (22, R)
ditions" " 1 2
2.6 Independent conserved functions
Having algorithm build conservation laws is nice, yet how 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 for each k E N but this does not add any other constraint), but rather 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ı,.... hx) conserved through W c C* (22 RP) is said to be independent if the vectors (0), (h) are linearly independent for all 0 E Ω.

---

An immediate upper bound holds on the largest possible number N of functionally independent func- tions h1.... ht....,
...,,,, n conserved through W: for E 52 S RD, the space span (Vh, (o)..... Vhv ( S RD is of dimension N (by independence) and (by Proposition 2.5 orthogonal to W()). Thus, it is necessary to have N S 1D-dim W(0), 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 recurring theme in control theory, since the existence of such surfaces is the basic obstruction of controllability of such systems [ The basic result to study these surfaces is the so-called Frobenius theorem from differential calculus (See Section 1.4 of [%] for a good reference for this theorem). 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, algebraic closure of We under Lie brackets. The fundamental object of interest is thus the Lie algebra generated by the Jacobian vector fields, that recall next. While this is only defined for vector fields with stronger smoothness assumption, the only consequence is that e is required to be infinitely smooth, unlike the loss f(. y) and the model g( z) that can be less smooth. All concretes examples of á in this paper are polynomial hence indeed infinitely smooth.
needs to consider the
Notations
Given vector subspace of infinitely smooth vector fields W C (()) = C (O, (O,
where θ is an open subset of RD, we recall (cf Proposition 2.5) that its trace at some @ is the subspace
W (0) == span x E W) C RR.
(11)
For each open subset ? C θ, we introduce the subspace of (I)): Win = (x)::
3.1 Background on Lie algebra
A Lie algebra A is a vector space endowed with a bilinear map I. called a Lie bracket, that verifies for all X, Y, Z E A: [X = 0 and the Jacobi identity: [X, z + [Y, [Z, X]] + [Z, YJ] 0.
For the purpose of this article, the Lie algebra of interest is the set of infinitely smooth vector fields (O), endowed with the Lie bracket [,,-] defined by
xx]: 0 θ + (0) == xxx (0)x2(0) - 0x2(0)xx
(12)
with (x) E RDx the jacobian of x 8. The space R*x of matrices is also 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. x0) A0.
Generated Lie algebra Let A be a Lie algebra and let W C A be a vector subspace of A. There exists 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 = (xx1, xa) X1 W, X2 W)), and W + W [xi X2 X1 E W, X2 E W').
Proposition 3.1. Given any vector subspace W C A we have Lic (W) U2 We where: W
{
Wo
We
= W (.o WA-1) for k 1.
We will see in Section 3.2 that the number of conservation laws is characterized by the dimension of the trace Lie(W,) & (0) defined in [.p. The following lemma (proved in Appendix D) gives a stopping criterion to algorithmically determine this dimension (see Section 3.3 for the algorithm).
Lemma 3.2. Given e O, for a given i, dimW,-1 (0') = dimw Baa.w,OL (Or every o' in a neighborhood of 0, then there exists a neighborhood Ω of 0 such that (0' (,(0') for all 0' E Ω and k > i, where the V are defined by Proposition 3.7 Thus Lie(W)(0") = W((0') for all 0' E St In particular, the dimension of the trace of Lic(W) is locally constant and equal to the dimension of W.(O).
7

---

3.2 Number of conservation laws
The following theorem uses the Lie algebra generated by Wa to characterize the number of conserva- tion 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 (li((Lic(W.))) (0)) is locally constant then each 0 E Ω S RD admits neighhorhood S such that there are D - dim (Lie( W. (0)) (and no more) independent conserved functions through W Wa Le., there are D (iim(Li (0) independent conservation laws of á on S
Remark 3.4. The proof of the Frobenius theorem (and therefore of our generalization Theorem 3.3 is actually constructive. From given 6, conservation laws obtained in the proof by integrating in time (i.e. solving an advection equation) the vector fields belonging to Ws. 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 Fourier expansions neural networks).
A fundamental aspect of Theorem is to rely only the dimension of the trace of the Lie algebra associated with the finite-dimensional vector space Wa- Yet, even if Wa is finite-dimensional, it might be the case that Lie(Wa) itself remains infinitedimmeniionall Nevertheless, what matters is not the dimension of Lie(W,), but that of its trace Lie) which is always finite (and potentially much smaller that dim 2 even when the latter is finite) and computationally tractable thanks to Lemma as detailed in Section 3.3 In section 4.1 we work out the example of matrix factorization, a non-urvvaal case where the full Lie algebra Lie () itself remains finite-dimensional.
Theorem 3.3 requires that the dimension of the trace at @ of the Lie algebra is locally constant. This is technical assumption, which typically holds outside 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 reparametrization ó for the architectures train, to determine the number of independent con- servation laws of é, we leverage the characterization 3.1 to algorithmically compute
Lie (0)) using an iterative construction of bases for the subspaces WA starting from Wo == W2. and stopping as soon as the dimension stagnates thanks to Lemma 3.2 Our open-sourced code is available at [ 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 = Wi = ()) == W (0) on Q. This is the case if and only if We satisfies that
[XI,X (0) == 0x (0)2 0x2
E (o),,, for all X1. X2 E We and all 0 E Ω. (13)
i.e, when Frobenius Theorem (see Theorem in Appendix B applies directly. The first example is
a follow-up to Example 2.11
Example 3.5 (two-layer ReLU networks without bias). Consider = VU,V) with U E Raxr V e (so that
Rmxr, n. m, r
(n + m))), and the reparametrization ((0) == (u.w. l= E
Raxmxr, where U = (u1: ... ; u,) and V (V::...:uu ;.)). As detailed in Appendix F.1 since ((0)
a collection of rank-one n x m matrices, (Wm on the domain & such that Wa A 0, and We, satisfies
rankdo(0) = (n + m Dr is constant
hence by Theorem 3.3 each 0 has
neighborhood S2' such that there exists r (and no more) independent conserved function through
Wa w The known conserved functions [9] given by h; (U,V)+ ||u,]| - ||:: || 1, ,
independent, hence they are complete.
Our Example where the iterations of Lemma 3.2 stop at the second step (but not the first one). primary example is matrix factorization, follow-up to Example 2.10 2.10
Example 3.6 (two-layer linear neural networks). With 0 = (U, V). where (U E RuXr, V E RR)x) the reparameterization (0))= = UVT E R=*m (here d = nm) factorizes the functions minimized during the training of linear two-layer neural networks (see Example 2.10p. As shown Appendix condition (3) is not satisfied when > 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 HL in Appendix
(Wo)1 W ) Lie(W,). o), hence the iterations of Lemma 3.2 stop at the second step.
8

---

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 this settings and extended ones is further studied in Section 4 and Appendix
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 dim (0))
dim(LieW, (0)) is constant on S2, i.e., ran)(a is constant on f2 and We satisfies 13 is to fully
rewrite the high-dimensional flow {(t) E RD as a now-dimensional flow on 2(t) = ((((t) R where this flow is associated with Riemannian metric tensor M that induced by ø and depends on the initialization min We insist on the fact that this is only possible in very specific cases, but this phenomenon underlying many existing works that aim at writing in closed form the implicit bias associated with some training dynamics (see Section D for some relevant literature. Our analysis sheds some light on cases where this possible, as shown in the next proposition.
Proposition 3.7. Assume that rank(Oo(()) is constant on Ω and that We satisfies
L If0(t) E R
satisfies the ODE ( where m S2 then there 0 < T Ti S Tomms such that x() ((((t)) E Rd satisfies the ODE
(t) = (x((t), m Vf(())(((() for all t t [o,Tomm) with z(0) = @00m)),
(14)
where M (z(t), 0mi) Raxd is a symmetric positive semi--efinite matrix.
See Appendix G a proof. Revisiting Example 3.5 leads to the following analytic example.
Example 3.8. Given the reparametrization o : (u E IR u E Rd)) 7 uv E Rd, the variable
== MU satisfies
14 with Mzz (mitm I + ||=||= 22 with | = ā + / |
5 = 1/2 init
Uimit P).
Another analytic example is discussed in Appendix G In light of these results, an interesting perspec- tive 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 (2, Omiis can have kernel. Indeed, in practice, while & is function from RD to Rd. the dimensions often satisfy (ankôø(0) < min(d, D). i... ((0) lives in manifold of lower dimension. The evolution (14 should then be understood as a flow on this manifold. The kernel of M (z, ma)) 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 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 0 = (U1, ... Ug) with U; E R= the weight matrices and we assume that 0 satisfies the gradient flow 2 In the linear case the reparametrization is Φlnn(0) (0) == U1 ... .U. For ReLU networks, we use the (polynomial) reparametrization daeta of [ Definition 6], which defined for any (deep) feedforward ReLU network, with or without bias. In the simplified setting of networks without biases it reads explicitly as:
@netu (U1, .Ug) = (U (U1.: (1 fal ... UU- 1 iq- 101 Jer" -e-1
(15)
with U [i the j-th entry of U. This covers ((0) (nzv, Di=1 Ruxmxx from Example 2.11
Some conservation laws are known for the linear case PLin I
2 and for the ReLu case dReta 9
Propasition 4.1 ( ( ). Iff (U1,...,U% satisfies the gradient Q for each 1.... .9- 1 the function 0 - U. U, U,41U (resp. the function + diag (U, U. U++U UT) defines n, (n: + 1)/2 conservation laws for dran (resp. n, conservation laws for @Raia).
Proposition defines Σ ni x (n. + 1)/2 conserved functions for the linear case. general they are not independent, and we give below in Proposition 4.2 for the case of 2, the exact

---

number of independent conservation laws among these particular laws. Establishing whether there are other (previously unknown) conservation laws an open problem for q > We already answered negatively to this question in the two-layer Rel lu case without bias (See Example Section (Corollary (.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
In the following
4.1 The matrix factorization case (q = 2)
To simplify the analysis when g = 2. we rewrite = (U, V) as a vertical matrix concatenation denoted (U; V) R(ntm)xx, and (0)) - Sr.sn (0) = UVT Raxm.
How many independent conserved fuuctioon 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. for proof.
- - E Rexy and assume that (U; V) has full rank noted rk. Then the function Va gives rk. (2r + r2)/2 independent conserved functions.
Proposition 4.2. Consider V 0 = (U;V) +
There exist no more independent conserved functions. We now come to the core of the analysis, which consists in actually computing Li) as well as its traces (0) in the matrix factor- ization case. The crux of the analysis, which enables us to fully work out theoretically the case 2, is that We is composed of linear vector fields (that are explicitly characterized in Proposition 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 sie(Ws)(U for any (U, V). We prove the following in Appendix H
in
Proposition 4.3. If (U;V) E R((+t))x has full rank noted rk, then: dim(Lie(W:)(U:: (n + m)r (2r + 1 r2)
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 V (U U U - 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 2. For deeper linear networks and for ReLU networks, the vector fields in We are non-linear polynomials, and computing Lie brackets of such fields can increase the degree, which could potentially make the generated Lie algebra
One can however use Lemma 3.2 and stop as soon as dim ((W)) o)xx (f)) stagnates. Numerically comparing this dimension with the number N of independent conserved functions known in the literature (predicted by Proposition [.J 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 ! for details). Our code is open-sourced and is available at [1]] 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 any space of displacements in place of Wg. 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. 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.
10

---

Acknowledgement
The work of G. Peyré was supported by the European Research Council (ERC project NORIA) and the French government under management of Agence Nationale de la Recherche as part of the "Investissements d' d'avenir" program, reference ANR-19-P3IA work of R. Gribonval was partially supported by the AllegroAssai ANR project ANR-19 19-CHIA--000 We thank Thomas Bouchet for introducing us to SageMath, as well as Léo Grinsztajn for helpful feedbacks regarding the numerics. We thank Pierre Ablin and Raphaël Barboni for comments draft of this paper. We also thank the anonymous reviewers for their fruitful feedback.
(PRAIRIE 31A Institute). The
References
[1] S. ARORA, N. COHEN, N. GOLOWICH, AND W. HU, convergence analysis of gradient descent for deep linear neural networks, arXiv preprint arXiv:1810.02230 (2018).
[2] 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. 244- 253.
[3] S. AZULAY, E. MOROSHKO, M. S. NACSON, B. E. WOODWORTHI N. SREBRO, A. GLOBER- SON, AND D. SOUDRY, On the implicit bias of initialization shape: Beyond infinitesimal mirror descent, in Proceedings of the 38th International Conference Machine Learning, M. Meila and T. Zhang, eds., vol. 139 of Proceedings of Machine Learning Research, PMLR, 18-24 Jul 2021, pp. 468-477.
[4] B. BAUL H. RAUHUT, U. TERSTIEGE, AND M. WESTDICKENBERG, Learning deep linear neural networks: Riemannian gradient and convergence to global minimizers, Information and Inference: A Journal of the IMA, 11 (2022), pp. 307-353.
[5] M. BELKIN, HSU, S. MA, AND S. MANDAL, Reconciling modern machine-learning practice and the classical 116 (2019), pp. 15849-15854,
trade-off, Proc. of the National Academy of Sciences,
[6] B. BONNARD, M. CHYBA, AND J. ROUOT, Geometric and Numerical Optimal Control - Appli- cation to Swimming at Low Reynolds Number and Magnetic Resonance Imaging, SpringerBrief's in Mathematics, Springer Int. Publishing, 2018.
[7] A. B. BROWN, Functional dependence, Transactions of the American Mathematical Society, 38 (1935), pp- 379-394.
[8] 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.
[9] 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).
[10] S. GUNASEKAR, LEE, D. SOUDRY, AND N. SREBRO, Characterizing implicit bias terms of optimization geometry, in Int. Conf. on Machine Learning, PMLR, 2018, pp. 1832-1841.
[11] S. GUNASEKAR, B. E. WOODWORTH, S. BHOIANAPALLI, B. NEYSHABUR, AND N. SREBRO, Implicit regularization in matrix factoriztion,,,, Adv. in Neural Inf. Proc. Systems, 30 (2017).
[12] G. GLUCH AND R. URBANKE, Noether: The more things change, the more stay the same, 2021.
[13] A. ISIDORI, Nonlinear system control, New York: Springer Verlag, 61 (1995), pp. 225-236.
[14] Z. J1. 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. 2109- 2136.
[15] Z. AND M. TELGARSKY, Gradient descent aligns the layers of deep linear networks, arXiv preprint arXiv:1810.02033 (2018).
11

---

[16] D. KUNIN, J. MEAASTTYYBBRNA S. GANGULI, D. L. YAMINS, AND H. TANAKA, Neural mechanies: Symmetry and broken conservation laws in deep learning dynamics, arXiv preprint Siii:2012.047255 (2020).
[17] Z. LI, T. WANG, J. D. LEE, AND S. ARORA, Implicit bias of gradient descent reparametrized models: On equivalence to mirror descent, in Advances 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. 16626-34640
[18] S. MARCOTTE, R. GRIBONVAL, AND G. PEYRÉ, Code reproducible research. Abide by the Law and Follow the Flow: Conservation Laws for Gradient Flows, Oct. 2023.
[19] H. MIN, S. TARMOUN, VIDAL, AND 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.
[20] W. F. NEWNS, Functional dependence, The American Mathematical Monthly, (1967), pp. 911-920.
[21] B. NEYSHABUR, Implicit regularization in deep learning, arXiv preprint arXiv: 1eXiv:17090 709900195. (2017).
[22] B. NEYSHABUR, R. TOMIOKA, AND N. SREBRO, In search of the real inductive bias: On the role of implicit regularization deep learning, arXiv preprint 141Yv:141..6614, (2014).
[23] E. NOETHER, Invariante variationsprobleme, Nachrichten von der Gesellschaft der Wis- senschaften zu Göttingen, Mathematisch--Phyikklliie Physikalissiei Klasse, 1918 (1918), pp. 235-257.
[24] A. M. SAXE, J. L. MCCLELLAND, AND S. GANGULI, Exact solutions to the nonlinear dynamics of learning deep linear neural networks, arXiv preprint arXiv: 11Xv:1312..6100, (2013).
[25] S. SHALEV-SHWARTZ AND algorithmss, Cambridge university press, 2014.
Understanding machine learning: From theory to
(26) D. SOUDRY, E. HOFFER, S. NACSON, S. GUNASEKAR, N. SREBRO. The implicit bias of gradient descent on separable data, The Journal of Machine Learning Research, 19 (2018), pp. 2822-2878.
[27] STOCK AND R. GRIBONVAL, An Embedding of ReLU Networks and an Analysis of their Identifiability, Constructive Approximation, (2022). Publisher: Springer Verlag.
[28] S. TARMOUN, G. FRANCA, B. D. HAEFTELE, AND R. VIDAL, Understanding the dynamics of gradient flow in overparameterized linear models, in Int. Conf. on Machine Learning, PMLR, 2021, pp. 10153-10161.
[29] THE SAGE DEVELOPERS, SageMath, the Sage Mathematics Software System (Version 9.7), 2022. https: //wp::://www... www.mathh org.
[30] C. ZHANG, S. BENGIO, M. HARDT, B. RECHT, AND O. VINYALS, Understanding deep learning requires rethinking generalization, Int. Conf. on Learning Representations, 2017.
[31] B. ZHAO, I. GANEV. R. WALTERS, YU, AND DEHMAMY, Symmetries, flat minima, and the conserved quantities of gradient flow, arXiv preprint arXiv:2210.1700 17216, (2022).
12


Reviews:
[{"review_id": "REZDJVQ72W", "reviewer": "Reviewer_n1Hr", "summary": "The 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.", "questions": "1. 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?", "limitations": "The applicability of the theory to practical neural networks is very restricted at the moment.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive comments and constructive suggestions.\n\n**Weaknesses and Questions addressed**\n\n> **W1a** The paper is highly technical and only provides generic mathematical tools without bringing additional insights over previous findings. Space permitting, it would be helpful to show more examples that go beyond what is already known.\n\nBeyond ReLU and linear networks, which are 1-homogeneous, our results also apply to networks using any positively $p$-homogeneous activation function. Specifically, for a 2-layer NN given by $g_{\\theta}(x) = U \\sigma(V^\\top x) $, the conserved quantities are $\\|u_i\\|^2 - \\|v_i\\|^2/p$, where $u_i$ and $v_i$ are the columns of $U$ and $V$, respectively. We plan to include this example in the supplementary material of the final version. While our framework and its generalization (refer to **W2a**) can cover more architectures, discussing them is beyond the scope of this paper.\n\n> **W2a** 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\nWe believe that determining this number is a fundamental question in the study of neural networks, as it can put an end to the “quest” for potential additional laws. We disagree with the perception of a “lack of comment on explicit constructions”: we provide explicit algorithms to compute both the number of conservation laws and the laws themselves, particularly when the models are polynomial. One significant application of these conservation laws is their ability to demonstrate that, in certain cases, high-dimensional flows can be recast in lower dimensions, see Section 3.4. Another application, which we will further emphasize in the paper, aids in convergence proofs; for instance, Theorem 5 of [6] and Section 2.5 of [*], utilize balancedness conditions which are conservation laws. While it is beyond the scope of this paper, our theory could easily be applied to other architectures, including residual connections, convolutional layers, and piecewise polynomial activations.\n\n[*] On the Convergence of Gradient Descent Training for Two-layer ReLU-networks in the Mean Field Regime, S. Wojtowytsch, 2020, preprint\n\n> **W3a** 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\nNumerical illustrations of conservation laws can be found, for instance, in Figure 1 of [7] for a 2-layer linear NN, in Figure 1c of [1] for a 3-layer linear NN, and in Figure 2 of [7] for a 3-layer ReLU NN. We will include these references. In the final version of the paper, we will incorporate such a figure into the supplementary material to further emphasize our main message.\n\n> **W4a** The definition of the main algorithm is obscure. A step-by-step implementation might help for clarity.\n\nThank you for the suggestion, we will add a pseudo-code in the final version to clarify it.\n\n> **Q1a** At line 120, if I am not mistaken, the fidelity function should be […] 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\nThank you, indeed there was a typo, we will correct this with a simpler (and correct!) expression using\n$ \\phi_j = \\phi_j(\\theta) := u_jv_j^\\top \\in \\mathbb{R}^{n \\times m}$, $\\phi(\\theta) = (\\phi_j )\\_{j}$ and \n$ f(\\phi) = \\sum\\_i \\ell( \\sum\\_j \\epsilon_{j,x_i} \\phi_j x_i, y_i)$.\n\n> **Q2a** 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\nIndeed it is a good idea to clarify by first writing that the gradient flow on $\\mathcal{E}$ is defined as $\\dot{\\theta}(t) = -\\nabla \\mathcal{E}(\\theta(t))$. Since we assume the factorization $\\mathcal{E} = f \\circ \\phi$ on a neighborhood of $\\theta_0$, for sufficiently small $t$ we deduce that the gradient flow satisfies Eq. 3. Our analysis of conservation laws is *local* to avoid considering what happens when we leave the domain.\n\n> **Q3a** The statement at line 168 is ambiguous to me; do you mean $\\nabla h \\perp \\chi, \\forall \\chi \\in V$?\n\nThe statement that we require is stronger: it means that $\\nabla h(\\theta) \\perp \\chi(\\theta), \\forall \\chi \\in V, \\forall \\theta \\in \\Omega$. In other words, it is a pointwise assumption: at every point $\\theta$, the vector $\\nabla h(\\theta) \\in \\mathbb{R}^D$ is orthogonal to the subspace $V(\\theta) \\subseteq \\mathbb{R}^D$. \n\n> **Q4a** 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\nAs summarized in ll 125–129, the main idea is indeed to decouple as much as possible the study of the conserved functions from the particularities of a dataset or a given loss. This is made possible when the factorization $\\mathcal{E} = f \\circ \\phi$ holds, and under some assumptions on the loss (see eq.7).\n\n"}, {"review_id": "qyzjyoMwWn", "reviewer": "Reviewer_wzMU", "summary": "This 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", "questions": "- 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? ", "limitations": "The authors included limitations by clearly stating the assumptions. There are no potential negative societal impacts of the work.\n", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- 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", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive comments and constructive suggestions.\n> **W1b** The paper’s contribution is overall limited in the aspect of applications [...]\n\nFor ReLU/linear networks *of any depth*, explicit conservation laws of $\\phi$ are known (Prop 4.1) and our algorithms allow us to verify on a number of *deep (q>2)* ReLU/linear networks that there are no other ones (section 4.2). The only ingredient of our framework that is not yet extended to deeper (q>2) *ReLU* architectures is Theorem 2.8, which ensures that the conservation laws of $\\phi$ (computable with our algorithms) are indeed exactly the conservation laws shared by all $\\mathcal{E}$, for any dataset and loss. For linear networks, Theorem 2.8 is valid irrespective of their depth.\n\nExperiments showing approximate conservation of these laws during gradient descent (instead of gradient flow) are given in figure 1 of [7] for 2-layer linear NN, figure 1c of [1] for 3-layer linear NN, and figure 2 of [7] for a 3-layer ReLU NN. We will add a more explicit comment on this fact. Regarding other optimization algorithms, while our framework leaves completely open the question of a similar analysis for stochastic algorithms, it seems feasible to adapt it to deterministic algorithms with momentum using their associated ODE. This is however out of scope and left for further work. \n\nBeyond standard multilayer linear/ReLU architectures covered in this work, it seems feasible but technical (and beyond the scope of this paper) to show that more general ReLU architectures covering e.g. residual connections and convolutive layers, satisfy such a factorization with polynomial $\\phi$. The main challenge for a follow-up is then to extend Theorem 2.8, which would make the whole framework applicable in such contexts.\n\nFor more general activation functions, our results directly apply when using any positively $p$-homogeneous activation function. For a 2-layer NN defined as $g_{\\theta} (x) = U \\sigma(V^\\top x)$, the conserved quantities are given by $\\|u_i\\|^2 - \\|v_i\\|^2/p$, where $u_i$ and $v_i$ are the columns of $U$ and $V$ respectively. We plan to incorporate this example in the supplementary material of the final version. A compelling direction for future exploration would be to extend these results to piecewise linear and piecewise-polynomial activations.\n\n> **W2b** The abstract vaguely mentions [...]\n\nIn general terms, \"desirable properties\" of initialization refer to characteristics that ensure convergence, and potentially faster convergence, to an optimal solution. For instance, the utilization of conservation laws in [6] and [*] demonstrates convergence, while [3] illustrates that initializing with certain values of the conserved function can lead to accelerated convergence. We will clarify this further in the final version of the paper. Also, refer to our response to **W1b**.\n\n[*] \"On the Convergence of Gradient Descent Training for Two-layer ReLU-networks in the Mean Field Regime\" by S. Wojtowytsch, 2020, preprint.\n\n> **W3b** The requirement to factor the cost in eq. 2 seems strict [...]\n\nEq. 2 is in fact not very demanding: we can always write $\\mathcal{E} = f \\circ \\phi$ with $f= \\mathcal{E}$ and $\\phi = id$. However the number of conservation laws of $\\phi= id$ is zero, and this trivial factorization fails to capture the existence and number of conservation laws as studied in this paper. This suggests that, among all existing factorizations $\\mathcal{E} = f \\circ \\phi$, there may be a notion of an optimal one, such that an equivalent of Theorem 2.8 holds. This is an interesting challenge for future work. Thanks for this opportunity to clarify this point that will be mentioned explicitly.\n\nRegarding the ability to handle other activation functions: piecewise linearity of $x \\mapsto g_\\theta(x)$ (as in the linear and ReLU cases) is not important, and extensions e.g. to positively $p$-homogeneous activations (e.g. the squared ReLU) can be achieved [see our response to **W1b**]. \n\n> **W4b** There are a few cases [...]\n\nThank you for this suggestion, we will keep it in mind for the final version.\n\n> **Q1b** The factorization for 2-layer ReLU network requires that $\\mathbb{1}(v\\_j^\\top x_i > 0)$ is constant. How likely does this condition hold throughout the gradient flow?\n\nThis condition will *not* be preserved throughout the gradient flow, however as soon as  $\\mathbb{1}(v\\_j^\\top x_i > 0)$ is *locally* constant we can conduct our analysis. Thus, apart from some instants along the trajectory where there are changes in these activations, the whole analysis is valid and allows to characterize the (number of independent) functions that are conserved. \n\n> **Q2b** Would it be possible to include a brief summary of what the Frobenius theorem is about? [...]\n\nThank you for the suggestion! In the final version's supplementary material, we'll include a summary as you've recommended. Additionally, we'll provide a section that clarifies the translation of notations and vocabulary between the theorem mentioned in the given reference and our paper. When we refer to a \"non-singular distribution\", it implies that the dimension of the associated trace remains constant (refer to the definition of \"non-singular\" on page 15 of [10]). Being \"involutively consistent\" directly relates to our second assertion using the Lie bracket (see eq. 1.13 on page 17 of [10]). Lastly, \"completely integrable\" aligns with our first assertion regarding orthogonality conditions (refer to eq. 1.16 on page 23 of [10]).\n\n> **Q3b** Conservation law is also an important concept in physics [...]\n\nOur theorem is indeed related to invariance in the model (each invariance such as scaling in ReLu is associated to a conserved quantity). We will mention and clarify this connexion in the revised version. This being said, we were not able to draw a precise connexion with Noether theorem, and the settings where Noether vs Frobenius apply seem rather different."}, {"review_id": "yVui3Eqyqu", "reviewer": "Reviewer_Tpzp", "summary": "The 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).", "questions": "- 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? ", "limitations": "There are no potential negative societal impacts of this work. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The 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. ", "weaknesses": "A 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive comments and constructive suggestions.\n\n**Weaknesses and Questions addressed**\n\n> **W1c** 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\nExperiments that display the approximate conservation of these laws during the process of gradient descent (as opposed to gradient flow) can be found in Figure 1 of [7] for a 2-layer linear NN, in Figure 1c of [1] for a 3-layer linear NN, and in Figure 2 of [7] for a 3-layer ReLU NN. We plan to include a more detailed commentary on this observation.\n\nIn addition to the standard multilayer linear/ReLU architectures discussed in this paper, preliminary studies, beyond the confines of this work, suggest that more diverse ReLU architectures, which encompass aspects like residual connections and convolutional layers, adhere to such a factorization with polynomial $\\phi$. With a suitable adaptation of Theorem 2.8, our entire framework should be applicable in these scenarios.\n\n> **W2c** 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\nGenerally speaking, the Lie algebra generated can indeed be of infinite dimension. This is particularly the case for deeper linear neural networks where $q > 2$, a point we intend to address in the final version. However, the *trace* of the generated Lie algebra is invariably finite-dimensional—it is bounded by $D$, the total number of parameters—and this trace is our focal point when considering the stopping criterion. As an analogy, the set of all smooth real-valued functions constitutes an infinite-dimensional Lie algebra, yet its trace at any given point corresponds to the finite-dimensional space $\\mathbb{R}$, thus having a dimensionality of one. Given that the trace has a finite dimension, we can deduce a basis for it. Consequently, the stopping criterion will be met in a maximum of $D$ steps. We will emphasize this distinction in our work.\n\n> **W3c**, **Q1c**, **Q2c** The above discussion is particularly important for hoping to apply these techniques on NNs with activation functions that are not piecewise linear. Can you add more discussion on the situation where the Lie algebra is infinite-dimensional?\nCan you add more discussion on what happens when the activation functions of the NN are not piecewise linear?\n\nOur theory readily accommodates an infinite-dimensional Lie algebra, given that its trace remains finite-dimensional. As an illustration, for deeper linear networks (where $q > 2$), the Lie algebra does become infinite-dimensional. Yet, our theory remains applicable, especially since theorem 2.8 is valid for linear networks irrespective of their depth. We will elucidate this point in the final version. When dealing with more intricate activation functions, our results are directly applicable for any positively $p$-homogeneous activation function. Specifically, for a 2-layer NN represented as $g_{\\theta} (x) = U \\sigma(V^\\top x)$, the conserved quantities are given by $\\|u_i\\|^2 - \\|v_i\\|^2/p$, where $u_i$ and $v_i$ denote the columns of $U$ and $V$ respectively. We plan to incorporate this example in the supplementary material of the final version. An intriguing avenue for exploration would be to extend these findings to encompass piecewise linear and piecewise-polynomial activations.\n"}, {"review_id": "p39xfhCQxD", "reviewer": "Reviewer_7TLM", "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.", "questions": "Do the results stated for ReLU networks hold for any network with homogeneous activation function?", "limitations": "authors discussed the limitation.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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.\n2. This paper has an in-depth discussion of the conservation law under the gradient flow on a class of loss functions.\n3. Analysis via Lie algebra that determines the maximum number of conservation laws.\n4. Showing existing conservation laws studied for linear and ReLU networks are maximal\n\n", "weaknesses": "none", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive comments and constructive suggestions.\n\n**Question addressed**\n\n> **Q1d** Do the results stated for ReLU networks hold for any network with homogeneous activation function?\n\nYes, our results also apply to networks using any positively $p$-homogeneous activation function. Specifically, for a 2-layer NN given by $g_{\\theta}(x) = U \\sigma(V^\\top x)$, the conserved quantities are $\\|u_i\\|^2 - \\|v_i\\|^2/p$, where $u_i$ and $v_i$ are the columns of $U$ and $V$, respectively. We plan to include this example in the supplementary material of the final version."}, {"review_id": "MrPPOtWAnD", "reviewer": "Reviewer_pz1Q", "summary": "The 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.", "questions": "\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. ", "limitations": "The 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive comments and constructive suggestions.\n\n**Weaknesses and Questions addressed**\n\n> **L1e** It is not clear to what extent the factorisation property the authors introduce is necessary to their analysis. Said differently, it is not clear whether 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\nThank you for the insightful observation. When you fix the dataset and the loss, you're essentially dealing with a singular vector field, specifically: $\\theta \\mapsto \\nabla \\mathcal{E} (\\theta)$. In this particular scenario, our framework is directly applicable. Given that we're examining a singular vector field, its associated Lie algebra has a dimension of one. This results in $D-1$ conserved quantities. However, it's worth noting that these quantities intricately depend on both the chosen dataset and the specified loss. Consequently, the utility of such an analysis, which is inherently data-dependent, may be somewhat limited. We'll be sure to include this point as a clarifying remark in our work.\n\n> **Q1e** It is not clear to me whether the requirement in Eq. 2 is vacuous, meaning that any architecture automatically 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\nThank you for raising this point! If you adopt this simplistic factorization, then $ \\partial \\phi (\\theta) = I_D$, leading to $V_\\phi(\\theta) = \\mathbb{R}^D$. This means that for this particular $\\phi$, there isn't any conservation law. In essence, such a parametrization doesn't have the requisite \"tightness\" to yield an equivalent of Theorem 2.8. This is a very good remark and we will add it to clarify!\n\n> **Q2e** 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 Ω 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\nIn the case of linear / ReLU networks, our analysis shows that there are no conservation laws beyond existing conservation laws, which are known to be global. This addresses the gluing issue for such cases. For more general settings compatible with our analysis, which is indeed only local, gluing conditions are an interesting but possibly difficult challenge, we will comment a bit on this in the revised version.\n\n> **Q3e** Is there a commonly used loss for which Eq. 7 is not satisfied?\n\nAs mentioned in Remark A.7, Eq. 7 is not satisfied for cross-entropy loss. However, it is possible to envision weaker assumptions on the span involved in Eq. 7 to extend the theory to such a loss. This is an interesting challenge left for further work.\n\n> **Q4e** 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\nIn the manuscript, Section 2.4 provides an algorithm (Algo<1>) that constructs directly all polynomial conservation laws. By comparing the number of independent polynomial conservation laws with $D - m$, with $m$ the dimension of the trace of the generated Lie algebra (whose algorithm (Algo<2>) is described in section 3.3), we found that the numbers match and that all polynomial conservation laws found by Algo<1> correspond to the ones already known by the literature. By Algo<2>, we know that there are no other conserved quantities.\n\nThe only “a posteriori” reasoning in our analysis is to say that all conserved quantities are in fact global as discussed in our answer to **Q2e**.\n\nRegarding the *discovery* of new conservation laws: for architectures involving piecewise polynomial activations, we expect that a polynomial $\\phi$ yielding the factorization $f \\circ \\phi$ can again be found, allowing to conduct the same analysis but different (polynomial) conservations laws. The main challenge, left to future work, would be to establish an equivalent of Th 2.8.\n\n> **L1e** The authors briefly discuss limitations in the main text. I 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.\n\nFor deeper architectures, we still have a local factorization $\\phi$ of the form Eq 2, as mentioned in our paper (ll 122-123, appendix A.2, section 4.2). It will be a good challenge to generalize theorem 2.8 to deeper cases, which we will do for further work."}]

Author rebuttal:

Meta-review: A fine work proposing a more systematic approach to conservation laws for gradient flows based on concepts and tools from differential topology, in particular ones related to Frobenius theorem.

The reviewers unanimously agreed that the paper forms a solid contribution to a growing body of work in this topic, appreciating the mathematically compelling perspective it offers on the problem of characterizing continuous conservation laws.
Decision: Accept (oral)