High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
Abstract
We study the ﬁrst gradient descent step on the ﬁrst-layer parameters W in a twolayer neural network: f ( x ) = 1 p N a >  ( W > x ) , where W 2 R d ⇥ N , a 2 R N are randomly initialized, and the training objective is the empirical MSE loss:n
P ni =
( f ( x i )  y i ) 2 .
In the proportional asymptotic limit where n, d, N ! at the same rate, and an idealized student-teacher setting where the teacher f ⇤ is a single-index model, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on W with learning rate ⌘ .
We consider two scalings of the ﬁrst step learning rate ⌘ .
For small ⌘ , we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random feature model, but cannot defeat the best linear model on the input.
Whereas for sufﬁciently large ⌘ , we prove that for certain f ⇤ , the same ridge estimator on trained features can go beyond this “linear regime” and outperform a wide range of (ﬁxed) kernels.
Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.
Introduction
We consider the training of a fully-connected two-layer neural network (NN) with N neurons, f NN ( x ) =
p N
N
X i = a i  ( h x , w i i )
[Equation]
where x 2 R d , W 2 R d ⇥ N , a 2 R N ,  is the nonlinear activation function applied entry-wise, and the training objective is to minimize the empirical risk.
Our analysis will be made in the proportional asymptotic limit , i.e., the number of training data n , the input dimensionality d , and the number of neurons N jointly tend to inﬁnity.
Intuitively, this regime reﬂects the setting where the network width and data size are comparable, which is consistent with practical choices of model scaling.
When the ﬁrst layer W is ﬁxed and the second layer a is optimized, we arrive at a kernel model, where the kernel deﬁned by features x 7!
( W > x ) (often called the hidden representation ) is referred to as the conjugate kernel (CK) [Nea95].
When W is randomly initialized, this model is an example of the random features (RF) model [RR08], the training and test performance of which has been extensively studied in the proportional limit [LLC18, MM22].
These precise characterizations reveal interesting phenomena also present in practical deep learning [BHMM19].
However, RF models do not fully explain the empirical success of NNs: one crucial advantage of deep learning is the ability to learn useful features [GDDM14, DCLT18] that “adapt” to the learning problem [Suz18].
In fact, recent works have shown that such adaptivity enables NNs optimized by gradient descent to outperform a wide range of linear/kernel estimators [AZL19, GMMM19].
While many explanations of this separation have been proposed, our starting point is the empirical ﬁnding that “non-kernel” behavior often occurs in the early phase of NN optimization, especially under large learning rates [JSF + 20, FDP + 20].
The goal of this work is to answer the following question:
Can we precisely capture the emergence of feature learning in the early phase of gradient descent, and demonstrate its improvement over the initial (ﬁxed) kernel in the proportional limit?
1.1 Contributions
Motivated by the above observations, we investigate a simpliﬁed scenario of the “early phase” of learning: how the ﬁrst gradient step on the ﬁrst-layer parameters W impacts the representation of the two-layer NN (1.1).
Speciﬁcally, we consider regression with the squared loss (MSE), and a studentteacher setting in the proportional asymptotic limit; we aim to characterize the prediction risk of the kernel ridge regression estimator on top of the ﬁrst-layer CK feature x 7!
( W > x ) , before and after one gradient descent step on the empirical risk (starting from Gaussian initialization).
Following prior works on the precise asymptotics of RF regression [GLK + 20, DL20], we focus on the setting where the input x is Gaussian and the teacher f
⇤ is a single-index model.
In this case, the prediction risk of a large class of
RF/kernel ridge regression estimators is lower-bounded by the L 2 -norm of the “nonlinear” component of the teacher k P >
f ⇤ k 2 L 2 , i.e., they only learn linear functions on the input.
After one gradient step on W , we compute the CK ridge estimator using separate training data, and compare its prediction risk against this linear lower bound.
Our analysis will be made under two choices of learning rate scalings
: • Small lr: ⌘ = ⇥ (1) .
In Section 4, we extend the Gaussian Equivalence Theorem (GET) in [HL20] to the updated feature map after one gradient descent step on W with learning rate ⌘ = ⇥ (1) ; this allows us to precisely characterize the prediction risk using random matrix theoretical tools.
We prove that after one gradient step, the ridge regression estimator on the learned CK features already exhibits nontrivial improvement over the initial RF ridge regression model (see pink curve in Figure 1), but it remains in the “linear regime” and cannot outperform the best linear estimator on the input (black dashed line). • Large lr: ⌘ = ⇥ ( p N ) .
In Section 5, we analyze a larger learning rate that coincides with the maximal update parameterization in [YH20].
For certain target functions f ⇤ , we prove that kernel ridge regression after one feature learning step can achieve lower risk than the lower bound k P >
f ⇤ k 2 L 2 ; thus, it outperforms a wide range of kernel estimators (see purple curve in Figure 1).
1.2 Related works
Asymptotics of kernel regression.
Recent works provided precise analysis of RF and kernel models in the proportional limit [GLK + 20, DL20, LCM20, AP20, MM22].
These results typically build upon analyses of the spectrum of kernel matrices, a key ingredient in which is the “linearization” of nonlinear random matrices via Taylor expansion [EK10] or orthogonal polynomials [CS13, PW17].
Consequently, a large class of kernel models are essentially linear in the proportional asymptotic limit [LR20, BMR21].
In the case of RF models, a similar property is captured by the Gaussian
Equivalence Theorem [GMKZ20, HL20, GLR + 21], which roughly states that RF estimators achieve the same prediction risk as a (noisy) linear model.
For inputs with unit norm, [GMMM21, MMM21] showed that sample size n = ⌦ ( d 2 ) is required to go beyond this “linear” regime.
As we will see in certain settings, such a limitation can also be overcome (in the n ⇣ d scaling) by training the feature map for one gradient step with a sufﬁciently large learning rate.
Advantage of NNs over ﬁxed kernels.
It is well-known that under a speciﬁc initialization, the learning dynamics of overparameterized NNs can be described by the neural tangent kernel (NTK) [JGH18].
However, the NTK description essentially “freezes” the model around its initialization [COB19], and thus does not explain the presence of feature learning in NNs [YH20].
In fact, various works have shown that deep learning is more powerful than kernel methods in terms of approximation and estimation ability [Bac17, Suz18, IF19, SH20, GMMM20].
Moreover, in some specialized settings, NNs optimized with gradient-based methods can outperform the
NTK (or more generally any kernel estimators) in terms of generalization error [AZL19, WLLM19,
GMMM19, LMZ20, DM20, SA20, AZL20, RGKZ21, KWLS21, ABAB + 21] (see [MKAS21, Table 2] for a survey).
These results often require a careful analysis of the landscape (e.g., properties of global optimum) or optimization dynamics; in contrast, our goal is to precisely characterize the ﬁrst gradient step and demonstrate a similar separation.
Early phase of NN optimization.
Recent empirical studies suggest that properties of the ﬁnal trained model is strongly inﬂuenced by the early stages of optimization [GAS19, LM20, PPVF21], and the NTK evolves most rapidly in the ﬁrst few epochs [FDP + 20].
Large learning rate in the initial steps can impact the conditioning of loss surface [JSF + 20, CKL + 21] and potentially improve the generalization performance [LWM19, LBD + 20].
Under structural assumptions on the data, it has been proved that one gradient step with sufﬁciently large learning rate can drastically decrease the training loss [CLB21], extract task-relevant features [DM20, FCB22], or escape the trivial stationary point at initialization [HCG21].
While these works also highlight the beneﬁt of one feature learning step, to our knowledge this advantage has not been precisely characterized in the proportional regime (where the performance of RF models has been extensively studied).
Problem setup and assumptions
Notations.
Throughout this paper, k · k denotes the ` 2 -norm for vectors and the ` 2 !
` 2 operator norm for matrices, and k · k F is the Frobenius norm.
For matrix M 2 R n ⇥ n , tr( M ) = 1 n Tr( M ) is the normalized trace.
O d ( · ) and o d ( · ) stand for the standard big-O and little-o notations, where the subscript highlights the asymptotic variable; we write ˜ O ( · ) when the (poly-)logarithmic factors are ignored.
O d, P ( · ) (resp.
o d, P ( · ) ) represents big-O (resp.
little-o) in probability as d ! 1 .
⌦ ( · ) , ⇥ ( · ) are deﬁned analogously.
is the standard Gaussian distribution in R d .
Given f : R d !
R , we denote its L p -norm w.r.t.
as k f k L p ( R d ,  ) , which we abbreviate as k f k L p when the context is clear.
2.1 Training procedure
Gradient descent on the 1st layer.
Given training examples { ( x i , y i ) } ni =1 , we learn the two-layer
NN (1.1) by minimizing the empirical risk: L ( f ) = 1 n
P ni
=
` ( f ( x i ) , y i ) , where ` is the squared loss ` ( x, y ) =
12 ( x  y ) 2 .
As previously remarked, ﬁxing the ﬁrst layer W at random initialization and learning the second layer a yields an RF model, which is a convex problem with closed-form solution.
In contrast, we are interested in learning the feature map (representation) ; hence we ﬁrst ﬁx a (at initialization) and perform gradient descent on W .
We write the initialized ﬁrst-layer as
W 0 , and the weights after one gradient step as W 1 .
The gradient update, which we refer to as the feature learning step , with learning rate ⌘ is given as: W 1 = W 0 + ⌘ p N · G 0 where
[Equation]
in which  is the Hadamard product,  0 is the derivative of  (acting entry-wise), and we denoted the input feature matrix X 2 R n ⇥ d , and the corresponding label vector y 2 R n .
We remark that the p N -scaling in front of ⌘ accounts for the 1 p N -prefactor in our deﬁnition of two-layer NN (1.1).
Ridge regression for the 2nd layer.
After obtaining the updated weights W 1 , we evaluate the quality of the new CK features by computing the prediction risk of the kernel ridge regression estimator on top of the ﬁrst-layer representation.
Note that if ridge regression is performed on the same data X , then after one feature learning step, W 1 is no longer independent of X , which signiﬁcantly complicates the analysis.
To circumvent this difﬁculty, we estimate the regression coefﬁcients ˆ a using a new set of training data { ˜ x i , ˜ y
i } ni
=1 , which for simplicity we assume to have the same size as the original dataset.
This can be interpreted as the representation being “pretrained” on separate data before the ridge regression estimator is learned.
Denoting the feature matrix on the fresh training set { ˜ X , ˜ y } as  := 1 p N  ( ˜ XW 1 ) 2 R n ⇥ N , the
CK ridge regression estimator can be obtained by solving ˆ a = argmin
a [Equation]
2.2 Student-teacher setting and main assumptions
Given a target function (teacher model) f ⇤ and a learned model ˆ f , we evaluate the model performance using the prediction risk: R ( ˆ f ) =
E x ( ˆ f ( x )  f ⇤ ( x ))
=
k
ˆ f  f ⇤ k 2 L 2 , where the expectation is taken over the test data from the same training distribution.
We utilize the orthogonal decomposition of the activation function  .
Deﬁne the coefﬁcients µ 0 =
E
[  ( z )] , [Equation]
[Equation] [Equation]
where z ⇠ N (0 , 1) .
This implies  ( z ) = µ 0 + µ 1 z
+  ?
( z ) , where E [  ?
( z )]
=
E [ z  ?
( z )]
= 0 , and E [  ?
( z ) 2 ] = µ 22 .
Similarly, for square integrable target function f
⇤ , we have the orthogonal decomposition
[Equation]
where P > 1 is the projector orthogonal to constant and linear functions in L 2 ( R d ,  ) , which implies that E [ P > 1 f ⇤ ( x )]
= 0 , E [ x P > 1 f ⇤ ( x )] = 0 .
As d ! 1 , quantities deﬁned in (2.3) satisfy k  ⇤ k
= 1 , k P > [-
-]1 f ⇤ k L 2 !
µ ⇤ 2 , where µ ⇤ 0 , µ ⇤ 1 , µ ⇤ 2 are bounded constants.
Intuitively, µ ⇤ 0 , µ ⇤ 1 , and µ ⇤ 2 can be interpreted as the “magnitude” of the constant, linear, and nonlinear components of f ⇤ , respectively.
Assumption 1.
Proportional limit.
n, d, N ! 1 , n/d ! 1 , N/d ! 2 , where 1 , 2 2 (0 , 1 ) .
Gaussian initialization.
p d · [ W 0 ] ij i .
i .
d .
⇠ N (0 , 1) , p N · [ a ] j i .
i .
d .
⇠ N (0 , 1) , for i 2 [ d ] , j 2 [ N ] .
Normalized activation
The activation function  has   bounded ﬁrst three derivatives almost
surely In addition,  satisﬁes µ 0 = and µ 1 , µ 2 6 = 0 deﬁned in (2.2) .
Single-index teacher.
Labels are generated as y i =
f ⇤ ( x i ) + " i , where x
i i i
d .
⇠ N (0 , I ) , and
i is i.i.d.
sub-Gaussian noise with mean and variance
The teacher f ⇤ ( x ) =  ⇤
( h x ,  ⇤ i ) ,
where  ⇤ 2 R d with k  ⇤ k =
and  ⇤ is Lipschitz with[-µ ⇤ 
-]
µ ⇤
0 = µ ⇤ 1 6 = 0 as deﬁned in
Remark.
We make the following comments on the above assumptions.
• Following [HL20], we assume smooth centered activation to simplify the computation; empirical evidence suggests that similar result holds beyond this condition (e
g. [LGC + 21]).
We also expect the Gaussian input assumption may be replaced by weaker orthogonality conditions as in [FW20].
• The single-index setting has been extensively studied in the proportional regime [GLK + 20, DL20, HL20].
However, prior works only considered training the coefﬁcients a on top of ﬁxed feature map and such RF models cannot efﬁciently learn a single-index f ⇤ in high dimensions [YS19].
Under Assumption 1, a relatively large sample size corresponds to larger 1 , and a relatively large network width corresponds to larger 2 .
The proportional scaling of n, d, N implies that the model width is not signiﬁcantly larger than the training set size, in contrast to the polynomial overparameterization often required in NTK analyses [DZPS19], which may be less realistic in practical settings.
Importantly, the initialization of our two-layer NN (1.1) resembles the mean-ﬁeld parameterization [MMN18, CB18]: the second layer is divided by an additional p N -factor compared to the kernel (NTK) scaling — this ensures
that f NN ( x ) =
o d, P (1) at initialization and enables feature learning (see [YH20, Corollary 3.10]).
As an illustrative example in Figure 2, we plot the gradient descent trajectory of the ﬁrst-layer parameters W in two coordinates.
Observe that under the meanﬁeld parameterization (main ﬁgure), the neurons travel away from the initialization and align with the target function (black dashed lines), whereas in the NTK parameterization (subﬁgure, which omits the 1 p N -prefactor), the parameters remain close to their initialization and hence do not learn useful features.
Preliminary results 3.1 Lower bound for kernel ridge regression
To illustrate the beneﬁts of feature learning , we compare the prediction risk of ridge regression on the trained CK (after one gradient step) against that on the initial RF and ﬁxed kernels.
Speciﬁcally, given training data { x i , y i } ni =1 , we consider the following classes of kernel models for comparison.
• Random features model.
We introduce two RF kernels associated with (1.1) at initialization: the conjugate kernel (CK) deﬁned by features  CK ( x ) = 1 p N  ( W > 0 x ) 2 R N , and the neural tangent kernel (NTK)
[JGH18] deﬁned by features  NTK ( x
[Equation]
Given
a feature map RF 2 { CK,NTK } , the RF ridge regression estimator can be written as
[Equation]
• Rotation invariant kernel model.
Consider the inner-product kernel: k ( x , y )
=
g ⇣ h x , y i d ⌘ , and the Euclidean distance kernel: k ( x , y )=
g
[Equation]
⌘ , where g satisﬁes the smoothness conditions in [EK10].
Denoting the associated RKHS with H
[Equation]
the kernel ridge estimator is given by
[Equation]
We write the prediction risk of the above kernel estimators as R CK (  ) , R NTK (  [Equation] , R ker (  ) , respectively.
The following lower bound on the prediction risk is a simple combination of existing results.
Proposition 1 ([HL20, MZ20, BMR21]) .
Under Assumption 1, we have
[Equation]
where P > 1 denotes the projector orthogonal to constant and linear functions in L 2 ( R d ,  ) .
This proposition implies that in the proportional limit, ridge regression on the RF or rotationally invariant kernels deﬁned above does not outperform the best linear estimator on the input – it cannot achieve vanishing risk unless the target function is linear ( k P > 1 f ⇤ k L 2 = 0 ).
In the following, we compare the prediction risk of the ridge estimator on trained features against this lower bound.
3.2 Almost rank-1 property of the gradient matrix
Before we analyze the prediction risk of the ridge regression estimator on the trained CK, we ﬁrst need to understand the gradient matrix G 0 in (2.1).
The following proposition shows that the ﬁrst gradient step on W can be approximated in operator norm by a rank-1 matrix under Assumption 1.
Proposition 2.
Deﬁne G 0 :
= 1 ⌘ p N ( W 1  W 0 ) and a rank-1 matrix A :
= µ 1 n p N X > ya > .
Given
Assumption 1, there exist some constants c, C > 0 such that for all large n, N, and d , we have
[Equation]
with probability at least 1  ne  c log 2 n .
Scaling of learning rate ⌘ .
Based on the above proposition, we can now specify an appropriate learning rate ⌘ such that the change in the ﬁrst-layer weights after one gradient descent step is neither insigniﬁcant nor unreasonably large.
Assumption 1 implies that, for proportional n, d, N , the initial weight matrix satisﬁes k W [-
-]0 k = [Equation] , P
[Equation] ) , k W 0 k [Equation] ⇥ d, P ( p d ) , and due to Proposition 2, the ﬁrst gradient step satisﬁes p N k G 0 k = ⇥ d, P (1) , p N k G 0 k F = ⇥ d, P (1) .
In light of the above scaling, if we write ⌘ = ⇥ ( N ↵ ) , then ↵  0 is required so that the change in the weight matrix is non-negligible (one may verify that for ⌘ = o d (1) , the test performance of kernel ridge regression remains unchanged after one GD step).
On the other hand, when ↵ > 1 / 2 , the gradient update “overwhelms” the initialized parameters W 0 , and the preactivation feature h x , w i i in the NN (1.1) becomes unbounded as N ! 1 .
This motivates us to consider the following two regimes of learning rate scaling.
[Equation]
In Section 4, we consider small step size ⌘ = ⇥ (1) , which is parallel to common practice in NN optimization 1 .
Whereas in Section 5, we analyze the larger step size ⌘ = ⇥ ( p N ) , which resembles the learning rate scaling in the maximal update parameterization in [YH20]; in particular, from
Lemma 10 in Appendix B.1, one can easily verify that given data point x ⇠ N (0 , I ) , the change in each coordinate of the feature vector is roughly of the same order as its initialized magnitude, that is, for i 2 [ N ] ,
[Equation]
⌘ = ⇥ (1) : improvement over the initial CK
From Proposition 2, we observe that the dominant rank-1 direction in the ﬁrst-step gradient matrix
G 0 contains information of the teacher model f ⇤ (through label vector y ).
Intuitively, this indicates that the learned feature map after one GD step x 7!
( W > 1 x ) can “adapt” to f ⇤ , and hence we may expect the ridge regression estimator on the trained CK to achieve better performance.
In this section, we precisely characterize the CK prediction risk under the small learning rate ⌘ = ⇥ (1) .
We ﬁrst introduce the Gaussian equivalence property which will be useful in the risk computation.
4.1 The Gaussian equivalence property
The Gaussian Equivalence Theorem (GET) states that the performance of a nonlinear kernel model is the same as that of a noisy linear model.
Speciﬁcally, for the ridge regression estimator, deﬁne
[Equation]
where F 2 { CK , GE } indicates the choice of feature map, which can be either the nonlinear
CK feature  CK ( x ) = [Equation] p N  ( W > x ) , or the linear Gaussian equivalent (GE) feature  GE ( x ) = 1 p N  µ
[Equation]
W >
[Equation] + [Equation]  where z ⇠ N (0 , I ) is independent of x , W .
In the following, for both  CK and  GE , we take W to be the updated weight matrix W 1 after one GD step.
The Gaussian equivalence refers to the universality phenomenon R CK (  ) ⇡ R GE (  ) .
For RF models (3.1), the GET has been rigorously proved in [HL20, MS22, MM22].
Furthermore, [GLR + 21, LGC + 21] provided empirical evidence that such equivalence holds for more general feature maps, including the representation of certain pretrained NNs (e.g., see [LGC + 21, Figure 4]).
Since our setting goes beyond RF models and cannot be covered by the prior results, we establish the GET for our trained feature map under small learning rate.
Theorem 3.
Suppose that Assumption 1 holds and the activation  is an odd function.
If the learning of W 1 in (2.1) and estimation of ˆ a  in (4.1) are performed on independent training data X and ˜ X , respectively, then the GET holds after the ﬁrst-layer weight is trained for one gradient step with learning rate ⌘ = ⇥ (1) ; that is, for the CK feature  CK ( x ) = 1 p N  ( W > 1 x ) , and  > 0 ,
[Equation]
This is to say, for learning rate ⌘ = ⇥ (1) , the Gaussian equivalent model provides an accurate description of the prediction risk of CK ridge regression after one feature learning step.
The important observation is that even though the trained parameters in W 1 are no longer i.i.d., the Gaussian equivalence property can still hold when W 1  W 0 remains “small” (in some norm, see (C.3) in
Appendix C.1 for details), which entails that the neurons remain nearly orthogonal to one another.
[Equation]
[Equation]
Implications of Gaussian equivalence.
Under the GET, we can alternatively compute R GE (  ) , the prediction risk of ridge regression on noisy Gaussian features  GE , which is much easier to analyze.
Theorem 3 is empirically validated in Figure 3(a)(b), where we observe an agreement between the experimental values and the analytic predictions 2 from Section 4.2.
On the other hand, the GET also implies that the kernel estimator is essentially “linear” in high dimensions.
For the squared loss, it is straightforward to verify that the Gaussian equivalent model cannot learn the nonlinear component of the target function P > 1 f ⇤ as follows.
Fact 4.
Under the same assumptions as Theorem 3, R GE (  )  k P > 1 f ⇤ k 2 L 2 for any 1 , 2 ,  > 0 .
Hence when ⌘ = ⇥ (1) , even though training the ﬁrst-layer W for one step can lead to non-trivial improvement over the initial RF model (which we precisely quantify in Section 4.2), the learned CK cannot outperform the best linear model on the input features.
In other words, to (possibly) learn a nonlinear f ⇤ , the trained feature map needs to violate the GET.
In the case of one gradient step on
W , this amounts to using a sufﬁciently large step size, which we analyze in Section 5.
4.2 Precise asymptotics of CK ridge regression
Having established the Gaussian equivalence property for the CK ridge estimator after one gradient step with ⌘ = ⇥ (1) , we can now compute the asymptotic prediction risk for the trained kernel and compare with the initialized RF.
To quantify the discrepancy in the prediction risk (4.1), we write
R 0 (  ) as the prediction risk of the initialized RF ridge regression estimator (on the feature map x 7!
( W > 0 x ) ), and R 1 (  ) as the prediction risk of the ridge estimator on the trained feature map after one feature learning step x 7!  ( W > 1 x ) .
Importantly, because of the dependency between the trained weights W 1 and the teacher model f ⇤ (due to the gradient update (2.1)), we cannot simply apply a rotation invariance argument (e.g., [MM22, Lemma 9.2]) to remove the dependency on the true parameters  ⇤ and reduce the prediction risk to the trace of certain rational functions of the kernel matrix.
In other words, knowing the spectrum (or the Stieltjes transform) of the CK is not sufﬁcient for these purposes.
Instead, we utilize the GET and the almost rank-1 property of G 0 in Proposition 2, which, in combination with techniques from operator-valued free probability theory [MS17, AP20], enables us to obtain the asymptotic expression of the difference in the prediction risk before and after one gradient step.
Theorem 5.
Under the same assumptions as Theorem 3 and ⌘ = ⇥ (1) , we have
[Equation]
where  ( ⌘ ,  , 1 , 2 ) is deﬁned by (C.19) in Appendix C.3.
is a non-negative function of ⌘ ,  , 1 , 2 2 (0 , + 1 ) with parameters µ ⇤ 1 , µ 1 , µ 2 , and it vanishes if and only if (at least) one of µ ⇤ 1 , µ 1 and ⌘ is equal to zero.
Remark.
Performance of the initial RF ridge estimator R 0 (  ) has been characterized by the prior works [GLK + 20, MM22]; hence, the precise asymptotics of  provided in Theorem 5 allows us to explicitly compute the asymptotic prediction risk of the CK model after one gradient step, i.e. R 1 (  ) .
Theorem 5 conﬁrms our intuition that training the ﬁrst-layer parameters improves the CK model, as shown in Figure 3(a)(b).
Remarkably, this improvement (when  > 0 ) holds for any 1 , 2(0 , 1 ) , that is, taking one gradient step (with learning rate ⌘ = ⇥ (1) ) is always beneﬁcial, even when the training set size n is small.
Moreover, we do not require the student and teacher models to have the same nonlinearity — a non-vanishing decrease in the prediction risk is present as long as µ 1 , µ ⇤ 1 6 = 0 .
On the other hand, the GET also implies an upper bound on the possible improvement:   R 0 (  )  µ ⇤ 22 as n, d, N ! 1 ; this is to say, the trained CK remains in the “linear” regime.
In addition, from inspecting the asymptotic risk formulae (C.19) in Appendix C.3.2, we can easily arrive at the following characterization of two special cases of interest
• Large sample regime ( 1 ! 1 ):  is increasing with respect to the learning rate ⌘ ; that is, taking a larger step results in greater decrease in the prediction risk, as shown in Figure 3(a). • Large width regime ( 2 ! 1 ):
In this case  ! 0 ; thus, the beneﬁt of one-step feature learning (with ⌘ = ⇥ (1) ) becomes less signiﬁcant as the width increases, as shown in Figure 3(b).
[Figure]
⌘ = ⇥ ( p N ) : improvement over the kernel lower bound
In this section, we consider a gradient step with large learning rate ⌘ = ⇥ ( p N ) , which matches the asymptotic order of the Frobenius norm of the gradient G {+
+} and that of the initialized weight matrix W 0 .
Note that after absorbing the prefactors, this learning rate scaling is analogous to the maximal update parameterization [YH20], which admits a feature learning limit.
More speciﬁcally, the change in each coordinate of the feature vector [  ( W > x )]
i is ˜ ⇥
d, P (1) , which has roughly the same order of magnitude as its value at initialization.
Due to the large step size, columns of the updated weight matrix W 1 are no longer near-orthogonal, which is an important property in existing analyses of the Gaussian equivalence (e.g., see Proposition 13 in Appendix C.1 or [HL20, Equation (66)]).
Indeed, we will see that in this regime, the ridge regression estimator on the trained CK features is no longer “linear” and can potentially outperform the kernel lower bound (3.3) in the proportional limit.
However, in the absence of GET, it is difﬁcult to derive the precise asymptotics of the CK model.
As an alternative, we establish an upper bound on the prediction risk R 1 (  ) , which we then compare against the kernel ridge lower bound.
Existence of a “good” solution.
Given the trained ﬁrst-layer weights W 1 , we ﬁrst construct a second-layer ˜ a for which the prediction risk can be upper-bounded.
For a pair of nonlinearities (  ,  ⇤ ) , we introduce a scalar ⌧ ⇤ which is the optimum of the following minimization problem:
[Equation]
where ⇠ 1 , ⇠ 2 i . i .
d .
⇠ N (0 , 1) .
We write  ⇤ as an optimal value at which ⌧ ⇤ is attained (when ⌧ ⇤ is not achieved by a ﬁnite  , the same argument holds by introducing a small tolerance factor ✏ > 0 in ⌧ ⇤ ; see Appendix D.2).
Roughly speaking, ⌧ ⇤ approximates the prediction risk of a speciﬁc student model which takes the form of an average over a subset of neurons (after one feature learning step).
In particular, the ﬁrst term on the RHS of (5.1) containing
⇤ corresponds to the teacher f ⇤ , and the second term E ⇠ 2 represents the constructed student model.
The following lemma shows that we can ﬁnd some ˜ a on the trained CK features whose prediction risk is approximately ⌧ ⇤ , under the additional assumption that the activation function  is bounded.
For more details, see Appendix D.
Lemma 6 (Informal) .
Suppose that Assumption 1 holds and  is bounded.
Then, after one gradient step on W with ⌘ = ⇥ ( p N ) , there exist some second-layer coefﬁcients ˜ a such that the constructed student model ˜ f
( x ) = 1 p N ˜ a >  ( W > 1 x ) achieves a prediction risk which is “close” to ⌧ ⇤ .
It is worth noting that the deﬁnition of ⌧ ⇤ does not involve the speciﬁc value of the learning rate ⌘ .
This is because for any choice of ⌘ = ⇥ ( p N ) , due to the Gaussian initialization of a i , we can ﬁnd a subset of weights that receive a “good” learning rate (with high probability) such that the corresponding neurons are useful for learning the teacher model.
In addition, observe that ⌧ ⇤ is a simple Gaussian integral which can be numerically or analytically computed (see Appendix D.2 for more examples).
For instance, when  =  ⇤ =
erf, one can easily verify that  ⇤ = p 3 and ⌧ ⇤ = 0 .
Prediction risk of ridge regression.
Since we have established the existence of a “good” student model ˜ f that can achieve a prediction risk close to ⌧ ⇤ (as deﬁned in (5.1)), in what follows, we prove an upper bound for the prediction risk of the ridge regression estimator on the trained CK features
R 1 (  ) in terms of the scalar ⌧ ⇤ .
The proof of the following result is shown in Appendix D.3.
Theorem 7.
Under the same assumptions as Lemma 6, after one gradient step on W with ⌘ = ⇥ ( p N ) , there exist constants C, ⇤ 1 > 0 such that for any n/d > ⇤ 1 , the ridge regression estimator (4.1) with regularization parameter n "  1 < N  1  < n  " for some small " > 0 satisﬁes
[Equation]
with probability 1 as n, d, N ! 1 proportionally.
While Theorem 7 does not provide exact expression of the prediction risk, the upper bound still allows us to compare the prediction risk of the CK ridge regression before and after one large gradient step.
In particular, if k P > 1 f ⇤ k 2 L 2  10 ⌧ ⇤ (the constant 10 is not optimized), we know that the trained CK can outperform the kernel lower bound (3.3) (and also the initialized CK) in the proportional limit, when the ratio 1 = n/d is sufﬁciently large.
The following corollary provides two examples of this separation (see Figure 3(c)).
Corollary 8.
Under the same conditions as Theorem 7, there exists a constant ⇤ 1 such that for any 1 > ⇤ 1 , the following holds with probability 1 when n, d, N ! 1 proportionally:
• For  =  ⇤ =
erf , we have R 1 (  )=
O ( d / n ) .
• For  =  ⇤ =
tanh , we have R 1 (  ) <
k P > 1 f ⇤ k 2 L 2 .
In the two examples outlined above, training the features by taking one large gradient step on the ﬁrst-layer parameters can lead to substantial improvement in the performance of the CK model.
In fact, the new ridge regression estimator may outperform a wide range of kernel models as described in Section 3.1, and as shown in Figure 3(c).
However, we emphasize that this separation is only present in speciﬁc pairs of (  ,  ⇤ ) for which the scalar ⌧ ⇤ is sufﬁciently small.
In general settings, learning a good representation would likely require a training procedure that takes more than one gradient step (even if f ⇤ is as simple as a single-index model, see Figure 4(c) in Appendix A.1).
Conclusion
We investigated how the conjugate kernel of a two-layer neural network (1.1) beneﬁts from feature learning in an idealized student-teacher setting, where the ﬁrst-layer parameters W are updated by one gradient descent step on the empirical risk.
Based on the approximate low-rank property of the gradient matrix, we quantiﬁed the improvement in the prediction risk of conjugate kernel ridge regression under two different scalings of ﬁrst-step learning rate ⌘ .
To the best of our knowledge, this is the ﬁrst work that rigorously characterizes the precise asymptotics of kernel models (deﬁned by neural networks) in the presence of feature learning.
We outline a few limitations of our current analysis as well as future directions.
• Dependence between W 1 and X .
One crucial assumption that we make is that the trained weight matrix W 1 is independent of the data ˜ X on which the CK is computed.
While this does not cover the important scenario where feature learning and kernel evaluation are performed on the same data, our setting is very natural in the analysis of pretrained models or transfer learning, which would be an interesting extension.
• Scaling of learning rate.
Our ﬁndings illustrate that different learning rate scalings such as ⌘ = ⇥ (1) and ⌘ = ⇥ ( p N ) result in drastically different behavior.
One natural question to ask is whether there exists a “phase transition” in between the two regimes that dictates whether the GET holds.
Interestingly, [RGKZ21] showed that instead of breaking the near-orthogonality of the weights W (via large gradient step), one can also introduce sufﬁciently large low-rank shifts to the input X to enable the initial RF model to ﬁt a nonlinear f ⇤ .
Intuitively, this may be due to the “dual” relation of the inputs X and the weights W in the CK model. Acknowledgement
The authors would like to thank (in alphabetical order)
Konstantin Donhauser, Zhou Fan, Hong Hu,
Masaaki Imaizumi, Ryo Karakida, Bruno Loureiro, Yue M. Lu, Atsushi Nitanda, Sejun Park, Ji Xu,
Yiqiao Zhong for discussions and feedback on the manuscript.
JB was supported by NSERC Grant [2020-06904], CIFAR AI Chairs program, Google Research
Scholar Program and Amazon Research Award.
MAE was supported by NSERC Grant [201906167], Connaught New Researcher Award, CIFAR AI Chairs program and CIFAR AI Catalyst grant.
TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST.
ZW was supported by NSF Grant DMS-2055340.
DW was partially supported by a Borealis AI Fellowship.
Part of this work was completed when DW interned at Microsoft Research (hosted by GY)