Transformative or Conservative? Conservation laws for ResNets and Transformers
Authors: Sibylle Marcotte
Rémi Gribonval
Gabriel Peyré
Venue: ICML 2025 oral (2025)
Keywords: Conservation laws
gradient flow
linear and relu neural networks
Convolutive ResNet
Transformer
SGD
Primary area: deep_learning->theory

Abstract: While conservation laws in gradient flow training dynamics are well understood for (mostly shallow) ReLU and linear networks, their study remains largely unexplored for more practical architectures. For this, we first show that basic building blocks such as ReLU (or linear) shallow networks, with or without convolution, have easily expressed conservation laws, and no more than the known ones. In the case of a single attention layer, we also completely describe all conservation laws, and we show that residual blocks have the same conservation laws as the same block without a skip connection. We then introduce the notion of conservation laws that depend only on *a subset* of parameters (corresponding e.g. to a pair of consecutive layers, to a residual block, or to an attention layer). We demonstrate that the characterization of such laws can be reduced to the analysis of the corresponding building block in isolation. Finally, we examine how these newly discovered conservation principles, initially established in the continuous gradient flow regime, persist under discrete optimization dynamics, particularly in the context of Stochastic Gradient Descent (SGD).

Transformative or Conservative?
Conservation laws for ResNets and Transformers
Sibylle Marcotte
Rémi Gribonval Gabriel Peyré
Abstract
studying conserved functions quantities that remain in- variant throughout the training process. These conserved functions reveal important geometric properties of the train- ing dynamics and serve a dual purpose. First, they provide valuable insights into the implicit bias induced by both the training algorithm and network architecture, by revealing properties that persist from initialization the final solu- tion (Saxe et al., 2013; Bah et al., 2022; Arora al., 201 Tarmoun et al., 2021; Min et al., 2021). Second, they have emerged as crucial tools in theoretical analyses, playing key role in convergence studies (Du et al., 2018: Arora et al., 2019; Bah et al., 2022; Chizat & Bach, 2020; Ji & Telgarsky, 2019a; Min et al., 2021). Understanding these conservation laws can also be applied to designing new opti- mization schemes, which longer preserve these laws but instead enforce them to reach desired value (e.g., bal- anced condition for ReLU networks) in order to potentially accelerate convergence (Saul, 2023; Stock et al., 2019).
While conservation laws in gradient flow training dynamics are well understood for (mostly shal- low) ReLU and linear networks, their study re- mains largely unexplored for more practical archi- tectures. This paper bridges this gap by deriving and analyzing conservation laws for modern archi- tectures, with focus convolutional ResNets and Transformer networks. For this, we first show that basic building blocks such as ReLU (or lin- ear) shallow networks, with or without convolu- tion, have easily expressed conservation laws, and more than the known ones. In the case of single attention layer, we also completely de- scribe all conservation laws, and we show that residual blocks have the same conservation laws as the same block without skip connection. We then introduce the notion of conservation laws that depend only on subset of parameters (cor- responding e.g. to pair of consecutive layers, to a residual block, or attention layer). We demonstrate that the characterization of such laws can be reduced to the analysis of the correspond- ing building block in isolation. Finally, we ex- amine how these newly discovered conservation principles, initially established the continuous gradient flow regime, persist under discrete opti- mization dynamics, particularly in the context of Stochastic Gradient Descent (SGD).
Conservation laws. In the context of Euclidean gradient flow training dynamics, conservation laws in the form of "balancedness conditions" have been established for ReLU and linear networks (Saxe et al., 2013; Du al., 2018; Arora et al., 2019). Subsequently, (Marcotte et al., 2023) demonstrated the "completeness" of these laws: no addi- tional conservation laws exist for these architectures in the shallow case under Euclidean gradient flows. For these net- work architectures, (Marcotte et al., 2024) unveiled novel conservation laws when considering alternative optimiza- tion algorithms particularly non-Fuclidean gradient flows, as employed in ICNN NMF, or momentum-based dynam- ics, also demonstrating their completeness. Furthermore, (Marcotte et al., 2024) revealed that conservation laws under momentum dynamics exhibit fundamentally different char- acteristics compared to simple gradient flows: these laws are time and velocity-dependent, and are generally fewer in number than in the gradient flow case. For feed-forward networks with single-channel convolutions, conservation laws were also identified under gradient flow dynamics (Du et al., 2018). While the investigation of conservation laws for more sophisticated neural architectures has remained largely unexplored, this paper addresses this gap by extend- ing the analysis to more complex network architectures.
1. Introduction
Understanding the behavior of neural networks during train- ing remains a fundamental challenge in deep learning. A particularly insightful approach to this challenge involves
'ENS-PSL Univ. Inria, CNRS, ENS de Lyon, Univer- sité Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France 'CNRS. Correspondence to: Sibylle Marcotte wsibylle.marootttecnss...
Proceedings of the 42a International Conference on Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Conservation laws for ResNets and Transformers
Residual networks (ResNets) constitute a fundamental
2. Conservation laws for Gradient Flows
class of deep learning architectures that revolutionized the field of computer vision through their groundbreaking per- formance (He et al., 2016). The distinguishing feature of these networks-the incorporation of skip connections- has since become a cornerstone principle in other deep learn- ing architectures, most notably exemplified in Transformer models (Vaswani. 2017). In (Marion et al., 2023), under spe- cific initialization assumptions and incorporating rescaling operation, the authors demonstrate that the solution reached during traing (i.e. the trained neural network) corresponds to discretization of a Neural ODE (Chen et al., 2018), thus revealing an implicit bias. This enables leveraging ODE (ordinary differential equation) theory to analyze the trained network.
In this paper, we consider learning problems where the features are represented as 2, R and the targets or labels are denoted by y E y. In regression tasks, y is typically defined as RR,, while in classification contexts, ys represents categorical labels. In scenarios involving unsupervised or self-supervised learning, yi can be treated as constant. We denote == (Ti, Mi) and Z = (Ii, Wi);
Predictions are generated through a parametric function T (such as a neural network). This func-
g0,, :) : I : - tion is trained by empirically minimizing cost function with respect to the parameters 0 E e EOCR. CR"
1 (z = N Σ, (gg(0, N
(I)
Transformers. Since their introduction (Vaswani, 2017), Transformers and their multi-head attention mechanism (Bahdanau et al., 2015) have achieved unprecedented per- formance across domains from natural language processing (Brown et al., 2020) to computer vision (Dosovitskiy et al., 2021). In (Vasudeva et al., 2024), the authors shows that when training self-attention layers with gradient descent, the key-query matrix naturally converges to a hard-margin SVM solution, revealing an implicit bias similar to that observed in linear logistic regression on separable data (Soudry et al., 2018: Ji & Telgarsky, 2019b).
with & a loss function. In practice, the training dynamics are realized using a gradient descent algorithm with weight decay:
0 0
0g
VLz(((k) - A
To facilitate mathematical analysis, the training dynamics is simplified by considering a gradient flow (GF) with weight decay (WD). This approach represents the continuous coun- terpart of Equation (2) as т approaches zero and can be expressed as the first-order ODE where ((t) 2 0:
0(t) + X(E)((t) = - V-VL((0(t)
(3)
Our contributions. After proving that conservation laws of gradient flows with weight decay are determined by their equivalent without weight decay (Theorem I), we dis- cover new conservation laws and show that these new laws complete for several basic building blocks of modern networks with skip connections: shallow multi-channel con- volution layers (Theorem 3.6), self self-attention layers (Corol- lary 3.9, Corollary 3. 10), cross-entropy classification layer (Proposition 3,11), and plain MLP with skip connections (Proposition 3.2). We then explain how these results can be used analyze deeper networks, via a generic analysis under the lens of the new taylored analysis of conserva- tion laws that only depend on a given subset of parameters (Proposition 4.3), Notable results (Theorem 4.6) include the formal proof that such laws exactly match the laws of the classical blocks considered in isolation. Besides, in the case of residual convolutional networks we show (Theorem 4.7) the absence of conservation law associated consecutive linear layers that "overlap" two residual blocks. To complete the theoretical analysis finally show that the conserva- tion laws of gradient flows are also approximately preserved during actual (discrete) SGD dynamics, under appropriate assumptions (Proposition 5.1), with an error bound (27) scaling as Ostep--size" that we showcase in our numerical experiments, Figure and Figure 5.
weight-
We aim to understand what quantities are preserved during the dynamic (3) for a variety of neural networks g. The mathematical study of what happens in discrete dynamics is done in Section 5, The analogue of the transition from (2) to (3) for a simplified version of Adam algorithm is analyzed in Section 2.3.
2.1. Conservation laws of neural networks
A function hitt, f) is conserved if for each solution @(t) of the ODE (3) with any initialization and any data-set, one has h(t, that remains constant over time.
-eieeddepndencc: with vs without weight-decay. Our first contribution is the following theorem, which clarifies fundamental distinction in the temporal dependence of con- served quantities of a dynamic system, whether it includes WD or not. It also demonstrates how the conservation laws with WD can be readily derived from those from the system without WD. See Appendix for a proof.
Theorem 2.1 (Structure theorem). Let h(t, 8) be a con- served function for the ODE (3). If for every O,
2

---

Conservation laws for ResNets and Transformers
there exists data-set Z such that (z)
0,
Corollary 2.5. Consider loss f(2, y) that satisfies Assump- tion 2.4. Then for all I E A:
then the function (()) == h(0,a) a) satisfies h(t, 0)
H (0 exp(Je {(s)ds).. V1,0. Thus, conserved functions can be expressed with D variables (instead of D + 1). More- over, h(t,0) 0) (0) is conservation law of (3) without WD (i.e. with x(()= 0).
Wa
span {0og (0,xxxx)'' w).
xEVo.wEV
Another useful assumption is the following.
Remark 2.2. In particular, when considering (3) with A(t)= 0, one has (4, 0) = k(0,0) 0) for all I. and 0.
Assumption 2.6 (Local reparameterization). There exists df and & E C² (O,R)) such that: for each parameter Og in the open set θ C IR for each E X such that 0 + g(0, x)) is C² in neighborhood of 0g' there is a neighborhood ft of Oo and f(,,x) E (o(( R²)) such that
Given this established correspondence between conserved functions with and without weight-decay, we can now re- strict our analysis to iine-independent conserved functions h()) in the case of gradient descent without WD:
AA E 2 g f()(0),
(5)
(tt) = V ((((t))).
(4)
Such factorization g(0.z) f(((0), is always pos- sible but never unique: = id and f - 9 yield a trivial factorization, and starting from an arbitrary factorization any diffeomorphism tp yields another one g(0, x) = f()(0), with f(a, x) = f(((a), x) and क = の Φ. The corre- sponding space W can be characterized with any such factorization (Marcotte et al., 2023, Proposition 2.12).
Definition and characterization of conservation laws. Here we recall the main ingredients of the framework for conservation laws from (Marcotte et al., 2023), introducing some formal definitions of intermediate objects and results that hopefully streamline the corresponding reasoning. A function (0) a conservation law for g if for each solution (tt of the ODE (4) with any initialization and any data-set, one has h(((t)) h((())), vt. The formal definition of conservation law in that case is given in (Marcotte et al., 2023, Definition 2.4) (corresponds to the notion of being locally conserved on θ for any dataset), and we recall an orthogonal characterization of a smooth conservation law (Marcotte et al., 2023, Corollary 2.6, Proposition 2.7):
Proposition 2.7. Assume that for every y the loss {(z, y) is C² di-fferrentiable with respect to 2. Under Assumption 2.6, for all E O:
W = do(0) Wil
(6)
6())
with @Φ(0)= E Rdx the Jacobian of 6 and
Proposition 2.3. Assume that for each y E y. the loss ((2, y) is @ddiiferenttable with respect
E R A func- tion h E (O, R) is conservation law for g with respect to the loss l if and only if Vh(0) 1 Wg for all 0 θ where:
Whe = span (0f)( V. 6(0) (x
where f²()) = f(,x).
Wg
span 2=(x,)))((Gx3)
[z)(
Under Assumption 2.6, Proposition 2.7 directly rewrites as:
span (0og(0,x) x) Tz
Corollary 2.8. Consider a loss f(z, v) that satisfies Assump- tion 2.4. Under Assumption 2.6, then for all 0 E θ
(xx))Xoo\)
with Xe the set of data points ze mh such that g(-,x) is C² -i-ifferenntiale in the neighborhood of f.
w). and, denoting Pc(o(0)) E Rdxd the matrix of the projection on this fmite-dimensional vector subspace, we have
WIe ((9)
span (2,0)
(7)
X
Assumption 2.4. Consider loss (zz.), We assume that for every y it is differentiable with respect to 2 E IR". We define for all E I the subspace
[z)2) = span, V (=((,,)) R..
Wg = raage (00(0) P(((()))
(8)
We also assume that Vz(z) does not depend on 2 E R,," so we rewrite V=(z)= V4.
With plain shallow ReLU and linear networks, when V/- R" (this holds with the Euclidean loss or the Kullback- Leibler loss), there happens to be a "distinguished" choice of & (Marcotte et al., 2023, Lemma 2.13) such that the projection matrix Pt is simply the identity, so that all the
In particular, this assumption is satisfied for all classical losses (c.g. mean-square error, Kullblack-eeillrr diver- gence, cross-entropy loss) as stated in (Marcotte et al., 2023, Lemma C.2, Remark C.3), and is known to imply the fol- lowing direct corollary:
'i... a belongs to the set Xo, as defined in Proposition 2.3.
3

---

Conservation laws for ResNets and Transformers
needed information about W is captured in Og. The
that is based solely on the Frobenius Theorem and the char- acterization of Proposition 2.9.
formalism with Pa enhances the flexibility of the framework to encompass the variety of possible factorizations via f and 6.
We now demonstrate why Definition 2.10 effectively elimi- nates all functional redundancies. The following proposition states (see Appendix D for a proof) that any conservation law can be expressed locally as a function of the independent reference conservation laws obtained in Theorem 2.11.
In light of (8) we introduce the vectors fields:
X 0 - P((())) isisd
(9)
Proposition 2.12. Consider a smooth conservation law h G + The of g with respect l, and 0 E O around is locally constant (equal to some
and the linear function space:
which cim((LeeWs::)) integer k). Then, on the neighborhood Uo of 0 given by Theorem 2.11, h can be expressed as function of the D - k smooth independent conservation laws hx+1," ' hD of g given by Theorem 2.11.
vor. = span(xí. .. .xa.
so that we have for any 0 E O, Wgee (). where the trace at 0 E θ of any set W c (0, of vector fields is defined as the linear space
Marcotte et al. (2023) developed a code (detailed in their Section 3.3) that calculates the dimension of the trace of the Lie algebra generated by finite set of vector fields.
(0) == span span x x W] C D
(10)
In particular i E C² is conservation law of g with respect the loss if, and only if, its gradient is orthogonal to (0) for every i and 0. This property is stable by Lie brackets (for self-containeddess see reminder on the underlying calculus in Appendix B), leading to a new orthogonal char- acterization of conservation laws (proved in Appendix B).
2.3. Conservation laws for Adam flows
A simplified version of Adam algorithm (Kingma & Ba, 2014) (full batch, without bias correction and e, Br, .2 0f updates the parameters 0, with estimate of the first and second moments me, using the following equations:
Proposition 2.9. If ((2,)) sutisfies Assumption 2.4 then h E (O, RR) is a conservation law for g with respect to if and only if Vh(0) 1 Lie(span, ( (0) all 0 E O.
my FeLz(0.),
0 = - (VeLz(0)))
(11)
2.2. Existence and "number" of conservation laws
0(41 02 11 mi 02 - risign nisign(VoLz (0,)
yur
Having characterized the conservation laws, we now seek to ascertain the quantity of such laws. However, to com- prehend the number that exists, it is essential to establish a notion of independence that eliminates all functional re- dundancies. We recall the definition of independency from (Marcotte et al., 2023, Definition 2.18):
where the square, square-root, and division are done wlement-wise.
The discrete dynamic (11) corresponds to the explicit Euler scheme of the ODE (informally corresponds to (11) when + 0):
Definition 2.10. A family of functions h;, 15 S is said to be independent if the vectors (O), 1 < i < are linearly independent for all 0 Ω.
i = -sin (VeLz(0)).
(12)
Remark 2.13. Connections between the Adam algorithm and variants of sign gradient descent (referred to as "variance-adapted sign descent") have been established in (Balles & Hennig, 2018) in the discrete dynamic.
The following fundamental theorem (Marcotte et al., 2023, Theorem 3.3) provides a formula for the exact number of independent conservation laws.
Thus one can adapt the results of Section 2.1 by considering the ODE (12) instead of (4). In particular under the same assumptions as in Proposition 2.3, a conservation law h for for the Adam flow (12) with respect to the loss is now characterized by Vh(0) 1 Wg 4O where
Theorem 2.11. If Wy.e S C P)) and
dim(Lie (W)..) is locally constant (equal to some k)) then each 9 E θ C ER" admits neighborhood Uo such that there are D - h smooth (Cx) independent conservation laws ha+1, ... hD of with respect to E on Uo-
wg
span (E X
(sign (VL2(0)) (13)
The proof of this theorem in (Marcotte et al., 2023, Theorem 3.3) relies on Frobenius Theorem (recalled in Theorem C. and necessitates a reasoning by contradiction, along with the use of an intermediate functional space. In this article, we present significantly simplified proof (see Appendix C)
"The continuous time version of the Adam algorithm (Kingma & Ba, 2014) has also been studied, including the hias correction steps, with any e, S1, 82 and any batch size in (Barakat & Bianchi, 2021).
4

---

Conservation laws for ResNets and Transformers
In particular, the associated space W9. is locally constant with respect to 0. Thus, the trace of the generated Lie alge- bra at 0 is directly equal to the one of wys at 0: by using Theorem 2. it suffices to determine the dimension of the trace of 3a. to know the number of independent conser- vation laws. In the case of a 2-layer linear neural network parameterized by é : V) E R x R=xrr + UVT (and similarly for an attention layer), we empirically find that there are no conservation laws, except in the special case = r- 1, where there is exactly one conservation law, given by [0| [V as detailed in Appendix R.
Proof. Use Proposition 2.3 and Corollary 2.5, and notice that since dog = Oog we have W w
Remark 3.3. It is worth noticing that this result does not require the assumption V R".
3.2. ReLU and linear networks: known results
Let us consider U R V R and we denote 0 == (U,V) the parameters and uy U2 the columns of U and V. In that case, the neural network writes:
y = Uo(VTa) x).
(15)
3. The case of shallow neural networks
where o = id (resp. ReLU) in the linear case case (resp. ReLU case). The number of parameters is D = (n + m)r.
Equipped with the general results of the previous section we now provide characterizations of the conservation laws of several elementary networks that serve as building blocks of standard modern network architectures such as ResNets and Transformers. After showing that a basic block has the same conservation laws with or without skip connections, we remind existing laws for shallow ReLU and linear networks, before providing our main contributions of this section characterization of the conservation laws of shallow ReLU networks with convolutions, of attention layers, and of cross- entropy classification layers.
We recall here the result demonstrated in (Marcotte et al., 2023, Lemma 2.13, Theorem 2.14) which shows that, through an appropriate parameterization Φ. the study of Wo. can be reduced to the study of Lie algebra generated by finite number of "well-behaved vector fields.
Theorem 3.4. Under Assumption 2.4, Vf = R", then con- sidering e = RD and ((0) == UV for linear neural net- works, one has: Wote) Rdd and Wg. = range()<<(0)
The same result holds for 2-layer ReLU networks (with or without bias by E R. 1 S k S r) with Q S & the set of all parameters O such that hidden neurons define pairwise distinct hyperplanes H2 := (x E Rd, 0g + by = 0}, and (0) == (ugu)) )a=1
Remark 3.1 (Conservation and invariances). In this section we characterize all conservation laws for elementary net- works (except multihead layers where the completeness of the identified laws is left to future work). It is noteworthy that these laws in most cases are intrinsically connected to network invariances (as detailed in Appendix J.1). When considering weight decay regularization and applying the structure thcorem Theorem 2.1, a particular consequence of our results is that the conservation laws from the non- regularized case (GF without weight decay) vanish at the optinum. This finding not only aligns with partially known "balancedness" properties ((Yang et al., 2022, Theorem 1) applies here in the case of a ReLU activation, as the associ- ated elementary networks go are homogeneous with respect to hidden neuron rescaling, and our results show that it is also true for linear networks and for cross-entropy classifi- cation layers (20)); it also provides additional insight on the dynamics of the convergence balanced parameters.
Indeed, by Theorem 3.4, computing the associated generated Lie algebra and finally applying Theorem 2.11, the authors are able to determine all conservation laws in these settings (Marcotte et al., 2023, Corollary 4.4):
Proposition 3.5. With the assumptions of Theorem 3.4, in the ReLU (resp. linear) case, we have: for any 0 θ (resp. @ O such that (%: has full rank), there is neighborhood of 0 in which all conservation laws for (15) are functions of |ux 1|2 - || 2 (resp. of (wi, un) (vx, vi)), S k, S
3.3. ReLU neural networks with convolutions
We now consider the case of a basic block of a one-hidden layer ReLU neural network with convolutions. This means that the input vector x E I is considered as the concatena- tion of channels (x(8) E TR., 1 S i S co each with p pixels (for images), or p samples (in case of time series), so that m = cop. Accordingly the output of the network is given by (15) where the matrices V and U are made of circulant blocks respectively corresponding to convolutions with fil- ters Uj.i, Uk.j. 1 S i < co, < j S Ci (c1 is the number of channels of the hidden layer), and 1 < k < C2 (C2 is the number of channels of the output y = g(0, x) E ER" considered as the concatenation of channels (k)) E R=1, that n = 02 X n1-
3.1. With ys without residual connections
We first establish a simple but generic result for conservation laws in the presence of skip connexions. Given any neural network g(0,) -) with = m, we consider the residual neural
network g(e, -) defined by:
g(0,-) x E Rd + x + g0, ER"
(14)
Propusition 3.2. With respect to any loss satisfying Assump- tion 2.4 g and i have the same conservation laws.
5

---

Conservation laws for ResNets and Transformers
More explicitly this corresponds to
We define <(0) (Φ1, 02) with oi Q K and 2 VTO the reparametrization such that (up to flat-
C1 U
02 UJ x(1)) Σ
Co
tening of matrices into vectors) g(0,z) f( x) (
g(0,z)= =
(16)
k-1
The following theorem (see Ap- pendix F for a proof) demonstrates that the parameterization 4 is, in a some sense, sufficiently rich and allows for reduc- tion to the study of a Lie algebra generated by the vector fields (0 + 06(())ex)))
Assuming that the filters all satisfy Weyd E RR (resp. V4 E I=")) the network parameters 0 = (iluk.))k.. (15,))ss), are of dimension D = C2 + Ci Co = ei (C2thu + com)).
Theorem 3.8. Under Assumption 2.4, if Ve R and N 5 2 then W Rd, and Wg range range(3<(()) }, v0 E Oatt- @(0)
We define Ceany as the set of all such that the matrix V has all its columns that define pairwise distinct hyperplanes.
A conservation law for this network has already been es- tablished in (Du et al., 2018, Theorem 2.3) for a single- channel networks (Cu = C1 = 2g = 1). We find that similar functions are preserved in the multi-channel case, and we demonstrate that there are no others by characterizing the trace of the associated Lie algebra well as its dimension. The following theorem presents all conservation laws of the network (16). A more general version of this result is proved in Appendix E, which allows for instance to consider strided convolutions (Zhang, 2019) (i.e. with zeros inserted in the filter) in order to define translation invariant CNNs.
Thanks to this theorem, the analysis boils down a similar problem as for Proposition 3.5 and allows us to determine all conservation laws. See Appendix G for proof.
Corollary 3.9. Under the assumptions of Theorem 3.8 for each 0 E Oat such that both horizontally concatenated matrices (Q, K) and (V. OV have full rank, there is a neigh- borhood of in which all conservation laws for (18) are functions of QQ KKT, and VVT- and vice- versa.
Theorem 3.6. Under Assumption 2.4, V = R", then in the neighborhood of each E θ Oconv there are exactly C₁ independent conservation laws for (16) given by
For more than one head, the neural network writes
(00,
Σ softmax(XQ) Knx KnX
)xVV XXV, On, (19)
6o
h0) | 2
|, Σ
isisa. (17)
up matrix flattening, with Qh. Kr, Vn, On E I タ dim In the case of multi-head attention, we can partially ex- tend our results to obtain set of conserved quantities (see Corollary 3.10, with proof detailed in Appendix G). How- ever, determining whether this set of conservation laws is complete remains an open problem.
Remark 3.7. The formulation of the neural network (16) in the multi-channel convolutive case generalizes the one (15) of simple 2-layer ReLU network without convolution. Specifically, setting p- = 1 and identifying co and n, with m, C1 with r, and 02 and n., with n, yields (15) and the conservation laws obtained in Proposition 3.5 coincide with the ones given in Theorem 3.6.
.., the functions Corollary 3.10. For any h = 1, 1,..., H.
QT - KuKT, VAVV - -O,O%
define conservation laws for (19) with respect to any loss l.
3.4. One attention layer
For an attention layer, the input X E TR. Nxdim is reshaped as the concatenation x E RT of N tokens (66 e didim with m = Ndim. The layer output is
3.5. Cross-Entropy Classification Layer
For classification tasks (which we consider in the numerical part Section 5.2), the final layer typically combines linear transformation with the evaluation of a cross-entropy. This corresponds to using a Kullback-Leibler loss over n classes Ely. == Σn= Yi log (v./v() ys + 2/. together with a soft-max layer -00,,) where I
g(0, x) softmax(X"' KX XXT)XX RN xdim (18)
Xxdim
(reshaped row by row as a n-dimensional vector with N dim fit our generic notations), and where:
e
exp(AA (A:)
g(0, x) s0ft softmax(2)) == Σ,e 2,
(20)
softmax(A), EN exp( (xx)
Note that this loss satisfies Assumption 2.4 and V2- = R". The following proposition (proved in Appendix H) shows that the softmax layer induces a new set of conservation laws with respect to this loss.
and with Q, K. E R4ixdim We assume that all the columns of T are non equal to zero. We consider Oatt the set of all parameters that satisfy this condition.
6

---

Conservation laws for ResNets and Transformers
Proposition 3.11. With respect to any loss E such that Ve= =
Considering 0 @ the parameters of a global network go. we focus on 0r subset of the parameter entries (typically we will consider 0r Or for some I, but other scenarios will also be covered), so we can write 0 = (0r, frr), where 0r gathers the remaining entries of the parameters 0.
R", there exoctly m independent conservation laws for the classification layer given by (20): he == E, 0zgij=
1. m.....m.
4. Deeper ResNets and Transformers
The following proposition (see Appendix 1) characterizes smooth conservation laws of g that only depend on 0r.
In this section, we examine conservation laws for deep net- works g(f,-) ) (denoted here as go go())) composed of q residual blocks. Specifically, each block I consists of parameters di such that 0 = (0,,..., .g) and corresponds to an elemen- tary network denoted dz. (z). Thus, the global network can be written as the composition of q elementary networks:
Proposition 4.3. Consider function h E ( RR) that only depends on the coordinates Or, and for each 0 e e denote
(0r, (0r,,))E0) E
(24)
Consider loss that satisfies the assumptions of Proposi- tion 2.3 as well as Assumption 2.4. The function is a conservation law of g with respect to if. and only if. for every E O one has Vo, h(6) 1 Ro,0...) where:
go : x E R
08, 9o, (x) E R...
In case of a classification task using cross-entropy loss, we consider a last block constitued of a linear (fully-comnected) layer and a softmax activation 9 as in (20), so that 0 ==
Rey (W19-)) ==
span 10,, (00r9( x)) Tww..
span
ge@ro (0r) To WEVi
(01, 04+1) and the global network then writes:
(25)
go x0 R
g+1 g0
9d,
gō, (x) R".
For concrete examples below, a technical challenge that arises when studying conservation laws that depend solely on Pr, and in comparing them with those of "internal" shal- low networks involving only 0r, is the analysis of the set Xa of input of the global network go around which it is smooth enough (cf the definition of X₂ in Proposition 2.3), rather than the set of inputs of the considered "internal" shallow network. Overall, the important property for our purposes is the density of this set, proved in Appendix K.
Example 4.1 (Convolutive ResNet). A convolutive ResNet corresponds to go : x E IR" + (x) E IR with residual blocks (we recall that here m = cop). Each block / < q has parameters g: v/,)) and consists of plain 2-layer convolutive ReLU network (16) with skip connection:
by + x U (V'z),
(21)
with matrices V and Lt made of circulant blocks respec- tively corresponding to convolutions with filters e', " Example 4.2 (Transformer). We consider a transformer architecture where go x E Thm + (x) E T2" m = N x dim) consists of alternating residual blocks with either an attention layer with one head, or 2-layer MLP. We notably omit normalization layers, and restrict to single head. In that case, up to appropriate harmless matricizations or flattening operations associating the vector to the token matrix X, each block 1 f q writes either:
Lemma 4.4. Denote ② 0 X ... O1 (or θ @q+1 x O, x.. O₁ with 0g+1 Raxm when there is last softmax layer) where for each layer S L g, O₁ is the set of parameters Or such that
1. gb, is an open map";
all the rows of the matrix V (resp. from (21) (resp- (22)) are nonzero in the convolutive ResNet case (resp. in the Transformer case).
gb. x + v T + (U)
(22)
For every 0 E ① we have X₂- = R."
where O (U;, Uf) corresponds to the two weight matrices of 2-layer MLP (15), or:
Remark 4.5. Obviously Item 2 only excludes a lower- dimensional set of parameters. We discuss in Appendix L why Item is also a generic condition on the parameters for the example of residual block associated with 2-layer ReLU network.
go, (x) = vec + softmax XI
KiX XXVV)XV O.J
(23)
where 0 = (Qi, K2, Vi, Ot) corresponds to the parameters of a single-head attention layer (18).
4.2. Block laws for natural residual blocks
4.1. Characterization of "block" conservation laws
Consider 0r = Ar where I E (1,... .q). The following theorem (See Appendix N for proof) demonstrates that
To analyze conservation laws in this context we focus on laws that depend only on a subset (or block) of parameters.
fi.e. it sends an open set to an open set
7

---

Conservation laws for ResNets and Transformers
the conservation laws of the global network g that depend exclusively on 0 are precisely those of the shallow network ob. of the (-th residual block.
O1+1 E @4+1, we further require that the rows of V are pairwise non-colinear. If = p then any conservation law of g with respect to the Euclidean loss that only depends on Bg is a constant function.
Theorem 4.6. With O as in Lemma 4.4, consider the I- th residual block of Example 4.1 (resp. Example 4.2), and denote Or- == 0t and Or the parameters of all other residual blocks. A function H : 0 = (Or, 0r.) E O + h(OT) that only depends on Q is a conservation law of with respect to a loss such that V2= = IR" if and only if h is a conservation law of the shallow residual network g' (0), == ge, (x) with respect to the Euclidean loss. The same result holds for 0r == 00+1 when considering last block (20).
4.4. Numerical confirmation
In the general case of residual network composed of q blocks, there could potentially exist conservation laws that depend on a larger subset of parameters than those previously analyzed, which helped us demonstrate that we recover exactly the same conservation laws as those of elementary blocks when considering the global net- work. However, by numerically computing the dimension of span Lz()) Z} C Wg with a sufficiently large batch size to adequately explore the space, we find that there is no additional conservation law when m > 1 for a ResNet with = 2 residual blocks (see code in our GitHub repository).
Thus by using Proposition 3.2, the conservation laws of g that only depends on O% are exactly the ones of (16) (resp. (18)), which are described in Theorem 3.6 (resp. Corol- lary 3.9) for a residual block defined with 2-layer ReLU networks (resp. with an attention layer).
When m = 1, numerical results suggest that there are addi- tional conservation laws. It might be that this specific case enables as in (Marcotte et al., 2024, Theorem 4.1) to shed the light on new invariances that give rise to new conserva- tion laws. This is confirmed by the following example: we exhibit a domain Ω where there are more conservation laws than the "block" ones.
4.3. Case of blocks overlapping a residual connection
This section focuses exclusively on the case of ResNet (with or without) convolutions as defined in Example 4.1. In the previous section Section 4.2, we examined the conserva- tion laws of the global network g that depend only on the parameters 0: of the 1-th block, and we just show that they exactly correspond to the ones of the shallow network o6, corresponding to the (-th block. However, it is also possible to recast the global network as composition of elementary networks that maintain parameter separation while involy- ing a subset of parameters located before and after a residual connection as described in Figure in Appendix O.
Example 4.8. Consider a ReLU neural network with two- residual blocks, g((u, v, 8, f) = x + (vz) + sa(t(x + uo (ez)) with (u, 8, E f C R and x E R. While there are two "block" conservation laws: u² =²- and s² -t?, there are three conservation laws on the set f2 of all parameters such that sgn(t) = sgn(v) = =gn(u). Indeed, for every x,
Specifically, let us consider / E (1..... . 1} and 07- (ef) 1+1 ")) In that case, 0r corresponds to two consec- utive parameter blocks that overlap a skip connection (as described in Figure 2 in Appendix O). Denoting 21 (resp. ) the input of the intermediate layer of the I-th (resp. 1-th) block, x2 (resp. 12) the copy of the input of the cor- responding block that is transferred via the skip connection, g can be written as a composition of 91 92 and Is with
1. either wx, to I 0, hence g(0,x) and Veg(.) = 0;
2. or ez, fo 2 0, hence g(0, = x + uex stx +txx+ttet stuva as wx, tz, tu 2 O, and thus 1 Vog(-) =: xx(-) is vector field that does not depend x.
x.
As the space W = Rx1 is spanned by a single non-null vector field, its Lie algebra is itself, and by Theorem 2. 1, there are exactly 4 3 conservation laws as claimed.
(%% = 92
02
x1
)) = x22
v1+1 ) (v' 1a) (a) Ia
(m) = go(0g, x) == yd
(ob__ 1-1 .... ......g))))))
5. Conservation laws for discrete dynamics
and gi (0., (1))) = 90. o...
..... o 1 (( ((u--1 12) (d)) (%)
In practice optimization is performed with a discrete dynam- ies associated to stochastic gradient descent. To what extent do conservation laws still apply in this context? This is the object of this section.
where 0, 0g gather all parameters appearing in and ga- The following proposition (See Appendix P for a proof) shows that there exist no conserved functions of the global network that depend exclusively on 02- 0r (V++ .
5.1. (Stochastic) gradient descent as a training dynamic
Theorem 4.7. Consider a layer index 1 S S 9 - and + defined as in Lemma 4.4 with the exception that for each
We consider the ERM problem (1). Instead of using gradient descent (2), we consider the stochastic gradient descent
8

---

Conservation laws for ResNets and Transformers
(SGD) method on mini-batches:
00000
0k+11
O2 Th VL (0k),
04.23
where the sequence of mini-batches (Zx)s is drawn i.i.d. from the data distribution. The following proposition (see Appendix Q for proof) shows that conservation law is approximately preserved by SGD.
0023
Proposition S.1. Let h()) be conservation law of the gra- dient flow with a bounded Hessian vB, [[0 S Ch. Suppose that the gradients remain bounded in expectation throughout the algorithm:
0021
Stops
Ez,,0 VL (0x)ll2 S CL.
Figure I. Tracking conserved function during ResNet- training (26)
on CIFAR-10.
C1CL
k-1
Then, we have E [1((0))- - n(0b)) 5
T
(27)
2
lines show the theoretical slopes Cr² derived from (27), confirming that the function is approximately conserved and that the slope coefficient maintains proportionality with 7² Our code is available at our GitHub repository.
i=0
For a constant step size Tk = т, this proposition shows that the conservation error grows as (Ox) - l(o))l = k))
which increases linearly with the number of iterations. If
In another experiment, we train transformer model on the IMDb sentiment analysis dataset (Maas et al., 2011) 1) using SGD optimization. We track the evolution of the Frobenius norm of the conserved matrix identified in Corollary 3.10, specifically examining the query and key matrices from the first attention head in the first layer. Consistent with our ResNet training results presented in Figure 1, we observe that the conservation error scales as O(step-size") through- out training, which confirms our theoretical bound (27). Furthermore, it is worth noting that the numerical behavior is unchanged whether masking is applied or not. See Ap- pendix for the associated figures. Our code is available our GitHub repository
one uses decaying step size, such as T /o/(k: + 1). which ensures convergence of the method, then the error re- mains bounded: [k(0)) (((x))((OO)) = O The requirement that h has a bounded Hessian holds in the case considered in this paper since the conservation laws we examine are quadratic. The key assumption for the result to hold is the bound on the gradient magnitude in (26). This condition met if the loss function i and the network ge are uni- formly Lipsehitz, though this is not generally true for deep networks. It also holds with explicit constant Cz, for smooth convex losses (Bach, 2024), but this setting is quite restrictive. More generally, such bounds hold for smooth losses when the variance of Lz, (f) is bounded (Garrigos Gower, 2023), because the iterates are bounded in ex- pectations, [[ < +00, though the constant C2 may not be explicitly determined. In the numerical experiments presented in Section 5.2, we empirically evaluate the con- stants to show that, in practice, they remain relatively small, ensuring approximate conservation.
Conclusion
This paper investigates conservation laws in deep networks (ResNet and Transformer architectures) within gradient flow dynamics and examines their behavior under discrete SGD dynamics. Our analysis does not currently account for trans- former normalization layers, and the multi-head attention mechanism is only partially addressed. The integration of these components as well as max pooolin layers presents promising avenues for future research.
5.2. Numerical experiments
In Figure 1, we train a ResNet-18 on CIFAR- 10 (Krizhevsky, 2009) while tracking the difference between the squared Frobenius norms of consecutive convolutional layers in the first residual block, considering the conservation law h(Or)= = ΣØ= hy (Br), where hy is defined in (17), with Or representing the parameters of the first residual block. We vary the learning rate between 10 and 5 10-3, using stochastic gradient descent (SGD) without momentum or weight decay. For each learning rate, we train 10 models for 50 steps with 10 different random seeds, recording both the loss evolution (bottom) and the evolution of the con- servation error (hx0)) h((fo)) (0()) (top). The dotted
A Acknowledgements
The work of G. Peyré was supported by the European Research Council (ERC project WOLF) and the French government under management of Agence Nationale de la Recherche as part of the "Investissements d' d'avenir" pro- gram, reference ANR-19-P3IA-0001 (PRAIRIE 31A Insti- tute). The work of R. Gribonval was partially supported by the AllegroAssai ANR project ANR- MRR-19-CHH 19-CHIA-0000 and the

---

Conservation laws for ResNets and Transformers
SHARP ANR Project ANR-23-PE1AA-0088 of the PEPR IA, funded in the framework of the France 2030 program.
Chizat, L. and Bach, Implicit bias of gradient descent for wide two-layer neural networks trained with the lo- gistic loss. In Conf. Learning Theory, pp. 1305-1338. PMLR, 2020.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenbor, D., Zhai, X., Unterthiner, T.., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
References
Arora, S., Cohen, N., and Hazan, E. On the optimization of deep networks: Implicit acceleration by overparameteri- zation. In International Conference on Machine Learning, pp. 244-253. PMLR, 2018.
Du, S. S., Hu, W., and Lee, J. D. Algorithmic regulariza- tion in learning deep homogeneous models: Layers are automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
Arora, S., Cohen, N., Golowich, N., and Hu, W. A conver- gence analysis of gradient descent for deep linear neural networks. In International Conference Learning Rep- resentations, 2019.
Garrigos, G. and Gower, R. M. Handbook of conver- gence theorems for (stochastic) gradient methods. arXiv preprint 2rXiv: 11235, 2023.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Bach, F. Learning theory from first principles. MIT press, 2024.
Bah, B., Rauhut, H., Terstiege, U., and Westdickenberg, M. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Information and Inference: Journal of the IMA, 11(1): 307-353, 2022.
Isidori, A. Nonlinear system control. New York: Springer Verlag, 61:225 236, 1995.
Ji, Z. and Telgarsky, M. Gradient descent aligns the layers deep linear networks. In International Conference on Learning Representations, 2019a.
Bahdanau, D.. Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations (ICLR), 2015.
Ji, Z. and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. Conference on learning theory, pp. 1772-1798. PMLR, 2019b.
Balles, L.. and Hennig, P. Dissecting adam: The sign, magni- tude and variance of stochastic gradients. International Conference on Machine Learning, pp. 404 413. PMLR, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14411690,, 2014.
Krizhevsky, A. Learning multiple layers of features from tiny images. 2009.
Barakat, A. and Bianchi, P. Convergence and dynamical behavior of the adam algorithm for nonconvex stochastic optimization. SIAM Journal on Optimization, 31(1):244- 274, 2021.
Maas, A. L., Daly, R. E., Pham, P. T., Huang, D., Ng. A. Y., and Potts, C. Learning word vectors for sentiment analy- sis. In Lin, D., Matsumoto, Y., and Mihalcea, R. (eds.), Proceedings of the 49th Annual Meeting of the Associ- ation for Computational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL https: /aa//claattttt ogy.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Marcotte, S., Gribonval, R., and Peyré, G. Abide by the law and follow the flow: Conservation laws for gradient flows. Advances in neural information processing systems, 36, 2023.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
10

---

Conservation laws for ResNets and Transformers
Marcotte, S., Gribonval, R., and Peyré, G. Keep the momen- tum: Conservation laws beyond euclidean gradient flows. In 41st International Conference on Machine Learning, 2024.
Zhang, Making convolutional networks shift-invariant again. In International conference on machine learning, pp. 7324-7334. PMLR, 2019.
Marion, P., Wu, Y.-H., Sander, M. E., and Biau, G. Implicit regularization of deep residual networks towards neural odes. arXiv preprint WiXiv2239.0012... 2023.
Min, H., Tarmoun, S., Vidal, R., and Mallada, E. On the plicit role of initialization on the convergence and implicit bias of overparametrized linear networks. In Interna- tional Conference on Machine Learning, pp. 7760-7768. PMLR, 2021.
Newns, F. Functional dependence. The Americun Mathe- matical Monthly, 74(8):911-920, 1967. ISSN 00029890, 19300972.
Saul, L. K. Weight-balancing fixes and flows for deep learning. Transactions on Machine Learning Research, 2023.
Saxe, A. M., McClelland, J. L., and Ganguli, S. Exact solutions the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 13/2.6120, 2013.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on sepa- rable data. The Journal of Machine Learning Research, 19(1):2822-288 2018.
Stock, P. Graham, B., Gribonval, R., and Jégou, H. Equi-nommalization of neural networks. In In- ternational Conference on Learning Representations, 2019. URL
id-r1gEqiC9FX.
Tarmoun, S., Franca, G., Haeffele, B. D., and Vidal, R. Understanding the dynamics of gradient flow overpa- rameterized linear models. In International Conference Machine Learning, pp. 10153 10161. PMLR, 2021.
Vasudeva, B., Deora, P., and Thrampoulidis, C. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738, 2024.
Vaswani, Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Yang, L., Zhang, J., Shenouda, J., Papailiopoulos, D.,, Lee, K., and Nowak, R. D. A better way to decay: Prox- imal gradient training algorithms for neural nets. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022. 2022. URL https: openreview. net/ forum? i 4y=411 xh
xh8 xhCchhhh....

---

Conservation laws for ResNets and Transformers
A. Proof of Theorem 2.1 (Structure theorem)
To prove Theorem 2.1, it will be essential to juggle with certain properties of conserved quantities when they depend on (t. Here we recall definitions and properties directly derived from (Marcotte et al., 2024), which are necessary to understand the proof of the structure theorem.
A.1. Definitions of conservation
We recall definitions and properties from (Marcotte al., 2024).
Definition A.1 (Conservation through a (family of) flow(s)), Consider an open subset Ω S R x RD and a function F C² (9, P). By the Cauchy-Lipschitz theorem, for each initial condition init (tinit: Qmitt) 52 there exists a unique maximal solution E (fiii Thin, linit Thai + (t,, init) t) of the ODE ((t) = F(t,0(t), with ((tmaa) 0 A function /: Q S IRD+1 - IR is conserved on Qi through the flow F if h(t, O(t, init)] = -o(fman mmi)) for each choice of init E Ω and every / E (linit V/init; liiit + Thmi). It is conserved on Q through a family of flows if h is conserved on Ω2 through all flows.
Proposition A.2 (Smooth functions conserved through a family of flows). Let I be any set. Given F, E C' C), R) for any i E I, function h E C (22 IR) is conserved through the family of flows induced by all F; if and only if (Vh(a), (1, )T) foror all a E Ω and for all i E I. Moreover by denoting Xi the vector field on St defined by
X x: DH
(F((a)) E ER"*.
(28)
and by denoting
W = span,(xi) C C² (St, Pt1).
(29)
this exactly means that for all a E Ω. (Vh(a), W(a) = O. where the trace of W is defined by
W(a) = span span(x(a) x E W) = span, (x:(a))
(30)
In the context of the dynamic (3), we consider a = (t, E R x θ C P++1 and we consider the flow Fz defined by
Ot = - M
X(1)0.
(31)
We keep in that section the matrix M()) (that allows to deal with non-euclidean metric as in (Marcotte et al., 2024)). However, we do not focus in that paper on the non-euclidean setting.
Definition A.3 (Conservation during the flow (3) with given dataset). Consider an open subset Ω = Ω2 x 2g x and a dataset 7 = (I;; yi such that L.z E CZ2,, IR). A function h C DR- + R is conserved on Ωt during the flow (3) if is conserved through the flow induced by F' defined in (31).
To identify conserved functions that do not depend on a specific dataset, we focus on a more precise class of conserved functions: those that remain conserved during all flows defined by the ordinary differential equation (ODE) (3). This leads us the following definition. The goal is to derive universal laws that hold true regardless of any particular dataset.
Definition A.4 (Conservation during the flow (3) with "any" dataset). Consider an open subset Q = 24 Qa C. R xxxC O and loss {(2, y) such that ((, y) is C2 -ddiffferentialee for all y Y. A function h Ω C P+1 - IR is conserved on Ω for any data set if, for each data set (X, Y) such that g(,, 22) E (? R²) for each i, the function h is conserved on L during the flow (3). This leads us to introduce the family of vector fields:
Wa = fx(+) 32, Vi o(-,x)) E (2gg R"), x() = (1, (L,F2()))) (
(32)
so that being conserved on Ω for any dataset is the same as being conserved on S2 through Wn.
The definitions provided above are local and contingent upon the selection of an open set of parameters Qg C O. However, our primary interest lies in functions defined across the entire parameter space O thus, we present the following definition.
Definition A.5 (Conservation law for neural network). function h : R x + I is a conservation law if for each open subset C R x O, h is conserved on Ω for any data set.
12

---

Conservation laws for ResNets and Transformers
Therefore, by using the "orthogonality" characterization of conserved function Proposition A.2 and Definition A.5, the object of interest study locally conserved functions is the union of the traces
W = U W (a)) C Rx θ with Ω neighborhood of a }.
(33)
Corollary A.6. A function h 18 θ E a conservation law if and only if V 1 W W. all a R x O
Proposition A.7. Assume that for each y E y the loss E(z,y) C C--differentiable with respect to RR.. For each
(1, R θ we have:
W
span (=(
fLl (10-000002/2)))))00000000
where Ng is the set of data points x such that x)) Ci-differentiable in the neighborhood of 0.
A.2. Proof of Theorem 2.1
Theorem 2.1 (Structure theorem). Let h(t,0) be a conserved function for the ODE (3). for every 0, there exists data-set Z such that (z) = O, then the function (a) = la(0, a) satisfies h(t, = H (0 exp(Je {(s)ds)), V1, 0. Thus, conserved functions be expressed with D variables (instead of D 1). Moreover, h(t,, 0) = H(0) is conservation law of (3) without WD (i.e. with X(t)= 0).
To manipulate the first assumption outline it as follows:
Assumption A.8. For all E O, there exists Z// such that (O) O, where we denote Z/J the collection of all data set Z ( yt); such that for all i, 9(-, 22) is C2 diffffeerente in the neighborhood of
Assumption A.8 is satisfied for all classical losses e.g. for the mean-squared the cross-entropy loss as shown in Lemma D.2 of (Marcotte et al., 2024).
Proof. Let htt, 9) be conservation law. Then by Corollary A.6, for each (t, R x θ we have: (Vh(a), Wg) 0. In particular, thanks to Proposition A.7 and Assumption A.8, one has (Vh(a), (1,-(()))) = 0. By Proposition A.2, is in particular conserved through the flow induced by F(t, 0) = -X(t00. For each initialization (0, 0o), the solution of the ODE ((t) Ftt, ((t)) is (tt) &n exp(- fo A(s)ds) for t R and by definition of conserved function, one has ht,,, 0())) = 1(0, 0(0)) 1x0, ((t) exp(_e X(s)ds) = (0(()) exp(Se A())ds)) Thus for all and 0, one has htt,, 0) cxp([ X(s)ds)).
Now, let show that (1,0) (0) is conservation law of the GF without WD scenario (A(I) 0). For any () R x θ, one has
[ahh(a)) (Voh(a))
( (O exp(Se ((t) ((s)ds)) (Neespth exp(Jr X(s)d:)VH((Oxxp(. A(s)ds) )
Vh(a)=
Let consider (t, R O. By taking Z E Z/ (NB: here we do not choose it such that LLz(0) 0) and using the characterization (Corollary A.6) of conservation law and Proposition A.7, one has:
0 = (ViCo), -
= -omp
capp / (())ds (Oxcp /' ((s)ds)))+
(VH(eexp() / ((s)ds)), -MO)) -
= -xexp( /" ((o)ds) ((o)ds)), M()))
13

---

Conservation laws for ResNets and Transformers
In partieular, this remains true for 0, and thus has:
= 0.
Finally by denoting 11 (t, + (0), one has
(H (-M(0)Vi(00)
and thus H is a conservation law for the GF without WD case by using Corollary A.6 and Proposition A.7.
B. Proof of Proposition 2.9 (Characterization of conservation laws via Lie algebras)
First, we recall that the Lie bracket [x,x2] of two vector fields X1, E C (Q,, is the vector field defined by:
xxxXel: E Ω [X1. (2) 0x1 -xxx))x0)0xx
(34)
Let us consider W C (9,, The generated Lie algebra W is the smallest space (S,, PR that contains W and that is stable by Lie brackets, i.e. for any X1, X2 A one has [1, A. We denote == Lie(W).
Proposition 2.9 is direct consequence of this lemma:
Lemma B.1. Let h be real-valued function defined on Zt and let consider smooth vector fields x1 and X2 defined Ω satisfying for all 0 E Qt:
(Vh(0),xi = and (Vh(0),x2(0) 0,
(35)
then
= 0.
(36)
Proof. Let us denote the Lie bracket of X1 and X2. By definition
x()) = xx (0) (z))) -
(37)
Differentiating (35) gives as (a)) is adjo-nd Using this relation (38), we have
xx (0) Vh(0) (O) = - (() x.(0) = (0
(38)
(VI(0). x(0)) = (0x: VA(0), (Vx - (0xx vh(0),x1(0) (00 (0(V))()) x2(0). 0.
C. new and simplified proof of Theorem 2.11 (Marcotte al., 2023, Theorem 3.3)
We first recall Frobenius theorem (see for example (Marcotte et al., 2023, Theorem E.1) with the same notations or see (Isidori, 1995, Section 1.4) for control theory practitioners).
Theorem C.1 (Frobenius theorem). Consider W C C (N,P") and assume that the dimension of its trace W is constant: dim (0) for every 0 E Q R. Then the two following assertions are equivalent:
1. each Ω admits neighborhood Uo such that there exists D k smooth (C real-valued functions hg41, hh Ug such that for all 0' E Uo:
span (he-1(0), =
(39)
2. the following property holds:
(2) X1,X2 W.
(40)
14

---

Conservation laws for ResNets and Transformers
By using Proposition 2.9 and Theorem C.I, derive the following corollary, which is already established in (Marcotte al., 2023), Theorem 3.3. Nevertheless, this new proof, by using Proposition 2.9, is more intuitive and straightforward. k is locally constant then each
Corollary C.2. Let assume that Wat S C (θ, P). If (immLie(Wss))00 A E O S RD admits neighborhood Ug such that there are D le smooth conservation laws ha+1,-' , of g on Uo and such that for all 0' E Uo, the vectors (Vha+1(0"),
Vh (0))) are linearly independent.
Thus by using Definition 2.10, Corollary C.2 can be reformulated as Corollary C.3. Let us assume that Wa. S C (o, ). If cli(Lie(Ws::)00) = is locally constant then each O E θ _ IR admits a neighborhood Ug such that there are smooth independent conservation laws ha+1,-' ,hD of g Uo-
D. Proof of Proposition 2.12 (role of independent conservation laws)
Proposition 2.12. Consider a smooth conservation law h θ + R of with respect to f, and 8 E O around which dim(Lie(Wa: is locally constant (equal to some integer k). Then, on the neighborhood Ug of given by Theorem 2. 1. can be expressed as a function of the D k smooth independent conservation laws het1,'" ,, ,bp of g given by Theorem 2.11.
Proof. We consider some 0 O such that dim (im(Lie(Ws::)(0)) is locally constant. By using Theorem 2.11, there exist h hD smooth conservation laws of g neighborhood Un. and thus by using Proposition 2.9 for any /' E Un:
(0'), .... (hn = Lie(w)>')
(41)
But as h is also a conservation law of 9. by using again Proposition 2.9, one has that for any 0' E Ug: Vh(0') e (f)) This exactly means that Vh is linearly dependent of all gradients Vh, everywhere on Ug. If all conservation laws are (C smooth on Uo. this corresponds to the fact (cf (Newns, 1967, Theorem 1)) that is then function of the h, on Ug.
span 1ha++(0'):
E. Proof of Theorem 3.6 (Conservation laws for convolutive two-layer ReLU networks)
We first recall all notations. Let be the number of channels for the input. Let cı be the number of channels for the hidden layer, and let 22 be the number of channels for the final output. Let n., (resp. n.)) be the size of the filters for the second (resp. first) convolution, and let p be the number of pixels of the input image. In that setting, we consider @ = ((ux.))k.s i (Ug,))..), so that D = (n + n,)0)) x ct, and the parameter vector consists of the collection of all filters. We write the input and the output:
(1)
y/
E R".
with co X
P
y
R", with C2 n2.
(co)
(02)
We denote G(O, x)) the convolutive 2-layer ReLU neural network defined by:
)
G(O,,) = (u )o
Σ Ca(oy.) 20(1)
(42)
where C1 :RT= H 121PP1 and C2 R" + RPXPI are lincar operators.
In particular, the formulation (42) is more general than (16). We explicit below the corresponding operators whenin the case of circular convolution.
Assumption E.1. We assume that write C1 (u) (Piu, ,Pp, u) (resp. C2(v) (Q1u, (Q,),) where all matrices P. E x (resp. Qi E B2xx:)) are injective.
Link with the two formulations. In case of a "full" circular convolution (i.e. n, = p). P. is the circular shift operator P,u = u. (mod. p). that corresponds to consider P; the i-cyclic shift matrix, and in particular P; is injective. In the case of filters of size n. $ p, then by considering I the canonical injection I R"* + R2 that adds zeros, then P.u = P,I(u), with P. the circular shift operator of the previous case n, p. in particular P: P.T is injective too.
15

---

Conservation laws for ResNets and Transformers
We now explain how (42) can be expressed with a structured 2-layer ReLU neural network (15), Obviously (42) can be rewritten
G 2) - Uo(V"a)
(43)
where U = (0) and V V(0) are defined blockwise as
U [Ci ((
V == (=(((
It will be convenient for further calculations to consider the vertical concatenation of blocks associated to given index & C1- For this reason we write w()) == ("k,i)k= E RR nاب = (Ug,)) E R and
1 S j
C1 Ci(u1s)
Ci(u1))=
R (mme))) so that U(0) = (u(()) ... Ci(u(a)) R(mee))((i))))
(44)
C (Wea.))
and
C2(U3a)
(2(1(0)
E (p so that V(0) - (O2((() ... (2((l()) E ((r0)))(11)),
(45)
C2(U),ou))
It will also be convenient to observe that by construction, using the notations of Assumption E.1 and the definition of the Kronecker product between matrices
(())0) = (P/ul)
pi u()
with P. = Tog 09 P.
(46)
C2(x01)= (0/2/1) 2 2(1i)
with Q Ic, @ Qi
(47)
Thus G(O, x)) g(U,, where g is the 2-layer ReLU network defined in (15). Lemma E.2. Consider G(O, x) as (42), g((U, v),x) from (15), and & the reparametrization of g defined Theorem 3.4 the (unstructured nonconvolutive) ReLU case. Denote (0) == ((T0), where To = vec(( ()) E poonectpcopppee with U (0) and V (0) defined blockwise as above. Denote @ the set of all parameters such that the columns of V (0) define pairwise disjoint hyperplanes. Under Assumption 2.4, Ve R", then for each 0 O have
W = =paan00 and WG, spant (pan[V ; :i= -.... ,d].
7ER"
Proof. We have G(O.z) = xT = f(d(TT), f(v(0),,) hence the Xe associated to the model G(,, x)) (see within Proposition 2.3) coincides with the set Xro associated to x.) and the general definition (7) yields Wafe (0)
Wifet o(Te)- Moreover, by definition of θ, each parameter 0 E θ is such To satisfies the assumptions of Theorem 3.4. = Whf
By Theorem 3.4 it follows that Wife 11() o(TO) Rd. Finally by using Proposition 2.7, one obtains: W
(0)(()) 7 and ypG.t = span (v:():i= L. .d.. yERd
Proposition E.3. Consider any structured two-layer ReLU network model as in (42) where the linear operators C1 Cz satisfy Assumption E. and & defined as in Lemma E.2. For each 0 E O there exists neighborhood of 0 in which there are exactly e independent conservation laws for (42). Such independent conservation laws are given, e.g. by - 11? ......1..
+ Σ
Before proving this proposition let us highlight that it yields Theorem 3.6 as a direct corollary by the very definition of Oconv and the fact that the matrices P.. Qi associated to the convolutive model indeed satisfy Assumption E.
Proof. Recall that P/ Icg @ P. and Q Ica @ Qi so that P/ (resp. Q' is injective) as P, (resp. Qi) is injective. It will
Pi
be useful denote P, P and P
in order get compact explicit expression of T (step below).
PMp
This will be used characterize the trace of the Lie algebra of WG, and its dimension.
16

---

Conservation laws for ResNets and Transformers
1st step: We first give an explicit expression of
(())
v())
Writing 0
where we recall that for j = 1, ,1, ut) = vec((uk,){k) E R and vا = vec((N)s)))
w
plex))
R we get using once again the definition of the Kronecker product,
(())) V()))
(44) (45) (Le, 2 P)
Te = vec
(48)
2d step: Characterization of WG. and of the trace of Lie algebra.
Recall that the reparameterization & of two-layer ReLU networks from Theorem 3.4 reads Φ: VU (((re2)((1ee))) ((P)))((1))) + (n \==,.-. .pic (I pr with u; v the columns of the matrices U, V. Let A = (Arg)is
(Rinexxmo))po poes and let us consider == vec(A).
SALL
Saz.
Straightforward computations show that 0o(u, () =
vec(U;
where: Sa..
0
SAp
(
Ai)) A ) Since (00) = ((T)) it follows that
(Sal.
0
(v) T 7 TT)TT))) TT
TO
Sam
Using Equation (48) thus obtain
Pt P. C.P.S.s SA.
(0) 7 = blockdiag, (p" (blockdiag(SAa.,. P) blockdiag,
Σ
The above is valid for any A, and we now exhibit particular choices of A to obtain particular constructions, in order to characterize W.. Given some (i.j) (1.... .PI x f1. .1 }, imposing Av,y for all (i',5") # (i.5) yields
0
0000) - 0 SA., P.
0.
(49)
As P. SAn. P.
0
P! AiJQi
E +h. + thwt with n. = n2 and n. = C, we obtain
Q Dij Tp
0
Ge Lemma
spant0 + (0)
7 7:77
A1
0
AT
span
0
0: A; XR. 15isa
(50)
0
Ac₁
0
A
T
0
We now show the converse inclusion. Let us consider some = 1. C1 and some A E Ru.x.. To show the converse inclusion, enough to show that there exists Ai, E Rrieexpos such that P/ A. A. As Pf. is injective by
17

---

Conservation laws for ResNets and Transformers
hypothesis, then P. is surjective. Thus, there exists B = (b: bū. (2 such that P. B A. Then as Q (91 (a.)) is injective, the vectors qj are linearly independent. us define Aig E xpen such that for all k == 1, ne. gk by- In particular, such A. satisfies P. Aij A.
Thus:
A1 0
0
A
span
A +
8:
A, Ruu sisia
C Espan(0 wo.. (51)
0
Aci
A
0
Finally by using (50) and (51) have characterized the space waae as
0
A1 0
0
Af
W.
span
0 +
0: A: E R
1
O
Ae
A
0
To conclude this step we need to compute the trace its Lie algebra
0
A
Ai
Lie = Lie
span
A. E R.uu
,1sia
Arg
0
A
this we can reuse computations from (Marcotte et al., 2023, Proposition H.3) showing that considered Lie algebra indeed space of linear operators that can identified with matrices:
A
0
A
Lie
span
A. E De XB
Acg
AT
0
(In
Thu
Mi
0
Mi
Ang tht S
Me.
where Aäuthy (fu)) is space skew symmetric matrices.
3d step: Let us show that there are exactly 61 independent conservation laws.
Let us consider E O (in particular& + 0). To conclude, by using Theorem 2.11, we only need to compute dim(Lie (()))) (as the computation is unchanged for 6' in neighborhood of 0, we will also get that this dimension is locally constant as needed).
(u(0))
For this, denoting g())
consider the lincar application:
Ta
0 In. )) x
Mi
((1)
P: (M,,.. Moj) E (An,+a,)" + La &
Me.
g(rz)
= Lie (VHG) (8), we only need to compute rank I as in (Marcotte al., 2023, Proposition H.5). Since, by the rank nullity theorem, we have dim ker rank I ei (ño + Re (fe + ñ. 1)/2, it equivalent to compute dim ker also easy to check that dim ker I = dim ker T1 + .... dim ker Te,, where
As range(T) =
P.: M:MM Anutno + (nn
MB0
-aa
18

---

Conservation laws for ResNets and Transformers
By (Marcotte al., 2023, Proposition H.5) (remember that g is vector in that case and a matrix), one has dim ker T, (Nu + Fy ( + P. - 1)/2, so that dim ker C1 (my fy (mu + Mg 1)/2, and thus
rank I = 01 (uu + ftg - 1) = c1
since the number of parameters is indeed D ei (Tt. + n.). By Theorem 2, we conclude that there are exactly en independent conservation laws.
4th step: Finally, ler us check that the claimed laws are indeed a set of independent conservation laws.
ci: . define for j 1,
by ( + Σ ||u
E 10,, || | 112 and hy : - hy h,
By Proposition 2.9 and the expression of the gradient of lig in the terms of Vhtj. the function hy is conservation law for (42) if and only the function hy satisfies for all and for E Anutha:
( (0) 1
Ia..
-la.) MB
(52)
Since hy (60) + TT (0 n,6) 2,00, by (Marcotte et al., 2023, Proposition H.1), this function indeed satisfies (52). Thus hy is conservation law for (42). Finally, it is straightforward to check that the gradient of the functions hj. S S vectors with disjoint supports, hence they are linearly independent. shows that considered conservation laws indeed independent.
F. Proof of Theorem 3.8 for attention layers
Theorem 3.8. Under Assumption 2.4, if V2 RR" and N 2 2 then
w (()
Rdd, and W range (06((())))))) AO Oat.-
Proof. Consider E aat.- By Proposition 2.7 it is sufficient to show that Wift Rd, since this implies that Wg d(0)
rang)( Since the considered model is smooth, here we have X' Rm which can be identified (with the reshaping of x E Rm into X E R.N xdin) R.ddi Recall for convenience that (00) (øi, Φb2) with Φ1 = Q K, 02 V O and
f(o, x) softmax(Xo1x* ) As assume V2= R" the space need to consider is thus
W o(0)
[[0](((0).x)] A}
span XERM AERNxdim
Our goal is to show that V.fe +())
(2dim)) As warmup given any H E Rdimxdim and A E Nxdim, we have:
S H, A) = (x0
(H, x (X (T))
(53)
hence
10p2 fl A = X (softmax(xx1x7]]
(54)
1st step: Considering any j,1 E (1,
,dim), we first show that
0
E
E
span XERN AERV (0) -
(55)
and
EEy}
VERNxdim ACR (0 -A..
span
(56)
Consider any i E (1..... ,N and X = E,g. By (54) we have
10,.fl A =
= =y [softsmxx EI
(57)
Moreover denoting ay ad) (e,, (1ez) have
exp(os, 1
softmax(E Ejsler= = 1.1 11
exp(a,)+N-11
N
Ci
19

---

Conservation laws for ResNets and Transformers
By straightforward calculus obtain that for any K E Rddmxdim
Do.s- K (ej, Ke))(( expla;)(N-1)
e (0) eT
[oxp(a))+N-1)
Thus for any A E R xdim one has:
O K, K,
(le, .) (B)e A) (Bee00)
hence
[0 A eje (3)
(58)
Combining (57) and (58) yields
Ge (8)
[0.1 flT A [0
[0, A
es (
expla;)
1 N) 6i)
expla,)+N-I
Specializing A Ex, for some & E {L. ,N3 we have
1+
exp(as) exp(as)+
N A
N
N
expla,)
i
exp(a,)+N-19
and
-0
if
87 (B = (e(()Te)) ( = (et (
((oTee)) ,e]
ifk=i
hence
when +
N
eje,
Oof
==
(59)
Ciese] (102 ej) .eT) CzereT
when k =
where C1 and C2 are nonzero. The case & * is possible since we assume N > and it yields (55) directly, while combining both cases yields
Ej,y ( eT)
(Ehr)
0
(0) A}. XERN AER
span
Finally as by hypothesis on Oatt all columns of OTV= = क are nonzero, we have Th ej * 0, hence there exists some l such that ((0) e;) ,e()) # O. This directly implies (56), which concludes this step.
2d step: Considering an arbitrary pair j, E (.....,,imm) we now show that
Ez41
span
(0,[()(),x) A).
(60)
XERN Nx dii
A M xdim.
This is of course a trivial from (56) when j = (which is always the case when dim 1) we focus on j * L. Considering = Fi.j + Ex,1, for some pair of indicex i < K E fI, (possible as N 2 2). we have
X) Softmax(XX,x1) XXo2 1)x02 ( + (a(δ))e + 3(0))en) (0) (T)e;)= + + 5(())ex) (0]e:))
20

---

Conservation laws for ResNets and Transformers
where
explle,, Φ1c,))
a(01) = N - 2+ exp(lep, 01e,)) - (epp:ee)))))) 01e)))
N
expl(er, (1e)))
8(0)) N - 2+ exp((ee,) 216,)) exp((er 01s;))
N
exp((e; 0le1))
1
7(11) =
N - 2+ exp( (1e))) + exp(
N
(xp((er, Φlei))
8(01) N - 2 (1e( exp((er,
N
Straightforward calculus yields that for any Iedimxdin
Dz.f f H = e [ai Hej) + by (e,, Hei)l (02 e4) + e, (ei Hei) + dr He,)l (0) et) + + (e2 Hej) + b2 (e,, Hei) (0) e;) + ea (ezle, He) d2 (e, Hey)l (0) e
with appropriate scalars as bis Gi di, where b, + 0. Thus for any A Nxdim and H Rdimxdim
(10,,11 A, H) = 1 H, = [ai (ej, Hey) bi (e Ha)) (e( e,)_ .)
+ |ele,, He) di (e)) (e:(0) ex)
,AA) A
laz(er, Hey b2 He (ex e,)
+ le2le, Hei) + (2 H (e
Since (ep; Heg) lepey H) (F H) for every p. q it follows that
[0s fl (a A)E,, (1 eg) A)E,, a ,)), (e e1) A)E + (4) T Et, + be e,)T A)E +c2(ex(0]e:) Et, (d A)Eij.
By (56), we have Ejg, Ei, E span ((o A'} hence we obtain
b (elelllee)) A) (e ( E,11
es)) A) (e TT) , Bis span
, v (l00) A'J.
To conclude it is enough exhibit choice of A such that + 0 and 3 O: this will indeed imply
E
Eja E spanx, (10,, ((6(0), X A',, hence the the existence of A E didmmdimm such that A
E
span a (0s X A'), and by (55) one will obtain (60) as claimed.
To ensure 3 = 0 is it enough have A e; Such choice also implies b,, which is indeed nonzero.
1 ex (possible as k + i), well as AdT 1 ex, (possible ( ( j)
3d step: Conclusion. Finally by combining (55) and (60), one obtains that
[0,)( XT] A} R did x (2dim)
span X xxdix x
which concludes the proof.
21

---

Conservation laws for ResNets and Transformers
Note that in the case when N 1, then (18) writes
(0,
(61)
with 0 (Q, K, In particular the parameters Q and K remain constant during (4), so 0 Q and 0 K conservation laws with respect to any loss &. Moreover (61) coincides with 2-layer linear network T(0, x) with 0 = (V, 0), whose exact independent conservation laws with respect to any loss such that V2 = IR" have been studied in Proposition 3.5, and are thus obtained via the set of conservation laws + VVT O0T
G. Proof of Corollaries 3.9 and 3.10 for attention layers
Corollary 3.9, Under the assumptions of Theorem 3.8 for each Oatt such that both horizontally concatenated matrices K) and (V, O) have full rank, there is a neighborhood of 0 in which all conservation laws for (18) are functions of QQ KK and VVT -00 and vice-versa.
Proof. We recall that (0)) (Q K2 VT, and we denote 81 = (Q, E Rd x x Rd xdim and 02 = (V,0)
Rdd xdim x Rd xdim and v (U,V) E Rdd xdim x Rd xdim + T E RRdimxdim Let us define: yh : + (0)) and 2/2 - + (02). Then (0) = (0), (0)), that to say ø can be decoupled into two functions depending on two separate blocks parameters. Its Jacobian and Hessian matrices are thus block-diagonal. Denoting Φs, w/. wp the coordinate functions of Φ, ui,, 2 the Lie bracket [V (+), V+)()) 0 thus vanishes for every (61, i2) and j = (21, E (l...., dim]² and as a consequence
(eesspan(VoblC k)) (Vvdspan(Vi) :1) @ (Vw)(() i)).
Since V R., by Theorem 3.8 and Proposition 2.7 we also have W/. = span : i), and thus for any 0:
dim(Lie(W9:))0)) dim(Lic(span (v (0) dim(Lic(span V V (0))
Let us consider 0 such that both horizontally concatenated matrices (Q, K) and (V,0) have full rank.
As K) (resp (V, has full rank, this remains locally the case. By Proposition 4.3 of (Marcotte et al., 2023), the dimension of Lie(span(V ()) (0) (resp. of (e(spannVee?() i (0)) is locally constant, denoted d, (0) (resp. (o)). Then, by Theorem 2.11, the exact number of independent conservation laws is equal to D - di (0) dz( Finally, by Proposition 3.5, all conservation laws are obtained by the set QQ - KKT, VV - -0 of conservation laws.
Corollary 3.10. For any h = =1,... H, the functions
QT KNKT, V.VV O,OO,
define conservation laws for (19) with respect to any loss &.
Proof. We recall that H is the number of attention heads and define 0, = (Qh, Kn, Vh, On) for any S h S H well as Sh + (.O)), where (0n)) = (Q), Kn, V, On). Then, we have using Proposition 2.7:
ws. C span { Vo (1 HJ.j (1,-.. 2dim dim}),
where (o)); denote the coordinate functions of の²_ particular by Corollary 3.9 (we do not need to assume the non- degeneracy condition on as Corollary 3.9 as we just use the direct inclusion the reparametrization d), the functions H. + Q K₁ and H² + V O₂OOT are conservation laws (for any l by using Corollary 3.9 and Proposition 2.7), hence by Proposition 2.3 they satisfy for all O for all h, H (0), V (0) 1. Váz (0) (as Hi and H& only depend on 0).. And thus (0), H (f) Ws = W By Proposition 2.3, we conclude that indeed H and H. are conservation laws for (19) with respect to any loss
22

---

Conservation laws for ResNets and Transformers
H. Proof of Proposition 3.11 about cross-entropy / classification layers
Proposition 3.11. With respect to any loss such that V = IR", there are exactly m independent conservation laws for the classification layer given by (20): h0) Σ, j 1, m.
Proof. Using Corollary 2.5 and straightforward calculus involving the Jacobian of the softmax (that are left the reader) compute
(0+ span #ER =EV;=R"
{0og(0, Tw]
(vec(wrr }) (vec(Z) Z 1 1T = 0}.
span ZER WER E. #1=0
This space is independent of 0, so that there are exactly m independent conservation laws. One verifies directly that for the functions h defined in Prop. 3.11, we have ( Z) = 0, VZ E (0)
= (0). By Theorem 2.11, since diim (im(Tie(Woo))()) -m m.
I. Proof of Proposition 4.3 on "block" conservation laws
Proposition 4.3. Consider a function h E (O, IR) that only depends on the coordinates Or, and for each 0 O denote
θr- = E (0T,n)= E }
(24)
Consider loss that satisfies the assumptions of Proposition 2.3 as well as Assumption 2.4. The function h is conservation law of g with respect to E if, and only if. for every E θ one has N())(0) 1 Ro, Ro,(W99 (WW..) where:
Roy R (1/9.)) =
span
span 9(00 n), x) Ww.
(25)
n (6T) ZEZ4r WEV,
Proof. By Corollary 2.5 and Proposition 2.3, h is conservation law of g if, and only for any 0 (0r, 0re) E A:
Vh(0) 1
span {ōog (0) Tw].
(62)
WEV
Since only depends on 0r, we have Vh(0)- = Vorh(0) Vor)((0) to:
-Vonb(et)
Var h(0))
hence the above orthogonality is equivalent
h(0) 1 span
100,9 x)) Tw), w) v
(63)
mEXo,
Given any and any 7) E Or- (0r), the vector 0' = (0r, n) E θ also satisfies h(0') = h(0) and vo, = h(,)
hence the orthogonality condition (63) also holds at 0', As a result,
V o (()) 1
span we (0r) (e XV;
span
(00,9((0 n), [w =: Ro, (Wa)).
(64)
Conversely, if h satisfies (64), then h also satisfies (63), which is equivalent to (62) and implies that h is conservation law of with respect to f.
J. About the invariances of the neural network
J.1. Conservation laws and invariances of the shallow case
Definition J.1 (Invariant transformation on the cost (1)). A cone-parameter transformation on an open set Q S C R map T IR Ω + RD such that T(, 0) is differentiable for each 0 E Ω and T (0, id. This transformation leaves invariant the cost (1) if for all 0 e Ω and for all 2 0, Lz(((c, Lz Lz (T((0))) Lz) When this holds, simple calculus yields for every A E Ω:
(Zz((),,
de T(e, 0) -=0/ O.
(65)
23

---

Conservation laws for ResNets and Transformers
We denote (r(+) == de T(, ===
From loss invariance to conservation laws. In the context of a gradient flow dynamic (4), implies
(o(t), ((((t))) = O.
(66)
Definition J.2 (conservation law obtained via an invariance). In particular, in light of Proposition 2.3, if function E (22, is such that for all I one has Vh(0)= == (r))), then h is a conservation law. We say in that case that h is a conservation law of g obtained via the invariant transformation T.
The case of a 2-layer linear network
Consider 0 == (U, V) E Ruxr x mxrr and define 7 as the linear transformation (A,,,, = (U exp(cA), V -c
(67)
Simple calculus yields A (U, = (UA, )V Considering g00,x) == UV I and any A E Rrxr T is an invariant transformation on (1): for all e and 8, ((, .) = -) hence Lz Lz(TTA( 0)) = Lz)) As a particular consequence, for gradient flows with linear networks, since (66) holds for T = T with any matrix A E R denoting (M, N) Tr(M N) we obtain 0- (U()),
(V()), (I)A) each time and for any A. Specializing this to any symmetric matrix, we obtain that 0 (U,UA) (V, VV) = 1/2 ((U,A)) (V, V A)). Thus for every symmetric matrix A, UA) VA) is conserved, which coincides with all conservation laws in that case (cf Proposition 3.5).
The case of a 2-layer ReLU network Consider 0= (U, v,) E Rext Rmxr
R lxr and define T as the linear
transformation
(c,,UVV = (U (xp((cA)) V "xxppee"")bepp(ced""") bexp(-edT))
Considering - Uo( bT) and any diagonal matrix A E Rexr, T is linear transformation that leaves (1) (UA, VA, -bA), as diagonal matrices are symmetric.
invariant. Moreover, (U, = (U A, VV -VA, - -bAT)
By restricting ourselves to elementary diagonal matrices A = Er; where Ei is the one-hot matrix in Brxr with the (i, i) entry being 1, i = 1,... ,r, we obtain that for all i, 0 (ia(t), u;()) (i)(1), (t) - (t)b =
d de ((u;, u,) - (u;, ui) - 62). Thus for all i, ||u,]12 || 1|2 63 is conserved, recovering all conservation laws (cf Proposition 3.5).
The case of a 2-layer ReLU convolutive network Consider 0 = ((ux,)) (1g,))ss) the collection of all filters and denote G(f, æ)) the convolutive 2-layer ReLU neural network defined by (42). As explained in (43), we express (42) as a 2-layer ReLU neural network (15) as follows. Let us write for j = 1, er ul)) = (ec((ux))&) E Ruu and v(5) vec((u;.))) E Ra.co, For any x E I one has:
(0,xx) = (u()))
(ulle))) of (CC(((()) ... (2(((t2))) s)),
where C1 (resp.C2) is a linear operator defined in (44) (resp. (45)). We rewrite 0 = (U, V), where U = (1(1) and V = (u(1) v(())) For any =I, C1, we define T, the linear transformation
uloi))
T,C,U,VV = (u (xxp(eE,,)) exp(-cEjj3) T))
The transformation T, leaves(I) invariant: for all and 6, ((T) = g(..-) as T1 and T2 are linear operators, hence (z, 0)) (0). Moreover, Ar, (U,V) V) = (U Eujs: -VE5) and thus (66) writes: 0 = (1), (1))
(16)(1) = d w()i) - (06), wb)))) Thus for all j. (ul)), (()) (v60), w()) is conserved, recovering all conservation laws (cf Theorem 3.6).
The case of an attention-layer We treat the case of the network (18) in the exact same way as for a 2-layer linear network by considering for any symmetric matrix A E Rd, xd, the linear transformations A and A defined by
K.V.K.VVVVOY = (x0((A1).ee((-)))x.000)
24

