On the Generalization Power of the Overﬁtted Three-Layer Neural Tangent Kernel Model
Abstract
In this paper, we study the generalization performance of overparameterized 3layer NTK models.
We show that, for a speciﬁc set of ground-truth functions (which we refer to as the “learnable set”), the test error of the overﬁtted 3-layer NTK is upper bounded by an expression that decreases with the number of neurons of the two hidden layers.
Different from 2-layer NTK where there exists only one hidden-layer, the 3-layer NTK involves interactions between two hiddenlayers.
Our upper bound reveals that, between the two hidden-layers, the test error descends faster with respect to the number of neurons in the second hidden-layer (the one closer to the output) than with respect to that in the ﬁrst hidden-layer (the one closer to the input).
We also show that the learnable set of 3-layer NTK without bias is no smaller than that of 2-layer NTK models with various choices of bias in the neurons.
However, in terms of the actual generalization performance, our results suggest that 3-layer NTK is much less sensitive to the choices of bias than 2-layer NTK, especially when the input dimension is large.
Introduction
Neural tangent kernel (NTK) models (Jacot et al., 2018) have been recently studied as an important intermediate step to understanding the exceptional generalization power of overparameterized deep neural networks (DNNs).
Deep neural networks (DNNs) usually have so many parameters that they can perfectly ﬁt all train data, yet they still have good generalization performance (Zhang et al., 2017;
Advani et al., 2020).
This seems contradicting to the classical wisdom of “bias-variancetradeoff” in the statistical machine learning methods (Bishop, 2006; Hastie et al., 2009; Stein, 1956;
James & Stein, 1992; LeCun et al., 1991; Tikhonov, 1943).
To understand this distinct behavior of DNNs, a recent line of work studies the so-called “double-descent” phenomenon, beginning with overﬁtted linear models.
These results on linear models suggest that the test error indeed decreases again in the overparameterized region, as the model complexity increases beyond the number of samples (Belkin et al., 2018, 2019; Bartlett et al., 2020; Hastie et al., 2019; Muthukumar et al., 2019;
Ju et al., 2020; Mei & Montanari, 2019).
However, these studies use linear models with simple features such as Gaussian or Fourier features, and hence they fail to capture the non-linearity in neural networks.
In contrast, NTK models adopt features generated by non-linear activation functions (i.e., neurons of DNNs), and thus they can be viewed as an intermediate step between simple linear models and DNNs.
Along this line, the work in Ju et al.
(2021) studies 2-layer NTK models, and shows that the 2-layer NTK model indeed exhibits better and different descent behavior in the overparameterized region, which might be closer to that of an actual neural network.
Motivated by Ju et al.
(2021), it is of great interest to understand whether similar insights extend to deeper NTK models.
In particular, in this paper we study NTK models with 3 layers.
Although both 2-layer and 3-layer NTK models share similar assumptions (e.g., trained weights do not change much from initialization, and features are linearized around the initial state), their difference in structure leads to completely different feature formation.
Compared with 2-layer NTK models that only contain one hidden-layer of neurons, 3-layer NTK models have two hidden-layers, which interact in more complex ways not observed in 2-layer NTK models.
Speciﬁcally, let p 1 and p 2 denote the number of neurons in the two hidden layers.
Then, the ultimate features of the 3-layer NTK models depend on both p 1 and p 2 .
This dependency leads to the following questions.
First, the width of which layer is more important in governing the descent behavior, p 1 or p 2 ?
Further, to get better descent behaviors, should p 1 and p 2 grow at the same speed, or should one of them grow faster than the other?
Second, do 3-layer NTK models have any performance advantage over 2-layer NTK models?
To answer these questions, in this paper we study the generalization performance of overﬁtted minℓ 2 -norm solutions for 3-layer NTK models where the middle layer is trained.
For a set of learnable functions (which we refer to as the “learnable set”), we provide an upper bound on the test error for ﬁnite values of p 1 and p 2 .
To the best of our knowledge, this upper bound is the ﬁrst result that can reveal the dependency of the descent behavior on p 1 and p 2 separately.
We then compare 3-layer NTK with 2-layer NTK with respect to the corresponding learnable set and the actual generalization performance.
Our comparison reveals several important differences between 3-layer NTK and 2layer NTK, in terms of the descent behavior, the size of the learnable set, and the sensitivity of the generalization performance to the choice of bias of the neurons.
Analyzing the Generalization Error:
First, we show that the generalization error (denoted by the absolute value of the difference between the model output and the ground-truth for a test input) is upper bounded by the sum of several terms on the order of O (1 / √ n ) ( n denotes the number of training data), O (1 /p 2 ) ( p 2 denotes the number of neurons in the second hidden-layer), O (p log p 1 /p 1 ) ( p 1 denotes the number of neurons in the ﬁrst hidden-layer), plus another term related to the magnitude of noise.
Similar to 2-layer NTK (Arora et al., 2019; Ju et al., 2021; Satpathi & Srikant, 2021), our upper bound suggests that when there are inﬁnitely many neurons, the generalization error decreases with the number of samples n at the speed of √ n and will approach zero when n → ∞ in the noiseless situation.
Further, the noise term will not explode when the number of neurons goes to inﬁnity, which is also similar to that for 2-layer NTK.
However, our upper bound also reveals new insights that are different from the results for 2-layer NTK.
Speciﬁcally, our upper bound decreases slower with respect to the number of neurons in the ﬁrst hidden-layer p 1 at the speed ofp (log p 1 ) /p 1 , and decreases faster with respect to the number of neurons in the second 1 hidden-layer p 2 at the speed of 1 / √ p 2 .
Further, our upper bounds hold regardless of how fast p 1 and p 2 increase relative to each other (e.g., they could increase at the same speed, or one could increase faster than the other).
Characterizing the Learnable Set: We then show that, even if we only train the middle-layer weights, the learnable set (i.e., the set of ground-truth functions for which the above upper bound holds) of the 3-layer NTK without bias contains all ﬁnite degree polynomials, which is strictly larger than that of the 2-layer NTK without bias and is at least as large as the 2-layer NTK with bias.
Recently, Geifman et al.
; Chen & Xu (2020) show that when all layers are trained, 3-layer
NTK leads to exactly the same reproducing kernel Hilbert space (RKHS) as 2-layer NTK with biased
ReLU (although they assumed an inﬁnite number of neurons, and did not characterize the descent behavior of the generalization error).
Combining with their results, we can draw the conclusion that training only the middle-layer weights is at least as effective as training all layers in 3-layer NTK, in terms of the size of the learnable set.
Sensitivity to the Choices of Bias: Even though a similar learnable set can be attained by 3-layer NTK (with or without bias) and 2-layer NTK (with bias), our results suggest that the actual generalization performance can still differ signiﬁcantly in terms of the sensitivity to the choice of bias, especially when the input dimension d is large.
One type of bias setting commonly used in literature (Ghorbani et al., 2021a; Satpathi & Srikant, 2021) is that the bias has a similar magnitude as each element of the input vector, which we refer to as “normal bias”.
However, we show that such a normal bias setting has a negative impact on the generalization error for overﬁtted 2-layer NTK when d is large.
To avoid this negative impact, it is important to use another type of bias setting where the bias has a similar magnitude as the norm of the whole input vector, which we refer to as
[Figure]
[Figure] “balanced bias”.
In contrast, for 3-layer NTK, different bias settings do not have an obvious effect on the generalization performance.
In summary, compared with 2-layer NTK, the use of an extra non-linear layer in 3-layer NTK appears to signiﬁcantly reduce the impact due to the choice of bias, and therefore makes the learning more robust.
Our work is related to the growing literature on the generalization performance of such shallow and fully-connected neural network.
However, most of these studies focus on 2-layer neural networks.
Among them, they differ in which layer to train.
For example, Mei & Montanari (2019); d’Ascoli et al. Mei et al.
(2022) consider the “random feature” (RF) model that only trains the top-layer weights and ﬁxes the bottom-layer weights, while 2-layer NTK trains the bottom-layer weights.
In contrast, our work on 3-layer NTK neither trains the bottom-layer or top-layer weights.
Instead, we train the middle-layer weights, since the middle-layer of a 3-layer model involves the interaction between two hidden-layers, which does not exist in 2-layer models.
The above studies of 2-layer network also differ in how the number of neurons/features p , the number of training samples n , and the input dimension d grow.
Mei & Montanari (2019); Mei et al.
(2022) study the generalization performance of the RF model where the number of neurons p , the number of training data n , and the input dimension grow proportionally to inﬁnity.
While Ghorbani et al.
(2021b) focuses on the approximation error (i.e., expressiveness) of both RF and NTK models, their analysis on generalization error is only on the limit n or p → ∞ .
All of these studies are quite different from ours with ﬁxed n and ﬁnite p .
Other works such as Arora et al.
; Satpathi & Srikant (2021); Fiat et al. (2019) study the situation where the number of training samples n is given and the number of neurons p is larger than a threshold, which is closer to our setup.
However, these studies usually do not quantify how the generalization performance depends on the number of neurons p .
Speciﬁcally, they usually provide an upper bound on the generalization error when the number of neurons p is greater than a threshold, while the upper bound itself does not depend on p .
Thus, such an upper bound cannot explain the descent behavior of NTK models.
The work in Ju et al.
(2021) does study the descent behavior with respect to p , and is therefore the closest to our work.
However, as we have explained earlier, there are crucial differences between 2 and 3 layers in both the descent behavior and the learnable set of ground-truth functions.
In addition to the above references, our work is also related to Allen-Zhu et al.
(2019) (which studies NTK without overﬁtting)
and Ji & Telgarsky (2019) (which studies classiﬁcation by NTK).
Their settings are however different from ours in that we consider overﬁtted solutions for regression.
In summary, our paper is the ﬁrst to provide a highprobability upper bound on the generalization error of the overﬁtted 3-layer NTK (where only its middle layer weights are trained), and to characterize how the generalization error decreases with the number of neurons p 1 and p 2 .
System Model
Let f : R d 7→ R denote the ground-truth function.
Let ( X i , f ( X i ) + ǫ i ) , i = 1 , 2 , · · · , n denote n pieces of training data, where X ∈ R n × d is the matrix, each column of which is the input of one training sample, ǫ
We use a [ j ] to denote the j -th part (sub-vector) of the vector a .
The part size depends on a .
Speciﬁcally, If a has p 1 elements, then each part has 1 elements.
If a has dp 1 elements, then each part has d elements.
If a has p 1 p 2 elements, then each part has p 1 elements.
We consider a fully-connected 3-layer neural network as illustrated in Fig.
1, which consists of normalized d -dimensional input x ∈ S d
− 1 (a unit hyper-sphere), p 1 ReLUs (rectiﬁer linear units max( · , 0) ) at the ﬁrst hidden-layer, p 2 ReLUs at the second hidden-layer, bottom-layer weights (between input and 1 st hidden-layer) V ∈ R ( p 1 d ) × 1 , middle-layer weights (between 1 st hidden-layer and 2 nd hidden-layer)
W 0 ∈ R ( p 1 p 2 ) × 1 , and top-layer weights w ∈ R p 2 × 1 (between 2 nd hiddenlayer and output).
2.1 Overﬁtted NTK solution
In this subsection, we will derive the overﬁtted solution for this 3-layer neural network using the
NTK approximation.
Let h RF V , x ∈ R p 1 × 1 denote the output of the ﬁrst hidden-layer.
We then have
[Equation]
(We use the superscript “RF” because h RF V , x is indeed the feature vector of a random feature model ( Mei & Montanari, 2019).)
After training the middle-layer weights, W 0 changes to W 1 : = W 0 + ∆ W .
Then, the change of the output is
[Equation]
The NTK model (Jacot et al., 2018) assumes that ∆ W is very small and thus the activation pattern does not change much.
In other words, we can approximate 1 { W 1 [ k ] T h RF V , x > 0 } by 1 { W 0 [ k ] T h RF V , x > 0 } .
Deﬁne ∆ W ∈ R ( p 1 p 2 ) × 1 as ∆ W [ k ] := w k · ∆ W [ k ] , k = 1 , 2 , · · · , p 2 .
Deﬁne h
Three V , W 0 , x ∈
R 1 × ( p 1 p 2 ) such that
[Equation]
where k
= 1 , 2 , · · · , p 2 .
Therefore, the change of the output can be approximated by
[Equation]
We thus obtain a linear model in ∆ W .
We provide an illustration of the formation and structure of these vectors in Fig.
4, Appendix A.1 in Supplementary Material.
Deﬁne the design matrix H ∈ R n × ( p 1 p 2 ) such that its i -th row is H i =
h Three V , W 0 , X i .
Notice that overﬁtted gradient descent on a linear model converges to the min ℓ 2 -norm solution 2 , which is denoted by
[Equation]
When H is full row-rank (which holds with high probability under certain conditions), the trained model is then
[Equation]
Notice that the trained model is determined by multiple random variables.
In order to analyze the generalization performance of the trained model, we have to make assumptions on the distribution of those random variables.
Let µ ( · ) , λ ( · ) , and γ ( · ) denote the probability density function of x , V [ j ] , and W
[ k ] , respectively.
For simplicity, we make the following assumption that all random variables follow uniform distribution
Assumption 1.
The input x and the bottom-layer initial weights V [ j ] ’s ( j = 1 , 2 , · · · , p 1 ) are i.i.d.
and uniformly distributed in S d − 1 .
In other words, µ ( · ) and λ ( · ) are both unif ( S d − 1 ) .
The middlelayer initial weights W
[ k ] ’s ( k = 1 , 2 , · · · , p 2 ) are i.i.d.
and uniformly distributed in S p 1 − 1 .
In other words, γ ( · ) is unif ( S p 1 − 1 ) .
The top-layer weights w are all non-zero 3 .
Remark 1 .
Readers may be curious why we only train the middle-layer weights.
Part of the reason is technicality: if the bottom layer is also trained, the aggregate output of the ﬁrst hidden-layer may have changed so much that the second hidden-layer’s inputs and ReLU activation patterns change signiﬁcantly from initialization, which may violate the NTK assumption.
The work in Geifman et al. ; Chen & Xu (2020) is not concerned about this difﬁculty, since they are mostly interested in the expressive power of the RKHS, assuming an inﬁnite number of neurons.
In contrast, we wish to capture the effect of ﬁnite width, and thus train only the middle layer to avoid this difﬁculty.
More importantly, the middle-layer weights interact with both the ﬁrst hidden layer and the second hiddenlayer, and are the major structural distinction compared with 2-layer NTK.
This setting thus helps us to answer the following interesting question: will training the middle layer alone already achieve the same (potential) beneﬁt as training all layers (especially given that the latter encounters more technical difﬁculty)?
Generalization Performance
In this section, we will show our main results about the generalization performance of the aforementioned 3-layer NTK model for a speciﬁc set of functions.
We ﬁrst introduce a set of ground-truth functions that may be learnable and then provide a high-probability upper bound on the test error.
We then discuss some useful implications of our upper bound.
A set of ground-truth functions that may be learnable
We deﬁne kernel functions K RF , K Two , and K Three : [ − 1 , 1] 7→ R as follows (whose meanings will be explained soon):
[Equation]
(Notice that 2 d · K RF ( a ) ∈ [0 , 1] for all a ∈ [ − 1 , 1] by Lemma 43 in Supplementary Material,
Appendix I, and hence K Three ( · ) is well deﬁned.)
We deﬁne a set F ℓ 2 (3) of ground-truth functions based on those kernels:
Deﬁnition 1 (learnable set of 3-layer NTK) .
[Equation]
where k g k ∞ : = sup z ∈S d − 1 | g ( z ) | .
To see why functions in F ℓ 2 (3) may be learnable, we can check what the learned result ˆ f ℓ 2 in Eq.
(3) should look like.
When there are inﬁnite number of neurons and there is no noise (i.e., ǫ = 0 ), what remains on the right-hand-side of Eq.
(3) can be viewed as the product of two terms, h Three V , W 0 , x H T and ( HH T ) − 1 F ( X ) .
For the ﬁrst term h Three V , W 0 , x H T , note that each row of H is given by h Three V , W 0 , X i for i = 1 , 2 , · · · , n .
Thus, when p 1 , p 2 → ∞ , the i -th element of h Three V , W 0 , x H T , which is the inner product between h Three V , W 0 , x and h Three V , W 0 , X i , converges in probability to K Three ( h Three V , W 0 , x ( h Three V , W 0 , X i ) T ) , which is exactly the kernel function of 3-layer NTK.
By representing the second term ( HH T ) − 1 F ( X ) with a certain g ( · ) , ˆ f ℓ 2 must then approach the form ∆ W [ k ] := w k · ∆ W [ k ] , k = 1 , 2 , · · · , p 2 . in Eq.
(See Supplementary Material, Appendix B for details.)
Intuitively, K Three can be thought of as the composition of the kernels of each of the two layers, which are K Two and K RF given in Eq.
(5) and Eq.
Speciﬁcally, suppose that we ﬁx the output of the ﬁrst hidden layer (i.e., h RF V , x ) and regard it as the input of a 2-layer NTK formed by the top two layers of the 3-layer neural network.
By letting p 2 → ∞ , we can show that the inner product between h Three V , W 0 , x and h Three V , W 0 , X i approaches K Two (( h RF V , x ) T h RF V , X i ) (with necessary normalization of h RF V , x and h RF V , X i ), where K Two is exactly the kernel of 2-layer NTK in Ju et al.
Second, when p 1 → ∞ , we can show that ( h RF V , x ) T h RF V , X i approaches K RF ( x T X i ) , where K RF is exactly the kernel of the random-feature model (Mei & Montanari, 2019).
In summary, we expect that functions in F ℓ 2 (3) can be approximated by ˆ f ℓ 2 ( · ) .
However, we note that the above deviation is only about the expressiveness of 3-layer NTK and it does not precisely reveal its generalization performance.
3.2 An upper bound on the generalization error
We now present the ﬁrst main result of this paper, which is an upper bound that quantiﬁes the relationship between the generalization performance and system parameters.
Theorem 1.
For any ground-truth function f ( x ) = f g ( x ) ∈ F ℓ 2 (3) , when d is ﬁxed and p 1 , p 2 are much larger than n , (with high probability) we have
[Equation]
A more precise version of the upper bound and the condition of Theorem 1 as well as its derivation can be found in Supplementary Material, Appendix C.
As we can see, Eq.
(8) captures how the test error depends on ﬁnite values of parameters n , p 1 , p 2 , k ǫ k 2 , and g .
Later in this section we will examine more closely how n , p 1 , and p 2 affect the value of the upper bound.
Regarding the dependency on g , Eq.
(8) works as long as k g [-(cid:107) 
-]k 1 and k g k ∞ are ﬁnite 4 .
Intuitively, the norm of g represents the complexity of the ground-truth function in F ℓ 2 (3) .
When the norm of g is larger, then the right-hand side of Eq.
(8) becomes larger, which indicates that such ground-truth function is harder to learn.
A simple example is that if we enlarge a ground truth function f g ∈ F ℓ 2 (3) by 2 times (which means g is 2 times larger), then since the model is linear, the test error | ˆ f ℓ 2 ( x ) − f ( x ) | will become 2 times larger.
We will discuss more about which types of functions satisfy the condition of ﬁnite norm of g in Section 4.
Next, we will discuss some implications of this upper bound of 3-layer NTK.
While some of them are similar to 2-layer NTK, others are signiﬁcantly different, revealing the complexity due to having more layers.
3.3 Interpretations similar to 2-layer NTK
Based on the upper bound in Theorem 3, we have the following insights for 3-layer NTK, which are similar to those for 2-layer NTK shown in Ju et al.
These similarities may reveal some intrinsic properties of the NTK models regardless of the number of layers.
Zero test error with n → ∞ in the ideal situation: In the ideal situation where there are inﬁnitely many neurons and no noise, the only remaining term in Eq.
(8) is Term A. Notice that Term A decreases to zero as n → ∞ , which indicates that the generalization error decreases to zero when [Figure] 2 p 1 = p [Figure]
[Figure] E [Figure] σ 2 =0.01 σ 2 = 1010 3 p10 − 10 − 10 − (b) fix p 2 =∞
σ 2 =10 − 1010 3 p10 − 10 − 10 − (c) fix p 2 = 1010 3 p 2 10 − 10 −
[Figure]
−
[Figure]
− (d) fix p
[Figure]
more training data are provided in the ideal situation.
Term A suggests that such decreasing speed is at least 1 / √ n .
Such result is consistent with that of 2-layer NTK, e.g., in Arora et al.
3.4 Insights that are distinct compared with 2-layer NTK
Compared with 2-layer NTK, an important difference for 3-layer NTK is that there are more than one hidden-layers.
Therefore, the speed of the descent of 3-layer NTK involves the interaction between two hidden-layers.
Descent with respect to the number of neurons: In Eq.
(8), Term B and Term C contain p 1 and p 2 , respectively.
For any given n and noise level k ǫ k 2 , Terms A and D do not change, and Term E decreases with p 1 and p 2 .
(More discussion about Term E can be found in Supplementary Material,
Appendix D, where we discuss the noise effect.)
Therefore, by increasing p 1 and p 2 , Term B and
Term C keep decreasing.
In summary, right-hand-side of Eq.
(8) decreases as the number of neurons p 1 and p 2 increases, which validates the descent in the overparameterized region of 3-layer NTK.
Different descent speed: As shown in Eq.
(8), p 1 and p 2 play different roles in the descent of the generalization error.
Comparing Term B and Term C of Eq.
(8), we can see that the upper bound of the test error | ˆ f ℓ 2 ( x ) − f ( x ) | decreases faster with respect to p 2 (at the speed of √ p 2 ) and slower with respect to p 1 (at the speed ofp p 1 / log p 1 ).
We emphasize that this difference is not due to the number of weights/parameters contributed by the number of neurons in each hidden-layer of p 1 and p 2 5 .
Instead, we conjecture that such difference in the speed of descent may be due to the different positions in this 3-layer neural network structure, where the second hidden-layer takes the trained middle-layer weights as its input (and thus utilizes the trained weights better than the ﬁrst hidden-layer).
We use numerical results to illustrate the different roles of p 1 and p 2 in reducing the generalization error.
We ﬁx p 2 = 200 and plot the MSE with respect to p 1 in Fig. 2(b).
Although the test error decreases when p 1 increases, the decreasing speed is slow, especially for the noisy situation.
Such a slow decreasing speed with p 1 remains even when p 2 is ﬁxed to a much higher value.
For example, in Fig. 2(b), we ﬁx p 2 = ∞ , we still observe the similarly slow decreasing speed with p 1 as shown by Fig. 2(c).
In contrast, the descent with respect to p 2 should be easier to observe and can reach a lower test MSE.
In Fig. 2(d), we ﬁx p 1 = 200 and increase p 2 (i.e., we exchange the values of pand p 2 in Fig. 2(c)(d)).
As we can see, all three curves in Fig. 2(d) have a more obvious descent and decrease to lower MSE compared with those in Fig. 2(c), which validates our conjecture that the descent speed with respect to the number of neurons of the second hidden-layer is faster.
Notice that our upper bound Eq.
(8) also suggests a descent when both p 1 and p 2 increase simultaneously.
We use simulation result by Fig. 2(a) to support this point.
We ﬁxed n = 200 and let p 1 = pincrease simultaneously.
The ground-truth model in this ﬁgure is f ( x ) =
( x T e )
+ ( x T e ) 3 where d = 3 .
The green, orange, and blue curves denote the situations of σ 2 =
(no noise), σ 2 = 0 .
01 , and σ 2 = 0 .
04 , respectively.
Every point in this ﬁgure is the median of 20 simulation runs.
We also provide the box plot 6 of the situation of σ 2 = 0 .
01 (correspond to the orange curve).
It is obvious that all three curves descend, which veriﬁes that the generalization error of the overﬁtted 3-layer NTK model decreases when p 1 and p 2 increase simultaneously at the same speed.
By observing the box plot for the situation σ 2 = 0 .
01 (the orange curve), we also notice that when p 1 = p 2 becomes large, the variance becomes small.
This is because all initial weights are i.i.d.
random and a large number of weights may reduce the variance of the model due to the law of large numbers.
Our upper bound in Theorem 3 also suggests such reduced variance as the probability in Theorem 3 increases as p 1 increases.
Types of Ground-Truth Functions
Are 3-layer (i.e., deeper) networks better than 2-layer networks in any way for generalization performance?
In the last section, we have seen that both 3-layer NTK and 2-layer NTK can achieve zero test error when n → ∞ in the ideal noiseless situation, when the ground-truth functions are in their respective learnable set 7 .
A natural question is then to compare the learnable sets between these two models, and to compare the generalization performance when the ground-truth function belongs to both learnable sets.
In this section, we provide some answers by studying various types of ground-truth functions and their effects on the generalization performance.
4.1 Size of the learnable set
For a 2-layer NTK, as shown in Ju et al.
(2021), when no bias is used in ReLU, the corresponding learnable set F ℓ 2 (2) contains all even polynomials and linear functions, but does not contain other odd polynomials.
In order to learn both even and odd polynomials, it is critical that bias is added to ReLU (Satpathi & Srikant, 2021; Ju et al., 2021).
In contrast, we prove the following result:
Proposition 2.
F ℓ 2 (3) (with unbiased ReLU, middle layer being trained) already contains all polynomials with ﬁnite degree (i.e., including both even and odd polynomials).
Further, the learnable set
F ℓ 2 (3) of 3-layer NTK is strictly larger than that of the 2-layer NTK with unbiased ReLU, and is at least as large as that of the 2-layer NTK with biased ReLU.
This independence to bias shown by Proposition 2 can be seen as one performance advantage of 3layer NTK compared to 2-layer NTK.
Details (including more precise statement) about this result is in Supplementary Material, Appendix J. Notice that Geifman et al.
; Chen & Xu (2020) show that when training all layers, 3-layer NTK leads to the same RKHS as 2-layer NTK with biased
ReLU.
However, it is unclear whether training one layer is already sufﬁcient for achieving the same
RKHS as training all layers.
Our result in Proposition 2 answers this question positively, i.e., only training the middle layer has already achieved all beneﬁts of training all layers in terms of the size of the learnable set.
(In other words, training all three layers will not expand the learnable set over training only the middle layer.)
4.2 Different bias settings with high input dimension
Even when a ground-truth function belongs to both F ℓ 2 (2) ,b and F ℓ 2 (3) , their generalization performance may still exhibit some differences.
In this subsection, we will show that when the input dimension d is high, some speciﬁc choice of bias of the 2-layer NTK has better generalization performance than others.
In contrast, the 3-layer NTK is less sensitive to different bias settings.
Notice that adding bias to each ReLU in 2-layer NTK is equivalent to appending a constant to x while still using ReLU without bias.
Speciﬁcally, the input vector for biased 2-layer NTK is
[Equation]
[Figure]
[Figure]
(cid:0) x T e(cid:1)
[Figure]
−
Every curve is the average of 10 random simulations.
where b ∈ (0 , 1) denotes the initial bias.
We also normalize the ﬁrst d elements of x b by √ 1 − bin Eq.
(9) to make sure that k x b k 2 = 1 .
Under this biased setting, the 2-layer NTK model has the learnable set F ℓ 2 (2) ,b := {
[Equation]
A common setup for the initial magnitude of the bias of each ReLU is to use a value that is close or equal to the average magnitude of each element of input x , e.g., Satpathi & Srikant (2021); Ghorbani et al. (2021a).
Speciﬁcally, we let b = 1 √ d
+
in Eq.
(9), and denote the corresponding learnable set by F ℓ 2 (2) , NLB .
We refer to this setting as the “normal-bias” setting.
Alternatively, the initial magnitude of the bias can be chosen to be close or equal to k x k 2 .
Speciﬁcally, we let b = 1 √in Eq.
(9) and denote the corresponding learnable set by F ℓ 2 (2) , BB .
We refer to this second setting as “balanced-bias”.
The speciﬁc expression of F ℓ 2 (2) , F ℓ 2 (2) , NLB , and F ℓ 2 (2) , BB can be derived by using similar methods shown in Section 3.1 (results are listed in Table 1).
We now discuss how the two different bias settings could affect the generalization performance when d is large.
For 2-layer NTK under the normal-bias setting, the kernel is K Two ( dd +1 x T z + 1 d +1 ) .
Although it contains both even and odd power polynomials, we notice that when d increases, K Two approaches its no-bias counterpart K Two ( x T z ) , which only contains even power polynomials and linear term.
Thus, we conjecture that, by increasing d , the generalization performance of 2layer NTK with normal-bias will deteriorate for those ground-truth functions inside F ℓ 2 (2) , NLB but far away from F ℓ 2 (2) (e.g., odd-degree non-linear polynomials).
In contrast, for 2-layer NTK under the balanced-bias setting, the kernel is K Two ( 12 x T z + 12 ) , which does not change with d .
Therefore, we expect that such deterioration should not happen.
Note that in 3-layer NTK, although normalbias setting still approaches no-bias setting when d increases, there does not exist such performance deterioration, because F ℓ 2 (3) (the learnable set of 3-layer NTK without bias) already contains both even and odd power polynomials.
These insights will be veriﬁed by the numerical results below.
We now use simulation results in Fig.
3 to validate the conjecture that 3-layer NTK models are less sensitive to different bias settings than 2-layer NTK models.
We let the ground-truth function be f ( x ) =
d + (cid:0) x T e(cid:1) − x T e 1 , which is orthogonal to F ℓ 2 (2) .
In Fig. 3(a) when d = 2 , all settings have similar performance except 2-layer NTK without bias, whose test error is always above the null risk.
In Fig. 3(b) when d = 15 , the purple curve of 2-layer NTK with normal bias gets closer to the red curve of 2-layer NTK without bias (and thus the generalization performance becomes worse), while other curves are still close to each other.
This validates our conjecture that 3-layer NTK models are less sensitive to different bias settings than 2-layer NTK models.
Further simulations can be found in Appendix A.2.
Conclusion
In this paper, we studied the generalization performance of overﬁtted 3-layer NTK models.
Compared with 2-layer NTK models, 3-layer NTK is less sensitive to different bias settings.
Further, training only the middle layer can get most of the performance advantage of 3-layer NTK, in terms of the learnable set.
Possible future directions include: (i) studying whether training other layers will get the same beneﬁt as training the middle layer; (ii) approximating the actual neural network where the learned result is far away from the initial state; (iii) investigating deeper network as well as other structures such as convolutional neural network (CNN) and recursive neural network (RNN).
Acknowledgments and Disclosure of Funding
This work has been supported in part by NSF grants: 2112471 (also partly funded by DHS), CNS2106933, CNS-1901057, CNS-2113893, and the Bilsland Dissertation Fellowship at Purdue University, and a grant from the Army Research Ofﬁce: W911NF-21-1-0244.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] E.g., see Remark 1.
(c) Did you discuss any potential negative societal impacts of your work?
[No] This work is on the learning theory and should have no potential negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] As this is mainly a theoretical paper.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Fig.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] As this is mainly a theoretical paper.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]