ReLU soothes NTK conditioning and accelerates optimization for wide neural networks

We thank the reviewer for the comments. We would like to address your concerns one by one below.

W1: “... it is more or less a direct consequence of [1] as [1] also shows the formula for the angle between post-activation.”

A: The formula for the angle between post-activation is merely the starting point of our analysis. Note that most of our analysis after this formula is very different from that of [1]. We will cite [1] around this formula in the revision.

We would like to draw the reviewer’s attention to our main contribution. The paper showed a new advantage of certain non-linear activations – decreasing the condition number of NTK. Before this, non-linear activations is just known to increase the expressivity. We view this result as an addition to our understanding of the effect of non-linear activation functions. In addition, this new advantage has important implications on the convergence rate.

W2: “the current results (Thm 4.2 and 4.4) are only considering the points that are very close … , if L is big, this quantity is very small.”

A: The small data input angle regime is the key and most interesting part. This is because the NTK condition number is mainly affected by its smallest eigenvalue, which in turn is closely related to the smallest angle between points. A tiny change in the smallest angle between points may lead to a non-negligible change in the NTK smallest eigenvalue $\lambda_{min}$, which may further lead to a significant change in NTK condition number $\kappa$ (as $\kappa \sim 1/\lambda_{min}$, and $\lambda_{min}$ is very small).

W3: *“As the authors mentioned, for small $z$, $g(z)$ behaves like identity. … Thus, the model gradient angles are nearly non-changing from the input angles. This is also why the improvement for the condition number in Theorem 5.2 can also be very small. It would be more interesting to see ReLU can improve data separability for input pair with small angle (but larger than the regime this paper presents).” *

A: In the small but positive angle regime, a tiny difference between $g(z)$ and $z$ could make a large difference in the NTK condition number. As discussed above, the NTK condition number scales roughly as $1/z$ for very small $z$. Please take a look at the Figure 2(a), for a numerical verification of this better separation property on real-world datasets. It shows that the ReLU non-linearity can significantly increase the smallest angle between points. Moreover, please take a look at Figure 2(b), the condition number improved significantly, often by about $10^2 \sim 10^4$.

We thank the reviewer for the constructive comments. We are grateful that the reviewer recognized the novelty of the work. We also understand your concerns. We would like to address your concerns and questions one by one below.

W1: “... However, it implicitly requires that the input dimension $d$ is bigger than the number of samples $n$. In fact, if that is not the case, $G$ is not full rank, and the statement becomes trivial …”

A: We apologize for not making this point clear in the submission. In short, we do not require $d>n$. Technically speaking, it is the “effective” NTK condition number (largest eigenvalue divided by the smallest non-zero eigenvalue) that controls the convergence rate. When $G$ is not full rank, this relates to the smallest non-zero eigenvalue of $G$. We would like to explain this in detail below:

Let’s consider the case $d<n$, where the concern arises. Why it is the smallest non-zero eigenvalue?

– In this case, the Gram matrix $G$ is just the NTK of the linear model, and is not full rank. Note that this linear model is in the under-parameterized regime. As we know, for the linear model, the NTK $K=XX^T$ has the same spectrum as the Hessian of least square loss $H=X^TX$, except the zero eigenvalues. This Hessian is expected to have full rank, and the least square loss is convex. As is well-known, the condition number is defined as $\lambda_{max}(H)/\lambda_{min}(H)$, which is equivalent to $\lambda_{max}(K)/\lambda_{min}^*(K)$, with $\lambda_{min}^*(K)$ as the smallest non-zero eigenvalue of NTK.

– How about the linear network and ReLU network? In the infinite width limit, they are essentially linear models, with the model gradient $\nabla f(x)$ (instead of original input $x$) as the feature. These linear models are over-parameterized, and have a hyper-plane as the solution set $S$. Intuitively, the optimization (by gradient descent) occurs in sub-spaces that are perpendicular to $S$. In addition, the zero eigenspace of NTK are also perpendicular to $S$. Hence, the zero eigenvalues of NTK contribute nothing to the optimization. Therefore, the condition number of optimizational interest is the “effective” condition number (i.e., ignoring the zero eigenvalues).

In summary, our results do not require $d>n$. Our results essentially rely on the “effective” condition number of NTK. We will clarify this point in the revision.

W2: “Theorem 5.3 holds for a general two layer network with the outer layer being fixed”.

A: This setting is commonly used in literature, including very good papers, for example Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. “Gradient descent provably optimizes over-parameterized neural networks”. We also observe that many novel findings originally come out in simple settings, which does not affect its completion by later works. We believe that this simple setting is a good starting point. We are aware that for deeper networks it requires much complicated analysis. It is believable the same results hold for deep networks, as partially evident by Figure 2 and Theorem 5.2.

W3: “the numbers reported in Figure 2(b) when $L=0$ are a bit suspicious” & Q1: “Can the authors comment on the points reported in Figure 2(b) when $L=0$?”

A: In the experiment to compute the numbers in Figure 2(b), we actually evaluated NTK based on batches of size $512$. This is due to the limitation by our computational resource, which could not compute and store larger Jacobian matrices. The reported numbers are averaged over the batches. In this case, the Gram matrix $G$ is always full rank, hence it is not surprising that $\kappa$ is finite. One can anticipate that the condition numbers $\kappa$ at $L=0$ may increase with larger batch size, and perhaps be infinite if it exceeds $d$. We will clarify this setting in the revision.

We would like to thank the reviewer for the follow up question. Sorry for the confusion.

