Better NTK Conditioning: A Free Lunch from (ReLU) Nonlinear Activation in Wide Neural Networks

We thank the reviewer for the comment. We address your concerns and questions below.

Re: “The main limitation of the paper lies in the narrow scope of the comparison. Specifically, the focus is on ReLU versus linear networks.… A more compelling direction would have been to explore and compare a broader range of activation functions.”

A: We’d like to point out that our goal is to extract the influence by the existence of non-linear activations. The only way to fulfil this is to directly compare non-linear activated NN with linear NN. Comparing among different non-linear activations does not work for this goal. Analogously, please think about another important property of non-linear activations – increasing expressivity of NNs: with non-linear activations NN can potentially approximate any continuous function; without it is basically a linear model.

In addition, we indeed considered other activation functions, as shown in Figure 2. Again, these need to be compared with linear NNs in order to see the influence of these activations.

Re: “... the impact of non-linear activations in neural networks is already well established, making this result less surprising.”

A: We didn’t find any similar work in the literature, and we don't think our result is previously known. We’d like to ask the reviewer to provide additional clarification and references.

Re: “... in Theorem 2.2 (linear network), the feature space is attributed to the gradient at initialization.”

A: Due to the constancy of NTK [Jacot et al. 2018] and transition to linearity [Liu et al. 2020], the gradient features during and after training keep the same as at initialization. Hence, an analysis at initialization is enough.

[Jacot et al. 2018] Jacot et al. Neural tangent kernel: Convergence and generalization in neural networks, NeurIPS, 2018.

[Liu et al. 2020] Liu et al. On the linearity of large non-linear models: when and why the tangent kernel is constant. NeurIPS, 2020.

Re: “In Theorem 4.2 (non-linear network), only the first layer is trainable, while all subsequent layers remain fixed.”

A: It is often an acceptable setting in theoretical works to fix upper layers to simplify analysis. See for example [Du et al. 2019]. Indeed, we have a section in Appendix F to relax this constraint on upper layers, and similar results also hold. In addition, in our numerical results, all layers are trainable and results are consistent with our theory. Hence, we believe the result should still hold when the constraint is removed.

[Du et al. 2019] Du et al. Gradient Descent Provably Optimizes Over-parameterized Neural Networks, ICLR, 2019.

Re: “In fact, for similar data points, increased separation may be undesired, especially when the labels are the same.” and Q2: “Could you clarify how the separation property of the data points relates to their corresponding labels?”

A: We don’t quite think increasing separation is undesired when the labels are the same. In fact, for the training data in the same class, whether there is more or less separation among them does not really matter. The classification task is to find decision boundaries between different classes. Therefore, as long as the feature separation is increased for those with different labels so that the decision boundaries need less fine-tuning, the optimization becomes easier.

Re: “The spectral properties of the NTK have been extensively studied in prior work. This manuscript focuses specifically on the condition number of the NTK, which is closely tied to its smallest eigenvalue.”

A: As we also mentioned in the related works, there are extensive studies on NTK spectrum and its smallest eigenvalues. Our work differs from the prior works in the sense that we are the first to investigate the influence of the existence of non-linear activations. We also have the better separation phenomenon (Section 3) which cannot be found anywhere else. Our analysis techniques are also different.

Re: “while the manuscript claims that the condition number is small (line 57), it appears to be quite large when applied to larger datasets”

A: We apologize for not making it crystal clear. Generally, the condition number is extremely large due to an infinitesimal smallest eigenvalue, and is far greater than the data size $n$. In Line 57, we say it is small, meaning it becomes much smaller. We will make this point clear in the revision.

Re: “Proposition 4.1 has limited applicability, as it is demonstrated only in the context of small datasets where the number of samples n is less than the dimensionality d.”

A: When $n$ is greater than $d$, the smallest eigenvalue of linear NTK becomes zero and its condition number is infinite. The better separation, Theorem 3.2, and better NTK conditioning, Theorem 4.2, still hold.

Q1: “What are the practical implications of your theoretical findings?”

A: We consider our paper as a theoretical work that provides insights and understandings of deep learning. Two possible practical implications include: (1) we may need to rethink the current NTK-based optimization theories which often assume a fixed condition number, our paper suggests the activation function can influence this condition number hence the convergence rate. (2) Another implication of our better separation theorem could relate to the origin of adversarial examples where the adversarial and clean examples can be considered as two similar inputs. Our theory indicates that, due to the non-linearity of activation function, they are quite largely separated in the gradient feature space which may attribute to the observed distinct class predictions between adversarial and clean examples.