---

Conservation laws for ResNets and Transformers
and
T C,K.K,KVVO, (Q, K, (epppee)V,xpp((eee )O))
and that leave (1) invariant. We conclude in the exact same way as for the 2-layer linear neural network case and we obtain that for every symmetric matrix A, (Q, AQ) (K,AK) AK) and (V, - (O, (O) are conserved, which coincides with all conservation laws that case (cf Corollary 3.9).
J.2. A conservation law sub-network obtained via an invariance gives conservation law the global network
Proposition J.3. Assume that θ = O x 2₂ x 03 and that for every = (2,,0,, 03) E we can write g(0, = 91 (.1. (2(02, feel..gss.sss. (03, 2))). If there a conservation law h O2 of 92 obtained via an invariance of g2 (cf Definition J.2). then H 0 E θ + f((02) conservation law of the global neural network g.
Proof. We assume there exists transformation T 02) E R O2 (e,, 02) that leaves ga invariant, i.e.
(2), 92 (T, .), Ve, 0g.
Let h 02 E O2 + 1(2)) IR be a conservation law of 92 obtained via the transformation T. This means that (cf Definition J.2)
Vh(02) (62), vOl O2.
(68)
We now define T: (e, f1, 02, 0z) E R 10, 02), (2)),,3)) 03). As T leaves invariant g2. T is transformation that leaves the global neural network g invariant: for all IR one has (T 0). x) g( 0)) Thus, T leaves (1) invariant too. Moreover by (68) we have for any 0 E O:
0
0 vh(02)
Ax(0) = Ar((2) =
(69)
Let show that H (01, 02 0z) E 11 x 2₂ x U3 + h(02) is conservation law of the global network Let us 0(t)
solution of (4). Then for all t such that it holds:
(I(O(()) (0(1), VH((())
di
@ (VLZ(0(1), VH(O())
0 VLZ(O(t)) (hh(este) )
6 (12() =0,
as T' leaves (1) invariant.
K. Proof of Lemma 4.4 on the density of X'
Remember that consider deep network g(8, x) composed of residual blocks denoted 9 (x), corresponding either to block of convolutive ResNet (21), residual MI.P (22) or an attention layer (23), see Section 4.2 for the notations.
Lemma 4.4. Denote θ O x O₁ (or θ = Og+1 x O, x O1 with Oyt1 Raxm when there is a last softmax layer) where for each layer 1 < < O; is the set of parameters O, such that
1. gl, is an open map*:
2. all the rows of the matrix y (resp. U!) from (21) (resp. (22)) are nonzero the convolutive ResNet case (resp. the Transformer case).
ie.. sends an open set to an open set
25

---

Conservation laws for ResNets and Transformers
For every 0 E O we have Xx- Rmm.
In particular, for an attention layer (23), Lemma 4.4 only requires the first assumption Item 1.
Proof. First we show that, without loss of generality, is sufficient to prove the result in the absence of last softmax layer (20). This is simple consequence of the fact that, in the presence of such a layer, since g 11 z softmax(O (0++12) is
infinitely smooth everywhere with respect to 0 E TR the set 2g (which we recall is defined as the collection of all data points x E I such that 0 + g(0,z) .r) is C2 in neighborhood of 0, of Proposition 2.3) for the model g+1 is the same as with the "truncated" neural network go (x) = 98. .... gè, (z).
99. .... 9b.
To prove the result without softmax layer, we define (z) dh, o.. gb, (a) the output after blocks with the convention (x) x.
Our goal is to show that Xy = R", and we proceed by contradiction, assuming that there exists To and r > 0 such that for any E B( r), there is no neighborhood of 0 in which 0' + go (z) is C².
Since each residual layer sb, is either defined as a convolutive 2-layer ReLU network (21), a dense 2-layer ReI.U network (22), or an attention layer (23), a straightforward induction on shows that given any x. the function + go (z) is infinitely smooth in a neighborhood of 0 unless there is at least one residual layer 1 S S q such that 0¹ + gb, (G/² (x)) fails to be infinitely smooth. Since attention layers (23) are always infinitely smooth, the lack of smoothness can only come from a 2-layer ReLU residual block, convolutive not. The lack of smoothness implies that the pre-activation of at least one of the hidden neurons of this layer must be zero, i.e., (Vi (z)) 0 and where we denote Vi2 either the matrix V from (21) or the matrix U/ from (22). Similarly we denote Uj either the matrix U from (21) or the matrix Uj from (22).
Thus, for any x E B(zo,r), r), there exists 1</ S such that gb. : a + a' + Uio(Via") corresponds to the 2-layer ReLU network (21) in the convolutive ResNet case (resp. (22) in the Transformer case), and such that there exists j satisfying = O.
In other words, we have proved that B(xa,r) S UHi where each of the finitely many sets
Lj
H4, = (x'ER"
= 0}
is a closed set of empty interior, as we now show. Indeed Hi. is closed as the reciprocal image of a closed set ({03) by a continuous function. To show that Hi, of empty interior we proceed by contradiction, assuming that there exists a non-empty open set B such that any x' E B satisfies (z') E (RVVLj Ç II" as : A O by Item 2 (0 E +). But as x' E B + Gz') = 1-1 a ..... 9è, (x) is open as composition of open maps (thanks to Item 1) and as B is open, the set (B) is an open set in RT: this is absurd as (B) C (IR [I) Ç R".
Overall, as (xr, r) is covered by finite union of closed sets with empty interior, by Baire's theorem B( r) also has empty interior: this is absurd, which concludes the proof.
L. Discussion on the genericity of Item in the assumptions of Lemma 4.4
We discuss here why Item 2 is a generic condition on the set of parameters, for the example of a residual block associated with a 2-layer ReL.U network:
gy x x (va()
where Uj and V are defined via 0: as in (21) or (22). Since g/و is a continuous, piecewise linear, and homogeneous function of x, we can write sh (z) = Jz(x)) x x. where Jı(z) represents the Jacobian matrix given by: Jı(x) = (x)) Here, Dz( denotes a diagonal matrix whose entries are binary values (either or 1) depending on the activation of the neurons at x. We denote by D the set of all such diagonal matrices. By construction, D is finite. Thus the Jacobian (zx belongs the finite set
J(H V1) = {1m + U,DV : D E D).
If each element of this set invertible (i.e. JU, Vi) C GT. (R), a generic condition), then g6, an open map.
26

---

Conservation laws for ResNets and Transformers
M. Finding parameters for simplified neural network forms
Lemma M.I. Let us consider consider some S. I S q and the associated O₁ defined in Lemma 4.4 and go, that can be either an attention layer with a skip connection (23), or a (convolutive (21) or not (22)) MLP. There exists E Oj such that for any mh one has ge. (x) Z.
Proof. 1st case: So. is defined by (21) (resp. (22)). Then by considering V (resp. Ui) such that all its rows are nonzero (so that Item 2 is satisfied) (it is possible by taking ( "j non equal to zero as the matrices Qi are injective), and by considering
T = 0 by taking all [ equal to zero (resp. U = 0), then for any x E Rm, ob. (2) = x (and particular sé, is open, so
Item 1 is satisfied and thus 0: (U", V*) (resp. 0: (U2, U1) is in θ1).
2d case: gb is defined by (23). Then by considering 0: = (Q), K1, VI, Oi) = (0, one has for any x E R" ob, (ar) x, and gb, is open, implying that 0: O2
N. Proof of Theorem 4.6 on block laws for natural residual blocks
Theorem 4.6. With O as in Lemma 4.4, consider the 1-th residual block of Example 4. (resp. Example 4.2), and denote
Or = O2 and Or the parameters of all other residual blocks. A function H 0 = (01, 0re) E θ + h(0r) that only depends on Or is conservation law of y with respect to loss such that V. = R if and only if is a conservation law of the shallow residual network g' 20), go, (z) with respect the Euclidean loss. The same result holds for 0r == 0y+1 when considering last block (20).
Proof. Case 1 = q 1. First treat the case where Or == 0gt1 In this case, one has for any 0 E O and for any R=:
g(0, x) softmax(0(+11(1)) o... gb. (x)) 90,11 (9%,
g+l
gb. (z)).
As consequence
Vx E No == I o (b,g(O,)) =00 00,+1 g0t 11 9d, ... 9è, (x)), 4 θ
(70)
Then by using Lemma M.I, one can choose 02= E O2 such that 8% o... o gbb (x) = x for any x E R", and thus one obtains Th = {08) c .... o gb, (x) : Or E Or,, x E R²). Therefore one has:
span 1o 19 n (Bg+2, =
span 10,, g(0,
(71)
Since we assume V; = IR", one has
Wgg Corollary 2.5
(00 (00 (0 w)
span (00,.19 g Tw]
span X 9-11 WER"
#ER" We "EER"
span (0o,g' T (0,
span
6 C we RER
Lemma 4.4
span EreEOre ReeerreeeRRE...
span paanfoo.90 9 (),,))
span span(Oe,.(e (0, x) Tw] =: Rey (wa..))
8ye g >E
where we recall that (WP9.) is defined in Proposition 4.3. This concludes the proof in the case Or := 0g+1.
Case </5 q. We now assume that 0r : 0z for some / = 1, .q. Let us consider h(0))) conservation law for gf. Then as all conservation laws of g¹ are obtained via the rescaling invariances of g' as discussed in Appendix J.1 and by using Proposition 3.2, the function H 0 E θ + h(O)) is a conservation law of by Proposition J.3.
We now show the converse result: considering a function H 8 E θ + h(())) that is a conservation law of we wish to show that is a conservation law g'
27

---

Conservation laws for ResNets and Transformers
A. We first consider the case without a last block (20), ie the case where:
g0,ar) =90, o Do. (3) IR...
(72)
It will be convenient rewrite this as the composition of three functions depending on three separate blocks of coordinates
of the variable 0- = (01, 02 03) E θ 62 O₂ (with θ2 = O₁ and θ Og given by Cartesian products of adequate }j. j + D), the middle one corresponding 02 0g:
0: x) gi (01: 92 (2 as (0):
(73)
If q (resp i 1) corresponds to the parameters of the last (resp. first) block, then we directly can write g(0; x) g2 (02; ga (03; :)) (resp. 20; gi (g2 (02; r))), but we will use the formulation (73) in all cases informally by simplicity.
With these notations our goal is to show that conservation law H for g in (73) that only depends on 0g is a conservation law for 92. A technical step, that we slightly postpone, is to prove that there are 01, 0g E θ1 x O such that for every 0g E O2
g(0, x) = g2(02, x), Var E RM
(74)
with 0 = =0:00,0000 00,,0,,00 02). Using this result we proceed using a density argument. Denoting X2₂ the set of all data points
E R such that 8/ + (2(()2,z) x) is C² in the neighborhood of 62, we observe that Xo _ Xx with No defined in Proposition 2.3. (Indeed if x E Ne. then 0' + g(0', x) is C² in the neighborhood of 0 = -01 02. 03), so in particular
8/ + g( L/2, x)). = 92(8), x) is C² in the neighborhood of (2.) By (74), for any 0₂ E O2 and any x E X/ 0g. (0,g(0;z) .) (0,92(02;2). 2². Moreover x + 0,,g(0, (= 92 x)) is continuous on Xg. Indeed either X22- = (when O2 is attention layer) Xú, is the whole space outside finite number of hyperplanes associated to the zeroes of the activation of each hidden neuron: in particular if x E X3, then 92 (02, x)) does not have any activation that vanishes, so it remains the case in a neighborhood B(r, r) of x. this implies in particular that x' E B(x, r) >> 92(02, x²) is continuous. Finally, as X3- = R by Lemma 4.4, for any x E X2, there exist (Ix) E (No) such that uN + x. Thus, as Xo S X2 and by continuity of 00.02((0: a)) on Nb.-
og 0.g(0:zN [N) = 00,92(0;xx) - 0g (0;2)),
Since V = Rm with é the Euclidean loss in RT, we obtain with the considered triplet of parameter 6 = (81,0,, 0s)
War.
Corollrr 2.5
span span (0 (02; WER"UENK
02
S span span(0/,g(0: x) w/ wER" rEA,
(as X23 CR ===xX, Xo)
= span span wER" x
8 90;; x)) mm
(as every finite Anite-dimensional space is closed)
C span C
span
(0) Tw) =: Rgg (379..))
22.0%)
where we recall again that Rg,, (W9.)) is defined in Proposition 4.3. Thus by Proposition 4.3, if a function ( h(22) that only depends on 82 is conserved for g with respect to the Euclidean loss, then for every 0g E é, the gradient Vg, (0))= = Vh(62) is orthogonal to R3, (W9.t)= W hence by Proposition 2.3 h is also conserved for g2 with respect to the Euclidean loss.
62
B. We finally consider the case with last block (20), ie the case where the neural network writes:
go : x E TTTH + g)0,x) =9 9ù, o... gb, (x))R
We recall that 1 A q 1. Denoting Ø- (.,,... 0g) and 0(θ, x)) = G the network (72), the analysis conducted in A above shows that there are 0, E Ō1,0g E ② such that for every 02 O₂ the parameter Ø == (Ø,, 02. 03) satisfies:
0(0,x)) = (2),
Yo E T
28

