Transformative or Conservative? Conservation laws for ResNets and Transformers

We would like to thank the reviewer for his positive comments and his insightful feedback.

Q1d "The Figure 1 should state which conservation law is being tested. Some additional numerical tests, especially regarding the new conservation laws (attention, residual blocks, etc), could improve the paper.” + “What conserved quantity is being tracked in Figure 1 and the Appendix Figures?”

Thank you for pointing out that point. The caption will be fixed. As detailed in section 5.2, the conserved quantity that is tracked is the sum of all conserved quantities of equation 14 associated to the first convolutive layer. In particular, Figure 1 focuses on some new conservation laws of our paper (ResNet architecture, thus skip connection and convolutive layers). Concerning numerical experiments regarding new conservation laws associated to attention layers, see our answer to your question Q3d.

Q2d Some informal discussions after each result could make the paper more readable

Thank you for the suggestion that will be taken care of in the final version.

Q3d In addition, it could improve the paper to test numerically some of the newer claims to demonstrate the power of this approach in the newer settings that go beyond Marcotte et al 2023.

Beside numerically testing our new theoretical founding in the case of a ResNet’s training (see our answer to your question Q1d) in our paper (cf Fig. 1), we propose to add this experiment https://anonymous.4open.science/r/Numerical_experiments-2730/Transformer_IMDb.ipynb that tracks the evolution of a conservation law during the training of a transformer. More precisely, we train a transformer on the sentiment analysis dataset (IMDb dataset) and we track the evolution of $| h(\theta_n) - h (\theta_0)|$ where $h (\theta) = \| \| Q_1 Q_1^\top - K_1 K_1^\top\| \|_p$ (ie the norm of Frobenius of the conservation law given by Corollary 3.10, where $Q_1$ and $K_1$ are the query and key matrices of the first head and from the first layer. Similarly to Figure 1, we obtain the same bound O (step * learning_rate^2). This new experiment also focuses on new conservation laws (attention block and skip connection).

Q4d The authors write that all conservation laws in the transformer case are functions of $QQ^\top−KK^\top$. Is this full matrix difference conserved or just specific nonlinear functions of it? I ask because in the deep linear case, the conservation laws take the form $W_l W_l^\top−W_{l+1}^\top W_{l+1}$ for all matrix entries. Is there a method to extract which nonlinear functions of this difference are conserved?

Yes, the full matrix difference is conserved, and this implies that any linear or nonlinear function of these matrix entries is also conserved: when a function $h(\theta)$ is conserved, any function$\Psi(h(\theta))$ is also conserved. This is exactly why we need to define a notion of independent functions (Definition 2.10), so that we remove all these functional redundancies (see also our Proposition 2.12). We will stress on that point on the final version of our paper.

Q5d Do the authors think that normalization layers (like layernorm) will make the number of conservation laws larger? If the gradient is independent of the scale / radial direction of the features, then it could be that the dynamics are confined to a lower-dimensional subspace

Since these layers are not parameterized and only marginally change the diversity of possible inputs to the subsequent layers, we expect them to have no effect on the analysis and, therefore, on the number of conserved laws. However, if you consider in other settings, another type of normalization that operates directly on the parameters, you should expect (as in the case of a cross-entropy layer) new conservation laws.

We thank the reviewer for the positive comments and insightful feedback.

Q1c about structure theorem

We stated theorem 2.1 right at the beginning to lighten the notation in the rest of the paper. It is new to the best of our knowledge.

Q2c CIFAR results on the Transformer architecture

In this experiment https://anonymous.4open.science/r/Numerical_experiments-2730/Transformer_IMDb.ipynb, we train a transformer and track the evolution of the conservation law given by Corollary 3.10, confirming the O (step * learning_rate^2) behaviour. See also our answer to Q3d of reviewer iC26 for more detail.

Q3c how this relates to conservation laws in PDEs and scientific machine learning literature (...)

