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

Summary

The authors study conservation laws in the gradient flow dynamics of neural networks. They 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$. Under this decomposition, the authors set out to characterise conserved quantities for fixed $\phi$ and for all $f$ giving rise to a proper ERM loss. They claim that under some assumption, and for ReLU networks, this is equivalent of characterising conserved quantities for fixed $\phi$ and for all smooth $f$. More 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$. The 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. They 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. Finally, 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. They conjecture, and briefly discuss this point in Appendix 9, that the same results hold for larger depths.

Strengths

- The article is well written, and guides the reader nicely through quite technical results without requiring too much previous knowledge. - 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). - 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

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

Questions

## Major questions - 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. - 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. - Is there a commonly used loss for which Eq. 7 is not satisfied? - 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? ## Suggestions for manuscript improvement line 50: missing closing bracket line 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. line 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

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.

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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

authors discussed the limitation.

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. The 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).

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: - 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 - 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. - The above discussion is particularly important for hoping to apply these techniques on NNs with activation functions that are not piecewise linear.

Questions

- Can you add more discussion on the situation where the Lie algebra is infinite-dimensional? - Can you add more discussion on what happens when the activation functions of the NN are not piecewise linear?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

There are no potential negative societal impacts of this work.

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.

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

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

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? - 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.”) - 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The authors included limitations by clearly stating the assumptions. There are no potential negative societal impacts of the work.

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.

Strengths

The paper is mathematically engaging, well-written and the content is presented clearly. While 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. I 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. 2. 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. 3. 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. 4. The definition of the main algorithm is obscure. A step-by-step implementation might help for clarity.

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. 2. 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)? 3. The statement at line 168 is ambiguous to me; do you mean $\nabla h \perp \chi, \forall \chi \in V$? 4. 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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The applicability of the theory to practical neural networks is very restricted at the moment.