Please note that the NTK (a $n\times n$ matrix) and loss Hessian (a $p\times p$ matrix) have different spaces, where $p$ is the number of parameters which is much larger than $n$ in the over-parameterized networks. Also note that it is this $p$-dimensional parameter space that the weights are updated and gradient descent resides in. Hence, we don't quite see how "If the NTK is low rank, then gradient descent can be stuck precisely in the span of the eigenvectors corresponding to the 0 eigenvalues", as gradient descent is not in the same space as the eigenvectors.

For ReLU network, the NTK is indeed of full rank, as shown by Du et al. However, no matter NTK is of full rank or not. As $p\gg n$ for over-parameterized models, the loss Hessian, $p\times p$, is low rank and have a lot of zero eigenvalues. Our statement that you quoted is actually talking about the gradient descent in this $p$-dimensional parameter space.

We hope this can address your concern. We are also happy to discuss if you have further questions.

We thank the reviewer for the comments. We would like to address your concerns one by one below.

W1: “The writing style and general clarity in some parts of the paper is lacking”

A: Could you provide more details? For example, which part of the paper is not clear to you. We are happy to revise wherever was not clear.

W2: “The main theoretical results compare with … a fully linear network … is not a very interesting comparison.”

A: First of all, we would like to restate our main contributions here: the paper showed a new advantage of certain non-linear activations – decreasing the condition number of NTK. (in parallel with the well-known advantage of increasing the expressivity of network functions).

To show this new advantage in decreasing condition number, it is necessary to compare with the linear network that has exactly the same architecture except the non-linear activation. It was this comparison that makes this new advantage evident. A comparison with other non-linear models (e.g., kernel machines) does not support such a claim (although we agree that this kind of comparison is an interesting future direction).

Analogously, think about the other well-known advantage of non-linear activation – increasing the expressivity of network function. This better expressivity was made clear by the comparison between non-linear network and linear network, but not between different non-linear activations.

We would like to point out that this paper is never meant to say the ReLU model is superior to any other models. Instead, we aim to deliver a message on understanding the role of non-linear activations.

W3: “It does not seem that surprising that a linear model cannot increase angle separation while the ReLU model can.”

A: We are not aware of any similar results in existing works.

W4: “The theory … is restricted to the infinite or large width regime… in practice the NTK changes significantly …”

A: We would like to point out that many recent theoretical works are conducted in this large width regime, for example [1,2,3,4]. This regime, although many people believe impractical, provides a lot of new understanding of modern deep learning. We believe this is a very good starting point.

[1]: Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. “Reconciling modern machine learning practice and the classical bias–variance trade-off”. PNAS 2019.

[2]: Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107– 115, 2021.

[3]: Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. “Gradient Descent Finds Global Minima of Deep Neural Networks”. In: International Conference on Machine Learning. 2019.

[4]: Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. “Stochastic gradient descent optimizes over-parameterized deep relu networks”. In: arXiv preprint arXiv:1811.08888 (2018).

[5]: Chaoyue Liu, Libin Zhu, and Mikhail Belkin. “Loss landscapes and optimization in over-parameterized non-linear systems and neural networks”. In: Applied and Computational Harmonic Analysis 59 (2022).

We thank the reviewer for the comments. After carefully reading the review, we think the reviewer’s concerns are largely based on a misunderstanding of the paper. We apologize for not making it crystal clear in the submission. We hope the following explanation can address all the concerns.

W1: “The major conclusions of this paper ……,as the advantages of non-linear activations over linear networks are fairly clear”.

A: We politely could not fully agree with this point. The advantage of non-linear activations was only partially clear. It was only known that non-linear activations increase the expressivity of the network function (i.e, can approximate more complicated functions). However, our paper showcases a new advantage: decreasing the NTK condition number (which in turn induces a faster worst-case convergence rate). This advantage was not known before.

W2: “The comparisons are only made between ReLU and linear networks …”

A: As we clarified above, our major contribution is to showcase an never-noticed advantage of certain non-linear activations.

Let’s make an analogy. Think about the advantage of non-linear activation in increasing expressivity. When talking about this advantage, it was the comparison between non-linear network and linear network that makes this advantage clear. Comparing between different activations does not support this known advantage of increasing expressivity.

Same thing here. To show this new advantage of decreasing condition number, it is necessary to compare with the linear network that has exactly the same architecture except the non-linear activation. Comparing different activations only implies which activation might be relatively better or worse (we agree this is an interesting point, but not the main scope of this paper).

W3: “The results of this paper may not be very surprising. … the conclusions of this paper are still saying the same thing. .. all seem to be more detailed descriptions of better expressivity”.

A: We could not agree with this point, either. This new advantage of certain non-linear activation (i.e., better separation, decreasing NTK etc) is absolutely different/independent from the better expressivity. We don’t see any connections between the two, and are not aware of any existing result making this connection. If you think the two advantages are the same, please provide details/evidence. We are happy to discuss this point.

W4: “It is not very clear what is the benefit to achieve better data separation in the feature space.”

A: As we elaborated in the paper, the better separation is deeply related to the better NTK conditioning, as well as the faster worst-case convergence rate. Loosely speaking, it is the better separation that leads to the better NTK conditioning. Intuitively, a better data separation in the feature space helps the learning of the model. Think about two similar data points with different labels, which is often hard to be distinguished by a model. With a better data separation in the feature space, it becomes easier for the model to distinguish. We will add a discussion of the intuition in the revision.

Overall, we think the newly discovered advantage of certain non-linear activation (i.e., better separation, better NTK conditioning, etc) is important for theoretical understanding of the neural network, and should be inspiring for further research. We hope the reviewer can recognize the novelty and importance of this discovery.