Noise Balance and Stationary Distribution of Stochastic Gradient Descent

Thank you for your feedback. We will answer the weaknesses and questions below.

Weaknesses:

The results of the paper are interesting and important,...The language should also be made more precise.

Thank you for your suggestions. We will do our best to refine our language in the revision.

Minor points: 1. The first paragraph in related works appears to overstate the novelty of the results...

Thanks for these references. We will include a discussion of these works. In the previous work Liu Ziyin 2021 (actually, Liu. et al 2021), the authors investigated the stationary distribution only near the local minima. In Eq. (13) in Sec. 4.6, they approximate the covariance matrix as a constant matrix near the minima which represents the strength of the noise. Hence, their results are approximate stationary distribution rather than an exact one. Corollary I.1 of arXiv:2306.04251 only gives a particular solution to the Fokker-Planck equation, assuming strong conditions on the initialization (A3-A4). In contrast, our solution is a general one, enumerating all possible solutions for the problem. In science and mathematics, there is a fundamental difference between the two. For example, the general solution to the Navier-Stokes equation is a difficult open problem in mathematics, whereas particular solutions to it are quite easy to find (e.g., https://en.wikipedia.org/wiki/Landau%E2%80%93Squire_jet). We will restate our contribution as the “first general exact solution to a deep-and-wide linear model with 1d input and output,” which we believe is accurate given the existing literature.

Also, the solution given in the paper is for a specific model. These should be made clear.

We agree that this point should be made more explicit and also in the abstract.

2. It seems that eq. 15 takes D=1, which has not been stated and thus is confusing.

Yes. Eq. (15) is for $D=1$. We will clarify this.

3. It is unclear why the left figure of Fig. 5 has only two theory lines instead of three.

The dashed lines in Fig. 5 show the upper and lower bounds of the tail of the stationary distribution around l.295-l.300, where the left dashed line corresponds to the case $D=\infty$ while the right line corresponds to the case $D=1$.

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

Thank you for pointing out this related work. This work studies discrete symmetry. Our focus is on the rescaling symmetry, which is a continuous symmetry. In addition, this work focuses on the structure of Hessian under the mirror symmetry, while we focus on the stationary solutions and distribution of the parameters. We will clarify this.

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.

Yes. Convergence to the stationary distribution is assumed. The problem of convergence itself is a difficult mathematical problem and beyond the scope of our work (See https://doi.org/10.1007/s10884-018-9705-8 for an example). We will add a discussion on this point. For the stationarity of Eq (3), see the next point.

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.

See the summary rebuttal. The solution to Eq (3) is unique in the sense that for fixed $u$ and $w$, there exists a unique $\lambda$ such that Eq (3) reaches stationarity. We also establish that if this $\lambda^*$ does not change in time, Eq (3) will converge to this fixed point.

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.

We meant that they are qualitatively similar because SGD bias is like an “adaptive” weight decay. We will clarify this.

Questions:

The entire paper is based on rescaling symmetry, but why tanh without rescaling symmetry is used in Fig. 3 and 5?

Here, we consider the tanh network to clarify the application of our insights to cases where the symmetry only approximately holds. For a small initialization, the tanh network can be approximated by a linear network since $\tanh x = x + O(x^2)$. Thus, the rescaling symmetry approximately holds.

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.

Thanks for this criticism. We would like to politely emphasize that we do have a nontrivial theory for nonlinear models (Theorem 3.2), and experiments on nonlinear nets (ReLU networks in Fig. 2 and the tanh network in Figs. 3 and 5). We also include a new set of experiments on ReLU net in the uploaded pdf.

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.

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)

Thank you for your constructive feedback. We will answer the weaknesses and questions below.

Weaknesses:

Important quantities are defined throughout the plain text. I understand that...

This is a good question. Theorem 3.1 holds generally for an arbitrary network with the rescaling symmetry. The only condition we need is the symmetry of the network regardless of the distribution of $x$ and the properties of $L$ and $C$. So mathematically, this is a strong result and our results may also be applicable to other fields. This is not a weakness but a strength of our work. In fact, solving ODE with the Lie group is well-known in mathematics literature, whereas using the Lie group to solve SDE is, to our best knowledge, novel even for mathematicians.

I have serious doubts about Theorem 3.1. It is derived in...

This is a misunderstanding. While we did not write it explicitly, the time derivative of each variable is actually stochastic. In the SDE limit, the explicit forms of the time derivatives (Eqs. (27) and (28)) are $\frac{du_i}{dt}=-\frac{\partial L}{\partial u_i}+\sqrt{T\mathrm{Var}(\partial\ell/\partial u_i)}dW_t, \frac{dw_j}{dt}=-\frac{\partial L}{\partial w_j}+\sqrt{T\mathrm{Var}(\partial\ell/\partial w_j)}dW_t,$ which can be used to to derive Eq. (29).

The authors do not pay any attention to...

This is a good point. In the revision, we will strengthen the theorem by proving the existence and uniqueness of the solutions to Eq. (2). See the summary rebuttal and the pdf.

For regularity, we do need some regularity conditions for the matrix $A$ as a function of the data in order for the matrices $C_1$ and $C_2$ to exist and be well-behaved. One sufficient condition would be that the matrix $A$ is a Lipshitz function of $x$. We will clarify this.

Questions:

The noise matrix $C$ defined below (1) is...

This is a typo and we will fix it.

In line 499, can you justify that $\ell$ is a function of...

Here, we need the loss function $\ell$ to be a differentiable function of $u$ and $w$, and the rescaling symmetry implies some additional smoothness properties. To be specific, we can equivalently write the derivative $\partial\ell/\partial u_i$ and $\partial\ell/\partial w_j$ into a composition of the derivatives $\frac{\partial\tilde{\ell}}{\partial(u_iw_j)}=\frac{\partial\ell}{\partial u_i}\frac{1}{2w_j}+\frac{\partial\ell}{\partial w_j}\frac{1}{2u_i}$ if we define $\tilde{\ell}(u_iw_j,u_i/w_j):=\ell(u_i,w_j)$. Hence, once the loss function is sufficiently smooth for the parameters $u_i$ and $w_j$ and the parameters are away from $0$, we can always have the smooth derivative $\partial\tilde{\ell}/\partial(u_iw_j)$.

The notation $\partial\tilde{\ell}/\partial(u_iw_j)$ is very unfortunate...

Here the introduction of $\tilde{\ell}$ facilitates us to derive the balancedness of the norm $\lambda$ and the stationary distribution of the 1d case (See Corollary 3.3). The more intuitive expression of Eq. (2) is provided in Eq. (30), which utilizes the gradient of the original loss $\ell$ to the parameters $u_i$ and $w_j$. The stationarity directly means the balancedness of gradient noises on different layers.

Where in the proof of...this indicates that 'gradient noise between the two layers is balanced'.

The proof of $tr(C(u))=tr(C(w))$ is provided in Eq. (30). Here the definition of $C(u)$ and $C(w)$ are given by $C(u_i):=\mathrm{Var}(\partial\ell/\partial u_i)$ and $C(w_j):=\mathrm{Var}(\partial\ell/\partial w_j)$. We apologize for the confusion.

What is the distribution of $x$ in the experiment of Figure 1?...

This experiment is easy to reproduce. Here, $x$ is Gaussian, $y=x +\epsilon$, and $\epsilon$ is an independent Gaussian. As long as the variance of $x$ and $\epsilon$ are nonvanishing, one will be able to reproduce this experiment.

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

When the r.h.s. of Eq (2) does not vanish, Eq. (2) has a unique fixed point in the degenerate valley of the rescaling transformation. Namely, there exists a unique $\lambda$ such that the transformation $(u,w) \to (\lambda u, \lambda^{-1}w)$ makes Eq. (2) vanish. This is made precise with our additional update Theorem 3.1. See the summary rebuttal for more detail.

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

Yes. For example, the loss function is given by $\ell=\|\theta_1 LReLU(\theta_2^Tx + \theta_3)-y\|^2$. The network has the rescaling symmetry: $\theta_1\to\lambda\theta_1,\theta_2\to\theta_2/\lambda,\theta_3\to\theta_3/\lambda$. Therefore, we still have the law of balance by defining $u:=\theta_1$ and $w:=(\theta^T_2,\theta_3)^T$ in Eq. (2), which is the same as the ReLU activation.

What do you mean by 'the difference between GD and SGD at stationarity is O(1)'?...

Here, the O(1) difference means the difference of the final values between the SGD and GD does not depend on the noise strength $T$. Therefore, the GD can only approximate SGD for a finite time horizon.

Should the minimizer of (10) be...

This is a typo. The global minimizer should be $v^{*}=\beta_2/\beta_1’$ while $\beta_1’:=\beta_1+\gamma$. We will complement it in the revision.

Below (10), what are the $\alpha_i$?...

Yes, the definitions of $\alpha_i$ are the same as in Eq. (8).

I find that...differs from the value proposed in the article by a factor of $4/\alpha_1$.

Thank you for pointing this out. Here the definition of $\Delta$ is actually given in Eq. (11). We would like to revise the formula in l.189 as $\Delta:=\alpha_1\min_vC(v)/4$. We apologize for this typo.

I cannot find the proof of (12)...

The definition of $\beta_1’$ is given by $\beta_1’:=\beta_1+\gamma$. Eq. (12) is easy to prove and we will include a proof of Eq. (12) in the revision.

In Figure 3, what is plotted for $v$ in the depth 1 case?...

Yes, the parameter $v$ is defined as the product of $u$ and $w$.

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.

Thank you for your feedback. We will answer the weaknesses and questions below.

The law of balance is an interesting phenomenon, yet...

Thanks for raising this point. We stress that the law of balance applies to high-dimensional problems as well. In the high-dimensional case, the convergence to the noise-balance point is no longer strictly exponential, but the following three strong properties still hold:

1. The fixed point (namely, the noise-balanced point) of Eq. (2) exists and is unique (among all degenerate solutions connected through the rescaling transformation). This is an insightful new theoretical result we will include that strengthens the law.

2. Assuming that this fixed point does not change in time, the dynamics in Eq. (2) converges to this fixed point. This is another new theoretical result we will include that strengthens the law.

3. Assuming that $C_1$ and $C_2$ are full-rank, convergence to an $O(1)$ neighborhood of the fixed point will be exponentially fast. This is a direct corollary of Theorem 3.1, which also enriches and strengthens our result.

We also perform an additional set of experiments to validate the corollary. See the rebuttal summary and the attached pdf.

I have to say that I was a bit bothered by the general overselling of the paper:

Thanks for this criticism. We will do our best to tune down our statements to ensure that they are as accurate as achievable.

As said before the law of balance is truly valid asymptotically in one-dimension

As argued above, we stress that our law of balance (theorem 3.1) is generally valid for an arbitrary-dimension network, and this message strengthens with the newly added part and corollary to Theorem 3.1.

The stationary distribution..., which is not standard and does not resemble a diagonal network!

We will make it explicit that the stationary distribution is for a special model. Also, We have never claimed the model to be a diagonal network (and diagonal networks are no more realistic than our model). We also never claimed the model to be “standard.” Our claim actually has a very restrictive qualifier: that it is the first “general exact” solution of this specific model. We will emphasize these points.

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

It is true that 1d distributions are easier to derive, but our key contribution is not to derive a 1d distribution, but to reduce an arbitrary dimensional problem to 1d, and this is a special property of the SGD noise. This is not trivial and not an overclaim. Also, none of these references derived exact solutions. To be specific, at the beginning of Sec. 2.3 in the reference arXiv: 2402.01382, the authors applied the decoupling approximation to approximate the covariance of the gradient noise near the minima. In Lemma 4.4 of the reference arXiv: 2407.02322, the authors assume the strength of the noise to be a constant near the minima. In comparison, we give an exact solution for arbitrary initial parameters under the law of balance, which is not restricted to the points near the minima. We will include these references and clarification in the revision.

Minor typos/flaws:

l.41: Fokker Planck is not inherently high-dimensional

When the dimension of the dynamical variable is higher than 1, the Fokker-Planck equation is inherently high-dimensional. We will clarify this.

l.165: The law of balance is not strictly applicable here since $\ell$ is not scale invariant because of the regularization.

What we want to derive is the dynamics of $C$, which can be decomposed into two parts: (a) contribution from symmetry loss, and (2) contribution from weight decay. These two parts can be analyzed separately. The contribution from the symmetry part directly follows from the law of balance.

Section 4.1 Depth - 0: I think that $\Delta>0$ is not currently the "most practical example" since it corresponds to a underparametrized model.

This is a misunderstanding and is based on the questionable assumption that overparametrized models reach a zero training loss and are more practical. For example, large language models often reach a training loss far above zero and are certainly underparametrized. For conventional CV models, it is also almost never the case that it reaches a zero training loss. As long as the training loss is above zero, it qualitatively corresponds to the case $\Delta >0$.

Questions:

l.65: Eq (1) is not the usual covariance matrix that is used for Stochastic modified equations.

This is a typo. It should be $C(\theta)=\mathbb{E}[\nabla\ell(\theta)^T\nabla\ell(\theta)]-\mathbb{E}[\nabla\ell(\theta)^T]\mathbb{E}[\nabla\ell(\theta)]$.

l.95: Eq (3) seems important but it is difficult to follow where this comes from without intermediate calculations, can the authors develop?

This is a rewriting of the right-hand side of Eq. (2) by letting it equal $0$.

l.110: Same thing for the equation $\lambda^4=\langle w^{\*},C_2w^{\*}\rangle/\langle w^{\*},C_1w^{\*}\rangle$ (3): can the authors develop?

This is also a type of rewriting of the right-hand side of Eq. (2). By following the steps in l.108 and l.109, we let $u$ be $\lambda u$ and $w$ be $\lambda^{-1}w$. Then, by letting Eq. (2) equals 0, we have l.110.

Thm 4.2 : Are they the only stationary distribution? How to know to which one it converges ?

Yes. Theorem 4.2 enumerates all the stationary distributions. The solution which the network converges to depends on the initial parameters. If we initially choose a set of parameters such that $v>0$, then the probability distribution of the parameters converges to the solution $p_{+}$. For the initial parameters with $v<0$, they converge to the solution $p_{-}$.