Noise Balance and Stationary Distribution of Stochastic Gradient Descent

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-08 02:51:21

Abstract: How the stochastic gradient descent (SGD) navigates the loss landscape of a neural network remains poorly understood. This work shows that the minibatch noise of SGD regularizes the solution towards a noise-balanced solution whenever the loss function contains a rescaling symmetry. We prove that when the rescaling symmetry exists, the SGD dynamics is limited to only a low-dimensional subspace and prefers a special set of solutions in an infinitely large degenerate manifold, which offers a partial explanation of the effectiveness of SGD in training neural networks. We then apply this result to derive the stationary distribution of stochastic gradient flow for a diagonal linear network with arbitrary depth and width, which is the first analytical expression of the stationary distribution of SGD in a high-dimensional non-quadratic potential. The stationary distribution exhibits complicated nonlinear phenomena such as phase transitions, loss of ergodicity, memory effects, and fluctuation inversion. These phenomena are shown to exist uniquely in deep networks, highlighting a fundamental difference between deep and shallow models. Lastly, we discuss the implication of the proposed theory for the practical problem of variational Bayesian inference.

Main paper: 
 introduction: 
 In natural and social sciences, one of the most important objects of study of a stochastic system is its stationary distribution, which is often found to offer fundamental insights into understanding a given stochastic process [36, 29] . Arguably, a great deal of insights into SGD can be obtained if we have an analytical understanding of its stationary distribution, which remains unknown until today.The stochastic gradient descent (SGD) algorithm is defined as ∆θ t = -η S ∑ x∈B ∇ θ ℓ(θ, x), where θ is the model parameter and ℓ(θ, x) is a per-sample loss whose expectation over x gives the training loss:L(θ) = E x [ℓ(θ, x)].B is a randomly sampled minibatch of data points, each independently sampled from the training set, and S is the minibatch size. Two aspects of the algorithm make it difficult to understand this algorithm: (1) its dynamics is discrete in time, and (2) the randomness is highly nonlinear and parameter-dependent. This work relies on the continuous-time approximation and deals with the second aspect.The main contributions are 1. the derivation of the "law of balance," which shows that SGD converges to a special subset of noised-balanced solutions when the rescaling symmetry is present;2. the first-of-its-kind solution of the stationary distribution of an analytical model trained by SGD;3. discovery of novel phenomena such as phase transitions, loss of ergodicity, memory effects, and fluctuation inversion, all implied by our theory.Organization. The next section discusses the closely related works. In Section 3, we prove the law of balance, the first main result of this work, and discuss its implications for common neural Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.networks. In Section 4, we apply the law of balance to derive the stationary distribution of SGD for a highly nontrivial loss landscape. The last section concludes this work. All proofs and derivations are given in Appendix A. 
related works: 
 Solution of the Fokker Planck (FP) Equation. The FP equation is a high-dimensional partial differential equation whose solution (and its existence) is an open problem in mathematics and many fields of sciences and only known for a few celebrated special cases [28] . Our solution is the first of its kind in a deep-learning setting. Stationary distribution of SGD. One of the earliest works that computes the stationary distribution of SGD is the Lemma 20 of Ref. [3] , which assumes that the noise has a constant covariance and shows that if the loss function is quadratic, then the stationary distribution is Gaussian. Similarly, using a saddle point expansion and assuming that the noise is parameter-independent, a series of recent works showed that the stationary distribution of SGD is exponential in the model parameters close to a local minimum: p(θ) ∝ exp[-aθ T Hθ], for some constant a and matrix H [21, 41, 19] . Assuming that the noise covariance only depends on the loss function value L(θ), Refs. [24] and [39] showed that the stationary distribution is power-law-like and proportional to L(θ) -c0 for some constant c 0 . A primary feature of these previous results is that stationary distribution does not exhibit any memory effect and also preserves ergodicity. Until now, no analytical solution to the stationary distribution of SGD is known, making it impossible to judge how good the previous approximate results are. Our result is the first to derive an exact solution to the stationary distribution of SGD without any approximation. We will see that in contrast to the approximate solutions in the previous results, the actual distribution of SGD has both a memory effect and features the loss of ergodicity.Symmetry and SGD dynamics. Also related to our work is the study of how symmetry affects the learning dynamics of SGD. A major prior work is [17] , which studies the dynamics of SGD when there is scale invariance, conjecturing that SGD reaches a fast equilibrium state at the early stage of training. Our result is different as we study a different type of symmetry, the rescaling symmetry. 
