Noise Balance and Stationary Distribution of Stochastic Gradient Descent

1. Does theorem 1 directly lead to balancedness? Notice that both $C_1$ and $C_2$ are functions of $u$ and $v$, and in (6) the matrices can be very ill-conditioned, which cannot directly imply the balance between $\|u\|$ and $\|w\|$. I believe in special cases like two-layer ReLU networks/diagonal linear nets, this can be correct as proved in the theorems. However, this paper states the "law of balance" as one of its main contributions. In the current form of this paper, the argument is limited to toy cases (both theory and experiments) and does not generalize to practical settings. To provide enough evidence, the authors should include more evidence (either theoretical or empirical) to verify the universality of this balance condition. For example, the authors could include more experiments with deep ReLU networks ($L\ge 3$) and check if the experiments also satisfy this balancedness. Otherwise, I suggest the authors put more focus on the theoretical results of the diagonal linear net and those empirical observations, and reconsider the paper's structure/message.

2. I don't think the edge of stability argument (Line 440-460) is valid. EoS is a phenomenon observed in full-batch GD cases without any gradient noise. However, the authors claim their deep model can explain the phenomenon by "the fluctuation of model parameters decreases as the gradient noise increases", which is clearly incorrect because there is no gradient noise for EoS with full batch GD. The authors should rethink what is the correct implication of this fluctuation inverse phenomenon instead of this argument.

3. Now the diagonal linear net only has 1-dimensional inputs, which is a minor limitation. I understand it could be too complicated to generalize to high-dimensional inputs so a minimal analytical model is more or less acceptable.

The most relevant missing piece of information is a paragraph describing the long line of work about the fact that the implicit regularization effect of mini batch SGD is not explainable through means of continuous stochastic differential equations. Starting from Yaida (2018), continuing with Smith et al., "Origin of..." (2021), and Li et al., "What ..." (2022), there has been a long line of works (40+) which deal with this problem. I believe discribing them and discussing the applicability to this setting it's absolutely necessary for a paper dealing exactly with this to be considered for publication.

See also Questions.

The analysis is limited to a simplified “linear network” setting, which, while ideal for capturing rescaling symmetry, may not fully translate to practical networks where architectures involve skip connections and layer normalization layers. The applicability of these findings to state-of-the-art (SOTA) models might be limited, given that current deep learning architectures often deviate from the minimal settings examined in the paper.

Inconsistency of the arguments. In the equation (1), the update rule is SGD and $T=\eta/S$. In Section 4.4, the discussion of "infinite D limit" shows that $d/D *S/\eta=const$ will give $D\propto 1/\eta$. However, in the Figure 7, the network is trained with Adam whose scaling behavior is much different from that of SGD.

In terms of presentation, there are many temporary notations $\alpha_1, \alpha_2, \alpha_3$, ..., which are hard to understand how these parameters changes when varying the problem parameters.

The current analysis focused on the pure theoretical. It is extremely interesting and helpful to show that whether the insights obtained from such solid analysis could be applied to complex task settings and practical architectures. If this could be made from either more theoretical or empirical justification, it would a much stronger paper. These apsects include whether stationary distribution could be obtained or not, whether the constant depth/width rule still holds or not, sign coherence happens or not, etc.

I think after many years of deep learning theory research, the entire ML community or industry call for a more insightful and practically useful results from theoretical understanding. This could benefit more and raise more impacts.

This was an emergency review, so unfortunately, I do not have as much time as I wish to parse it. However, the main weakness I find is that the writing, claims, and flow are confusing. Theorems in the paper are opaque: setting, preciseness, and assumptions are often unclear. I looked at the appendix, and it did not help much.

I will place here my questions, which highlight weaknesses:

1. Theorem 3.1, especially in combination with Fig.1 is very confusing. To study the quantity you have, you use Ito's Lemma - and I see no problem with this. It is however highly unclear which noisy gradient model you choose (in the paper and even in the appendix). The usual model of gradient noise is additive Gaussian noise, and from what I understood in the appendix, you also use this. What is the difference between "SGD" and "Gaussian" in Figure 1? If you have an additive gaussian model, then these should be the same. If not, then you need a section to explain what you are doing!

2. Thm 3.1 Is stated as a differential equation driven by population quantities. I believe this equation only holds in expectation, correct? Even from the appendix, I find this hard to parse; I believe (related to point (1)) that you are claiming in some way that noisy components cancel? How can C1 and C2 then be expectation quantities?

3. Still about Thm 3.1. I find it very unclear how to interpret this: What does the equation converge to? Are C1 and S2 positive and definite? I guess it depends, and the equation alone cannot predict if $u$ and $v$ get close together. Yet I sense that this is later assumed in further statements.

4. Thm 4.1: the proof of this statement is very confusing: I cannot track what the authors assumed from the appendix. Do you presume balance? I am sorry that I could not get much out of this.

5. From Section 4.3 I got lost. I agree that formula (15) is potentially interesting, but there are so many open points left that I could not continue confidently.