Thank you for these references! The common technical aspect to our work and these references is the link between symmetries and conservation laws. However, while we uncover conservation laws satisfied by the parameters of a neural network during the flow associated to their training, these references design neural networks implementing functions that satisfy certain physics-driven/PDE-driven conservation laws.

Q4c in Essential References Not Discussed

Thank you for these references! The 'gradient flow' ODE (3) describes parameter dynamics during training, while the Neural ODEs governs the input transformation $x$ through an infinitely deep network, not the parameter evolution $\theta$. Connections to Neural ODEs demonstrate the broader utility of Lie algebraic techniques for universality proofs in ML, see e.g. arXiv:1908.07838, arXiv:2404.08289. We will add this discussion to our related work section.

Q5c How to support Adam as future work

Analyzing the original Adam algorithm is challenging due to the presence of two momentum parameters. However, a simplified version can be readily analyzed using our framework, and we will update the manuscript accordingly. A more detailed analysis of Adam itself remains an interesting direction for future work.

The continuous limit of Adam yields an ODE where the gradient is replaced by its sign. $W(\theta)$ is spanned by signed gradients, it satisfies the assumptions of our results. In the simple case of a $2$-layer linear neural network we easily relate it with $\phi(\theta) = U V^\top$ and deduce in this code https://anonymous.4open.science/r/Numerical_experiments-2730/Adam.ipynb that $dim Lie (W(\theta)) = D$, so that there is no conservation law (except when $U$ and $V$ are scalars, in which case $|U|- |V|$ is conserved). For more detail, see our answer to Q3b of reviewer RSFD. We will mention this at the end of our paper.

Q6c (...) connected to PDEs with a spatial component (...)?

This is a good question that relates closely to your earlier question Q3c. In response to Q3c, we mentioned that our analysis focuses on conservation with respect to the parameters $\theta$, rather than the input $x$ of the networks. However, we would like to point out that it is indeed possible to apply our approach to networks that include layers acting spatially on a signal or image $x$ (such as differential operators in the context of PDEs). In particular, one can design weight matrices in the network to operate directly on the spatial domain, typically corresponding to a discretization of a differential operator. This approach could provide a useful starting point for studying architectures such as PINNs or Neural Operators. After discretization, translation-invariant differential operators (e.g. Laplacians) become convolutions, which align naturally with our current methodology. That said, extending our results to infinite-dimensional operators remains an open and promising direction for future work. We will include these possibilities in the final version of the manuscript. If the reviewers have specific architectures or PDEs in mind, we would be very interested to hear more about them.

Q7c (...) extended to time series data?

