📊 statistics

Conservation Laws from Data Symmetry in Neural Networks

This paper investigates whether intrinsic data symmetries lead to conserved quantities during neural network training, proving that such symmetries generally do not induce additional integrals of motion for analytic non-polynomial losses, while demonstrating that mean squared error loss combined with data augmentation can yield extra conserved quantities in tensorizable networks.

Original authors: Jakob Galley, Vahid Shahverdi, Axel Flinth

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jakob Galley, Vahid Shahverdi, Axel Flinth

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to recognize patterns, like identifying different types of fruit. You show it pictures of apples, oranges, and bananas. Now, imagine you have a special rule: no matter how you rotate an apple picture, it's still an apple. This rule is a symmetry.

This paper asks a fascinating question: If we teach the robot using these symmetrical rules, does the robot's learning process get "stuck" in a specific way?

In the world of physics and math, when something stays the same while things change around it, we call it a conserved quantity (like energy in a closed system). In machine learning, finding these "conserved quantities" is like finding a hidden map that tells us exactly how the robot's brain will change as it learns.

Here is the breakdown of what the authors found, using simple analogies:

1. The Two Types of "Loss" (The Scorecard)

To learn, the robot tries to minimize its "mistakes." This mistake score is called the Loss. The paper looks at two very different ways of calculating this score:

The "Strict" Scorecard (Analytic Non-Polynomial Loss): Think of this like a very complex, smooth curve where every tiny change in the input creates a unique, slightly different penalty.
- The Finding: The authors proved that for this type of scorecard, symmetries in the data do NOT create new hidden maps. Even if you rotate your apple pictures, the robot's learning path remains just as flexible and unpredictable as if you hadn't rotated them. The symmetry doesn't "lock" the robot into a specific behavior.
- The Analogy: Imagine walking through a dense, foggy forest (the learning path). If you have a rule that "trees look the same from any angle," it doesn't actually help you find a hidden, straight path through the fog. You still have to wander randomly.
The "Simple" Scorecard (Mean Squared Error - MSE): This is a simpler, more rigid way of calculating mistakes (like measuring the straight-line distance between the guess and the truth). It's like a polynomial equation (think $x^2$ or $x^3$ ).
- The Finding: Here, the story changes! If the data has symmetries (like rotating images), new hidden maps DO appear. The robot's learning path gets constrained. It can't just go anywhere; it is forced to stay on a specific track.
- The Analogy: Now imagine the forest has a hidden river (the conserved quantity). If you know the trees are symmetrical, you realize the river must flow in a straight line. The symmetry of the trees forces the water to follow a specific, unchangeable path. The robot is now "surfing" this river; it can't leave the water, no matter how hard it tries.

2. The "Tensorizable" Networks (The Special Architecture)

The authors found that these new "rivers" (conserved quantities) only appear in a specific type of robot brain architecture they call Tensorizable Networks.

What is it? Imagine a factory assembly line.
- Step 1: The raw material (the input data, like an image) is pre-processed into a standard "kit" of parts.
- Step 2: The robot's brain (the parameters) simply assembles these pre-made kits.
Why it matters: Because the data is pre-packaged into these kits, the symmetry of the data (like rotation) gets "lifted" into the assembly process. If you rotate the input, it's equivalent to rotating the assembly instructions. This creates a new, continuous symmetry in how the robot can adjust its brain, leading to those "conserved quantities."

3. The "Lightning Attention" Example

The paper uses a specific, modern type of neural network called Lightning Attention as a real-world example.

They showed that if you train this network with rotated images (data symmetry), the network's internal "knobs" (parameters) are forced to move in a way that keeps a specific geometric shape (the "column space") constant.
The Result: Even though the network is learning, a specific part of its structure remains frozen in place, like a spinning top that wobbles but never falls over.

Summary

The Big Question: Does teaching a robot with symmetrical data (like rotated images) force it to learn in a predictable, constrained way?
The Answer: It depends on the "scorecard" (Loss function).
- If the scorecard is complex and smooth: No. The symmetry doesn't create new rules.
- If the scorecard is simple (MSE) and the robot has a specific "assembly line" structure: Yes. The symmetry creates a hidden rule (a conserved quantity) that the robot must follow, restricting its learning path.

The authors didn't just guess this; they used advanced math (like Noether's theorem from physics, which links symmetry to conservation) to prove it, and they ran computer experiments to show that these "hidden rivers" actually exist in practice.

Technical Summary: Conservation Laws from Data Symmetry in Neural Networks

Problem Statement

The paper investigates whether intrinsic symmetries present in training data induce new conserved quantities (integrals of motion) during the gradient-flow training of neural networks. While classical Noether's theorem links continuous symmetries of a dynamical system to conservation laws, prior work in deep learning has largely focused on symmetries within the parameter space (e.g., weight permutations or rescalings) that leave the realized function invariant. This work addresses a distinct question: Can symmetries in the data (e.g., invariance under permutations, rotations, or group actions) generate new conservation laws that constrain the optimization path, even when the network architecture itself does not explicitly encode these symmetries?

Methodology

The authors employ a theoretical framework grounded in dynamical systems and representation theory to analyze gradient flow dynamics defined by $\dot{\theta} = -\nabla E_\pi(\theta)$ .

