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 Nare randomly initialized, and the training objective is the empirical MSE loss:1 n
P ni =
( f ( x i )  y i ) 2 .
In the proportional asymptotic limit where n, d, N ! 1at the same rate, and an idealized student-teacher setting where the teacher f ⇤ isa single-index model, we compute the prediction risk of ridge regression on theconjugate kernel after one gradient step on W with learning rate ⌘ .
We considertwo scalings of the ﬁrst step learning rate ⌘ .
For small ⌘ , we establish a Gaussianequivalence property for the trained feature map, and prove that the learned kernelimproves upon the initial random feature model, but cannot defeat the best linearmodel 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 evenone gradient step can lead to a considerable advantage over random features, andhighlight 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
[Equation] [Equation]
[Equation] [Equation]
( W >
[Equation]
) , (1.1) where x 2 R d , W 2 R d ⇥ N , a 2 R N ,  is the nonlinear activation function applied entry-wise andthe training objective is to minimize the empirical risk.
Our analysis will be made in the proportionalasymptotic limit , i.e., the number of training data n , the input dimensionality d , and the number ofneurons N jointly tend to inﬁnity.
Intuitively, this regime reﬂects the setting where the networkwidth 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 ) isreferred to as the conjugate kernel (CK) [Nea95].
When W is randomly initialized, this model is anexample of the random features (RF) model [RR08], the training and test performance of which hasbeen extensively studied in the proportional limit [LLC18, MM22].
These precise characterizationsreveal interesting phenomena also present in practical deep learning [BHMM19].
However, RF models do not fully explain the empirical success of NNs: one crucial advantage ofdeep learning is the ability to learn useful features [GDDM14, DCLT18] that “adapt” to the learningproblem [Suz18].
In fact, recent works have shown that such adaptivity enables NNs optimized bygradient 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 ﬁndingthat “non-kernel” behavior often occurs in the early phase of NN optimization, especially underlarge 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” oflearning: how the ﬁrst gradient step on the ﬁrst-layer parameters W impacts the representation of thetwo-layer NN (1.1).
Speciﬁcally, we consider regression with the squared (MSE) loss, and a studentteacher setting in the proportional asymptotic limit; we aim to characterize the prediction risk of thekernel ridge regression estimator on top of the ﬁrst-layer CK feature x 7!
( W > x ) , before andafter 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 wherethe input x is Gaussian and the teacher f
⇤ is a single-indexmodel.
In this case, the prediction risk of a large class of
RF/kernel ridge regression estimators is lower-bounded bythe L 2 -norm of the “nonlinear” component of the teacherk 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 CKridge estimator using separate training data, and compare itsprediction risk against this linear lower bound.
Our analysiswill be made under two choices of learning rate scalings
:• Small lr: ⌘ = ⇥ (1) .
In Section 4, we extend the GaussianEquivalence 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 improvementover the initial RF ridge regression model (see pink curve in Figure 1), but it remains in the “linearregime” 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 themaximal 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 boundk 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.
A plethora of recent works provided precise performance 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 random matrices, a key ingredient in which is the “linearization” of nonlinear random matrices via Taylor expansion [EK10] ororthogonal polynomials [CS13, PW17].
Consequently, a large class of kernel models are essentially linear in the proportional asymptoticlimit [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 achievethe 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 incertain settings, such a limitation can also be overcome (in the n ⇣ d scaling) by training the featuremap 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, thelearning 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 interms 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, Table2] for a survey).
These results often require a careful analysis of the landscape (e.g., properties ofglobal optimum) or optimization dynamics; in contrast, our goal is to precisely characterize the ﬁrstgradient step and demonstrate a similar separation.
Early phase of NN optimization.
Recent empirical studies suggest that properties of the ﬁnaltrained 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 initialsteps can impact the conditioning of loss surface [JSF + 20, CKL + 21] and potentially improve thegeneralization performance [LWM19, LBD + 20].
Under structural assumptions on the data, it hasbeen proved that one gradient step with sufﬁciently large learning rate can drastically decrease thetraining loss [CLB21], extract task-relevant features [DM20, FCB22], or escape the trivial stationarypoint at initialization [HCG21].
While these works also highlight the beneﬁt of one feature learningstep, 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 104 norm for matrices, and k · k F is the Frobenius norm.
For matrix M 2 R n ⇥ n , tr( M ) = 1 n Tr( M ) isthe normalized trace.
O d ( · ) and o d ( · ) stand for the standard big-O and little-o notations, where thesubscript highlights the asymptotic variable; we write ˜ O ( · ) when the (poly-)logarithmic factors areignored.
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 , wedenote 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 squaredloss ` ( x, y ) =
12 ( x  y ) 2 .
As previously remarked, ﬁxing the ﬁrst layer W at random initializationand learning the second layer a yields an RF model, which is a convex problem with closed-formsolution.
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 thefeature 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 denotedthe input feature matrix X 2 R n ⇥ d , and the corresponding label vector y 2 R n .
We remark thatthe 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 thequality of the new CK features by computing the prediction risk of the kernel ridge regressionestimator on top of the ﬁrst-layer representation.
Note that if ridge regression is performed onthe same data X , then after one feature learning step, W 1 is no longer independent of X , whichsigniﬁcantly complicates the analysis.
To circumvent this difﬁculty, we estimate the regressioncoefﬁcients ˆ a using a new set of training data { ˜ x i , ˜ y
i } ni
=1 , which for simplicity we assume to havethe 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 n 1 n k ˜ y
a k 2 +  N k a ko .
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 expectationis 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
[Equation]
( z
[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 impliesthat 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
= 141 1 , k P >
f ⇤ k L 2 !
µ ⇤ 2 , where µ ⇤ 0 , µ ⇤ 1 , µ ⇤ 2 are bounded constants.
Intuitively, µ ⇤ 0 , µ ⇤ 1 , and µ ⇤ 2 can be 142 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 almostsurely.
In addition,  satisﬁes µ 0 = 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 0 and variance  2 " .
The teacher f ⇤ ( x ) =
⇤ ( h x ,  ⇤ i ) ,
where  ⇤ 2 R d with k  ⇤ k
= 1 , and  ⇤ is Lipschitz with µ ⇤
= 0 , µ ⇤ 1 6 = 0 as deﬁned in (2.3) .
Remark.
We make the following comments on the above assumptions.
• Following [HL
, 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 +
, DL
HL20].
However, prior works only considered training the coefﬁcients a on top of ﬁxed featuremap, 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 thetraining set size, in contrast to the polynomial overparameterization often required in NTK analyses [DZPS19], which maybe less realistic in practical settings.
Importantly, the initialization of our two-layer NN (1.1) resembles the mean-ﬁeld parameterization [MMN18, CB18]: thesecond layer is divided by an additional p N -factor comparedto 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 Figure2, 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 awayfrom the initialization and align with the target function (blackdashed lines), whereas in the NTK parameterization (subﬁgure, which omits the 1 p N -prefactor), the parameters remainclose 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 onthe 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.
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]
Nd Vec
[Equation]
( W
[Equation]
0 x ) x >  2 R Nd .
Givena 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 ,
the
kernel ridge estimator is given by
[Equation]
We write the prediction risk of the above kernel estimators as R CK (  ) , R NTK (  ) , 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 rotationallyinvariant kernels deﬁned above does not outperform the best linear estimator on the input – it cannotachieve vanishing risk unless the target function is linear ( k P > 1 f
⇤ k L 2 = 0 ).
In the following, wecompare 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 ﬁrstneed to understand the gradient matrix G 0 in (2.1).
The following proposition shows that the ﬁrstgradient 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, 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 appropriatelearning rate ⌘ such that the change in the ﬁrst-layer weights after one gradient descent step is neitherinsigniﬁcant nor unreasonably large.
Assumption 1 implies that, for proportional n, d, N , the initialweight matrix satisﬁes k W
k = ⇥ d, P (1) , k W 0 k F = ⇥ d, P ( p d ) , and due to Proposition 2, the ﬁrstgradient 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 theweight 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 , thegradient update “overwhelms” the initialized parameters W 0 , and the preactivation feature h x , w i iin the NN (1.1) becomes unbounded as N ! 1 .
This motivates us to consider the following tworegimes of learning rate scaling.
[Equation]
In Section 4, we consider small step size ⌘ = ⇥ (1) , which is parallel to common practice in NNoptimization 1 .
Whereas in Section 5, we analyze the larger step size ⌘ = ⇥ ( p N ) , which resemblesthe learning rate scaling in the maximal update parameterization in [YH20]; in particular, from
Lemma 10 in Appendix one can easily verify that given data point x ⇠ N (0 , I ) , the change in eachcoordinate of the feature vector is roughly of the same order as its initialized magnitude, that is, fori 2 [ N ] ,
[Equation]
⇥
[Equation]
) with probability 1 as 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 indicatesthat the learned feature map after one GD step x 7!
( W > 1 x ) can “adapt” to f ⇤ , and hence wemay expect the ridge regression estimator on the trained CK to achieve better performance.
In thissection, 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 modelis 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 ) = 1 p N  ( W > x ) , or the linear Gaussian equivalent (GE) feature  GE ( x ) =1 p N  µ
[Equation]
W > x + µ 2 z  where z ⇠ N (0 , I ) is independent of x , W .
In the following, for both  CKand  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 237 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 generalfeature 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 establishthe 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 learningof 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 withlearning 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 accuratedescription 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 Gaussianequivalence property can still hold when W 1  W 0 remains “small” (in some norm, see (C.3) in
Appendix for details), which entails that the neurons remain nearly orthogonal to one another.
[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 toanalyze.
Theorem 3 is empirically validated in Figure 3(a)(b), where we observe an agreementbetween 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 thesquared loss, it is straightforward to verify that the Gaussian equivalent model cannot learn thenonlinear 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-trivialimprovement over the initial RF model (which we precisely quantify in Section 4.2), the learned CKcannot outperform the best linear model on the input features.
In other words, to (possibly) learn anonlinear 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 gradientstep with ⌘ = ⇥ (1) , we can now compute the asymptotic prediction risk for the trained kernel andcompare 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 mapx 7!
( W > 0 x ) ), and R 1 (  ) as the prediction risk of the ridge estimator on the trained feature mapafter one feature learning step x 7!
( W > 1 x ) .
Importantly, because of the dependency between the trained weights W 1 and the teacher modelf ⇤ (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 predictionrisk to the trace of certain rational functions of the kernel matrix.
In other words, knowing thespectrum (or the Stieltjes transform) of the CK is not sufﬁcient for these purposes.
Instead, weutilize the GET and the almost rank-1 property of G 0 in Proposition 2, which, in combination withtechniques from operator-valued free probability theory [MS17, AP20], enables us to obtain theasymptotic 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.28) 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) oneof µ ⇤ 1 , µ 1 and ⌘ is equal to zero.
Remark.
Performance of the initial RF ridge estimator R 0 (  ) has been characterized by the priorworks [GLK + 20, MM22]; hence, the precise asymptotics of  provided in Theorem 5 allows us toexplicitly 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 2(0 , 1 ) , that is, taking one gradient step (with learning rate ⌘ = ⇥ (1) ) is always beneﬁcial, evenwhen the training set size n is small.
Moreover, we do not require the student and teacher models tohave 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.28), 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 matchesthe asymptotic order of the Frobenius norm of the gradient G 0 and that of the initialized weightmatrix W 0 .
Note that after absorbing the prefactors, this learning rate scaling is analogous to themaximal 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 thesame order of magnitude as its value at initialization.
Due to the large step size, the columns of the updated weight matrix W 1 are no longer nearorthogonal, which is an important property used in existing analyses of the Gaussian equivalence(e.g.
, see Proposition 14 or [HL20, Equation (66)]).
Indeed, we will see that in this regime, the ridgeregression estimator on the trained CK features is no longer “linear” and can potentially outperformthe kernel lower bound (3.3) in the proportional limit.
However, in the absence of GET, it is difﬁcultto derive the precise asymptotics of the CK model.
As an alternative, we establish an upper boundon 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 asecond-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 notachieved 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 studentmodel 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 ⇤ , andthe second term E
⇠ 2 represents the constructed student model.
The following lemma shows that wecan ﬁnd some ˜ a on the trained CK features whose prediction risk is approximately ⌧ ⇤ , under theadditional assumption that the activation function  is bounded.
Lemma 6 (Informal) .
Suppose that Assumption 1 holds and  is bounded.
Then, after one gradientstep on W with ⌘ = ⇥ ( p N ) , there exist some second-layer coefﬁcients ˜ a such that the constructedstudent 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 thecorresponding neurons are useful for learning the teacher model.
In addition, observe that ⌧ ⇤ is asimple Gaussian integral which can be numerically or analytically computed (see Appendix D.2 forsome 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” studentmodel ˜ f that can achieve a prediction risk close to ⌧ ⇤ (as deﬁned in (5.1)), in what follows, we provean upper bound for the prediction risk of the ridge regression estimator on the trained CK features
R 1 (  ) in terms of the scalar ⌧ ⇤ .
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 stillallows 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 thatthe trained CK can outperform the kernel lower bound (3.3) (and also the initialized CK) in theproportional limit, when the ratio 1 = n/d is sufﬁciently large.
The following corollary providestwo 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 any1 > ⇤ 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.
Infact, the new ridge regression estimator may outperform a wide range of kernel models as describedin Section 3.1, and as shown in Figure 3(c).
However, we emphasize that this separation is onlypresent 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 onegradient step (even if f ⇤ is as simple as a single-index model).
Conclusion
We investigated how the conjugate kernel of a two-layer neural network (1.1) beneﬁts from featurelearning in an idealized student-teacher setting, where the ﬁrst-layer parameters W are updated byone gradient descent step on the empirical risk.
Based on the approximate low-rank property ofthe gradient matrix, we quantiﬁed the improvement in the prediction risk of conjugate kernel ridgeregression 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ﬁnedby 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, [RGKZ
] showed that instead of breaking the near-orthogonality ofthe weights W (via large gradient step), one can also introduce sufﬁciently large low-rank shiftsto the input X to enable the initial RF model to ﬁt a nonlinear f ⇤ .
Intuitively, this may be due tothe “dual” relation of the inputs X and the weights W in the CK model.