Paper: Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization (ICLR-2024)

Reviewer comment (Theory):
The setting of the neural network is unconventional. It requires the second layer weights to depend on the first layer weights in a way as shown in Eq.(3), instead of being independently initialized. This setting is used neither in practice nor in most theoretical analysis. According to the analysis, there is doubt that the results of this paper hold without this restriction on the second layer weights.

Author response:
**Clarification on balanced initialization**: We would like to point out that our balanced initialization assumption is made for theoretical analysis, and we believe it is not necessary for the dynamics behavior described by our Theorem 1 to happen in practice: All experiments in our manuscript use Gaussian initialization with small variance $\alpha^2$ with some $\alpha<<1$ (please see highlighted text on Page 9 for the details), and the numerical results can be well explained by our Theorem 1, despite not having a balanced initialization. 

**What happens if initialization is not balanced**: Having random Gaussian initialization (not balanced) for both layers only introduces additional technical challenges for the analysis, and we would like to discuss them here in detail: Even if the initialization is not balanced, the imbalance is preserved, i.e., one has $v_j^2(t)-\Vert w_j(t)\Vert^2=v_j^2(0)-\Vert w_j(0)\Vert^2=\mathcal{O}(\epsilon^2)$, thus $v_j(t)=sign(v_j(t))\Vert w_j(t)\Vert+\mathcal{O}(\epsilon^2)$. The last equation is ALL we use in the analysis regarding the last layer. In the balanced case, there is no $\mathcal{O}(\epsilon^2)$ term, and we also have $sign(v_j(t))=sign(v_j(0))$ fixed. If the initialization is not balanced, there is an additional $\mathcal{O}(\epsilon^2)$ error term in all our analyses, but it can be well controlled when $\epsilon$ is sufficiently small. The only nontrivial part is that the $sign(v_j(t))$ is no longer fixed. Thus, one needs to keep track of the changes in $sign(v_j(t))$ over the GF trajectory, which is an interesting future research direction. Although there are additional technical challenges, NONE of them substantially change the GF dynamics: one still has the two-phase training (first alignment, then final convergence), and the neuron's directional dynamics still follow equation (4), with $sign(v_j(0))$ replaced by $sign(v_j(t))$. Therefore, our theoretical results would still hold without a balanced initialization assumption, and this is exactly what we show in the numerical experiments.

Reviewer comment (Theory):
The paper has some discussion on this unconventional network weight setting. However, it does not justify the validity of this setting. The fact that it is commonly assumed in other papers does not directly justify it. There should be some analysis, or at least some intuition, on why the results would hold in the natural setting.

Author response:
**Clarification on balanced initialization**: We would like to point out that our balanced initialization assumption is made for theoretical analysis, and we believe it is not necessary for the dynamics behavior described by our Theorem 1 to happen in practice: All experiments in our manuscript use Gaussian initialization with small variance $\alpha^2$ with some $\alpha<<1$ (please see highlighted text on Page 9 for the details), and the numerical results can be well explained by our Theorem 1, despite not having a balanced initialization. 

**What happens if initialization is not balanced**: Having random Gaussian initialization (not balanced) for both layers only introduces additional technical challenges for the analysis, and we would like to discuss them here in detail: Even if the initialization is not balanced, the imbalance is preserved, i.e., one has $v_j^2(t)-\Vert w_j(t)\Vert^2=v_j^2(0)-\Vert w_j(0)\Vert^2=\mathcal{O}(\epsilon^2)$, thus $v_j(t)=sign(v_j(t))\Vert w_j(t)\Vert+\mathcal{O}(\epsilon^2)$. The last equation is ALL we use in the analysis regarding the last layer. In the balanced case, there is no $\mathcal{O}(\epsilon^2)$ term, and we also have $sign(v_j(t))=sign(v_j(0))$ fixed. If the initialization is not balanced, there is an additional $\mathcal{O}(\epsilon^2)$ error term in all our analyses, but it can be well controlled when $\epsilon$ is sufficiently small. The only nontrivial part is that the $sign(v_j(t))$ is no longer fixed. Thus, one needs to keep track of the changes in $sign(v_j(t))$ over the GF trajectory, which is an interesting future research direction. Although there are additional technical challenges, NONE of them substantially change the GF dynamics: one still has the two-phase training (first alignment, then final convergence), and the neuron's directional dynamics still follow equation (4), with $sign(v_j(0))$ replaced by $sign(v_j(t))$. Therefore, our theoretical results would still hold without a balanced initialization assumption, and this is exactly what we show in the numerical experiments.

