Title: Noise Balance and Stationary Distribution of Stochastic Gradient Descent
Venue: NeurIPS 2024
Accepted: no
Keywords: stochastic gradient descent

stationary distribution

stochastic differential equation

phase transition
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.

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.

1 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

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;

the first-of-its-kind solution of the stationary distribution of an analytical model trained by SGD;

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 35 law of balance, the first main result of this work, and discuss its implications for common neural 36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

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.

2 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 ( θ ) -c 0 for some constant c 0 . Aprimary 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.

3 Noise Balance

We consider the continuous-time limit of SGD [15, 16, 18, 32, 8, 11]:

where 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].

3.1 Rescaling Symmetry and Law of Balance

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,

where C 1 = E [ A T A ] -E [ A T ] E [ A ] , C 2 = E [ AA T ] -E [ A ] E [ A T ] and A ki = ∂ ˜ ℓ / ∂ ( u i w k ) with 80 ˜ ℓ ( u i w k , v, x ) ≡ ℓ ( u i , w k , v, x ) . 1 81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

Here, v stands for the parameters that are irrelevant to the symmetry, and C 1 and C 2 are positive semi-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.

While the matrices C 1 and C 2 may not always be full-rank, we emphasize that in common deep-learning settings with rescaling symmetry, the law of balance is almost always well-defined and applicable. In Appendix A.4, we prove that under very general settings, for all active hidden neurons of a two-layer ReLU net, C 1 and C 2 are always full-rank. Equation (2) is the law of balance, and it implies two different types of balance. The first type of balance is the balance of gradient noise. The proof of the theorem shows that the stationary point of the law in (2) is equivalent to

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.

2.5

minimum

2.0

balance

conservation

1.5

u

SGD

1.0

0.5

0.0

0.0

0.5

1.0

1.5

2.0

2.5

W

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 C 2 w ∗ ( u ∗ ) T C 1 u ∗ . 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:

where λ 1 m ( 2 m ) and λ 1 M ( 2 M ) 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)

Thus, 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 as

where u i ∈ R d u , w i ∈ R d w and b i is a scalar with i being the index of the hidden neuron. For each 121 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 122 law of balance to each neuron separately. The per-sample loss function is 123

Here, x has a full-rank covariance Σ x , and y = g ( x ) + ϵ for some function g and ϵ is a zero-mean 124 random vector independent of x and have the full-rank covariance Σ ϵ . The following theorem shows 125 that for this network, C 1 and C 2 are full rank unless the neuron is 'dead'. 126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

1.4

layer 1

15.0

1.2

layer 2

12.5

actor

1.0

10.0

0.8

2

ance

7.5

0.6

E

5.0

va

0.4

2.5

layer 1

0.2

layer 2

0.0

0.0

0

2000

4000

6000

8000

0

2000

4000

6000

8000

t

t

Theorem 3.2. Let the loss function be given in Eq. (6) . Let C ( i ) 1 and C ( i ) 2 denote the corresponding noise matrices of the i -th neuron, and p i ∶= P ( w T i x + b i > 0 ) . Then, C ( i ) 1 and C ( i ) 2 are full-rank for all i such that p i > 0 .

See Figure 2. We train a two-layer ReLU network with the number of neurons: 20 → 200 → 20 . 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.

3.2 1d Rescaling Symmetry

The theorem greatly simplifies when both u and w are one-dimensional.

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 TC 0 to TC 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 that

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

where we have introduced the parameters

When α 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 (-4 T ( α 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 , /uni22EF , d } at stationarity, in agreement with the Corollary.

4 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:

The 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.

4.1 Depth0 Case

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 . Defining the 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 ) . Note that the loss landscape L only depends on β 1 and β 2 , and the gradient noise only depends on α 1 , α 2 and, α 3 . It is thus reasonable to call β the landscape parameters and α the noise parameters. Both β and α appear in all stationary distributions, implying that the stationary distributions of SGD are strongly data-dependent. Another relevant quantity is ∆ ∶= min v C ( v ) ≥ 0 , which is the minimal level of noise on the landscape. It turns out that the stationary distribution is qualitatively different for ∆ = 0 and for ∆ > 0 . For all the examples in this work,

When 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)

in 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.

4.2 An Analytical Model

Now, we consider the following model with a notion of depth and width; its loss function is

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

2.5

2.5

D=0

D=0

2.0

D =1

2.0

8

D =1

1.5

1.5

6

P(v)

P(v)

(n)d

1.0

1.0

4

0.5

0.5

2

0.0

0.0

0

0

1

2

3

0

1

2

3

-0.2

0.0

0.2

0.4

V

V

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].

We introduce v 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):

Theorem 4.2. Let δ ( x ) denote the Dirac delta function. For an arbitrary factor z in [ 0 , 1 ] , an invariant solution of the Fokker-Planck Equation is p ∗ ( v ) = ( 1 -z ) δ ( v ) + zp ± ( v ) , where

where 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 236 at zero, they will remain there with probability 1 . This means that the stationary distribution is not 237 unique. Since the result is symmetric in the sign of β 2 = E [ xy ] , we assume that E [ xy ] > 0 from 238 now on. 239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

Also, we focus on the case γ = 0 in the main text. 4 The distribution of v is

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 α 3 T -3 / 2 integrates to give v -β 2 / 2 α 3 T -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 α 3 T -3 / 2 , which integrates to v β 2 / 2 α 3 T -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 in Ref. [45]. See Figure 3 for illustrations of the distribution across different values of T . We also compare with the stationary distribution of a depth0 model. Two characteristics of the two-layer model appear rather striking: (1) the solution becomes a delta distribution at the sparse solution u = w = 0 at a large learning rate; (2) the two-layer model never learns the incorrect sign ( v is always non-negative). Another exotic phenomenon implied by the result is what we call the 'fluctuation inversion.' Naively, the variance of model parameters should increase as we increase T , which is the noise level in SGD. However, for the distribution we derived, the variance of v and u both decrease to zero as we increase T : injecting noise makes the model fluctuation vanish. We discuss more about this 'fluctuation inversion' in the next section.

Also, while there is no other phase-transition behavior below T c , there is still an interesting and practically relevant crossover behavior in the distribution of the parameters as we change the learning rate. When training a model, The most likely parameter we obtain is given by the maximum likelihood estimator of the distribution, ˆ v ∶= arg max p ( v ) . Understanding how ˆ v ( T ) changes as a function of T is crucial. This quantity also exhibits nontrivial crossover behaviors at critical values of T .

When T < T c , a nonzero maximizer for p ( v ) must satisfy

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 σ 2 > 8 21 k 2 , there always

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

exists a transitional learning rate: T ∗ = 4 k + √ 42 σ 4 ( 21 σ 2 -8 k 2 ) . Obviously, T c / 3 < T ∗ . One can characterize the learning of SGD by comparing T with T c and T ∗ . For this simple example, SGD can be classified into roughly 5 different regimes. See Figure 4.

4.3 Power-Law Tail of Deeper Models

An interesting aspect of the depth1 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). The d -dependent term vanishes only if D = 1 . Another intriguing aspect of the depth1 distribution is that its tail is independent of any hyperparameter of the problem, dramatically different from the linear regression case. This is true for deeper models as well.

Since d only affects the non-polynomial part of the distribution, the stationary distribution scales as p ( v ) ∝ 1 v 3 ( 1 -1 /( D + 1 )) ( α 1 v 2 -2 α 2 v + α 3 ) . Hence, when v → ∞ , the scaling behaviour is v -5 + 3 /( D + 1 ) . The tail gets monotonically thinner as one increases the depth. For D = 1 , the exponent is 7 / 2 ; an infinite-depth network has an exponent of 5 . Therefore, the tail of the model distribution only depends on the depth and is independent of the data or details of training, unlike the depth0 model. In addition, due to the scaling v 5 -3 /( D + 1 ) for v → ∞ , we can see that E [ v 2 ] will never diverge no matter how large the T is.

An intriguing feature of this model is that the model with at least one hidden layer will never have a divergent training loss. This directly explains the puzzling observation of the edge-of-stability phenomenon in deep learning: SGD training often gives a neural network a solution where a slight increment of the learning rate will cause discrete-time instability and divergence [40, 4]. These solutions, quite surprisingly, exhibit low training and testing loss values even when the learning rate is right at the critical learning rate of

100

Ib

1

Tc

Te/3

IIb

T*

10-1

Ha

Ia

T

10-2

IⅢI

b

Ib

10-3

100

101

a

instability. This observation contradicts naive theoretical expectations. Let η sta denote the largest stable learning rate. Close to a local minimum, one can expand the loss function up to the second order to show that the value of the loss function L is proportional to Tr [ Σ ] . However, Σ ∝ 1 /( η sta -η ) 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 by Ref. [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.

InfiniteD limit. As D tends to infinity, the distribution becomes

where k 1 = d ( α 3 β 1 -2 α 2 β 2 )/( TDα 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.

Now, if we T , there are three situations: (1) d = o ( D ) , (2) d = c 0 D for a constant c 0 , (3) d = Ω ( D ) . 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

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

100

0.5

D=0

98.50

0.4

D=1

98.25

Nsta

10-1

98.00

(n)d

0.3

2

97.75

10-2

D =1

0.2

U

U

97.50

D = 4

97.25

D =10

0.1

97.00

nα1/D

10-3

standard

0.0

96.75

100

101

0.00

0.25

0.50

0.75

1.00

5

10

15

V

n

D

zero. Therefore, we should increase the model width as we increase the depth. In the second case, d / D is a constant and can thus be absorbed into the definition of T and is the only limit where we obtain a nontrivial distribution with a finite spread. If d = Ω ( D ) , the distribution becomes a delta distribution at the global minimum of the loss landscape, p ( v ) = δ ( v -β 2 / β 1 ) and achieves the global minimum.

4.4 Implication for Variational Bayesian Learning

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 of

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 .

5 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.

370

References

John C Baez and Brendan Fong. A noether theorem for markov processes. Journal of Mathe371 matical Physics , 54(1):013301, 2013. 372

Rapha¨ el Berthier. Incremental learning in diagonal linear networks. Journal of Machine Learn373 ing Research , 24(171):1-26, 2023. 374

Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational infer375 ence, converges to limit cycles for deep networks. In 2018 Information Theory and Applica376 tions Workshop (ITA) , pages 1-10. IEEE, 2018. 377

Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gra378 dient descent on neural networks typically occurs at the edge of stability. arXiv preprint 379 arXiv:2103.00065 , 2021. 380

Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gra381 dient descent at the edge of stability. arXiv preprint arXiv:2209.15594 , 2022. 382

Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homoge383 neous models: Layers are automatically balanced. Advances in neural information processing 384 systems , 31, 2018. 385

Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal 386 linear networks: Implicit regularisation, large stepsizes and edge of stability. arXiv preprint 387 arXiv:2302.08982 , 2023. 388

Xavier Fontaine, Valentin De Bortoli, and Alain Durmus. Convergence rates and approxi389 mation results for sgd and its continuous-time counterpart. In Mikhail Belkin and Samory 390 Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory , volume 134 391 of Proceedings of Machine Learning Research , pages 1965-2058. PMLR, 15-19 Aug 2021. 392

Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? 393 Advances in neural information processing systems , 31, 2018. 394

Boris Hanin and David Rolnick. How to start training: The effect of initialization and archi395 tecture. Advances in Neural Information Processing Systems , 31, 2018. 396

Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of 397 nonconvex stochastic gradient descent. arXiv preprint arXiv:1705.07562 , 2017. 398

Laurent Valentin Jospin, Hamid Laga, Farid Boussaid, Wray Buntine, and Mohammed Ben399 namoun. Hands-on bayesian neural networks-a tutorial for deep learning users. IEEE Com400 putational Intelligence Magazine , 17(2):29-48, 2022. 401

Kenji Kawaguchi. Deep learning without poor local minima. Advances in Neural Information 402 Processing Systems , 29:586-594, 2016. 403

Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori 404 Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dy405 namics. arXiv preprint arXiv:2012.04728 , 2020. 406

Jonas Latz. Analysis of stochastic gradient descent in continuous time. Statistics and Comput407 ing , 31(4):39, 2021. 408

Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of 409 stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning 410 Research , 20(40):1-47, 2019. 411

Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with tra412 ditional optimization analyses: The intrinsic learning rate. Advances in Neural Information 413 Processing Systems , 33:14544-14555, 2020. 414

Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochas415 tic differential equations (sdes), 2021. 416

Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and fluctuation of finite learning rate 417 stochastic gradient descent, 2021. 418

Agnieszka B Malinowska and Moulay Rchid Sidi Ammi. Noether's theorem for control prob419 lems on time scales. arXiv preprint arXiv:1406.0705 , 2014. 420

Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as ap421 proximate bayesian inference. Journal of Machine Learning Research , 18:1-35, 2017. 422

John C Mauro, Prabhat K Gupta, and Roger J Loucks. Continuously broken ergodicity. The 423 Journal of chemical physics , 126(18), 2007. 424

Tetsuya Misawa. Noether's theorem in symmetric stochastic calculus of variations. Journal of 425 mathematical physics , 29(10):2178-2180, 1988. 426

Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Power-law escape rate of sgd. In 427 International Conference on Machine Learning , pages 15959-15975. PMLR, 2022. 428

Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias 429 of the step size in linear diagonal neural networks. In International Conference on Machine 430 Learning , pages 16270-16295. PMLR, 2022. 431

Richard G Palmer. Broken ergodicity. Advances in Physics , 31(6):669-735, 1982. 432

Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for di433 agonal linear networks: a provable benefit of stochasticity. Advances in Neural Information 434 Processing Systems , 34:29218-29230, 2021. 435

Hannes Risken and Hannes Risken. Fokker-planck equation . Springer, 1996. 436

Tomasz Rolski, Hanspeter Schmidli, Volker Schmidt, and Jozef L Teugels. Stochastic pro437 cesses for insurance and finance . John Wiley & Sons, 2009. 438

Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear 439 dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 , 2013. 440

N. Shirish Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On Large441 Batch Training for Deep Learning: Generalization Gap and Sharp Minima. ArXiv e-prints , 442 September 2016. 443

Justin Sirignano and Konstantinos Spiliopoulos. Stochastic gradient descent in continuous 444 time: A central limit theorem. Stochastic Systems , 10(2):124-151, 2020. 445

Hidenori Tanaka and Daniel Kunin. Noether's learning dynamics: Role of symmetry breaking 446 in neural networks, 2021. 447

DThirumalai and Raymond D Mountain. Activated dynamics, loss of ergodicity, and transport 448 in supercooled liquids. Physical Review E , 47(1):479, 1993. 449

Christopher J Turner, Alexios A Michailidis, Dmitry A Abanin, Maksym Serbyn, and Zlatko 450 Papi´ c. Weak ergodicity breaking from quantum many-body scars. Nature Physics , 14(7):745451 749, 2018. 452

Nicolaas Godfried Van Kampen. Stochastic processes in physics and chemistry , volume 1. 453 Elsevier, 1992. 454

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, 455 Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information 456 processing systems , 30, 2017. 457

Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homo458 geneity: Convergence and balancing effect, 2022. 459

Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type part ii: 460 Continuous time analysis. Journal of Nonlinear Science , 34(1):1-45, 2024. 461

Lei Wu, Chao Ma, et al. How sgd selects the global minima in over-parameterized learning: 462 A dynamical stability perspective. Advances in Neural Information Processing Systems , 31, 463 2018. 464

Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dy465 namics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint 466 arXiv:2002.03495 , 2020. 467

Sho Yaida. Fluctuation-dissipation relations for stochastic gradient descent. arXiv preprint 468 arXiv:1810.00004 , 2018. 469

Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. In 470 Advances in Neural Information Processing Systems , 2022. 471

Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in 472 SGD. In International Conference on Learning Representations , 2022. 473

Liu Ziyin and Masahito Ueda. Exact phase transitions in deep learning. arXiv preprint 474 arXiv:2205.12510 , 2022. 475

A Theoretical Considerations 476

A.1 Background 477

A.1.1 Ito's Lemma 478

Let us consider the following stochastic differential equation (SDE) for a Wiener process W ( t ) : 479

We are interested in the dynamics of a generic function of X t . Let Y t = f ( t, X t ) ; Ito's lemma states 480 that the SDE for the new variable is 481

Let us take the variable Y t = X 2 t as an example. Then the SDE is 482

Let us consider another example. Let two variables X t and Y t follow 483

The SDE of X t Y t is given by 484

A.1.2 Fokker Planck Equation 485

The general SDE of a 1d variable X is given by: 486

The time evolution of the probability density P ( x, t ) is given by the Fokker-Planck equation: 487

where J ( X,t ) = µ ( X ) P ( X,t ) + 1 2 ∂ ∂X [ B 2 ( X ) P ( X,t )] . The stationary distribution satisfying 488 ∂P ( X,t )/ ∂t = 0 is 489

which gives a solution as a Boltzmann-type distribution if B is a constant. We will apply Eq. (25) 490 to determine the stationary distributions in the following sections. 491

A.2 Proof of Theorem 3.1 492

Proof. We omit writing v in the argument unless necessary. By definition of the symmetry 493 ℓ ( u , w , x ) = ℓ ( λ u , w / λ, x ) , we obtain its infinitesimal transformation ℓ ( u , w , x ) = ℓ (( 1 + ϵ ) u , ( 1 -494 ϵ ) w / λ, x ) . Expanding this to first order in ϵ , we obtain 495

The equations of motion are 496

Using Ito's lemma, we can find the equations governing the evolutions of u 2 i and w 2 j : 497

where C u i = Var [ ∂ℓ ∂u i ] and C w j = Var [ ∂ℓ ∂w j ] . With Eq. (26), we obtain 498

Due to the rescaling symmetry, the loss function can be considered as a function of the matrix uw T . 499 Here we define a new loss function as ˜ ℓ ( u i w j ) = ℓ ( u i , w j ) . Hence, we have 500

We can rewrite Eq. (30) into 501

where 503

The proof is thus complete. 504

A.3 Second-order Law of Balance 505

Considering the modified loss function: 506

In this case, the Langevin equations become 507

Hence, the modified SDEs of u 2 i and w 2 j can be rewritten as 508

where 502

In this section, we consider the effects brought by the last term in Eqs. (39) and (40). From the 509 infinitesimal transformation of the rescaling symmetry: 510

we take the derivative of both sides of the equation and obtain 511

where we take the expectation to ℓ at the same time. By substituting these equations into Eqs. (39) 512 and (40), we obtain 513

Then following the procedure in Appendix. A.2, we can rewrite Eq. (44) as 514

where 515

For one-dimensional parameters u, w , Eq. (45) is reduced to 516

Therefore, we can see this loss modification increases the speed of convergence. Now, we move 517 to the stationary distribution of the parameter v . At the stationarity, if u i = -w i , we also have the 518 distribution P ( v ) = δ ( v ) like before. However, when u i = w i , we have 519

Hence, the stationary distribution becomes 520

where 521

From the expression above we can see K 1 ≪ 1 + β 2 / 4 Tα 3 and K 2 ≪ ( α 3 β 1 -α 2 β 2 )/ 2 Tα 3 √ ∆ . 522 Hence, the effect of modification can only be seen in the term proportional to v . The phase transition 523 point is modified as 524

Compared with the previous result T c = β 2 α 3 , we can see the effect of the loss modification is α 3 → 525 α 3 + β 2 2 , or equivalently, Var [ xy ] → E [ x 2 y 2 ] . This effect can be seen from E 1 and E 2 . 526

527

A.4 Proof of Theorem 3.2

Proof. For any i , one can obtain the expressions of C ( i ) 1 and C ( i ) 2 from Theorem 3.1 as 528

where 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 ] ∶= 529 E [ O ∣ w T i x + b i > 0 ] . 530

531

where the last inequality follows from 532

Here we denote that Var i [ 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 ] . 533 For C ( i ) 2 , we let the vector ˜ n ∶= ( n T , n f ) T be a unit vector in R d w + 1 , yielding 534

Note that this quantity can be decomposed as 535

The covariance term vanishes because 536

Therefore, 537

where the penultimate inequality follows from the fact that ϵ is independent of x . Hence, both the 538 matrices C ( i ) 1 and C ( i ) 2 are full-rank. The proof is completed. 539

A.5 Derivation of Eq. (4) 540

541

We here prove inequality (4). At stationarity, d (∥ u ∥ 2 -∥ w ∥ 2 )/ dt = 0 , indicating

The first inequality in Eq. (63) gives the solution 542

The second inequality in Eq. (63) gives the solution 543

Combining these two results, we obtain 544

which is Eq. (4). 545

546

A.6 Proof of Theorem 4.1

Proof. This proof is based on the fact that if a certain condition is satisfied for all trajectories with 547 probability 1, this condition is satisfied by the stationary distribution of the dynamics with probabil548 ity 1. 549

Let us first consider the case of D > 1 . We first show that any trajectory satisfies at least one of 550 the following five conditions: for any i , (i) v i → 0 , (ii) L ( θ ) → 0 , or (iii) for any k ≠ l , ( u ( k ) i ) 2 -551 ( u ( l ) i ) 2 → 0 . 552

The SDE for u ( k ) i is 553

where v i ∶= ∏ D k = 1 u ( k ) i , and so v = ∑ i v i . There exists rescaling symmetry between u ( k ) i and u ( l ) i for 554 k ≠ l . By the law of balance, we have 555

556

557

558

559

560

561

562

563

564

565

566

where 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: for D = 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 is

Comparing the dynamics of v i and v j for i ≠ j , we obtain 567

By condition (iii), we have ∣ u ( 0 ) i ∣ = /uni22EF = ∣ 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 = 568 ( v 2 i ) ( D -1 )/( D + 1 ) . 5 Therefore, we obtain 569

We first consider the case where v i and v j initially share the same sign (both positive or both nega570 tive). When D > 1 , the left-hand side of Eq. (72) can be written as 571

which follows from Ito's lemma: 572

Substitute in Eq. (72), we obtain Eq. (73). 573

Now, we consider the right-hand side of Eq. (72), which is given by 574

Combining Eq. (73) and Eq. (75), we obtain 575

By defining z i = v 2 /( D + 1 )-1 i , we can further simplify the dynamics: 576

Hence, 577

Therefore, if v i and v j initially have the same sign, they will decay to the same value in the long578 time limit t →∞ , which gives condition 2-(b). When v i and v j initially have different signs, we can 579 write Eq. (72) as 580

Hence, when D > 1 , we simplify the equation with a similar procedure as 581

Defining z i = ∣ v i ∣ 2 /( D + 1 )-1 , we obtain 582

which implies 583

From this equation, we reach the conclusion that if v i and v j have different signs initially, one of 584 them converges to 0 in the long-time limit t → ∞ , corresponding to condition 1 in the theorem 585 statement. Hence, for D > 1 , at least one of the conditions is always satisfied at t →∞ . 586

Now, we prove the theorem for D = 1 , which is similar to the proof above. The law of balance gives 587

We can see that ∣ u ( 1 ) i ∣ → ∣ u ( 2 ) i ∣ takes place unless Var [ ∂ℓ ∂ ( u ( 1 ) i u ( 2 ) i ) ] = 0 , which is equivalent to 588 L ( θ ) = 0 . This corresponds to condition (ii). Hence, if condition (ii) is violated, we need to prove 589 condition (iii). In this sense, ∣ u ( 1 ) i ∣ →∣ u ( 2 ) i ∣ occurs and Eq. (72) can be rewritten as 590

591

When v i and v j are both positive, we have

With Ito's lemma, we have 592

593

599

600

601

602

603

604

Therefore, Eq. (85) can be simplified to

which indicates that all v i with the same sign will decay at the same rate. This differs from the case 594 of D > 2 where all v i decay to the same value. Similarly, we can prove the case where v i and v j are 595 both negative. 596

Now, we consider the case where v i is positive while v j is negative and rewrite Eq. (84) as 597

Furthermore, we can derive the dynamics of v j with Ito's lemma: 598

Therefore, Eq. (88) takes the form of

In the long-time limit, we can see log ( v i ∣ v j ∣) decays to -∞ , indicating that either v i or v j will decay to 0 . This corresponds to condition 1 in the theorem statement. Combining Eq. (87) and Eq. (90), we conclude that all v i have the same sign as t → ∞ , which indicates condition 2-(a) if conditions in item 1 are all violated. The proof is thus complete.

A.7 Proof of Theorem 4.2

Proof. Following Eq. (70), we substitute u ( k ) i with v 1 / D i for arbitrary k and obtain 605

With Eq. (78), we can see that for arbitrary i and j , v i will converge to v j in the long-time limit. In 606 this case, we have v = dv i for each i . Then, the SDE for v can be written as 607

If v > 0 , Eq. (92) becomes 608

Therefore, the stationary distribution of a general deep diagonal network is given by 609

If v < 0 , Eq. (92) becomes 610

The stationary distribution of ∣ v ∣ is given by 611

Thus, we have obtained 612

Especially when D = 1 , the distribution function can be simplified as 613

where we have used the integral 614

Furthermore, we can also see that p ( v ) = δ ( v ) is also the stationary distribution of the Fokker-Planck 615 equation of Eq. (93). Hence, the general stationary distribution of v can be expressed as 616

The proof is complete.

A.8 Analysis of the maximum probability point

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 that

617

618

619

620

621

622

623

624

625

The second case is α 2 2 < α 1 α 3 < 25 21 α 2 2 . In this case, we need to further compare the value between 626 5 α 2 β 1 and 7 α 1 β 2 . If 5 α 2 β 1 < 7 α 1 β 2 , we have A > 0 , which indicates that the maximum point 627 exists. If 5 α 2 β 1 > 7 α 1 β 2 , we need to further check the value of minimum of A , which takes the 628

is 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.

form of 629

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 630 exists. However, if β 1 β 2 ≥ 5 α 2 + √ 25 α 2 2 -21 α 1 α 3 3 α 3 , there is always a critical learning rate T ∗ . If 631

β 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 632 T c / 3 < T < T ∗ , there is a solution to the maximum condition, while there is no solution when 633 T > T ∗ . If β 1 β 2 > 5 α 2 + √ 25 α 2 2 -21 α 1 α 3 3 α 3 , there are two critical points: 634

For T < T 1 and T > T 2 , there exists a solution to the maximum condition. For T 1 < T < T 2 , there 635 is no solution to the maximum condition. The last case is α 2 2 = α 1 α 3 < 25 21 α 2 2 . In this sense, the 636 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 637 critical learning rate and the maximum always exists. Nevertheless, when β 1 β 2 > 7 α 1 5 α 2 , there is always 638 a critical learning rate as T ∗ = β 2 1 20 α 2 β 1 -28 α 1 β 2 . When T < T ∗ , there is a solution to the maximum 639 condition, while there is no solution when T > T ∗ . 640

641

642

643

644

645

A.9 Other Cases for D = 1

The 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:

Since the analysis of the Fokker-Planck equation is the same, we directly begin with the distribution 646 function in Eq. (15) for u i = -w i which is given by P (∣ v ∣) ∝ δ (∣ v ∣) . Namely, the only possible 647 weights are u i = w i = 0 , the same as the non-interpolation case. This is because the corresponding 648 stationary distribution is 649

For the case u i = w i , the stationary distribution is given from Eq. (15) as 651

The integral of Eq. (105) with respect to ∣ v ∣ diverges at the origin due to the factor ∣ v ∣ 3 2 + β 1 2 Tα 1 k . 650

Now, we consider the case of γ ≠ 0 . In the non-interpolation regime, when u i = -w i , the stationary 652 distribution is still p ( v ) = δ ( v ) . For the case of u i = w i , the stationary distribution is the same as 653 in Eq. (15) after replacing β with β ′ 2 = β 2 -γ . It still has a phase transition. The weight decay 654 has the effect of shifting β 2 by -γ . In the interpolation regime, the stationary distribution is still 655 p ( v ) = δ ( v ) when u i = -w i . However, when u i = w i , the phase transition still exists since the 656 stationary distribution is 657

where θ 2 = 1 2 Tα 1 k ( β 1 -γ k ) . The phase transition point is θ 2 = 1 / 2 , which is the same as the noninterpolation one.

658

659

660

661

662

663

664

665

666

667

668

669

670

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 :

The corresponding stationary distribution is

where ϕ 2 = 1 2 Tα 1 k ( β 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 is

It also contains a critical point: -3 2 + ϕ 2 = -1 , or equivalently, c = β 1 -α 1 kT . There are two cases. 671 When c < 0 , the probability density only has support for ∣ v ∣ > k since the gradient always pulls the 672 parameter ∣ v ∣ to the region ∣ v ∣ > k . Hence, the divergence at ∣ v ∣ = 0 is of no effect. When c > 0 , 673 the probability density has support on 0 < ∣ v ∣ < k for the same reason. Therefore, if β 1 > α 1 kT , 674 there exists a critical point c = β 1 -α 1 kT . When c > β 1 -α 1 kT , the distribution function P (∣ v ∣) 675 becomes δ (∣ v ∣) . When c < β 1 -α 1 kT , the integral of the distribution function is finite for 0 < ∣ v ∣ < k , 676 indicating the learning of the neural network. If β 1 ≤ α 1 kT , there will be no criticality and P (∣ v ∣) 677 is always equivalent to δ (∣ v ∣) . The effect of having weight decay can be similarly analyzed, and 678 the result can be systematically obtained if we replace β 1 with β 1 + γ / k for the case u i = -w i or 679 replacing β 1 with β 1 -γ / k for the case u i = w i . 680

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We believe that the abstract and introduction reflect the contributions and scope of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is 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 of the work performed by the authors?

Answer: [Yes]

Justification: We have discussed the limitations of our work in the Discussion session at lines 369-371.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a 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-specification, asymptotic approximations only holding locally). The authors should reflect on 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, e.g., if the approach was only tested on a few datasets or with a 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 approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it 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 address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to 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 a complete (and correct) proof?

Answer: [Yes]

Justification: We believe that the assumptions are clarified and complete proofs are provided 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented 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 conclusions 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 to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via 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 a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce 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 authors 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: The code or data of the experiments are simple and easy to reproduce following the description in the main text.

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.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have specified the training and test details in the captions of the experiments in Figs. 2,4, and 5.

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.

7. 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.

8. 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).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have confirmed that the research is conducted with the NeurIPS Code of Ethics.

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).

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 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).