Q3: “Figure 2 provides numerical evidence that the separation phenomenon also holds for various non-linear activation functions. However, the behaviour differs across ReLU², GeLU, and Tanh. Can your theoretical framework be extended to accommodate other non-linear activations and account for the observed differences?”

A: The main message we had in Figure 2 is that, similar to the ReLU case, for small input angles (i.e., left end of each picture), non-linearly activated NNs are better separated; and when depth increases, the separation is larger. Therefore, the behavior is not significantly different from ReLU.

As for the extension beyond ReLU, the technical difficulty is to carry out a few integrations for other activations (similar to the ones in Line 652 and Line 661). In the paper, we numerically estimated them and showed the results in Figure 2. Therefore, if there exist close-form expressions for these integrations, our theory will extend to those non-ReLU activation functions.

We thank the reviewer for the comment. We address your concerns and questions below.

W1: “The title should be specific to ReLU activations as the main contribution of this paper is the theoretical findings, which is only conducted on ReLU.”

A: We thank the reviewer for this important suggestion. We will update the title in the revision to include ReLU. Please also note that we present numerical results for other types of activation functions in section 3 (seen in Figure 2).

W2: “This paper demonstrates that as depth increases, the condition number of the NTK matrix with ReLU improves. However, many previous works [1] [2] show that as depth increases, the NTK degenerates. That is, the infinite depth limit of the NTK regime is trivial.”

A: Yes. As our Theorem 4.3 suggests, the NTK becomes trivial for the infinite depth limit, which is consistent with [1,2]. But our main focus is on the finite depth scenario, where NTK and gradient features may not be trivial. At a practical depth, a better feature separation can make classification tasks easier as there is less need to fine-tune the decision boundaries. We focus on how different depths lead to different levels of gradient feature separation and how it affects the NTK matrix.

W3: “.... If the analysis could be extended to arbitrary non-linear activations, it could offer valuable guidance for the design of new activation functions.” and Q1: “What are the specific technical challenges in extending the theoretical analysis to other activation functions?... have you considered analyzing the Hermite coefficients? Perhaps the higher-order coefficients could be bounded or even omitted.”

A: Yes, the technical challenge is to carry out a few integrations for other activations (similar to the ones in Line 652 and Line 661). Instead, we numerically estimated them and showed the results for non-ReLU functions in Figure 2.

We appreciate your valuable suggestion on bounding or omitting higher-order Hermite coefficients. We will try it.

We thank the reviewer for the comment. We address your concerns and questions below.

W1: “The theory depends on infinite width of the network, and is restricted to ReLU activations.” and Q1: “Can the authors comment on finite width effects?”

A: Our theory is conducted in a sufficiently large but finite width setting (See for example Theorem 3.2). Basically, our theory is based on the well-established large-width NN analysis (such as [8,9,15,17,19] in the reference list of our paper), where the NN deviates from an infinitely wide NN slightly (in the order of $O(1/\sqrt{m})$). While there is not much theory for smaller network width, we temporarily only consider this sufficiently large but finite width setting.

In fact, we considered other activation functions including ReLU², GeLU and Tanh (See Figure 2 and its discussion). Numerical evaluations show that similar results hold for these activation functions. The difficulty for theoretical analysis of these non-ReLU activations is to carry out the integrations similar to those in Line 652 and Line 661. Nevertheless, these integrals can be estimated numerically.

W2: “The section on optimization acceleration could benefit from more quantitative analysis. For example, a quantitative convergence improvement as a function of depth, informed by the main result would be insightful. … Currently, the faster convergence for higher depths is qualitative and does not build on the main result in a theoretical capacity.”

A: We agree that a quantitative result on the optimization acceleration would be more interesting. While we can get a quantitative result in the simple setting of data size equal $2$ based on the quantitative result of feature separation – Corollary 3.3, the major difficulty for general case is that the connection between the smallest eigenvalue of NTK and data feature separation is still not quantitatively clear, therefore, the quantitative result of feature separation downgrades to a qualitative result of NTK. We conjecture that a sophisticated analysis involving random matrix theories can achieve this goal. We consider it as a future work.