noise balance: 
 We consider the continuous-time limit of SGD [15, 16, 18, 32, 8, 11] :EQUATIONwhere C(θ) = E[∇ℓ(θ)∇ T ℓ(θ)] is the gradient covariance, dW t is a stochastic process satisfying dW t ∼ N (0, Idt) and E[dW t dW T t ′ ] = δ(t-t ′ )I, and T = η/S. Apparently, T gives the average noise level in the dynamics. Previous works have suggested that the ratio T is a main factor determining the behavior of SGD, and using different T often leads to different generalization performance [31, 19, 44] .Due to standard architecture designs, a type of invariance -the rescaling symmetry -often appears in the loss function and it is preserved for all sampling of minibatches. The per-sample loss ℓ is said to have the rescaling symmetry for all x if ℓ(u, w, x) = ℓ (λu, w/λ, x) for a scalar λ ∈ R + . This type of symmetry appears in many scenarios in deep learning. For example, it appears in any neural network with the ReLU activation. It also appears in the self-attention of transformers, often in the form of key and query matrices [37] . When this symmetry exists between u and w, one can prove the following result, which we refer to as the law of balance.Theorem 3.1. Let u, w, and v be parameters of arbitrary dimensions. Let ℓ(u, w, v, x) satisfy ℓ(u, w, v, x) = ℓ(λu, w/λ, v, x) for arbitrary x and any λ ∈ R + . Then,EQUATIONwhereC 1 = E[A T A] -E[A T ]E[A], C 2 = E[AA T ] -E[A]E[A T] and A ki = ∂ l/∂(u i w k ) with l(u i w k , v, x) ≡ ℓ(u i , w k , v, x). 1semi-definite by definition. The theorem still applies if the model has parameters other than u and w. The theorem can be applied recursively when multiple rescaling symmetries exist. See Figure 1 for an illustration the the dynamics and how it differs from other types of GD. (3) where C(w) and C(u) are the gradient covariance of w and u, respectively. Therefore, SGD prefers a solution where the gradient noise between the two layers is balanced. Also, this implies that the balance conditions of the law is only dependent on the diagonal terms of the Fisher information (if we regard the loss as a log probability), which is often well-behaved. As a last caveat, we emphasize that the fact that the noise will balance does not imply that either trace will converge or stay close to a fixed value -it is also possible for both terms to oscillate while their difference is close to zero.The second type is the norm ratio balance between layers, though the norm ratio may not necessarily be finite. Equation (2) implies that in the degenerate direction of the rescaling symmetry, a single and unique point is favored by SGD. Let u = λu * and w = λ -1 w * for arbitrary u * and w * , then, the stationary point of the law is reached atλ 4 = (w * ) T C2w * (u * ) T C1u * .The quantity λ can be called the "balancedness" of the norm, and the law states that when a rescaling symmetry exists, a special balancedness is preferred by the SGD algorithm.When C 1 or C 2 vanishes, λ or λ -1 diverges, and so does SGD. Therefore, having a nonvanishing noise actually implies that SGD training will be more stable. For common problems, C 1 and C 2 are positive definite and, thus, if we know the spectrum of C 1 and C 2 at the end of training, we can estimate a rough norm ratio at convergence:-T (λ 1M ||u|| 2 -λ 2m ||w|| 2 ) ≤ d dt (||u|| 2 -||w|| 2 ) ≤ -T (λ 1m ||u|| 2 -λ 2M ||w|| 2 ),where λ 1m(2m) and λ 1M (2M ) represent the minimal and maximal eigenvalue of the matrix C 1(2) , respectively. Thefore, the value of ||u|| 2 /||w|| 2 is restricted by (See Section A.5)EQUATIONThus, a remaining question is whether the quantities u T C 1 u and w T C 2 w are generally well-defined and nonvanishing or not. The following proposition shows that for a generic two-layer ReLU net, u T C 1 u and w T C 2 w are almost everywhere strictly positive. We define a two-layer ReLU net asEQUATIONwhere u i ∈ R du , w i ∈ R dw and b i is a scalar with i being the index of the hidden neuron. For each i, the model has the rescaling symmetry:u i → λu i , (w i , b i ) → (λ -1 w i , λ -1 b i ).We thus apply the law of balance to each neuron separately. The per-sample loss function isℓ(θ, x) = ∥f (x) -y(x, ϵ)∥ 2 .(6) Here, x has a full-rank covariance Σ x , and y = g(x) + ϵ for some function g and ϵ is a zero-mean random vector independent of x and have the full-rank covariance Σ ϵ . The following theorem shows that for this network, C 1 and C 2 are full rank unless the neuron is "dead". The dataset is a synthetic data set, where x is drawn from a normal distribution, and the labels: y = x+ϵ, for an independent Gaussian noise ϵ with unit variance. While every neuron has a rescaling symmetry, we focus on the overall rescaling symmetry between the two weight matrices. The norm between the two layers reach a state of approximate balance -but not a precise balance. At the same time, the model evolves during training towards a state where u T C 1 u and w T C 2 w are balanced.Standard analysis shows that the difference between SGD and GD is of order T 2 per unit time step, and it is thus often believed that SGD can be understood perturbatively through GD [11] . However, the law of balance implies that the difference between GD and SGD is not perturbative. As long as there is any level of noise, the difference between GD and SGD at stationarity is O(1). This theorem also implies the loss of ergodicity, an important phenomenon in nonequilibrium physics [26, 34, 22, 35] , because not all solutions with the same training loss will be accessed by SGD with equal probability.The theorem greatly simplifies when both u and w are one-dimensional.Corollary 3.3. If u, w ∈ R, then, d dt |u 2 -w 2 | = -T C 0 |u 2 -w 2 |, where C 0 = Var[ ∂ℓ ∂(uw) ].Before we apply the theorem to study the stationary distributions, we stress the importance of this balance condition. This relation is closely related to Noether's theorem [23, 1, 20] . If there is no weight decay or stochasticity in training, the quantity ||u|| 2 -||w|| 2 will be a conserved quantity under gradient flow [6, 14, 33] , as is evident by taking the infinite S limit. The fact that it monotonically decays to zero at a finite T may be a manifestation of some underlying fundamental mechanism. A more recent result in Ref. [38] showed that for a two-layer linear network, the norms of two layers are within a distance of order O(η -1 ), suggesting that the norm of the two layers are balanced. Our result agrees with Ref. [38] in this case, but our result is stronger because our result is nonperturbative, only relies on the rescaling symmetry, and is independent of the loss function or architecture of the model. It is useful to note that when L 2 regularization with strength γ is present, the rate of decay changes from T C 0 to T C 0 + γ. This points to a nice interpretation that when rescaling symmetry is present, the implicit bias of SGD is equivalent to weight decay. See Figure 1 for an illustration of this point.Example: two-layer linear network. It is instructive to illustrate the application of the law to a two-layer linear network, the simplest model that obeys the law. Let θ = (w, u) denote the set of trainable parameters; the per-sample loss is ℓ(θ, x) = (∑ d i u i w i x -y) 2 + γ||θ|| 2 . Here, d is the width of the model, γ||θ|| 2 is the L 2 regularization term with strength γ ≥ 0, and E x denotes the averaging over the training set, which could be a continuous distribution or a discrete sum of delta distributions. It will be convenient for us also to define the shorthand: v ∶= ∑ d i u i w i . The distribution of v is said to be the distribution of the "model." Applying the law of balance, we obtain thatEQUATIONwhere we have introduced the parametersEQUATIONWhen α 1 α 3 -α 2 2 or γ > 0, the time evolution of |u 2 -w 2 | can be upper-bounded by an exponentially decreasing function in time:|u 2 i -w 2 i |(t) < |u 2 i -w 2 i |(0) exp (-4T (α 1 α 3 -α 2 2 )t/α 1 -4γt) → 0.Namely, the quantity (u 2 i -w 2 i ) decays to 0 with probability 1. We thus have u 2 i = w 2 i for all i ∈ {1, ⋯, d} at stationarity, in agreement with the Corollary. 
stationary distribution of sgd: 
 As an important application of the law of balance, we solve the stationary distribution of SGD for a deep diagonal linear network. While linear networks are limited in expressivity, their loss landscape and dynamics are highly nonlinear and exhibits many shared phenomenon with nonlinear neural networks [13, 30] . Let θ follow the high-dimensional Wiener process given by Eq.( 1). The probability density evolves according to its Kolmogorov forward (Fokker-Planck) equation:EQUATIONThe solution of this partial differential equation is an open problem for almost all high-dimensional problems. This section solves it for a high-dimensional non-quadratic potential of a machine learning relevance.Let us first derive the stationary distribution of a one-dimensional linear regressor, which will be a basis for comparison to help us understand what is unique about having a "depth" in deep learning.The per-sample loss is ℓ(x, v) = (vx -y) 2 + γv 2 . DefiningEQUATIONthe global minimizer of the loss can be written as:v * = β 2 /β 1 .The gradient variance is also not trivial: C(v) ∶= Var[∇ v ℓ(v, x)] = 4(α 1 v 2 -2α 2 v + α 3 ). NoteEQUATIONWhen is ∆ zero? It happens when, for all samples of (x, y), xy + c = kx 2 for some constant k and c. We focus on the case ∆ > 0 in the main text, which is most likely the case for practical situations.The other cases are dealt with in Section A.For ∆ > 0, the stationary distribution for linear regression is (Section A)EQUATIONin agreement with the previous result [24] . Two notable features exist for this distribution: (1) the power exponent for the tail of the distribution depends on the learning rate and batch size, and(2) the integral of p(v) converges for an arbitrary learning rate. On the one hand, this implies that increasing the learning rate alone cannot introduce new phases of learning to a linear regression; on the other hand, it implies that the expected error is divergent as one increases the learning rate (or the feature variation), which happens at T = β ′ 1 /α 1 . We will see that deeper models differ from the single-layer model in these two crucial aspects.Now, we consider the following model with a notion of depth and width; its loss function is where D can be regarded as the depth and d 0 the width. When the width d 0 = 1, the law of balance is sufficient to solve the model. When d 0 > 1, we need to eliminate additional degrees of freedom. We note that this model conceptually resembles (but not identical to) a diagonal linear network, which has been found to well approximate the dynamics of real networks [27, 25, 2, 7] .EQUATIONWe introducev i ∶= ∏ D k=0 u (k)i , and so v = ∑ i v i , where we call v i a "subnetwork" and v the "model."The following proposition shows that independent of d 0 and D, the dynamics of this model can be reduced to a one-dimensional form by invoking the law of balance.Theorem 4.1. For all i ≠ j, one (or more) of the following conditions holds for all trajectories at stationarity: (1)v i = 0, or v j = 0, or L(θ) = 0; (2) sgn(v i ) = sgn(v j ). In addition, (2a) if D = 1, for a constant c 0 , log |v i | -log |v j | = c 0 ; (2b) if D > 1, |v i | 2 -|v j | 2 = 0.This theorem contains many interesting aspects. First of all, the three situations in item 1 directly tell us the distribution of v if the initial state of of v is given by these conditions. 2 This implies a memory effect, namely, that the stationary distribution of SGD can depend on its initial state. The second aspect is the case of item 2, which we will solve below. Item 2 of the theorem implies that all the v i of the model must be of the same sign for any network with D ≥ 1. Namely, no subnetwork of the original network can learn an incorrect sign. This is dramatically different from the case of D = 0. We will discuss this point in more detail below. The third interesting aspect of the theorem is that it implies that the dynamics of SGD is qualitatively different for different depths of the model.In particular, D = 1 and D > 1 have entirely different dynamics. For D = 1, the ratio between every pair of v i and v j is a conserved quantity. In sharp contrast, for D > 1, the distance between different v i is no longer conserved but decays to zero. Therefore, a new balancing condition emerges as we increase the depth. Conceptually, this qualitative distinction also corroborates the discovery in Ref. [43] , where D = 1 models are found to be qualitatively different from models with D > 1.With this theorem, we are ready to solve the stationary distribution. It suffices to condition on the event that v i does not converge to zero. Let us suppose that there are d nonzero v i that obey item 2 of Theorem 4.1 and d can be seen as an effective width of the model. We stress that the effective width d ≤ d 0 depends on the initialization and can be arbitrary. 3 Therefore, we condition on a fixed value of d to solve for the stationary distribution of v (Appendix A): EQUATIONwhere p -is the distribution on (-∞, 0) and p + is that on (0, ∞), and g∓ (v) = α 1 |v| 2 ∓ 2α 2 |v| + α 3 .The arbitrariness of the scalar z is due to the memory effect of SGD -if all parameters are initialized at zero, they will remain there with probability 1. This means that the stationary distribution is not unique. Since the result is symmetric in the sign ofβ 2 = E[xy], we assume that E[xy] > 0 from now on.Also, we focus on the case γ = 0 in the main text. 4 The distribution of v isEQUATION)This measure is worth a close examination. First, the exponential term is upper and lower bounded and well-behaved in all situations. In contrast, the polynomial term becomes dominant both at infinity and close to zero. When v < 0, the distribution is a delta function at zero: p(v) = δ(v). To see this, note that the term v -β2/2α3T -3/2 integrates to give v -β2/2α3T -1/2 close to the origin, which is infinite. Away from the origin, the integral is finite. This signals that the only possible stationary distribution has a zero measure for v ≠ 0. The stationary distribution is thus a delta distribution, meaning that if x and y are positively correlated, the learned subnets v i can never be negative, independent of the initial configuration.For v > 0, the distribution is nontrivial. Close to v = 0, the distribution is dominated by v β2/2α3T -3/2 , which integrates to v β2/2α3T -1/2 . It is only finite below a critical T c = β 2 /α 3 . This is a phasetransition-like behavior. As T → (β 2 /α 3 ) -, the integral diverges and tends to a delta distribution.Namely, if T > T c , we have u i = w i = 0 for all i with probability 1, and no learning can happen.If T < T c , the stationary distribution has a finite variance, and learning may happen. In the more general setting, where weight decay is present, this critical T shifts to T c = β2-γ α3 . When T = 0, the phase transition occurs at β 2 = γ, in agreement with the threshold weight decay identified inRef. [45] . See Figure 3 When T < T c , a nonzero maximizer for p(v) must satisfyEQUATION)The existence of this solution is nontrivial, which we analyze in Appendix A.8. When T → 0, a solution always exists and is given by v = β 2 /β 1 , which does not depend on the learning rate or noise C. Note that β 2 /β 1 is also the minimum point of L(u i , w i ). This means that SGD is only a consistent estimator of the local minima in deep learning in the vanishing learning rate limit. How biased is SGD at a finite learning rate? Two limits can be computed. For a small learning rate, the leading order correction to the solution is v = β2 β1 + ( 10α2β2β 2 1 - 7α1β 2 2 β 3 1 -3α3 β1 )T . This implies that the common Bayesian analysis that relies on a Laplace expansion of the loss fluctuation around a local minimum is improper. The fact that the stationary distribution of SGD is very far away from the Bayesian posterior also implies that SGD is only a good Bayesian sampler at a small learning rate.Example. It is instructive to consider an example of a structured dataset: y = kx + ϵ, where x ∼ N (0, 1) and the noise ϵ obeys ϵ ∼ N (0, σ 2 ). We let γ = 0 for simplicity. If An interesting aspect of the depth-1 model is that its distribution is independent of the width d of the model. This is not true for a deep model, as seen from Eq. ( 14). should be a very large value [42, 19] , and therefore L should diverge. Thus, the edge of stability phenomenon is incompatible with the naive expectation up to the second order, as pointed out byRef. [5] . Our theory offers a direct explanation of why the divergence of loss does not happen: for deeper models, the fluctuation of model parameters decreases as the gradient noise level increases, reaching a minimal value before losing stability. Thus, SGD always has a finite loss because of the power-law tail and fluctuation inversion. See Figure 5 -mid.Infinite-D limit. As D tends to infinity, the distribution becomesp(v) ∝ 1 v 3+k 1 (α1v 2 -2α2v + α3) 1-k 1 /2 exp ( - d DT ( β2 α3v + α2α3β1 -2α 2 2 β2 + α1α3β2 α 2 3 √ ∆ arctan( α1v -α2 √ ∆ )) ),wherek 1 = d(α 3 β 1 -2α 2 β 2 )/(T Dα 2 3). An interesting feature is that the architecture ratio d/D always appears simultaneously with 1/T . This implies that for a sufficiently deep neural network, the ratio D/d also becomes proportional to the strength of the noise. Since we know that T = η/S determines the performance of SGD, our result thus shows an extended scaling law of training:d D S η = const.The architecture aspect of the scaling law also agrees with an alternative analysis [9, 10] , where the optimal architecture is found to have a constant ratio of d/D. See Figure 5 . If d = o(D), k 1 → 0 and the distribution converges to p(v) ∝ v -3 (α 1 v 2 -2α 2 v + α 3 ) -1, which is a delta distribution at 0. Namely, if the width is far smaller than the depth, the model will collapse to One of the major implications of the analytical solution we found for machine learning practice is the inappropriateness of using SGD to approximate a Bayesian posterior. Because every SGD iteration can be regarded as a sampling of the model parameters. A series of recent works have argued that the stationary distribution can be used as an approximation of the Bayesian posterior for fast variational inference [21, 3] , p Bayes (θ) ≈ p SGD (θ), a method that has been used for a wide variety of applications [12] . However, our result implies that such an approximation is likely to fail. Common in Bayesian deep learning, we interpret the per-sample loss as the log probability and the weight decay as a Gaussian prior over the parameters, the true model parameters have a log probability ofEQUATION)This distribution has a nonzero measure everywhere for any differentiable loss. However, the distribution for SGD in Eq.( 14) has a zero probability density almost everywhere because a 1d subspace has a zero Lebesgue measure in a high-dimensional space. This implies that the KL divergence between the two distributions (either KL(p Bayes ||p SGD ) or KL(p SGD ||p Bayes )) is infinite. Therefore, we can infer that in the information-theoretic sense, p SGD cannot be used to approximate p Bayes . 
discussion: 
 In this work, we first showed that SGD systematically moves towards a balanced solution when rescaling symmetry exists, a result we termed the law of balance. Applying the law of balance, we have characterized the stationary distribution of SGD analytically, which is an unanswered fundamental problem in the study of SGD. This is the first analytical expression for a globally nonconvex and beyond quadratic loss without the need for any approximation. With this solution, we have discovered many phenomena that could be relevant to deep learning that were previously unknown.We found that SGD only has probability of exploring a one-dimensional submanifold even for a very-dimensional problem, ignoring all irrelevant directions. We applied our theory to the important problem of variational inference and showed that it is, in general, not appropriate to approximate the posterior with SGD, at least when any symmetry is present in the model. If one really wants to use SGD for variational inference, special care is required to at least remove symmetries from the loss function, which could be an interesting future problem. Our theory is limited, as the model we solved is only a minimal model of reality, and it would be interesting to consider more realistic models in the future. Also, it would be interesting to extend the law of balance to a broader class of symmetries.A.1 BackgroundLet us consider the following stochastic differential equation (SDE) for a Wiener process W (t):EQUATIONWe are interested in the dynamics of a generic function of X t . Let Y t = f (t, X t ); Ito's lemma states that the SDE for the new variable isEQUATIONLet us take the variable Y t = X 2 t as an example. Then the SDE isEQUATIONLet us consider another example. Let two variables X t and Y t followEQUATIONThe SDE of X t Y t is given byEQUATIONThe general SDE of a 1d variable X is given by: dX = -µ(X)dt + B(X)dW (t).The time evolution of the probability density P (x, t) is given by the Fokker-Planck equation:EQUATIONwhere J(X, t) = µ(X)P (X, t)EQUATIONwhich gives a solution as a Boltzmann-type distribution if B is a constant. We will apply Eq. ( 25)to determine the stationary distributions in the following sections.A.2 Proof of Theorem 3.1Proof. We omit writing v in the argument unless necessary. By definition of the symmetry ℓ(u, w, x) = ℓ(λu, w/λ, x), we obtain its infinitesimal transformation ℓ(u, w, x) = ℓ((1 + ϵ)u, (1ϵ)w/λ, x). Expanding this to first order in ϵ, we obtainEQUATION)The equations of motion areEQUATION)EQUATION)Using Ito's lemma, we can find the equations governing the evolutions of u 2 i and w 2 j :EQUATIONwhere C u i = Var[ ∂ℓ ∂ui ] and C w j = Var[ ∂ℓ ∂wj ]. With Eq. ( 26), we obtainEQUATIONDue to the rescaling symmetry, the loss function can be considered as a function of the matrix uw T .Here we define a new loss function as l(u i w j ) = ℓ(u i , w j ). Hence, we haveEQUATION)We can rewrite Eq. ( 30) intoEQUATIONwhereEQUATIONEQUATIONwhereEQUATION)The proof is thus complete.Considering the modified loss function:EQUATION)In this case, the Langevin equations becomeEQUATION)EQUATIONHence, the modified SDEs of u 2 i and w 2 j can be rewritten asEQUATION)EQUATION)In this section, we consider the effects brought by the last term in Eqs. ( 39) and (40) . From the infinitesimal transformation of the rescaling symmetry:EQUATIONwe take the derivative of both sides of the equation and obtainEQUATIONEQUATIONwhere we take the expectation to ℓ at the same time. By substituting these equations into Eqs. (39) and ( 40), we obtainEQUATION)Then following the procedure in Appendix. A.2, we can rewrite Eq. ( 44) asEQUATIONwhereEQUATIONEQUATIONEQUATIONEQUATIONFor one-dimensional parameters u, w, Eq. ( 45) is reduced toEQUATIONTherefore, we can see this loss modification increases the speed of convergence. Now, we move to the stationary distribution of the parameter v. At the stationarity, if u i = -w i , we also have the distribution P (v) = δ(v) like before. However, when u i = w i , we havedv dt = -4v(β 1 v-β 2 )+4T v(α 1 v 2 -2α 2 v+α 3 )-4β 2 1 T v(β 1 v-β 2 )(3β 1 v-β 2 )+4v √ T (α 1 v 2 -2α 2 v + α 3 ) dW dt .(51) Hence, the stationary distribution becomesEQUATIONEQUATION)From the expression above we can seeK 1 ≪ 1 + β 2 /4T α 3 and K 2 ≪ (α 3 β 1 -α 2 β 2 )/2T α 3 √ ∆.Hence, the effect of modification can only be seen in the term proportional to v. The phase transition point is modified asT c = β 2 α 3 + β 2 2 .(54)Compared with the previous result T c = β2 α3 , we can see the effect of the loss modification is α 3 →α 3 + β 2 2 , or equivalently, Var[xy] → E[x 2 y 2 ]. This effect can be seen from E 1 and E 2 .A.4 Proof of Theorem 3.2Proof. For any i, one can obtain the expressions of C (i)1 and C (i)2 from Theorem 3.1 asEQUATIONEQUATIONwhere we use the notation r α ∶= ∑d j=1 u α j v T j x -y α , x ∶= (x T , 1) T , v i = (w T i , b i ) T and E i [O] ∶= E[O|w T i x + b i > 0].We start with showing that C(1) 1is full-rank. Let m be an arbitrary unit vector in R du . We have thatEQUATIONwhere the last inequality follows fromEQUATION)Here we denote that Vari [O] ∶= E i [O 2 ] -E i [O] 2 and Cov i [O 1 , O 2 ] ∶= E i [O 1 O 2 ] -E i [O 1 ]E i [O 2 ].ForC (i)2 , we let the vector ñ ∶= (n T , n f ) T be a unit vector in R dw+1 , yielding ñT C (i)EQUATIONNote that this quantity can be decomposed asEQUATION)The covariance term vanishes becauseEQUATIONwhere the penultimate inequality follows from the fact that ϵ is independent of x. Hence, both the matrices C(i)1 and C (i)2 are full-rank. The proof is completed.A.5 Derivation of Eq. (4)We here prove inequality (4). At stationarity, d(∥u∥ 2 -∥w∥ 2 )/dt = 0, indicatingEQUATION)The first inequality in Eq. ( 63) gives the solutionEQUATION)The second inequality in Eq. ( 63) gives the solutionEQUATION)Combining these two results, we obtainEQUATIONwhich is Eq. ( 4).Proof. This proof is based on the fact that if a certain condition is satisfied for all trajectories with probability 1, this condition is satisfied by the stationary distribution of the dynamics with probability 1.Let us first consider the case of D > 1. We first show that any trajectory satisfies at least one of the following five conditions: for any i, (i) v i → 0, (ii) L(θ) → 0, or (iii) for any k ≠ l, (u(k) i ) 2 - (u (l) i ) 2 → 0.The SDE for uEQUATIONwherev i ∶= ∏ D k=1 u (k)i , and so v = ∑ i v i . There exists rescaling symmetry between u EQUATION)whereVar ⎡ ⎢ ⎢ ⎢ ⎢ ⎣ ∂ℓ ∂(u (k) i u (l) i ) ⎤ ⎥ ⎥ ⎥ ⎥ ⎦ = ( v i u (k) i u (l) i ) 2 (α 1 v 2 -2α 2 v + α 3 ) (69) with v i /(u (k) i u (l) i ) = ∏ s≠k,l u (s)i . In the long-time limit, (u(k) i ) 2 converges to (u (l) i ) 2 unless Var [ ∂ℓ ∂(u (k) i u (l) i ) ] = 0, which is equivalent to v i /(u (k) i u (l) i ) = 0 or α 1 v 2 -2α 2 v + α 3 = 0.These two conditions correspond to conditions (i) and (ii). The latter is because α 1 v 2 -2α 2 v + α 3 = 0 takes place if and only if v = α 2 /α 1 and α 2 2 -α 1 α 3 = 0 together with L(θ) = 0. Therefore, at stationarity, we must have conditions (i), (ii), or (iii). Now, we prove that when (iii) holds, the condition 2-(b) in the theorem statement must hold: forD = 1, (log |v i | -log |v j |) = c 0 with sgn(v i ) = sgn(v j ).When (iii) holds, there are two situations.First, if v i = 0, we have u ( i k) = 0 for all k, and v i will stay 0 for the rest of the trajectory, which corresponds to condition (i).If v i ≠ 0, we have u(k) i≠ 0 for all k. Therefore, the dynamics of v i isdv i dt = -2 ∑ k ⎛ ⎝ v i u (k) i ⎞ ⎠ 2 (β 1 v-β 2 )+2 ∑ k ⎛ ⎝ v i u (k) i ⎞ ⎠ 2 √ η(α 1 v 2 -2α 2 v + α 3 ) dW dt +4 ∑ k,l ⎛ ⎝ v 3 i (u (k) i u (l) i ) 2 ⎞ ⎠ η(α 1 v 2 -2α 2 v+α 3 ).(70) Comparing the dynamics of v i and v j for i ≠ j, we obtainEQUATIONBy condition (iii), we have |u(0) i | = ⋯ = |u (D) i |, i.e., (v i /u (k) i ) 2 = (v 2 i ) D/(D+1) and (v i /u (m) i u (l) i ) 2 = (v 2 i ) (D-1)/(D+1). 5 Therefore, we obtaindv i /dt (D + 1)(v 2 i ) D/(D+1) - dv j /dt (D + 1)(v 2 j ) D/(D+1) = ⎛ ⎝ v i D(v 2 i ) (D-1)/(D+1) 2(v 2 i ) D/(D+1) -v j D(v 2 j ) (D-1)/(D+1) 2(v 2 j ) D/(D+1) ⎞ ⎠ η(α 1 v 2 -2α 2 v+α 3 ).(72) We first consider the case where v i and v j initially share the same sign (both positive or both negative). When D > 1, the left-hand side of Eq. (72) can be written as 5 Here, we only consider the root on the positive real axis.EQUATIONwhich follows from Ito's lemma:EQUATIONSubstitute in Eq. ( 72), we obtain Eq. ( 73). Now, we consider the right-hand side of Eq. ( 72), which is given byEQUATIONCombining Eq. ( 73) and Eq. ( 75), we obtain1 1 -D dv 2/(D+1)-1 i dt - 1 1 -D dv 2/(D+1)-1 j dt = -2D(v 1-2/(D+1) i -v 1-2/(D+1) j )η(α 1 v 2 -2α 2 v + α 3 ).(76) By definingz i = v 2/(D+1)-1 i, we can further simplify the dynamics:EQUATIONHence,EQUATIONTherefore, if v i and v j initially have the same sign, they will decay to the same value in the longtime limit t → ∞, which gives condition 2-(b). When v i and v j initially have different signs, we can write Eq. (72) asEQUATIONHence, when D > 1, we simplify the equation with a similar procedure asEQUATIONwhich impliesEQUATIONFrom this equation, we reach the conclusion that if v i and v j have different signs initially, one of them converges to 0 in the long-time limit t → ∞, corresponding to condition 1 in the theorem statement. Hence, for D > 1, at least one of the conditions is always satisfied at t → ∞. Now, we prove the theorem for D = 1, which is similar to the proof above. The law of balance givesEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONTherefore, the stationary distribution of a general deep diagonal network is given byp(v) ∝ 1 v 3(1-1/(D+1)) (α 1 v 2 -2α 2 v + α 3 ) exp (- 1 T ∫ dv d 1-2/(D+1) (β 1 v -β 2 ) (D + 1)v 2D/(D+1) (α 1 v 2 -2α 2 v + α 3 )) .(94)If v < 0, Eq. ( 92) becomesEQUATIONThe stationary distribution of |v| is given byp(|v|) ∝ 1 |v| 3(1-1/(D+1)) (α 1 |v| 2 + 2α 2 |v| + α 3 ) exp (- 1 T ∫ d|v| d 1-2/(D+1) (β 1 |v| + β 2 ) (D + 1)|v| 2D/(D+1) (α 1 |v| 2 + 2α 2 |v| + α 3 )) .(96) Thus, we have obtainedp ± (|v|) ∝ 1 |v| 3(1-1/(D+1)) (α 1 |v| 2 ∓ 2α 2 |v| + α 3 ) exp (- 1 T ∫ d|v| d 1-2/(D+1) (β 1 |v| ∓ β 2 ) (D + 1)|v| 2D/(D+1) (α 1 |v| 2 ∓ 2α 2 |v| + α 3 )) .(97) Especially when D = 1, the distribution function can be simplified asEQUATIONwhere we have used the integral∫ dv β 1 v ∓ β 2 α 1 v 2 -2α 2 v + α 3 = α 3 β 1 -α 2 β 2 α 3 √ ∆ arctan α 1 |v| ∓ α 2 √ ∆ ± β 2 α 3 log(v)± β 2 2α 3 log(α 1 v 2 -2α 2 v+α 3 ).(99) Furthermore, we can also see that p(v) = δ(v) is also the stationary distribution of the Fokker-Planck equation of Eq. (93). Hence, the general stationary distribution of v can be expressed asEQUATION)The proof is complete.To investigate the existence of the maximum point given in Eq. ( 16), we treat T as a variable and study whether (β 1 -10α 2 T ) 2 + 28α 1 T (β 2 -3α 3 T ) ∶= A in the square root is always positive or not.When T < β2 3α3 = T c /3, A is positive for arbitrary data. When T > β2 3α3 , we divide the discussion into several cases. First, when α 1 α 3 > 25 21 α 2 2 , there always exists a root for the expression A. Hence, we find thatEQUATIONis a critical point. When T c /3 < T < T * , there exists a solution to the maximum condition. When T > T * , there is no solution to the maximum condition.The second case is α 2 2 < α 1 α 3 < 25 21 α 2 2 . In this case, we need to further compare the value between 5α 2 β 1 and 7α 1 β 2 . If 5α 2 β 1 < 7α 1 β 2 , we have A > 0, which indicates that the maximum point exists. If 5α 2 β 1 > 7α 1 β 2 , we need to further check the value of minimum of A, which takes the form ofEQUATION) If 7α1 5α2 < β1 β2 < 5α2+ √ 25α 2 2 -21α1α3 3α3, the minimum of A is always positive and the maximum exists. However, if β1 β2 ≥5α2+ √ 25α 2 2 -21α1α3 3α3, there is always a critical learning rate T * . If without weight decay with weight decay single layer(α 1 v 2 -2α 2 v + α 3 ) -1- β 1 2T α 1 α 1 (v -k) -2- (β 1 +γ) T α 1 non-interpolation v β 2 /2α 3 T -3/2 (α1v 2 -2α2v+α3) 1+β 2 /4T α 3 v S(β 2 -γ)/2α 3 λ-3/2 (α1v 2 -2α2v+α3) 1+(β 2 -γ)/4T α 3 interpolation y = kx v -3/2+β 1 /2T α 1 k (v-k) 2+β 1 /2T α 1 k v -3/2+ 1 2T α 1 k (β 1 - γ k ) (v-k) 2+ 1 2T α 1 k (β 1 - γ k ) exp(-βγ 2T α1 1 k(k-v)) Table 1 : Summary of distributions p(v) in a depth-1 neural network. Here, we show the distribution in the nontrivial subspace when the data x and y are positively correlated. The Θ(1) factors are neglected for concision.β1 β2 = 5α2+ √ 25α 2 2 -21α1α3 3α3, there is only one critical learning rate as T c = 5α2β1-7α1β2 2(25α 2 2 -21α1α3) . When T c /3 < T < T * , there is a solution to the maximum condition, while there is no solution whenT > T * . If β1 β2 > 5α2+ √ 25α 2 2 -21α1α3 3α3, there are two critical points:EQUATION)For T < T 1 and T > T 2 , there exists a solution to the maximum condition. For T 1 < T < T 2 , there is no solution to the maximum condition. The last case is α 2 2 = α 1 α 3 < 25 21 α 2 2 . In this sense, the expression of A is simplified asβ 2 1 + 28α 1 β 2 T -20α 2 β 1 T .Hence, when β1 β2 < 7α1 5α2 , there is no critical learning rate and the maximum always exists. Nevertheless, when β1 β2 > 7α1 5α2 , there is always a critical learning rate asT * = β 2 120α2β1-28α1β2 . When T < T * , there is a solution to the maximum condition, while there is no solution when T > T * .A.9 Other Cases for D = 1The other cases are worth studying. For the interpolation case where the data is linear (y = kx for some k), the stationary distribution is different and simpler. There exists a nontrivial fixed point for ∑ i (u 2 i -w 2 i ): ∑ j u j w j = α2 α1 , which is the global minimizer of L and also has a vanishing noise. It is helpful to note the following relationships for the data distribution when it is linear:EQUATIONSince the analysis of the Fokker-Planck equation is the same, we directly begin with the distribution function in Eq. ( 15) for u i = -w i which is given by P (|v|) ∝ δ(|v|). Namely, the only possible weights are u i = w i = 0, the same as the non-interpolation case. This is because the corresponding stationary distribution isEQUATION)The integral of Eq. ( 105) with respect to |v| diverges at the origin due to the factor |v|3 2 + β 1 2T α 1 k .For the case u i = w i , the stationary distribution is given from Eq. (15) asEQUATIONNow, we consider the case of γ ≠ 0. In the non-interpolation regime, when u i = -w i , the stationary distribution is still p(v) = δ(v). For the case of u i = w i , the stationary distribution is the same as in Eq. ( 15) after replacing β with β ′ 2 = β 2 -γ. It still has a phase transition. The weight decay has the effect of shifting β 2 by -γ. In the interpolation regime, the stationary distribution is still p(v) = δ(v) when u i = -w i . However, when u i = w i , the phase transition still exists since the stationary distribution isp(v) ∝ v -3 2 +θ2(v -k) 2+θ2 exp (-EQUATIONwhereθ 2 = 1 2T α1k (β 1 -γ k ).The phase transition point is θ 2 = 1/2, which is the same as the noninterpolation one.The last situation is rather special, which happens when ∆ = 0 but y ≠ kx: y = kx -c/x for some c ≠ 0. In this case, the parameters α and β are the same as those given in Eq. ( 104) except for β 2 :EQUATIONThe corresponding stationary distribution isEQUATIONwhereϕ 2 = 1 2T α1k (β 1 -c).Here, we see that the behavior of stationary distribution P (|v|) is influenced by the sign of c. When c < 0, the integral of P (|v|) diverges due to the factor |v| -3 2 -ϕ2 < |v| -3/2 and Eq. (109) becomes δ(|v|) again. However, when c > 0, the integral of |v| may not diverge.The critical point is 3 2 + ϕ 2 = 1 or equivalently: c = β 1 + T α 1 k. This is because when c < 0, the data points are all distributed above the line y = kx. Hence, u i = -w i can only give a trivial solution.However, if c > 0, there is the possibility to learn the negative slope k. When 0 < c < β 1 + T α 1 k, the integral of P (|v|) still diverges and the distribution is equivalent to δ(|v|). Now, we consider the case of u i = w i . The stationary distribution isEQUATIONIt also contains a critical point: - Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on 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 new proposed method and baselines. If only a 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 to the paper) is recommended, but including URLs to data and code is permitted. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• 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 (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?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 the type of compute workers CPU or GPU, internal cluster, or 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 in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 a fundamental research on the dynamics of SGD and hence it does not have direct positive or negative societal impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• 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 specific 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 the 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 mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?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 a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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.This result also holds using the modified loss (See Appendix A.3).L → 0 is only possible when ∆ = 0 and v = β2/β1.One can initialize the parameters such that d takes any value between 1 and d0. One way to achieve this is to initialize on the stationary points specified by Theorem4.1 at the desired d.When weight decay is present, the stationary distribution is the same, except that one needs to replace β2 with β2 -γ. Other cases are also studied in detail in Appendix A and listed in Table.1. 
 
 Appendix:  


