Parameter Symmetry and Noise Equilibrium of Stochastic Gradient Descent
Authors: Liu Ziyin
Mingze Wang
Hongchao Li
Lei Wu
Venue: NeurIPS 2024 poster (2024)
Keywords: SGD
fixed point
symmetry
Primary area: neuroscience_and_cognitive_science

Abstract: Symmetries are prevalent in deep learning and can significantly influence the learning dynamics of neural networks. In this paper, we examine how exponential symmetries -- a broad subclass of continuous symmetries present in the model architecture or loss function -- interplay with stochastic gradient descent (SGD). We first prove that gradient noise creates a systematic motion (a ``Noether flow") of the parameters $\theta$ along the degenerate direction to a unique initialization-independent fixed point $\theta^*$. These points are referred to as the noise equilibria because, at these points, noise contributions from different directions are balanced and aligned. Then, we show that the balance and alignment of gradient noise can serve as a novel alternative mechanism for explaining important phenomena such as progressive sharpening/flattening and representation formation within neural networks and have practical implications for understanding techniques like representation normalization and warmup.

Parameter Symmetry and Noise Equilibrium of Stochastic Gradient Descent
Liu Ziyin
Mingze Wang
Massachusetts Institute of Technology,
Peking University
NTT Research ziy y m1 t. edu
nigewangestu pku.edu.co
edu.
Hongchao Li
Lei Wu
The University of Tokyo
Peking University
1hclcat. phys. wt..okko..cc.jj5
Leivußmath. pku, aaaa.edu.comm
Abstract
Symmetries are prevalent in deep learning and can significantly influence the learning dynamics of neural networks. this paper, we examine how expo- nential symmetries model architecture or loss function interplay with stochastic gradient descent (SGD). We first prove that gradient noise creates a systematic motion (a "Noether flow") of the parameters along the degenerate direction to a unique initialization- independent fixed point 0". These points referred to the noise equilibria because, at these points, noise contributions from different directions are balanced and aligned. Then, show that the balance and alignment of gradient noise serve as a novel alternative mechanism for explaining important phenomena such progressive sharpening//lattennng and representation formation within neural networks and have practical implications for understanding techniques like repre- sentation normalization and warmup.
broad subclass of continuous symmetries present in the
Introduction
Stochastic gradient descent (SGD) and its variants have become the cornerstone algorithms used in deep learning. In the continuous-ting limit, the algorithm be written as [19, 13, 21, 32, 9]:
d0, = - V dt + (203I((,) dW,,
(1)
where (0) is the covariance matrix of gradient noise (Section 3) with the prefactor a² = n/ (2S) modeling the impact of finite learning rate 7) and batch size S; W, denotes the Brownian motion.
When Æ 0, Eq. (1) corresponds to gradient descent (GD)². However, SGD and GD can exhibit significantly different behaviors, often converging to solutions with significantly different levels of performance [31, 39, 44, 22, 49]. Notably, even when σ² < 1, where we expect close resemblance between SGD and GD over finite time [19], their long-time behaviors still differ substantially [26). These observations indicate that gradient noise can bias the dynamics significantly. and revealing its underlying mechanism thus crucial for understanding the disparities between SGD and GD.
Contribution. In this paper, study the how of SGD noise biases training through the lens of symmetry. Our key contributions are summarized as follows. We show that
1. when symmetry exists in the loss function, the dynamics of SGD can be precisely characterized and is different from GD along the degenerate direction;
38th Conference on Neural Information Processing Systems (NeurIPS 2024).
"Gradient descent" and "gradient flow" are used interchangeably work in the continuous time limit.

---

2. the treatment of common symmetries, including the rescaling and scaling symmetry in deep learning, can be unified in single theoretical framework that we call the exponential symmetry;
3. for any 0, every exponential symmetry implies the existence of a unique and attractive fixed point along the degenerate direction for SGD;
4. symmetry and balancing of noise can serve novel mechanisms for important deep learning phenomena such as progressive sharpening/lattennng and latent representation formation.
See Figure 1 for an illustration of how symmetry leads systematic flow of SGD. This work is organized as follows. We discuss the most relevant works in Section 2. The main theoretical results are presented in Section 4. We apply our theory to understand specific problems and present numerical results in Section 5. The last section concludes this work. All the proof' are presented in the Appendix.
θ ++
ve
2 Related Works
The dynamics of SGD in the degenerate directions of the loss landscape is a poorly understood problem. There are two closely related prior works. Ref. [48] studies the dynamics of SGD when there is a simple rescaling symmetry and applies it derive the stationary distribution of SGD for linear net- works. Our result is more general because rescaling symmetry is the simplest case of exponential symmetries Another re- lated work is Ref. [20]. which studies a different special case of exponential symmetry, the scale invariance, and in the pres- ence of weight decay. Their analysis assumes the existence of the fixed point of the dynamics, which we proved to exist. Also related is the study of conservation laws under gradient flow [34, 18, 24, 43, 40], which we will discuss more closely in Section 4. However, these works do not take the stochasticity of training into account. Comparing with these results that as- sume no stochasticity, our result could suggest that SGD con- verges to iiniipaaiionn-dddeeednentttt solutions, whereas the GD finds solutions are strongly inirialization-deeedde tion, Section extends our main result to discrete-time SGD.
An example of a 2d loss function with scale invariance: (()) ((A)) for scalar À and @ ε R². Be- cause of the symmetry, the gradient V must be tangential to the circles whose center is the origin. This implies that the norm |0] does not change during gradient flow training. However, when the training is stochastic or discrete- time, SGD must move outward. If the model starts at 9,, it must move to a larger circle. As an illustrative exam- ple, this loss function has unique and attractive fixed point: |0] will diverge after training under scale invariance. Also, see Remark 4.4 for discussion of the difference between discrete- time and continuous-tine dy- namics.
Figure I:
SGD
In addi-
3 Preliminaries
Setup and Notations. Let e Ω x Z + R denote the per-sample loss, with and Z denoting the parameter and sample space, respectively. Here, 2 é Z includes both the input and label and accordingly. We use E, E denote the expectation over a given training set. Therefore, L(0)= E. [_[ =)] is the empirical risk function. The covariance of gradient noise is given by
E(O) R (=_[V((())V(00..))))0--0) VVV
VV.L())V.L(()) to denote the covari- ance of gradient noise impacting on the subset of parameters u. Denote by (0,) the trajectory of SGD or GD. For any h : + I, we write h₁ = h(O.)) and h(0.)= = 4 (00.) for brevity. When the context is clear, we also use {(f) to denote {(0, 2).
Additionally, we use (+) :- (..
Symmetry. The par-sample loss E(.,-) is said to possess the Q-symmetry if E(0,2) (Q,, =), p TR,
where (Qo) is a set of continuous transformation parameterized by P E R. Without loss of generality, we assume Qo id. The most common symmetries exist within the model f. namely fa is invariant under certain transformations of 0. However, our formalism is slightly more general in the sense that it is also possible for the model to be variant while the per-sample loss remains unchanged, which appears in self-supervised learning [50], for example.
'These are known as "continuous symmetries." " Prior works also studied SGD training under discrete sym- metries [45. 3], which different from continuous symmetries.

---

4 Continuous Symmetry and Noise Equilibria
Taking the derivative with respect to p at = in Eq. (2), we have
= ())
where J(0)= dQ,(0)
dp p--0- Denote by C the antiderivative of J. that is, VC(0) = J().. Then, laking the expectation over 2 in (3) gives the following conservation law for GD solutions (0z)
C(0,) = 0.
(4)
Essentially, this is a consequence Noether's theorem [25], and will be called a "Noether charge" in analogy to theoretical physics. The conservation law (4) implies that the GD trajectory is con- strained on the manifold {0 ((0) C(6))). We refer to Ref. [18] for study of this type of conservation law under the Bregman Lagrangiam [16].
4.1 Noether Flow in Degenerate Directions
In this paper, we are interested in how C(0+) changes, if it changes at all, under SGD. By Ito's lemma, we have the following Noether flow (namely, the flow of the Noether charge):
C(0,) = of [E())
(5)
where V²C denotes the Hessian matrix of C. The derivation is deferred to Appendix B. By def- inition, E(O,) is always positive semidefinite (PSD). Thus, we immediately have: if V is PSD throughout training, C(,,)) is a monotonically increasing function of time. Conversely, if VăC is negative semidefinite (NPD), C(0,) is monotonically decreasing function of time.
The existence of symmetry implies that (with suitable conditions of smoothness) any solution resides within a connected, loss-invariant manifold, defined as Mg := (Q)(#) P e IR.. We term directions within this manifold as "degenerate directions" since movement along them does not change the loss value. Notably, the biased flow (5) suggests that SGD noise can drive SGD to explore within this manifold along these degenerate directions since the value C(0) for 0 Ma can vary.
4.2 Exponential symmetries
Now, let us focus on a family of symmetries that is common in deep learning. Since the correspond- ing conserved quantities are quadratic functions of the model parameters, we will refer this class of symmetries as exponential symmetries.
Definition 4.1. (Qp), is said to be exponential symmetry if J(0) (o)lo=0 A0 for symmetric matrix A.
This implies when p < 1, Q, -ddidd+ pA op). In the sequel, we also use the words " "'A-yymmerry" and "Q-symmetry" interchangeably since all properties of Q, we need can be derived from A. This definition applies to the following symmetries that are common in deep learning:
Rescaling symmetry: (a b) = (c(p+1),b/ppt 1)), which appears in linear and ReLU networks [7, 48). In this symmetry, A diag(Ta, -Io), where I is the identity matrix with dimensions matching that and b.
Scaling symmetry: 2 (p 1)8, which exists whenever part of the model normalized using techniques like batch normalization [14], layer normalization [2], or weight normalization [29].
In this case, A
Double rotation symmetry: This symmetry appears when parts of the model involve a matrix factorization problem, where for arbitrary invertible matrix B f ((UW) (UBBB-IW) Writing the exponential symmetry for this case is a little cumbersome. We need first to view U and W as a single vector, and the exponential transformation given by block-wise diagonal matrix diag( B, B.... B, B-1 "..,,,,,",,,.,, B-1). ). See Section 5.1 for more detail.
possible for only a subset of parameters to have a given symmetry. Mathematically, this cor- responds to the case when A is low-rank. It is also common for & to have multiple exponential symmetries at once, often for different (but not necessarily disjoint) subsets of parameters. For example, ReL.U network has a different rescaling symmetry for every hidden neuron.

---

