High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
Authors: Jimmy Ba
Murat A Erdogdu
Taiji Suzuki
Zhichao Wang
Denny Wu
Greg Yang
Venue: NeurIPS 2022 Accept (2022)
Keywords: random matrix theory
two-layer neural network
kernel method
feature learning

Abstract: We study the first gradient descent step on the first-layer parameters $\boldsymbol{W}$ in a two-layer neural network: $f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^\top\sigma(\boldsymbol{W}^\top\boldsymbol{x})$, where $\boldsymbol{W}\in\mathbb{R}^{d\times N}, \boldsymbol{a}\in\mathbb{R}^{N}$ are randomly initialized, and the training objective is the empirical MSE loss: $\frac{1}{n}\sum_{i=1}^n (f(\boldsymbol{x}_i)-y_i)^2$. In the proportional asymptotic limit where $n,d,N\to\infty$ 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 $\boldsymbol{W}$ with learning rate $\eta$. We consider two scalings of the first step learning rate $\eta$. For small $\eta$, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large $\eta$, 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.   

High-dimensional Asymptoties of Feature Learning: How One Gradient Step Improves the Representation
Jimmy Bad, Murat A. Erdogdu', Taiji Suzuki², Zhichao Wang", Denny Wu', Greg Yang'
l University of Toronto and Vector Institute, University of Tokyo and RIKEN AIP, University of California, San Diego, *Microsoft Research AI
(jba, erdogdu, demnywu toronto.roo......ooooooo edu, wwt.:miiisiiuuttiiiie .cc ac. zhw0360ucad.cm edu, gregyangomicrosoff
Abstract
We study the first gradient descent step on the first-layer parameters W a two- layer neural network: f(x) a a a (o ax), where W E Rd a R are randomly initialized, and the training objective is the empirical MSE loss: >>=(((( 22) In the proportional asymptotic limit where n, d, N - xo at the same rate, and an idealized student-teacher setting where the teacher /* is single-index model, we compute the prediction risk ridge regression on the conjugate kernel after one gradient step on W with learning rate 1). We consider two scalings of the first step learning rate n). For small 7), we establish 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 1), we prove that for certain f" the same ridge estimator on trained features can go beyond this "linear regime" and outperform wide range of (fixed) kernels. Our results demonstrate that even one gradient step can lead to 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,
1 a (W x),
[NN(W))
a ((x, w:))
(1.1)
where x E Rdd W E B a RN a is the nonlinear activation function applied entry-wise, and the training objective is to minimize the empirical risk. Our analysis will made the proportional asymptotic limit, i.e., the number of training data n. the input dimensionality d, and the number of neurons N jointly tend 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 kernel model, where the kernel defined by features z + o ( T x) (often called the hidden representation) is referred to as the conjugate kernel (CK) [Nca95]. 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 the ability learn useful features [GDDM14, DCLT18] that "adapt" to the learning problem [S..18]. In fact, recent works have shown that such adaptivity enables NNs optimized by
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

gradient descent to outperform a wide range of linear/kereel estimators [AZL19, GMMM19). While many explanations of this separation have been proposed, our starting point the empirical finding that "non-kernel" behavior often occurs in the early phase of NN optimization, especially under large learning rates [JSF+20, 20, 20]. The goal of this work 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 (fixd 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 student- teacher setting in the proportional asymptotic limit; we aim characterize the prediction risk of the kernel ridge regression estimator on top of the first-layer CK feature a o(W ar), before and after one gradient descent step on the empirical risk (starting from Gaussian initialization).
Following prior works on the precise asymptotics of RF re- gression [GLK 20, DL20J, we focus on the setting where the input Φ is Gaussian and the teacher f* is single-index model. In this case, the prediction risk of large class of
10P
predictton risk
RF/kernel ridge regression estimators is lower-bounded by the L of the "nonlinear" component of the teacher
10
initialzed n e(1)
IP i... they only learn linear functiion on the in- put. After one gradient step 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:
n o119)
P.
16a
16
semple size
Figure 1: Prediction risk of ridge re- gression on trained CK features (crf) af- ter one feature learning step. Markers represent empirical simulations and solid curves predicted asymptotic values; red line indicates ((%nn) rate.
Small Ir: 7 = ((). In Section 4, we extend the Gaussian Equivalence Theorem (GET) in [HL20] to the updated fea- ture map after one gradient descent step W with learn- ing rate 1/ = (1); this allows us to precisely characterize the prediction risk using random matrix theoretical tools. We prove that after 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 "lincar regime" and cannot outperform the best linear estimator on the input (black dashed line).
Large Ir: 1 ((VN). In Section 5, we analyze larger learning rate that coincides with the maximal update parameterization in [YH20]. For certain target functions f³, we prove that ker- nel ridge regression after one feature learning step can achieve lower risk than the lower bound ||P>1f" 122 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 mod- els in the proportional limit [GLK+20, DL.20, LCM20, AP20, MM22]. These results typically build upon analyses of the spectrum of kernel matrices, a key ingredient which is the "linearization" of nonlinear random matrices via Taylor expansion [EK10] or orthogonal polynomials [CS13, PW17].
Consequently, large class of kernel models are essentially linear in the proportional asymptotic limit [LR20, BMR21J. In the case of RF models, similar property is captured by the Gaussian Equivalence Theorem [GMKZ20, HI.20, GLR 22], which roughly states that RF estimators achieve the same prediction risk as (noisy) linear model. For inputs with unit norm, [GMMM21, MMM21] showed that sample size n = ((d2) is required to go beyond this "lincar" regime. As we will see in certain settings, such limitation can also be overcome (in the d scaling) by training the feature map for one gradient step with a sufficiently large learning rate.
Ad Advantage of NNs over fixed kernels. It is well-known that under 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 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 kemel methods in terms of approximation and estimation ability [Bac17, Suz18, IF19. SH20, GMMM20J. More- over, 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 [AZL,19, WLIL19, GMMM19, 19, LMZ200, DM20, SA20, AZI.20, 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 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 [GASI9, LM20, PPVF21], and the NTK evolves most rapidly in the first few epochs PFD.+20]. Large learning rate in the initial steps can impact the conditioning of loss surface [JSF 20, CKL+211 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 rell-rell 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).
2 Problem setup and assumptions
Notations. Throughout this paper, || denotes the fn-nrrm for vectors and the Eg -> fg operator norm for matrices, and +IIF is the Frobenius norm. For matrix M E Rax=, tr(M) 1/2 -r is the normalized trace. Oa(-) and oa(-) stand for the standard big-O and little-o notations, where the subscript highlights the asymptotic variable; we write (-)) when the (poly-)logaritmmic factors are ignored. Oa,(--) (resp. o-)= represents big-O (resp. little-o) in probability as + 3o. 22()), (-)) are defined analogously. I' is the standard Gaussian distribution in Rd. Given f IRd R, we denote its L--norm w.r.t. r as IIffl which we abbreviate as [l when the context is clear.
2.1 Training procedure
Gradient deseent on the 1st layer. Given training examples ((rr, y=) we learn the two-layer NN (I.1) by minimizing the empirical risk: C(f) = 1. n E.= (z, yt), where é is the squared loss f(x,y) 3) = {x - y))² As previously remarked, fixing the first layer W at random initialization and learning the second layer 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 (at initialization) and perform gradient descent W. We write the initialized first-layer as Wo. and the weights after one gradient step as W1. The gradient update, which we refer to as the feature learning step, with learning rate 1) is given as: W1 = Wo + TVNN Go where
[ y
(XW o)a aT) o od(Xwo)) 1
G =
(2.1)
VN
in which @ is the Hadamard product, o' is the derivative of a (acting entry-wise), and we denoted the input feature matrix X E Rx and the corresponding label vector y F We remark that the /N Scsliing in front of 17 accounts for the -reefacoo in our definition of two-layer NN (1.1).
Ridge regression for the 2nd layer. After obtaining the updated weights W,, 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, W1 is no longer independent of X, which significantly complicates the analysis. To circumvent this difficulty, we estimate the regression coefficients à using new set of training data (ē,, ūi ki 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, } as 4 = 1 (X E BOXN the VN
CK ridge regression estimator can be obtained by solving à argmin, a { $ + A ||a 2]
3

---

2.2 Student-feacher setting and main assumptions
Given a target function (teacher model) f and a learned model f. we evaluate the model perfor-
mance using the prediction risk: R(f) = Ex (f(x)) (x))² = Il - f',l. where the expectation
taken over the test data from the same training distribution.
We utilize the orthogonal decomposition of the activation function o. Define the coefficients
Ho []]( A1 = [=))
112 - v [|o()) - ng r where 2 NN(0,)) 1). (2.2)
This implies (2) Mo + 2 + ao (2), where Fl (=)] B| O. and B 123 Similarly, for square integrable target function f", we have the orthogonal decomposition
(ax) 16 + / (z. (_)) (>))"(a), Mi3.
(2.3)
where >1 is the projector orthogonal to constant and linear functions in L (2( LT), which implies that E| (a)] = 0, =| (x=] = O. As d + oc, quantities defined in (2.3) satisfy || 1, || ll La + μdg where (4 ( μu² are bounded constants. Intuitively, 14 μū, and μáz can be interpreted the "mmanitude" of the constant, linear, and nonlinear components of f", respectively.
Assumption 1.
1. Proportional limit. n, n,d, N + OC, n/d- 1h N/d + 2/2, where 1n 4/2 E (0)
2. Gaussian initialization. /d-[Wolij Lid. ~ 1),, [al; Lid. 1)0, [d.j E [N]
3. Normalized activation. The activation function has bo-bounded first three derivatives almost surely. In addition, a satisfies μ = 0 and p1, P2 F defined in (2.2).
4. Single-index teacher. Labels are generated as yi (x;) + Es. where F; Ltd. NO, I), and
E is i.i.d. sub- Gaussian noise with mean 0 and variance o²- The teacher f*(x)= = o² ((x, B,)), where B, E Rdd with || 1, and a* is Lipschitz with μ = μ + 0 as defined in (2.3).
Remark. We make the following comments on the above assumptions.
Following [HI.20], 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-inden innlx-iidxxxxe setting has been extensively studied in the proportional regime [GLK 20, DL20, HL20]. However, prior works only considered training the coefficients on top of fixed feature map, and such RF models cannot efficiently learn single-index f* in high dimensions [YS19].
Under Assumption 1, a relatively large sample size corre- sponds to larger 11, and a relatively large network width cor- responds to larger 4/z The proportional scaling of n, d, N im- plies that the model width is not significantly larger than the training set size, in contrast the polynomial overparameter- ization often required in NTK analyses [DZPS19] which may be less realistic in practical settings.
NTK parameterizaton on
Importantly, the initialization of our two-layer NN (1.1) resem- bles the mean-field parameterization [MMN18, CB18]: the second layer is divided by additional / -factor compared to the kernel (NTK) scaling - this ensures that JNN = (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 pa- rameters W in two coordinates. Observe that under the mean- field parameterization (main figure), the neurons travel away from the initialization and align with the target function (black dashed lines), whereas in the NTK parameterization (subfig- ure, which omits the VN prefactor), the parameters remain close to their initialization and hence do not learn useful features.
Teacher 1*
Inttialiee
Trained
0
Figure 2D visualization opti- mization trajectory under mean-fict field (main) and NTK (subfigure) parame- terizations. f consists of two ReL.U neurons and the student is two-layer ReLU neural network. Darker color indicates earlier in training. and vice versa. We set d = 512, wi = doz = 10; both models are optimized until training losses are below 10

---

3 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 { Vi )i=1 we consider the following classes of kernel models for comparison.
Random features model. We introduce two RF kernels associated with (1.1) initialization: the conjugate kernel (CK) defined by features (xk VN (W a) E RRV and the neural tangent
kernel (NTK) [JGH18] defined by features @NTK (x) vNd a feature map RF E [CK,NTK), the RF ridge regression estimator can be written as
Vee (W' x)xx) R.dd Given
n { Σ (y (2))(())))))))))))))))) a)) N lall" (3.1)
[RF(#) (2)) a), á argmin aERN
71
Rotation invariant kernel model. Consider the inner-product kernel: k(z, y) = 9 (x.z)) and
the Euclidean distance kernel: k(x,,, y) g [ (==-vl")) ) where y satisfies the smoothness conditions
d
in [EKI0]. Denoting the associated RKHS with H. the kernel ridge estimator is given by
fecr argmin JEN
{ -( -(.))()))
+1I
faa(x)) k X K AT) y. (3.2)
We write the prediction risk of the above kernel estimators as Rck( RNT (A), Rker (A), respec- tively. The following lower bound the prediction risk is a simple combination of existing results.
Proposition 1 ([HL200, MZ20, BMR21)). Under Assumption 1, we have
inf min RRc (A), RNTK (A), Rier Rxer(A)) 2 | I' + oar(1),
(3.3)
A50
where P> denotes the projector orthogonal to constant and linear functions (2R
This proposition implies that in the proportional limit, ridge regression the RF 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 PI 0). In the following, 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 Go in (2.1). The following proposition shows that the first gradient step W can be approximated in operator norm by rank-1 1 matrix under Assumption I.
Proposition 2. Define Go , ( Wo) and rank-1 matrix nVN
X ya T. Given
Assumption 1, there exist some constants c, C > 0 such that for all large n, N, and d, we have
C log" n
Goo A6l-A11- <
|Goll.
with probability at least 1 ne - -e
Scaling of learning rate 7). Based on the above proposition, we can now specify an appropriate learning rate 7) such that the change in the first- weights after one gradient descent step is neither insignificant nor unreasonably large. Assumption implies that, for proportional d, N, the initial weight matrix satisfies [ oll Oar(11, [W oll @ar((VV (V)), and due to Proposition 2, the first gradient step satisfies v Goll Oar(1), VN |olll (1))
In light of the above scaling, if we write 7 = @(N*), then a 2 is required so that the change the weight matrix is non-negligible (one may verify that for 17 (1) the test performance of kernel
5

---

ridge regression remains unchanged after one GD step). On the other hand, when a > 1/2. the gradient update "overwhelms" the initialized parameters Wo. and the preactivation feature (z, w;) in the NN (1.1) becomes unbounded as N + co. This motivates us to consider the following two regimes of learning rate scaling.
Small Ir: 7= (11 I Woll oll Woll
Large Ir: @(v + || - Woll x ||oll
In Section 4, we consider small step size 1 (1), which is parallel to common practice in NN optimization'. Whereas in Section 5, we analyze the larger step size
O(vN), which resembles the learning rate scaling in the maximal update parameterization in [YH20]; in particular, from Lemma 10 in Appendix B.I, one can easily verify that given data point x ~ (O, I), the change in each coordinate of the feature vector is roughly of the same order as its initialized magnitude, that is, for i E [N]. (W x) - ( a)); x ( 0 x)), 6(1) with probability 1 as N + 00.
4 7 (1): improvement over the initial CK
From Proposition 2, we observe that the dominant rank-1 direction in the first-step gradient matrix Go contains information of the teacher model f* (through label vector y). Intuitively, this indicates that the learned feature map after GD step æ ト> oW æ) can "adapt" to f³, and hence may expect the ridge regression estimator on the trained CK to achieve better performance. In this section, precisely characterize the CK prediction risk under the small learning rate y ((1). We first introduce the Gaussian equivalence property which will be useful the risk computation.
4.1 The Gaussian equivalence property
The Gaussian Equivalence Theorem (GET) states that the performance of a nonlinear kernel model the same as that of a noisy linear model. Specifically, for the ridge regression estimator, define
R (A) E.z (())))), - (x)),
a₂- = argmin { H - (0r(2,) (o)(z))))))+ X N [la]]]]
(4.1)
where F {CK, GE} indicates the choice of feature map. which can be either the nonlinear CK feature (x) o ax), the linear Gaussian equivalent (GE) feature ΦG
VN (N1 W x p2z) where z ~ NN0 I) independent of æ. W. In the following. for both ΦCK and ΦGE, we take W to be the updated weight matrix W1 after one GD step.
The Gaussian equivalence refers to the universality phenomenon Rok(A) = RGE For RF models (3, the GET has been rigorously proved in [HL.20, 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 [L.GC+11, 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 I holds and the activation is an odd function. the learning of W1 in (2.1) and estimation of â in (4.1) are performed on independent training data X and X, respectively, then the GET holds after the first-layer weight trained for one gradient step with learning rate r = (();; that is, for the CK feature &CK (x) VN o x). and X O,
|Rek (A) RGE(A)) (1.
(4.2)
This is to say, for learning rate 77 = ()), the Gaussian equivalent model provides an accurate description of the prediction risk of CK ridge regression after one feature learning step. The impor- tant observation is that even though the trained parameters in W are no longer i.i.d., the Gaussian equivalence property can still hold when W - Wo remains "small" (in some norm, see (C.3) in Appendix C.I for details), which entails that the neurons remain nearly orthogonal to one another.
Heeriistcclll, the updated NN under 7 = @(1) remains close to the "kernel regime" in the sense that each neuron is close to initialization, ie., as N + co. [W Wolui < |[o]ys] with high probability.

---

Implications of Gaussian equivalence. Under the GET, we can alternatively compute RGB(A), 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)))), where we observe an agreement between the experimental values and the analytie predictions' 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 as follows.
Fact 4, Under the same assumptions as Theorem 3, RGE(A) 2 || for any th1, 1/2, À >>> 0.
Hence when 77 (1), even though training the first-layer W for one step can lead to non-trivial improvement over the imitial 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 nonlinear f", the trained feature map needs to violate the GET. In the case of one gradient step on W, this amounts 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 17 = (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 (A) as the prediction risk of the initialized RF ridge regression estimator (on the feature map + or (W ))), and Ri (A) the prediction risk of the ridge estimator on the trained feature map after one feature learning step x + oW x).
Importantly, because of the dependency between the trained weights W and the teacher model f (due to the gradient update (2.1)), we cannot simply apply rotation invariance argument (e.g., [MM22, Lemma 9.2l) remove the dependency on the true parameters B3, 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 not sufficient for these purposes. Instead, we utilize the GET and the almost rank-1 property of Go in Proposition 2, which, in combination with techniques from operator-valuce free probability theory [MS17, AP20J, 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 (1), we have
Ro(A) - (A) P 8(1, A, 1/, (2) 0,
where 8(1, X, wh, v2) is defined by (C.19) in Appendix C.3. Here, 5 is non-negative function of 7) A, th, 12 E (0, +00) with parameters μi, P1, P2, and it vanishes if and only if (at least) one of U. M1 and 17 is equal to zero.
Remark. Performance the the initial RF ridge estimator Ro(X) has been characterized by the prior works |GLK+20, M; hence, the precise asymptotics of provided in Theorem 5 allows to explicitly compute the asymptotic prediction risk the
the CK model after one gradient step, Lee. RR (A).
Theorem 5 confirms our intuition that training the first-layer parameters improves the CK model, as shown in Figure (b) Remarkably, this improvement (when 8 > 0) holds for any 1/1, d e (0, oo). that is, taking one gradient step (with learning rate 77 = (()) 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 u1, M1 o O. On the other hand, the GET also implies an upper bound on the possible improvement: S (A) - 11²2 as n, d, N oo; 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 (v)1 + oo): increasing with respect the learning rate 77; that is, taking a larger step results in greater decrease in the prediction risk, as shown in Figure 3(a).
Large width regime (1/2 + so): In this case 6 + 0: thus, the benefit of one-step feature learning (with 7 = (())) becomes less significant as the width increases, as shown in Figure 3(b).
"We note that when the first-layer weights are trained for more gradient steps, the GET (Theorem 3) will likely fail eventually, empirically demonstrated in Appendix A.
7

---

10
10P
Intiiiized
rribllzed OK
102
g=2
n
prediction rack
prediction ride
predictio
Posf R
10-1
7 9(1)
10
10-2
sample
(a) (1) (risk vs. sample size).
(b) 11 (1) (risk vs. width).
(c) 17 N (risk sample size).
Figure 3: Prediction risk of CK ridge regression on trained features: dots represent empirical simulations
(d = 512, averaged over 50 runs) and solid curves are asymptotic predictions; dashed black line corresponds to the kemel lower bound (3.3). (a) 1) @(1). a tanh, o² = SoftPlus: we set the = 2, À 10
0.25. (b) ( tanh, o² ReLU; we set e = 5, À 10 0.1. (c) y No for [0, 1/2]; brighter color represents larger step size. We choose
erf, 1/2
10 and oz 0.1.
5 77 - O improvement over the kernel lower bound
In this section, we consider gradient step with large learning rate = O(vN), which matches the asymptotic order of the Frobenius norm of the gradient Go and that of the initialized weight matrix Wo. 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 [o(W" x)]; is (1), which has roughly the same order of magnitude its value initialization.
Due to the large step size, columns of the updated weight matrix W are no longer near-orthogonal, which is an important property in existing analyses of the Gaussian equivalence (e.g.. see Proposi- tion 13 in Appendix C.1 [HH.20, 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 difficult to derive the precise asymptotics of the CK model. As alternative, we establish upper bound on prediction risk R₂ (A), which we then compare against the kernel ridge lower bound.
Existence of "good" solution. Given the trained first-layer weights W we first construct second-layer à for which the prediction risk can be upper-bounded. For a pair of nonlinearities (a,a") we introduce scalar 7" which is the optimum of the following minimization problem:
== inf Egg (o0(E)) - (KEE + +2)) REB
(5.1)
where &2 1iddd (0, 1). We write n² as optimal value at which r" is attained (when r* is not achieved by a finite K, the same argument holds by introducing a small tolerance factor 0 in r": see Appendix D.2). Roughly speaking, r" approximates the prediction risk of a specific student model which takes the form an average over a subsel of neurons (after one feature learning step). In particular, the first term on the RHS of (5.1) containing o" corresponds to the teacher fr, and the second term Eg, represents the constructed student model. The following lemma shows that we can find some à the trained CK features whose prediction risk is approximately T*, under the additional assumption that the activation function a is bounded. For more details, see Appendix D.
Lemma 6 (Informal). Suppose that Assumption I holds and is bounded. Then, after one gradient step on W with n (N)N), there exist some second-layer coefficients à such that the constructed student model f(x) VN ā [W x) achieves prediction risk which is "close"
worth noting that the definition of 7* does not involve the specific value of the learning rate 7). This is because for any choice of r @(vN), due to the Gaussian initialization a, we can find subset of weights that receive "good" learning rate (with high probability) such that the corresponding neurons are useful for learning the teacher model. In addition, observe that 7 is simple Gaussian integral which can be numerically analytically computed (see Appendix D.2 for more examples). For instance, when o² erf, can easily verify that x" v and 7" 0.

---

Prediction risk of ridge regression. Since we have established the existence of a "good" student model if that can achieve prediction risk close to r" (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 (A) in terms of the scalar r" 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 =
O (VN), there exist constants C, v > 0 such that for any n/dd >> v1, the ridge regression estimator (4.1) with regularization parameter n -==1_ < N-NNNN
A < n ===fooo .for some small & > 0 satisfies
Ri (A) < 10 + +c(VF V +
with probability as n. d, N + do proportionally.
While Theorem 7 does not provide exact expression of the prediction risk, the upper bound still allows us compare the prediction risk of the CK ridge regression before and after one large gra- dient step. In particular, if ||P>1, the trained CK can outperform the kernel lower bound (3.3) (and also the initialized CK) in the proportional limit, when the ratio th1 n/d is sufficiently large. The following corollary provides two examples of this separation (see Figure 3(c)).
ll ff 2 10r* (the constant 10 not optimized), we know that
Corollary 8. Under the same conditions as Theorem 7, there exists a constant ot such that for any 1 > vi. the following holds with probability when n, d, N - oo proportionally:
For o -0 =erf, we have R, (A) (.)m)) For tanh, we have Ri (A) < |P>1S*Ii..
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 only present in specific pairs of (a, o²) ) for which the scalar r" 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 single-index model, Figure 4(c) in Appendix A.
6 Conclusion
We investigated how the conjugate kernel of 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 few limitations of our current amalysis well future directions.
Dependence between Wi and X. One crucial assumption that we make is that the trained weight matrix Wi 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 transfer learning, which would be an interesting extension.
Scaling of learning rate. Our findings illustrate that different learning rate scalings such as 1) (1) and 7 ((VN) result in drastically different behavior. One natural question ask is whether there exists "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 the input X to enable the initial RF model 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 AT Chairs program, Google Research Scholar Program and Amazon Research Award. MAE was supported by NSERC Grant [2019- 06167J. 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 sup- ported by NSF Grant DMS-2055340, DW was partially supported by Borealis AI Fellowship. Part of this work was completed when DW interned at Microsoft Research (hosted by GY).
References
ABBA++21]
Emmanuel Abbe, Enric Boix Adsera, Matthew Brennan, Guy Bresler, and Dheeraj Nagaraj, The staircase property: How hierarchical structure can guide deep learning, Advances in Neural Information Processing Systems 34 (2021).
[ABAM22]
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz, The merged- staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks, arXiv preprint arXiv:2202.08658 (2022).
[Ada15]
Radoslaw Adamezak, A note on the hanson-wright inequality for random vectors with dependencies, Electronic Communications in Probability 20 (2015), 1-13.
[ADH* 19] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang, On exact computation with an infinitely wide neural net, Advances in Neural Information Processing Systems 32 (2019).
[AP20]
Ben Adlam and Jeffrey Pennington, The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization, International Conference on Machine Learning, PMLR, 2020, pp. 74-84.
[AZL19]
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels?, Advances Neural Information Processing Systems 32 (2019).
[AZL20]
Backward feature correction: How deep learning performs deep learning, Xiv preprint arXiv:2001.0444 133441111113311011 (2020).
[AZLL19]
Zeyuan Allen-Zhus Yuanzhi Li, and Yingyu Liang, Learning and generalization in overparameterized neural networks, going beyond two layers, Advances in neural in- formation processing systems 32 (2019).
[Bac17]
Francis Bach, Breaking the curse of dimensionality with convex neural networks, The Journal of Machine Learning Research 18 (2017), no. 1, 629-681.
[Bac23]
Learning theory from fromfirst principles, MIT Press, 2023.
[BES+22]
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang, High--dimensional asymptotics of feature learning in the early phase of neural network training, In Preparation (2022).
[BHMM19]
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal, Reconciling modern machine-learnine practice and the classical bias-variance trade-off, Proceedings of the National Academy Sciences 116 (2019), no. 32, 15849-15854.
[BL20]
Yu Bai and Jason D. Lee, Beyond linearization: On quadratic and higher-order ap- proximation of wide neural networks, International Conference on Learning Repre- sentations, 2020.
[BM21] Antoine Bodin and Nicolas Macris, Model, sample, and epoch-wise descents: exact solution of grudient flow in the random feature model, Advances in Neural Information Processing Systems 34 (2021).
[BM221] Peter I. Bartlett, Andrea Montanari, and Alexander Rakhlin, Deep learning: a statis- tical viewpoint, Acta numerica 30 (2021), 87-201.
10

---

[BP21]
Lucas Benigni and Sandrine Péché, Eigenvalue distribution of some nonlinear models of random matrices, Electronic Journal of Probability 26 (2021), 1-37.
[BP22]
Largest eigenvalues the conjugate kernel of single-layered neural net- works, arXiv preprint arXiv:2201 220022200004753 (2022).
[BS98]
Zhi-Dong Bai and Jack W Silverstein, No eigenvalues outside the support of the lim- iting spectral distribution of large-dimensional sample covariance matrices, The An- nals of Probability 26 (1998), no. 1, 316-44..
[BS10]
Zhidong Bai and Jack Silverstein, Spectral analysis of large dimensional random matrices, vol. 20, Springer, 2010.
[CB18]
Lenaic Chizat and Francis Bach, On the global convergence of gradient descent for over-paraun over-paraaeterized models using optimal transport, Advances in neural information processing systems, 2018, pp. 3036-3046.
[CB20]
Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss, Conference Learning Theory, PMLR, 2020, pp. 130S-1338.
[Chi22]
Lénaïc Chizat, Mean-field langevin dynamics: Exponential convergence and anneal- ing, Xiv preprint arXiv:2202.01000 (2022).
[CKI,+21] Jeremy Cohen, Simran Kaur, Yuanzhi Li, 7ico Kolter, and Ameet Talwalkar, Gradi- ent descent neural networks typically occurs the edge of stability, International Conference on Learning Representations, 2021.
[CLB21] Niladri S Chatterji, Philip M Long, and Peter L Bartlett, When does gradient descent with logistic loss find interpolating two-layer networks?,Joee Journal of Machine Learning Research 22 (2021), no. 159, 1-48.
[COB19]
Lenaic Chizat, Edouard Oyallon, and Francis Bach, On lazy training in differentiable programming, Advances in Neural Information Processing Systems 32 (2019).
[CS13]
Xiuyuan Cheng and Amit Singer, The spectrum of random inner-product kernel ma- trices, Random Matrices: Theory and Applications 2 (2013), 04, 1350010.
[DCLT18]
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova, Bert: Pre- training of deep bidirectional transformers for language understanding, arXiv preprint arXiv:18 10.040005 (2018).
[DGA20]
Ethan Dyer and Guy Gur-Ari, Asymptotics of wide networks from feynman diagrams, International Conference on Learning Representations, 2020.
[DL20]
Oussama Dhifallah and Yue M Lu, A precise performance analysis of learning with random features, arXiv preprint arXiv:2008. 11904 (2020).
[DM20]
Amit Daniely and Eran Malach, Learning parities with neural networks, Advances in Neural Information Processing Systems 33 (2020), 20356-20366
[DV13]
Yen Do and Van Vu, The spectrum of random kernel matrices: universality results for rough and varying kernels, Random Matrices: Theory and Applications (2013), no. 03, 1350005.
[DW18]
Edgar Dobriban and Stefan Wager, High-dimensional asymptotics of prediction: Ridge regression and classification, The Annals of Statistics 46 (2018), no. 1, 247- 279.
[DWY21]
Konstantin Donhauser, Mingqi Wu, and Fanny Yang, How rotational invariance of common kernels prevents generalization in high dimensions, International Conference on Machine Learning, PMLR, 2021, pp. 2804-2814,
[DZPS19]
Simon Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh, Gradient descent provably optimizes over-parameterized neural nelworks, International Conference Learning Representations, 2019.
[ЕК10]
Noureddine El Karoui, The spectrum of kernel random matrices, The Annals of Statis- tics 38 (2010), no. 1, 1-50.
[EK18]
On the impact of predictor geometry on the performance on high- dimensional redglrreuulriieed generalized robust regression estimators, Probability Theory and Related Fields 170 (2018), no. 1, 9575
11

---

[FCB22] Spencer Frei, Niladri S Chatterji, and Peter L Bartlett, Random feature ampli- fication: Feature learning and generalization in neural networks arXiv preprint arXiv:2202.0768 27662 (2022).
[FDP*20] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli, Deep learning versus kernel learning: an em- pirical study of loss landscape geometry and the time evolution of the neural tangent kernel, Advances in Neural Information Processing Systems 33 (2020), 5850-5861.
[FM19]
Zhou Fan and Andrea Montanari, The spectral norm of random inner-product kernel matrices, Probability Theory and Related Fields 173 (2019), no. 1-2, 27-85.
[FOBS06]
Reza Rashidi Far, Tamer Oraby, Wlodzimierz Bryc, and Roland Speicher, Spectra of large block matrices, arXiv preprint cs/0610045 (2006).
[FW20]
Zhou Fan and Zhichao Wang, Spectra of the conjugate kernel and neural tangent kernel for linear-width neural networks, Advances in neural information processing systems 33 (2020), 7710-7721.
[GAS19] Aditya Sharad Golatkar, Alessandro Achille, and Stefano Soatto, Time matters in reg- ularizing deep networks: Weight decay and data augmentation affect early learning dynamics, matter little near convergence, Advances in Neural Information Processing Systems 32 (2019).
[GDDM14] Ross Girshick, Jeff Donahue, Trevor Darrell, and Jitendra Malik, Rich feature hier- archies for accurate object detection and semantic segmentation, Proceedings of the IEEE conference on computer vision and pattern recognition, 2014, pp. 580-587.
20] Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mézard, and Lenka Zde- borová, Generalisation error in learning with random features and the hidden mani- fold model, International Conference on Machine Learning, PMLR, 2020, pp. 3452- 3462.
[GLR+21] Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Mare Mézard, and Lenka Zdeborová, The gaussian equivalence of generative models for learning with shallow neural networks, Proceedings of Machine Learning Research vol 145 (2021),
1-46.
[GMKZ20] Sebastian Goldt, Mare Mézard, Florent Krzakala, and Lenka Zdeborová, Modeling the influence of data structure on learning in neural networks: The hidden manifold model, Physical Review X 10 (2020), no. 4, 041044.
[GMMM19] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari, Limita- tions of lazy training of two-layers neural network, Advances in Neural Information Processing Systems 32 (2019).
[GMMM20]
When do neural networks oulperform kernel methods?, Advances in Neural Information Processing Systems 33 (2020), 14820-14830.
[GMMM21]
Linearized two-layers neural networks in high dimension, The Annals of Statistics 49 (2021), no. 2, 1029-1054.
[GS/W20] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart, Disentangling fea- ture and lazy training in deep neural networks, Journal of Statistical Mechanies: The- ory and Experiment 2020 (2020), no. 11, 113301.
[HCG21] Karl Hajjar, Lénaïc Chizat, and Christophe Giraud, Training integrable param- eterizations of deep neural networks in the infinite-width limit, arXiv preprint MrXiv:2110.155% (2021).
[HFS07] J William Helton, Reza Rashidi Far, and Roland Speicher, Operator-valued semicir- cular elements: solving a quadratic matrix equation with positivity constraints, Inter- national Mathematics Research Notices 2007 (2007), no. 9, ram086-mm086.
[HL20]
Hong Hu and Yue LL, Universality laws for hign-dimensoonal learning with random features, arXiv preprint (xXix:2000907666 (2020).
[HMS18]
J William Helton, Tobias Mai, and Roland Speicher, Applications of realizations (aka linearizations) to free probability, Journal of Functional Analysis 274 (2018), no. 1, 1-79.
12

---

[HY20]
Jiaoyang Huang and Horng-Tze Tzer Yau, Dynamics of deep neural networks and neu- rul tangent hierarchy, International conference on machine learning, PMLR, 2020, pp. 4542-4551.
[1119]
Masaaki Imaizumi and Kenji Fukumizu, Deep neural networks learn non-smooth functions effectively, The 22nd international conference artificial intelligence and statistics, PMLR, 2019, pp. 869-878.
[JGH18]
Arthur Jacot, Franck Gabriel, and Clément Hongler, Neural tangent kernel: Conver- gence and generalization in neural networks, Advances in neural information process- ing systems, 2018, pp. 8571-8580.
[JSF+20)
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras, The break-even point on optimization trajecto- ries of deep neural networks, International Conference on Learning Representations, 2020.
[JT20]
Ziwei Ji and Matus Telgarsky, Polylogarithmic width suffices for gradienl descent to achieve arbitrarily small test error with shallow relu networks, International Confer- ence on Learning Representations, 2020.
[KWLS21]] Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh, Local signal adaptivity: Provable feature learning in neural networks beyond kernels, Advances in Neural Information Processing Systems 34 (2021).
[LBD+20] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur- Ari, The large learning rate phase of deep learning: the catapult mechanism, arXiv preprint arXiv:2003.0m 2022:2000.002218 (2020).
[LCM20]
Zhenyu Liao, Romain Couillet, and Michael W Mahoney, random matrix analysis of random fourier features: beyond the gaussian kernel, precise phase transition, and the corresponding double descent, Advances in Neural Information Processing Systems 33 (2020), 13939-13950.
[LGC*21] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mezard, and Lenka Zdeborová, Learning curves of generic features maps for realistic datasets with a teacher-student model, Advances in Neural Information Processing Systems 34 (2021).
[LLC18) Cosme Louart, Zhenyu Liao, and Romain Couillet, A random matrix approach to neural networks, The Annals of Applied Probability 28 (2018), no. 1190-1248.
[LM20]
Guillaume Leclere and Aleksander Madry, The two regimes of deep network training, arXiv preprint arXiv:2002. 10376 (2020).
[LMZ200
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang, Learning over-parametrized two- layer neural networks beyond nık, Conference on learning theory, PMLR, 2020, pp. 2613-2682.
[LR20] Tengyuan Liang and Alexander Rakhlin, Just interpolate: Kernel "ridgeless" regres- sion can generalize, The Annals of Statisties 48 (2020), no. 3, 1329-1347.
[LWM19] Yuanzhi Li, Colin Wei, and Tengyu Ma, Towards explaining the regularization effect of initial large learning rate in training neural networks, Advances in Neural Infor- mation Processing Systems, 2019, pp. 11674-11685
[MKAS21] Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro, Quantifying the benefit of using differentiable learning over tangent kernels, International Conference on Machine Learning, PMLR, 2021, pp. 7379-7389.
[MM22] Song Mei and Andrea Montanari, The generulization error of random features regres- sion: Precise asymptotics and the double descent curve, Communications on Pure and Applied Mathematics 75 (2022), no. 4, 667-766.
[MMM21] Song Mei, Theodor Misiakiewicz, and Andrea Montanari, Generalization error of random feature and kernel methods: hypercontructivity and kernel matrix concentra- tion, Applied and Computational Harmonic Analysis (2021).
[MN118] Song Mei, Andrea Montanari, and Phan-Minh Nguyen, A mean field view of the land- scape of two-layer neural networks, Proceedings of the National Academy of Sciences 115 (2018), 10), no. 33, E7665-E7771..
13

---

[MS17]
James A Mingo and Roland Speicher, Free probability and random matrices, vol. 35, Springer, 2017.
[MS22]
Andrea Montanari and Basil N Saeed, Universality of empirical risk minimization, Conference on Learning Theory, PMLR, 2022. pp. 4310-4312.
[MZ20]
Andrea Montanari and Yiqiao Zhong, The interpolation phase transition in neu- ral networks: Memorization and generalization under lazy training, arXiv preprint arXiv:2007. 12t (2020).
[Nea95]
Radford M Neal, Bayesian learning for neural networks, vol. 118, Springer Science Business Media, 1995.
[Ngu21]
Phan-Minh Nguyen, Analysis of feature learning in weight-ried autoencoders via the mean field lens, arXiv preprint arXiv:2 1tt.:21022.0373 (2021).
[NS17]
Atsushi Nitanda and Taiji Suzuki, Stochastic particle gradient descent for infinite en- sembles, arXiv preprint arXiv: 1eXXi:11120 1712.05438 (2017).
[NWS22]
Atsushi Nitanda, Denny Wu, and Taiji Suzuki, Conver analysis of the mean field langevin dynamics, arXiv preprint arXiv:2201.1044 (2022).
[P6c19]
S Péché, A note on the pennington-word distribution, Electronic Communications in Probability 24 (2019), 1-7.
[PPVF21] Scott Pesme, Loucas Pillaud- Fillaud-iivien, and Nicolas Flammarion, Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity, Advanees Neural Information Processing Systems 34 (2021).
[PW17] Jeffrey Pennington and Pratik Worah, Nonlinear random matrix theory for deep learn- ing, Advances in Neural Information Processing Systems, 2017, pp. 2637-2646.
[RGKZ21]
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová, Classifying high-on gigh-dimensoonal gaussian mixtures: Where kernel methods fail and neural networks succeed, International Conference on Machine Learning, PMLR, 2021, pp. 8936- 8947.
[RR08]
Ali Rahimi and Benjamin Recht, Random features for large-scale kernel machines, Advances in neural information processing systems, 2008, pp. 1.p.1177 1177-1184.
[SA20]
Taiji Suzuki and Shunta Akiyama, Benefit of deep learning with non-convex noisy gradient descent: Provable excess risk bound and superiority to kernel methods, arXiv preprint arXiv:2012.03224 (2020).
[SH20]
Johannes Schmidt-Hieber, Nonparametric regression using deep neural networks with relu activation funcctonon The Annals of Statisties 48 (2020), no. 4, 1875-1897.
[Suz18] Taiji Suzuki, Adaptivity of deep relu network for learning in besov and mixed smooth besov spaces: optimal rate and curse of dimensionality, arXiv preprint arXiv:1810.08033 (2018).
[TAP21]
Nilesh Tripuraneni, Ben Adlam, and Jeffrey Pennington, Covariate shift in high- dimensional random feature regression, arXiv preprint arXiv:211 1t0i::2111.00220 11.02334 (2021). Roman Vershynin, High--imenssionul probability: An introduction with applications in data science, vol. 47, Cambridge university press, 2018
[Ver18]
[VSL+22] Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zde- borová, Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks, arXiv preprint arXiv:2202.00 002292 (2022).
[WGL+20] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro, Kernel and rich regimes in overparametrized models, Conference on Learning Theory, PMLR, 2020, pp. 3635- 3673.
[WLLM19] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma, Regularization mallers: Gener- alization and optimization of neural nets vs their induced kernel, Advances in Neural Information Processing Systems, 2019, pp. 9712-9724.
[WX20]
Denny Wu and Ji Xu, On the optimal weighted & regularization in overparameter- ized linear regression, Advances in Neural Information Processing Systems 33 (2020), 101 101123....
14

---

[WZ21]
Zhichao Wang and Yizhe Zhu, Deformed semicircle law and concentration of nonlin- ear random matrices for ultra-wide neural networks, arXiv preprint arXiv:2109.09304 (2021).
[Yan20]
Greg Yang, Tensor programs ii:: Neural matrix laws, arXiv preprint arXiv:2009.10683 (2020).
[YH20]
Greg Yang and Edward Hu, Feature learning in infinite-width neural networks, arXiv preprint arXiv:2011 1. 14522 (2020).
[YS19]
Gilad Yehudai and Ohad Shamir, On the power and limitations of random features for understanding neural networks, Advances in Neural Information Processing Systems 32 (2019).
Checklist
1. 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 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 reproduce the main experi- mental results (either in the supplemental material or as 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 experi- ments 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 (c.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 URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'l re using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifi- able information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "lDkSKrtcn3X", "reviewer": "Reviewer_268n", "summary": "The authors analyze the behavior of a two-layer network in the proportional asymptotic regime. The authors consider a training procedure whereby a single gradient step is taken for the weights of the initial layer, which then fixes the features for a kernel ridge regressor. The authors analyze two distinct regimes: 1) a small step size regime (step size $O(1)$), and 2) a “large” step size regime (step size $O(\\sqrt{N})$). In the small step size regime, the authors leverage a perturbative argument to derive a gaussian equivalence theorem, through which they are able to give a detailed analysis of the prediction risk. In the large step size regime, the authors establish an upper bound on the prediction risk, which for some problems can be contrasted with known lower bounds for random features model to demonstrate a separation for the 1-step training procedure.", "questions": "1) Given that the assumption of independent training sets for $W$ and $a$ is one of the main limitations of the theoretic results, would it be possible to provide empirical or heuristic evidence as to its impact (or lack thereof)? In other words - is this limitation due to the analysis or is the behavior fundamentally different?\n\n2) Is it possible to carry out the analysis of theorem 5 by taking different sample sizes for the estimation of $W$ and $a$? This could provide interesting insights into scenarios inspired by e.g. pre-training where the number of samples for $W$ may be much larger than that for $a$. This could also enable one, in the context of applying the current result to a bona-fide sample using data splitting, to tune the ratio to optimize the prediction risk.\n", "limitations": "The limitations of the work are well addressed by the authors. I do not believe this work has any particular negative social impact.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I found the paper well-written and easy to follow. The paper presents interesting insights in bridging between recent advances in characterizing random features models and more accurate models of neural network training. The results obtained by the authors are as sharp as can be reasonably expected given the problem under consideration.\n\nIn my view, the main weakness of the theory presented in this paper is that the analysis performed by the authors requires independence between the data used for the gradient step in $W$ and the data used to estimate $a$. In particular, the obvious to obtain this empirically is to split the data, causing an effective change in the ratio $N / d$ by a factor of 2. In the proportional asymptotic, this may qualitatively change the interpretation of the obtained result.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VS2CC2lct2B", "reviewer": "Reviewer_nHmM", "summary": "This paper studies the exact asymptotics of  feature learning in a two-layer neural network. The training procedure includes two steps: 1) fix the second layer and train the frist layer using gradient descent; 2) fix the first layer and optimize the second layer. The data samples used in the two steps are independent. Under the Gaussian assumption on the initialized weights and data input, a Gaussian equivalence result is established, based on which the exact formula of the generalization error is obtained. Specifically, it is rigorously shown that 1) after one small gradient step, the generalization error is smaller than the random feature model; 2) after one large step, it is possible to obtain a better estimate than the linear model. These results demonstrate the benefits of feature learning in a precise way.", "questions": "1. Should $\\boldsymbol{W}_1$ in Eq. (2.1) be $\\boldsymbol{W}_0$?\n\n2. Does the upper bound on $\\delta$ (line 292) directly follow from the form of $\\delta(\\cdot)$ in Theorem 5? If so, it is better to make an explicit statement.\n\n3. Can the authors give some intuition on how the specific model in Lemma 6 is constructed and where Eq. (5.1) come from?\n\n4. Can the also comment on the impact of the choice of $\\sigma(x)$? For example, which $\\sigma(x)$ can maximize $\\delta(\\cdot)$ in Theorem 5 and what is a good $\\sigma(x)$ for large gradient step? \n\n5. Can the authors comment on the exact asymptotics after multiple gradient steps?", "limitations": "NA", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The exact asymptotics of learning in the linear (or NTK) regime has been studied in several works before. However, what is much less unknown is the exact chacracterization beyond the linear regime, where feature learning is possible. This is also one of the key reasons why neural network can outperform traditional methods such as kernel method. This paper take an initial step towards this direction and the results are quite new and interesting. For example, a quantitative result (Theorem 5) is given on the improvement after one gradient step. Besides, the presentation is clear and the proof parts are well-written.\n\nWeaknesses: The model (single-index) and training algorithm (data-spliting) considered in this paper is relatively special and it will be interesting to further extend to more general settings.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "80tXsJ7Edbw", "reviewer": "Reviewer_SfUC", "summary": "This paper characterizes the asymptotic risk of a single hidden layer neural network after one gradient descent step. The asymptotic theoretical analysis reveals two regimes of operation depending on the scaling of the learning rate: a regime where the scaling is of order 1 where the learning is better than random weights (random features) but worse than the best linear classifier. In the second regime where the rate is $\\mathcal{0}(\\sqrt{N})$ with $N$ the number of neurons, the learning not only outperforms a random feature but also any linear classifier. Theoretical analysis is performed using tools from random matrix theory and opens the perspective of a better understanding of the feature learning process in a neural network at least after one learning stage. Some insights from the theoretical formulas are also derived to explain the interplay between the number of samples and the number of neurons in the hidden layer.", "questions": "1- In Equation (2.1) line 119, the gradient should be taken in the initial weights $W_0$ and not $W_1$. In other words, $G_0$ should be function of $W_0$ and not $W_1$?\n\n2- In the line 123 about the discussion of the ridge regression for the 2nd layer, the authors might have discussed based on one experiment how the assumption of using fresh dataset for learning $a$ affects the insights drawn in the paper. The explanation of the theoretical motivation of such assumption is great but it needs a better study of the impact of such assumption.\n\n3- Line 226 whence --> when ?\n\n4 - As mentioned earlier, some experiments on real data would have been welcome (at least not for comparing the close fit between theory and practice which will likely be not good but at least the intuitions drawn).", "limitations": "1- The authors have discussed quite thoroughly the limitations of their works. But I think applications to real data might give more strengths to the theoretical analysis. Also some experiments about assessment of the strong assumption of using different dataset for learning $W$ and $a$ would have been a great addition to the experimental part.\n\n2- The setup of the experiments should also be thoroughly discussed in the main paper for the sake of reproducibility and assessment of the pertinence of the model used.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*\n\n1- Theoretical understanding of the learning features process of neural network is a very important and challenging issue in machine learning and is one of the motivations of this work.\n\n2- The paper is very well written and the authors have done a substantial job in making it pedagogical despite the technical difficulties inherent to the tackled problem and the heaviness of the theoretical analysis.\n\n3 - Although the model studied is simple, the insights derived are non-trivial and the technical approach is original. \n\n*Weaknesses*\n\n1- In an empirical or intuitive approach to explaining neural networks, one would be tempted to think that the strong interdependence of the joint optimization of the hidden layer weights $W$ and the regression vector $a$ is the key factor that explains the high performance of neural networks. It is tempting to say that this work, by breaking this dependency (by taking another independent set of data for the second learning phase) misses a fundamental understanding of the learning process. It would have been very interesting to compare, at least empirically, the strong intuitions that are lost when this kind of hypothesis is made, for example by adding a global learning (using the same dataset) after a gradient descent step in an experiment that would not break this dependence, in order to at least quantify empirically the impact of this strong hypothesis.\n\n2- Although the experiments conducted on synthetic data are very well done, it would have been interesting to run simulations on some real data sets to test the robustness of the conclusions in a real application context.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept