Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization

We thank the reviewer for the comments, please see our response to your questions/concerns:

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.

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.

We thank the reviewer for the positive review. Please see our response to your questions/concerns:

Assumption on the data and initialization: We agree with the reviewer that our assumptions that data is positively correlated and initialization is balanced are often unrealistic. However, we believe they are not necessary for the dynamics behavior described by our Theorem 1 to happen in practice. We would like to point out that our numerical experiments enforce NONE of these assumptions: 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), and we initialize the weights by Gaussian initialization with small variance (not balanced). Despite not strictly satisfying these assumptions, 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.

We hope our response addresses your concerns. If you have additional questions/concerns, please feel free to post another comment. We will respond to them before the discussion period ends.

We thank the reviewer for acknowledging the strengths of our paper. However, we strongly (yet respectfully) disagree with the reviewer's comment on the weaknesses. Please see our response to your questions/concerns:

Goal/Contribution of this paper: We believe the reviewer's assessment of the contribution of our paper is inaccurate. The well-separated dataset studied in this paper is indeed easy and a linear classifier can achieve perfect accuracy. However, the point of this paper is NOT to find a classifier that fits the dataset, but rather to understand (a) how one can successfully train a ReLU network to fit the data, without getting stuck at some local minimum (as Reviewer Qhf5 pointed out, this can happen when training ReLU on linearly separable data); (b) why using a small initialization in the training of a two-layer ReLU network will lead to a classifier with special property (in our case, low-rank), while a network trained using a large initialization scale does not have such property [Chizat'19]. Through a complete, non-asymptotic, quantitative analysis of the GF on two-layer ReLU networks with small initialization, our results show that one can achieve the global minimum of the loss (with a characterization of the convergence rate), and provide a detailed analysis of how network weights evolve during training, along with its connection to the low-rank property of the weight matrix. As such, our paper makes a significant contribution toward understanding the implicit bias of small initialization in training a ReLU network.

Assumption on the data: Although we assumed separable data for our analysis, we believe it is not necessary for the dynamics behavior described by our Theorem 1 to happen in practice. 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 only 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.

We hope our response addresses your concerns and the reviewer can re-evaluate our manuscript based on our response. If you have additional questions/concerns, please feel free to post another comment. We will respond to them before the discussion period ends.

Reference:

[Chizat'19] Chizat, L., Oyallon, E., and Bach, F. On lazy training in differentiable programming. NeurIPS, 2019.

We thank the reviewer for the constructive comments, please see our response to your questions/concerns:

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.

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.

Assumption on the data: We refer the reviewers to the global response we posted for a detailed discussion about our assumption on the data.

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.

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.

Reference:

[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