Q2: “Further, instead of increasing depth, could similar results be derived as a function of width for a fixed depth?”

A: Lemma 3.1 and Theorem 3.2 show that, with a fixed depth, the width contributes minimally compared to a width-independent term (see for example Eq. (7)). Our paper currently focuses on the leading term (i.e., the width-independent term). As for the width-dependent term, we believe it also depends on the initialization seed, in which case we only obtain an upper bound.

We thank the reviewer for the insightful and detailed comments. We’d like to address your concerns and questions below.

W1a: “Reference [1] shows that a transition occurs for the NTK between a “freeze” (order) phase ... and a “chaos” phase. The current paper … looking specifically at what happens at the boundary, but does not explicitly mention this connection.”

A: We thank the reviewer for providing this related work. We will cite it in the revision and clarify this connection. Basically, our theory only considered the bias-free setting, which is widely adopted by many prior theoretical works, such as [Du et al. 2019]. After reading [1], we agree that the bias makes a big difference (bigger than we previously thought) and it is interesting to extend the analysis to include the bias term. We also realize that this extension is technically non-trivial (just like observed in [1], different scaling $\beta$ of the bias leads to distinct results for NTK). We would like to consider it as a future work.

[Du et al. 2019]: Du et al. Gradient Descent Finds Global Minima of Deep Neural Networks. ICML 2019.

W1b: “I believe that the linear NTK has been studied before or is classical knowledge (it results directly from the recurrence relation in Theorem 1 of [2]).”

A: While the linear NTK result, i.e., Corollary 2.3, can be a direct consequence of Theorem 1 of [2], the feature angle separation theorem, Theorem 2.2, does not follow [2]. We didn’t find any prior work having the result of Theorem 2.2, although it is not technically hard. Hence, we put it in the Preliminary section and do not claim it as a contribution. In addition, we will clarify the connection of Corollary 2.3 and Theorem 1 of [2] in the revision.

W1c: “The improved conditioning property…between linear and non-linear activation were already known. Specifically, the linear NTK has infinite condition number when the number of data points exceeds the dimension. Whereas … the NTK restricted to the sphere is positive-definite as soon as the activation function is non-polynomial.”

A: The known result is only for the special case of an infinite linear NTK condition number. In addition, it also requires the data lies in a unit sphere. Our result not only applies for more general settings (e.g., linear NTK is also positive-definite), but, more importantly, focuses on the better separation of gradient features, which is not known from prior works. We’d like to cite [2] and add a discussion about the connection.

W1d: “Could you discuss also relation with [3].”

A: The paper [3] provides both upper and lower bounds for the smallest eigenvalues of NTKs for both shallow and deep NN. First, while these bounds are insightful and valuable, it is different from our results as it is not clear in [3] what effect the activation function has on the smallest eigenvalue. In addition, [3] obtained a linearly increasing upper bound for the smallest eigenvalue: $\lambda_{min}< L$. However, note that when depth $L$ increases, NTK accumulates (an addition of $L$ positive definite terms), also resulting in an linear increasing top eigenvalue, hence condition number is still independent of depth $L$. Therefore, we also don’t see a depth-dependent NTK condition number in [3], which is also different from ours. Moreover, we conjecture that the bounds obtained in [3] is not tight enough to extract the effect of non-linear activations and depth. As we argued above, upper bounds of [3] would leads to a depth independent bounds for condition number, which hides the depth dependent nature of the true value of condition number, due to a upper bound that is not tight enough.

W2: “It would be interesting to explicitly discuss the role of bias, as they can significantly affect the behaviour of the NTK.”

A: Yes, we agree that it is interesting to discuss the role of bias. With a quick trial, we found that the analysis with different scaling $\beta$ of the bias is technically non-trivial. Similar to the distinct results observed in [2] for different $\beta$, we believe the analysis of bias could potentially lead to new observations, which will be quite interesting. We consider it as a future work.

W3: “The paper … does not discuss the generalization performance... Empirical results exploring generalisation behaviour would be give valuable additional insights.” and Q1: “Fast convergence can occur with a decrease of generalization…. Could you discuss this trade-off explicitly? Specifically, how does generalization behave in Figure 4 (the experiments that illustrate the increase of training speed) ?”

A: We agree there is often a trade-off between optimization and generalization, especially when depth is large. We will add discussions about this trade-off in the revision.

