Source-linked AI summary
Abide by the Law and Follow the Flow: Conservation Laws for Gradient Flows
Sibylle Marcotte, Rémi Gribonval, Gabriel Peyré
TL;DR
The paper asks which properties of initialization persist during training of over-parameterized models, a question tied to understanding their observed generalization. It formalizes conservation laws for model architectures and uses Lie-algebra dimensions to characterize their maximal number, reporting that known laws are maximal in its studied settings.
Problem
The question is which initialization properties over-parameterized models preserve during training, a route to understanding their observed generalization despite near interpolation.
Source: §1 · p. 1Method
The paper defines conservation laws as quantities preserved through all gradient flows for a model architecture and uses Lie-algebra dimensions to count independent laws.
Source: Contributions · p. 2, §3.2 · p. 8Results
Known conservation laws are maximal in the studied cases, and numerical checks on sampled small depths and widths found no additional laws for deeper linear and ReLU networks.
Source: Contributions · p. 3, §4.2 · p. 10Takeaways & Limitations
The framework can also study displacement spaces beyond gradient flows, including momentum after lifting the flow to a higher-dimensional phase space.
Source: Conclusion · p. 10The theory covers continuous-time gradient flow; finite-step gradient descent disrupts exact conservation, while approximate conservation remains future work.
Source: Conclusion · p. 11
Abstract
from arXiv · showhide
Understanding the geometric properties of gradient descent dynamics is a key ingredient in deciphering the recent success of very large machine learning models. A striking observation is that trained over-parameterized models retain some properties of the optimization initialization. This "implicit bias" is believed to be responsible for some favorable properties of the trained models and could explain their good generalization properties. The purpose of this article is threefold. First, we rigorously expose the definition and basic properties of "conservation laws", that define quantities conserved during gradient flows of a given model (e.g. of a ReLU network with a given architecture) with any training data and any loss. Then we explain how to find the maximal number of independent conservation laws by performing finite-dimensional algebraic manipulations on the Lie algebra generated by the Jacobian of the model. Finally, we provide algorithms to: a) compute a family of polynomial laws; b) compute the maximal number of (not necessarily polynomial) independent conservation laws. We provide showcase examples that we fully work out theoretically. Besides, applying the two algorithms confirms for a number of ReLU network architectures that all known laws are recovered by the algorithm, and that there are no other independent laws. Such computational tools pave the way to understanding desirable properties of optimization initialization in large machine learning models.
1 Introduction
The paper makes the study of initialization properties preserved during training precise by defining conservation laws and linking them to algebraic computations. It proposes algorithms to identify such laws and their maximum number, and shows that in certain settings they support a low-dimensional description of over-parameterized flows.
1 Introduction
- Over-parameterized networks can generalize well despite interpolation, and the optimization algorithm can induce an implicit bias in the trained model.The paper frames understanding which initialization properties persist during training as a way to analyze this behavior.Source: §1 · p. 1
Contributions
- The paper formalizes conservation laws as quantities preserved through gradient flows for a given architecture, and links their maximal number to Lie-algebra computations.It treats these laws as a weak form of implicit bias that can reveal properties of initialization retained at convergence.Source: Contributions · p. 2, §1 · p. 2
- For deep linear and shallow ReLU networks, characterizing conservation laws for several classical losses reduces to analyzing a finite-dimensional space of vector fields.The paper states that the proposed framework covers linear and ReLU networks, shallow or deep, with quadratic and other losses.Source: Contributions · p. 2
- The SageMath algorithms identify polynomial conservation laws and compute the maximum number of independent laws; on selected examples, they identify known laws and show those laws are maximal.The paper reports this completeness result for matrix factorization and certain ReLU architectures.Source: Contributions · pp. 2–3
- In certain settings, the results allow an over-parameterized flow to be rewritten as an intrinsic low-dimensional flow.The paper says this occurs only in certain settings.Source: Contributions · p. 2
2 Conservation Laws for Gradient Flows
The section formalizes quantities preserved during gradient flows and develops a reparametrization-based framework for characterizing them. It uses finite-dimensional vector fields to construct polynomial laws and define the maximum number of independent laws.
2.1 Gradient dynamics
- Gradient-flow conservation is defined for functions preserved along the optimization dynamics, with the paper distinguishing a fixed dataset from conservation for any dataset.Stochastic optimization and discrete gradients generally preserve these quantities only approximately.Source: §2.1 · p. 3, §2.2 · pp. 3–4
2.2 Conserved functions
- A function is conserved through a vector-field space exactly when its gradient is orthogonal to that space at every parameter value.For conservation under any dataset, the relevant vector fields are characterized through the model and loss; a scalar linear network example has the law h(u, v) = u^2 − v^2.Source: §2.2 · p. 4
2.3 Reparametrization
- Factoring a model as f(φ(θ), x) makes the parameterization φ central to analyzing conservation laws; linear-network factorizations can be global, while ReLU factorizations are local.For ReLU models, the local factorization uses regions where activation patterns remain constant.Source: §2.3 · pp. 4–5
2.4 From conserved functions to conservation laws
- Under the stated loss and model assumptions, conservation for any dataset is equivalent to conservation through the finite-dimensional space Wφ for linear and two-layer ReLU networks.For the ReLU result, hidden neurons must have pairwise distinct hyperplanes; whether the corresponding conclusion holds for deep ReLU networks remains open.Source: §2.4 · pp. 5–6
2.5 Constructibility of some conservation laws
- Polynomial constraints let the authors compute independent polynomial conservation laws of a chosen degree, recovering all known laws in selected linear- and ReLU-network examples.The equations are implemented in SageMath.Source: §2.5 · p. 6
2.6 Independent conserved functions
- Conserved functions are independent when their gradients are pointwise linearly independent, and their number cannot exceed D − dim W(θ).The section relates tightness of this bound to additional conditions on W involving Lie brackets.Source: §2.6 · pp. 6–7
3 Conservation Laws using Lie Algebra
The section characterizes the maximum number of conservation laws through the Lie algebra generated by the model’s Jacobian vector fields. It gives an iterative computation and, under specific conditions, recasts the high-dimensional parameter flow as a lower-dimensional Riemannian flow.
3 Conservation Laws using Lie Algebra
- Frobenius theory connects trapping surfaces to Lie brackets, but the relevant vector-field space may not satisfy its condition, so the generated Lie algebra must be considered.The section allows the model’s reparameterization φ to be infinitely smooth while the loss and model may be less smooth.Source: §3 · p. 7
3.2 Number of conservation laws
- When dim(Lie(Wφ)(θ)) is locally constant, the number of independent local conservation laws is exactly D − dim(Lie(Wφ)(θ)).The local-constancy assumption is technical and typically holds outside pathological points; it holds generically in the matrix-factorization example.Source: §3.2 · p. 8
3.3 Method and algorithm, with examples
- The algorithm builds bases for Wk iteratively and stops when dimension stagnates; for r rank-one n×m matrices, the r independent known conserved functions hi(U,V)=||ui||²−||vi||² are complete.Source: §3.3 · p. 8
3.4 Application: recasting over-parameterized flows as low-dimensional Riemannian flows
- If rank(∂φ(θ)) is constant and Wφ satisfies condition (13), the parameter flow can be rewritten as a lower-dimensional flow with a Riemannian metric induced by φ and dependent on initialization.The metric can have a kernel when φ maps onto a lower-dimensional manifold; the kernel is orthogonal to that manifold’s tangent space.Source: §3.4 · p. 9
4 Conservation Laws for Linear and ReLU Neural Networks
The section tests whether known conservation laws are complete for linear and ReLU networks, proving completeness for two-layer linear networks and reporting numerical checks for deeper linear and ReLU cases.
4 Conservation Laws for Linear and ReLU Neural Networks
- For q-layer models, the section uses the product reparameterization for linear networks and a polynomial reparameterization for ReLU networks, including networks with biases.The simplified ReLU expression is specified for networks without biases.Source: §4 · p. 9
- Previously known laws compare adjacent layers’ weight matrices; their stated counts are not generally independent.For q > 2, whether other, previously unknown laws exist remains open.Source: §4 · pp. 9–10
4.1 The matrix factorization case (q = 2)
- For a concatenated parameter matrix (U; V) of rank rk, the function Ψ(U,V) = UᵀU − VᵀV gives rk · (2r + 1 − rk)/2 independent conservation laws.When (U; V) has full rank, all conservation laws are given by Ψ, with no more independent conservation laws.Source: §4.1 · p. 10
- The q = 2 proof computes the dimension of the Lie-algebra trace, matching the count of known laws and establishing their completeness.The trace dimension is (n + m)r − rk(2r + 1 − rk)/2 for full-rank concatenated parameters.Source: §4.1 · p. 10
4.2 Numerical guarantees in the general case
- For deeper linear and ReLU networks, numerical tests found no more independent laws than the known set, while the authors leave completeness for q > 2 open.The tests used small samples of depths and widths; the authors also report finding maximal sets of polynomial laws algorithmically.Source: §4.2 · p. 10, §4 · p. 10
Conclusion
The conclusion presents the work as a constructive program for counting conservation laws and points to broader architectures and non-gradient flows as future applications. Its current theory is limited to continuous-time gradient flow.
Conclusion
- The authors propose extending the analysis to deep convolutional networks, normalization, and attention layers.They also note that the theory can apply to other displacement spaces, including momentum-based flows lifted to a higher-dimensional phase space.Source: Conclusion · p. 10
- The theory is restricted to continuous-time gradient flow; finite-step gradient descent disrupts exact conservation, leaving approximate conservation for future study.Source: Conclusion · p. 11
A Proof of Proposition 2.5
The appendix establishes a derivative test for whether a function is conserved by a flow, then relates model-gradient vector fields to data points and targets under smoothness conditions.
A Proof of Proposition 2.5
- A C1 function is conserved through a flow exactly when its derivative along the vector field vanishes everywhere.If the derivative is nonzero at a point, continuity gives a short interval on which the function changes along the flow from that point.Source: A Proof of Proposition 2.5 · p. 13
B Proof of Proposition 2.7
- The proof of Proposition 2.7 uses C2 smoothness of the model outputs near the parameter to characterize the gradient fields generated by datasets.It shows inclusion in both directions, with a single feature-target pair sufficing to represent an individual sample gradient field.Source: B Proof of Proposition 2.7 · p. 13
C Proof of Lemma 2.13 and Theorem 2.14
The appendix proves the gradient-field reductions used for linear and two-layer ReLU networks under smoothness and loss conditions. It verifies those conditions for mean-squared error and logistic loss, but notes that cross-entropy does not satisfy one of them.
C Proof of Lemma 2.13 and Theorem 2.14
- For the reparameterized model, the span of gradient fields is generated by Jacobian-transpose vectors over smooth inputs and output-space vectors.This characterization is stated under the local reparameterization and loss conditions.Source: C Proof of Lemma 2.13 and Theorem 2.14 · pp. 13–14
- Mean-squared error and logistic loss satisfy condition (9), whereas cross-entropy does not.The appendix identifies extending the relevant lemma under weaker assumptions to cover cross-entropy as an open challenge.Source: C Proof of Lemma 2.13 and Theorem 2.14 · p. 14
- For linear networks of any depth, the proof reduces the model’s gradient fields to those of its product reparameterization when the loss satisfies condition (9).The argument uses the full span of matrix fields wx⊤ over inputs x and vectors w.Source: C Proof of Lemma 2.13 and Theorem 2.14 · p. 15
- For two-layer ReLU networks, the theorem assumes condition (9), C2 loss differentiability, and pairwise distinct neuron hyperplanes.The proof treats both the no-bias and bias cases.Source: C Proof of Lemma 2.13 and Theorem 2.14 · pp. 15–16
- The appendix describes a possible extension to deeper ReLU networks using a rescaling-invariant polynomial reparameterization and activation-status-dependent matrices.Source: C Proof of Lemma 2.13 and Theorem 2.14 · p. 17
D Proof of Lemma 3.2
If the spaces W_i stop growing in dimension at a point throughout a neighborhood, then all later spaces coincide there. Consequently, the trace of Lie(W) has locally constant dimension.
D Proof of Lemma 3.2
- When dim(W_i+1(θ′)) = dim(W_i(θ)) throughout a neighborhood, every later W_k(θ′) equals W_i(θ′) there, so Lie(W)(θ′) equals W_i(θ′).The proof establishes the equality first at i+1 using a locally persistent basis, then by induction using Lie-bracket closure.Source: D Proof of Lemma 3.2 · p. 18
- The trace of Lie(W) therefore has locally constant dimension, equal to dim(W_i(θ)).The induction bounds each bracket [W_0, W_k] within W_i+1, which equals W_i on the neighborhood.Source: D Proof of Lemma 3.2 · p. 18
E Proof of Theorem 3.3
This proof connects Frobenius integrability to the dimension of the Lie algebra generated by Wφ. When that dimension is locally constant, it gives the exact local number of independent conserved functions.
E Proof of Theorem 3.3
- Frobenius supplies D − dim(W(θ)) local independent conserved functions when W is a Lie algebra with locally constant dimension.The proof uses the equivalence between bracket closure and W′ being a Lie algebra to apply Frobenius.Source: E Proof of Theorem 3.3 · pp. 18–19
- For a locally constant-dimensional distribution W, Frobenius bracket closure is equivalent to the space W′ of vector fields taking values in W being a Lie algebra.W′ is defined pointwise by membership in W, although it is a priori infinite-dimensional.Source: E Proof of Theorem 3.3 · pp. 18–19
- If dim(Lie(Wφ)(θ)) is locally constant, exactly D − dim(Lie(Wφ)(θ)) independent conserved functions exist locally through Wφ.Frobenius provides existence; a contradiction argument rules out any larger independent family.Source: E Proof of Theorem 3.3 · p. 19
F.1 Proof of the result given in Example 3.5
The examples establish Frobenius conditions and count independent conserved quantities for factorized network parametrizations. They cover matrix-factorized outputs and a further positively homogeneous activation setting.
F.1 Proof of the result given in Example 3.5
- For the factorization into r separate bilinear blocks, Wφ has dimension (n + m − 1)r and satisfies Frobenius’s condition.The proof reduces both claims to one block, since the Jacobian and Hessian are block-diagonal and Wφ is a direct sum.Source: §F.1 · pp. 19–20
F.2 An additional example beyond ReLU
- For a positively one-homogeneous activation, including ReLU, leaky ReLU, and linear activations, the stated nonzero-input setting gives r and no more local conserved functions.Here dim(Wφ(θ)) = r(d − 1), and the functions h_i(u_i,v_i) = u_i^2 − ∥v_i∥^2 are conserved through all Wφ.Source: §F.2 · pp. 20–21
G Proof of Proposition 3.7 and additional example
Under constant-rank and Frobenius conditions, the paper characterizes the evolution of the model output using its value and initialization. It also counts independent conserved quantities for a matrix-valued invariant under full-rank assumptions.
G Proof of Proposition 3.7 and additional example
- With constant rank and Frobenius’s condition, z(t) = φ(θ(t)) follows an ODE whose matrix M depends only on z(t) and θ_init.The proof uses local conserved functions and local inversion to express θ(t), and hence M(t), in terms of the output and initialization.Source: G Proof of Proposition 3.7 and… · p. 22
H.1 Proof of Proposition 4.2
- Under full rank, Ψ(U,V) = U⊤U − V⊤V yields (n + m)(r − 1/2(n + m − 1)) independent conserved functions when n + m ≤ r, and r(r + 1)/2 when n + m > r.The two counts apply in the respective dimension regimes stated for Proposition H.1.Source: §H.1 · p. 22
H.2 Proofs of other results
For full-rank (U; V), the generated Lie-algebra dimension yields the complete set of conserved functions. In Example 3.6, the stated condition fails when max(n, m) > 1 and r > 1.
H.2 Proofs of other results
- The generated Lie algebra is finite-dimensional, and its matrix Lie brackets provide the route to computing its trace dimension.The proof characterizes Lie(Wφ) using skew-symmetric matrices.Source: §H.2 · p. 23
- When n + m ≤ r, dim(Lie(Wφ)(U; V)) = (n + m)(n + m − 1)/2; when n + m > r, it equals (n + m)r − r(r + 1)/2.These dimensions are established for full-rank (U; V).Source: §H.2 · p. 24
- For full-rank (U; V), all conserved functions are given by Ψ(U, V) = U^⊤U − V^⊤V, with no others.The Lie-algebra characterization, together with Proposition 4.2, establishes that Proposition 4.1 exhausts the independent laws.Source: §H.2 · p. 25
I About Example 3.6
- For full-rank (U; V), Wφ does not satisfy condition (13) when max(n, m) > 1 and r > 1.The proof compares dim(Wφ(U; V)) with dim(Lie(Wφ)(U; V)) and shows the former is smaller in the stated setting.Source: I About Example 3.6 · pp. 25–26
- For m < r, the dimension difference is increasing in r and is positive for r in {m + 1, …, n + m − 1}.The proof gives f(r) = (n + m)r − r(r + 1)/2 − nm and shows f(r) > f(m) ≥ 0.Source: I About Example 3.6 · p. 26
J Details about experiments
Across tested linear and ReLU architectures, the computed Lie-algebra dimensions matched the counts implied by known conservation laws. A separate degree-2 search recovered all known laws on the tested architectures.
J Details about experiments
- The computed Lie-algebra dimension matched D − N in every tested architecture, where N counts independent conserved functions known from the literature.Tests covered linear and ReLU architectures, with and without biases, across varied depths and widths.Source: J Details about experiments · p. 26
- The dimension comparison used 50 random architectures per network type, with depths from 2 to 5 and widths drawn from 2 to 10 for ReLU or 2 to 6 for linear networks.ReLU architectures included biases with probability 1/2.Source: J Details about experiments · p. 26
- A degree-2 polynomial search recovered all known conservation laws in 10 randomly selected ReLU and linear architectures.These architectures had depths from 2 to 4 and layer widths drawn from 2 to 5.Source: J Details about experiments · p. 27