11. 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.

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

975

976

977

978

979

980

981

982

983

984

985

986

987

988

989

990

991

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.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 not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 paper include the full text of instructions given to 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 subjects is included in our work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

992 993 994 995 996

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe paper contributes to the understanding of SGD properties. The noise balance theorem is novel and important. The analytical solution as well as the interpretation is interesting and insightful.\n\nWEAKNESSES:\nThe 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.\n\nMinor points:\n1.\tThe 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.\n2.\tIt seems that eq. 15 takes D=1, which has not been stated and thus is confusing.\n3.\tIt is unclear why the left figure of Fig. 5 has only two theory lines instead of three.\n\nMajor points:\n1.\tThe related works on symmetry and SGD dynamics are insufficient. There are a few related works that are missing, e. g. arXiv:2309.16932 (2023).\n2.\tThe 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.\n3.\tThe 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.\n4.\tThe 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.\n\nQUESTIONS:\n1.\tThe entire paper is based on rescaling symmetry, but why tanh without rescaling symmetry is used in Fig. 3 and 5?\n\nLIMITATIONS:\nThe authors have listed limitations at the end. The major limitations are the simplicity of the model and lack of experiments on deep neural networks.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFor 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.\n\nSTRENGTHS:\nThe topic is well-chosen and the results - if correct - are very interesting.\n\nWEAKNESSES:\n* 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).\n\n* 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.\n\n* 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.\n\n* 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.\n\nQUESTIONS:\n* 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).\n\n* 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.\n\n* 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$?\n\n* 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'.\n\n* 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.\n\n* How does (2) imply that 'a single and unique point is favored by SGD'?\n\n* Would Theorem 3.1 hold for *all* neurons with leaky ReLU activation?\n\n* 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?\n\n* Should the minimizer of (10) be $\\beta_2 /(\\beta_1+\\gamma)$ or should $\\beta_1 = \\mathbb E[x^2] + \\gamma$?\n\n* Below (10), what are the $\\alpha_i$? Is this the same as in (8) for a different model?\n\n* I find that \n$$\n\\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)\n$$\nis achieved for $v= \\alpha_2/\\alpha_1$. This differs from the value proposed in the article by a factor of $4/\\alpha_1$. \n\n* 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.\n\n* In Figure 3, what is plotted for $v$ in the depth 1 case? Is it the product of parameters?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe idea that a conservation law for the gradient flow implies an asymptotic balancedness condition for the stochastic flow is a good and striking idea.\n\nThe one-dimensional examples that are given in the text are very pleasant to follow and they are good exercices to display the ability of the stochastic flow to diverge from the gradient flow.\n\nThe example given in Eq.(13) is thoroughly analyzed.\n\nWEAKNESSES:\nThe paper present the following weaknesses:\n\n- 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.\n\n- I have to say that I was a bit bothered by the general overselling of the paper : \n     - As said before the law of balance is truly valid asymptotically in one-dimension\n     - 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! \n     - 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).\n\nMinor typos/flaws:\n\n- l.41: Fokker Planck is not inherently high-dimensional\n- l.44: Go to the line for new paragraph \n- l.165: The law of balance is not strictly applicable here since $\\ell$ is not scale invariant because of the regularization.\n- Section **4.1 Depth - 0**: I think that $\\Delta > 0$ is not currently the \"most practical example\" since it corresponds to a underparametrized model.\n\nQUESTIONS:\nOn top of the questions raised by the weakness section, here a some (more minor) questions:\n\n- l.65: Eq (1) is not the usual covariance matrix that is used for Stochastic modified equations.\n- l.95: Eq (3) seems important but it is difficult to follow where this comes from without intermediate calculations, can the authors develop?\n- 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? \n- Thm 4.2 : Are they the only stationary distribution ? How to know to which one it converges ?\n\nLIMITATIONS:\nAs said before, all conclusion are drawn for models that live intrinsically in one dimension.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}