Definitions and Setup:
- Integrals of Motion: A smooth function $I(\theta)$ is an integral of motion if it remains constant along gradient flow trajectories for any data distribution $\pi$ (or specifically symmetric distributions).
- Data Symmetry: The training data distribution $\pi$ is symmetric under a compact group $G$ if applying group actions to inputs and outputs leaves the distribution invariant.
- Group-Augmented Risk: The authors analyze the "group-augmented" loss $E^G_\pi$ , obtained by averaging the loss over the group orbit. By Proposition 1, a function is an integral of motion for symmetric data if and only if it is preserved by the augmented gradient flow.
Analytical Approach:
- The core condition for a new integral of motion to exist is that the span of the augmented gradients must be strictly smaller than the span of the standard gradients (Equation 9). If the augmented gradients span the same subspace as the standard gradients, no new conservation laws emerge.
- The authors distinguish between two classes of loss functions:
  - Analytic, Non-Polynomial Margin Losses: (e.g., logistic, exponential).
  - Polynomial Losses: Specifically Mean Squared Error (MSE).
- Tensorizable Networks: For the polynomial case, the authors introduce a class of architectures called "tensorizable networks." These are networks where the dependence on parameters $\theta$ and inputs $x$ can be separated via a lifted feature space $H$ such that $f_\theta(x) = M(\theta)T(x)$ . This class includes linear networks, polynomial networks, and specific attention mechanisms like Lightning Attention.

Key Contributions and Results

1. Absence of New Conservation Laws for Analytic Non-Polynomial Losses

Theorem 3 establishes that for finite groups and analytic, non-polynomial margin losses, data symmetries do not generically induce new integrals of motion.

Mechanism: The proof relies on the Taylor expansion of the analytic loss function. Since the loss is non-polynomial, its derivative contains infinitely many non-zero terms. Under mild regularity assumptions (Assumption 1, which excludes degenerate equivariant initializations), these infinite terms act as a full-rank linear system. This allows the individual gradient directions of the group orbit to be linearly reconstructed from the augmented (averaged) gradient.
Implication: The span of the augmented gradients equals the span of the standard gradients. Consequently, any quantity conserved under symmetric data training is already conserved under arbitrary data training; no new constraints are introduced by the data symmetry.

2. Emergence of Conservation Laws for MSE Loss in Tensorizable Networks

Theorem 4 demonstrates that for tensorizable networks trained with MSE loss, data symmetries can create new integrals of motion.

Mechanism: Unlike analytic non-polynomial losses, MSE is determined only by its quadratic and linear terms. When augmented over a group $G$ , these terms impose equivariance constraints on the lifted feature space $H$ . This creates a new symmetry group $H$ acting on the lifted space that preserves the augmented loss.
Continuous Symmetries: Crucially, even if the data symmetry group $G$ is discrete, the induced symmetry group $H$ on the lifted feature space can be continuous (e.g., orthogonal, unitary, or symplectic groups).
Result: If the network parametrization realizes this continuous symmetry (specifically, if rotating a parameter block $P$ is equivalent to acting on the lifted space by an element of $H$ ), then the column space (range) of $P$ becomes an integral of motion. The range of $P$ is preserved along the gradient flow.

3. Characterization of the Symmetry Group $H$

The paper provides a detailed algebraic characterization (Appendix D) of the group $H$ that preserves the augmented MSE loss.

$H$ is the intersection of orthogonal matrices that commute with all $G$ -equivariant symmetric operators and $G$ -equivariant linear maps.
The structure of $H$ depends on the decomposition of the lifted feature space into irreducible representations of $G$ . Depending on the type of these irreps (real, complex, or quaternionic), $H$ can contain continuous factors such as $O(n)$ , $U(n)$ , or $Sp(n)$.

Illustrative Examples

The authors validate their theoretical findings with two examples:

Linear Model with Cyclic Symmetry: A linear model trained on data augmented by a cyclic group $C_3$ . The theory predicts that the direction of the weight matrix projected onto the orthogonal complement of the all-ones vector is conserved. Experiments confirm that this quantity remains constant under augmented gradient descent, whereas it drifts significantly without augmentation.
Lightning Attention: A single-head Lightning Attention model trained with right-action $O(d)$ augmentation. The theory predicts that the column space of the concatenated parameter block $P = [Q, K, V]$ is preserved. Numerical experiments show that the distance between the column spaces of the parameters at different training steps remains near zero for augmented training, while it grows for standard training.

Significance and Claims

The paper claims to provide a rigorous theoretical boundary for when data symmetry leads to new optimization constraints:

Generality vs. Specificity: It clarifies that for the broad class of analytic non-polynomial losses (common in classification), data symmetry does not create new conservation laws. The phenomenon is specific to polynomial losses (like MSE) and architectures that allow for a "tensorizable" separation of inputs and parameters.
Novelty: The work identifies a mechanism where discrete data symmetries can induce continuous symmetries in the parameter space, leading to conserved quantities (specifically, the preservation of subspaces) that are not present in the standard training dynamics.
Limitations: The authors modestly note that their framework relies on specific technical assumptions (analyticity, tensorizability) and does not claim to cover all possible new conserved quantities from symmetric data. They also acknowledge that not all $H$ -symmetries are realizable in every parametrization.

In summary, the paper establishes that while data symmetry generally does not alter the fundamental conservation laws of gradient flow for standard classification losses, it can fundamentally constrain the trajectory of MSE-trained tensorizable networks by preserving the geometric structure (column spaces) of specific parameter blocks.