Transformative or Conservative? Conservation laws for ResNets and Transformers

Transformative or Conservative? Conservation laws for ResNets and Transformers

Speakers: Sibylle Marcotte, Rémi Gribonval, Gabriel Peyré

Conference: ICML 2025

YouTube: https://slideslive.com/39043897

Overview

This detailed technical article explores the groundbreaking work presented at ICML 2025 by Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré on conserved quantities during the training dynamics of neural networks, with a particular focus on ResNets and Transformers. The talk delves into the fundamental mathematical properties that remain invariant as a neural network undergoes training via gradient flow, and how these properties translate to the more practical Stochastic Gradient Descent (SGD) setting. By identifying and characterizing these conservation laws, the research provides crucial insights into the implicit bias of training dynamics and offers powerful mathematical tools for analyzing the convergence and behavior of complex deep learning models.

The motivation behind this investigation is multi-faceted. Firstly, conservation laws inherently link the initial state of a neural network to its final trained state, revealing shared properties that persist throughout the optimization process. This connection is vital for understanding the implicit bias – the phenomenon where optimization algorithms, despite being capable of finding multiple global minima, consistently converge to solutions with specific desirable properties. Secondly, these laws serve as potent analytical instruments, simplifying the study of complex, high-dimensional training dynamics and potentially enabling more rigorous convergence proofs.

Historically, the study of conservation laws in neural networks has been confined to highly simplified architectures. This talk addresses a critical gap by extending this analysis to Convolutional Neural Networks (CNNs), Attention Layers, deep ResNets, and Transformers, which are the backbone of modern AI. The findings demonstrate that even in these realistic and complex settings, underlying conservation principles exist and can be systematically identified, offering a new lens through which to understand the often mysterious behavior of deep learning training.

Background

▶ Watch: Introduction: Conserved quantities in neural network training dynamics (0:00)

The core problem addressed by this research is the minimization of empirical risk, a ubiquitous task in machine learning. Given a dataset $\mathcal{D} = \{(X_i, Y_i)\}_{i=1}^N$, the objective is to find the parameters $\theta$ of a neural network $G$ that minimize the average loss $L(G(X_i, \theta), Y_i)$. The training dynamic considered initially is the gradient flow, a continuous counterpart to gradient descent. This is described by the differential equation $\frac{d\theta}{dt} = -\nabla L(\theta)$, where $\theta_0$ is the initialization and the dataset is fixed. The goal is to identify functions $H(\theta)$ that remain constant along any trajectory defined by this gradient flow, irrespective of the specific dataset or initialization. Such a function $H$ is termed a conservation law.

The existence of conservation laws is not a new concept in physics, where Noether's theorem elegantly links symmetries in a system to conserved quantities. In the context of neural networks, these laws provide insights into the invariant structures within the parameter space during training. Prior work has identified such laws for very simple neural networks, such as two-layer linear or ReLU networks. For instance, in a 2D to 1D linear neural network, functions like $U^2 - ||V||^2$ (where $U, V$ are weight matrices) are known to be conserved. However, extending this analysis to deeper and more complex architectures like Convolutional Neural Networks, Attention Layers, ResNets, and Transformers presents significant challenges due to their intricate parameter spaces and non-linearities.

A crucial prerequisite for discussing conservation laws is a rigorous definition of their "independence." If $H$ is a conservation law, then any arbitrary function of $H$, say $f(H)$, is also technically a conservation law, leading to functional redundancies. To counter this, the authors define independence by requiring that the gradients of independent conservation laws are pointwise linearly independent. This ensures that each identified law provides distinct information about the system's invariants. The primary challenge lies in characterizing the space $W_\theta$, which is defined as the span of all possible gradients of the empirical risk in $\theta$ for any possible dataset. This space is generally difficult to manipulate directly, necessitating a more tractable approach to identify and count conservation laws.

Key Findings

▶ Watch: Method: Reparameterization simplifies finding conservation laws (3:20)

The research makes several pivotal contributions, significantly expanding our understanding of conservation laws in deep learning:

1. Exhaustive Conservation Laws for Shallow Architectures: For a range of shallow neural network architectures, including linear networks, ReLU networks, Convolutional Neural Networks, and Attention Layers, the authors have not only discovered specific conservation laws but also rigorously proved that these identified laws are exhaustive. This means there are no other independent conservation laws for these specific architectures under gradient flow. This result extends previous findings for linear and ReLU networks to more complex, practically relevant single-layer modules.

1. Extension to Deep ResNets and Transformers via Skip Connections: The study successfully extends the analysis to deep architectures like ResNets and Transformers. A key insight is that these deep networks can be viewed as compositions of two-layer modules (e.g., convolutional ReLU blocks or attention layers) combined with skip connections. The authors demonstrate that conservation laws depending only on the parameters of a specific residual block are precisely those of that shallow block taken in isolation. Furthermore, they establish that the full set of conservation laws for a network with skip connections is identical to the set of conservation laws for the same network without skip connections. This surprising result simplifies the analysis significantly, as skip connections, rather than adding complexity to the conservation landscape, often aid in its mathematical characterization.

