🤖 machine learning

Conservation Laws for Modern Neural Architectures

This paper establishes a unified theoretical framework to characterize conservation laws in gradient flow for modern neural architectures—including models with GELU, SiLU, SwiGLU activations, various attention mechanisms, and Mixture-of-Experts designs—validating these findings through experiments to better explain the implicit bias of over-parameterized models.

Original authors: Viet-Hoang Tran, Vinh Khanh Bui, Tan Lai Ngoc, Nam Nguyen, Tuan Dam, Tan M. Nguyen

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Viet-Hoang Tran, Vinh Khanh Bui, Tan Lai Ngoc, Nam Nguyen, Tuan Dam, Tan M. Nguyen

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 watching a massive, complex machine being built and adjusted in real-time. This machine is a modern AI model (like the ones that write stories or recognize images). As the machine learns, its internal knobs and dials (called "parameters") are constantly turning to minimize mistakes.

This paper is like a detective story about what stays the same while everything else is changing.

The Big Idea: The "Unchanging Rules" of Learning

Usually, we think of training an AI as a chaotic process where numbers go up and down randomly. But the authors discovered that even in this chaos, there are strict Conservation Laws.

Think of these laws like the conservation of energy in physics. If you roll a ball down a hill, its potential energy turns into kinetic energy, but the total energy stays the same. Similarly, as an AI learns, certain mathematical combinations of its internal settings remain perfectly constant, no matter how much data it sees or how long it trains.

The paper's main goal was to find these "hidden invariants" for the most popular, modern AI architectures used today.

The Detective Work: How They Found the Rules

The authors didn't just guess; they used a mathematical "magnifying glass." They asked: "If we change the data or the starting point, what mathematical formulas involving the AI's settings will never change?"

They treated the AI's learning process like a flowing river. Even though the water (the data and the specific path) changes, the shape of the riverbed (the architecture) forces the water to follow specific patterns. They mapped out these patterns for three major types of modern AI components:

1. The "Smooth" Feedforward Networks (The Brains)

Modern AI uses special "activation functions" (mathematical switches) to decide what to pay attention to.

GELU and SiLU: These are like smooth, gentle switches. The authors found that for networks using these, nothing interesting stays constant other than the fact that the system is just doing its job. It's like a smooth river with no whirlpools; the water just flows.
SwiGLU: This is a more complex switch used in top-tier models (like LLaMA). Here, they found a hidden balance. Imagine two weights, A and C, acting like a seesaw. If A gets heavier, C must get lighter in a very specific way to keep the total "balance score" constant. The paper proved exactly how this balance works.

2. The Attention Mechanism (The Focus)

This is the part of the AI that decides which words in a sentence are important.

Standard Multi-Head Attention: The authors solved a puzzle that previous researchers couldn't finish. They found that for each "head" (a sub-processor), there are two pairs of weights (Query/Key and Value/Output) that must maintain a specific difference. Think of it like a tug-of-war: the strength of the "pull" team must always match the strength of the "push" team in a precise mathematical way.
Rotary Positional Encoding (RoPE): This is a fancy way of telling the AI where words are located in a sentence (e.g., "first word" vs. "last word"). The authors found that this changes the rules entirely. Instead of a simple tug-of-war, the weights now have to rotate in a specific circle to keep the "rotational balance" constant. It's like a dancer spinning; their speed and position change, but their angular momentum stays fixed.

3. Mixture-of-Experts (The Team of Specialists)

Imagine an AI that doesn't use one giant brain, but a team of 100 smaller experts, where only a few are called to work on each problem.

The Gating Mechanism: This is the "manager" that decides which experts work. The authors found that the manager's decisions and the experts' settings are linked.
The Discovery: Whether the manager picks the top 2 experts or the top 10, or whether it uses a "soft" vote (softmax) or a "hard" vote (sigmoid), the total "weight" of the team's decisions remains constant. The individual experts might change, but the sum of their influence follows a strict rule.

The Experiment: Does it Work in Real Life?

The authors didn't just do math on paper. They built small AI models and trained them on real data (like images and text).

They tracked these "conserved quantities" over thousands of steps.

The Result: Just like a pendulum swinging, the values stayed incredibly stable.
The Catch: When they used a very fast learning rate (a big step size), the values wobbled a little bit, but the wobble was predictable and small. If they slowed down the learning, the values stayed almost perfectly still. This proved that these laws aren't just theoretical; they are real, physical constraints of how these AI models learn.

Summary

In simple terms, this paper says: "Modern AI models are not chaotic messes. They are governed by hidden, unbreakable mathematical rules."

Just as a car engine has rules about how fuel turns into motion, these AI models have rules about how their internal numbers must balance each other out. The authors successfully wrote down the rulebook for the most advanced AI architectures currently in use, showing us exactly what stays the same while the AI learns.

Technical Summary: Conservation Laws for Modern Neural Architectures

Problem Statement
Understanding the training dynamics of over-parameterized neural networks is central to explaining their success. A key mechanism in this understanding is the concept of conservation laws—quantities that remain invariant along optimization trajectories (specifically gradient flow). While these laws are well-characterized for linear and ReLU networks, they remain largely unexplored for modern architectures. Existing frameworks struggle to provide complete characterizations for contemporary models utilizing smooth activations (e.g., GELU, SiLU), complex attention mechanisms (Multihead Attention with positional encodings), and Mixture-of-Experts (MoE) designs. Furthermore, previous work on Multihead Attention (MHA) was limited to single-head settings, leaving the multi-head case as an open problem.

