Paper: ReLU soothes NTK conditioning and accelerates optimization for wide neural networks (ICLR-2024)

Reviewer comment (Theory):
Theorem 5.3 holds for a general two-layer network with the outer layer being fixed, but it implicitly requires that the input dimension \( d \) is bigger than the number of samples \( n \). If this is not the case, \( G \) is not full rank and the statement becomes trivial, as the smallest eigenvalue of \( G \) is 0, and it is well known that the smallest eigenvalue of the NTK is bounded away from 0. This requirement \( d > n \) is violated in all the experiments of Section 5.

Author response:
**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.

Reviewer comment (Evaluation):
The numbers reported in Figure 2(b) when \( L = 0 \) are suspicious. One would expect the condition number to be \(\infty\) since \( G \) has at most rank \( d \). This raises the question of whether something is missing in the interpretation.

Author response:
**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.

Reviewer comment (Evaluation):
Can the authors comment on the points reported in Figure 2(b) when \( L = 0 \)?

Author response:
**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.

Reviewer comment (Novelty):
There are existing works that study the impact of the activation function on the NTK, such as [R1] and [R2]. How do the results of the manuscript compare to such existing works?  

[R1] Panighahi et al., "Effect of Activation Functions on the Training of Overparametrized Neural Nets", ICLR 2020.  
[R2] Zhu et al., "Generalization Properties of NAS under Activation and Skip Connection Search", NeurIPS 2022.

Author response:
**Q3:** *“There are some works that study the impact of the activation function on the NTK, see [R1], [R2]. How do the results of the manuscript compare to such existing works?”*

**A:** The most relevant part of [R1] and [R2] is their theoretical bounds on the smallest eigenvalue $\lambda_{min}$ of NTK. [R1] provides an $\Omega$-type bound on $\lambda_{min}$ (in their Theorem 4.1 and 4.2). [R2] gives a lower-bound on $\lambda_{min}$ by Hermite coefficient (see Theorem 1 therein). In addition, both work studied the dependence of the lower-bound on the activation type. These results are surely interesting, we will cite them in the revision.  
However, based on these analyses, it is hard to judge whether any of these non-linear activations helps to decrease the condition number, because we are not sure about the tightness of the lower-bound. (Although [R2] provides an upper bound, it seems loose, for example for ReLU network, the upper bound for $\lambda_{min}$ is the number of layers $L$, see Theorem 1). On the other hand, we directly compare the condition numbers, not their bounds. 

In this paper, we showed a new effect for certain non-linear activations – decreasing the condition number of NTK. Before this, non-linear activations were known to increase the expressivity. We view this result as an addition to our understanding of the effect of non-linear activation functions. We don’t see similar results in either [R1] or [R2].

Reviewer comment (Writing):
The writing style and general clarity in some parts of the paper is lacking.

Author response:
**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.

Reviewer comment (Experiments):
The main theoretical results compare with the baseline of a fully linear network, which is not a very interesting comparison. It does not seem that surprising that a linear model cannot increase angle separation while the ReLU model can. A more interesting comparison would be with other non-linear models like kernel machines for example.

Author response:
**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.

Reviewer comment (Theory):
The theory of the improved convergence rates of gradient descent is restricted to the infinite or large width regime where the NTK is constant. However, it is observed in practice that the NTK changes significantly in most interesting cases [1], and this change corresponds to important feature learning. Per my understanding, the theory in this work fails to extend to this changing NTK regime.

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

Reviewer comment (Theory):
Does the conditioning result extend to the finite-width empirical NTK?

Author response:
First of all, we would like to point out that nearly all optimization theories for neural networks so far are either in the infinite width limit or with super large width (or more precisely, require sufficiently large network width) (see, e.g., [1-5]). To the best of our knowledge, there is no general theory for convergence of GD/SGD outside of large width regimes. The main technique in those works is the infinitesimal or sufficiently small change of NTK during training. 

Using the same technique of small NTK change, Theorem 5.2 (NTK for deep networks with two data samples) can be easily extended to the setting of finite but large width (as can be seen from its proof, the eigenvalues have closed form expressions, and can be quantatively analyzed). As for Theorem 5.3, in order to extend to finite width, we need a bit of extra analysis that can quantitatively control the difference of NTKs between with and without ReLU. But as the NTK change is small, we believe it is not a hard extension.

Indeed, we don’t view the infinite width limit as a significant limitation. Please see the original NTK paper [6], as well as [7] for example. The analyses of these works are conducted in this infinite width limit. It turns out that these results have brought significant insights to neural networks.



[1] Allen-Zhu, Li, Song. A Convergence Theory for Deep Learning via Over-Parameterization. ICML. 2019.

[2] Du, Zhai, Poczos, Singh. Gradient Descent Provably Optimizes Over-parameterized Neural Networks. ICLR, 2018.

[3] Zou, Cao, Zhou, Gu. Gradient descent optimizes overparameterized deep ReLU networks. Machine Learning, 2020.

[4] Oymak, Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. : IEEE Journal on Selected Areas in Information Theory, 2020.

[5] Liu, Zhu, Belkin, Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. ACHA, 2022.


[6] Jacot, Gabriel, Hongler. “Neural tangent kernel: Convergence and generalization in neural networks”. NeurIPS. 2018.


[7] Lee, Xiao, Schoenholz, Bahri, Novak, SohlDickstein, Pennington. “Wide neural networks of any depth evolve as linear models under gradient descent”. NeurIPS 2019

Reviewer comment (Experiments):
The major conclusions of this paper (about the advantages of ReLU) are only demonstrated in comparison with linear networks. This makes the results not very strong, as the advantages of non-linear activations over linear networks are fairly clear.

Author response:
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.**

Reviewer comment (Experiments):
Since the comparisons are only made between ReLU networks and linear networks, the results of the paper may not be very comprehensive. For example, in the title, “ReLU soothes NTK conditioning” may give readers the impression that ReLU activation has a unique property when compared with other activation functions, which is not really the case. The results would be more comprehensive if the authors can extend the comparison between ReLU and linear activation functions to a comparison between a general class of non-linear activations and the linear activation.

Author response:
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).

Reviewer comment (Novelty):
The results of this paper may not be very surprising. As the authors mentioned, the major known advantage of non-linear activation is to improve the expressivity of neural networks. It seems that the conclusions of this paper, to a large extent, are still saying the same thing. Better data separation, better NTK conditioning, and faster convergence to zero training loss, all seem to be more detailed descriptions of better expressivity.

Author response:
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.

Reviewer comment (Evaluation):
The impact of the results are not sufficiently demonstrated. For example, it is not very clear what is the benefit to achieve better data separation in the feature space.

Author response:
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.