Source-linked AI summary
On Parameter Symmetries and Conservation Laws in Gradient Flow
Khang Nguyen, Guido Montúfar
TL;DR
The relationship between parameter symmetries and conservation laws under gradient flow remains nuanced. The paper develops a geometric framework and an inheritance principle, showing that symmetries and conservation laws correspond under specified conditions, with limits on applying inheritance to deep linear networks.
Problem
The precise relationship between symmetries and conservation laws under gradient flow, including whether all conserved quantities arise from known symmetries, remains unclear.
Source: §1 · p. 1Method
The paper develops a unified vector-field framework and uses compositional identifiability to establish inheritance and completeness results for multilayer networks.
Source: §1.1 · p. 1, §4 · p. 7Results
Independent conservation laws induce independent loss symmetries, but not vice versa; correspondence with conservation laws requires closedness of the generating vector field.
Source: §1.1 · p. 1, §3 · p. 5Takeaways & Limitations
The framework characterizes conservation laws for multi-head and grouped-query attention, polynomial neural networks, and square deep linear networks.
Source: §1.1 · p. 1Compositional identifiability does not hold for matrix multiplication, so the inheritance result cannot be applied to deep linear networks.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Parameter space symmetries and conservation laws play an important role in understanding the loss landscapes and implicit biases of neural networks. Inspired by Noether's theorem in physics, prior works have sought to derive conservation laws under gradient flow from parameter symmetries, but the scope and limitations of this connection remain unclear. We develop a unified geometric framework that clarifies the precise relationship between the two notions, including the conditions under which symmetries correspond to conservation laws. We introduce a notion of compositional identifiability and use it to establish a general inheritance principle for complete characterizations of symmetries and conservation laws in multilayer networks. We apply the framework to multi-head and grouped-query attention, polynomial neural networks, and square deep linear networks.
1 Introduction
The paper addresses the nuanced relationship between parameter symmetries and conservation laws under gradient flow. It develops a geometric framework, inheritance results, and architecture-specific characterizations to clarify that relationship.
1 Introduction
- Symmetries characterize loss-landscape redundancy, while conservation laws constrain optimization trajectories and characterize resulting optimization bias.Source: §1 · p. 1
1.1 Main contributions
- The framework formalizes the gap between differentiable symmetries and conservation laws: independent conservation laws induce independent loss symmetries, but not vice versa.Source: §1.1 · p. 1, §1 · p. 1
- The framework establishes inheritance and completeness results for compositional models and characterizes conservation laws for attention, polynomial, and square deep linear networks.Source: §1.1 · p. 1
1.2 Related works
- Parameter symmetries have been studied across architectures and linked to optimization-landscape understanding and practical convergence and generalization techniques.Source: §1.2 · p. 2
- Prior frameworks sought to derive neural-network conservation laws from symmetries, but did not completely characterize their relationship for gradient flows.Source: §1.2 · p. 2
2 Preliminaries
The preliminaries specify the gradient-flow setting and define conservation laws and symmetries. They also establish the distinction between differentiable symmetries and independent conservation laws through a two-layer linear-network example.
2 Preliminaries
- Under the stated smoothness assumptions, gradient flow has a unique solution for all time for any finite nonempty dataset and initialization.Source: §2 · pp. 2–3
2.1 Conservation laws
- A conservation law remains constant along every gradient-flow trajectory, and a complete set consists of functionally independent laws with no additional independent conservation laws.The paper also gives a local factorization result for functions dependent on an independent family of functions.Source: §2.1 · p. 3
2.2 Loss symmetries and functional symmetries
- Functional symmetries preserve the represented model function, whereas loss symmetries preserve loss values; every functional symmetry is a loss symmetry, but not conversely in general.Source: §2.2 · p. 3
2.3 A gap between differentiable symmetries and conservation laws
- In a two-layer linear network, four independent one-parameter functional symmetries coexist with a complete set of only three conservation laws, so one symmetry yields no new independent law.Source: §2.3 · p. 4
3 Loss Symmetries, Conservation Laws, and Vector Fields
The section characterizes loss symmetries and conservation laws through vector fields orthogonal to the gradient-flow distribution. It gives precise correspondences and conditions, including when their counts agree.
3 Loss Symmetries, Conservation Laws, and Vector Fields
- A connected Lie-group action is a differentiable symmetry of a dataset loss exactly when it satisfies the stated necessary-and-sufficient condition.Source: §3 · pp. 4–5
- Vector fields in the orthogonal complement of the gradient-flow distribution correspond one-to-one with partial loss symmetries, which are global when their generating fields are complete.Source: §3 · p. 5
- Closed vector fields in that orthogonal complement correspond one-to-one with local conservation laws up to an additive constant; star-shaped parameter spaces make the laws global.Source: §3 · p. 5
- Under constant-rank assumptions on the gradient-flow distribution and its Lie completion, the local symmetry–conservation-law gap is nonnegative and vanishes exactly when the distribution is involutive.Source: §3 · p. 6
- When the distribution is not involutive, initialization-time conservation laws cannot determine the specific local minimum, though they provide an initialization-dependent implicit bias.Source: §3 · p. 6
4 Inheritance Principle for Symmetries and Conservation Laws
The section develops inheritance and completeness results for symmetries and conservation laws in multilayer networks, connecting them to functional equivalence and compositional identifiability. Under the stated conditions, larger-network symmetries and conservation laws are inherited from subnetworks.
4 Inheritance Principle for Symmetries and Conservation Laws
- Separability of predictions lets partial loss symmetries imply partial functional symmetries, supporting the framework’s inheritance and completeness results.The paper describes this as a natural condition satisfied by many loss functions.Source: §4 · pp. 6–7
- A general inheritance result extends subnetwork partial symmetries and conservation laws to a larger network when subnetwork functional equivalence is preserved, covering many compositional structures.Examples include serial composition, sums, concatenations, Hadamard products, outer products, and graph aggregation.Source: §4 · p. 7
- When functional equivalence of the larger network forces equivalence of its subnetworks, its symmetries and conservation laws are infinitesimally generated by natural subnetwork extensions.Those extensions need not themselves be partial symmetries or conservation laws of the larger network.Source: §4 · p. 7
- If the larger network is compositionally identifiable, all its symmetries and conservation laws are inherited from its component subnetworks, yielding a completeness result.The paper combines the preceding inheritance results to establish completeness for multilayer networks.Source: §4 · p. 7
5 Application to Neural Networks
The framework applies to attention, polynomial, and square deep linear networks, establishing complete symmetry and conservation-law characterizations where supported and bounds on independent conservation laws for square deep linear networks.
Grouped query attention
- GQA admits complete local characterizations of functionally equivalent parameters and conservation laws, with standard MHSA covered as the case nH = nG.The attention result assumes full-rank matrices and pairwise distinct Qᵢʲ(Kʲ)⊤ matrices at the point considered.Source: §5 · p. 8
- The standard multi-head self-attention case is recovered by setting nH = nG.The attention results continue to hold in this case.Source: §5 · p. 8
SGD experiment
- During prolonged constant-step-size SGD in a language model, approximate conservation coexists with falling training and validation loss, while preservation for some parameter blocks may deteriorate.This appendix experiment examines conservation laws derived in Proposition 18.Source: §5 · p. 8
Polynomial neural networks
- For PNNs finite-to-one at a generic parameter, the framework gives a complete local set of conservation laws under Assumption 13.The paper identifies this as the first deep network architecture known to have a complete conservation-law characterization.Source: §5 · p. 8
Deep square linear networks
- For a deep square linear network, if W_L...W_1 is full rank, then locally around its parameters, there are at most (L − 1)D^2 independent conservation laws for any loss satisfying Assumption 13.Source: §5 · p. 9
Other applications
- If trainable matrices A and B enter a multilayer network only through AB, the network inherits conservation laws given by the entries of A⊤A − BB⊤.The paper gives low-rank adaptation (LoRA) as an example.Source: §5 · p. 9
6 Conclusion
The paper develops a geometric framework for symmetries and conservation laws in neural network gradient flow, with stated scope limits and directions for further work.
Conclusion
- The framework characterizes the relationship between symmetries and conservation laws, and its inheritance result cannot be applied to deep linear networks because compositional identifiability fails for matrix multiplication.The authors suggest investigating inheritance under weaker compositional identifiability conditions.Source: §6 · p. 9
- The authors identify extending the results to non-gradient-flow and data-dependent dynamics as future directions.They also propose investigating conservation-law-informed initialization and training.Source: §6 · p. 9
AI use statement
The supplied sections state the authors’ AI-use boundaries, introduce vector-field and distribution concepts, and contrast Noether’s theorem with gradient-flow dynamics.
AI use statement
- The authors report that generative AI assisted with claims, hypotheses, methods, experiments, code, literature, and manuscript readability, while they independently checked the mathematical content and verified AI-assisted code.They state that AI did not develop theoretical models or conceptual frameworks, supply critical proof ingredients, or write proofs.Source: AI use statement · p. 9
A Background on Vector Fields
- A vector field assigns each point in an open parameter domain a direction, while a covector field assigns a linear functional represented as a row vector.A distribution collects subspaces spanned pointwise by smooth vector fields; its annihilator consists of covectors that vanish on those subspaces.Source: A Background on Vector Fields · p. 14
B.1 A typical derviation of Noether’s theorem
- Noether’s theorem links Lagrangian symmetries to conserved charges, but it does not directly apply to gradient flow because its dynamics are specified directly rather than derived from a Lagrangian.The passage describes the Noether charge as a conserved quantity along an optimal path.Source: §B.1 · p. 15
B.2 Comparison with other frameworks connecting symmetries and conservation laws under gradient flow
The comparison explains how the framework relates symmetry generators to conserved potentials, situates it against prior approaches, and covers one-parameter symmetries and global existence of gradient flow.
B.2 Comparison with other frameworks connecting symmetries and conservation laws under gradient flow
- The cited prior approaches impose additional conditions or lack a general criterion for when a prescribed generator admits a conserved potential.Kunin et al. (2021) and Zhao et al. (2023) require an additional condition for their construction; Tanaka and Kunin (2021) recover an equivalent identity in a massless limit but not a general criterion.Source: §B.2 · p. 16
- For the scalar network G(w1, w2) = w1e^2w2x, the framework finds a local conservation law from a closed generator that the cited Kunin et al. (2021) and Zhao et al. (2023) conditions do not cover.The generator is not linear in the parameters, and the additional Kunin et al. (2021) condition is not satisfied.Source: §B.2 · p. 16
- The framework recovers sufficiently smooth conservation laws derived by Kunin et al. (2021), Głuch and Urbanke (2021), and Zhao et al. (2023).It makes explicit the local integrability criterion underlying the potential-based construction.Source: §B.2 · p. 17
B.3 Recent works connecting conservation laws and symmetries via vector fields
- A conservation law yields a partial symmetry whose generator is its gradient, and a generator yields a local conserved potential exactly when it is closed.The framework presents this as a systematic characterization of the local relationship between loss symmetries and conservation laws.Source: §B.3 · p. 17, §B.2 · p. 17
C Justification for 1D Symmetries
- For a connected Lie group, its one-parameter subgroups can reconstruct the multidimensional symmetry, justifying the focus on one-dimensional symmetries.Products of exponentials from a basis of its Lie algebra generate the group.Source: C Justification for 1D Symmetries · pp. 17–18
D.1 Proof for Proposition 1
- With a lower-bounded loss and continuously differentiable per-example losses having locally Lipschitz gradients on R^D, gradient flow has a unique solution for all t ≥ 0.The result applies to every finite nonempty dataset and every initialization.Source: §D.1 · p. 18
D.2 Proof for Proposition 4
These proofs connect factorization through a map to local symmetry and conservation-law characterizations. They also extend the vector-field relationship to partial symmetries and relate conservation laws to their infinitesimal generators.
D.2 Proof for Proposition 4
- Locally, h factors through H exactly when its variation is confined to the independent directions represented by the gradients of H's components.The proof flattens H by local diffeomorphisms and shows h is independent of the remaining coordinates.Source: §D.2 · pp. 18–19
E.1 Proof for Proposition 5
- A connected Lie-group action is a differentiable loss symmetry exactly when the loss is invariant under that action.The proof uses the vanishing derivative along the group action and connectedness to obtain invariance throughout the group.Source: §E.1 · pp. 19–20
E.2 Proof for Corollary 6
- For the additive real-group action, Corollaries 6 and 7 specialize the symmetry criterion to a dataset loss and to the overall loss, respectively.Corollary 7 uses Wθ, the span of gradient-flow directions for datasets of size one.Source: §E.2 · p. 21, §E.3 · p. 21
E.4 Proof for Proposition 8
- A conservation law has a gradient that is a closed vector field orthogonal to W; conversely, such a field has a local potential, and a global one on star-shaped Θ.The potential is unique up to an additive constant.Source: §E.4 · p. 21
E.5 Proof for Proposition 9
- A loss symmetry generates a smooth vector field in W⊥; conversely, a smooth field there generates a unique maximal local flow, which is a loss symmetry if the field is complete.The flow obeys the group-action properties wherever defined.Source: §E.5 · p. 21
E.6 Extension for Corollary 6, Corollary 7 and Propostion 9 to partial symmetries
- The vector-field criterion and flow construction extend to partial symmetries, which need only be defined locally; completeness still yields a loss symmetry.The flow preserves the loss wherever defined.Source: §E.6 · p. 22
E.7 Proof for Corollary 10
- A conservation law determines a unique associated partial loss symmetry; conversely, a partial symmetry with a closed infinitesimal generator determines a local conservation law up to a constant.The correspondence uses the conservation-law gradient as the symmetry generator.Source: §E.7 · p. 23
E.8 Proof for Theorem 12
The theorem counts local symmetries and conservation laws under constant-rank assumptions, then the appendix results connect loss symmetries to functional symmetries and establish inheritance from subnetworks.
E.8 Proof for Theorem 12
- Under the constant-rank assumptions, local symmetry and conservation-law counts differ by the rank gap between W and Lie(W), and equality holds exactly when W is involutive.The theorem gives D−r partial symmetries and D−r conservation laws, respectively, using the ranks of W and Lie(W).Source: §E.8 · p. 23
F.1 Loss functions that satisfy Assumption 13
- Lemma 26 lists losses satisfying Assumption 13: for Z=Y, losses zero on equal inputs and positive otherwise; metrics such as Euclidean or Lp distances; quadratic and absolute losses; and Lp losses for p∈(0,∞).Source: §F.1 · p. 24
F.2 Proof for Proposition 14
- Under Assumption 13, every partial loss symmetry is also a partial functional symmetry of the model.The proof transfers invariance of the loss to invariance of the model output.Source: §F.2 · p. 25
F.3 Proof for Proposition 15
- A subnetwork's partial symmetry or conservation law extends to the larger network under the stated identifiability condition; the extended conservation law holds for any loss.The extension leaves the other parameters unchanged and extends the law as a function of the subnetwork parameters.Source: §F.3 · pp. 25–26
F.4 Proof for Proposition 16
The subnetwork results characterize larger-network symmetries and conservation laws through extensions of the corresponding subnetwork objects, with a caveat about whether individual extensions are themselves valid objects.
F.4 Proof for Proposition 16
- Under the stated conditions, every larger-network partial symmetry or conservation law is infinitesimally generated by extensions of complete subnetwork sets.The individual extensions need not themselves be partial symmetries or conservation laws of the larger network.Source: §F.4 · pp. 26–28
F.5 Proof for Theorem 17
- For multiple subnetworks, the extensions of their partial functional symmetries and conservation laws form complete sets for the larger network.The theorem assumes the larger network's loss also satisfies Assumption 13.Source: §F.5 · p. 28
- The extended objects act on disjoint parameter spaces across subnetworks, supporting their independence in the larger network.The proof establishes this independence to complete the characterization.Source: §F.5 · p. 28
G.1 Proof for Proposition 18
The proofs characterize symmetries and conservation laws for grouped-query attention and polynomial networks, and bound the number of laws for square deep linear networks. In a scalar scaling example, the conservation law determines the convergent minimizer in the positive quadrant.
G.1 Proof for Proposition 18
- Grouped-query attention has a complete local characterization of its functional symmetries and conservation laws under the proposition’s assumptions.The proof combines attention identifiability with complete conservation-law characterizations for its subnetworks.Source: §G.1 · pp. 29, 31
G.2 Proof for Proposition 19
- For a polynomial network with monomial activations, finite-to-one architecture at a generic θ0 yields a complete set of conservation laws locally.The result holds for losses satisfying Assumption 13.Source: §G.2 · p. 31
G.3 Proof for Proposition 20
- When the product of weights is full rank, a deep square linear network has at most (L − 1)D^2 independent conservation laws locally.The bound holds for losses satisfying Assumption 13.Source: §G.3 · p. 32
H.1 The scalar case
- In the scalar two-layer scaling problem, the symmetry-law gap is zero, and the conserved value determines the unique positive-quadrant minimizer if gradient flow converges.The uniqueness follows from solving the conservation-law constraint together with the scaling-function condition.Source: §H.1 · pp. 32–33, H Learning to Scale · p. 32
H.2 The 2x2 matrix case
In the 2 × 2 two-layer linear case, conservation laws do not identify a unique minimizer: solutions form an O(2) family. They still determine the minimizer’s singular values and condition numbers.
H.2 The 2x2 matrix case
- The complete set of independent conservation laws for the 2 × 2 case is given by a matrix function.The conserved matrix value is denoted H0 in the solution characterization.Source: §H.2 · p. 33
- The symmetry-law gap is one, so conserved values leave a one-dimensional manifold of local minimizers rather than identifying a unique point.The solution set is induced by the orthogonal group O(2).Source: §H.2 · p. 33
- Despite the remaining O(2) freedom, the conservation laws determine the minimizer’s singular values and condition numbers.These quantities are invariant under the orthogonal action on the solution family.Source: §H.2 · pp. 33–34
- Solving for the positive definite Gram matrix reduces to a quadratic equation for each diagonal entry, with a unique positive diagonal solution.The resulting solution set is recovered by reversing the diagonalization and applying the orthogonal action.Source: §H.2 · p. 34
I Empirical Examination of Conservation Laws in a Language Model
The language-model experiment examines approximate conservation during constant-step-size SGD while training losses decrease. Conservation-law drift is generally smaller than control drift, but preservation varies across blocks and can deteriorate over time.
I Empirical Examination of Conservation Laws in a Language Model
- The modified six-layer nanoGPT model was trained from scratch on Tiny Shakespeare for 20,000 updates using plain SGD at a constant 10^-2 learning rate.Layers 1 and 4 use grouped-query attention; the other four use standard multi-head self-attention.Source: I Empirical Examination of… · p. 34
- Training and validation losses decrease, and fixed-prompt samples after training contain recognizable words and dialogue-like formatting.The sample comparison is before training versus after 20,000 updates.Source: I Empirical Examination of… · p. 34, Figure 2 · p. 35, Figure 3 · p. 35
- Conservation-law matrices often drift less than non-conserved Gram-matrix sums, especially in Q, K blocks; in V, O blocks of layers 1 and 2, this holds only through 5,000 updates.Figures 4 and 5 compare normalized conservation-law drift (solid curves) with normalized control drift (dashed curves).Source: I Empirical Examination of… · p. 35, Figure 4 · p. 36, Figure 5 · p. 36
- The experiment illustrates approximate conservation during SGD, with preservation varying across layers and blocks and potentially deteriorating over time.For constant step size, the cited theoretical basis states that conservation-law error grows linearly with the number of steps.Source: I Empirical Examination of… · pp. 34–35