1. Approximate Conservation under Stochastic Gradient Descent (SGD): While the primary theoretical framework relies on continuous gradient flow, the talk also investigates how these identified conservation laws behave under the discrete, noisy dynamics of SGD, which is the standard optimization algorithm in practice. Through both theoretical bounds and empirical validation, the research shows that these conservation laws remain approximately conserved during SGD training. Specifically, the error in conservation is bounded by a term proportional to $K\tau^2$, where $K$ is the number of SGD steps and $\tau$ is the learning rate, indicating that for small learning rates, the laws hold reasonably well.

1. Link to Underlying Symmetries: During the Q&A, it was clarified that the discovered conservation laws are indeed linked to underlying symmetries of the neural network architectures, akin to Noether's theorem. While not explicitly detailed in the main talk, the speakers mentioned "rescaling symmetries" as an example, reinforcing the deep mathematical foundation of their findings. This connection highlights the elegance and fundamental nature of these conservation principles.

Technical Deep Dive

▶ Watch: Shallow Networks: Exhaustive conservation laws for convolution, attention (3:50)

The technical core of this research revolves around a sophisticated characterization of conservation laws and the use of reparameterization and Lie algebra computations.

A function $H(\theta)$ is a conservation law if its value remains constant along any gradient flow trajectory, irrespective of the dataset or initialization. Mathematically, this implies that the gradient of $H$ must be orthogonal to the gradient of the empirical risk for all possible datasets. More formally, for all $\theta$, $\nabla H(\theta)$ must be orthogonal to every vector in $W_\theta$, where $W_\theta$ is defined as the space spanned by all possible gradients of the empirical risk $\nabla L(\theta)$ for any possible dataset $(X_i, Y_i)$. Initially, $W_\theta$ is a high-dimensional and complex space, making direct manipulation challenging.

To simplify the characterization of $W_\theta$, the authors introduce the concept of reparameterization $\phi$. The neural network $G(\theta)$ is factorized as $G_\phi(\phi(\theta))$, where $\phi(\theta)$ is a set of intermediate parameters. This reparameterization is chosen such that it captures the essential structural properties of the network. For various architectures, specific reparameterizations are employed:

• For linear neural networks, $\phi(\theta)$ is simply the product of all weight matrices.
• For ReLU networks, $\phi(\theta)$ involves pairs of rank-one matrices, e.g., $U_i V_i^T$.
• For convolutional layers, the reparameterization utilizes circulant matrices generated by the filters.
• For attention layers, $\phi(\theta)$ groups together the query and key matrices, as well as the value and output matrices.

The crucial technical advancement is demonstrating that for all these examples, there is an equality between $W_\theta$ and $W_\phi(\theta)$, where $W_\phi(\theta)$ is a finite-dimensional functional space spanned by the gradients of the coordinates of the reparameterization $\phi$. This transformation significantly simplifies the problem, as $W_\phi(\theta)$ is much easier to manipulate computationally. This allows for a much simpler characterization of conservation laws: $H(\theta)$ is a conservation law if and only if $\nabla H(\theta)$ is orthogonal to $W_\phi(\theta)$.

With this simplified characterization, the authors employ Lie algebra computations on the space $W_\phi$ to systematically discover and prove the exhaustiveness of conservation laws for shallow networks. For instance, in linear and ReLU networks, previously known laws like $U^2 - ||V||^2$ are rediscovered. For convolutional and attention layers, novel conservation laws are identified and proven to be exhaustive using this methodology.

Extending this analysis to deep ResNets and Transformers requires a different approach. A ResNet can be conceptualized as a composition of two-layer convolutional ReLU neural networks with skip connections. Similarly, a Transformer can be seen as a composition of either a two-layer ReLU network or an attention layer, also with skip connections. Instead of seeking conservation laws that depend on the entire parameter set, the researchers first focus on laws that depend only on the parameters associated with a single residual block. They prove that these block-specific conservation laws are precisely those identified for the shallow (isolated) residual block. A significant finding is that the presence of skip connections, rather than complicating the analysis, actually simplifies the constructive proof for identifying the space $W_\theta$. The set of all conservation laws for a network with skip connections is shown to be identical to that of the same network without skip connections. Numerical experiments further suggest that all conservation laws in deep networks are combinations of these block-specific conservation laws.

Experimental Setup & Results

▶ Watch: Deep Networks: Block-specific laws, skip connections' role (4:50)

While the core of the talk is theoretical, the research includes empirical validation to demonstrate the practical relevance of these conservation laws, particularly in the context of Stochastic Gradient Descent (SGD). The main objective of the experiments was to track how well the identified conservation laws hold when the training dynamics are discrete and noisy, as is the case with SGD.