Decision: This paper investigates the behavior of SGD, demonstrating its convergence to a noise-balanced solution. Subsequently, it discusses detailed properties of the stationary distribution of SGD. Generally, reviewers recognize the correctness and theoretical rigor of this paper, yet express concerns about the clarity and limitations of the theoretical models, and thus do not strongly support it. The authors are encouraged to substantially revise the paper to improve clarity and extend the theoretical results to a more general setting.

Reviews: review: paper_summary: paper_summary: This paper tries to analyze the specific features that carry the noise of SGD (through a continuous model). The authors show that there is a certain 'law of balance' across the layers when some invariance is assumed. Going further, they derive a toy model to push their study, showing that there is an analytic stationary solution to it. They finally propose a phenomenology related to the role of the noise of SGD when analyzing this precise stationary distribution.
strength_weakness: strength_weakness: The paper present the following weaknesses:

- The law of balance is an interesting phenomenon, yet considering it with a closer look, it seems that not much can be said generally and that one has to understand it case by case. In one dimension, sure, it is possible to conclude that balancedness will occur at exponential speed, yet in dimension more than $2$, it seems impossible to predict it surely.

- I have to say that I was a bit bothered by the general overselling of the paper : 
     - As said before the law of balance is truly valid asymptotically in one-dimension
     - The stationary distribution that the authors claim to be the first to derive is for a very specific model, which is not standard and does not resemble a diagonal network! 
     - The fact that the stationary distribution can be computed is also very inherent to $1d$ calculation and is simply a recognition of a Pearson diffusion that already made in way in ML (at least in https://arxiv.org/pdf/2402.01382 and https://arxiv.org/pdf/2407.02322).

Minor typos/flaws:

- l.41: Fokker Planck is not inherently high-dimensional
- l.44: Go to the line for new paragraph 
- l.165: The law of balance is not strictly applicable here since $\ell$ is not scale invariant because of the regularization.
- Section **4.1 Depth - 0**: I think that $\Delta > 0$ is not currently the "most practical example" since it corresponds to a underparametrized model.
questions: questions: On top of the questions raised by the weakness section, here a some (more minor) questions:

- l.65: Eq (1) is not the usual covariance matrix that is used for Stochastic modified equations.
- l.95: Eq (3) seems important but it is difficult to follow where this comes from without intermediate calculations, can the authors develop?
- l.110: Same thing for the equation $\lambda^4 = \frac{\langle w^*, C_2 w^* \rangle}{\langle w^*, C_1 w^* \rangle}$ (3): can the authors develop? 
- Thm 4.2 : Are they the only stationary distribution ? How to know to which one it converges ?
limitations: limitations: As said before, all conclusion are drawn for models that live intrinsically in one dimension.
paper_summary: paper_summary: For ReLU networks trained by gradient flows, it is classical that a type of Minkowski inner product between the coefficients of consecutive layers is preserved. The authors demonstrate a monotonicity of the same quantity for stochastic gradient descent in continuous time. They use this to study the invariant distribution of parameters trained by (continuous time) SGD.
strength_weakness: strength_weakness: * In its current form, I find the article a bit unpolished and the results not easy to access. Many questions remained unanswered when I tried reading the article (see questions).

* Important quantities are defined throughout the plain text. I understand that reading as a reviewer under time pressure is different from normal reading, but for instance in Theorem 3.1, I would have hoped for a more self-contained statement on relations and properties $L, C, \ell$ and the distribution of $x$ have to satisfy. As far as I can tell, the statement is fairly general and not specific to machine learning.

* I have serious doubts about Theorem 3.1. It is derived in Appendix A from Itô's Lemma without the diffusion term. This is valid *in expectation over $\theta$*, but not pointwise in $\theta$. Pointwise in $\theta$, there should be white noise in the 'time derivatives', i.e. the ODE identity should be written as an SDE. In the proof, equations (27) and (28) appear to be wrong.

* The authors do not pay any attention to whether solutions to the evolution equations exist (or are unique). Problems with regularity can sometimes be alleviated if the distribution in $x$ is sufficiently regular, but I would appreciate a short discussion.
questions: questions: * The noise matrix $C$ defined below (1) is an uncentered covariance matrix, while the centered version appears to be used below (and I believe that this would be correct).

* In line 499, can you justify that $\ell$ is a function of only the rank one matrix, and is the function sufficiently smooth to take derivatives? Is it defined on the space of matrices, or at least in an open neighborhood of the set of rank 1 matrices so that the derivatives are well-defined? This is easy to settle in the linear settings in the main article, but a more thorough consideration is required in general.

* The notation $\partial\tilde \ell /\partial(u_iw_j)$ is very unfortunate - for a single derivative, there should not be two arguments. This would be a derivative of $\tilde\ell$ in direction $z_{ij}$ on a space of matrices... Can you explain why the modification is needed? And could it be avoided in the statement of Theorem 3.1 to focus on the more intuitive quantities involving the original loss $\ell$?

* Where in the proof of Theorem 3.1 is there any indication that $tr(C(u)) = tr(C(w))$? I have looked fairly closely and cannot find a justification for (3). I also assume that the covariance matrix of $u$ depends on $v, w$ and the covariance matrix of $w$ equally depends on all three sets of variables? I am also unsure how precisely this indicates that 'gradient noise between the two layers is balanced'.

* What is the distribution of $x$ in the experiment of Figure 1? Please include code for reproducibility or describe the experiments which illustrate your point.

* How does (2) imply that 'a single and unique point is favored by SGD'?

* Would Theorem 3.1 hold for *all* neurons with leaky ReLU activation?

* What do you mean by 'the difference between GD and SGD at stationarity is O(1)'? GD does not have an invariant distribution and the terminal state depends on (random) initialization. Is the message that a gradient flow analysis is only valid for a fixed finite time horizon?

* Should the minimizer of (10) be $\beta_2 /(\beta_1+\gamma)$ or should $\beta_1 = \mathbb E[x^2] + \gamma$?

* Below (10), what are the $\alpha_i$? Is this the same as in (8) for a different model?

* I find that 
$$
\Delta = \min_v C(v) = 4\min_v \left( \alpha_1v^2 - 2\alpha_2v + \alpha_3\right) = 4\left(\alpha_3 - \frac{\alpha_2^2}{\alpha_2}\right)
$$
is achieved for $v= \alpha_2/\alpha_1$. This differs from the value proposed in the article by a factor of $4/\alpha_1$. 

* I cannot find the proof of (12) in the appendix anywhere, and I am unsure what $\beta_1'$ is. More concrete references may be useful.

* In Figure 3, what is plotted for $v$ in the depth 1 case? Is it the product of parameters?
limitations: limitations: Yes.
paper_summary: paper_summary: The paper studies the effect of rescaling symmetry in SGD and shows SGD tends to favor solutions with balanced gradient noises. The authors then derive an exact solution of the stationary distribution of a toy model trained by SGD.  The derived solution shed lights on problems observed in deep learning such as fluctuation inversion and edge of stability.
strength_weakness: strength_weakness: The results of the paper are interesting and important, the writing needs refinement to improve clarity and precision. The conditions under which the results hold is sometimes omitted, leading to confusion. The language should also be made more precise.

Minor points:
1.	The first paragraph in related works appears to overstate the novelty of the results. Specifically, our result is the first to derive an exact solution to the stationary distribution of SGD without any approximation. (Line 55-56) This is a strong claim, but it seems inaccurate. There are previous results showing exact solution of stationary distribution of SGD (e. g. Liu Ziyin 2021). Corollary I.1 in arXiv:2306.04251 (2023) also states the stationary distribution on a deep learning setup similar to the D=1 model discussed in this paper. Also, the solution given in the paper is for a specific model. These should be made clear.
2.	It seems that eq. 15 takes D=1, which has not been stated and thus is confusing.
3.	It is unclear why the left figure of Fig. 5 has only two theory lines instead of three.

Major points:
1.	The related works on symmetry and SGD dynamics are insufficient. There are a few related works that are missing, e. g. arXiv:2309.16932 (2023).
2.	The paper has not discussed convergence to the stationary distribution. The authors seem to assume convergence to stationary distribution and use interchangeably the SGD properties and the stationary solution properties (e. g. line 97-98). However, the properties of SGD can be very different from the properties of stationary solutions unless convergence to the stationary solutions is guaranteed. The authors should clarify this.
3.	The authors fail to discuss uniqueness of the stationary solutions. For example, it is unclear to me why eq (3) is a necessary and efficient condition for stationarity. Eq (3) is a critical result in the paper, and it would be better to make it a theorem or corollary. However, since eq (2) cannot be interpreted as a deterministic ODE. The unique condition for a stationary distribution should be justified, especially considering that C1 and C2 are not constant but depend on u and w.
4.	The equivalence of SGD bias and weight decay is not rigorous. (line 155-158) The C0 term is not constant but depends on u and w, while the weight decay rate is constant.
questions: questions: 1.	The entire paper is based on rescaling symmetry, but why tanh without rescaling symmetry is used in Fig. 3 and 5?
limitations: limitations: The authors have listed limitations at the end. The major limitations are the simplicity of the model and lack of experiments on deep neural networks.

Comments: title: reply
Thank you for the reply
reply
reply
reply
After Rebuttal
Reply
reply
comment: Thanks for the reply. We will refine our statements. What you point to is a corner case. Technically speaking, $\lambda^*$ is indeed allowed to be infinity. The meaning is clear. When $\lambda^* =\infty$, the system diverges with probability 1 (here, an $O(1)$ neighborhood of infinity simply means infinity, and thus, divergence). This is not uncommon when the model has some sort of continuous symmetry. For example, in case of scale invariance, it is well-understood that SGD diverges along the degenerate directions (namely, the radial direction) when there is no weight decay (e.g., see https://arxiv.org/pdf/2102.12470)
I thank the authors for the clarifications! I am unable to arrive at Corollary D.2 without seeing the proof. Overall, I think the paper has novel contribution and the results are interesting. However, I feel like it still needs some refinement, especially for the language precision. For example, in the rebuttal pdf, $\lambda^*$ can take $\infty$, and it is unclear to me for Corollary D.2 what it means for an $O(1)$ neighborhood of $\lambda^*$ when $\lambda^*$ takes $\infty$.
Hi! We noticed that you have not responded to our previous clarification. It would be really helpful to us when revising and improving our manuscript if we could hear more from you. We would be happy to hear any additional thoughts or feedback you have on our work and on our previous reply!
Thanks for the reply. We will do our best to incorporate the feedback from the reviewers to improve the manuscript.
I will increase my score. I still feel that this paper is not fully ready for publication and requires significant revision, but my concerns about mathematical correctness have been addressed. I thank the authors for engaging with the questions and providing thoughtful replies.
Thank you very much for this very good question.

It is true that our definition only defines a rank-1 subspace of the gradient vector of $\tilde{\ell}$, but this is all that is required for us. Since the same argument applies to $C_1$ and $C_2$, we focus on $C_1$.

Essentially, we only need the quantity $u^T C_1 u$ to exist, and this only requires a rank-$1$ subspace of $C_1$ to be defined, which is defined as: $u^T C_1 u = \sum_{i,j,k}u_i u_j\mathbb{E}[\partial\tilde{\ell}/\partial (u_i w_k) \partial\tilde{\ell}/\partial(u_j w_k)] =  \sum_{i,j,k}\mathbb{E}[ u_i \partial\tilde{\ell}/\partial(u_i w_k) \partial\tilde{\ell}/\partial(u_j w_k) u_j] $ (ignoring the second term of $C_1$ as it follows from the same argument). Meanwhile, the quantity $ u_i \frac{\partial\tilde{\ell}}{\partial(u_i w_k)} = u_i \frac{\partial{\ell}}{\partial(u_i)} \frac{1}{2 w_k} +  \frac{1}{2}\frac{\partial{\ell}}{\partial(w_k)} $ (see our very first rebuttal), which is well-defined whenever the gradient with respect to the original loss is well-defined. Since each term of the sum is well-defined, the quantity $u^T C_1 u$ is also well-defined. Similarly, $w^T C_2 w$ is always well-defined. 

Lastly, the fact that we only require a rank-$1$ condition for the gradient of $\tilde{\ell}$ does not imply that $\nabla \tilde{\ell}$ when viewed as a matrix is rank-$1$. As an example, consider the case $\ell(u_i, w_i)= \sum_i u_i w_i$, and $\tilde{\ell}(Z) = \sum_i Z_{ii}$, such that $Z_{ii}(u, w) = u_i w_k$. Thus, we have that $\nabla_Z \tilde{\ell}(Z) = {\rm identity}$, which is full-rank and so its largest and smallest eigenvalues can still be well-defined.
Thanks for your reply and update. We would really appreciate it if you can be more specific regarding "convincing examples." To be specific, which claim in the paper do you think requires more example to strengthen? Knowing this would really help us improve our work.
**I looked at the proof of Theorem 3.1 and I am unable to fill in the gaps where diffusion terms should be.**

I thank the authors for the explanation. It makes a lot more sense now and I appreciate the detailed presentation. I believe that integrating it into the article will help future readers as well. 

**It is still not entirely clear to me how to show that the loss only depends on the rank 1 matrix and how the 'lifted' function would be defined.**

I still do not follow. Let me explain my dilemma: The authors define the function $\tilde \ell$ on the space of rank 1 matrices by $\tilde \ell (u\otimes v) = \ell(u, v)$. In my understanding, they argue that it does not depend on the scaling between the vectors $u, v$ which are used to decompose $u\otimes v$ (unique up to scaling). Where I do not follow is when they take the derivative with respect to the ij-entry of the matrix, corresponding to a difference quotient
$$
\frac{\tilde \ell (u \otimes v + h e_{ij}) - \tilde \ell(u\otimes v)} h.
$$
However $u\otimes v + h e_{ij}$ is not necessarily a rank 1 matrix, meaning that $\tilde\ell$ is not defined there, as far as I can tell. This issue also arises in their previous response. I think this is the result of formal manipulations which may not be applicable here.

If the authors can provide an explanation or alternative argument, I will be happy to increase my score.
I thank the authors for the good rebuttal that has answered partially to my concern. For this reason I decide to increase my score by one, still thinking that the article lacks some convincing example to strenghten their claim.
Thank you for your additional questions.

**For the author’s response to my question, I would suggest replacing the experimental results with the ones run with ReLU activation. ReLU is widely used, and it still makes little sense to me that the author decided to run the experiments with tanh while the entire paper is based on rescaling symmetry.**

Thank you for your suggestion. We will move the results of tanh nets to the appendix, and use the ones for ReLU nets in the main text to avoid confusion.

**For the supplement pdf, I have two follow up questions.**

**1. From my understanding, Equation 124 should still be considered a stochastic differential equation by nature. Then, how can the author arrive at Corollary D.2, which is a deterministic statement without out any probability condition?**

This is a misunderstanding. There is no stochasticity in the evolution equation in Eq. 124. The diffusion terms in the dynamics of $\|u\|^2$ and $\|w\|^2$ cancel with each other due to the rescaling symmetry – even if $u$ and $w$ themselves are random. This (rather surprising fact) is essentially what the theorem is trying to prove. Also, please see our last reply to reviewer (56YC) for an alternative and more technical proof of theorem 3.1, which may be easier to understand for some readers.

Now, because Eq. (124) has no noise term, one can further construct a deterministic dynamics that strictly upper bound Eq. 124, which is how one can arrive at the corollary.

**2. In the right panel of the new Figure 6, why does the difference in the norm grows up again later in time?**

First of all, Figure 6-right can be a little visually misleading. It looks like the neurons first have a decreasing $|\|u\|^2-\|w\|^2|$ and then an increasing $|\|u\|^2-\|w\|^2|$, but the actual situation is simply that roughly half of the signs of $\|u\|^2-\|w\|^2$ are changing during learning, and so in a log scale, it looks like they are first decreasing and then increasing, even though $\|u\|^2-\|w\|^2$ are quite monotonic in reality.

Secondly, $\|u\|^2-\|w\|^2=0$ is not necessarily the stationary point and Corollary D.2 only guarantees a convergence to a neighborhood $B$ close to a stationary point, and so there is no reason to expect a convergence to any specific value, especially when the model is nonlinear. What happens the most often is that the parameters fluctuate around some small neighborhood, which is exactly what Figure 6-right shows.

Please ask for additional clarification for any point that is not yet clear.
**I appreciate the authors' response and explanation.**

**I looked at the proof of Theorem 3.1 and I am unable to fill in the gaps where diffusion terms should be. It is possible that the authors chose an unfortunate simplifying notation and that this is correct, but I cannot verify the result right now. I would also be surprised at true monotonicity for a stochastic quantity. In Corollary 3.3, the authors seem to use the result in an ODE fashion to conclude that |u-v|^2 in fact decays to zero.**

Thanks for this question. We clarify the details below. We chose a simplified notation in the original proof to avoid confusing the readers with too much technical details. More technically, the proof should start with Eq. (1), and the standard SDE formulation. The following is how one would arrive at the result in this route.

Also, a core aspect of the result of Theorem 3.1 is exactly that Eq. (2) is an ODE! This is because the diffusion term cancel each other for the time evolution of $\|u\|^2-\|w\|^2$. The mechanism for this cancellation is essential to the proof: the gradient noice covariance is low-rank due to symmetry.

Now, we outline the proof of Theorem 3.1 in a more technical and standard manner. By defining $\theta=(u^T,w^T)^T$, we can rewrite the quantity $\|u\|^2-\|w\|^2$ as $\theta^TB\theta$, where $B:=\quad
\begin{pmatrix} I_u & O \\\\ O & -I_w \end{pmatrix}
\quad$. The dynamics of $\theta$ is described by Eq. (1): $d\theta/dt=-\nabla_{\theta}L+\sqrt{TC(\theta)}dW/dt$, where $C(\theta):=\mathbb{E}[\nabla_{\theta}\ell\nabla_{\theta}^T\ell]-\mathbb{E}[\nabla_{\theta}\ell]\mathbb{E}[\nabla_{\theta}\ell]$. To obtain the dynamics of $\theta^TB\theta$, we use Ito's lemma: If the vector $X_t$ satisfies the stochastic process $dX_t=\mu_tdt+D_tdW_t$, then the dynamics of a function $f(X_t)$ of $X_t$ can be written as $d f(X_t) = \left(\nabla_X^T f \mu_t + \frac{1}{2}\mathrm{Tr}[D_t^T \nabla_X^2 f(X_t) D_t] \right)dt + \nabla_X^Tf(X_t) D_t dW_t$. Hence, the dynamics of $G(\theta)=\theta^TB\theta$ is given by $dG(\theta)/dt=-\theta^TB\nabla_{\theta}L+\theta^TB\sqrt{TC(\theta)}dW/dt+T\mathrm{Tr}[C(\theta)B]$. To simplify this, we use the infinitesimal form of the rescaling symmetry, which can be expressed as $\ell(\theta,x)=\ell(A\theta,x)$ with the matrix $A:=\quad
\begin{pmatrix} (1+\epsilon) I_u & O \\\\ O & (1-\epsilon)I_w \end{pmatrix}
\quad$. We expand the equation to first order in $\epsilon$ and obtain $\theta^TB\partial\ell/\partial\theta=0$. By taking expectation to both sides, we have $\theta^TB\partial L/\partial\theta=0$. In addition, $\theta^TBC(\theta)=\mathbb{E}[\theta^TB\nabla_{\theta}\ell\nabla_{\theta}^T\ell]-\mathbb{E}[\theta^TB\nabla_{\theta}\ell]\mathbb{E}[\nabla_{\theta}\ell]=0$. Therefore, we can see $\theta^TB\sqrt{C(\theta)}=0$ since $C(\theta)$ and $\sqrt{C(\theta)}$ share the same null space. By substituting $\theta^TB\partial L/\partial\theta=0$ and $\theta^TB\sqrt{C(\theta)}=0$ into the evolution equation of $G(\theta)$, we have $dG(\theta)/dt=T\mathrm{Tr}[C(\theta)B]$, which is nothing but Eq. (30). Here we can see the diffusion terms related to $dW$ vanish due to the rescaling symmetry.






**It is still not entirely clear to me how to show that the loss only depends on the rank 1 matrix and how the 'lifted' function would be defined on the space of matrices (not just the set of rank 1 matrices, which does not form a linear subspace). An explicit construction or formula would be appreciated.**


Due to the rescaling symmetry, the loss function only depends on this rank 1 matrix.  To show this, we redefine a new function $\tilde{\ell}(u_iw_j,w_j,x):=\ell(u_i,w_j,x)$ to reorganize the parameters. Then the derivatives of $\ell$ can be rewritten as $\frac{\partial \ell}{\partial w_j}=\sum_{i}u_i\frac{\partial \tilde{\ell}}{\partial(u_iw_j)}$ and $\frac{\partial\ell}{\partial u_i}=\sum_{j}w_j\frac{\partial \tilde{\ell}}{\partial(u_iw_j)}+\frac{\partial \tilde{\ell}}{\partial(w_j)}$. However, due to the rescaling symmetry $\tilde{\ell}(u_iw_j,w_j,x)=\tilde{\ell}(u_iw_j,\lambda w_j,x)$, we have $\sum_jw_j\partial\tilde{\ell}/\partial(w_j)=0$. Since this equality holds for an arbitrary $w_j$, we always have $\partial\tilde{\ell}/\partial(w_j)=0$ for all $j$. Hence, the loss function only relies on the rank 1 matrix $u_iw_j$. We will clarify this in the revision.


Please ask for additional clarification for any point that is not yet clear.
I appreciate author’s response to my review! 

For the author’s response to my question, I would suggest replacing the experimental results with the ones run with ReLU activation. ReLU is widely used, and it still makes little sense to me that the author decided to run the experiments with tanh while the entire paper is based on rescaling symmetry.

For the supplement pdf, I have two follow up questions.

1.	From my understanding, Equation 124 should still be considered a stochastic differential equation by nature. Then, how can the author arrive at Corollary D.2, which is a deterministic statement without out any probability condition?

2.	In the right panel of the new Figure 6, why does the difference in the norm grows up again later in time?
I appreciate the authors' response and explanation.

* I looked at the proof of Theorem 3.1 and I am unable to fill in the gaps where diffusion terms should be. It is possible that the authors chose an unfortunate simplifying notation and that this is correct, but I cannot verify the result right now. I would also be surprised at true monotonicity for a stochastic quantity. In Corollary 3.3, the authors seem to use the result in an ODE fashion to conclude that |u-v|^2 in fact decays to zero.

* It is still not entirely clear to me how to show that the loss only depends on the rank 1 matrix and how the 'lifted' function would be defined on the space of matrices (not just the set of rank 1 matrices, which does not form a linear subspace). An explicit construction or formula would be appreciated.

I currently maintain my score and certainty. If the authors can provide a convincing answer to the two questions, I will most likely increase the score but drop the certainty, since I do not feel like I will be able to adequately assess its merit during the discussion period.

For the record, I find the work interesting and I believe that it deserves to be published, but I have reservations about the correctness (or the presentation, if this is really a question of notation).