It is obvious that under this Q symmetry, the Noether charge has a simple quadratic form:
C(0) A0.
(6)
Moreover, the interplay between this symmetry and weight decay can be explicitly characterized in our framework. To this end, we need the following definition.
Definition 4.2. For any 7 E I. we say (0, x) {(0,x)) >110= has the Q symmetry long as {(0,z) x)) has the Q symmetry.
For the SGD dynamics that minimizes (0) E., (0) follows from (5) that =
Thus, positive y always causes |C(6t)] to decay, and the influence of symmetry is determined by the spectrum of A. Denote by A Σ, my the eigendecompositiin of A. Then,
Tr{E(0+)A] Σ p.m. ((
This gives clear interpretation of the interplay between SGD noise and the exponential symmetry: the noise along the positive directions of A causes C(0)) to grow, while the noise along the negative directions causes C(6,) to decay. In other words, the noise-induced dynamics of C(O,) is deter- mined by the competition between the noise along the positive- and negative-cigenvalue directions of A.
Time Scales. The above analysis implies that the dynamics of SGD can be decomposed into two parts: the dynamics that directly reduce loss, and the dynamics along the degenerate direction of the loss, which is governed by Eq (5). These two dynamics have essentially independent time scales. The first part is independent of the o² in expectation, whereas the time scale of the dynamics in the degenerate directions depends linearly on o².
The first time scale ferm is due to the dynamics of empirical risk minimization. The second time scale tequi is the time scale for Eq. (5) to reach equilibrium, which is irrelevant to direct risk minimization. When the parameters are properly tuned, ferm is of order 1, whereas tequi proportional to o² (2S).. Therefore, when 2 is large, the parameters will stay close to the equilibrium point early in the training, and one can expect that C(0.) is approximately zero after tequi. In line with Ref. [20], this can be called the fast-equilibrium phase of learning. Likewise, when o² 1, the approach to equilibrium will be slower than the actual time scale of risk minimization, and the dynamics in the degenerate direction only take off when the model has reached local minimum. This can be called the alow-equilibrium phase of learning.
4.3 Noise Equilibrium and Fixed Point Theorem
important and practically relevant to study the stationary points of dynamics in Eq. (7). For- mally, the stationary point is reached when (>C(()++T=(((())) no assumption about (()) and E()), one might feel that is impossible to guarantee the existence of a fixed point. Remarkably, we prove below that a fixed point exists and is unique for every connected degenerate manifold.
0. Because we make essentially
To start, consider the exponential maps generated by A:
AA p0
lin + PA o(p))
which applics the symmetry transformation to for A/p times. Then, it follows that if we apply Qp transformation to 0 infinitely many times and for perturbatively small p.
((0) = E (c)A)).
(8)
Thus, the exponential symmetry implies the symmetry with respect to an exponential map. fun- damental element of Lie groups [11]. Note that exponential -app symmetry is also an exponential symmetry by definition. For the exponential map, the degenerate direction is clear: for any X, o connects to e AA0 without any loss function barrier. Therefore, the degenerate direction for any ex- ponential symmetry is unbounded. Now, we prove the following fixed point theorem, which shows that for every exponential symmetry and every there one and only one corresponding fixed point in the degenerate direction.

---

Theorem 4.3. Let the per-sample loss satisfy the A-exponential symmetry and 0x == exp[AA]e.
Then, for any 9 and any y 0,2
(1) G(Ox) (Eq. (7)) and -C(0x) are monotonically decreasing functions of X;
(2) there exists a X E IR u } such that G(Ox,) )= = =0;
(3) in addition, if G(xx) + 0, X is unique and G())) is strictly monotonic:
(4) in addition to (3), if (() is differentiable, X (f) is a differentiable function of B.
Remark 4.4. It is now worthwhile to differentiate gradient flow (GF), GD, SGD, and stochastic gradient flow (SGF). Technically, one can prove that the same result holds for discrete-time GD and SGD in expectation, and GF is the only of the four algorithms that do not obey this theorem (See Section D), and so one could argue that the discrete step size is the essential cause of noise balance. Mathematically, the SGF can be seen size and thus also share this effect (remember that the Ito Lemma contains a second-order term in d0). That being said, there is a practical caveat: in practice, reach these fixed points with SGD than with GD, and so it is fair to say that this effect is the most dominant when gradient noise is present.
model of the leading order effect of having a finite step
find it much easier for models to
Part (I), together with Part (2), implies that the unique stationary point is essentially attractive. This is because C' decreases with À while C' increases with it. Let C* = C(Ox-) Thus, C(0)- C* always have the opposite sign of X", while a C(0) will have the same sign. Conceptually, this means that C' will always move to reduce its distance to C*. Assuming that C* is a constant in time (or close to constant, which is often the case at the end of training), Part (I) implies that d (C(0) C*) oc -Cg)(C( C*). signaling a convergence to C(f) = =C". In other words, SGD will move to restore the balance behaved, one can indeed establish the convergence to the fixed point in the relative distance even if C* mildly divergent due to diffusion. Because this part strongly technical and our focus is on the fixed points, we leave the formal statement and its discussion to Appendix B.3.
dt
perturbed away from X = 0. If the matrix ΣA well-
Theorem 4.5. (Informal) Let C* follow a drifted Brownian motion and E.A satisfy two well-behaved conditions. Then, either C C + 0 in L22 or (C-C*) + 0 in probability.
Parts (2) and (3) show that unique fixed point exists. We note that it is more common than not for the conditions of uniqueness to hold because there is generally no reason for Tr Tr[E(0) A] or TH[00TA) A] vanish simultaneously, except in some very restrictive subspaces. One major (perhaps the only) reason for the first trace to vanish is when the model is located at an interpolation mini- mum. However, interpolation minima are irrelevant for modern large-scale problems such as large language models because the amount of available text for training far exceeds the size of the largest models. Even when the interpolation minimum exists, the unique fixed point should still exist when the training is not complete. See Figure 1. Part (4) means that the fixed points of the dynamics is behaveee If the parameter 0 has a small fluctuation around given location, C' will also have a small fluctuation around the fixed point solution. This justifies approximating C' by constant value when ! changes slowly and with small fluctuation.
Fixed point as Noise Equilibrium.
Let 0" be a fixed point of (7). It must satisfy ("Tr))))(T(((")))))
(9) Hence, a large weight decay leads to small IC(O*)|, whereas large gradient noise leads to large IC(0*)). When there is no weight decay, we get a different equilibrium condition: Tr[E(0") ) = O. which can be finite only when A contains both positive and negative eigenvalues. This equilibrium condition is equivalent to Σen.s0 [n,, E(0")n, = -Egr, 11,n,(((' )n,. Namely, the overall gradient fluctuation in the two different subspaces specified by the symmetry A must balance. We will see that the main implication of this result is that the gradient noise between different layers of deep neural network should be balanced at the end of training. Conceptually, Theorem 4.3 suggests the existence of a special type of fixed point for SGD, which the following definition formalizes.
Definition 4.6. f is noise equilibrium for nonconstant function C(0) if Č(0)= = 0 under SGD.
5 Applications
Now, we analyze the noise equilibria of a few important problems. These examples are prototypes of what appears frequently in deep learning practice and substantiate our arguments with numerical
'A similar result can be proved for the tiscrete-time SGD. See Section D.
5

---

10
theary
[UT.U- WT
2000
20
A12
GO
A14
1500
SGD
615
10:
1000
-20
A17
500
-40
A19
10a
1000
2000
1000 2000
3000
x
3xx10 4 10
[WTW
Figure 2: Comparison between GD and SGD for matrix factorizations. Left: Example of learning trajectory. The convergence speed is almost exponential like in experiments. Mid: evolution of 10 individual elements of Ag := (U T'U WT' w As the theory shows, they all move close zero and fluctuate with small variance. Right: Converged solutions SGD agree with the prediction Theorem 5.2, but are an order of magnitude away from the solution found by GD, even they start from the same init.
examples. In addition, an experiment with the scale invariance in normalized tanh networks is presented in Appendix A.1.
5.1 Generalized Matrix Factorization
Exponential symmetry is also observed when the model involves a (generalized) matrix factoriza- tion. This occurs standard matrix completion problems [33] or within the self-attention of trans- formers through the key and query matrices [35]. For a (generalized) matrix factorization problem, we have the following symmetry in the objective:
E(U, W, -
(10) parameters O', We consider matrices A that
for any invertible matrix A and symmetry-irrelevaae close to identity: = + pB + O("²), and A = 1 -ppp-BB-- + O(p²). Therefore, for an arbitrary [UUBU"]= -Tr[W"] BW]. This
symmetric B, have conserved quantity for GD: Cn(0) conservation law can also be written the matrix form, which a well-known result for GD [8, 24]:
WW.WT- U, U.) (W.W. Ug).).
(11)
For SGD, applying (5) gives the following proposition.
Proposition 5.1. Suppose the symmetry (10) holds. Let U = (ūI.-.;ūa)) e Rdadd W (wi,-., wd,) E Rdxdo where üi úy E IR. Let (0) TU -T[W"" for any symmet- ric matrix B E pdddd Then, for SGD, we have
de Tr (0.) Σ Tr[Ea, (8)) J=1
ČB(0,)= = 2² Σ 1=1
This dynamics is analytically solvable when lxd and W Rdx1, this case, taking B = Ex, + E,, where E,, denotes the matrix with entries of 1 at (i.j) and zeros elsewhere. For this choice of B, we obtain that (f) = W UgU, and applying the results we have derived, it easy to show that for some random variable r: Cn((,) = [48]. Since the choice of B is arbitrary, we have that
which signals an exponential
decay. For common problems, Var[r(0,)] W&W U2Ui for all & and [. The message is rather striking: SGD automatically converges to solution where all neurons output the same sign (sgn(U,) egn(U,)) at exponential rate.
5.2 Balance and Stability of Matrix Factorization
As concrete example, let us consider two-layer linear network (this can also be seen as a variant of standard matrix factorizations):
4 UWW JUl
(12)
where x e Rd is the input data, and y c Rdy is a noisy version of the label. The ground truth mapping is linear and realizable: y' U="= The second moments of the input and noise are denoted as Σ,, B[rx"] and Σ, = E[ec"], respectively. Note that this problem essentially identical to matrix factorization problem, which is not only theoretical model of neural networks but also an important algorithm frequently in use for recommender systems [41]. The following theorem gives the fixed point of Noether flow.

---

theory SGD WD
95.0
1.4
GD+WD
07.5
90.0
theory init.
07.5
1.0
85.0
0.8
0.6
0.8
1.0
1.2
Figure 3: A two-layer linear network after training. Here, the problem setting the same as Figure The theoretical prediction is computed from Theorem S.2. Left: balance of the norm is only achieved when da 1, namely, when the data has an isotropic covariance. We also test SGD with a small weight decay (10"), "), which is sufficiently small that the solution we obtained for SGD without SGD still holds approximately. In contrast, training with GD + WD always converges to norm-balanced solution. Right: the sharpness of the converged model trained with SGD. We see that for some data distributions, SGD converges to a sharper solution, whereas it converges to flatter solutions for other data distributions. This flattening and sharpening effect are both due to the bise-balance effect of SGD. Here, we find that the systematic error between experiment and theory is due the of a finite learning rate and decreases we decrease 1).
Theorem 5.2. Let r = UWr x y be the prediction residual. For all symmetric B, CB = Of WTwW" TUU
(13)
where Tw [[[r||3xx"] 2,,, To -[|=]|/rr*]++++ 2
See Figure for the convergence of SGD this solution under different learning rate, batch size and width. The equilibrium condition takes a more suggestive form when the model is at the global y= = Assuming that and x are independent and that there is weight minimum, where U"***== x-y==== decay, we have:
W W U.E
(14)
Here, the bar over the matrices indicates that they have been normalized by their traces: Σ 2/Tr[E]. The matrices I'w and Io simplifies because the global minimum, т; E, and so [[[ = Tr[E,] [E, and E[ [[[r|]] "=xx]== Tr[E,]]x. This condition should be compared with the alignment condition for GD in Eq. (11), where the alignment is entirely determined by the initialization and perfect alignment is achieved only if the initialization is perfectly aligned. This condition simplifies further if both E, and Σ, are isotropic, where the equation simplifies to WWT = Udy.. Namely, the two layers will be perfectly aligned. and the overall balance depends only on input and output dimensions. Figure 2-Left shows an experiment that shows that the two-layer linear net is perfectly aligned after training. Here, every point corresponds to the converged solution of an independent run with the same initialization and training procedures but different values of Σ., In agreement with the theory, the two layers are aligned according to The- orem 5.2 under SGD, but not under GD. In fact, GD finds solutions that are more than an order of magnitude away from SGD.
Noise Driven Progressive Sharpening and Flattening. This result implies a previously unknown mechanism of progressive sharpening and flattening, where, during training, the stability of the algorithm steadily improves (during flattening) or deteriorates (during sharpening) [39, 15, 6]. To see this, we first derive metric of sharpness for this model.
Proposition 5.3. For the per-sample loss (12), les S(0) TH(VL((O)) Then, S(0) dy 1/2 + T|
The trace of the Hessian is good metric of the local stability of the GD and SGD algorithm be- cause the trace upper bounds the largest Hessian eigenvalue. Let us analyze the simplest case of an autoencoding task, where the model is at the global minimum. Here, Ex Ldz. Σ. la,- For random Gaussian initialization with variance ow and où the trace at initialization is, in expec- tation, Sinit d,,'Tit d'Tr[ [x,, + of). At the end of the training, the model is close to the global minimum and satisfies Proposition 5.3. Here, the rank of U and W matters and is upper bounded by min (d, d, and at the global minimum, U and W are full-rank (equal to min dd, d, )), and all the singular values are 1. Thus,
Smit d,d(o + ow Tr|[
(15)
Sond 2 22 (d, d.. Tr[2..].
7

---

The change in the sharpness during training thus depends crucially on the initialization scheme. For Xavier init, of (d, + d)-1 and oi (d + dz)= and so Smitt Sand (but Siit slightly smaller). Thus, for the Xavier init., the sharpness of loss experiences a small sharpening during training. For Kaiming init. of = and of d=¹ Therefore, always holds that Sint 2 Sand. and so the stability improves as the training proceeds. The only case when the Kaiming init. does not experience progressive flattening is when d dz = dy, which agrees with the common observation that training is easier if the widths of the model are balanced [12]. See Figure 4 for an experiment. In previous works, the progressive sharpening happens when the model is trained with GD [6]; our theory suggests an alternative mechanism for it.
A practical technique that the theory explains is using warmup to stabilize training in the early stage. This tech-
2000
1750
SGD Kalming
waamup
nique was first proposed in Ref. [10] for training CNNs, where it was observed that the training is divergent if we start the training at a fixed large learning rate Thaxx- How- ever, this divergent behavior disappears if we perform a warmup training, where the learning rate is increased gradually from minimal value 7/max. Later, the same technique is found to be crucially useful for training large language models [27]. Our theory shows that the gradient noise can drive Kaiming init. to a stabler status where a larger learning can be applicd.
SGD Kaiming
1500
SGD Xevier
1250
Kaiming
1000
750
500
250
2000
4000
6000
8000
10000
Figure 4: Dynamics of the stability condi- tion S during the training of rank-1 ma- trix factorization problem. The solid lines show the training of SGD with Kaiming init. When the learning rate (1) 0.008) is too large, SGD diverges (orange line). However, when one starts training at a small learning rate (0.001) and increases 1) to 0.008 after 5000 iterations, the training remains stable, This is hecause SGD training improves the stability condition during training, which is in agreement with the theory. In contrast, the stability condition of GD and that of SGD with Xavier init increases only slightly. Also, note that both Xavier and Kaiming init. under SGD converges to the same sta- bility condition because the equilibrium is unique.
Flat or Sharp? Prior works have often argued that SGD prefers flatter solutions to sharper ones (e.g., see Ref. [42]). The exact solution we found, however, implies subtle picture: for some datasets, SGD prefers sharper solutions, while for others, SGD prefers flatter solutions. Therefore, there is no causal relationship between SGD training and the sharpness of the found solution. See Fig- ure 3 for the dependence of the flatness on the data distri- bution. A related question is whether SGD noise creates similar effect weight decay training. The answer is also negative: weight decay always prefers smaller norms and, thus, norm- warm-balanced solutions, which are not neces- sarily noise-aligned solutions. Figure 3 shows that SGD can also lead to unbalanced solutions, unlike weight de- cay.
5.3 Noise-Aligned Solution of Deep Linear Networks
Here, we apply our result to derive the exact solution of an arbitrarily deep and wide deep linear network, which has been under extensive study due to its connection in loss landscape to deep neural networks [4, 17, 23, 47, 38]. Deep linear networks have also been major model for understanding the implicit bias of GD [1]. The per-sample loss for a deep linear network can be written as:
(16) where W; is an arbitrary dimensional matrix for all i. The global minimum is realizable: y= Vx+c, for i.i.d. noises e. Because there is a double rotation symmetry between every two neighboring matrices, the Noether charge can be defined with respect to every such pair of matrices. Let B3 be symmetric matrix; we define the charges to be CB. W! B,W, The noise equilibrium solution given by the following theorem.
{(0) = IW y11
Theorem 5.4. Let Wp...W; = V. Let = VV.VVE, such that V LS'R is its SVD and d = rank(V')) (V). Then, for all and all B,, noise equilibrium for CB, at the global minimum is
VE.WD= = LEDUB- D-1: W. = U,E,UL,1 W WIVE, U.E,R,
(17)
for = 2.....D-11. 1. U₁ are arbitrary matrices satisfying U U, = Laxdi and Dn are diagonal matrices such that
d
(D-2)/2D
TriS' 1/D
Σn ΣD
S ΣS
1axd.
(18)
T
d
8

---

input
input
output
output
intermediate
intermediato
3-
2000 4000
6000
2000
4000
6000
iteration
iteration
Figure 5: Norms of weights of multilayer deep linear network during training MNIST without weight decay. We see that the intermediate layers converge to the layers are different because they are determined by the input and output noise. This effect is robust against different initializations. This agrees with our analysis for deep linear nets (Theorem 5.4). Left: initializing all layers with the same norm. Right: initializing all layers at randomly different norms.
norm during training, whereas the input and output
This solution has quite a few striking features. Surprisingly, the norms of all intermediate layers are balanced:
Tr (
(19)
All intermediate layers are thus rescaled orthogonal matrices aligned with the neighboring matrices and the only two matrices that process information are the first and the last layer. See Figure for an illustration of this effect. This explains an experimental result first observed in Ref. [30], where the authors showed that the neural networks find similar solutions when the model is initialized with the standard init., where there is no alignment at the start, and with the aligned init. Thus, the balance and alignment between different layers in the neural networks can be attributed to the rescaling symmetry between each pair of matrices.
5.4 Approximate Symmetry and Bias of SGD
Lastly, let us consider what happens if the loss function only has an approximate symmetry. As a minimal model, let us consider the following loss function: = & (0.x) + (l2(0). Here, & has the A-symmetry, whereas (0) has no symmetry nor randomness and so & does not affect Σ all. Ç determines the relative strength between the two terms. In totality, l no longer has the A-symmetry.
As before, let CA = A0. Then, CA( = =-((v()) AO" + +o TTTr(E())))], whose fixed point is <(V(2) A0" [O)))
(20)
This equilibrium condition thus depends strongly on how large is. When ( is small, we see that SGD still favors the fixed point given by Theorem 4.3, but with first-order correction in Ç.
Conversely, if Ç is large and o² is small, we can expand around local minimum of the loss function 8", and so the fixed point becomes
<(1 0') H(O') A0 = TT
0- 0-
(21)
where H is the Hessian of fz. Certainly, this implies that SGD will stay around point that deviates from the local minimum by an O(a²) amount. This stationary point potentially has many solutions. For example, class of solution is when 0 0" is eigenvector of H with eigenvalue h- 0 and eigenvector n, we can denote = (0 - 8*) and obtain a direct solution of s:
o (T((E(('))
(22)
Ch AO
This deviation disappears in the limit 0² + 0. Therefore, this implicit regularization effect is only consequence of SGD training and is not present under GD. With this condition, can obtain clear expression of the deviation of the quantity C from its local minimum value C" := (0*) A0*. We have that
C(0)=C"+2(0-01))
+ T 101000000000000 = C* + THEA Çh+
(23)
Thus, our results in the previous section still apply. The quantity C' will be systematically larger than
the local minimum values of if the approximate symmetry matrix is PD. It is systematically smaller if A is ND. When A contains both positive and negative eigenvalues, the deviation of C

---

33
35
30
25
25
20
20
20
15-
15
15-
15
10
5
10
20
20
inturion Indexx
noturon Indixx
nduron indix
nsturnn Endlex
The latent representations of two-layer tanh net trained under SGD (left) are similar acToss different layers, in agreement with the theory. However, the learned representations are dissimilar under GD (right). Here, plot the matrices W.:WW W (first and third plots) and US.U (second and fourth plots). Note that the quantity W.., W is equal to the covariance of the preactivation representation of first layer. This means that SGD and GD learn qualitatively different features after training. Also, see Appendix A.4 for other activations. This mechanism also complements the recent result in Ref. [46], which proposes inysscs-inspired theory showing that gradient noise is key factor in determining the latent representation of neural networks.
Figure 6:
depends on the local gradient fluctuation balancing condition. When the smallest eigenvalue of H is close zero (which true for common neural networks), the dominant factor that biases C occurs in this space. Therefore, it is not bad to approximate the deviation as C(0) C* +ro"Tr[EA]] //ainmiin where min is the smallest eigenvalue of the Hessian at the local minimum. In reality, Ç Çis neither too large nor too small, and one expects that the solution favored by SGD is an effective interpolation between the true local minimum and the fixed point favored by the symmetries.
set of experiments is shown in Figure 6, where we compare the latent representation of a two- layer tanh net with the prediction of 5.2. This is natural example because fully connected networks are believed to be approximated by deep linear networks because they have the same connectivity patterns. We thus compare the prediction Theorem 5.2 with the experimental results of nonlinear networks. Here, the task is a simple autoencoding task, where E R240 and y + e. x is sampled from isotropic Gaussian, and is an independent non-isotropic (but diagonal) Gaussian noise such that Var[e]] = 5 and Var[e] ] = 1 for i * 1. We train with SGD or GD for 10 iterations. The experimental results show that if trained with SGD, the learned representation agrees with the prediction of Theorem 5.2 well, whereas under GD, the model learned a completely different rep- resentation. This suggests that our result may be greatly useful for understanding the structures of latent representations of trained neural networks because the quantity W W has a clean interpre- tation as the normalized covariance matrix of pre-activation hidden representation. Also, this result is not special feature of tanh networks. Appendix A.4 also shows that the same phenomenon can be observed for swish [28], ReLU, and leaky-ReLU nets.
6 Conclusion
In this work, have studied how continuous symmetries affect the learning dynamics and fixed points of SGD. The result implics that SGD converges to crtializationinnepeeeede solutions at the end of training, in sharp contrast to GD, which converges to strongly initialization-deeende solutions. We constructed the theoretical framework of exponential symmetries to study the special tendency of SGD to stay close to special fixed point along the constant directions of the loss landscape. We proved that every exponential symmetry leads to a mapping of every parameter to unique and essentially attractive fixed point. This point also has clean interpretation: is the point where the gradient noises of SGD in different subspaces balance and align. Because of this property, termed these fixed points the "noise equilibria." The advantage of result is that it only relies on the existence of symmetries and is independent of the particular definitions of model architecture or data distribution. A limitation of our work is that we only focus on the problems that exponential symmetries can describe. It would be important to extend the result to other types of symmetries in the future. Another interesting future direction is study these noise equilibria of more advanced models, which may deepen both our understanding of deep learning and neuroscience.
Acknowledgement
Lei Wu is supported by the National Key R&D Program of China (No. 2022YFA 1008200) and Na- tional Natural Science Foundation of China (No. 22881011 101). Hongchao Li is supported by Forefront Physics and Mathematics Program to Drive Transformation (FoPM), a World-leading Innovative Graduate Study (WINGS) Program, the University of Tokyo. Mingze Wang is supported in part by
10

---

the National Key Basic Research Program of China (No. 2015CB856000), We thank anonymous reviewers for their valuable comments.
References
[1] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1oX1v:1607.06450, 2016.
[3] Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks. arXiv preprint arXiv:2306.044 04251, 2023.
[4] Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun. The loss surfaces of multilayer networks. In Artificial Intelligence and Statistics, pages 192- 204, 2015.
[5] Anna Choromanska, Yann LeCun, and Gérard Ben Arous. Open problem: The landscape of the loss surfaces of multilayer networks. In Conference on Learning Theory, pages 1756-1760. PMIR, 2015.
[6] Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gra- dient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103. 00065, 2021.
[7] L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp Minima Can Generalize For Deep Nets. ArXiv e-prints, March 2017.
[8] Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homoge- neous models: Layers are automatically balanced. Advances in neural information processing systems, 31, 2018.
[9] Xavier Fontaine, Valentin De Bortoli, and Alain Durmus. Convergence rates and approxi- mation results for sgd and its continuous-timg time counterpart. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 1965-2058. PMLR, 15-19 Aug 2021.
[10] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: I 7766.02677 2017.
[11] Brian C Hall and Brian C Hall. Lie groups, Lie algebras, and representations. Springer, 2013.
[12] Boris Hanin and David Rolnick. How to start training: The effect of initialization and archi- tecture. Advances Neural Information Processing Systems, 31, 2018.
[13] Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of nonconvex stochastic gradient descent. arXiv preprint arXiv:1705.0756: 2017.
[14] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167. 2015
[15] Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2019.
[16] Michael Jordan. Dynamical, symplectic and stochastic perspectives on gradient- gradient-based based op- timization. In Proceedings of the International Congress of Mathematicians: Rio de Janeiro 2018, pages 523-549. World Scientific, 2018
[17] Kenji Kawaguchi. Deep learning without poor local minima. Advances Neural Information P Processing Systems, 29:586-594, 2016.

---

[18] Daniel Kunin, Javier Breastuy-Brens Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanies: Symmetry and broken conservation laws deep learning dy- namics. arXiv preprint arXiv:2012.04728, 2020.
[19] Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40):1-47 2019.
[20] Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with tra- ditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33:14544-14555, 2020.
[21] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochas- tic differential equations (sdes), 2021.
[22] Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and fluctuation of finite learning rate stochastic gradient descent, 2021.
[23] Haihao Lu and Kenji Kawaguchi. Depth creates no bad local minima. arXiv preprint arXiv: 085888 2017.
[24] Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré. Abide by the law and follow the flow: Conservation laws for gradient flows. 2023.
[25] Emmy Noether. Invariante variationssrobllmee. Königlich Gesellschaft der Wissenschaften Göutingen Nachrichten Mathematii- physiik Klasse, 2:235-267, 1918.
[26] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for di- agonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, 11-29218-29236 2021.
[27] Martin Popel and Ondfej Bojar. Training tips for the transformer model. arXiv preprint arXiv: 18X4v:180. 2018.
[28] Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions, 2017.
[29] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to ac- celerate training of deep neural networks. Advances in neural information processing systems, 29, 2016.
[30] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning deep linear neural networks. arXiv preprint 132vv1312.66120, 2013.
[31] N. Shirish Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On Large- Batch Training for Deep Learning: Generalization Gap and Sharp Minima. ArXiv e-prints, September 2016.
[32] Justin Sirignano and Konstantinos Spiliopoulos. Stochastic gradient descent in continuous time: A central limit theorem. Stochastic Systems, 144-1112 2020.
[33] Nathan Srebro, Jason Rennie, and Tommi Jaakkola. Maximum-margin matrix factorization. Advances in neural information processing systems, 17, 2004.
[34] Hidenori Tanaka and Daniel Kunin. Noether's learning dynamics: Role of symmetry breaking in neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wort- man Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 25646-25660. Curran Associates, Inc., 2021.
[35] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[36] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learn- ing dynamics of normalized neural network using sgd and weight decay. Advances Neural Information Processing Systems, 34:6380-6391, 2021.
12

---

[37] Ming Chen Wang and George Eugene Uhlenbeck. On the theory of the brownian motion ii. Reviews of modern physics, 17(2-3):323, 1945.
[38] Zihao Wang and Liu Ziyin. Posterior collapse of a linear latent variable model. Advances in Neural Information Processing Systems, 35:37537-37548, 2022.
[39] Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Sys- tems, 31, 2018.
[40] Yizhou Xu and Liu Ziyin. When does feature learning happen? perspective from an analyti- cally solvable model. arXiv preprint arXiv:2401.07088 2024.
[41] Hong-Jian Xue. Xinyu Dai, Jianbing Zhang, Shujian Huang, and Jiajun Chen. Deep matrix factorization models for recommender systems. In IJCAI, volume 17, pages 3203-3209. Mel- bourne, Australia, 2017.
[42] Ning Yang, Chao Tang, and Yuhai Tu. Stochastic gradient descent introduces an effec- tive landscape-dependen regularization favoring flat solutions. Physical Review Letters, 130(23)223701, 2023.
[43] Bo Zhao, Nima Dehmamy, Robin Walters, and Rose Yu. Symmetry teleportation for accel- erated optimization. Advances in Neural Information Processing Systems, 35:16679-16690, 2022.
[44] Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint 18Xiv: 00195, 2018
[45] Liu Ziyin. Symmetry induces structure and constraint of learning. Forty-first International Conference Machine Learning.
[46] Liu Ziyin, Isaac Chuang, Tomer Galanti, and Tomaso Poggio. Formation of representations in neural networks. arXiv preprint MxXiv:2410.03006 2024.
[47] Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. In Advances in Neural Information Processing Systems, 2022.
[48] Liu Ziyin, Hongchao Li, and Masahito Ueda. Law of balance and stationary distribution of stochastic gradient descent. arXiv preprint arXiv:2308.066. 000000006661... 2023.
[49] Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in SGD. In International Conference on Learning Representations, 2022.
[50] Liu Ziyin, Ekdeep Singh Lubana, Masahito Ueda, and Hidenori Tanaka. What shapes the loss landscape of self supervised learning? In The Eleventh International Conference on Learning Representations, 2023.
13

---

16
16
Met
Net
14
Nett
14
Net-2
12-
12
10-
8-
6-
0
1000
2000
3000
4000
5000
1000 2000 3000
4000
5000
1000 2000
3000
4000
5000
Figure 7: When there is scaling symmetry, the norm of the parameters increases monotonically under SGD but remains unchanged under GD. Left: evolution of the total model norm for two-layer nonlinear networks where there is a rescaling symmetry. Mid: evolution of the second layer. Right: evolution of the first layer. This shows that the evolution of cach layer can be vastly different, but the total norm of the parameters with the scaling symmetry is always monotonically increasing. Also, note that for net B, each layer also has the rescaling symmetry, and the norm of each layer for net-B is also increasing. In contrast, net-A does not have layer-wise symmetry, and the individual norms be either increasing or decreasing.
A
Additional Experiments
A.1 Scale Invariance
The scale invariance appears when common normalization techniques such batch normalization [14] and layer normalization [2] are used, Let ((0,z) denote a per-sample loss such that for any P e IR: x(),z) = xp0,x), where e Rd. For this symmetry, A = 1. Thus, by Eq. (7), we have during SGD training that
10,11 = --1001, -
(24)
dt
Thus, without weight decay, the parameter norm increases monotonically and even diverges, partic- ularly for ander-parameterize models where the gradient noise is typically non-degenerate.
Here, numerically compare two networks trained
GD and SGD: Net-A: f(x) dhe tanh(w].xr) w/ Here, w and u are ma- Σj [w tanh(w tanh(w]a//lul]> [u][F; and Net-B: f(x) = E; |ul
trices and wy denotes the j-th row of w and My denotes the j-th column of u. The two networks and w. However, both networks have the global scale invariance: if we scale both U and W by arbitrary positive scalar p. the network output and loss function remain unchanged for any sample x. We train these two networks on simple linear Gaussian data with GD or SGD. Figure 7 shows the result. Clearly, for SGD, both networks have a monotonically increas- ing norm, whereas the norm remains unchanged when the training proceeds with GD. What's more, Net-B has two additional layer-wise scale invariances where one can scale only (or only w) by while keeping the loss function unchanged. This means that both layers will have a monotonically increasing norm, which is not the case for Net-A.
are different functions
Recent works have studied the dynamics of SGD under the scale-invariant models when weight decay is present [36, 20). Our result shows that the model parameters will diverge without weight decay, leading to potential numerical problems. Combining the two results, the importance of having weight decay becomes clear: it prevents the divergence of models.
A.2 Experiment Detail for Alignment Dynamics Matrix Factorization
Here, we give the details for the experiment in Figure 2. We train two-layer linear with do = d2 = 30 and = 40. The input data is x ~N 1), and y = + where e is i.i.d. Gaussian with unit variance. At the end of SGD training, every element of the matrix U is close that of WTw WT, and they are therefore very well aligned. Such phenomenon does not happen for GD.
A.3
Convergence of Matrix Factorization the Noise Equilibrium
See Figure 8.
14

---

125
125
n= 0.04
d-500
5
100
WM WT.WT
d-700
WTru WT wwww'
= 0.1
0.08
75
d-900
d-1100
5 140
50
d-1300
50
11 0.14
d=1500
25
d=1700
dm 1900
200
600
1000
200
400
600
800
1000
400 800
Figure 8: The convergence matrix factorization to the noise equilibria robust against different hyperpa- rameter settings. The task is autoencoding task where y
10000 The distribution of is controlled by a parameter dz: #1:50 ~ (0, Φx). 251-100 ~ (0, dz). This directly controls overall covariance of x. The output noise covariance is set to be identity. Unless it is the independent variable, 1), and are set to be 0.1, 100 and 2000, respectively. Left: using different learning rates. Mid: different data dimension: d, dy d. Right: different batch size S.
A.4 Alignment in Nonlinear Networks
Here, we complement the experiment in the main text with other types of activations. The experi- mental setting exactly the same except that we switch the activation to swish, ReLU, and Leaky- ReLU. This shows that the prediction of Proposition 13 may have surprisingly wide applicabilityy See Figure 9.
15

---

35
35
30
30
25
25
20
20
15
15
10
10
5
5
0-
0
10
20
30
neuron index
neuron index
35
35
30
30
25
25
20
20
15
10
5
5
0
10
20
30
10
20
30
neuron index
neuron index
5
10
20
30
10
20
30
neuron index
neuron index
Figure 9: Activation patterns of nonlinear networks trained with SGD (Upper to lower: ReLU, leaky-ReLU, swish). Left: WTwW' WT Right: U UTyUU The similarity between the matrices is striking.
16

---

B Proofs
B.1 Ito's Lemma and Derivation of Eq. (5)
Let a vector X, follow the following stochastic process:
dX, = 14 dt + G,WW;
(25)
for matrix G.. Then, the dynamics of any function of X₁ can be written as (Ito's Lemma)
af(x,) = V " ((x))) dt V.(X.''((")).)v d,W,
(26)
Applying this result to quantity C(0) under the SGD dynamics, we obtain that
x² of (E di V [
dC v CVV+
(27)
where we have used
V7, G. - ( By Eq. (3), we have that
V O,
(28)
and
v C'CEEE H CV 4] =0. O.
(29)
Because (0) and v (1) share eigenvectors, we have that
V'CVo22(0)00.
(30)
Therefore, we have derived:
d² dC
2 T[
(31)
B.2 Proof of Theorem 4.3
We first prove a lemma that links the gradient covariance the gradient covariance at 0x. Lemma B.1.
Tr])(0)) = Trle 22A ())
(32)
Proof. By the definition of the exponential symmetry, have that for an arbitrary X,
{(0) (())A))
(33)
Taking the derivative of both sides, we obtain that
Vol(O) e Vo, {(A)),
(34)
The standard result of Lie groups shows that An is full-rank and symmetric, and its inverse is e Therefore, we have
AAA
Vv((() Ve, Ve,e(O.) ((0x).
(35)
Now, we apply this relation to the trace of interest. By definition,
E(Ox) E[Vo, ((O) {(b)) e -\A) AA
(36)
(37)
Because AA is function of A, it commutes with A. Therefore,
[r] Tr[e A M -M A] Trie 22 2A4 (]]
(38)
(39)
Now, we ready prove the main theorem.
17

---

Proof. First of all, it is easy to see that C(Ox) is a monotonically increasing function of X. By definition,
C(O2) = e A =0" (Z+ + Z-)0.
(40)
(41)
where we have decomposed the matrix e^AAc A Z+ + Z_ into two symmetric matrices such that Z, only contains nonnegative eigenvalues, and Z- only contains nonpositive eigevalues. Because AA commute with A, they share the eigenvectors. Using elementary Lie algebra shows that the eigenvalues of 74 a, exa and that of 7_ a_e where a, 0 and dg O. This implies that o² Z_0 0 and 0 Z_0 0 are monotonically increasing functions of A.
Now, by Lemma B.I, we have
Tr -r[ ((0)))
(42)
Similarly, the regularization term is
=0AA0x=>Tr(00 >TTr00" Ac.2AA T.
(43)
A] Tr]][))A] = 8 Ae = 0, we have already proved item (2) of the theorem. Therefore, let us consider the case when either (or both) T+) # 0T A0 0
Now, by assumption, if G(Ox) O, we have either T+] 0 or Tr[00T
Without loss of generality, we assume y O, and the case of y < follows an analogous proof. In such a case, we can write the trace in terms of the eigenvectors of A:
A0 (THE((0)))] Σ #450
14,00 + Σ
7 Σ
uulot +7 Σ 2 11,1 03
4i50
Mic0
M.50
11(2)
1((A)
=: I(A),
= where is the i-th eigenvalue of A, 0, = =n(((0); 2 0 is the norm the projection of Σ in this direction.
By definition, 1, is either zero function or strictly monotonically increasing function with (- +o, (+0o) 0 Likewise, 12 is either zero function strictly monotonically increas- ing function with [2(-00) ) = 0, I'2( (+00) +oo. By the assumption Tr(E(A) # Oor Tr(60T A) + O, have that least of and I2 must be strictly monotonic function.
If I or is zero, we can take À to either +oo or oo to satisfy the condition.
If both J and In are nonzero, then I J1-12 is a strictly monotonically decreasing function with I(-oo) such that I(>*) O.
and I(+00) -00. Therefore, there must exist only unique X IR
For the proof of (4), we denote the multi-variable function J(0; G(Ox). Given that Σ(0) is differentiable, 8J 38 exists.
It is easy to see that àJ is continuous. Moreover, for any 0 and >= xx))
OJ = Σ e- 2-2Nillp 20X
22 Σ 22Nle oz +y Σ
2MApi (allellp,11 |2,1007 + 0.
4,20
P:50
M.<0
m.20
Consequently, according to the Implicit Function Theorem, the function X0 is differentiable. Additionally, 2A
50
18

---

B.3 Convergence
First of all, notice an important property, which follows from Theorem 4.3: X 0 if and only C- = 0.
Lemma B.2. For all (tt)
Ō(0) [2() (=) 20 Tx
ff*0 0: if** <...
2
(44)
Proof. As the main text, let denote 0x-, C C(0) and C' C(0*). Thus, dC
dt TT)))((())]]
(45)
o (T(E)(")) A].
(46)
where the second equality follows from Lemma B.I. One can decompose as a sum of two sym- metric matrices
A= +=QEQ0+00 + QE_Q",
(47)
=A
==
where Q is an orthogonal matrix, E, (2.) is diagonal and contains only non-negative (non-positive) entries. Note that by the definition of X", we have Tr E(0 )A] 0 and, thus,
Tr(E(0") ( -r(=
(48)
Thus,
)
(49)
= "IT
(50)
-Tr[I(")(e"" -1)AA4++ TH[I("" -)A.].
(51)
Using the inequality I+As eA (namely, that IA-I is PSD), we obtain lower bound 2X A -T-E-( ) - A
)A_]
(52)
X > O, - -1A_]
TT{ 2 2x(") )
(53)
Likewise, there is an upper bound, which simplifies to the following form if X* < 0:
Tr{E(0)A] s (2TT)<E(0)
(54)
This finishes the proof.
Lemma B.3. For any 0
C-C C*
A2 if > 0; A TA
(55)
if <.
Proof. The proof is conceptually similar to the previous one. By definition, we have
C-C* = A00" ( A0*
(56)
= =00') A(e-2x'___ 1)0*
(57)
=(0')T A, A2 A. T)0* -0') A I)0"
(58)
By the inequality / + +AAS e,, we have an upper bound
C-C² 2-C"2-2X* --2" A.00 + (0*) A A_ -I
(59)
-1)0" If (00)0 A- 2 - 2 0.
C-C² 2 -2X -2X (0*)² AZ00.
(60)
Likewise, if X* < 0, one can prove a lower bound:
C C* < 2X* (8*) A20-
(61)
19

---

Combining the above two lemmas, one can prove the following corollary. Corollary B.4.
o 1A (+ AN TT((( (0+) A28
Č C S
ifC C* 0;;
(62)
i >0 0.
Now, one can prove that as long as C is not moving too fast, C converges to C in mean square.
Lemma B.5. Let the dynamics of C be drifted Brownian motion: dC putt + sdW, where
W is a Brownian motion with variance s². If there exists co such that (0*) A20" T(E((0))A] (6
"[(((")AA Co and
Co.
2p 0(1)
E(C - -0)-
(63)
20 2
Proof. By assumption,
Č S -o'cc
(64)
C C
Let us first focus on the case when C C* > 0. By the definition of C and Ito's lemma,
d(c c*) - co (C- C²)dt= pdt sdW.
(65)
Let 7 eos (C C*), we obtain that
dz = aos ((C (
(C-CC")e Cd
(66)
S - o'oo dt se cWW dW.
(67)
Its solution is given by
Z meo
8 f coo dW.
(68)
0²co
Alternatively, if C C* < O, we let d'col cot (C+ C)2 and obtain
eett
Z s 0²co + ʃ cfe dW.
(69)
Thus,
2 200000t cot ot μ¹² 22220' out
8² f 20tcott dt
(70)
5² 20200est
(71)
20200
where we have used Ito's isometry in the first line. By construction,
24² + 82 - [C)
(72)
20 2
The proof is complete.
Let - denote convergence in probability. One prove the following theorem, the conver- gence of the relative distance to zero in probability.
Theorem B.6. Let the assumptions be the same as Lemma. (B.5). Then, if s = / 0, E[(C C* C²) O. Otherwise,
(C- C*)²
(C*)² 7p
(73)
20

---

Proof. By Lemma B.5 and Markov's inequality:
Pr([C(() - (t))> ((1))>//))))) 0.
(74)
Now, consider the distribution of (C*)². Because C is a Gaussian variable with mean pul and variance tt, we have that
Pr(IC*|> vt) 1
(75)
Now,
Pr((C(t)-C +*00) Vt) t 2 max (0, (*
(76)
(77)
+ 0,
(78)
where we have used the Frechet inequality in the second line. This finishes the proof.
B.4 Proof of Proposition 5.1
= Proof. Recall that = "", "d,)" E Bdaxd == (111"",,aa) E Radddo where ū, Ws E IRd.
0= = vec(U, W) ("],"" '", udz 2i m, w))) & R(dz+dodd
Č= = C=T((E(())
For C. C, it holds that
Vi C= =
B, i=
ve V2 C=
-B,
i=j:
v2
C 0.
0, otherwise."
>
0,
otherwise.
Therefore,
de Σ
de Σ
TH Σ Tr Vol C7
Σ
LV,
=1
k-
Σ TY Σ B). B]. 1=1
1-1
The proof is complete.
B.5 Proofs of Proposition 5.3
Proof. The loss function is
UU y/12 + 7(1171) + [W]))).
Let us adopt the following notation:
U (u,, .., ūd.) Rddxd W = (wi, Wa,) RRdxd, where ) (dd
W4 Wi IRd, = vec(U, W) (i,,"., 07 wi, w de
Due to
Va.t= Wx(u (ū) Wx vi) 2yū;, Vie [dy];
dy Va Σ d, (1.) Wx v:) + 2y05,
Vj [dz]:
the diagonal blocks of the Hessian Vigl have the following form:
2 u War T 2 Vie [dy]:
dy l= xÿ Σ ma|
22 uinno
प + 2
Vj [dg].
21

---

The trace of the Hessian is good metric of the local stability of the GD and SGD algorithm because the trace upper bounds the largest Hessian eigenvalue. For this loss function, the trace of the Hessian of the empirical risk is
S(U,W) = Tr -2y1]
Trl Wx WT] + Σ Tr| u.u
/=1
W + | I [+ d 1
+10
where Eg T....
C Proof of Theorem 5.2
Proof. First, split U and W like U (u...., Ud) Rd,.dd and W (wI,..,,?)) "T) e Rdxd, The quantity under consideration is CB ]U[U Tr[WT BW] for an arbitrary symmetric matrix B. What will be relevant us is the type of B that is indexed by two indices & and I such that
B (k..) BOkk =0
(h.1)
=1 if k and = or i = and k;
(79)
is
otherwise.
Specifically, for k, [d], we select B/ 80 de.0ge in Cg. With this choice, for an arbitrary pair of h and 1,
Cgaun uiu w.ul.
and
W B(h. lwwm + w
(80)
(B( = -usl + u
(81)
Therefore,
E Σ [E():) [(k]
E Σ a. Wx y. Tr [ wax T W B(0]
(82)
#1
E [|r|] an W B(,D]]
(83)
Tr [B[]]]] 2x I W]
(84)
-Tr[Ew wlmm)
(85)
Ew
(86)
where we have defined r, = 07 Wr ys and E'w E[[r|"r>"]
Likewise, we have that
dy
E Σ Tr [T()))))
(k)
= E Σ
2
Σ
is (i Wx ui)
ie
N [ Wx u:) 18(K)
J-1
j-1
Tr U
= =2u]E', u.
where we have defined E, 1 Therefore, we have found that for arbitrary pair of l: and
Čga. 2x w.w[)) + 2n E ut 2.4)
(87)
The fixed point of this dynamics is:
(88)
22

---

where En 17 + I and Ey = 1/ +YI. Because this holds for arbitrary k and L, the equation
can be written matrix form:
WT U.U.
(89)
Let V = UW. To show that a solution exists for an arbitrary V. Let W = W and U' = /E,U, which implies that
U'W" - VVEVVV =
(90)
and
-
(91)
Namely, U and (W")T must have the same right singular vectors and singular values. This gives us the following solution. Let V = LSR be the singular value decomposition of V", where L and R are orthogonal matrices an a positive diagonal matrix. Then, for an arbitrary orthogonal matrix F, the following choice of U' and W" satisfies the two conditions:
U LVSF: I VR.
(92)
This finishes the proof.
D Discrete-Time GD and SGD
In fact, our results hold in similar form for discrete-time GD and SGD. Let us focus on the expo- nential symmetries with the symmetric matrix A.
The following equation holds with probability 1:
0 vol(0.z)
(93)
For discrete-time SGD, it is notationally simpler and without loss of generality to regard {(0) as the minibatch-average loss, which is the notation we adopt here. This because a symmetry holds for every per-sample loss, then it must also hold for every empirical average of these per-sample losses.
The dynamics of SGD gives
A0; = -VVo(0)((0),,)))
(94)
This means that
A()- =..
(95)
Therefore, we have that
AC = 220, A0, AC, AA0, ADO.
(96)
(97)
Therefore,
AC - TT)
(98)
where (d) (o(( is by definition PSD. Already, note the similarity between Eq. (98) and its continuous-timg version. The qualitative discussions carry over: if PSD, C, increases monotonically.
Now, while the first-order terms in y also vanish in the r.h.s, the problem is that the r.h.s. becomes stochastic because Ea(0,) is different for every time step. However, one can still analyze the ex- pected flow and show that the expected flow (over the sampling of minibatches) is zero a unique point in a way similar to the continuous- citin limit of the problem. Therefore, we define
Ga(O.) E_[AC.]]
(99)
(o.) = 1E, [Ea]
(100)
We prove the following theorem. Note that this theorem applies for any batch size, and so applies to both SGD and GD.
Theorem D.I. (Discrete-time fixed point theorem of SGD.) Let the per-sample loss satisfy the A exponential symmetry and 0g == exp[XA]0. Then, for any 0 and any 7 RR
23

---

(1) Ga(Ox) is monotonically decreasing function of X: there exists a X* IR U {+00} such that G0(())) O:
(3) in addition, if (T))(
0 or Tr[00' A] * 0, X is unique and Ga(Ox) is strictly
monotonic;
(4) in addition to (3), if (0) is differentialee, X (0) is differentiable function of 0.
Proof. Similarly, let us establish the relationship between Ve(0) and vε(exp(AA)). By the defini- tion of the exponential symmetry, we have that for arbitrary X,
((0) ((eAA) 0).
(101)
Taking the derivative of both sides, we obtain that
Vol(O) e Ve, ((02),
(102)
The standard result of Lie groups shows that eM is full-rank and symmetric, and its inverse is e eA Therefore, we have
AA (((( Ve, ((0x).
(103)
Now, we apply this relation to the trace of interest. By definition,
Ex) E[Vo, ((x)) {(x))
(104)
E -AA- EA (0)e AA.
(105)
Because AA is function of A, it commutes with A. Therefore,
(O) Tr(e- (0)e AA A]
(106)
Tr.e 22 (0)
(107)
Similarly, the regularization term is
A0 A0 = > -T Ae 2M1
(108)
Now, if (r())(())A] = 0 AO 0, we have already proved item (2) of the theorem. Therefore, let us consider the case when either (or both) Tr[Ea(0)
0 or 6 "A0 #
Without loss of generality, we assume 0, and the case of 7 < O follows an analogous proof. In such a case, we can write the trace in terms of the eigenvectors n, of A:
A0x (T)() = Σ
pslo +7 Σ 22eullll 7 Σ
2
pulo +7 E
2
8420
4,<0
1,00
(2(
I2(A)
=: I(A), where i4 is the i-th eigenvalue of A, 0, (n]0,) =n] is the norm of the projection of Eg in this direction.
By definition, In is either zero function or strictly monotonically increasing function with [- + +oo, I2 = 0 Likewise, In is either a zero function or strictly monotonically increasing function with I200) 0, 12(+ (00) +oo. By the assumption Tr(Ea(0)A)) 0 or Tr(00 A) 0, we have that at least one of J and T2 must be a strictly monotonic function.
* If Tı or Tz is zero, we can take À to be either + or -o to satisfy the condition.
If both J, and 12 are nonzero, then / = J J1 is strictly monotonically decreasing function with T(-oo) such that T(x*) = O.
and I(+00) -00. Therefore, there must exist only unique X IR
The proof of (4) follows from the Implicit Function Theorem, as in the continuous-time case.
The final question is this: what does mean for to reach a point where G.(0) = 0? An educated guess can be made: the fluctuation in C does not vanish, but the flow takes C' towards this vanishing flow point- - something like a Brownian motion trapped in local potential well [37]. However, difficult to say more without specific knowledge of the systems.
24

---

E Proof of Theorem 5.4
We first prove the following theorem, which applies to an arbitrary parameter that are not necessarily local minima of of the loss.
Theorem E.1. Let i. y, Ei.1 W...WW. and h; := W.-.WW For all layer the equilibrium is achieved al
W W+14 67 Ci W W.N.CH
(109)
where C = 111 J, == B8| r1 _ J. Or equivalently,
cl C/fs = h h
(110)
Proof. By Proposition 5.1,
d C
de
de SW...
) ET
[ (
oow,) )
de
B de aw;
0²
ETr
B
(111)
dt
@W
The derivatives are
0e rW:,t(
(112)
aW.+1 i+1
0
(113)
oW;
Therefore, the two terms on R.H.S of Eq. (111) are given by
de B
de
) (WWher)))))))))) B(W.h)) (..],
ETr
@W
TE r
( PTr[haaa" W. BW,]
(114)
( Oe (aw.)
)
B 0e
ETr
Tr[W El. r(,,2) B(harrr (+ [
ow;
EE[ h Tr W rr rr E W]]
(115)
Because the matrix arbitrary, we can let B = &r, + 8,, Then, the two terms become
ae
OE
ETr
21 1 1 WT har h1 w..h.
(116)
OW
@WW.,
ETr ( de
(ow,)
B 0e
2 2 1, sla rr (--WWer..
(117)
aw,
Here, we define the vectors W, (w., "., ".))T and W = (Wi+1,1 w..1.)) Because Eq. (116) and (117) hold for arbitrary k; I, we have
Oe
ae (WWe1 0e OW;
ETr
2W 2W.4, 2ax 22x WT
(118)
aW
de
ETT
B
2W 2W B1 1r Wai.
(119)
ow,
For Eq. (111) to be 0, we must have
W W. W. = h,ær Prr "rrr]e W
(120)
which is Eq. (109). The proof is complete.
We now ready to prove Theorem 5.4.
25

---

Proof. It suffices to specialize Theorem E.1 to the global minimum. At the global minimum, we can define
D y=c
(121)
Then, Eq. (109) can be written as
W W.. D., D...WW. W
W41 W; Wi-1 W1 WT W....
W.
(122)
Tr Tr[WT D..
TT[W Wi E WT 7-1
To solve Eq. (122), we substitute WD and W1 with Wi = W1 WivVg and W/ = W. which transform Eq. (122) into
WT -. W. D W W.1 W. W W...-wW.T.WWW. W7.
W 1+2
(123)
T W h.......
TH WW
The global minimum condition be written
W-------W VV =V
(124)
Then, we can decompose the matrices
as
WD LEPU D-1 = (i= W = UUER,
(125)
where ".. E1 Badd, L E Rdxxd U; RRddxd R Rddd with d rank(V") and arbitrary d. The matrices Uj satisfy U,U: Laxa By substituting the decomposition into Eq. (123), we have
Ex.1 DD Tr 2 En E.n]
EE EE
(126)
Since these diagonal matrices commute with each other, we can see D; claxd. Then we move on to fix the parameter By taking i and i D - in Eq. (126), we obtain
Σg Tr 2
Σ
Σ
(127)
T
D]
d
>>
>
Σ Tr[2
(128)
d
Tr
D-1
where d represents the dimension of the learning space. By taking trace to both sides Eqs. (127)
and (128), can Tr[E]] Tr[E}]] and hence, Σ₁ Ep. The parameter c is given by
Tr E
c=
(129)
d
With the SVD decomposition V' LSR,, R, we have
D-2 8¹
(130)
Therefore, the solutions for and En are
T 1/D
S
(D-2)/2D
Σ
VSa.
(131)
d
(D-2)/2
T
The scaling of the diagonal matrices are shown as
T+ d' 2/- 2/
Tr
(I
2 2-2/
Tr[
(132)
The proof is complete.
26

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We believe that the abstract and introduction reflect the contributions and scope the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of our work in the conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specificatton, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [es]]
27

---

Justification:
We believe that the assumptions are clarified and complete proofs are pro- vided for the theoretical parts.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We believe that all of the experimental results are reproducable in our work.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
28

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
The code data of the experiments are simple and easy reproduce follow- ing the description in the main text.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https: //n::/nnnn /ttp//ips..o cc/ ssionP cesssonooocyyy for more details.
public/ gwwdgegie//eegg Seuidee/ooeeeee
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. cc/ ppublic gu i des Codeeuum ssi oooooooooooooo icyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have specified the training and test details of the experiments in the cap- tions or the corresponding explanations in the main text for Figs. 2, 3, 4, 5 and 6.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should presented in the core the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification: Here the dynamics is deterministic and there is no need to consider the error bars here.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
29

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative еπτог rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification:
The experiments can be simply conducted on personal computers.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:/ hete:///nnuut /neuripsss../// ca http/pbbbc//ticaudeelll
Answer: [Yes]
Justification:
We have confirmed that the research conducted with the NeurIPS Code of
Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is fundamental research on the learning dynamics of SGD and hence does not have direct positive or negative societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
30

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
We believe there is no risks for misuse for the data and models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions to access the model or implementing safery filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer:[NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
31

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw thoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [No]
Justification: Nothing introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: We believe that neither the crowdsourcing nor the research with human sub- jects is included in our work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects. and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not contain crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
32

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "nTOcLnzV0F", "reviewer": "Reviewer_rhMv", "summary": "The paper takes important steps toward showing that the noise in SGD pushed the dynamics along symmetry directions toward a unique fixed point. This differs in an important way from GD and gradient flow (GF) where one can show that in the presence of continuous symmetries (called exponential symmetries in the paper) there will be conserved quantities $C$, such as layer imbalance. These conserved quantities ensure that runs starting with different values of $C$ end up on different points along a minimum manifold (with some caveats). \nThis paper shows that SGD noise causes drift toward a fixed point on a minimum manifold.", "questions": "There are many details int he computations and proofs which I think need to be clarified. One aspect which I think is an error is related to the Lie algebra of the symmetries discussed in the paper. Most derivations in the paper hold for symmetric Lie algebra elements (i.e. generators of scaling and hyperbolic directions in the group) but fail for rotations. A relevant work which studies the parameter space more thoroughly is:\n[3] Zhao, Bo, Iordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy. \"Symmetries, flat minima, and the conserved quantities of gradient flow.\" ICLR (2023). \n\nMy questions are as follows:\n\n\n\nline 49: ref [40] may be wrong (it doesn't discuss conservation laws) maybe another paper by Zhao is meant: Zhao et al ICLR 2023 \"Symmetries, flat minima, and the conserved quantities of gradient flow\". \nsec 4: important caveat about conserved quantities is that the antiderivative C may not always exist. An important case is when the symmetry is a rotation. (see Zhao et al \"symmetries, flat minima ...\" ICLR 2023) In this case, the Noether charge is zero. \n\n1. line 78: eq 5: to clarify, $\\nabla^T C \\nabla \\ell(\\theta, z) = 0$ is always zero due to eq 2 and 3, right? Is that why both terms in eq 30 (App B) vanish? \nHowever, there are some statements in Appendix B which seem inaccurate to me: \n  a. when $\\Sigma  = GG^T$, the matrix $G$ need not be square and $G = \\sqrt{\\Sigma}$ is just one of the possible $G$. What is your argument for this $G$ being the correct $G$? \n  b. I don't agree with your statement that \"Because $\\Sigma$ and $\\sqrt{\\Sigma}$ share eigenvector...\" eq 31 would vanish. Trivial example: the 2x2 matrices $A=diag(1,-1)$ and the identity matrix share eigenvectors, but the vector $v=(1,1)$, $Av = 0$, but $Iv \\ne 0$. I think the real reason is that $\\Sigma = GG^T$. Doing SVD, $G = USV^T$ and $\\Sigma = US^2U^T$. Then eq 30 yields $\\nabla^T C U S= 0$. Assuming the covariance is nondegenerate (i.e. $S_{ii}\\ne 0$ for all $i$), we must have  $\\nabla^T C U = 0$ which also implies $\\nabla^T C G = 0$. \n\n2. line 83: at a minimum, some symmetry orbits (i.e. invariant manifold) may collapse to a point, e.g. rotation symmetry in a quadratic loss as in linear regression $\\|\\theta x -y \\|^2$.    \n\n3. line 104: Instead of \"double rotation\" this should be called double linear transformation, becuase the group is not just $O(h)$, but rather $GL(h)$\n\n4. eq 6: As noted above, this Noether charge vanishes for rotations $SO(h)$ because $A$ is anti symmetric and hence $\\theta^T A \\theta=0$ (see Zhao ICLR 2023).   \n\n5. line 114: why? Why the explicit weight decay? Why in defining the symmetry do you ignore $\\|\\theta\\|^2$ which doesn't have rescaling or $GL(n)$ symmetry, only rotation symmetry $O(n)$?\n\n6. Eq 7: what happens to the cross terms of $\\theta^t \\nabla \\ell$ in $\\Sigma_\\gamma$? \n\n7. line 118: eigenvalues and vectors of A can be complex (e.g. for Lie algebra of rotation), so $n^T$ needs to be $n^\\dagger$. \n\n8. line 214-215: I don't understand how you arrive at $\\dot{C} = \\mathrm{Var}[r]C$. I think you mean when the hidden dimension is 1, not the output, right? I hidden is $d>1$, let $\\sigma=\\Sigma_{u_i}$, then for $B= E_{kl}+E_{lk}$ we have  $\\mathrm{Tr}[\\Sigma_{u_i}B] = \\sigma_{lk}+\\sigma_{kl}$, which are off-diagonal terms in the covariance matrix, unless $k=l$. \n\n9. eq 13: When $\\gamma\\ne 0$, eq 13 does not have the symmetry (11) for symmetric Lie algebra element $B$. The reason is $\\|W\\|_F$ is invariant under rotations where $A^{-1} = A^T$. The Lie algebra of rotations (infinitesimal $A = I+\\rho B$) has anti-symmetric $B$. \n\n10. line 235: under what condition will the covariances $\\Sigma_x$ and $\\Sigma_\\epsilon$ be isotropic? I don't think this happens often in SGD. I think you need adaptive gradient optimizers or Newton's method to have isotropic $\\Sigma_\\epsilon$. \n\n11. eq 39,40: $\\nabla^T_{\\theta_\\lambda} \\ell = [\\nabla_{\\theta_\\lambda} \\ell]^T e^{-\\lambda A^T}$, so we don't have (40)  and in (41) there should $e^{-\\lambda A^T}A e^{-\\lambda A}$. \n\n12. line 531: $Z_++Z_-$ will have complex eigenvalues when $A$ is not symmetric. Diagonalizing $A$, your arguments about $Z_\\pm$ only hold for symmetric $A$.", "limitations": "There are some gaps in the derivations (see questions) which I hope the authors can fill. \nAdditionally, the experiments seem to only look at two aspects: alignment of weights in matrix factorization (which is already known under weight decay, afaik), and bias toward less sharp minima (again, known). I would have liked to see a more direct test of the mechanism the paper claims is driving this movement toward flatter minima. At least, the balancing of the weights or some other direct result from the paper would have been desirable.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This is potentially a very significant work. This is one of the few works I know which focuses on the effect of SGD noise on the symmetry directions and derives a concrete dynamics for this. I think the approach is at least partly original. Although there are many related works (see weaknesses), this paper is the only one which derives concrete solutions for the noisy dynamics along the symmetry directions. \nThere are some corners that need to be ironed out. But if the results survive, this work could answer many questions about the implicit bias and generalization in SGD.", "weaknesses": "There exist multiple other works addressing the same questions. Two particular ones with very similar claims that the authors may have missed are: \n[1] Chen, Feng, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. \"Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks.\" Advances in Neural Information Processing Systems 37 (2023).\n[2] Yang, Ning, Chao Tang, and Yuhai Tu. \"Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions.\" Physical Review Letters 130, no. 23 (2023): 237101.\n\nI ask the authors to clarify the distinction between their work and the above. [1] is also about the effect of SGD noise on the symmetry directions, afaik. \n[2] also derives the effective dynamics of SGD as a Langevin equation and shows a result similar to yours, namely that along symmetric minima, the noise drives the system toward the least sharp minimum.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review. Please first see our summary rebuttal above. We will answer concisely due to space constraint. However, please ask us to elaborate on every point that you need more detail.\n\n**Ref [1]...**\n\nRef. [1] shows the existence of an invariant set, mainly due to the permutation symmetry. Our work focuses on continuous symmetry. We will include a reference to clarify this.\n\n**Ref [2] ...**\n\nThis is a misunderstanding of our result. Our result shows that SGD can bias the solutions towards both sharper and flatter solutions, and this depends strongly on the data distribution. See our answers below on this point and the new experiment in Figure 7. We will add this reference and discussion to the manuscript.\n\n**Questions:**\n\n**Most derivations hold for symmetric Lie algebra elements...**\n\nThis is due to your misreading. In the starting Definition 4.1, the exponential symmetry is explicitly defined with respect to a symmetric A. The reason is exactly what you point out -- if A is antisymmetric, the existence of Noether charge is questionable. We will add a note to clarify this below Def 4.1.\n\nAlso,  we point out that antisymmetric matrices can have noether charge -- they are currently very poorly understood and hidden from simple analysis. As a counterexample, consider a rotation symmetry for (u,w) in 2d: the symmetry implies the following dynamics for GD: $wdu - udw = 0$. Simplifying the relation leads to (ignoring the sign problem for now) d(log(u) - log(w)) = 0. We have thus constructed a Noether charge for an antisymmetric-matrix symmetry: log(u) - log(w). \n\n**line 49: ref [40] wrong...**\n\nWe will fix this.\n\n**Existence of C**\n\nHere, a technical remark on the antiderivative C is that one needs to assume that J is conservative, which is well known and so we omitted its discussion. We will add a note on this in the revision.\n\n**line 78: eq 5:... a. What is your argument for this G being the correct G?**\n\nYes. You are correct.\n\nFor point a, this is a good point that needs to be clarified. We made the specific choice that $G=\\sqrt{\\Sigma}$ in Eq. (1), but the goal was only to avoid introducing unnecessary technical details. Generally, one should use a general $G$ in the place of $\\sqrt{\\Sigma}$ such that $\\Sigma = GG^T$, and, had we made this choice in Eq (1), what you suggested in point b will be the proof and so the result is the same. We will clarify this point.\n\n**\"I don't agree with your statement...\"**\n\nThanks for this very deep question and suggestion (also, see the answer to point a). What you suggested is a more technical (and more correct) way of proving the relation, but there are multiple routes to get to Eq 31, and our original proof is correct. We also used a hidden condition that is rather obvious: the eigenvalues of $\\sqrt{\\Sigma}$ are the square roots of the corresponding eigenvalues of $\\Sigma$. Thus, a zero eigenvalue of $\\Sigma$ must correspond to a zero eigenvalue of $\\sqrt{\\Sigma}$. We will add a note on this. \n\n**line 83...**\n\nOur statement in line 83 does not contradict this statement. A point is still a \"connected manifold\" but of dimension zero.\n\n**line 104...**\n\nWe will fix it.\n\n**eq 6...**\n\nA is defined to be symmetric.\n\n**line 114: why weight decay...**\n\nThanks for this question. Here, the (fundamental) question is not \"why\" we include weight decay, but \"what is the most general setting for which Theorem 4.3 holds correct?\" It turns out that Theorem 4.3 holds true for any nonnegative $\\gamma$, and this is the main reason why Def 4.2 allows for a weight decay.\n\n**Eq 7: the cross term...**\n\nWe are not sure if we understand this question. There is no contribution from the weight decay term to $\\Sigma$ because $\\Sigma$ is a covariance of the gradient, whereas the weight decay only affects the expectation of the gradient, not its variance.\n\n**line 118...**\n\nBy definition, A is symmetric. See our answer above.\n\n**line 214-215:**\n\nThanks for this question. This example assumes the standard MSE loss for matrix factorization and is misplaced. We will move it below Eq. (13), and include a derivation of it.\n\n**eq 13: When $\\gamma \\neq 0$...**\n\nYes, but this is not a problem. Because (please first see our reply above) Eq 13 can be decomposed into a sum of a symmetry term and weight decay (which does not have the A symmetry). So, this loss function obeys Def 4.2 and so Theorem 4.3 is applicable.\n\n**line 235...**\n\nThese are not functions of the learning algorithm -- they are simply properties of the data points. \n\n**eq 39,40...**\n\nMatrix A is by definition symmetric. \n\n**line 531...**\n\nA is by definition symmetric. \n\n\n**Limitations:**\n\n**...alignment of weights in matrix factorization is already known under weight decay**\n \nWe are afraid to point out that you have misunderstood our result, and our insight is not known. \nThere are two aspects of alignment:\n1. Alignment in norm: under SGD, the norm ratio is determined by the input and label covariance -- and so the balance in norm in general will NOT happen (even with weight decay), and this is a crucial new message we offer. Please examine Eq (14): only for very special $\\Sigma_x$ and $\\Sigma_\\epsilon$, a balance in norm will happen. \n2. Alignment in eigenvectors. Our crucial result is again that using SGD, the alignment will not occur even if weight decay is used.  \n\nSee Figure 7 in the pdf for additional experiments.\n\n**\"and bias toward less sharp minima (again, known). \"**\n\nThis is also a misunderstanding. See the summary rebuttal and Figure 7-right for an experiment.\n\n**I would have liked to see a more direct test of the mechanism the paper...**\n\nBesides the result in Figure 7, another new insight is precise balance in the norm only happens in networks with at least two hidden layers, and, for a two-hidden layer network, only the two intermediate layers will always have balanced norms -- See Figure 8.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "PCu2trivVK", "reviewer": "Reviewer_HyRa", "summary": "This paper considers the relationship between the behavior of SGD for solving empirical risk minimization and the exponential symmetry.", "questions": "- SGD for minimizing $L(\\theta) = E_z [\\ell (\\theta,z)]$ is defined by (1): $d \\theta_t = - \\nabla L(\\theta_t) dt + \\sqrt{2 \\sigma^2 \\Sigma (\\theta_t)} d W_t$ using $\\sigma^2 = \\eta/(2S)$, where $\\eta > 0$ is a step-size and $S > 0$ is a batch size. My first question is \"Which is a problem considered in this paper: the expected risk minimization or the empirical risk minimization?\"  $L$ is defined by using the expectation $E$ (Lines 56-57), however, Line 57 says \"the empirical risk function.\"\n- Line 19 says that SGD with $\\sigma = 0$ implies GD. I do not understand it, since $\\sigma^2 = \\eta/(2S) = 0$ iff $\\eta = 0$, which contradicts $\\eta > 0$.\n- Line 21: I do not understand what $\\sigma^2 \\ll 1$ is. \n- Line 58: Please define $\\Sigma_v (\\theta)$ explicitly. \n- Line 93: a --> an \n- Eq (7): Is $-4 \\gamma$ in (7) replaced with $-2 \\gamma$?\n- Line 142: I do not understand the definition of \"fixed point.\" Please define \"fixed point\" explicitly. A fixed point $x$ of a mapping $T \\colon R^d \\to R^d$ is defined by $x = T(x)$. The definition of \"fixed point\" considered in this paper seems to be different from a fixed point of a mapping (see Theorem 4.3).\n- Eq (8): I do not understand the relationship between $\\ell (\\theta, z)$ and $\\ell$ defined by (8). Why do we consider $\\ell$ in (8) not $\\ell (\\theta, z)$? I think that the main objective of this paper is to minimize $L(\\theta) = E_z [\\ell (\\theta, z)]$. \n- Eq (9): We cannot define (9) when $\\gamma = 0$. Please define $C(\\theta^\\star)$ explicitly. \n- Eq (11): $\\theta'$ is not defined. \n- Section 5.1: This section considers a matrix factorization: Given $A$, find $U$ and $W$ such that $A = UW$. The assumption of $A$ seems to be strong such that $A$ is close to the identity matrix $I$ (Lines 204 and 205). \n- Proposition 5.1: The assumption \"the symmetry (11) holds\" (based on $A \\approx I$) seems to be strong. Please provide practical examples of empirical risk minimization and SGD satisfying (11). \n- Figure 2: There are not definitions of the matrix size, a batch size $S$, $\\eta$, and $\\ell (\\theta,x)$.\n- Line 222: $\\epsilon$ is not defined. \n- Theorem 5.2: Is the theorem based on (11)? If so, please provide some examples of $\\dot{C}_B = 0$ under (11). \n- Figure 3: There are not definitions of $d_x$, $d_y$, $\\epsilon$, $S$, and so on. Moreover, I do not understand why a warm-up with $\\eta_\\min = 0.001$, $\\eta_\\max = 0.008$, and the switching steps $5000$ is used. Please justify the setting from the theoretical results.  \n- ...\n- Additional comments: I would like to know new insights that come from the theoretical results in this paper. This is because I am not surprised at the results. For example, can we apply the results to training more practical deep neural networks on practical datasets? Then, can we say new insights for training practical deep neural networks? Moreover, can we estimate appropriate step-size $\\eta$, batch size $S$, and hyperparameter $\\gamma$ using the theoretical results before implementing SGD? Anyway, I would like to know contributions of this paper.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The strength of this paper is to study the relationship (Theorem 4.3) between the behavior of SGD for solving empirical risk minimization and the exponential symmetry (Definition 4.1). Moreover, it considers a noise equilibrium (Definition 4.5) and provides explicit solutions of the noise equilibriums for some applications (Section 5).  In addition, it provides some numerical results to support the theoretical analyses presented in this paper.", "weaknesses": "The weaknesses are the following. Please see Questions for details.  \n- The mathematical preliminaries are insufficient. \n- The assumptions seem to be strong. \n- The explanations of contributions are insufficient.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your criticisms. We believe that your criticisms are due to misunderstanding of mathematics and the related literature and not paying attention to the details given in the immediate context. We will answer concisely due to space constraints. However, please ask us to elaborate on every point that you need more detail.\n\n**Questions:**\n\n**the expected risk minimization or the empirical risk minimization...**\n\nThanks for this question, and we apologize for the confusion. In our formalism, it is not essential to differentiate between the population risk and the empirical risk. The loss function can be either the population risk (if the training proceeds with online SGD) or the empirical risk (if the training proceeds with a finite-size dataset). This is because $\\mathbb{E}$ denotes the expectation over the distribution of the training data, which can be arbitrary and is problem/algorithm-dependent.\n\n**Line 19...**\n\nWe are referring to the limit $S\\to \\infty$. Within the SDE framework, this limit is the GF/GD algorithm.\n\n**Line 21: I do not understand what σ2≪1 is.**\n\nThis is a common notation meaning that \"$\\sigma$ is far smaller than 1,\" for example, $\\sigma = 0.1$, $0.01$, etc.\n\n\n**Line 58: Please define Σv(θ)...**\n\nThis is defined and explained in line 58. Colloquially, $\\Sigma_v$ is a shorthand for the gradient covariance for a subset of all parameters $v$.\n\n**Eq (7)...**\n\n Yes. We will fix it.\n\n**Line 142: I do not understand the definition of \"fixed point...**\n\nIn the study of dynamical systems (DS), the fixed point is defined differently. In DS, a fixed point is defined for differential equations of the form: $\\dot{x} = f(x)$ as the set of $x$ such that $f(x) =0$. See standard textbooks of this field (e.g., Introduction to the modern theory of dynamical systems, by Katok et al.). In this terminology, Def 4.5 is an explicit definition of what \"fixed point\" means in our context. This definition is related to the fixed point of mappings. Here, the underlying mapping is one step of GD, which maps $\\theta_t \\to \\theta_{t+1}$.\n\n\n**Eq (8): relationship between ℓ(θ,z) and ℓ...**\n\n$\\ell$ and $\\ell(\\theta,z)$ are always the same thing. The $z$ argument is always abbreviated for concision. This has been explained in line 61.\n\n**Eq (9): We cannot define (9) when γ=0.**\n\nThis is a good point. It is better to write Eq (9) as $2\\gamma C =\\sigma^2 Tr[\\Sigma A]$, which is now well-defined. We will fix this.\n\n**Please define $C(θ^\\*)$ explicitly.**\n\nThis quantity is defined in the line 183. $C(\\theta^*)= (\\theta^*)^T A \\theta^*$, and $\\theta^*$ is whatever quantity that satisfies $G(\\theta^*) = 0$ (whose existence is guaranteed by Theorem 4.3).\n\n\n**Eq (11): θ′ is not defined.**\n\nThanks for pointing this out. $\\theta'$ are the other parameters that are irrelevant to the double rotation symmetry. We will include this in the revision.\n\n**Section 5.1: The assumption of A seems to be strong such that A is close to the identity matrix I....**\n\nWe are afraid to say that this is a major misunderstanding of the problem setting of Section 5.1. Here, the goal is to find $U$, $W$ such that $UW=C$ for a fixed matrix $C$ (note that this $C$ is not $A$). This problem has a lot of degeneracy in the solution. For any invertible matrix $A$: $UA$, $A^{-1}W$ is also a solution to $UW=C$. Hence, this problem has \"symmetries.\" The matrix \"A\" in our original text has to do with these arbitrary invertible transformations, not with the target matrix $C$. \n\nThere is no assumption at all. To derive proposition 5.1, one focuses on the type of symmetries that are close to identity: $A \\approx I$; this is because these matrices are the generator of the Lie group for these symmetries. This is not an assumption because $A$ can be an arbitrary matrix.\n\n**Proposition 5.1: The assumption \"the symmetry (11) holds\" (based on A≈I) seems to be strong. Please provide practical examples of empirical risk minimization and SGD satisfying (11).**\n\nSee the answer above. This is not an assumption. Eq (11) is satisfied by any (deep) matrix factorization problem.\n\n**Figure 2...**\n\nThe experimental details are given in appendix A2. The loss function is MSE and given by Eq (13). The matrix size is 30 by 30, the learning rate is 0.1, and the batch size is 1. We did not specify learning and batch size because the result is highly insensitive to them (as any learning rate and batch size would work and give a similar figure). See Figure 9 of the pdf. We will include these details in the revision.\n\n**Line 222**\n\n$\\epsilon$ is a noise on the label. We will fix it.\n\n**Theorem 5.2: Is the theorem based on (11)? If so, please provide some examples of C˙B=0 under (11).**\n\nWe are not sure if we understand this question. Theorem 5.2 (and Eq. (14)) enumerates all examples where $\\dot{C}_B=0$. \n\n**Figure 3**\n\nThey are the same as in Figure 2. $d_x=d_y=30$ and $\\epsilon \\sim \\mathcal{N}(0, I_{30})$ is an independent Gaussian noise. We did not specify learning and batch size because the result is highly insensitive to them. We will clarify this in the revision.\n\n**Moreover, I do not understand why a warm-up with $η_{min}=0.001, η_{max}=0.008$, and the switching steps 5000 is used...**\n\nLike the previous experiment, the result is robust to these choices, and the specific values of these quantities are not essential at all. The specific values are just chosen out of convenience. For example, a switching step of $5000$ is arbitrary, and we observe that choosing it to be any number larger than 1000 will give a similar result. $η_{min}=0.001$ is also arbitrary and nonessential. Choosing it to be any value such that the model does not diverge at initialization will also work.\n\n**...Anyway, I would like to know contributions of this paper.**\n\nThanks for this important question. This is a good question but admits many different answers. Due to space constraints, we will have to answer this in the comment.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "6CvLYc1IA6", "reviewer": "Reviewer_QY9h", "summary": "The authors study the effect of exponential symmetries on the learning dynamics of SGD. They establish the theorem that every exponential symmetry implies the existence of a unique and attractive fixed point in SGD dynamics, which could explain phenomena in machine learning such as matrix factorization and progressive sharpening or flattening.", "questions": "1.\tIn line 124-128, the authors state that the SGD dynamics can be decomposed into two independent parts. Why does equation (5) have no effect on the loss? I would expect the noise will result in fluctuation of the loss at the end of training, which means that the decomposition should not hold.\n2.\tFor the double rotation symmetry, why it is hard to write the matrix A explicitly? It seems that A can just be any symmetric matrix as discussed in section 5.1. Also, why do you require A to be symmetric when defining the exponential symmetry? I am wondering if the results can be generalized to a general matrix A.\n3.\tI don’t understand why one needs to assume C* to be a constant in line 163. By definition, C is a deterministic function that only depends on $\\theta$.", "limitations": "The authors have included the limitations. For additional points, please see my comments above.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "The paper studies the dynamics of SGD. The authors prove a novel theorem that relates continuous symmetry to the attractive fixed point in SGD. The theorem is quite general and can have significant potential impacts on the field. The authors apply their theorem to matrix factorization, offering strong analytical and experimental support. And by my knowledge, this is a novel contribution.", "weaknesses": "1.\tC is not guaranteed to exist. (Line 70)\n2.\tUsually, the stability of the learning is determined by the largest eigenvalue of the Hessian instead of the trace as discussed in the paper. The trace may not be an upper bound on the largest eigenvalue as it can contain negative eigenvalues.\n3.\tI think it’s missing dependence on $\\gamma$ in proposition 5.3.\n4.\tThe authors proved that in deep linear networks, the noise will lead to a balanced norm for all intermediate layers. This is an important theorem, but it lacks empirical verification. It is important to verify that the empirics actually match the theorem, which would strengthen the paper by a lot. Such an experiment should not be hard to run.\n5.\tThere is a lack of clarity in how the analyses in Section 5.4 (Equation 24) relate to Figure 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the feedback. Please first see our summary rebuttal above.\n\n**Weaknesses:**\n\n**C is not guaranteed to exist. (Line 70)**\n\nThanks for this comment. This is a well-known fact in the study of symmetries and the Noether theorem. It is required that $J$ needs to be conservative. We will add a note on this. That being said, for the type of symmetries (Def 4.1) we focus on in our work, $C$ always exists.  \n\n**Usually, the stability of the learning is determined by the largest eigenvalue of the Hessian... The trace may not be an upper bound on the largest eigenvalue as it can contain negative eigenvalues.**\n\nThanks for this question. For (deterministic) GD, the training stability is indeed determined by $\\lambda_{\\max}(Hess)$. However, this does not apply to SGD. In fact, [1][2] demonstrates that the stability of SGD is impacted by $Tr(Hess)$, rather than $\\lambda_{\\max}(Hess)$. [1][3] suggest that $Tr(Hess)$ significantly determines the properties of SGD’s solutions.\n\n\n[1] Wu and Su. The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent. (ICML 2023)\n\n[2] Ma and Ying. On the linear stability of SGD and input-smoothness of neural networks. (NeurIPS 2021)\n\n[3] Li et al. What Happens after SGD Reaches Zero Loss? –A Mathematical Framework. (ICLR 2022)\n\n\n**I think it’s missing dependence on γ in proposition 5.3.**\n\nWe will fix it in the revision.  In the experiments, we let $\\gamma=0$ and so we did not use the weight decay term.\n\n\n**The authors proved that in deep linear networks, the noise will lead to a balanced norm for all intermediate layers. This is an important theorem, but it lacks empirical verification....**\n\nThanks for this very constructive feedback. We have added an extensive set of experiments for deep linear nets. We show that for MNIST, a deep linear net (with five hidden layers), the intermediate layers converge to the same norm, whereas the input and output layers are special and move towards different norms. This effect is also shown to be robust across different types of initializations. See Figure 8. In addition, see the additional experiments in Figure 7 and Figure 9.\n\n\n**There is a lack of clarity in how the analyses in Section 5.4 (Equation 24) relate to Figure 4.**\n\nThanks for this criticism. Section 5.4 (and Eq 24) shows that if a model approximately has a symmetry (for example, when it can be approximated by a model with symmetry), then its solution will be biased towards the noise equilibrium of this symmetry model. \n\nFigure 4 shows exactly this. A tanh network can be approximated by a linear network of the same depth and connectivity (because $tanh(wx) \\approx wx$ to leading order). For a linear network, it is true that the noise equilibrium will satisfy $W\\Sigma_xW^T = U^T \\Sigma_\\epsilon U$, and so one would expect that under SGD training, a tanh network will also be biased towards satisfying this relation, and this qualitative prediction is confirmed by Figure 4. We will add this explanation to the revision.\n\n**Questions:**\n\n**In line 124-128, the authors state that the SGD dynamics can be decomposed into two independent parts. Why does equation (5) have no effect on the loss?**\n\nThis is a good question. This statement is made assuming that the fluctuation is not significant. This is supported by the experiment in Figure 3  and the mid panel of Figure 2, where, at the end of training, the fluctuation in the converged solution is small.\n\nAlso, there is a deeper mathematical reason for this statement. The fluctuation can also be decomposed into the parts in the subspace parallel to the gradient and orthogonal to it. The fluctuation in the orthogonal subspace (which is our main focus of study) will not affect the loss. \n\n\n**For the double rotation symmetry, why it is hard to write the matrix A explicitly?... I am wondering if the results can be generalized to a general matrix A**\n\nThanks for this question. It is just cumbersome and space-taking to write. First of all, to write it in the form 4.1, one needs to view both matrices (U and W) as vectors, and then the corresponding symmetry matrix will be composed by a sequence of block matrices: A=diag(B,...,B,B^{-1},...,B^{-1}), where each block corresponds to a column or row of $U$ or $W$.\n\nRequiring $A$ to be symmetric is actually to ensure that a Noether charge of the form $\\theta^T A \\theta$. Extending the result to antisymmetric matrices is highly nontrivial, and we know very little about this type of symmetry. On the one hand, noether charge does exist for some antisymmetric matrix. For an example, consider rotation symmetry for (u,w) in 2d. The antisymmetric matrix is ((0,-1),(1,0)). The symmetry implies the following dynamics for GD: $wdu - udw = 0$. Simplifying the relation leads to (ignoring the sign problem for now) d(log(u) - log(w)) = 0. We have thus constructed a Noether charge for an antisymmetric-matrix symmetry: log(u) - log(w); however, it is no longer in a quadratic form.\n\nOn the other hand, we do not know a general formula for them or whether they exist for an arbitrary antisymmetric symmetry, which is an important open problem in the field.\n\n**I don’t understand why one needs to assume $C^\\*$ to be a constant in line 163...**\n\nThis is a good question. This is because $C^*$ depends on the value of $\\theta$ in the nondegenerate subspace, and $\\theta$ has a non-zero time derivative in general. To see this mathematically, note that by definition $C^*= C(\\theta_{\\lambda^*})$, but $\\lambda^*=\\lambda^*(\\theta)$ is a function of $\\theta$ as well, and so $C^*$ sometimes depends on time through $\\theta$'s dependence on time (which is often non-zero due to gradient or random noise).\n\nTherefore, a good mental and intuitive picture of the process is that there is a run-and-catch game between $C^*$ and $C$: $C^*$ moves around randomly (due to the motion of $\\theta$), and $C$ chases it systematically due to symmetry.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
# Summary Rebuttal


We thank all the reviewers for their constructive feedback, which we will rely on to further improve our manuscript. We are encouraged to see that both reviewers **rhMv** and **QY9h** rate our contribution as "excellent." We believe some questions and criticisms may arise from misunderstandings or misinterpretations, which we aim to clarify below and in the revised manuscript.

To address the questions and criticisms from reviewers, we will include the following changes and additions in the revision. See the attached pdf for the additional experimental results.

1. **Balanced and Unbalanced norms**. An additional experiment in Figure 7-left shows that the norm balance and imbalance can be controlled by controlling the data distribution, which is in agreement with our theory. This shows one experiment that directly validates our theory and clarifies a misunderstanding of **rhMv** (namely, our theory shows that SGD can lead to both balanced and unbalanced norms of training, not just balanced ones)
2. **Convergence to flatter and sharper solutions**. An additional experiment in Figure 7-right shows that sharpness can be controlled by controlling the data distribution, which is in agreement with our theory. This shows one experiment that directly validates our theory and clarifies another misunderstanding of **rhMv** (namely, our result shows that SGD can lead to both flatter and sharper solutions, not just sharper ones). Both points 1 and 2 are novel insights, and it also novel (and perhaps surprising) that they can be explained by a single theory
3. **Deep linear nets**. Two additional experiments in Figure 8 on deep linear nets trained on MNIST (**QY9h, rhMv**). This directly confirms a theoretical prediction we made: the intermediate layers of a deep linear net converge to the same norm, while the first and last layers will be different
4. **Robustness Across Hyperparameters**. Three sets of additional experiments in Figure 9 show the robustness of the prediction to different hyperparameter choices such as the data dimension, learning rate, and batch size (**HyRa**)
4. **Technical Clarifications**. Additional clarification of technical and mathematical details (**rhMv, QY9h, HyRa**)
5. **Experimental Detail**. Additional detail and clarification of experimental setting (**HyRa**)
6. **References**. Discussion of missed references (**rhMv**)

HyRa named three weaknesses that we believe to be due to misunderstanding: 

1. **Insufficient mathematical preliminaries**: We believe this is due to potential misreading of our equations and unfamiliarity with related mathematical tools and literature within the ML community.
2. **Strong assumption**: This stems from a misunderstanding of the problem setting of section 5.1 -- as we make no assumption in section 5.1. 
3. **Explanation of Contributions**: This may result from unfamiliarity with related fields such as deep learning theory, neuroscience, and stochastic differential equations. Our work is interdisciplinary, requiring a broader perspective beyond practical deep learning to appreciate its contributions fully.

We have addressed these points from HyRa in detail below.

We look forward to your feedback on these revisions and are open to further discussion to enhance the manuscript.

Meta-review: The paper studied the learning dynamics of SGD when there is a symmetry in the loss function (called exponential symmetry).
Reviewer rhMv and QY9h appreciate the contributions of the paper. Reviewer HyRa offered a bolderline review, but agreed upon positive ratings. 

Overall, I recommend accept.
Decision: Accept (poster)