Methodology
The paper develops a unified theoretical framework to characterize conservation laws for modern deep learning models. The approach builds upon the work of Marcotte et al. (2023; 2025) but diverges in its analytical strategy:

Problem Formulation: The authors define a conservation law $h(\theta)$ as a function invariant under the Euclidean gradient flow $\dot{\theta} = -\nabla L_D(\theta)$ . To make the problem tractable, they adopt the assumption that conservation laws must hold for any dataset, isolating the implicit bias induced solely by the architecture and optimization dynamics.
Reduction to PDEs: Instead of relying on Lie-theoretic machinery, the authors reduce the conservation condition to a system of partial differential equations (PDEs). Specifically, they require that the gradient of the conservation law $\nabla_\theta h$ be orthogonal to the span of the model's Jacobian vectors for all inputs.
Complex-Analytic Techniques: The core of the proof strategy involves viewing the orthogonality conditions as complex-valued functions. By restricting these functions to specific subspaces and utilizing the identity theorem for meromorphic functions, the authors derive input-independent constraints on the gradients of $h$ .
Characterization via Invariants: The resulting PDE constraints are solved to identify "characteristic invariants." The conservation laws are shown to be constant on the connected components of the level sets defined by these invariants.

Key Contributions and Results
The paper provides complete characterizations of conservation laws for three primary classes of modern architectures:

Feedforward Networks (FFNs):
- GELU and SiLU: The authors prove that for FFNs using GELU or SiLU activations, the only conservation laws are constant functions. No non-trivial invariants exist.
- SwiGLU: In contrast, SwiGLU architectures (which involve a multiplicative interaction between parameter matrices) exhibit non-trivial conservation laws. The invariants are determined by the quantities $\|A_{:,i}\|^2 - \|C_{i,:}\|^2$ for each hidden unit $i$ , where $A$ and $C$ are the weight matrices involved in the gating mechanism.
Multihead Attention (MHA):
- Vanilla MHA: The authors resolve the open problem regarding multi-head attention, confirming that conservation laws are characterized by the invariants $Q_i^\top Q_i - K_i^\top K_i$ and $V_i^\top V_i - O_i^\top O_i$ for each head $i$ .
- Sinusoidal Positional Encoding (PE): Since sinusoidal PE acts as a fixed, bijective shift to the input, the conservation structure remains identical to the vanilla MHA case.
- Rotary Positional Encoding (RoPE): RoPE fundamentally alters the internal structure of attention. The authors derive a distinct set of invariants: $\|Q_i^{(j)}\|_F^2 - \|K_i^{(j)}\|_F^2$ (where $Q$ and $K$ are decomposed into $2 \times 2$ blocks) and $V_i^\top V_i - O_i^\top O_i$ . This demonstrates that RoPE introduces a different class of conservation laws compared to absolute positional encodings.
Mixture-of-Experts (MoE):
- Dense and Sparse MoE: The study covers both dense gating and Sparse MoE (SMoE) with Top- $k$ routing. The results show that conservation laws for SMoE coincide with those of the corresponding dense MoE.
- Gating Mechanisms: The invariants are localized at the expert level (following the SwiGLU invariants for each expert) and include an additional invariant for the gating weights: the sum of the gating weight vectors $\sum W_i$ .
- Sigmoid Gating: The authors demonstrate that replacing softmax gating with normalized sigmoid gating does not alter the set of conservation laws; the invariants remain identical to the softmax-gated case.

Experimental Validation
The theoretical findings are validated empirically across language modeling (WikiText-103, Penn Treebank) and computer vision (CIFAR-10, ImageNet-1K) tasks.

Setup: Experiments utilize both full-batch gradient descent (to approximate continuous-time flow) and mini-batch SGD on small and large-scale datasets.
Metrics: The authors track the "conservation error" (deviation of invariants from their initial values) over training steps.
Findings: The experiments confirm that conservation laws remain bounded under training. Specifically, the error scales with the learning rate $\tau$ and the number of steps $k$ as predicted by theory ( $O(\tau^2 k)$ for constant step sizes). The results hold consistently across MHA, SwiGLU FFNs, RoPE, and MoE gating mechanisms, regardless of the dataset or training regime.

Significance and Claims
The paper claims to advance the theoretical understanding of training dynamics by providing the first complete characterization of conservation laws for a broad spectrum of modern neural architectures.

Theoretical Unification: It establishes a unified framework that extends previous results from linear/ReLU networks to smooth activations and complex attention mechanisms.
Resolution of Open Problems: It explicitly solves the open problem of characterizing conservation laws for multi-head attention and extends the analysis to include positional encodings and MoE gating.
Implicit Bias: By identifying these invariants, the work sheds light on the geometric constraints inherent to modern architectures, revealing properties that persist from initialization to convergence.
Limitations: The authors modestly note that their reliance on complex-analytic techniques (meromorphic extensions) currently limits the framework's application to architectures involving non-meromorphic components, such as layer normalization or deep compositions involving functions like $\sqrt{x}$ or $e^{1/x}$ . They also acknowledge that the dataset-independence assumption, while principled for isolating architectural bias, may be viewed as restrictive in broader contexts.