High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
Abstract
We study the first gradient descent step on the first-layer parameters W in a twolayer neural network: f ( x ) = 1 √ N a ⊤ σ ( W ⊤ x ) , where W ∈ R d × N , a ∈ R N are randomly initialized, and the training objective is the empirical MSE loss:n (cid:80) 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 first 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 sufficiently 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 (fixed) 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 ) = 1 √
N
N (cid:88) i = a i σ ( ⟨ x , w i ⟩ ) =
[Equation]
where x ∈ R d , W ∈ R d × N , a ∈ 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 infinity.
Intuitively, this regime reflects the setting where the network width and data size are comparable, which is consistent with practical choices of model scaling.
When the first layer W is fixed and the second layer a is optimized, we arrive at a kernel model, where the kernel defined by features x (cid:55)→ σ ( 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 finding 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 (fixed) kernel in the proportional limit?
1.1 Contributions
Motivated by the above observations, we investigate a simplified scenario of the “early phase” of learning: how the first gradient step on the first-layer parameters W impacts the representation of the two-layer NN (1.1).
Specifically, 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 first-layer CK feature x
(cid:55)→ σ ( W ⊤ x ) , before and after one gradient descent step on the empirical risk (starting from Gaussian initialization).
[Figure] 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 ∥ P > [-
-]1 f ∗ ∥ 2 L 2 , i.e., they only learn linear functions on the in- put.
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: η = Θ( √ 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 ∥ P > [-
-]1 f ∗ ∥
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 sufficiently large learning rate.
Advantage of NNs over fixed kernels.
It is well-known that under a specific 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 first gradient step and demonstrate a similar separation.
Early phase of NN optimization.
Recent empirical studies suggest that properties of the final trained model is strongly influenced by the early stages of optimization [GAS19, LM20, PPVF21], and the NTK evolves most rapidly in the first 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 sufficiently 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 benefit 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, ∥ · ∥ denotes the ℓ 2 -norm for vectors and the ℓ 2 → ℓ 2 operator norm for matrices, and ∥·∥ F is the Frobenius norm.
For matrix M ∈ 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 → ∞ .
Ω( · ) , Θ( · ) are defined analogously.
Γ is the standard Gaussian distribution in R d .
Given f : R d → R , we denote its L p -norm w.r.t.
Γ as ∥ f ∥ L p ( R d , Γ) , which we abbreviate as ∥ f ∥ 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 (cid:80) ni [-=
 `-]=1 ℓ ( f ( x i ) , y i ) , where ℓ is the squared loss ℓ ( x, y ) =
12 ( x − y ) 2 .
As previously remarked, fixing the first 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 first fix a (at initialization) and perform gradient descent on W .
We write the initialized first-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 + η √ N · G 0 where [Equation]
[Equation]
) in which ⊙ is the Hadamard product, σ ′ is the derivative of σ (acting entry-wise), and we denoted the input feature matrix X ∈ R n × d , and the corresponding label vector y ∈ R n .
We remark that the √ N -scaling in front of η accounts for the 1 √ N -prefactor in our definition 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 first-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 significantly complicates the analysis.
To circumvent this difficulty, we estimate the regression coefficients ˆ 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 √ N σ ( ˜ XW 1 ) ∈ R n × N , the
CK ridge regression estimator can be obtained by solving ˆ a = argmin a (cid:110) 1 n ∥ ˜ y − Φ a ∥ 2 + λN ∥ a ∥ 2 (cid:111) .
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 )) [-
-]2 = ∥
ˆ f − f ∗ ∥ 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 σ .
Define the coefficients µ [Equation]
E
[Equation] where z ∼
[Equation] [Equation]
[Equation]
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 → ∞ , quantities defined in (2.3) satisfy ∥ β ∗ ∥ = 1 , ∥ P > 1 f ∗ ∥
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 → ∞ , n/d → ψ 1 , N/d → ψ 2 , where ψ 1 , ψ 2 ∈ (0 , ∞ ) .
Gaussian initialization.
√ d · [ W
] ij i .
i .
d . ∼ N (0 , 1) , √ N · [ a ] j i .
i .
d . ∼ N (0 , 1) , for i ∈ [ d ] , j ∈ [ N ] .
Normalized activation.
The activation function σ has λ σ -bounded first three derivatives almost surely.
In addition, σ satisfies µ 0 = 0 and µ 1 , µ 2 ̸ = 0 defined 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 ) =
σ ∗ ( ⟨ x , β ∗ ⟩ ) , where β ∗ ∈ R d with ∥ β ∗ ∥ = 1 , and σ ∗ is Lipschitz with µ ∗ 0 = 0 , µ ∗ 1 ̸ = 0 as defined in (2.3) .
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 coefficients a on top of fixed feature map, and such RF models cannot efficiently 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 significantly 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-field parameterization [MMN18, CB18]: the second layer is divided by an additional √ 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 first-layer parameters W in two coordinates.
Observe that under the meanfield parameterization (main figure), the neurons travel away from the initialization and align with the target function (black dashed lines), whereas in the NTK parameterization (subfigure, which omits the 1 √ 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 benefits 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 fixed kernels.
Specifically, 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) defined by features ϕ CK ( x ) = 1 √ N σ ( W ⊤ 0 x ) ∈ R N , and the neural tangent kernel (NTK)
[JGH18] defined by features ϕ NTK ( x ) =
[Equation]
∈ R Nd .
Given
a feature map RF ∈ { 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 (cid:16) ⟨ x , y ⟩ d (cid:17) , and the Euclidean distance kernel: k
( x , y ) =
g (cid:16) ∥
[Equation]
(cid:17) , where g satisfies 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 rotationally invariant kernels defined above does not outperform the best linear estimator on the input – it cannot achieve vanishing risk unless the target function is linear ( ∥ P > 1 f ∗ ∥ 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 first need to understand the gradient matrix G 0 in (2.1).
The following proposition shows that the first gradient step on W can be approximated in operator norm by a rank-1 matrix under Assumption 1.
Proposition 2.
Define G 0 :
= 1 η √ N ( W 1 − W 0 ) and a rank-1 matrix A := µ 1 n √ 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 first-layer weights after one gradient descent step is neither insignificant nor unreasonably large.
Assumption 1 implies that, for proportional n, d, N , the initial weight matrix satisfies ∥
[Equation]
d, P ( √ d ) , and due to Proposition 2, the first gradient step satisfies √ N ∥ G 0 ∥ = Θ d, P (1) , √ N ∥ G 0 ∥ 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 ⟨ x , w i ⟩ in the NN (1.1) becomes unbounded as N → ∞ .
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 η = Θ( √ 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 ∈
[ N ] , (cid:12)(cid:12) σ ( W ⊤ 1 x ) − σ ( W ⊤ 0 x ) (cid:12)(cid:12) i ≍ (cid:12)(cid:12) σ ( W ⊤ 0 x ) (cid:12)(cid:12) i = ˜Θ(1) with probability 1 as N → ∞ .
η = Θ(1) : improvement over the initial CK
From Proposition 2, we observe that the dominant rank-1 direction in the first-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 (cid:55)→ σ ( 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 first 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.
Specifically, for the ridge regression estimator, define
[Equation]
where F ∈ { CK , GE } indicates the choice of feature map, which can be either the nonlinear
CK feature
ϕ CK ( x ) = 1 √ N σ ( W ⊤ x ) , or the linear Gaussian equivalent (GE) feature ϕ GE ( x ) = 1 √ N (cid:0) µ 1 W ⊤ x + µ 2 z (cid:1) 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 first-layer weight is trained for one gradient step with learning rate η = Θ(1) ; that is, for the CK feature ϕ CK [Equation] x [Equation]
[Equation]
N σ
[Equation]
⊤
[Equation]
and λ [Equation] [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.
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 ( λ ) ≥ ∥ P > 1 f ∗ ∥ 2 L 2 for any ψ 1 , ψ 2 , λ > 0 .
Hence when η = Θ(1) , even though training the first-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 sufficiently 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
(cid:55)→ σ ( 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 (cid:55)→ σ ( 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 sufficient 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 defined by (C.19) in Appendix C.3.
Here, δ is a non-negative function of η, λ, ψ 1 , ψ 2 ∈ (0 , + ∞ ) 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 confirms our intuition that training the first-layer parameters improves the CK model, as shown in Figure 3(a)(b).
Remarkably, this improvement (when δ > 0 ) holds for any ψ 1 , ψ 2 ∈ (0 , ∞ ) , that is, taking one gradient step (with learning rate η = Θ(1) ) is always beneficial, 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 ̸ = 0 .
On the other hand, the GET also implies an upper bound on the possible improvement: δ ≤ R 0 ( λ ) − µ ∗ 22 as n, d, N → ∞ ; this is to say, the trained CK remains in the “linear” regime.
Fore the details of Theorem 5, see Appendix C.3.2. Additionally, from inspecting the asymptotic risk formulae (C.19), we can arrive at the following characterization of two special cases of interest
• Large sample regime ( ψ 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 → ∞ ):
In this case δ → 0 ; thus, the benefit of one-step feature learning (with η = Θ(1) ) becomes less significant as the width increases, as shown in Figure 3(b).
[Figure]
η = Θ( √ N ) : improvement over the kernel lower bound
In this section we consider a gradient step with large learning rate η = Θ( √ N ) , which matches the asymptotic order of the Frobenius norm of the gradient G [-
-]0 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 specifically, 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 difficult 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 first-layer weights W 1 , we first 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]
(5.1) 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 finite κ , the same argument holds by introducing a small tolerance factor ϵ > 0 in τ ∗ ; see Appendix D.2).
Roughly speaking, τ ∗ approximates the prediction risk of a specific student model which takes the form of an average over a subset of neurons (after one feature learning step).
In particular, the first 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 find 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 η = Θ( √ N ) , there exist some second-layer coefficients ˜ a such that the constructed student model ˜ f
( x ) = 1 √ N ˜ a ⊤ σ ( W ⊤ 1 x ) achieves a prediction risk which is “close” to τ ∗ .
It is worth noting that the definition of τ ∗ does not involve the specific value of the learning rate η .
This is because for any choice of η = Θ( √ N ) , due to the Gaussian initialization of a i , we can find 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 κ ∗ = √ 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 defined 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 η = Θ( √ 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 satisfies
[Equation]
with probability 1 as n, d, N → ∞ 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 ∥ P > 1 f ∗ ∥ 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 sufficiently 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 → ∞ proportionally:
• For σ = σ ∗ =erf , we have R 1 ( λ )=
O ( d / n ) .
• For σ = σ ∗ =
tanh , we have R 1 ( λ ) <
∥ P > 1 f ∗ ∥ 2 L 2 .
In the two examples outlined above, training the features by taking one large gradient step on the first-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 specific pairs of ( σ, σ ∗ ) for which the scalar τ ∗ is sufficiently 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) benefits from feature learning in an idealized student-teacher setting, where the first-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 quantified the improvement in the prediction risk of conjugate kernel ridge regression under two different scalings of first-step learning rate η .
To the best of our knowledge, this is the first work that rigorously characterizes the precise asymptotics of kernel models (defined 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 findings illustrate that different learning rate scalings such as η = Θ(1) and η = Θ( √ 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 sufficiently large low-rank shifts to the input X to enable the initial RF model to fit 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). Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[No] (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)?
[N/A] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[N/A] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A] (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)?
[N/A] 4.
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 identifiable 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]