Transformers are indeed well-suited for time series data (as well as for NLP tasks involving next-token prediction). The only difference compared to the setup presented in our paper is the inclusion of masking to make the architecture causal. This modification does not affect our theoretical framework—the conservation laws remain unchanged. Additionally, the numerical code used to test the approximate conservation law under SGD can be applied in the causal setting simply by enabling the use_mask =True flag (cf https://anonymous.4open.science/r/Numerical_experiments-2730/Transformer_IMDb.ipynb). In practice, we observe that the predictions of our theory remain accurate even when accounting for discrete stepping in the algorithm. We will update the manuscript to reflect this. Extending our methods to other types of causal architectures, such as recurrent neural networks and Selective State Space Models, represents a promising direction for future work, which we will also mention.

We would like to thank the reviewer for his constructive feedback and positive comments.

Q1b challenges in proving the completeness of laws beyond blockwise laws

From a theoretical perspective, the Lie algebra for deeper cases becomes infinite dimensional, which makes the analysis significantly more involved than in block setting. Numerical experiments briefly described in Section 4.4 (see also the supplementary material), however, suggest that combining the laws associated with each block yields a complete set of laws for the deeper case. We will emphasize that theoretical challenge at the end of section 4.

Q2b leveraging newly discovered conservation principles in practice

A first key application of conservation laws is to understand how initialization influences optimization dynamics—for example, by deriving an implicit Riemannian metric parameterized by the initialization. The results of arXiv:1910.05505 (in the case of linear neural networks) show that initializing with certain values of the conserved function can lead to accelerated convergence. Another application lies in the so-called “Edge of Stability” analysis, which, in the linear case, explains how increasing the step size in gradient descent leads to minimizers with small values of the conservation law (“balanced condition”). Extending this analysis beyond the linear case, as initiated in arXiv:2502.20531, is an interesting direction for future research.

We will insist more on these practical applications in the introduction of our paper.

Q3b Could the authors briefly outline their planned approach (...)

Regarding normalization layers, they can be treated as a new block and we expect much of the analysis from Appendix M to be adaptable. Then, the main challenge to fully address multihead attention is to prove the ``minimimality’’ of the associated reparameterization, which requires more geometric insight than the proof in Appendix F for the case of a single head. Finally, combining our findings on blockwise laws and the non-Euclidean setting of Marcotte et al. 2024 is likely to yield advances regarding other optimizers. Indeed Marcotte et al. 2024 can deal with mirror gradient flow or with natural gradient flow. Analyzing the original Adam algorithm is challenging due to the presence of two momentum parameters. However, a simplified version can be readily analyzed using our framework, and we will update the manuscript accordingly. A more detailed analysis of Adam itself remains an interesting direction for future work.

To make the analysis more tractable, we consider a limiting case of Adam where both momentum parameters tend to infinity. In this regime, the algorithm approaches sign gradient descent, and its continuous-time limit is given by the ODE: $\dot \theta (t) = - \text{sgn} (\nabla L_Z(\theta(t))) \ (1).$ In our framework, this corresponds to considering the linear space of signed gradients: $W(\theta) := \text{span} \\{ \text{sgn} (\nabla L_Z(\theta)) : Z \\}$ for any parameter $\theta$. An interesting property of $W(\theta)$ is that, due to the quantization introduced by the sign function, $W(\theta)$ is locally constant around generic values of $\theta$. This leads to the identity $\text{Lie}(W)(\theta) = W(\theta)$. By Theorem 2.11, this implies that in a neighborhood of $\theta$, there are exactly $D - \text{dim} (\text{Lie}(W)(\theta)) = D - \text{dim} (W(\theta))$ independent conservation laws. In particular, if $\text{dim} (W(\theta)) = D$, then no conservation law exists in that region.

In the simple case of a $2$-layer linear neural network $g((U, V), x) = U V^\top x$, we easily show (direct with Theorem 3.4) that $W(\theta) = \text{span} \\{ \text{sgn} ( \partial \phi (\theta)^\top w): w \\}$ where $\phi(\theta) = U V^\top$, which is easily computationaly tractable. In this code https://anonymous.4open.science/r/Numerical_experiments-2730/Adam.ipynb, we test it on different dimension settings and we obtain that, except in the case where $U$ and $V$ are both scalars, we have $\text{dim} W(\theta) = D$, so that there is no conservation law. In the special case where $U$ and $V$ are both scalars, we obtain that $\text{dim} W(\theta) = 1$, and thus there is only $2-1= 1$ conservation laws: indeed, $|U|- |V|$ is conserved.

While we only provide numerical explanations in the $2$-layer linear case, we expect that it holds for all other architectures. This indicates a radically different geometric behavior of the Adam optimizer. It is worth recalling that adding a momentum term also leads to a loss of the total number of conservation laws (in most settings (cf Marcotte et al. 2024), there is no conservation law at all when considering a momentum term). A better understanding of these phenomena is an interesting future direction. We will mention this discussion at the end of our paper.

We thank the reviewer for his constructive feedback.

Q1a the abstract claims (...)

Thank you for pointing this out, you’re right. Since this is not the focus of our paper as we don't have a detailed theoretical analysis of this phenomenon, we will remove this sentence from the abstract. Instead, we will add the following simple example at the end of section 4.4.

Example: Consider a ReLU neural network with two-residual blocks, $g((u, v, s, t), x) = x + u \sigma (v x) + s \sigma(t (x + u\sigma( v x))$ with $(u, v, s, t) \in \Omega \subseteq \mathbb{R}^4$ and $x \in \mathbb{R}$. While there are two "block" conservation laws: $u^2 - v^2$ and $s^2 - t^2$, we exhibit a domain $\Omega$ where there are three conservation laws. For this, consider $\Omega$ the set of all parameters such that $sgn(t) = sgn(v) = sgn(u).$ Then, for any $x$:

• either $vx, sx \leq 0$ and then $g(\theta, x) = x$ and thus $\nabla_{\theta} g = 0$;
• or $vx, sx \geq 0$ and then $g(\theta, x) = x + uv x + stx +stuv x$ as $vx, sx, tu \geq 0$, and then $1/x \nabla_{\theta} g=: \chi$ is a vector field that does not depend on $x$.

Thus the space $W_{\theta}^{g, l} = \mathbb{R} \chi(\theta)$ (cf Proposition 2.3) is made by only one non-null vector field. Its associated Lie algebra is thus itself, and by Theorem 2.11, there are exactly 3 conservation laws as claimed.

This phenomenon seems however restricted to the scalar case, as we empirically show in Section 4.4.

Q2a h cannot depend on the dataset (...)

Thank you for this insightful comment. We view the dataset-independence condition as a deliberate design choice rather than a weakness, as it yields general insights on the implicit bias stemming solely from the architecture and optimization scheme.

While we agree that introducing dataset dependency could enrich the analysis, it also opens the Pandora’s box of how to constrain this dependency to ensure the problem remains nontrivial and the conservation law remains informative. For instance, allowing unrestricted dataset dependency would make the problem on the number of independent conservation laws trivial, as we now explain.

Indeed, the Lie algebra framework of Marcotte et al. can also be applied to vector fields associated with specific datasets (see Definition A.3). For a fixed dataset, the associated functional space (Definition A.4) becomes one-dimensional, and its generated Lie algebra is simply itself. By Frobenius' theorem, this yields exactly D-1 conserved functions which typically exhibit non-trivial dataset dependencies. Your example of overparameterized linear regression fits perfectly within this framework.

Thus, finding the right constraints to keep the problem well-posed and meaningful becomes challenging. It seems feasible to adapt the algorithm that builds polynomial conservation laws (section 2.5 of Marcotte et al. 2023) to settings with fixed datasets by identifying polynomial conservation laws that are also polynomial in the given dataset parameters using formal calculus.

We appreciate this suggestion and will add a discussion in the conclusion highlighting this as a promising direction for future research.

Q3a (...) conservation laws for Transformers (...)

Indeed, this is a direct consequence of invariances, however, since we could not find any reference clearly stating this phenomenon for transformers, we believe it is worth highlighting. More importantly, our main contribution is, however, the nontrivial proof that there is no other conservation law that only depends on the parameters of such attention blocks. In particular, while finding the reparameterization $\phi(\theta)$ is relatively easy, proving its "minimality" (cf Appendix F) is much more involved. We will reword to insist on these facts.

Q4a Many of the results in Section 2 and 3 seem to be restated results (...)

All results in section 3 are new except the one recalled in section 3.2 “ReLU and linear networks: known results”. In Section 2, novelties are notably Theorem 2.1, Corollary 2.8, Propositions 2.9 and 2.12. While Theorem 2.11 is not new, we provide a new simplified proof. We will clarify this.

Q5a in which ways does their paper generalize the results of these prior works?

The only result that can be seen as a generalization of previous work (all other results are entirely new) is Theorem 3.6, as detailed in Remark 3.7. In particular, the function $f$ used to factorize the neural network $g_{conv}$ defined in equation (13) $g_{conv} = \psi \circ f$ is the same as for the 2-layer ReLU neural network $g$ defined in equation (12): one has $g = \phi \circ f$. Thus, to show Theorem 3.6, you can first directly use the result of Theorem 3.4 as it only depends on $f$ (see Lemma E.2). However, to complete the proof of Theorem 3.6, you also need to compute the generated Lie algebra associated to $\psi$, which is pretty involved and differs from Lie algebra computations associated to $\phi$.