---

Conservation laws for ResNets and Transformers
so that with 0 (0, 0g+1) have
g(0,x) = softmax(00+10(0..0 = =01tmax(00-122000.00)) Va E R
As result, reasoning as in A above, for any x E R3 one has:
Moreover have
[00s w
span span
0et1 Z WER"
000.92002,0) 150 00 w
= sspa span
1 ZE WERn
02
(04t:00oofmaax00004:440 w)
span #,
span 0gt WER"
span span 10g.92 (02; =: Woak E
In the last line used that
span (0 10-=10000ttmmax(0++100.0. [w -R", 00F1,
(75)
a property that we now show. Indeed, given any E R consider g.t1 == ciy where er E R" the first canonical vector. Simple calculus yields 0 of+1 (a.ls0sotmmx((0+110 x))] yz where
2T 7 (xp(A)(n - (exp(A) +n + (exp(A) + 1
exp(A)
exp(A) (exp(A) + n - 1)2 * 0, with A = x)).
With w = 2/1==111 we get 0 10stttmax 10 x))] w = y. Since this holds for any choice of y, we get (75).
Finally, we can use the exact same previous proof by using Lemma 4.4 and conclude the proof.
C. We finally prove (74). By Lemma M.1, can choose for all layer p * I with p S 9 some parameters Op such that for any x I m one has Jh, (x) z. Thus with 01 (01+1.-... 0g) θ and 0g (0,, 01-1) Og. one has 9i (01,z) 22 and g3(0g, z) for any x e R" Finally one has g(0, x) g2(0g, v) for all æ e R which concludes the proof.
O. Additional figures
We illustrate the notion of blocks overlapping or not a residual connexion in Figure 3-Figure 2, and display experimental results on conserved functions during ResNet training in Figure 4 and during Transformer training in Figure
gy( (x..-
layer
2oke consilutiv natsork
R
g(d.x) R
Figure 2. Block overlapping residual connection
29

---

Conservation laws for ResNets and Transformers
82( (
3-laper convoietvv
R
g(8.x) e R
Figure 3. Block sharing residual connection
Proof of Theorem 4.7 no block conservation laws for blocks overlapping a residual connexion
Theorem 4.7. Consider layer index 1 S S 9-1 and O defined as Lemma 4.4 with the exception that for each 0141 E O we farther require that the rows of Vl+1 are pairwise non-colinear. n,e p then any conservation law of g with respect to the Euclidean loss that only depends 0g is constant function.
While Theorem 4.7 is stated in the main text for the specific case of ResNets with residual blocks that involve 2-layer ReLU network with matrices U, V that are either dense with a convolutive structure, we actually prove a more general version where these matrices have the generic structure associated to (42), with C1, C2 satisfying Assumption E. and two additional hypotheses on the matrices P. and Q, from Assumption E.
Assumption P.I. First we assume that rangee (Qi) + {0), where Q₁ is defined Assumption E.I.
Link with the convolutive case. Assumption P.1 (as well as Assumption E.1) is satisfied for the matrices Qi defined just after Assumption E, that associated the convolutive model (16) when n, p. Indeed in this setting, all Qi correspond to full circular shift operators, and notably 1 Q,1, where denotes the all-ones vector.
(V++1, ,U:) corresponds two consecutive parameter blocks before and after skip connection and that can be written as the composition of g1. 92 and Is where
Proof of Theorem 4.7. Here 02 Or
91
92
(V
21 222)
(1(1)
(Ut
la)
σ id
(76)
//2
Ia
x2
x1
gz(03.2)= =
V
(gbi)) o... 9b, (x),
(77)
x2
and gi
0,
1/1
(g) ... )
(114
La)
y/
(78)
U/
id
y2
The parameters 0,, fg gather all relevant parameters involved in the definitions of 91 and 93. We define d1, d2 such that vl+ RRd:xn 77 Raxdz and will painlessly alternate between matrix and vector representations (in IRnd.) of such parameters. The core of the proof is to show that
Roy (W921))= ==
span
span
(0o,g ((Or, n): x)w} = (dd +d2jn
(79)
7 (6, (0r
(x R
with & the Euclidean loss. By Proposition 4.3 this will imply claimed that the only conservation laws h of g with respect to the Euclidean loss that only depend on 0r are the constant ones.
To prove (79) we proceed two steps, proving separately
(0amn } Rad C Ror (37790.)
(80)
Rdi x {0dgn} C Rez (W)9..))
(81)
By the definition of Re, (W9..)) C (dd as a linear span, this indeed implies (79).
As a warmup, let us characterize Or-(Or) {n E RTD (Or,n) Q) as defined in (24), where we recall that θ = O, x , with the O defined in Lemma 4.4, and where O is also such that Vi+1 has all its rows that define distinct hyperplanes. One then has Or (fr) θ θg where
30

---

Conservation laws for ResNets and Transformers
Evonution conservation
1-100
0.611
0002
0.53
Thinning Loos
9-0.031
25
34.
22
Figure Tracking conserved function and loss during ResNet-18 training on CIFAR-10. each learning rate between 1e-3 and 5e-3, we train 10 models for 50 steps using SGD without momentum weight decay, with 10 different random seeds. For each configuration, we record both the loss evolution (hottom) and the evolution of the conservation error [(0c)-h(9p) (top) with hy defined A(Pa)
(17), where Or represents the parameters associated with the first residual block. The dotted lines show theoretical slopes derived from bound (27), which scale quadratically with Cr The empirical results confirm that the function is approximately conserved and that the slope coefficient maintains proportionality with 72, For non-conserved functions, the evolution is 7 rather than 72, since the first-order term in the Taylor expansion (used in the proof of Proposition 5.1) does vanish that case.
θ Og
O142 x 1441 with θ {U++] E Raxds (U+)), v(+1) (U++I, E O4t1]:
θg = θ x O-11 x x O with θ = (V' Rd (U', v\) E θ.}.
It will also be useful to denote G (resp. 55) the set of all 0 (resp. of all 02) such that for each k: 2 1 + (resp. all
1 1) we have: a) U (by taking all u equal zero), so that (U*, V*) satisfies Item 1 : and b) V satisfies (k) Item 2 of Lemma 4.4 (possible by taking uj non equal to zero as the matrices Q, are injective).
For every 0 62 and 0g Fg we have
91
01:
91
=
Ui+1 (11) + 92-
(82)
U2
V'xx
g3
(3:.z)
=
(83)
With the above notations, highlight that if 0 E O2 is such that U++1 E θ:41 then 0 Ô,, and if 0g Gg is such that yi E Ô then 0g E ②g. In particular such a choice of 0 and 0g implies that (0,, 02) E Or: In the rest of the proof we consider 0,, 03 with these properties, this will enable us to leverage (82)-(83).
1st step: proof of (80).
In that part, we consider 0, E 1 such that T1+1 0, and thus (cf (21)-(22)-(33) we have ger id and the block
31

---

Conservation laws for ResNets and Transformers
Evolution the conservation error (Tremsformer on IMDb without mas attention)
1o-7
k=0.00
4-2.002
--2.00
0.3
Nohnolibo conservtion aret
06-
0.
10
10
Scops
Evclution of the conservation emor (Transformr an IMDs with masked attention)
8-20001
R-2.003
k-2.003
Nomalizdd conservatoo enor
04
10
50
Steps
Figure 5. Tracking conserved function and the loss during Transformer training IMDb dataset. For each learning rate between 1e-3 and Se-3, we models for 50 steps using SGD without momentum or weight decay, with different random seeds. each configuration, record both the evolution of the Frobenius norm of the conserved matrix identified in Corollary 3.10, specifically examining the query and key matrices from first attention head in the layer where masking applied (hottom) or not (top). The dotted lines show the theoretical slopes derived from bound (27), which scale quadratically with Cr² The empirical results confirm that the function is approximately conserved and that the slope coefficient maintains proportionality with 7
32

---

Conservation laws for ResNets and Transformers
parameter 0141 (U++1, V satisfies Item of Lemma 4.4 (i.e. U++ 1441 and thus 0₃ θτ). Since Ul+1 0. combining (82)-(83) with the expression of 92 (76) we have
s(0; x)) = x(0x))=U'a -V' + x.
(84)
Since the above expression of g is independent of VL+1 and oy = 02 = (Vit),U U²) E R4x0 x Ruxds for any
K) E Red xx x Innxds any w E IR." and any x E "R,", at (0), 02 0z) = (Or,0rc) we have:
(0,0(0;)) (H, w) = (V) w) = ((H, K): (0, wo (V'x)
that is to say
de ( x)) ~= (vec(wo((V=2)]
In light of the definition (79) of Re, (W.,)), since 17 (01,0s) Or (fr) for each of the parameters Q 0g considered in this part of the proof, sufficient condition to establish (80) is thus
vec (V'x) Rd
span
V WER
or equivalently (by simple linear algebra)
span o(V'x) R24 V'
For each V E Ôr by definition the corresponding 0 = (01, OT, 02) belongs to θ hence by Lemma 4.4 X2 is dense. Since + oV'zz) is continuous, the above condition is also equivalent to
span VEE
o(V'r) (V' R
(85)
To prove (85) we simply exhibit for each canonical vector e, E R a matrix V E Q and x E R such that (V'xx) es.
For this simply consider À > 0, x = Xe E Rmm, V matrix with its first column equal to e, and V = V/X. If A > 0 is large enough then one can easily check that (U',V) E O, hence V E Q, as claimed.
2d step: proof of (81).
We denote , Ud, the d, rows of the matrix V Rddxn, and for each j ",,,, d1 we denote
Hj = {ar E R oJ x = 0}.
Combining (82)-(83) with the expression of g2 (76) have for any Vt:
g(0; = ul+2 (v x) +x, for each x such that o o(V'x) = =0.
(86)
As in the first step, since this expression of g is independent of U., and as 0r = 02 = (Vi+1, (U(+),U)= E xd x Ruxd2 we obtain that for any V and any w E IR," at f = (01, 0z 0₂) = (0r, 0re), denoting D(x) == diag( ( (199)))))))
00 )]] w (v=c(D(=)) wr )) for each & E R - UH, such that a oV' x) = 0.
(87)
The claim (81) is a direct consequence of the following two inclusions that we prove below
{
(eo(D(2)0)+1 wa ))² C Roy (Woz.))
span 0
span
(88)
wER" =E(R*-UH)
0
Bdi x {0dm) S span span o
{ (vee(D(a))) ww })}
(89)
wx2" -UH)))
Proof of (88). The main difference between the left-hand-side of (88) and the definition of Rø, in (79) is that in (88) æ can be freely chosen in the set R. U, independently of 01 (and of fz)). To prove (88) we exhibit below two vectors 0 E such that of 0- - Or and that for every x R U,H, and w E R", either have
0+ 0+,0-
100-9(0 ;x)] w (vec(D(a)))++1 wx T)) E
;x)]
span x'E wER"
((00,0(0::=')) w]
(90)
33

---

Conservation laws for ResNets and Transformers
or the equivalent with 0 instead of 0t. By the closedness of finite-dimensional spaces and (79) we get
(aac(D(a))"+1T w=x )) E span [10,,,(0+;;))) w) S Røy (eerr(WW) (W/...),
0
a'Ex, WERR
or of course the equivalent with -- instead of 8*, yielding the desired conclusion (88).
We now proceed to the construction of 0+, 00= is where we use Assumption P.1: this enables us to consider an arbitrary
nonzero an E n rangeQ, _ IR" and to denote a = ap
ac
E R
= RR" as well as
A (x' RD <0), A (o' E
0), and A= {a' E R. To' 0}
We construct o+ (resp. 0-) ) so that (90) or its equivalent with -- holds for every x in the intersection of IR" UHg with A with A4, or with A
Recall that we consider 0, 1 such that U++1 θj and that this implies 0 E θ₁. We consider 0z G3 such that Vi has all its rows that are equal² to ea in particular V4 satisfies Item in the assumptions of Lemma 4.4- where e O is small enough to ensure that (U", Vi) E O₁ and hence V E é, and 03 E θg. We denote B+ = (01, 0r, 0z). The construction 0- is similar, except that V has all its rows equal to -ea for small enough e > O.
Considering x E (TR" - -UH n+ our goal is now to establish (90). As Xz- = R. by Lemma 4.4 (as 0+ E +), there exists a sequence of vectors IN E Xg+ such that xN - x. By construction AA is open, and Xg1 C R" LH, (Indeed E No+. then 6' + g(0', x) is C² in the neighborhood of 6 == (01, 0r, 0z). so in particular 0' + g( f 0z), Ul+1 ((V+)))z) (by (86)) is C² in the neighborhood of 0r.) hence for N large enough we have IN E (R" -UH, As A+ S fa' (V = 03. the expression (87) of [0,,,(0+;; [0} is valid on (R" nAt Since it is continuous (because D(-) locally constant on FT -UHz) we get
(vac(D())) mnr )) 100r ((t ;x)]] w-w= lim 100rg(0* :xN)]] TT w span x'EX WER""
{ ((0t)(()==']]]
(91)
N-000
therefore establishing (90) as claimed. The same reasoning holds for x E (R" -HH, - using 0- instead of 0+.
Finally for E (K" _H, n A one has in particular Vir Vi = As we have seen, Xo+ Xy- R", hence there exists a sequence of vectors æN E Xgt U Xg- such that NN - x. By extracting a subsequence if necessary can assume that or (V[ax) = 0 for every N (or that aVV = 0 for every N). Without loss of generality we assume [V ) 0 for all N (the other option is treated similarly). As proven above, we have Xg: U %o S R² - UUH, hence
EN E (R" - UHl,) n fa': (V(a'' = 0} for every N. Thus (91) remains valid and yields (90).
Proof of (89). First, observe that it is enough prove that for each k E (1,-.. (1,.), } there exists 0, Ô 0 (0,, 0r, 02) 9, and w E R such that
vec (ex.² T
span (ōe, g (0: x" (0:z")w, Vz E R
(92)
of U,
Indeed, by (87) and the closedness of umite-dimensional spaces this implics (89).
To prove (92), given an arbitrary k E (1,....dd)) J the core of the proof is to exhibit below some 0, E o5 such that the matrix UL+1 O (recall that this implies 0, ②)) has nonzero k-th column, and that for every E R
vec
"ax" TauT,0 ,0) E span Twa
(0org' 9 (0;x')w), Vx E R".
(93)
ot ECERR-U,,,
Since the k-th column of U1+1 is nonzero, there exists an index l such that UL+1 H 0. Setting w == eg we have Eg. (U,+UUI) = ac,x" T, hence (93) implies (92).
Given an arbitrary A E (1,.... , di }. we consider k'' = mod Pl, by assumption Assumption E, one has range Par) + {0). hence one can build (1+1 as (44) such that its k-th column non equal zero. To ensure that (U++1, vitl) satisfies Item 1 (hence U++ θ and hence 0 +) rescale U++1 such that its norm is small enough.
'The assumptions of Theorem 4.7 the matrix forbid yl+ to have colinear rows, but there is no such constraint on V
34

---

Conservation laws for ResNets and Transformers
With this choice of U prove (92) for every IR" we first show for living in Hi. = Hr - - (note that this non empty set the hyperplanes are pairwise distinct by definition of )) before proving it for x: in the complementary direction. We denote
A = f= E RR
th > 0}, and Ag fxx E IR. : 1 < <0}
where we recall that 01,... Vd, are the rows of V
We first show that (92) is satisfied on H'g-
Consider x' E Hig- Denote B(c, n) the open Euclidean ball of radius vj > centered at c E R". Given any 1/ > 0, by continuity of x E T + Vi+lz (of x. ... Va x) E Rdh there exists o. E B(x', n) n and "n E B 17) n n such that 1 = sign(e) x * sign(e, )= - -1 while for all j & k, sign(e) xx)) = sign(y) oj x') 0 (as the hyperplanes are pairwise distinct by definition of (). It follows that z E I2" LHg. As consequence, we have D(at) ((,,) = diag(ex) hence for every w E IR:
vec (Ex.(U(t1) "ox","" .) vec ((D(x)) - ((2,)) (U++1)T wxw",,))
lim vec (D(x;)) (w((t))((((tt)) w(x,)) D(r,, (( w(, ,0.)
1-00
(0o,9 g (0; 2):]. lim (0o,g) (0; )w G T (0: )w) E span -ER"-UHA
9-00
This establishes (92) for any x' Hy..
We finally show that (92) is satisfied for any = 2k.
With x' E H4 above the continuity of x E IR" Vi+1x 10Tx, ,,T, "0"T,,,"",1 such that the vectors
E Rdd also implies the existence of 0
FK =
= x'' + K Uky
Ke (-2 11,1
satisfy 1y xx F O while for all j + k, s xx) = sig x') (as the hyperplanes are pairwise distinct by definition of θ), so that xx E R - LLH and we similarly obtain
7 vec (F wow
(U uo (U*=))(U (T
vec
=7 7 vec
-((x))
=D(e_2)
vec (D(12) wan ((a))( woT - (D(x wx - D( ((Utt))Tu:) 1x ,0)
= 100, 22)l_ w (0, (00, ())
E span {0u.g] [00,)T(00) (0;x)u} =ER"-UH,
which gives (92) for any Rek-
We have thus proved that (92) holds for any x E H'' U Rug. Since R= span(H'' U Reg) this establishes (92) for every x E _R. Since this holds for any k this completes the proof of (89) and therefore of the theorem.
Q. Proof of Proposition 5.1 (approximate conservation of laws under discrete dynamics)
Proposition 5.1. Let h()) be conservation law of the gradient flow with a bounded Hessian YO, |1th(0)] < Cr. Suppose that the gradients remain bounded in expectation throughout the algorithm:
Ez,,0, (Bx) < Cl.
(26)
C₂Cr T2
Then, we have
E E(((k)) h(0b)l <
(27)
1=0
35

---

Conservation laws for ResNets and Transformers
Proof. Since h is conservation law, by Proposition 2.3 have (Vh(O*), V LZz. expansion yields
0 for every k. Besides, Taylor
VLz, (0x)] (h(E)V (0s) h(0x+1) h(O,)) -r
for some & in the segment [0k, 0x+1]. Applying the bounds on h()) and V Lz. (fn), we get
~ - C.CL
2
Summing over k completes the proof.
R. Conservation laws for Adam Flow
We recall that for Adam flow the space w is defined by
W
span
{sign (VLz(0)))
(94)
2 2= ((xx,
By directly adapting the results of Section 2.1 with Wg defined in (94), this leads to the direct corollary:
Corollary R.1. Consider a loss f(z, y) that satisfies Assumption 2.4. Under Assumption 2.6, then for all 0 O
Wa span {sign (0o(0)T ww)
(95)
wEW 610)
In particular if Wf. 6(0)
Bd one has
Wa. span (sign (sign (de(0)) w)}
(96)
wERd
Thus by Theorem 3.4 and Corollary R.1 one directly obtains:
Theorem R.2. Under Assumption 2.4, Vi RR, then considering θ TR and (00) UVT for linear neural networks,
one has: W.f.t RRd and W = span (sign (T T w) for Adam flows.
6(8)
wER
Similarly, by applying Theorem 3.8 and Corollary R.1 one directly obtains the following theorem for an attention layer: Theorem R.3. Under Assumption 2.4, V. R." and N 2 2 then
W.fef 6(0)
Rdd and Wa span (sigu (8e(0)Tw)) v0 Oatt. wERd
for attention layer and for Adam flows, where recall that the reparametrization defined by x(0) (1h,, (22) with Φ1 = K and x2 V VOO.
Moreover as the parametrization & and d2 are separable, one has under the same assumptions:
{
(iggn (deh (Q, K)T w1) sign (V, w2) }} 40 E Oatz.
K)TT)
Wa.
span W2))
Numerical experiments The code provided in our GitHub repository numerically investigates the dimension of the space w defined in (96), i.e. the space spanned by the sign vectors of gradients dØ(0) 2, where @U, V) UVT and V) Raxr, x mmxr For various choices of n, m, r, the code: 1) generates random parameter matrices and 2) samples random vectors E Rnxm; 3) computes the gradients of with respect U and and 4) collects the sign patterns of the projected gradients 0o(0) It then estimates the dimension of the linear span of these sign vectors by calculating the rank of the resulting sign matrix. This provides a lower bound the dimension of the sign- gradient space W In all tested configurations, this lower bound consistently equals the total parameter dimension D (n m)r, indicating that the sign vectors span the full parameter space, expect in the case n m = r = 1 where in that case the dimension the space is equal to 1. Consequently, this numerical observation suggests that there are no conservation laws for the mapping @UU V), except in the case where n = m = r = 1. In the latter case, there is indeed exactly one independent conservation law given by: [ |V| as we now show.
36

---

Conservation laws for ResNets and Transformers
The special case m Adam flow (12).
We show that 0 (u, R2
|u|
|o]]
is conservation law for g for the
By characterization of conservation laws we only need to show that for any 0 E (R,)²
()( 1 (sign (do(0)Tw)) w)} Rsiig
wER
sign(u) -iign(u)
which is always true as Vh(0)
37


Reviews:
[{"review_id": "QPfqaoxK63", "reviewer": "Reviewer_zuVu", "summary": "This paper studies conservation laws for gradient flow on a variety of modern neural network architectures, including ResNets and Transformers. The paper provides a characterization of all possible conservation laws for both shallow and deep architectures. Moreover, the paper quantifies the extent to which SGD satisfies an approximate conservation law.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his constructive feedback.\n\n> **Q1a** the abstract claims (...)\n\nThank 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.\n\nExample: 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$:\n\n- either $vx, sx \\leq 0$ and then $g(\\theta, x) = x$ and thus $\\nabla_{\\theta} g = 0$;\n- 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$.\n\nThus 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. \n\nThis phenomenon seems however restricted to the scalar case, as we empirically show in Section 4.4. \n\n> **Q2a** h cannot depend on the dataset (...)\n\nThank 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.\n\nWhile 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. \n\nIndeed, 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.\n\nThus, 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. \n\nWe appreciate this suggestion and will add a discussion in the conclusion highlighting this as a promising direction for future research.\n\n> **Q3a** (...) conservation laws for Transformers (...)\n\nIndeed, 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.\n\n> **Q4a** Many of the results in Section 2 and 3 seem to be restated results (...) \n\nAll 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.\n\n> **Q5a** in which ways does their paper generalize the results of these prior works?\n\nThe 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$.", "claims_and_evidence": "For the most part, the main claims of the paper (a characterization of all possible conservation laws) are well supported.\n\nOn a more minor note, the abstract claims that \"some deeper networks can [have] more conservation laws than intuitively expected.\" I do not see where this claim is proven in the paper.", "methods_and_evaluation_criteria": "Yes, the methods in this paper make sense.", "theoretical_claims": "I skimmed the proof of the main claims, and they appear to be sound.", "experimental_designs_or_analyses": "The experimental validation in Figure 1 appears to be sound.", "relation_to_broader_scientific_literature": "This paper builds on prior works establishing a \"balancedness\" condition for ReLU and linear neural networks (Du et al., 2018; Arora et al., 2019), along with previous work on conservation laws of gradient flow (Marcotte et al., 2023, 2024).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "- One weakness of this paper is that the conservation law $h$ must hold for \"any initialization and any dataset.\" In particular, $h$ cannot depend on the training dataset. This seems to be an overly strict condition. For example, in overparameterized linear regression (i.e L(w) = \\|y - Xw\\|^2), the quantity $P_{X^TX}^\\perp w$, which does depend on the dataset, is conserved, and this is useful for characterizing the implicit bias of gradient descent.\n\n- Next, while the paper builds significant technical machinery, the specific conservation laws for Transformers derived in Corollary 3.9 and 3.10 are already known in the matrix factorization literature; indeed, given any model with a parameter $W$ parameterized as $W = Q^TK$, it is well known that the quantity $QQ^T - KK^T$ will be conserved throughout the gradient flow trajectory. It is thus not clear to me what additional insights this paper provides.\n\n- On the clarity front, I find Section 2 to be quite technically dense and hard to follow. The paper would improve with more exposition providing intuition for why each of these technical results are needed.", "other_comments_or_suggestions": "Typos:\n- abstract: “can more” → “can have more”\n- abstract: “the introduction the” → “the introduction of the”\n- line 102, left column: “that “overalapping” two residual blocks”", "questions_for_authors": "- Can the authors please comment on my concerns in the \"Strengths and Weaknesses\" section?\n- Many of the results in Section 2 and 3 seem to be restated results from (Marcotte et al., 2023, 2024). It is thus not clear to a reader unfamiliar with these prior works what the additional contribution of this current paper is. Could the authors please clarify in which ways their paper generalizes the results of these prior works?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9eR1Z1ib3t", "reviewer": "Reviewer_RSFD", "summary": "**Summary of Contributions:**\n\nThe paper makes significant contributions by **deriving and analyzing conservation laws for modern deep learning architectures, specifically convolutional ResNets and Transformer networks**. It extends the understanding of conservation laws, which were previously mostly studied in shallow ReLU and linear networks, to these more practical architectures. The authors first show the conservation laws for basic building blocks and then analyze deeper networks, introducing the notion of conservation laws that depend only on a subset of parameters (e.g., a layer, a residual block, or an attention layer). A key finding is that **residual blocks have the same conservation laws as the block without the skip connection**. The paper also demonstrates that these conservation laws, initially established in the continuous gradient flow regime, are approximately preserved under discrete Stochastic Gradient Descent (SGD) dynamics.\n\n*   **Literature:**\n    *   The paper **appropriately cites relevant works** in the field of conservation laws in neural networks and the theory of deep learning. It clearly positions its contributions in relation to existing results, such as the completeness of conservation laws for shallow ReLU and linear networks and the analysis of implicit biases in ResNets and Transformers. The references are comprehensive and include both foundational works and recent advancements.\n\n**3. Pros and Cons:**\n\n*   **Pros:**\n    *   **Novel and significant theoretical contributions** extending the understanding of conservation laws to modern deep architectures (ResNets and Transformers).\n    *   **Rigorous mathematical framework** with formal definitions, theorems, and detailed proofs in the appendices.\n    *   Introduction of the concept of **conservation laws depending on a subset of parameters** (\"block\" conservation laws).\n    *   Important findings such as **residual blocks having the same conservation laws as their non-residual counterparts**.\n    *   Theoretical analysis showing that conservation laws are **approximately preserved under SGD dynamics** with an error bound.\n    *   **Empirical validation** of the approximate conservation under SGD through numerical experiments on ResNet-18.\n    *   Completeness results for conservation laws in several basic building blocks with skip connections.\n\n**Overall Assessment:**\n\nThis paper presents a **highly valuable and significant contribution** to the theoretical understanding of deep learning by extending the concept of conservation laws to modern neural network architectures. The rigorous mathematical derivations and the analysis of the behavior under SGD provide important insights into the training dynamics of ResNets and Transformers. While some limitations exist in terms of the architectures and optimizers covered, the novelty and soundness of the presented work are strong.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for his constructive feedback and positive comments.\n\n> **Q1b** challenges in proving the completeness of laws beyond blockwise laws\n\nFrom 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.\n\n> **Q2b** leveraging newly discovered conservation principles in practice\n \nA 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.\n\nWe will insist more on these practical applications in the introduction of our paper.\n\n> **Q3b** Could the authors briefly outline their planned approach (...) \n\nRegarding **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. \n\nTo 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:\n$ \\dot \\theta (t) = - \\text{sgn} (\\nabla L_Z(\\theta(t))) \\ (1).$\nIn our framework, this corresponds to considering the linear space of signed gradients:\n$\nW(\\theta) := \\text{span} \\\\{ \\text{sgn} (\\nabla L_Z(\\theta)) : Z \\\\}\n$\nfor 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\n$\nD - \\text{dim} (\\text{Lie}(W)(\\theta)) = D - \\text{dim} (W(\\theta))\n$\nindependent conservation laws. In particular, if $\\text{dim} (W(\\theta)) = D$, then no conservation law exists in that region.\n\nIn 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.\n\nWhile 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.", "claims_and_evidence": "The paper \"Transformative or Conservative?\" claims to support its contributions with clear evidence through **mathematical derivations and proofs** presented within the main text and in the appendices. Several key findings are explicitly linked to theorems, propositions, and corollaries, with their proofs detailed in the appendices. For instance:\n\n*   The structure theorem (Theorem 2.1) demonstrating the relationship between conservation laws with and without weight decay has its proof in Appendix A.\n*   The conservation laws for shallow multi-channel convolution layers are presented in Theorem 3.6, with a more general proof in Appendix E.\n*   The characterization of conservation laws for self-attention layers is given in Corollary 3.9 and 3.10, with proofs in Appendix G.\n*   The existence of independent conservation laws for the cross-entropy classification layer is shown in Proposition 3.11, with a proof in Appendix H.\n*   The proposition that residual blocks have the same conservation laws as their non-residual counterparts (Proposition 3.2) includes a proof in the main text.\n*   Theorem 4.6, stating that conservation laws depending only on a block's parameters match those of the isolated shallow block, has its proof in Appendix M.\n*   The absence of conservation laws for overlapping blocks in ResNets (Theorem 4.7) is proven in Appendix O.\n*   The approximate preservation of conservation laws under SGD (Proposition 5.1) is supported by a proof in Appendix P.\n\nFurthermore, the paper includes **numerical experiments** in Section 5.2 and Figure 1 that empirically evaluate the behavior of a conserved function during ResNet training with SGD. The results are presented as confirming the approximate conservation and aligning with theoretical expectations derived from Proposition 5.1.", "methods_and_evaluation_criteria": "Here's a breakdown:\n\n*   **Proposed Methods:** The core method employed in the paper is **mathematical derivation and analysis** of conservation laws. This is a fundamental approach for studying invariant quantities in dynamical systems, which aligns directly with the definition of conservation laws in the context of gradient flow training dynamics. By first establishing conservation laws for basic building blocks (like shallow networks and attention layers) and then extending the analysis to deeper architectures and considering subsets of parameters, the paper adopts a **structured and analytical approach** to a complex problem.\n\n*   **Evaluation Criteria:** The evaluation strategy consists of two main components:\n    *   **Mathematical Proofs:** The paper claims to support its theoretical contributions with **rigorous mathematical proofs** provided in the main text and, more extensively, in the appendices [the previous response]. Deriving these proofs is the primary way to validate the existence and form of the identified conservation laws. The presentation of findings as theorems, propositions, and corollaries indicates a commitment to mathematical rigor [the previous response].\n    *   **Empirical Validation:** To examine the persistence of conservation laws under practical training scenarios with discrete dynamics like Stochastic Gradient Descent (SGD), the paper includes **numerical experiments**. These experiments track the evolution of a specific conserved quantity during the training of a ResNet-18 on the **CIFAR-10 dataset**. CIFAR-10 is a standard benchmark dataset in computer vision, making it a relevant choice for evaluating the behavior of ResNet architectures, which have been highly successful in this domain. The paper analyzes the conservation error in relation to the learning rate, comparing it with theoretical bounds derived for SGD. This empirical approach helps to bridge the gap between the theoretical continuous gradient flow regime and real-world discrete optimization.\n\nIn summary:\n\n*   The use of **mathematical derivations** is a natural and appropriate method for investigating conservation laws, which are inherently mathematical properties of the training dynamics.\n*   The **evaluation based on mathematical proofs** is the cornerstone of validating the theoretical claims made in the paper [the previous response].\n*   The inclusion of **numerical experiments on a relevant benchmark dataset (CIFAR-10) with a widely used architecture (ResNet-18) and optimizer (SGD)** provides empirical support for the theoretical findings, particularly concerning the behavior of conservation laws in discrete optimization settings. Tracking a theoretically conserved quantity and comparing its empirical behavior with theoretical predictions is a sensible evaluation criterion.", "theoretical_claims": "I haven't checked all the proofs of the paper, but some of them were inspired by previous literature.", "experimental_designs_or_analyses": "Yes Experimental design is well versed.", "relation_to_broader_scientific_literature": "*   **Extending the Study of Conservation Laws to Modern Architectures:** The paper explicitly states that while conservation laws in gradient flow training dynamics are relatively well-understood for shallow ReLU and linear networks, their study has been **\"largely unexplored for more practical architectures\"** like convolutional ResNets and Transformer networks. This directly addresses a gap in the existing literature. Prior works like **Saxe et al. (2013)**, **Du et al. (2018)**, and **Arora et al. (2019)** established conservation laws primarily for these simpler architectures. This paper aims to bridge this gap by focusing on more complex, state-of-the-art models.\n\n*   **Building Upon Completeness Results for Shallow Networks:** The paper mentions that **(Marcotte et al., 2023) demonstrated the \"completeness\" of conservation laws for shallow ReLU and linear networks** under Euclidean gradient flows, meaning no additional conservation laws exist in those cases. The current paper extends this line of inquiry to more complex architectures, not just characterizing laws but also considering their completeness for certain building blocks.\n\n*   **Connecting to Work on Conservation Laws with Different Optimizers:** The paper cites **(Marcotte et al., 2024)** for unveiling novel conservation laws under alternative optimization algorithms like non-Euclidean gradient flows and momentum-based dynamics, highlighting their different characteristics compared to simple gradient flows. While the current paper primarily focuses on Euclidean gradient flows, it acknowledges this related work, suggesting a broader interest in the role of the optimization algorithm in determining conservation laws.\n\n*   **Generalizing Prior Findings on Convolutional Networks:** The paper notes that **(Du et al., 2018) identified conservation laws for feed-forward networks with single-channel convolutions**. Theorem 3.6 in the current paper **generalizes these findings to multi-channel convolutional ReLU networks**, demonstrating that similar types of functions are preserved in this more complex setting.\n\n*   **Investigating the Impact of Residual Connections:** Given the fundamental role of skip connections in ResNets **(He et al., 2016)**, the paper specifically examines their effect on conservation laws. Proposition 3.2, which shows that residual blocks have the same conservation laws as their non-residual counterparts, is a novel finding that directly relates to the architectural innovation of ResNets. This also contrasts with work like **(Marion et al., 2023)** which explores the implicit bias of ResNets towards Neural ODEs but doesn't focus on conservation laws in the same way.\n\n*   **Analyzing Conservation Laws in Transformer Architectures:** Transformers **(Vaswani, 2017)** have become dominant in various domains. The paper's derivation and analysis of conservation laws for **single attention layers (Corollary 3.9) and partially for multi-head attention (Corollary 3.10)** contribute significantly to understanding the training dynamics of these complex models. This builds upon work like **(Vasudeva et al., 2024)** which focuses on the implicit bias and fast convergence of self-attention but through a different lens (convergence to SVM).\n\n*   **Relating Conservation Laws to Network Invariances:** Remark 3.1 and Appendix J.1 discuss the intrinsic connection between the identified conservation laws and network invariances. This links the paper's findings to a broader theme in machine learning where symmetries and invariances of the model and loss function play a crucial role in the learning process.\n\n*   **Examining the Persistence of Conservation Laws Under Discrete Optimization:** Section 5 investigates whether the conservation laws derived in the continuous gradient flow regime persist under **Stochastic Gradient Descent (SGD)**, a widely used discrete optimization method. Proposition 5.1 and the numerical experiments provide insights into the approximate nature of conservation in this practical setting, connecting the theoretical findings to real-world training scenarios. This relates to the vast literature on the behavior and convergence of SGD in deep learning **(e.g., Bach, 2024; Garrigos & Gower, 2023)**.\n\nIn summary, the paper's key contributions are deeply intertwined with the existing scientific literature by **extending established concepts of conservation laws to modern deep learning architectures**, building upon prior theoretical results for simpler models, and investigating the practical implications of these laws in the context of widely used optimization techniques and network design principles. The specificity of analyzing ResNets and Transformers, including the impact of residual connections and attention mechanisms, positions this work as a significant advancement in understanding the training dynamics of contemporary neural networks.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "*   **Strengths**\n    1.  Novel theoretical framework for analyzing conservation laws in ResNets and Transformers.\n    2.  Rigorous mathematical derivations and completeness results for key building blocks.\n    3.  Insightful findings regarding the impact of residual connections and the behavior under SGD.\n    4.  Introduction of \"block\" conservation laws, offering a new perspective on analyzing deep networks.", "other_comments_or_suggestions": "**Detailed Evaluation:**\n\n*   **Novelty, Relevance, and Significance:**\n    *   **Novelty**: The paper is **highly novel** as it tackles a largely unexplored area of conservation laws in modern, deep architectures like ResNets and Transformers. The introduction of \"block-wise\" conservation laws and the analysis of their persistence under SGD are new contributions. The paper also provides new completeness results for several basic building blocks with skip connections.\n    *   **Relevance**: Understanding the training dynamics of deep networks is a fundamental challenge in deep learning. Conservation laws offer valuable insights into the implicit bias of training algorithms and network architectures, as well as playing a role in theoretical convergence analyses. This research is therefore **highly relevant** to the deep learning research community. Furthermore, the potential application of these laws in designing new optimization schemes adds to the practical relevance.\n    *   **Significance**: By extending the theory of conservation laws to ResNets and Transformers, the paper bridges an important gap in our theoretical understanding of these widely used architectures. The findings, such as the invariance of conservation laws with respect to skip connections and the approximate preservation under SGD, are **significant** for both theoretical analyses and potentially for guiding the development of new training techniques. The completeness results for basic building blocks provide a solid theoretical foundation.\n\n*   **Soundness:**\n    *   The paper appears to be **theoretically sound**. It builds upon existing frameworks for analyzing conservation laws in gradient flows. The paper provides formal definitions and propositions, with many proofs detailed in the appendices. The structure theorem (Theorem 2.1) establishing the relationship between conservation laws with and without weight decay seems crucial. The characterization of conservation laws for various building blocks (shallow convolutional ReLU networks, attention layers, cross-entropy classification layers) is supported by theorems and corollaries.\n    *   The extension to deeper networks through the analysis of \"block\" conservation laws (Proposition 4.3 and Theorem 4.6) provides a rigorous way to connect the conservation laws of individual components to the global network behavior. The negative result regarding conservation laws overlapping residual connections (Theorem 4.7) is also an interesting and potentially important finding.\n    *   The analysis of the persistence of conservation laws under discrete SGD dynamics (Proposition 5.1) provides a theoretical justification for the empirical observations. The error bound derived scales with the square of the step size, which is supported by the numerical experiments.\n    *   The assumptions made throughout the paper (e.g., Assumption 2.4, 2.6, A.8, E.1) are clearly stated, and their implications are discussed. The paper acknowledges limitations, such as not covering Adam optimizer or normalization layers.\n\n*   **Quality of Writing/Presentation:**\n    *   The writing is generally **clear and well-structured**. The paper begins with a concise introduction and clearly outlines its contributions. The use of definitions, propositions, theorems, and corollaries provides a formal and logical flow.\n    *   The organization of the paper, moving from basic building blocks to deeper networks and then to discrete dynamics, is logical. The appendices provide necessary details for the proofs, allowing the main body of the paper to remain focused on the key concepts and results.\n    *   The inclusion of numerical experiments (Section 5.2 and Figure 1/4) helps to illustrate and validate the theoretical findings regarding the approximate conservation under SGD.\n    *   Potential improvements could include:\n        *   More intuitive explanations of the implications of the discovered conservation laws for the training process and the properties of the learned models.\n        *   Further discussion on the intuition behind the absence of conservation laws overlapping residual connections.\n        *   A more detailed explanation of the conditions under which the key assumptions (e.g., bounded Hessian, bounded gradients) hold in the context of deep networks.", "questions_for_authors": "1.  **Completeness of Conservation Laws for Key Architectures:** The paper establishes conservation laws for several building blocks, including single attention layers, and shows that residual blocks inherit the laws of their non-residual counterparts. However, the completeness of conservation laws for deeper ResNet and Transformer architectures, as a whole, is not fully addressed, and completeness for multi-head attention remains an open problem.\n    *   **Question:** Could the authors elaborate on the challenges in proving the completeness (or demonstrating the existence of additional laws) for deeper ResNet and Transformer architectures? What are the key obstacles compared to the analysis of the individual building blocks?\n    \n2.  **Practical Exploitation of Conservation Laws for Training:** The paper touches upon the potential for designing new optimization schemes that enforce or deviate from these conservation laws to potentially accelerate convergence. The numerical experiments demonstrate the approximate preservation under SGD.\n    *   **Question:** Beyond observing their approximate preservation under SGD, do the authors have further insights or preliminary ideas on how the newly discovered conservation principles could be actively leveraged in practice? For example, could these laws inform initialization strategies, regularization techniques, or the design of more efficient optimizers for ResNets and Transformers?\n   \n3.  **Addressing Limitations in Future Work:** The conclusion explicitly mentions several limitations, including the absence of analysis for Adam optimizer, normalization layers, and the full complexity of multi-head attention in Transformers.\n    *   **Question:** Could the authors briefly outline their planned approach or initial thoughts on how they intend to address these limitations in future research? What are the main challenges they foresee in extending their analysis to these more complex scenarios?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PXsnK89V8K", "reviewer": "Reviewer_Y1a9", "summary": "The authors study conservation in more general networks than the previously studied ReLU and linear networks. In particular, they focus on convolutional ResNets and Transformers.\n\n## Update after rebuttal\nI appreciate the authors' detailed rebuttal, which has resolved my concerns. As a result, I have raised my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments and insightful feedback.\n\n>  **Q1c** about structure theorem\n\nWe 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.\n\n> **Q2c** CIFAR results on the Transformer architecture\n\nIn 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.\n\n> **Q3c** how this relates to conservation laws in PDEs and scientific machine learning literature (...)\n\nThank 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.* \n\n> **Q4c** in Essential References Not Discussed\n\nThank 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.\n\n> **Q5c** How to support Adam as future work\n\nAnalyzing 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. \n\nThe 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.\n\n> **Q6c** (...) connected to PDEs with a spatial component (...)?\n\nThis 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.\n\n> **Q7c** (...) extended to time series data? \n\nTransformers 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.", "claims_and_evidence": "Yes the claims that the paper is the first to analyze conservation in Transformers is true to the best of my knowledge and the proofs are rigorous of the proposed theorems.", "methods_and_evaluation_criteria": "Yes, the methods used in the proofs make sense for the gradient flow problem at hand. This work is mainly theoretical.", "theoretical_claims": "The authors prove that conservation laws of gradient flows with weight decay are determined by their equivalent without weight decay in Theorem 2.1 and list this as a contribution. The emphasis on weight decay is unclear at this point in the introduction. The authors then discover new conservation laws for more modern architectures with the assumption that they have skip connections, i.e., shallow multi-channel convolution layers (Theorem 3.6), self-attention layers (Corollaries 3.9-3.10), cross-entropy classification layer (Prop 3.11) and MLP (Prop 3.2). Theorem 4.6 is the proof that these conservation laws match the laws of blocks in isolation. They also prove that the error bound scales as the step-size squared. The proofs seem correct to me.", "experimental_designs_or_analyses": "The paper mainly consists of theoretical results. There are numerical results that demonstrate the the error scales numerically as the $O(step-size^2)$ in Figure 1. This numerical experiment studies ResNet-18 on CIFAR-10 dataset. It may be interesting to test different architectures on additional datasets as well so that the numerical results can be more thorough. I would be interested in seeing the CIFAR results on the Transformer architecture.", "relation_to_broader_scientific_literature": "- The paper is well-connected within the gradient flows literature and relevant to extend the work there to more modern architectures including transformers. The authors in the introduction describe conservation laws as a \"balancedness condition\". I think it would be interesting to see how this relates to conservation laws in PDEs and scientific machine learning literature, where a conservation law in integral form states that the rate of change of total mass of the system is balanced by the change in fluxes in the domain, which seems to be a related condition (see ansen et al., \"Learning Physical Models that Can Respect Conservation Laws\", ICML, 2023 and Richter-Powell et al., \"Neural Conservation Laws: A Divergence-Free Perspective\", NeurIPS, 2022).", "essential_references_not_discussed": "- Other good ODE-based references to add include  Ott et al., \"ResNet After All: Neural ODEs and Their Numerical Solution\", ICLR, 2021 and Krishnapriyan et al., \"Learning continuous models for continuous physics\", Nature Communications, 2022, which views a ResNet as a discrete Forward Euler discretization of a ODE and shows similar performance to the \"continuous\" NeuralODE. Onken et al., \"Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows\", 2020 shows that it may be preferable to first discretize and then optimize.", "other_strengths_and_weaknesses": "## Weaknesses\n- The authors mention a clear limitation in the conclusion that their work does not dynamics from Adam optimizer, which is the most commonly used optimizer for training Transformer and hence limits the practical aspects of the work. It would be interesting for the authors to discuss ideas of how to support Adam as future work.\n### Minor\n- Typo in \"traing\" on line 065 left column for training.\n- \\citep is also incorrect on line 062\n- “overalapping” on line 102\n- Increase x and yaxis label size, legend and titles in Figure 1\n\n## Strengths\n- Literature is well-summarized and reviewed.\n- Section 2 provides a clear problem definition\n- The authors are very clear on their contribution.\n- Clear limitations and directions for future work are discussed in the conclusion.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. There is a clear connection to ODEs in the paper. Could the method also be connected to PDEs with a spatial component potentially in the image domain?\n2. The problem definition defines the problem for tabular data with respect to classical regression or classification tasks. How could this method be extended to time series data?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1UkuLnfDea", "reviewer": "Reviewer_iC26", "summary": "This paper studies gradient flow conservation laws in neural networks of many architectures including linear networks, ReLU networks, networks with attention layers, and networks with residual skip connections. Using ideas from Lie Algebra (specifically Frobenius' theorem), they are able to characterize the number of conservation laws and in some cases prove that their provided conservation laws are complete. For transformers such conservation laws take the form of functions of $Q Q^\\top - K K^\\top$ and $VV^\\top - OO^\\top$, similar to the conservation laws of weights of deep linear networks.  They also characterize gradient flow conservation laws for subsets of the parameters such as blocks in a deep residual network. Lastly, they provide an error analysis for discrete time step GD and gradient flow, showing an error in the conservation law which accumulates over time on the order of $\\text{learning rate}^2$.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for his positive comments and his insightful feedback.\n\n> **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?”\n\nThank 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**.\n\n> **Q2d** Some informal discussions after each result could make the paper more readable\n\nThank you for the suggestion that will be taken care of in the final version.\n\n> **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.\n\nBeside 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).\n\n> **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?\n\nYes, 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.\n\n> **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\n\nSince 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.\nHowever, 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.", "claims_and_evidence": "This paper is primarily theoretical and provides many proofs and references to the very relevant prior work of Marcotte et al 2023.  The claim of stability of the conservation law at finite learning rate is verified in an experiment in Figure 1, where for very small learning rate, the conservation law holds well and for larger learning rates the error accumulates.", "methods_and_evaluation_criteria": "The paper is primarily theoretical so there is not much to check. Perhaps more numerical experiments verifying different conservation laws in the transformer could be interesting as these results appear to be novel contributions of this work.", "theoretical_claims": "I checked and read Appendix A, B and C for the theoretical derivations. I also read some of Marcotte et al 2023 which contains some background material on this approach. Appendix B and C contains much of the technical machinery to prove properties and exhaustiveness of the conserved quantities. The result relies on the Frobenius theorem which is a fundamental tool in integrable systems theory. I also briefly skimmed Appendices E-I which derive the results for the attention layer which appeared correct.", "experimental_designs_or_analyses": "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.", "relation_to_broader_scientific_literature": "This paper provides a nice analytical toolkit to examine conservation laws in different neural network architectures. This is fundamental to deep learning theory since it aids our understanding of the implicit bias of gradient flows and the types of solutions networks converge to. While I was familiar with the two layer linear and ReLU conservation laws, I like that this approach is flexible enough to provide some insights into the residual networks and models with attention layers.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Another strength of the paper is its rigor and precision. However, the current writing could potentially be challenging to a machine learning theorist with less background in Lie Algebras and integrable systems. Some informal discussions after each result could make the paper more readable. 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.", "other_comments_or_suggestions": "See questions", "questions_for_authors": "1. 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_{\\ell} W_{\\ell}^\\top - W_{\\ell+1}^\\top W_{\\ell+1}$ for **all matrix entries**.  Is there a method to extract which nonlinear functions of this difference are conserved?\n2. What conserved quantity is being tracked in Figure 1 and the Appendix Figures?\n3. 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.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper extends an existing framework for studying conservation laws to key components of modern neural architectures, focusing in particular on convolutional ResNets and Transformers

The reviewers commended the paper for its thorough derivation of conservation laws in these settings, addressing an important gap in the literature, as well as for its detailed examination of how these theoretical principles manifest in practice, accounting for discrete effects and supported by empirical validation.
Decision: Accept (oral)