Reviewer comment (Theory):
The assumption on the data (Assumption 1) is strong, as it is not met by almost all real data.

Author response:
**Assumption on the data**: Although we assumed separable data for our analysis, we would like to point out that our numerical experiments do not enforce this assumption: For the two-MNIST-digits experiment, we show that centered data approximately satisfies our Assumption 1 (please see the left-most plot in Figure 5). Despite not strictly satisfying Assumption 1, the numerical results can be well explained by our Theorem 1. That being said, we are working diligently to extend our results to more general/relaxed assumptions in our future work. Moreover, we would like to refer the reviewer to the global response we posted for a detailed discussion about our assumption on the data.

Reviewer comment (Novelty):
The significance of the results is limited, as it is an extension of similar results from $\epsilon \to 0$ to the finite but small $\epsilon$.

Author response:
**significance of our results**: Our results are NOT some simple extension to prior results that study $\epsilon\rightarrow 0$. First of all, $\epsilon\rightarrow 0$ is impractical because all-zero initialization keeps the weights fixed at zero throughout training (or, equivalently, the origin is a saddle point of the GF dynamics). For non-zero $\epsilon$, if one makes the initialization scale, $\epsilon$, smaller and smaller, the GF dynamics will further slow down. This can be theoretically shown by our Lemma 3 in Appendix B.1: Given an $\epsilon$-small initialization, all weights are close to zero for at least $T=\Theta(\log\frac{1}{\epsilon})$ amount of time. When $\epsilon\rightarrow 0$, we have $T\rightarrow \infty$, meaning the weights do not move away from zero for an infinite amount of time! Therefore, any analysis with $\epsilon\rightarrow 0$ can not be implemented in practice. Our results work for a finite $\epsilon$, a choice that can be used in practice (We note that our analysis is for GF; there needs to be a choice of step size when applying our results to gradient descent, which can be obtained from [Elkabetz'21]). Secondly, it requires much more technical analyses for deriving quantitative bounds under a finite $\epsilon$, as we explained in "Weakness in prior analyses" on Page 4 and our proof sketch on Page 7, and our results make significant technical contributions to understanding the GF dynamics with small initialization.

Reviewer comment (Theory):
In Eq.(5), why is the R.H.S. independent of the network output $f(x_i)$? Or, why is there no term $y_i-f(x_i)$ which usually appears in the expression of gradients?

Author response:
**RHS of equation (5)**: During the alignment phase (first $\Theta(\log\frac{1}{\epsilon})$ time of the GF trajectory), we can formally show that the network's output is small for all input $x_i$, i.e. $\max_i|f(x_i)|\leq \mathcal{O}(\epsilon)$ (please see our Lemma 3 in Appendix B.1). Therefore, in the approximation we have in equation (4)-(5), we can approximate $y_i-f(x_i)$ by $y_i$, with an error of $\mathcal{O}(\epsilon)$. One minor point is that one would get $y_i-f(x_i)$ in the gradient if the $l_2$ loss is used. Here, we are using exponential/logistic loss, the gradient term has $y_i\exp(y_if(x_i))$, and we approximate it by $y_i$.

Reviewer comment (Experiments):
Do the results/analysis extend to other activation functions?

Author response:
**Do the results extend to other activation functions**: Yes. Please find a generalized version of Lemma 1 in Appendix B, where we extend the analysis of neuron directional dynamics to any two-layer network with a positive homogenous activation function of degree 1 (including all $\alpha$-leaky ReLU activations, $\alpha\in[0,1]$). However, we note that our main results (Theorem 1) will not directly apply as the neuron directional dynamics have changed as we consider an activation function different than ReLU, and additional technical analysis is required to establish the directional convergence for general leaky-ReLU functions.

Reviewer comment (Theory):
How small the initialization shall be to ensure the two-phase convergence? In general, since the training is nonconvex even with correlated/separable data due to the nonlinear relu activation function in the los. SGD/GD converges to a local minimum and indeed, this has been observed and numerically validated in the literature; see also [R1] Brutzkus et al. 2018. SGD learns over-parameterized networks that provably generalize on linearly separable data. ICLR. [R2] Wang et al. 2019. Learning ReLU networks on linearly separable data: Algorithm, optimality, and generalization. IEEE TSP. In [R1], SDG for two-layer ReLU networks under separable data converges to local minimum; yet for leaky ReLU networks, it finds global minimum. In [R2], it also shows that plain SGD on ReLU networks using separable data converges to local minimum numerically. Yet, a bit modification on the SGD helps SGD converge to a global minimum but with random initialization. It would be great if a comparison can be made between the approach in [R2] and the small-initialization SGD for training two-layer ReLU networks on e.g., MINIST. Moreover, is there any transition for such initialization value to go from the two-phase to single-phase gradient flow convergence to local minimum?

Author response:
**How small $\epsilon$ should be for two-phase convergence**: On page 8, we have a remark "choice of $\epsilon$" answering this question (we highlighted the text in the revision). In short, the initialization scale $\epsilon$ needs to be made small so that (a) the approximation of neuron angular dynamics in equation (4) is good enough (having specific error bound), and (b) the alignment phase lasts long enough so that directional convergence described in Theorem 1 can be achieved. The quantitative analysis in the appendices shows that (a) requires $\epsilon=\mathcal{O}(\frac{\sqrt{\mu}}{\sqrt{h}n})$ and (b) requires $\epsilon=\mathcal{O}(\frac{1}{\sqrt{h}}\exp(-\frac{n}{\sqrt{\mu}}\log n))$. For two-phase convergence, one needs $\epsilon$ to be sufficiently small to satisfy both.

**Comparison with some prior work using SGD**: We thank the reviewer for the references and added them to our comparison with prior work. However, we first would like to point out the key difference in the problem settings: In [R1, R2], the second-layer $v$ is **fixed** and not updated throughout the training, whereas in our manuscript, both the second-layer $v$ and the first-layer $W$ are updated by gradient flow dynamics. We took the reviewer's suggestion and did additional numerical experiments on our two-MNIST-digits classification problem (Please check the highlight section in Appendix A.5) by plotting the quantitative changes in neuron norms and directions separately. We showed that two-phase training (align then fit) happens when weights in both layers are trained, starting from a small initialization (the scenario studied in our manuscript). In contrast, there is no two-phase training when only the first-layer weights are trained (the scenario in [R1, R2]), even from a small initialization.

Reviewer comment (Experiments):
If the second-layer weights are initialized not in a balanced manner (although not initialized according to (3)), I understand it would also work and guess that the imbalance between positive v_j and negative v_j values only influences the time it takes for the two-phases. More balanced initialization, faster convergence. It would be interesting to numerically validate if this is the case.

Author response:
**Clarification on balanced initialization**: We would like to point out that our balanced initialization assumption is made for theoretical analysis, and we believe it is not necessary for the dynamics behavior described by our Theorem 1 to happen in practice: All experiments in our manuscript use Gaussian initialization with small variance $\alpha^2$ with some $\alpha<<1$ (Please see highlighted text on Page 9 for the details), and the numerical results can be well explained by our Theorem 1, despite not having a balanced initialization. Using random Gaussian initialization for both layers only introduces additional challenges for the analysis. For example, an important one among these technical challenges is that the $sign(v_j(t))$ is no longer fixed if the initialization is not balanced, thus one needs to keep track of the changes in $sign(v_j(t))$ over the GF trajectory, which is an interesting future research direction.

Reviewer comment (Writing):
There are some grammar issues and typos. Please correct.

Author response:
We hope our response addresses your concerns (we will also proofread our manuscript to minimize grammar issues and typos). If you think there are additional experiments we can do, or clarifications we can make, please feel free to post another comment. We will respond to them before the discussion period ends.