As our theory suggests in Theorem 3.5 and Theorem 4.3, in the extreme case of infinite depth, any input pairs become equally separated in gradient features regardless of their original similarity. Even though not mutually orthogonal, this could also result in a trivial generalization: close to random guess for unseen data. This somehow can also be obtained from [Wang et al. 2020], where they dropped the initial random guess value and obtained a zero prediction for unseen data. We will cite this paper in the revision.

As for finite depth, there is no theory that clearly suggests at what depth this trade-off starts to happen. In our experiments in Figure 4, we can slightly see this trade-off. For example, for MNIST, the test accuracies for different depths $L$ are:

 95.98% (L=1),  97.43% (L=3),  97.57% (L=6),  97.52% (L=8),  97.39% (L=10),  97.19% (L=12)

Above, we can see a decreased generalization starting from a certain depth (in this case around $L=8$).

[Wang et al. 2020]: Wang et al. Why Do Deep Residual Networks Generalize Better than Deep Feedforward Networks? — A Neural Tangent Kernel Perspective. NeurIPS. 2020.

Q2: “... in the ReLU setting, angles converge to around 75.5 degrees, whereas with other activation functions used in Figure 2, it seems that the angles converges to around 90 degrees, and this could potentially reduce generalization properties. Could you provide the generalisation for these activation functions in the setup of Figure 2 and discuss this depending on what you obtain? Also can you report the limiting angles somewhere?”

A: As mentioned above, very deep ReLU NNs already have a reduced generalization performance and converge to random guess at infinite depth, hence we don’t think the difference between $75.5$ degree for ReLU and $90$ degree for others make much difference in generalization.

We will also add the (numerical-based) observations about the limiting angles of other activation functions in the revision.

Q3: “In corollary 3.3, I suggest you could explicitly mention an insightful consequence: $L$ should be of order $1/\theta_{in}$ for $\phi$ to be non trivial. Do you agree with this statement?”

A: Yes. We agree. We will update it in the revision.

Q4-1: “Isn’t Lemma 4.1 a consequence of the Law of Large Numbers …”

A: There is no Lemma 4.1 in our submission. We’d like to ask which lemma you refer to. We are happy to make further clarification after.

Q4-2: “I believe it’s inaccurate to claim ‘solving the above equation (40), we found that …’ …Do you agree?”

A: Thanks for pointing out this issue. Yes, you are correct. We will update our proof in the revision. Basically, based on Eq. (40), we can see that $1/4$ is an attraction point for the series $\cos \phi_L$ (look at the sign of $Delta \cos \phi_L$). With a bit of effort we can show that the series converges. Then, suppose $\cos \phi_L$ converges to a point not equal to $1/4$, the l.h.s. of Eq, (40) will be at the order of $1/L$, which results in a non-converging $\cos \phi_L$, contradictory. Hence, we can show that the series converges to $1/4$.

We thank the reviewer for the comment. We address your concerns and questions below.

W1 : “Focus solely on ReLU activations without considering a broad array of possible functions around may miss some insights.” and Q1: “... can the theoretical framework be extended to other activation functions and their derivatives? Would the conclusions (and numerical limits) still hold?"

A: In fact, we considered other activation functions including ReLU², GeLU and Tanh (See Figure 2 and discussions at Line 238). Numerical evaluations show that similar results hold for these activation functions. We believe the results holds more generally. The difficulty for theoretical analysis of these non-ReLU activations is to carry out the integrations similar to those in Line 652 and Line 661. Nevertheless, these integrals can be estimated numerically, as we did in Figure 2.

Q2: “What are the practical takeaways for model design or training procedures, how can your findings be used in certain settings?”

A: We consider our paper as a theoretical work that provides insights and understandings of deep learning. Two possible practical implications include: (1) we may need to rethink the current NTK-based optimization theories which often assume a fixed condition number, our paper suggests the activation function can influence this condition number hence the convergence rate. (2) Another implication of our better separation theorem could relate to the origin of adversarial examples where the adversarial and clean examples can be considered as two similar inputs. Our theory indicates that, due to the activation function, they are quite largely separated in the gradient feature space which may attribute to the observed distinct predictions between adversarial and clean examples.

Dear Reviewer ndzW:

We would like to reach out to see if your concerns have been resolved or whether you have further questions. We are happy to address any remaining concerns/questions. If no further concerns, we kindly ask the reviewer to consider updating their overall rating. Thank you very much!