The theoretical analysis provides a bound on the "conservation error" for a given conservation law $H$ under SGD. Specifically, the error, defined as the difference between the conservation law evaluated after $K$ SGD steps and its value at initialization, is bounded by $K\tau^2$, where $\tau$ is the learning rate. This bound suggests that for small learning rates, the conservation laws should remain approximately constant over a significant number of training steps.

To test this, the authors conducted experiments on two representative deep learning models:

1. ResNet-18 trained on the CIFAR-10 image classification dataset.
2. A Transformer model trained on the IMDB sentiment analysis dataset.

During the training of these models using SGD, the "conservation error" for one of the identified conservation laws was tracked over time. The results consistently showed that the observed increase in the conservation error was indeed proportional to $\tau^2$, as predicted by the theoretical bound. The slope coefficients of the error accumulation closely matched the expected quadratic dependence on the learning rate.

These experimental findings are crucial because they bridge the gap between continuous gradient flow theory and practical deep learning optimization. They demonstrate that despite the inherent noise and stochasticity of SGD, the conservation laws derived from the continuous setting still provide a meaningful and approximately preserved structure, reinforcing their utility for understanding real-world training dynamics.

Practical Implications

▶ Watch: Discrete Dynamics: Conservation laws' behavior under SGD (5:50)

The discovery and characterization of conservation laws in deep neural networks carry significant practical implications for various roles within the ML ecosystem, from model builders to infrastructure teams.

Firstly, these laws offer a deeper understanding of implicit bias. The fact that conservation laws hold throughout training means they impose constraints on the possible trajectories and, consequently, on the properties of the final converged solution. By linking the initialization to the final state, these laws can help explain why optimization algorithms, despite the vastness of the parameter space, often converge to solutions with specific desirable characteristics (e.g., good generalization, specific types of sparsity or low-rank structures). This understanding can inform the design of better initialization schemes or regularization techniques.

Secondly, conservation laws serve as powerful mathematical tools for convergence analysis. As mentioned in the Q&A, they can be used to prove convergence properties of gradient flow dynamics. In certain scenarios, particularly when conservation laws are "balanced" at initialization (i.e., equal to zero), they can allow for the re-expression of a high-dimensional gradient dynamic in a lower-dimensional space. This dimensionality reduction can drastically simplify theoretical analysis and potentially lead to more tractable proofs of convergence or stability.

For practitioners and model builders, understanding these invariants can provide a new perspective on debugging and diagnosing training issues. Deviations from expected conservation behavior could signal problems with the model, data, or optimization setup. While not directly offering immediate hyperparameter tuning advice, this foundational knowledge contributes to a more principled approach to model development.

For infrastructure teams and deployers, while less direct, the insights into the fundamental dynamics of neural network training can inform long-term strategies for hardware and software optimization. Knowledge of inherent structural properties can guide the development of more efficient numerical methods or specialized hardware architectures that implicitly respect these conservation principles.

Tradeoffs and Limitations: The primary theoretical framework is based on gradient flow, a continuous idealization. While the work demonstrates approximate conservation under SGD, the exactness of these laws is lost in the discrete setting. The tightness of the $K\tau^2$ bound depends on the learning rate $\tau$, meaning conservation might degrade significantly with larger learning rates or very long training horizons. Furthermore, the analysis assumes arbitrary datasets. As noted in the Q&A, restricting the dataset to specific distributions or structures might alter the number or form of conservation laws, suggesting that the identified laws are universal across any dataset, but more specific laws might exist for specialized data contexts. The talk also highlighted that in certain degenerate cases, like a three-layer MLP where all weights become negative, leading to zero output, a different and larger set of conservation laws emerges, indicating the sensitivity of these laws to the specific functional behavior of the network.

Key Takeaways

• Conservation laws exist and can be systematically identified in the training dynamics of deep neural networks, providing fundamental invariants.
• The research establishes exhaustive sets of conservation laws for shallow architectures including linear, ReLU, convolutional, and attention layers under gradient flow.
• For deep ResNets and Transformers, conservation laws primarily depend on individual residual blocks, and skip connections simplify their mathematical characterization without altering the fundamental set of laws.
• Despite the stochasticity of SGD, these conservation laws remain approximately conserved in practice, with the conservation error bounded by $K\tau^2$ (number of steps $\times$ learning rate squared).
• Conservation laws are crucial for understanding the implicit bias of optimization algorithms and serve as powerful mathematical tools for convergence analysis and potentially reducing the dimensionality of training dynamics.
• These laws are linked to underlying symmetries of the neural network architectures, providing a deeper theoretical foundation for deep learning.

About the Speaker(s)

The work presented was a collaborative effort by Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré. Sibylle Marcotte presented the talk, detailing their joint research on conserved quantities in neural network training dynamics. While specific institutional affiliations were not explicitly stated during the presentation, their collective work significantly contributes to the theoretical understanding of deep learning, focusing on the mathematical properties that govern the behavior of complex models like ResNets and Transformers during optimization. Their research bridges theoretical machine learning with practical applications, providing rigorous analysis for widely used architectures.