Learning Hierarchical Polynomials with Three-Layer Neural Networks
Authors: Zihao Wang
Eshaan Nichani
Jason D. Lee
Venue: ICLR 2024 poster (2024)
Keywords: Hierarchical polynomials
feature learning
three-layer networks
sample complexity
gradient descent
Primary area: learning theory
TL;DR: We show that three-layer neural networks learn hierarchical polynomials of the form $h = g \circ p$, where $p : \mathbb{R}^d \rightarrow \mathbb{R}$ is a degree $k$ polynomial, in $\widetilde O(d^k)$ samples.

Abstract: We study the problem of learning hierarchical polynomials over the standard Gaussian distribution with three-layer neural networks. We specifically consider target functions of the form $h = g \circ p$ where $p : \mathbb{R}^d \rightarrow \mathbb{R}$ is a degree $k$ polynomial and $g: \mathbb{R} \rightarrow \mathbb{R}$ is a degree $q$ polynomial. This function class generalizes the single-index model, which corresponds to $k=1$, and is a natural class of functions possessing an underlying hierarchical structure. Our main result shows that for a large subclass of degree $k$ polynomials $p$, a three-layer neural network trained via layerwise gradient descent on the square loss learns the target $h$ up to vanishing test error in $\widetilde O(d^k)$ samples and polynomial time. This is a strict improvement over kernel methods, which require $\widetilde \Theta(d^{kq})$ samples, as well as existing guarantees for two-layer networks, which require the target function to be low-rank. Our result also generalizes prior works on three-layer neural networks, which were restricted to the case of $p$ being a quadratic. When $p$ is indeed a quadratic, we achieve the information-theoretically optimal sample complexity $\widetilde O(d^2)$, which is an improvement over prior work (Nichani et al., 2023) requiring a sample size of $\widetilde\Theta(d^4)$. Our proof proceeds by showing that during the initial stage of training the network performs feature learning to recover the feature $p$ with $\widetilde O(d^k)$ samples. This work demonstrates the ability of three-layer neural networks to learn complex features and as a result, learn a broad class of hierarchical functions.

Published as a conference paper at ICLR 2024
LEARNING HIERARCHICAL POLYNOMIALS PHREE-LAYER NEURAL NETWORKS
WITH
Zihao Wang
Eshaan Nichani & Jason D. Lee Princeton University
Peking University
itaowang@stu. pku edu.
(eshnich, jasonaaaeonenoeeneeeee www.:neeeeeepriietooooooooom
ABSTRACT
We study the problem of learning hierarchical polynomials over the standard Gaussian distribution with three-layer neural networks. We specifically consider target functions of the form h = g o p where p : IRd - R is a degree & polyno- mial and g I2 - I is degree q polynomial. This function class generalizes the single-index model, which corresponds to k = 1, and a natural class of functions possessing underlying hierarchical structure. Our main result shows that for large subclass of degree polynomials p. a three-layer neural network trained via layerwise gradient descent on the square loss learns the target h up to vanishing test error in (dd) samples and polynomial time. This strict improvement over kernel methods, which require @(da)) samples, as well as existing guaran- tees for two-layer networks, which require the target function to be low-rank. Our result also generalizes prior works on three-layer neural networks, which were restricted the of p being quadratic. When p is indeed quadratic, we achieve the information iooommation-hhooretcllllly optimal sample complexity Odd)), which is an improvement over prior work (Nichani et al., 2023) requiring a sample size of (dd)). Our proof proceeds by showing that during the initial stage of training the network performs feature learning to recover the feature p with Ō(d)) sam- ples. This work demonstrates the ability of three-layer neural networks to learn complex features and as result, learn a broad class of hierarchical functions.
INTRODUCTION
Deep neural networks have demonstrated impressive empirical successes across a wide range of domains. This improved accuracy and the effectiveness of the modern pretraining and finetuning paradigm often attributed to the ability neural networks to efficiently learn input features from data. On "real-world" learning problems posited to be hierarchical in nature, conventional wisdom is that neural networks first learn salient input features to more efficiently learn hierarchical functions depending these features. This feature learning capability is hypothesized to be key advantage of neural networks over fixed-feature approaches such as kernel methods (Wei et al., 2020; Allen- Zhu & Li, 2020b; Bai Lee, 2020).
Recent theoretical work has sought to formalize this notion of a hierarchical function and understand the process by which neural networks learn features. These works specifically study which classes of hierarchical functions can be efficiently learned via gradient descent on neural network, with sample complexity improvement over kernel methods or shallower networks that cannot utilize the hierarchical structure. The most common such example is the multi-index model, in which the target f" depends solely on the projection of the data onto low-rank subspace, i.e f*((x) = g(Ux) for projection matrix U E igrxd and unknown link function g : RT - R. Here, hierarchical learning process simply needs to extract the hidden subspace U and learn the r-dimensional function g. Prior work (Abbe et al., 2022; 2023; Damian al., 2022; Bietti al., 2022) shows that two-layer neural networks trained via gradient descent indeed learn the low-dimensional feature Ux and thus learn multi-index models with improved sample complexity over kernel methods.
Beyond the multi-index model, there is growing work on the ability of deeper neural networks to learn more general classes of hierarchical functions. (Safran & Lee, 2022; Ren al., 2023;

---

Published as a conference paper at ICLR 2024
Nichami et al., 2023) show that three-layer networks trained with variants of gradient descent can learn hierarchical targets of the form h = gop. where p is a simple nonlinear feature such as the norm p(z) = ||>|| quadratie p(z) x Ar. However, it remains an open question to understand whether neural networks can more efficiently learn a broader class of hierarchical functions.
1.1 OUR RESULTS
In this work, we study the problem of learning hierarchical polynomials over the standard d- dimensional Gaussian distribution. Specifically, we consider learning the target function h IR - F. where h is equipped with the hierarchical structure h g o p for polynomials g + I and p R= + IR of degree q and k respectively. This class of functions is a generalization of the single-index model, which corresponds to k 1.
Our main result, Theorem 1, is that for a large class of degree k polynomials p, a three-layer neural network trained via layer-wise gradient descent can efficiently learn the hierarchical polynomial = go p in (dd)) samples. Crucially, this sample complexity is a significant improvement over learning h via kernel method, which requires ((dsk) samples (Ghorbani et al., 2021). Our high level insight is that the sample complexity of learning o is the same as that of learning the feature as p can be extracted from the low degree terms of gop. Since neural networks learn in increasing complexity (Abbe et al., 2022; 2023; Xu, 2020), such learning process is easily implemented by GD on three-layer neural network. We verify this insight both theoretically via our layerwise training procedure (Algorithm 1) and empirically via simulations in Section A.
Our proof proceeds by showing that during the initial stage of training the network implements kernel regression in d-dimensions to learn the feature p even though it only sees o p, and in the next stage implements 1D kernel regression to fit the link function g. This feature learning during the initial stage relies on showing that the low-frequus low-frequency component of the target function o is approximately proportional to the feature p. by the "approximate Stein's 's Lemma" stated in Lemma 2, which is our main technical contribution. This demonstrates that three-layer networks trained with gradient descent, unlike kernel methods, do allow for adaptivity and thus the ability to learn features.
1.2 RELATED WORKS
Kernel Methods. Initial learning guarantees for neural networks relied on the Neural Tangent Kernel (NTK) approach, which couples GD dynamics those of the network's linearization about the initialization (Jacot et al.. 2018; Soltanolkotabi et al., 2018; Du et al., 2018; Chizat et al., 2019). However, the NTK theory fails to capture the success of neural networks in practice (Arora et al., 2019; Lee et al., 2020; E. et al., 2020). Furthermore, Ghorbani et al. (2021) presents lower bound showing that for data uniform on the sphere, the NTK requires 12(d²) samples to learn any degree polynomial in dimensions. Crucially, networks in the kernel regime cannot learn features (Yang & Hu, 2021), and hence cannot adapt low-dimensional structure. An important question is thus to understand how neural networks are able to adapt to underlying structures in the target function and learn salient features, which allow for improved generalization over kernel methods.
Two-layer Neural Networks. Recent work has studied the ability of two-layer neural networks to learn features and as a consequence learn hierarchical functions with a sample complexity improve- ment over kernel methods. For isotropic data, two-layer neural networks are capable of efficiently learning multi-index models, i.c. functions of the form f (x) g(Zz. Specifically, for Gaussian covariates, Damian al. (2022); Abbe et al. (2023); Dandi et al. (2023) show that two-layer neu- ral networks learn low-rank polynomials with a sample complexity whose dimension dependence does not scale with the degree of the polynomial, and Bietti et al. (2022); Ba et al. (2022) show two-layer networks efficiently learn single-index models. For data uniform on the hypercube, Abbe et al. (2022) shows learnability a special class of sparse boolean functions in (d) steps of SGD. These prior works rely layerwise training procedures which learn the relevant subspace in the first stage, and fit the link function g in the second stage. Relatedly, fully connected networks trained via gradient descent on standard image classification tasks have been shown to learn such relevant low-rank features (Lee et al., 2007; Radhakrishman et al., 2022).
Three-layer Neural Networks. Prior work has also shown that three-layer neural networks can learn certain classes of hierarchical functions. Chen et al. (2020) shows that three-layer networks can
2

---

Published as a conference paper at ICLR 2024
more efficiently learn low-rank polynomials by decomposing the function 2" as (2P/2)2. Allen-Zhu et al. (2019) uses a modified version of GD to improperly learn a class of three-layer networks via second-order variant of the NTK. Safran & Lee (2022) shows that certain ball indicator functions of the form 1 | are efficiently learnable via GD on a three-layer network. They accompany this with a lower bound showing that such targets are not even approximatable by polynomially-sized two-layer networks. Ren et al. (2023) shows that multi-layer mean-field network can learn the | Our work considers broader class of hierarchical functions and features.
target ReLU((
Our work is most similar to Allen-Zhu & Li (2019; 2020a); Nichani et al. (2023). Allen-Zhu & Li (2019) considers learning target functions of the form p + ag о p with a three-layer residual network similar our architecture (1). They consider a similar hierarchical learning procedure where the first layer learns p while the second learns g. However Allen-Zhu & Li (2019) can only learn the target up to O(a²) error, while our analysis shows learnability of targets of the form p. corresponding to a = @(1) up to (1) error. Allen-Zhu & Li (2020a) shows that a deeper network with quadratic activations learns similar class of hierarchical functions up to arbitrarily small error, but crucially requires a to be (1). We remark that our results do require Gaussianity of the input distribution, while Allen-Zhu & Li (2019; 2020a) hold for more general class of data distributions. Nichani et al. (2023) shows that a three-layer network trained with layerwise GD, where the first stage consists of single gradient step, efficiently learns the hierarchical function g o p when p is quadratic, with width and sample complexity (dd)). Our Theorem I extends this result to the case where p is a degree k polynomial. Furthermore, when p is quadratic, Corollary 1 shows that our algorithm only requires width and sample complexity of (d?)), which matches the information- theoretic lower bound. Our sample complexity improvement for quadratic features relies on showing that running gradient descent for multiple steps can more efficiently extract the feature p during the feature learning stage. Furthermore, the extension to degree h polynomial features relies on a generalization of the approximate Stein's 's lemma, key technical innovation of our work.
1.3 NOTATIONS
We let Ei, denote the sum over increasing sequences fs),... i.e Eircise..ei." We use X S Y denote X & CY for some absolute positive constant C and X Σ Y is defined analogously. to denote a quantity that depends on 21..... polynomially. We also use the standard big-O notations: ( (.)) and 52(.) to only hide absolute positive constants. In addition, we use O and Ω to hide higher-order terms, e.g., ((log dlog log d) = Ö(log d) and Oddlog d) (d). Let Ab = min(a, b), [k] 21, k} for A: N. For vector v, denote by ||u||p == C |e; 1)/// the S norm. When p 2, we omit the subscript for simplicity. For matrix A, let |A]] and ||A|| be the spectral norm and Frobenius norm, respectively. We use Amax and Amin(.) (.) to denote the maximal and the minimal eigenvalue of a real symmetric matrix. For a vector w E IR and & R. we use Wggk E IR* to denote the first & coordinates of w and w>s to denote the write (Wek, w>k))
We use poly(
last R k coordinates of That is to say,
2 PRELIMINARIES
2.1 PROBLEM SETUP
Our aim is to learn the target function h IRdd IR. where Rdd is the input domain equipped with the standard normal distribution ;= NO 1a). We assume our target has a compositional structure, that is say, h =g g oop p for some g: IR + R and p 12d R.
Assumption 1. p is a degree i polynomial with le > 2, and g a degree q polynomial.
The degree of h is at most r := kg. We treat k, q as absolute constants, and hide constants that depend only on k, q using big-O notation. We require the following mild regularity condition on the coefficients of g.
Assumption 2. Denote g(z) Eosis giz We assume sup, 19.1 = 0(1.
Three Layer Network. Our learner is a three-layer neural network with a bottleneck layer and residual link. Let ma, m2 be the two hidden layer widths, and a (-), ()), be two activation func-
3

---

Published as a conference paper at ICLR 2024
(x)
Ax(
+
+b
Figure 1: Three-layer network with bottleneck layer and residual link, defined in (1).
tions. The network, denoted by he. formally defined as follows:
m2
ho(x) = (2),,v(z) + (a o2(aqu. v (2) + b) gu.e, (2) + Σ (o b.)
(1)
where a, c E 2 u, E RT and V E Rmixd xd. Here, the intermediate embedding Su.s.V is two-layer neural network with input and width mı. while the mapping Ju.s,V H ho is another two-layer neural network with input dimension 1, width m2, and a residual connection. We let 0 = (a, b, c, u, V) be aggregation of all the parameters. We remark that the bottleneck layer and residual connection are similar to those in the ResNet architecture (He et al., 2016), as well as architectures considered in prior theoretical work (Ren al., 2023; Allen-Zhu & Li, 2019; 2020a). See Figure I for diagram of the network architecture.
The parameters ((0) = (a ((0), c()), (0)) s(0), V(O)) are initialized as c()) = 0, ( = O.
(!0) rid N(o, 1/2), and of() mid (L/V() the sphere of radius 1/v2, where {uz (0) helmu] are the rows of ((0), Furthermore, we will assume b/0) riid
(0) -iid Unif{-1, a
where T5 a distribution with density (-))). We make the following assumption on uay:
Assumption 3. pss(1= Z (14) + P) for absolute constant p > 0, and Es-ja, [b*] 1
Remark 1. For example, we can choose To to be the Student's t-distribution with degree of free- dom larger than 8. Student's t-distribution has the probability density function (PDF) given by
r [++
1
1 (v+1)/22
(t((t)
VPTl
2/
where y is the number of degrees of freedom and I' is the gamma function.
Training Algorithm. The network (1) is trained via layer-wise gradient descent with sample split- ting. We generate two independent datasets D1, D2, each of which has n independent samples (xr h())) with 7. We denote L.D. (f) as the empirical square loss on D,, i.e
Lp. (0)
Σ (ho(x) (ho(z)-h(z)) h(z))² TED;
T
In our training algorithm, we first train via gradient descent for Ti steps on the empirical loss LD, (0), then train via gradient descent for T2 steps on (0). the whole training process, b. V are held fixed. The pseudocode for this training procedure is presented in Algorithm 1.
2.2 HERMITE POLYNOMIALS
Our main results depend the definition of the Hermite polynomials. We briefly introduce key properties the Hermite polynomials here, and defer further details to Appendix E. Definition 1 (ID Hermite polynomials). The k-th normalized probabilist's Hermite polynomial, hy : + R, is the degree polynomial defined as
(-1) d' Hg (x)
hx(x))
(2)
Vk!
(a(
where (=(x)= exp(-x² /2) 2 the density of the standard Gaussian.
4

---

Published as a conference paper at ICLR 2024
Algorithm Layer- Training Algorithm
Input: Initialization g(0), learning rate 7), 7/2 weight decay Eι, E2, time Ti T2.
fort ,,..., Ti do
u(t) u( 711 V LLD (0(t-1)) fu((t11) g(t) (a b d0), u(t) s(0), )(0)
end for
for 1
c() Ti 1, ,.... Ti T2 do
(t-1)) 122 V L.D (0((-1)) + (2elt-1)) ((0) (a!0) ((o), cll) u(ll) s00) v(o))
end for <- ((T++t))
Output: 9.
The first such Hermite polynomials are
22
23 -3
ho(z) = 1, =2,
V2 hx((2)
V6
Denote 3 = (0, 1) to be the standard Gaussian in 1D. A key fact is that the normalized Hermite polynomials form an orthonormal basis (8)); that is (h
The multidimensional analogs of the Hermite polynomials Hermite tensors:
Definition 2 (Hermite tensors). The k-th Hermite tensor in dimension d, Heg Rd - (TR is defined as
(-1)* V (z()) Hel(x)
My (z)
where (z(x)) = ( is the density of the d-dimensional standard Gaussian.
The Hermite tensors form orthonormal basis of (3()); that is, for any L3((),, one write the Hermite expansion
f(x) = E ((CU))
where C*(J) == (a) E (Rd)
1320
such, for any integer & > 0 we can define the projection operator Ps: I2((7) - 1,((7) onto the span of degree k Hermite polynomials as follows:
(Paf) = (e)))).
Furthermore, denote P&R- = Eusigh P: and Pek == Eogiek P, as the projection operators onto the span of Hermite polynomials with degree no more than ki and degree less than k, respectively.
3 MAIN RESULTS
Our goal is to show that the network defined in (1) trained via Algorithm can efficiently learn hierarchical polynomials the form g o p.
First, we consider restricted class of degree & polynomials for the hidden feature p. Consider with the following decomposition:
1 VL
Σ Exevi(o))
p(x)
Assumption 4. The feature p can be written in the form (3). We make the following additional assumptions on pc

---

Published as a conference paper at ICLR 2024
There is set of orthogonal vectors (Ui, (eet),at)) satisfying J, & k and | 1, such that (x(x) only depends on "...,,...., U.
For each i, Paats = th. Equivalently, t/s lies in the span of degree Hermite polynomials.
E [+)(()) = and E [p(x)² =
The X; balanced, i.e sup; D7rl (1), and L (d))
Remark 2. The first assumption tells us that each th depends on different rank & k subspace, all of which are orthogonal each other. As consequence of the rotation invariance of the Gaussian, the quantities (z) are thus independent when we regard z as random vector. The second assumption requires to be a degree & polynomial orthogonal to lower-degree polynomials, while the third is normalization condition. The final condition requires
be sufficiently spread out, and depend on many 1b, Our results can easily be extended to any = Wa(1), at the expense of a worse error floor.
Remark 3. Since P b/ th for each have PxP p. We can thus write p(x) (A, He (x)) for some A E (R2))4 There are two important classes of A which satisfy Assumption 4:
First, let A be orthogonally decomposable tensor
A
Σ Xeu 8k
where (vi, uz) = dig- Using identities for the Hermite polynomials (Appendix E. one can rewrite the feature p
(())
Σ Aclut Hea(()))
Σ Acha(u
(4)
VL
VI
P thus satisfies Assumption with J, for all assuming the regularity conditions hold. Next, we show that Assumption 4 is met when p is sum of sparse parities, i.e.,
Σ A 1i @ ... @ Ui
A
where (i Vi2.ja) = di, Sj ja- that case, feature p can be rewritten as
(vag. 2) )
p(x)
1
Ae via. Her(s))
Σ X
VL
VL
For example, taking L d/k and choosing Uh.s Ek(i-1)+j- the standard basis elements in IRd, the feature p becomes
pΦ#)
Vdtk (A12122 ... IK + Tk41 ... U22 + Aasktta -..Ia)))
and hence the name "sum of sparse parities." This feature satisfies Assumption 4 with J; = k for all assuming that the regularity conditions hold.
We next require the following mild assumptions on the link function and target h. The assumption on purely for technical convenience and can be achieved by a simple pre-processing step. The assumption g, in the single-index model literature (Arous et al., 2021), is referred to as y having an information exponent of
Assumption 5. Ex-y [h(=)] and (g~)((,,)) Ig (=)] (())
Finally, we make the following assumption on the activation functions σ1, U2. Assumption 6. We assume an is k degree polynomial, Denote (1) = Eogige 0,24, we further
assume sup; jo,1 = O(1) and [oyl = @(1). Also, set 02(2) = max{2, (z, 0}, i.e., the ReLU activation.
With assumptions in place, ready state main theorem.

---

Published as a conference paper at ICLR 2024
Theorem 1. Under the above assumptions, for any constant (0, 1), any mı > d and any n 7 d set M2 = d". Ti = -olyydd (d, m1, n.). T2 = poly (d, my, m2: n.), 7 = poly(d.m1,x) Then, for any absolute constant 8 (0, 1), with probability at least 1 å over the sampling of initialization and the sampling of training dataset D1, D2, the estimator 9 output by Algorithmm 1 satiifies
2m1 and &2 m2 polyid,m my,") E1 de-
1bg - h M|120000 O(d
Theorem states that Algorithm can learn the target h gop in n Odtt) samples, with widths my = {(d)), m2 = (1). Up to log factors, this is the same sample complexity as directly learning the feature p. On the other hand, kernel methods such as the NTK require n = (2(d")) samples to learn h, and are unable to take advantage of the underlying hierarchical structure.
A simple corollary of Theorem follows when = 2. In this case the feature is quadratic polynomial and can be expressed the following for some symmetric A RRdxd
p()) = (A, xx 1) = x Ar tr(A).
Taking tr(A) = 0, and noting that since A always has an eigendecomposition, Assumption 4 is equivalent to ||A] = and ||Allop O(1/vd), one obtains the following:
Corollary 1. Let h(x) g(x Ax) where tr tr(A) = 0. |lI 1, and || = ((1vVd) Then under the same setting of hyperparameters as Theorem 1, for any sample size n 7 d2+3a with probability at least - 8 over the initialization and data, the estimator satisfies
hba - >1 ÖO ()
Corollary states that Algorithm can learn the target g T Ax) in Odd)) samples, which matches the information-theorrtiialll theoretically optimal sample complexity. This improves over the sample complexity of the algorithm in Nichani et al. (2023) when is polynomial, which requires @(d*) samples. See Section 5.1 for discussion on why Algorithm is able to obtain this sample complexity improvement.
PROOF SKETCH
The proof of Theorem proceeds by analyzing each of the two stages of training. First, we show that after the first stage, the network learns to extract the hidden feature out (Section 4.1). Next, we show that during the second stage, the network learns the link function g (Section 4.2).
4.1 STAGE 1: FEATURE LEARNING
The first stage of training the feature learning stage. Here, the network learns to extract the degree k polynomial fcature so that the intermediate layer satisfies Ju.a.V (up to scaling constant).
At initialization, the network satisfies he Iu.s.V. Thus during the first stage of training, the network trains u to fit I to the target h. Since the activation d1 is degree k polynomial with O (1). we can indeed prove that at the end of the first stage Ju.s.V will learn to fit the best degree k polynomial approximation to h, Peu (Lemma 9). During the first stage the loss is convex u, and thus optimization and generalization can be handled via straightforward kernel arguments. The following lemma formalizes the above argument, and shows that at the end of the first stage the network learns to approximate Pεkh.
Lemma For any constant a (0, 1), any mi > tl+a and any n 2 dlo+3o, set Ti = 2m1 poly(n, mi, d)),), Th poly(n.m..d) and & g=+a. Then, for any absolute constant E (0,1), with probability at least 1 - 8/2 over the initialization V. and training data DD, we have
[ - -PEhillLA)) = (dd-)).
It thus suffices to analyze the quantity Pεth. Our key technical result, and main innovation of our paper, is Lemma 2. It shows that the term Pgsh is approximately equal to P h, and furthermore, up to scaling constant, Pah is approximately equal to the hidden feature p:
7

---

Published as a conference paper at ICLR 2024
Lemma 2. Under the previous assumptions. we have
Pkh (0.) /9 (=))))))))> -
O(d and |P<xhllea()) - 0dd
proof sketch of Lemma is deferred to Section 4.3, with the full proof in Appendix B.
Combining Lemma and Lemma 2, we obtain the performance after the first stage:
Corollary 2. Under the setting of hyperparameters in Theorem for any constants ô E (0, 1). with probability 6/2 over the initialization and the data D,, the network after time satisfies
ligerr) 1 /o' (=(2))))]l (=))) a)
Proofs for stage are deferred to Appendix C.
4.2 STAGE 2: LEARNING THE LINK FUNCTION
After the first stage of training. Vu,,.V is approximately equal to the true feature up to scaling constant. The second stage of training uses this feature to learn the link function g. Specifically, the second stage aims to fit the function g using the two-layer network 2 ト 2 + c c2 0z + b). Since only is trained during stage 2, the network is random feature model and the loss is convex in c.
Our main lemma for stage 2 shows that there exists c* with low norm such that the parameter vector ((())
0r = (a("), b(0),'', , (T:), 80) v(0)) satisfies hg- as h. Let ý be an arbitrary degree k polynomial satisfying - 1 (0.1) lo' (=(=)] satisfies this condition with high probability). The main lemma the following.
O(log dy (and recall that after stage Ju.s.V
Lemma 3. Let m d With probability at least 5/4 over the sampling of a, there exists some cº such that || = O(log di M)ttt)ddd and
L(00)
p(x) Σ co(o,b(z) + +.) h(x)
O(loo /
LP(7)
Since the regularized loss is strongly convex in c, GD converges linearly some â with (2( z (2()) and ||ål|2 $ || 112- Finally, we invoke standard kernel Rademacher arguments to show that, since the link function g is one-dimensional,
Õ(1) sample suffice for generalization in this stage. Combining everything yields Theorem 1. Proofs for stage 2 are deferred to Appendix D.
4.3 THE APPROXIMATE STEIN'S LEMMA
To conclude the full proof of Theorem 1, suffices to prove Lemma 2. Lemma can be inter- preted approximate version of Stein's 's lemma, generalizing the result in Nichani et al. (2023) to polynomials of degree lo > 2. To understand this intuition, we first recall Stein's 's lemma:
Lemma (Stein's 's Lemma). For any : + R and g E C²,, one has
E 12- [=g[)) E2 (=)).
Recall that the feature is of the form p(x) - (r)) Since each s/s depends only on the projection of x onto {.i. Us,,, }, and these vectors are orthonormal, the individual terms th (x) independent random variables. Furthermore they satisfy [[()(()) = 0 and ((ys(x)) Since L = (d), the Central Limit Theorem tells us that in the d + oo limit
Arh + N)
when the A are balanced. The distribution of the feature p is thus "close" to Gaussian. As consequence, one expects that
Evrr ((<))((()(())) R [._ = Ez~w(a,))l'
(5)
8

---

Published as a conference paper at ICLR 2024
Next, let q be another degree k polynomial such that || = 1 and (p, 1)()))))))))))))))()) = 0. For most q,
we can expect that (p. q) is approximately jointly Gaussian. In this case, p and q are approximately independent due to (p: 9) (v) = 0, and as a consequence
E E ( = Be-
0.
(6)
(5) and (6) imply that the degree & polynomial g o has maximum correlation with p. and thus Pa p) ~ (n).)ro.,))b'==)]
Similarly, if is degree < k polynomial, then since PkP p one has (p. (p.g) (2))) 0. Again, we can expect that p. q are approximately independent, which implies that (h, L2()) R2 0.
We remark that the preceding heuristic argument, and in partieular the claim that p and q are ap- proximately independent, is simply to provide intuition for Lemma 2. The full proof of Lemma 2, provided in Appendix B, proceeds by expanding the polynomial g o p into sums of products of monomials, and carefully analyzes the degree k projection of each of the terms.
5 DISCUSSION
5.1
COMPARISON TO NICHANI ET AL. (2023)
In the case where k and the feature quadratic, Corollary tells us that Algorithm I re- quires O(d²) samples to learn h, which matches the improvement over Nichani et al. (2023), which requires ((d') samples.
lower bound. This is an
The key to this sample complexity improvement is that our algorithm runs GD for many steps during the first stage to completely extract the feature p(z), whereas the first stage in Nichani et al. (2023) takes a single large gradient step, which can only weakly recover the true feature. Specifically, Nichani et al. (2023) considers three-layer neural networks of the form ha(r)) a 02 (Wo1 01 (Vx) + b), and shows that after the first large step of GD on the population loss, the network satisfies wt/ 01 (Vx) z d pxx).. As consequence, due to standard 1/vn concentration, n (2(d)) samples are needed concentrate this term and recover the true feature.
On the other hand, the first stage of Algorithm 1 directly fits the best degree 2 polynomial to the target. It thus suffices to uniformly concentrate the loss landscape, which only requires (cd)) samples as the learner is fitting a quadratic. Running GD for many steps is thus key to obtaining this optimal sample complexity. We remark that Nichani et al. (2023) handles a slightly larger class of link functions g (1-Lipschitz functions) and activations 01 (nonzero second Hermite coefficient).
5.2 LAYERWISE GRADIENT DESCENT ON LAYER NETWORKS
Algorithm takes advantage of the underlying hierarchical structure in h to learn in ((d*) samples. Regular kernel methods, however, cannot utilize this hierarchical structure, and thus require θ ((dha) samples to learn up to vanishing error. Each stage of Algorithm implements a kernel method: stage uses kernel regression to learn p ((dh) samples, while stage uses kernel regression to learn in Õ(1) samples. Crucially, however, our overall algorithm is not a kernel method, and can learn hierarchical functions with significantly improved sample complexity over naively using a single kernel method to learn the entire function. It is a fascinating question to understand which other tasks can be learned more efficiently via such layerwise GD. While Algorithm is layerwise, and thus amenable to analysis, it still reflects the ability of three-layer networks in practice to learn hierarchical targets; see Appendix A for experiments with more standard training procedures.
5.3 FUTURE WORK
In this work, we showed that three-layer neural networks are able to efficiently learn hierarchical polynomials of the form h g p. for large class of degree l polynomials direction is to understand whether our results can generalized to all degree k polynomials. We conjecture that our results should still hold as long p is homogeneous and close in distribution to a Gaussian, which should be true for more general tensors A. Additionally, the target functions consider depend only single hidden feature p. It is interesting to understand whether deep networks can efficiently learn targets that depend on multiple features, ie. of the form h(x) = g(pi (x)())),Pr((;: for some g : RR R.
An interesting

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
Zihao Wang is partially supported by the elite undergraduate training program of School of Math- ematical Sciences at Peking University. Eshaan Nichani acknowledges support from National Defense Science & Engineering Graduate Fellowship. Eshaan Nichani and Jason D. Lee acknowl- edge support of the ARO under MURI Award 1 NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994.
the Sloan Research Fellowship.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theoor, pp. 4782-4887. PMLR, 2022.
Emmanuel Abbe, Enric Boix Adserå, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-sadde dynamics. In The Thirty Sixth Annual Conference on Learn- ing Theory, pp. 2552-2623, PMLR, 2023.
Zeyuan Ahu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Ad- vances Neural Information Processing Systems, 32, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.0444 2020a.
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels?, 2020b.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
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 neural information processing systems, 32, 2019.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1):4788-488 2021.
Jimmy Ba. Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37942 2022.
Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks, 2020.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768-9783, 2022.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004. doi: 1010177CB0998061188441
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne,, and Qiao Zhang. JAX: composable transforaations of Python+NumPy programs, 2018. URL http:
//gi //github. com/ googl e/ jax.
Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang. Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances in Neural Information Processing Systems, 33:22134-22145, 2020.
10

---

Published as a conference paper at ICLR 2024
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory. pp. 5413-5452. PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two- layer neural networks, one (giant) step at time. arXiv preprint arXiv:2305, 18270, 2023.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes pwer-parameterized neural networks. arXiv preprint 1rXiv:1810002054 2018.
Weinan E, Chao Ma, and Lei Wu. A comparative analysis of optimization and generalization prop- erties of two-layer neural network and random feature models under gradient descent dynamics. jan 2020. doi: 1010077/11425--1--1622 URL
Science China Mathematics,
htttss://di.... .og. 100 10w11110022111
01 9 -
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz. and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029 1054, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Honglak Lee, Chaitanya Ekanadham, and Andrew Ng. Sparse deep belief net model for visual area v2. volume Vol 20, 01 2007.
Jachoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 33:15156-1511 2020.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference Learning Theory, pp. 3351-3418. PMLR, 2021.
Eshaan Nichani, Alex Damian, and Jason Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. arXiv preprint arXiv:2305.06988 2023.
Ryan O'Donnell. Analysis of boolean functions. Cambridge University Press, 2014.
Giuseppe Da Prato and Luciano Tubaro. Wick powers in stochastic pdes: an introduction. 2007.
URL https://////
/ttrps:://phh /api. semant csccolar....... www.llar.rrroooouuoeeee :ISIID::54911 55493217.
Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Feature learning neural networks and kernel machines that recursively learn features. arXiv preprint arXiv:2212. 13881, 2022.
Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. arXiv preprint arXiv:2304.01063 2023.
Itay Safran and Jason Lee. Optimization-base based separations for neural networks. In Conference on Learning Theory, pp. 3-64. PMLR, 2022.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of avrr-parameterized shallow neural networks. IEEE Transactions Information Theory, 65(2):742-769, 2018.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 10Xiv:101 2010.
Roman Vershynin. High- dinge-himmiional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018.
11

---

Published as a conference paper at ICLR 2024
Martin J. Wainwright. High-Dimensional Statisties: A Non-Asymptooic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019.
Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel, 2020.
Zhi-Qin John Xu. Frequency principle: Fourier analysis sheds light on deep neural networks. Communications in Computational Physics, 2865):1746-1767, jun 2020. doi: 10.4208/cicp. 2a-2020-00885 URL
org/ 10. 4208 w2208222ccccpppc ca-2020-0085.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11733 PMLR, 2021.
12

---

Published as a conference paper at ICLR 2024
Appendix
Experiments
13
Proof of Lemma
14
B.1 Results for General Features
14
B.2 Special Cases
19
C Proof of Lemma 1
20
C.1 Approximation
20
C.2 Empirical Performance
24
C,3 Uniform Generalization Bounds
26
Proof of Theorem
30
D.1 Approximation
D.2 Empirical Performance
D.3 Uniform Generalization Bounds
36
Technical Background
38
E.1 Hermite Polynomials
38
E.2
Gaussian Hypercontractivity
39
E.3
Polynomial Concentration
40
E,4
Uniform Generalization Bounds
40
E,5
Convex Optimization
E,6
Univariate Approximation
42
A
EXPERIMENTS
We empirically verify Theorem 1, and demonstrate that three-layer neural networks indeed learn hierarchical polynomials g by learning to extract the feature p.
Our experimental setup is as follows. The target feature is of the form h g p, p(z) Ede \,hz(a;), where the X; are drawn ii.d from { uniformly, and the link function is g(=) = C2z3, where C2 is a normalizing constant chosen so Rz (xx)2) 1. Our architecture the same ResNet-like architecture defined in (1), with activations 01 = 23 and a₂ = ReLU.. We additionally use the μP initialization (Yang & Hu, 2021). For chosen input dimension d and sample size n,, we choose hidden layer widths m, = d² and m2 = 1000. We optimize the empir- ical square loss convergence by simultaneously training all parameters (u, s, V, b, c) using the Adam optimizer. We then compute the test loss of the learned predictor, as well as the correlation between the "learndd feature" (defined to be 9u,a,v) and the "true feature" these test points.
In Figure/ 2, we plot both the test loss and feature correlation as function n, for d e (16, 24, 32, 40). We observe that, across varying values of depth, roughly d3 samples are needed to
13

---

Published as a conference paper at ICLR 2024
3,d mi 25,,,g mg-
3,d=24, =576,mg mg
Test Loss
Correlation of Fealures
Test Loss
Correlatin of Features
betwers uard
betwoee sast
X 3,d m2
1024,m:
1000
k 3,
40
1600,m m2
1000
benvee Buot ard
Tex
Figure 2: We train the ResNet architecture (1) to learn the hierarchical polynomial h = o p when the degree of p is = 3. We observe that the network learns the true feature as measured by the correlation between Ju.s.V and p (right panel of each figure). As consequence, the network learn h in d samples (left panel of each figure).
learn up to near zero test error. Additionally, we observe that as n grows past d², the correlation between the true feature and learned feature approaches 1. This demonstrates that the network is indeed performing feature learning, and learns to fit using Ju,a,V in order to learn the entire func- tion. Overall, this demonstrates that our high-level insight that the sample complexity of learning is equal to the sample complexity of and that three-layer neural networks implement the more efficient algorithm of learning to first extract p out of g p, holds in the more realistic setting where all parameters of the network are trained jointly.
Experimental Details. Our experiments were written in JAX (Bradbury et al., 2018) and run on single NVIDIA RTX A6000 GPU.
B PROOF OF LEMMA
B. RESULTS FOR GENERAL FEATURES
In this subsection, will consider the following feature class
p(x)
(x))
VI.
Recall our assumptions on Assumption 4. The feature p be written in the form (3). We make the following additional assumptions on p:
There is set of orthogonal vectors {Vi, heee))eu satisfying J, & k and |e 1, such that (z(x) only depends on T..... 1,
For each i, Patts = th.. Equivalently, tia lies in the span of degree Hermite polynomials.
E [1) = and E [px))2] = 1.
The X; balanced, i.e sup; |>, (1), and L Odd).
Next, recall that the link function g()) = Eosien 9:z" satisfies sup, |gi| = O(1) by Assumption 2. Denote = g Due to Assumption 5, we naturally have Poh = Ez-r [=x]]] = 0. Next, we will prove the following two Lemmas, which directly implies Lemma 2.
14

---

Published as a conference paper at ICLR 2024
Lemma 5. Under all the assumptions above, have
||Psh- Ez--((o.)) 19 (z)) pll (2))) =O(
Lemma 6. Under all the assumptions above, for any m E k 1 we have
|| (2(()) = O(L 112)
Proof of Lemma 5. Firstly, we will compute the Hermite degree i components of p(x)" w 3 2 From the definition of Pxx and multinomial expansion theorem, we know
Pa(p(z))) Twf2
1
Σ (a)
Σ
w!
1 U/i (x) Wry (1,
2q
2122.q.21
1,
w!
+
Pk.
1/
Ao thi (x)) ... (1,
L
21
Zq..
2.22.g.21
w.
by expanding (vz EESEEL (x) and computing the projection for each term. The key observation that leads to (7) is the following:
Lemma 7. Let Φ1, 02 E (1) be two functions such that 01 lies in the span of degree ky Hermite polynomials and t22 lies in the span of degree k2 Hermite polynomials. That is to say, PE, = 6; for 1,2
If $1, 62 only depend on the projection of x onto subspaces V1, V2 respectively. and V1, V/ are orthogonal to each other, i.e V V.V. = 0, then Pky 1kzzzz0102) = Φ1
Lemma 7 follows directly from the fact that the d-dimensional Hermite basis is formed from taking products of the 1-dimensional Hermite basis elements.
In the above expansion, if there are two indices i,, i2 each with exponent 1, then we get 4) (x) (/2 (z) I1 4h, (z:)= term. By Lemma 7, this term is polynomial with Hermite degree least 2k: Equivalently
Pk
Ui (2)vr, (x)
II bi, (x))² 0. 333
This is because (xr) only depends on " ".......................... " and ( (ele) are orthogonal vectors. Similarly, for terms of the form ehy (x) 11,22 ,h, (r)* we have that
Pk
Wi 2
II Vi, (x)3² = th, (x)Po
II 4h (x))3) /22
132
Altogether, this gives (7) above.
Let us firstly compute the Po terms in the above equation (7).
Case I. Firstly consider the case that w is odd and 2s + 1. Then have
w!
X Vi, (x) v1r, (z) = Σ w!
Σ
2 X1
X2,
21 2q
2,32g.44 18878
4,7i
Σ
w!
1 Ni (x) ..., (x)
21
2
8, 2.22,68,81 +2,
For the first term, we have
Po
Σ w!
2a Ao ...
.2 (x) .s. (x)²
= Σ
w!
X2
A.A2, (x)
...
Σ X2 ...
2
2
(8)
15

---

Published as a conference paper at ICLR 2024
For the second term, count the number of monomials to get
Po
Σ
w!
X5 ehe (2) 2/,, (2)
-
1,6
21
mq
w!
Σ
x x
EE | (x) ... 0r (z)
21 mq.
(9)
3/ 2
2q =
1,8986
Σ
X
2: 2.
x L
In the second inequality, use Gaussian hypercontractivityy Lemma 31. Combining equation (8) and (9) together, and noticing that
Σ A X2 Σ A ....., .
< SA Σ
X ... Xo_ S L
which can help us substitute Σ,
X? ... X2 for Σ
Es X3 ... ...AP, we can have
Σ
w!
Ur (z)
Sr (x)
L
21 21!...
2.32,9.81
w!
Σ 2
..... ...,
Σ
Aevr(())
(1+
L
2
where sup; K $ 1/L.
Case II. observe that
Secondly we will consider the case that w even and denote 2s. In that case, we
Σ
w!!
Uh (x)
Di, (x)
Zg.
2
LigHt
w!
1 N (z) ... (I,()))
21!... 2gg
By similar argument like equation (9),
sup 1EEL
Σ
S T
2. + - 21 2g!
Therefore, we have the following bound for the Po terms in our equation (7).
1 ... ao (z)) ... Sh (2))2) ))
w! 21!...
Σ Axda(x)P
Σ
Lw/2
2g!
5.22..
(KKv((z)
where sup, [K|] & 1/L.
Then let us compute the Ps terms. Firstly. we divide the monomials into two groups
w!
Σ
*n (1r) vi, (,, (x))
21!
zg!
s. 2.22.9.21 -
Σ
w!
X tu (x:) (2)
w!
X (x) ... ...vi,(1))
(x:2
Ai
21
2g
2n
16

---

Published as a conference paper at ICLR 2024
For the first group, we have the following
2
Pk
Σ
w!
2i (2) ... Vi (2))g)
21!...
2g
Ri 2,20 11
(7(7)
2
w!
N7. e/yi (zz) ... 1r, (2)
_1!... Zgg
2:322245..11 wiyy
17(7))
< (W.) (w/21-1
Σ
w!
81
1.5 t/i, (2))
Wi, (2))²
2
2,22 <w
2gg
L7(7))
3 2
In the second equality we use Gaussian hypercontractivity, Lemma 31. For the second group, we have that
(((()) (x) - Σ Pke
/2
Pk
w!
b
(x)
w!
2 No (0) ... U2 (x)
2
81
(2(7))
12(7)
w!!
2
Σ Σ
(kk (A ... (T), vr Px (A), .A, (,, (2()))
2
11
w!!
(Pa (A ... ".) Ps (A), .A,
13) ... (0..,())))))) (,))))
2
hnch (1) 7f
w!!
s² L 28-1
sup P& (A2 ... (x) ... v (z) (1)
2
From the second line to the third line, we use the fact that if Nún Ps (X,, (x)?.... vi, (x) and Pa (x2 mean-zero random variables. Also, the third line to the fourth line is just counting the number of pairs of tuples with nonempty intersections. The fourth line to the fifth line is using gaussian hypercontractivity, Lemma 31, to bound the moments.
Ø, then ... ... ,) (x)) ... ty (z)²) are two independent
In word, we have derived for any k > 2, and any that
w!
Ps
th (z) ... the (
0(2-1/2)
I L
21!
Zqq
2:22.0.21
m.
12(())
Sum up all the derivations above, and we get the following conclusion.
Lemma 8, Given l > 2,
When w = 2s + 1 with > 1, we have
Pk(p(z))m)) L
w!
Σx p(x)
= O(L-1/2)
19(7))
When w = 2s with s > 1, we have
17

---

Published as a conference paper at ICLR 2024
Recall our g(2) = Eosiegiet. After the projection, the feature that get is approximately [. G (2x 1)1g20+1 (Eis .2.....P.) xP.) p. Precisely speaking, we have
Pkh Σ 20 (2s 1)1024+1
Σ Ae..... EAR......) ig
(10)
12(N)
Let's recall E; X² L, so that informally speaking, we expect p(z) ~ (O, 1) in a limiting sense due central limit theorem when L is large and X; somehow balanced. Again, from the main text, it is tempting to conjecture some kind of approximated Stein's Lemma like
Pa p) 12- [z)(
Now we will verify this is indeed right. In our case, the derivative of y is is (z) = + + 99g==-1, and we can compute that (v,1( [s) E. 920+1 (2s + 1)!!. Furthermore, we have
L
Σ X
1I, s!
E9,, ...2? ig
And as a direct consequence, we have
2 (2s + 1)lg2++1
Σ 22 .... 13
-2s + +
+ O(L
Simply plugging the above equation in equation (10), we get final result.
Proof of Lemma 6. Firstly, we compute the hermite degree m components of p(z)", w 2 2. From the definition of Pm and multinomial theorem, we know
(m) = L Pm
Σ 4 =
wo!
(1)= (2))²
21
2q
by expanding ( (Eisis Antu(r)) and computing the projection for each term. In the above expansion, if there one index i with exponent 1, then we get a vh,, (ar) II 122 Vi, (x) term. By Lemma 7, this term a polynomial with Hermite degree least As result,
Pm
thi, (x) II (1)) O.
J22
This is because (z)) only depends on v. v. x........ "r and (Ve.sheeLLIJEE are orthogonal vectors. Firstly, notice that
w!!
Σ
21
A (x) ... (,)
21! !!!...
su 2.22.9.21 t +3 =
Σ
w!
A1
X
wi (x:)
Wi (x)= +
2a A
w!
X R,,
...Vi,(.7) (x)) (
2g
2,32 2.2222gmmm..m
w
18

---

Published as a conference paper at ICLR 2024
For the first term, we have the following estimation
Pm
Σ
Ui, (z)
Vi (z)
21
21 2132,20510,,
12(7)
2
Σ
w!
rbe (30)
1/4, (z)
21 !!!...
b1
20, 10,3
12((7)
d 1w/21-1
Σ
w!
N t/ (z) ... vi. (z)
21
Zq
2 20 21 W..yy
(P(T)
S 88811/21-2
From the third line to the fourth line we use Gaussian hypercontractivity, Lemma Appendix E.2 to bound the high order moments of hermite polynomials. And for the second term, we only need to consider the case that w 2s is even. In that case,
2
P
wl
( (1))
w! E
2 M
X2 (z) Us (x)3
2
LP(G)
1.(7)
w!!
Pm .2. (x) U., (x) Pm (A) (x) U5,5 (x) (2)))) (2
2
41
w!!
(Pm .A, (x)² (x))) Pm (2), (x) 33. 53) (x) 17
2
S sup Pm (A2, 12, (.)...... (x)))) d
From the second line to the third line, we the fact that if (in Pm (A ... X² vi, (z)²... th; (z)²) and Pm. (X2, ... X3 thi (z))² ... Uú (z) are two independent mean-zero random variables. From the third line to the fourth line, we are just counting the number of pairs of tuples with nonempty intersection which is O(ds-1)).
Ø, then
In a word, we have derived that
L m
1
Σ
w!
dh (z) ... Uie (z)x
(LL-1/2)
21
2132,4,21
w,yy
LP(1)
Write g(2) Euciey g and sum over all the terms, and we get the desired result.
B.2 SPECIAL CASES
Orthogonal Decomposable Tensors. and A is an orthogonal decomposable tensor
Firstly, we will consider the case that p(x) = (A, He)(x))
A
Σ Aeu
VL
where (v,, uz) = 8zj. Using identities for the Hermite polynomials (Appendix E. one can rewrite the feature as
L Extrueep x))
p(x)
Σ Au Hee(z))) = VT
VI
This kind of feature satisfies Assumption 4 with J, = 1 for all i, if we further assume the regularity conditions sup, [aj] = O(1) and Σ, X L.
19

---

Published as a conference paper at ICLR 2024
Sum of Sparse Parities.
Secondly, we will consider the case that
Σ A 1i 80 ... @ Vi.k
iigin Jiiza In that case, feature be rewritten where (Vi Vb2.d2?
p(z) VL
A (or, @ ... ....8.............. u2 Hen Hen(
Σ
A
II
(ur.s x)
VT
This kind of feature also satisfies Assumption 4 with J4 & for all i, if we further assume the regularity conditions sup; |Ad] = (1) and Σ, X2
For concrete example, when vi.j Ex(i-1)+) and L d/k.
( vd/k (A14132.... ... 3s + + Ad .-... ...
and hence the name "sum of sparse parities".
C PROOF OF LEMMA
The goal in this appendix is to prove Lemma 1, which is restated below:
Lemma 1. For any constant a E (0, 1), any mi 2 d and any n > dl+30, set T₁ = poly(n, m1, d), 1n poly(n.mu,.d) .) and & = 2m1 with probability at least 1 6/2 over the initialization V, and training data D1, we have
Then, for any absolute constant E (0, 1),
2.1a.
[ PechillE = O)
Proof Outline. Throughout the first stage of Algorithm 1, remains at 0. Consequently, during this stage, the network is given by
Ju.s. (x) u (V
where di is degree k polynomial. Given that V, is kept constant and only is trained, the network equivalent to random feature model with the random feature (V(r s).
The proof proceeds in three steps:
First, we show that there exists u² such that v approximates Pah, the degree k com- ponent of the target.
Next, leverage strong convexity of the empirical loss minimization problem show that GD can find an approximate global minimizer in polynomial time.
Finally, we invoke kernel Rademacher complexity argument to bound the test perfor- mance.
In this section, we may use a--) to refer (1(,), and m to refer my due to notation simplicity.
C.I APPROXIMATION
First, we show that when is a degree polynomial, the random feature model can and only can approximate the degree & k part of the target function.
Lemma 9, For any R", we have the following equality for any function h E L2
19 All 1lgg...v Pechil ||P - 1al---hhi 11ll2
Remark From Lemma 9, we can see when we try to approximate h using Ju.s.V, we are actually trying our best to approximate Pgah. That is to say.
argmin Il hilliatn)) argmin Il Pe
20

---

Published as a conference paper at ICLR 2024
Proof. By direct computation, have
2
11 - 21))
T VoooV +8) - Σ (H Hez(z))
17(7))
T VTo(V 5) E He)(z))
+
Σ
(1z, Hey(()
jEk
L2(x)
s3k+1
LP(7)
1l Pst Pexlll + || P (12))))))))))) 12(x))
(11)
where H, = Es (e) Here we use the hermite expansion which we state in Appendix E.I.
We next show that Pah h can be expressed by an wintt network by the following three lem- mas.
Lemma 10. There exists f gd-1 - R such that B.lf(e)ha(v T x)] (P)h)(z) and =
where u obeys the uniform distribution on Sd-1
Proof. Recall that (Pth)(x) can be represented as (A, Hex for some symmetric tensor (R4) Furthermore, observing that
By [f(o)he(u x)] = (E. [f(v)) Hex(z))
by Lemma 28, it suffices to solve for u(-) such that E, [(t) = A.
Let Vec: (RRd)0 gd be the unfolding operator. We claim that one solution for f is
f(v) Ve)( (E,, Vec(u Ve Vec(A).
First, by Corollary 42 in Damian al. (2022), we have
TT Exnry [Veo(= @k Vec(z T] 2 k!ll (Rd))
(12)
where (R is the projection operator onto symmetric & tensors. Since A is symmetric, we indeed see that
Vec (B, [f()) = E, Vcc(u Vk) T (E, Vec(2 Vc Vec(A) = Vee(A).
Plugging this back to E. [()]]] and applying the Cauchy inequality, we get
E. [f())]] & Amax ((F, Vec( Vec(u(8)7)) || Vec(
(13)
Therefore, to estimate the I norm of f() we only need look at the spectrum of the matrix above.
For X ~ (o Ia), it is clear that YZ shares the same distribution with X, where Y ~ x(d) and ~ Uni and Y, are independent. Therefore,
Ex (Ve((X Vee(X T| = Ey 2 Ez VVec(z*h V < d Ez [Vec 2h Vec( ]
due to Lemma 44 in Amax (EEx [Vec(X V TT)"
Damian et al. (2022).
Furthermore, we get < k by equation (12). Plugging this back to equation
(13), will have
E, [/(x]]]
d (eo( $ dd || E [(Pa))(=))] 0(1).
where we used the fact that Vee( Vec(A)]lB
21

---

Published as a conference paper at ICLR 2024
Lemma 11. Ler ~ (O 1). Then, there exists IR R with E, [w({)) O(1) and
2 2
Eg
w(s)
= hx(2).
V2
Proof.
One has the following Hermite addition formula:
h4
+ 8
2-21/2
22 1/2
Σ
()
1/2
h (
2
==
Thus writing (2) = Eizo (2)) we have
1/2
+
42 1/22/1
h (+)(())
V2
130 j=U
022 2/2-1/2 (1)
1/2
hz(z))
hx-)()).
130
i=1
Define wo. ..... wh recursively by
Ck 2²/2
1/2
11-
k + 5) 1/2
" ck 22/2
Σ
20t1-,2- (k+1-
As consequence, for j > 1, we have
2 + 1)
1/2
W..
Therefore for all 0 S i < k - 1, we have
k j
(i+j 1/2
city 2+4)2(6+51/2
(6+1)/2
Wi
1=0
1/2
c2 1/2
wi-j
Setting wss) = Ei= (s),(s), we thus have that
V2 (2t+) )1
(
11)
1/2
E, w(s)
(z()))
c2
1/2
W
j30
R-1
k
(
1/2
2 k/2 -
Σ
h2(2))
Σ
22
21221/2
wi-j
520
i=)
14 (=),
as desired. Since we regard constant, and we have sup, |c|| ((1) and c O(1) due to Assumption 6, the norm bound follows.
Lemma 12. There exists Sd-1 R such that
1 TTT +x*+ = (P)h)(x) and Eyia [u(u, = O(ah)
Ey,9
8)0
/2
22

---

Published as a conference paper at ICLR 2024
Proof. By Lemma 11, we get Es (s)) a degree polynomial. Substitute with o' I, and then use Lemma 10, we have
2+t v2 hy (z) for some Es [wo(s)]] = O(1) and w(.)
Ess (())o(()) (u (vx + (
E. [f(e)he(v (Ta)] (Pxh)(=)
V2
Set u(uu,) f(u)w(s). We next bound the L norm of u(u, s) by the independence between v and
E [u(v, s) IE [(()) = E [S(v)]] [wo() d
Remark 5. In the above lemma, our feature is (TTats) with uniformly sampled from the unit
sphere and sampled from (0, 1). This is equivalent with our feature ( xo"'x+-0 s) in the main text, with uniformly sampled from the sphere of radius 1 and 8 sampled from (0,11/2) We will use
vI
the a o+t
formulation in the remainder of the section without loss of generality.
Next, we show that we can use this infinite width construction to construct a finite-width network that approximates Psh.
Lemma 13. For any absolute constant ô E: (0, 1) and m N+, with probability least 1 8/8 over the sampling of 8, there exists ut such that
a Pechil2a()) Om and ||u 112 Om
Remark 6. Due to Lemma 2 and utilizing the lemma above, have
a PEnhil S d + m de
Proof of Lemma 13. We use Monte Carlo sampling to help us construct the u². Let u(.:)) be the function from Lemma 12, so that (Pah)(z) Ev., ju(us)or e ( ). We sample θ {1,, 8: )) i.i.d. and set of = 1 m u(u,, 8:). As such, one has that
m ulu;,
1 x + 83
EgEz 19, - (Pah)(z))2 = E, Ee
83)0
- (Pah))
V2
j=1
0 x + 8j
))
(u(e), sija
1jj ooT x + 81
Eg
Ee
u(u,, sy)0
m²
v2
v2
3.1m1
27
0..Ir + 8j
E.B B.rBuyy
22
j=1
1
Ey,s [f(v) S2h]]
m
1
x
Ea,s [u(u,)) s)
m
(14)
and
Ee
Σ u(u,,
Ev,s [u(v, s)
Therefore, from Markov inequality, we can derive that for any constant K >
we have
IP
E
Pa4l2 > O(1) e [u(u, s)31 < K
(15)
23

---

Published as a conference paper at ICLR 2024
and
m
Peo
83) KE [u(u, s)3] ==1
E K
for some (1). Setting 1/K = 8/16, plugging in the bound on E [u(u, s)²] from Lemma 12 and noting that ||=112 m² >x= (V 8,)² yields the desired result.
el = ((1) K [u(e, s) for notation simplicity where the ((1) is from equation (15). Since we see K as absolute constants, we have é1 = /m).
Throughout the remainder of this section, we el
C.2 EMPIRICAL PERFORMANCE
Next, we focus on the concentration over the population loss given by
(uu) |lgu.e. Al] (2())
evaluated the point =", which is defined in Lemma 13. Our primary tool for this concentration is Corollary 3. For the sake of notational clarity, let define Í(u)
Ei= (Iu,s,V (are) - h(xx))² to represent the empirical loss based on the initial dataset Dj. Lemma 14. Under the setup and the results Lemma 13, we will have with probability at least 1 1- 5/4,
L(u))L((y")
x Vn
Proof. By Corollary 3, for any B >> O, have
L(u)) -(w)))
[E(w")
Var ((gur, -
-))))
< 2exxp (-0(1) min(3? ,in()??
Moreover,
Var ( & Ey
S (En ( S (EL + Ex [h(x)²1)² < 1,
where the second inequality relies on Gaussian hypercontractivity (Lemma 31). and the final step sets m > dk+a so that 6 $ Plugging this back and choosing some B ((1) finishes the proof.
Observe that during the first stage of Algorithm 1, solving the following minimization prob- lem:
1 5 ||u|?
min L(w)
(16)
Since this problem strongly convex and smooth, plain GD can converge to an approximate mini- mizer exponentially fast. The next lemma bounds the time needed to obtain a small empirical loss:
2m Lemma 15. Set क For any e2 (0, 1). let Ti 2 m(log m) log(m/e2). Then, when m,n n
are larger than some absolute constant, with probability at least 35/8, the predictor == (T)) satisfies
i(0) < + //- -h--Pen PePhhlll (26) +0 - 0(1) + E
and || < ghta m
24

---

Published as a conference paper at ICLR 2024
Proof. If u is an ég-minimizer of (16), then we have
1
L(ū) + 56 || < L(u")
E 11 t < L(d") + 25 11"'| + +11)
+ E2
Vn
By choosing E
2m g=+..
we gel
1 +
dkta ||اه Se 1|h Pexh (2(x) Vn
At the same time, we will also have
1 Vni
L(a) < c1 + +(d
1|A 1 PEah 11--salll)) ((1) + E2
thus suffices to analyze the optimization problem (16).
Clearly, this convex optimization problem is at least 2-strongly convex. To estimate the time com- plexity, also need estimate the smoothness of our optimization objective.
Lemma 16. With probability least 1 O(1/m).
-
- vi(uz) $ m(logm)* 1 w2ll
Proof. We calculate the gradient out
vi(u)
1
2
u
Vx; +
h(xx.)
Vx; + v2
n
2
and then bound the Lipschitz constant of the gradient
vi(un) vi(22)
=
2
U1
u2 o
V2 ('/2
Vx; +
10
Vx; +
V22
1=1
2
Vz; + s
12)
|/u1
-wall
2
Using Corollary 3, we have the following concentration inequality for any
2
27
2)
1
+
83
x + 83
1 + 8j
< 2e -(1)3//
P
Ex
3
Var
71
2
v2
v2
1=1
27
(
2
Furthermore, estimating Ey
a
Var
and doing union bound over
all Uj we get the following inequality with probability least 1 2me 0(1)at/K
)[
8
S 1 B nn
Σ (1 +
2
By Corollary 3 agaan,, we can concentrate at least 1 - -
m
(1 + 22k) and get the following with probability
Σ
(1+ 2b
S 1
j=1
In that case, choose 3 =((1) m for some large O(1) and the lemma is proved.
25

---

Published as a conference paper at ICLR 2024
Having derived the above Lemma, using Lemma 36 in Appendix E.5, can choose the learning rate 7/1 m(log (I) and have
(
< (
(1))mlo m)* )
Uope 1|
Hopt
where Hopt is the unique optimal solution for that optimization problem. In addition, in order bound the empirical performance, we also need upper bound the gradient.
2
Vx; +
)||
Vx, + ++) h(xs))
sup |
vi(u)= 2u
< 2R
V2
1
v2
Viy +
0
(Ving + 8 >)""
$ 2R
h((s))
[uul]
V2
v2
< 2R
(1 R)
a
Vaz +
+ +h(z:)?)
v2
2
$ (1 + Rlogmmmmm)
h(1;)²
with probability at least 1 - ((1/m). In order to bound following for any >
E: h((x;)², by Corollary 3, we have the
P Σ h(x;)² - Ey h(x) > Var((((x))
< 2e -(1)at
vn
B = (11)(lgnn) n with some large ((),, with probability at least 1 1/n, we have 1 ". - h(x;) 2 1. In that case, we have
Therefore, by choosing
((u()) | 12 (uopt) + || 11? +
vi(u) + 2ul| |1(0) Mopt
sup |
Since ||uopel| = =1),, sup u|l uf
Vi(u)= + 2n = Omm[oo m)*), ), if we want
vi(u)= + 2u
|100) - uope 6
sup ||a|| Mop!l
it is sufficient to have Ti 2 m(log m log(m/e2)
C.3 UNIFORM GENERALIZATION BOUNDS
To conclude, need do union bound over for our population loss I1g (>)) We first consider truncated version of population loss, which allows us invoke standard Rademacher complexity generalization bounds. We conclude by properly handling the truncation.
Proof of Lemma 1. Let us denote e, (z, = (2- A². Via standard Rademacher complexity generalization bounds, detailed in Lemmas 33, 34 and 35, recall that see 8 as absolute constant, when m. n, d are larger than some absolute constant, we have that with probability at least 5/16
Σ e, (o....v(Z)) h(z:)) Ex ( hxx)] 2=1
S 2 Rad, (F) + 72
sup |M
1
< 4T Rad,,()) + -²
m + - $ MT
1
26

---

Published as a conference paper at ICLR 2024
where G (ou...V ||u|| < Mu) and F = [L (.) |al|| & Mu)- The first step just standard uniform generalization bounds for bounded function class. The second step is via contraction lemma to compute the Rademacher complexity, and the third step is a direct calculation. So, by that bound, we can see Eg [er (ga,s, v (ar), h(x))] is well controlled for moderate large т. Combining this with Lemma 15, with probability 76/16, have
Ex ( (x), - ||h -PPPhlllaa & STaa m + V + a=
+ vn
Dealing with the Truncation. Based on the above arguments, bound the L generalization error, it suffices to control the quantity
Ex ((nnav(a) 1 (19)..
This is done in the following lemma, whose proof is deferred Appendix C.3.1 Lemma 17. With probability at least 1 - 6/32, for any T Z 1, we have
Ex ((00...(())-(())))) 1 (s) S -(1)
Altogether, when m, n, d are larger than some absolute constant, with probability at least 1 - -022. have the following inequality
I1aa... -1l126) - 11/- - Pextilli2())
< E, (e.......v h(x)]] - ||h -PeuhllEsm)) + E. ((aas)-(()-((())))) (10)
/ where we recall 6 ddt).
5 MM uu v m + I - aa +
1 - + + exp (-0(1)72(r)
Vn
For any a E (0, select 62 = a Clearly we have Ti = poly(n, m, d) and 7h poly(a.m.d) in that case. Recall that we have chosen the width m > dhta the sample size > + and choose the truncation level to be (log d and M2 θ yields
Plugging those in 6² 1a
11 - PEx & llgā, - 11())) 11 - Pex 1am + O(1d))
S (log d2 + (log d)" d'" 4/222/33///
O)
as desired.
C.3.1 PROOF OF LEMMA 17
Proof of Lemma 17. We will first Cauchy inequality, then estimate the moments.
(Ez ((0.... - 1(x))²) 1 11n.)
1) 2 < E. [( h(x)) PP (()),. h(x)) > TT)
$ (Ex ( + +B [h(x)*]) P (19g,.. - [((x)1 > T)
E (x) +I [Mxx){?]
P (1ya.v.(x) - h((x)) > 7)
(17)
The last step is by Gaussian hypercontractivityy Lemma 31. Recall (,,,,V()) TTo
Vits /2
Notice that
T7
Vx + s
Vx + 8 22
E, Ez (x) = TE, 0
u
(18)
/2
27

---

Published as a conference paper at ICLR 2024
Therefore, we just need to give a tight bound for a Ex
Vz+a V2
o
( Vutt
v2
For notation
simplicity, in this proof, we will temporarily denote 2, = a Eg [2
Veite
Z =
Vz+t v2
:=
v2
Noticing that we have
n Σ (xi)
2 n
(g.,) - Mxx))) -
2
n Σ
h(az)²
n
with probability at least 5/64, due to the small training loss and some standard concentration for 1 >, h. That to say,
1 D
T
Z
Zj)
Ia v $
n 1m
Next, we bound the difference between ( Ei=1 Z,Z,) á and Σá. To this end, we or-
thogonally decompose Σ as Σ = K OK, where O is diagonal matrix and K is or- thogonal matrix. Write O diag(7 (), ..... nt. ......0. 0} for some integer
rank(E), where > 0 for E [t]. Notice that 01/2 diag(1) 1/2 ne 1/2 .......), 0}, and we formally denote O-1/2 diag(>ī .....: 7t -1/2 .......0 0). Due to the fact that E, [KZZ] K = O, we know KZ lies in the span of ex) Therefore, we have
T
1
E
à = à T 01/2 TT
2 0 1/2 KKZZZ,KT--H
(1, .) 0 Kü
i=1
i=1
K 0-122 -(t
< à
Σú
1/2 K7
n2 i=1
Denote W; := 0-1/2K2, Z, and W = 0010011222
We see that the second moment of W is equal It. That is to say, Wa is isotropic. Next, we
to identity matrix in f dimensions: Ez [W W will bound the following operator norm
gt
Ie
))
Σ
W Wen. W
Et. 4e
by the following concentration lemma.
Lemma 18. Let W ((x) E Ih be random vector which is a function of 7. Assume for each i [m), the i-th coordinate W, is k degree polynomial w.r.t Also assume E, [WWT] I. Let Wi...., Wn be i.i.d. generated samples. Then with probability at least 1 - 6/64, we have
maxicjem (s vnl V / m(logn)
where W = (W,,..., WT and 8j is the singular value.
Proof. For any v Var(() we have the following estimation for the tail probability
P maxigisn 2 > 2 m)
nP (II 2 > 2 a m
(II 1 E,, ||/l|]] > 22)
< nP
1/k
< 2n exp -0(1)
28

---

Published as a conference paper at ICLR 2024
due to polynomial concentration, Corollory 3, where
m
Var(|W]]3> < Ex
m
E
[
S m
E [W?]
2
m²
Therefore,
estimate
maxxe [W
we
can
choose
truncation
level
(11g n**
Var ([])) m with large (1).
E maxigien
||31|"] 2 m(logn)* + Er maxisien [Wil maaxegs//lWWW]222e0000 wi* Var([|
1/k
to0 S mmlogn) e(1)(logn)* / Var([[W]]])
2exp -0(1)
+ log dz
/ ((1)logn
to0
$ Sm
exp(-O(1)E + (ogn)=*-1
S m)
We will use the above estimation and the following Lemma from Theorem 5.45, Vershynin (2010) to estimate the singular values of W.
Lemma 19. Let be N x matrix whose rows A; are independent isotropic random vectors in In Let E maxig N || Then
E max JEn
85(A)
N
x V m log min(N, n)
Therefore, combining that lemma and Markov inequality to gain a high probability bound, with probability least 1 - 6/64, we have
maxisjsm (s Vnl 2 mmllo m(log n
Applying Lemma 18 to Ws,, we have
a 1
log t(logn)
W WT
Ets - 4
n
n
with probability at least 1 5/44. Next, we give an upper bound over I, the rank of our kernel matrix Σ. Using the Hermite addition formula, we have
Vx +
Σ
(V=)
V2
3=0
where A, E Th some vector that only depends on a-)), and Plugging that in our Σ, have the following decomposition
Vxr +
(Vo+ 8 22
Vx
2)
Σ (V
k
E
a
Ex
22
1=
Σ
Ex [())(V)) A =
Σ Mi
=
1=0
==
For each 0 & < kt, we have
Ey(p, q) Aig ,
01) (Agnu )
29

---

Published as a conference paper at ICLR 2024
where A is the 1-th element of A,, and Ey q) is the (p.q) element of our matrix Σy. Therefore, define M2- A, ..... A u E Rd xm and we have Σj M2 Mj and thus rank(E,) < d'. Therefore, rank(_) & Σラ= rank (2n) S dh and $ d"
Therefore, we have
1
/ log (log ni*
de log d(log n)*
W
2
Eri
72
n
n
and
11 ú
1 a
11/
d log d(log n) Eú.
Z Σ
Eú
Σ
W
W
2
E7,
71
As consequence, we have
1 E Zizi
E [ = Ed au T
á ú
n
when d is larger than some absolute constant. Recall that Bg [M(x)²] back into equation (17), we have
0(1) and plug everything
RR
- )(())) 1 (oo 1)² ((ga,,,.vv)))
- M(x) > r)
Therefore, we only need to bound the IP (1)),, > r) by polynomial concentration. From Lemma 32, we get
P (19....(=) h(x)1 > BV Var(on,,, (x))))()))))))))))))))))))))){))
< <2ex((-((1)21n)
for any >>1 1. Furthermore, notice that Var (x) h(z)) Eg { S [g,,,,
+ [ $1 $1
-
which is from the arguments above. Thus, for every 2 1. we have
P ((ga,,.v M(x)) 7) < 2exp (-0(1)221)
and the proof is complete.
D PROOF OF THEOREM
At the end of the first stage, our learner is horr = Ja.... In the second stage our training algorithm, letting = Vá,s.V, the network becomes
ma
ho(x) = ((x)(())))))) + Σ (o2((ap(z) bb)
with a; by random and fixed and e; trainable. The network thus implements 1-D kernel regression over the new input ý in the second stage of our training algorithm.
By Corollary 2, with probability -8/2 we have
||ga...v - P = O d and ||Pah - Ez-N( Is (z))) 22)) = O(d
For notational convenience, the remainder of this section we let á be an arbitrary degree poly- nomial satisfying the following assumption:
Assumption 7. We have a k-degree polynomial p which satisfies
- Eg [g'(2)) (LAm,) = ((log d)
where a E (0, 1). Also, recall that we have assumed Eg~ (x)) = e(1) and we denote this quantity as Cg-
30

---

Published as a conference paper at ICLR 2024
To prove Theorem 1, condition the event that 9a.v satisfies this assumption, which occurs with probability 1 5/2.
In the following may use a-)) to denote x2(:), and use to refer m2. for notation simplicity. The proof strategy will be very similar with the proof in Appendix C. We begin by constructing a low- norm solution that obtains small loss. Next, show GD converges an approximate minimizer. We conclude by invoking Kernel Rademacher arguments to show generalization.
D.1 APPROXIMATION
Define ō(z) = o( á(, =). The target can thus be represented as ((cp(())) We will proceed using the following two steps to bound the approximation error in L?))).
Step Bound the difference between o " and g (Cap).
Step II. Using 1-D two-layer neural network to approximate the 1-D link function ğ.
For step we have the following simple Lemma.
Lemma 20. Under the assumptions above, l o p nIlE2(0)) ((log dy
Proof of Lemma 20. We have that
l o p- - go (Cop)IBE
(f(x))) - (Cop(z))*|
k=1
Ex (0(x)) C,p(x)) +
+ +.... + ((Cg(()))*
Ev [((xx) Cxp(z))4]BB E. - + +... (Cop(z)))
En [(p(x)) - C2p(x)) E. +000) +.... (C2
k=1
1 ( S (log d2-
where the fourth inequality and the fifth inequality are due to Lemma 31, Gaussian hypercontractiv. ity. We implicitly use lI (0)) (1) and [|C,pllze67) (1) in the fifth inequality, too.
Step II relies Lemma 3. which is restated below: Lemma 3. Let d With probability least 8/ over the sampling of a, b, there exists some oº such that ||o"|lo0 O(loo df) dfMotte and
L(0') = p(x) + Σ ca(aap())+ h(x)
= O 22222+2))+0 aa)
L2(3)
Proof of Lemma 3. We will firstly control the typical value of p. From Lemma 32, we have
P [[((x)) > BV Var(ū(x)) < 2eep (-0(1) min (2k,))2))
for any B > 0. That is say, when 8 1, with probability least 2e - have [K(z) S B. We implicitly use (1 = (1) in this argument to bound Var Vār()(z)).
Next, we will use Lemma 39 to give a representation for ğ in the bounded domain. There exists .()))) supported on {-1,1} [0, 2033 such that for any satisfying [p(x) CB,
E.g [u(a, (ap(()) + b)] 7((())) ((x)
31

---

Published as a conference paper at ICLR 2024
where U and has density (t()). Furthermore, recall that we have assumed (t) 2 (1+ + [d] P. and we have the following estimation sup,,6 [v(a, b)l = ((B
Next, we will do Monte Carlo sampling to approximate the target.
Ea, Ey
Σ
o(ai, (.)o((;;(())) + b.) - (T((())) - p(x)))
bia, + -b) p(x3)) )² 1200)203
< Bo, Ex
1
Σ
m
(19)
m
1
+
D Ex
oaa (a) (T((())) {(x)) 1 pio)scs
m
For the second term, we have
E,,b Ex
m Σ u(a,, (()(a(()) b,) (T(((x))
1
(0(((=))-((x)) -b()))
1 Σ u(as, (.))aa)))( b,) E., (al b)]
Ea,s Es
(20)
1 Ex En (v(a, b)o((p(()) + +)))
m
0 20) (E,r p(x)) + Eg b2) m
1 m
((300+22))
Here we implicitly use the fact that Eg b² = (1) which is from our assumptions on (t))). For the first term, by Cauchy inequality,
Ea, Eg
Σ 1=1
o(a by (1)(((()()))) + b.) - (T((())) - )(a)) 1 1(0)))008
m
m
ola;, (a. + +b) (T((())) - -(o))) ((((()) CB)
Bag
m
m
$ -(11)82/k
u(a, b))(a((() + (0((())))))))))))))))))))((()))) p(x)) -
E Ea,b,z
m
4 raa b.)o((,(())) +b,) + ((((())) -(()))
0(1)820k
m
Ea,d,
m
3 -(1))2
V BaN (v(a,b)a(apttl + b)) + 011)
$ Se- -(1))) g20+29
Here we implicitly use the fact that E, 6² O(1) which is again from assumptions on (t).).
We also use gaussian hypercontractivity, Lemma 31 to show (g((((())) - p(z)) = O(1). Since p(z:) is & degree polynomial with Gaussian input distribution, its higher order moments can be bounded by a polynomial of its second moment which is clearly ((1).
From the above arguments, already derive
Eg,b Ez (
+ bb) (T(((z)) -p(o))
- x(1)) + e -
) gapt2g
32

---

Published as a conference paper at ICLR 2024
Therefore, for any absolute constant 8 (0, 1), with probability least 1 8/4 over the sampling of the random features ag b, using Markov inequality, we have
m
v(a,, by) b.)o((;((())) b.) (T((()) {(oi) )(o))) 2
( m -
Ex
g2p+21
m
Combining this with our previous result, Lemma 20, with probability at least 1 - 8/4 over the sampling the random features, we can find the parameters e* the third layer with sup, |c = O /m), such that
-(07)- p(x) + Σ Go(a,(()) +b)) - h(zz)
1 +e - ((0spr))) ga (1ogd)
(2(7))
where 0r = (a) _,O0, , û, V Let us further set 8 = (log d) where ((1) is some large
absolute constant. Set d= In this case, will have
I(0*) $ (d=a + e log" (log d)2k(o-s) + (log (logd)/22 d) S (log 2
D.2 EMPIRICAL PERFORMANCE
Next will show the existence of good estimators in our empirical landscape. Firstly, we need to concentrate the landscape at the special point e" we constructed. With a little abuse of notations, denote the empirical version of the square loss as
2
L(0)
1
n2
p(z,)
m Σ co((p((), + b,) h(2z))
where we recall that zj D2 is newly generated data which is independent of D1. Lemma 21. With probability least 1 35/8 O(1)d we will have
1 L(0") < - ((logd)))tttt))) + O(log (1((oog)//2+22200 2)22++26000))))) Vn
Proof of Lemma 21. In the following, we compute the variance term.
((
E, (L((7) I(87))²
1
m (a(aa(())) (h(x) -(x()))p(()))
Var
n
E2
(a((((())+ (h(x)) -b)((()))))))))))))))))))))))))) ((x))))()))
17
2/
E (a(a)(() + Ez (x(x))) + E. $(x))
7
3
E2 ci (a,((x) + bb)* + +0(1))
m
71
1
+4
m Σ
1
(b) + E. b(x)
1
1
34p+44
1
b
n
m
33

---

Published as a conference paper at ICLR 2024
Here some technical arguments to bound 1 Σ :=1
b/.
We have
m
Es
D
84
Eg b
m
i=1
and
2
1
1 m
Pb
Σ b
Eg ba
21
Es
bf
Eg
8
Eg
8
778
Therefore, recall that Es 6 (11) based on assumption pb(t), we will have with probability 1 - ((1)d
== b1 S 1. In that case, we have
Ex (((0)-((07) [((0) - L(B*)]
2 2 g4p444
-
1
1 (log d) (
11
n
Therefore, by Markov inequality, we have least 1 5/8. In this case, we have
2 Va (log d) 2k(pt() with probability at
i(0*) x Tn (logd))2tpp)) (log (0gg))/++222(t0))g)
In the second stage of our training algorithm, we are doing the following minimization problem
min i(()) + 26 ||e]|2
via vanilla GD, where 0 = (o ((0), C, ú, v Since this problem is strongly convex and smooth, the optimization problem can easily solved by plain GD.
Lemma 22. Set &a = 2. For any (0, 1), let T2 Z m mlog(m/c). Then, when m, n, d are larger than some absolute constant, with probability least 76/16, the predictor (2) and i (a ((0), è, ú, V(0) satisfies
L(0) x vn
(log d) 2k(p+4) + (log 222
and
||a|| x - vn
flog d) 2k(ppt) +e (log 2)2g))////2ttt))ada-
Proof. For any given threshold E (0, 1), assuming is an minimizer the optimization problem, then will have
L + 2 E ||ब|? < =(0)) 20 1e 1|e
1
//2
(log d) 2k(p+)) +c+ E2) (log d) /2+22k(p+))))) d (log
1
vn
Plug & = 2 in, then we will have
1 i(0)) x
vn (log a) 20(0+0) + + (log (1ogd))/++2(0000))
and
||a|
1
(log dl 2k(ppt)) + (log 2)/2222k00t
Vn
It thus suffices to analyze the optimization problem.
Clearly, this convex optimization problem is at least 2-strongly convex. To estimate the time com- plexity, we also need to estimate the smoothness of our optimization objective.
Lemma 23. With probubility at least 1 0(1)d 2e 1/2
have
vi(a)) - vi(ez) S m
34

---

Published as a conference paper at ICLR 2024
Proof. We first calculate the gradient
VL( = Σ (K(zz) + c a(ap(z)) +) - -h(2)) a(op(z)) + b))
then bound the Lipschitz constant for the gradient
vila)-vile)
2 Σ -) +b)
[ - 22ll (1)(((1)))++0
j=1
2 72
< ||0 c2ll
(1 + b,)²
4m
< [
c2ll
p(z)) + _ b2
n
Here are some technical arguments to estimate
Σ,
b7.
We have
m
Es
6
Eg
6²
Eg Egb4 ba
771
m
and
2 E& b²
2
m m Σ
1
m
1 m
Pb
21
Es
b
Eg b²
E, 6
771
Therefore, recall that a and also E, b2 ((1) due to our assumption on (4),), we will have with probability 1 - O J. Σ 63 S Morcover, we can use Corollary 3 to con- centrate >3 {,,)), More concretely, we will have 1 )(,)) I with probability at least 1 - 2e (01))))) /22 since p(z)² is degree 21: polynomial and Var(f(z)²) x 1 via Gaussian hy- percontractivity, Lemma 31. Therefore, with probability at least 1 O 2e 10--2---1000 0(1)n'/2k have
we
vi(a)) vi(co) $1
Having derived the above Lemma, using Lemma 36 in Appendix E.5, we can choose the learning rate 1/1 (()) and have
(00)) Copt ||
1 ('<(1_emm))
[[ooeel
((m)
where Copt is the unique optimal solution for that optimization problem. Furthermore, we have the following
sup |
Vi(c)
+ 20
< sup ||
vi(e)| + 2R
2
[|α(ας(;;)) + b)ll (1ā(x;) M(x))) + R[(o((p(25) + + b)1l + 2R
2
(R (R++)]l((ap(ss) + +]]1 + (B(x,) - h(xz))²) 2R
j=l
2
(R+ (1++1)((m) +
Σ (p(z,) h(z;))² +
n
35

---

Published as a conference paper at ICLR 2024
with probability at least - 0(1)d=a 2e 2 The last inequality follows from the same gument in Lemma 23. Moreover, we can use Corollary 3 to concentrate E,((ar;) - h(x;))² More
concretely, we will have 3 ()(z)))))))) h(z,))² 2 with probability at least 2e @(1)=* 1/2
since (f(z) h(x))² degree 2r polynomial and Var ((f(z) h(z)) ) 2 via Gaussian hy- percontractivity, Lemma 31. Therefore, with probability at least 1 ((1)d 2e 8(1))//P have
Vi(c)+ ( 20 2 S 1)m
sup |
Utilizing that fact, we have
_(c(0) + 1|/0001/ < [(Cope) + ||Copul1 2
vi(c) + 20
[
sup ||
Copt
Since |[Copell 0()),
vi((e) + 2c
O(m), if we want
vi(c) + 2c
|000) - cope E2
sup ||||| ccllopel
is sufficient to have T2 2 mlog(m/ee).
In addition, for any truncation level > O, we will also have
1 n Σ ((hg(z)), < L(0) x n j=1
1 - (log 2f(k((tt) +et (log 220222220002
vn
which we will use later. Here we recall (xr,,) = (x 1ス3
D.3 UNIFORM GENERALIZATION BOUNDS
To conclude, we need uniform generalization bound over e for our population loss L( |/, /1l /1|2260) [P(7)) As in Appendix C, we bound the truncated loss via a Rademacher complexity argument, and deal with the truncation term later.
Proof of Theorem I. Recall that (x, = yx² ^.². From Lemma 33 and 34, with probability at least 1 - 8/32, we will have
4 (ho(T::) Ep (ho(z), < 47 Rad, + 7 OO()
sup M
1=1
where H = (hg ||l|| < Mc). Then we will compute Rad, (H).
Lemma 24. With probability at least 1 - ((1)d= over the sampling of a, b, we have
Rad, S Me
m
36

---

Published as a conference paper at ICLR 2024
Proof.
m
Rad, (71) = =Ex Rgg
sup
Σ E 5=1
Go (a.((1))) + b;)
|ell & M
1
m
Ry l8g
sup |ell & Mo i=1
C/
(ax(()))
71
1=1
M
Ep Eg
{5o(a,p(zz) +b)
72
2
M
E2 Be
(aa(ab(zz)) +b)
i=1
M
m
72
E2
(a(ap(13)) b,))
n
i=
j=
M
m
m Eg p(z)²
Σ by
vn
Here are some technical arguments to estimate
E, b2. We have
E,
Σ
63
EB
13
EB b²
i=1
and
2
2
m Σ
m 1 m
1 m
Pp
6
Eg 62
>1
Eg
Eg 6²
<
Eg
6a
1=1
Therefore, recall that m = do, and also Eg ba S 1 due to our assumption on (t),), we will have with probability 1 O(1)d~ Σm by 1. In that case, plugging that in, get our Lemma.
As consequence, with probability least 5/32 - O(1)d-
1 n Σ (+hh(zzs) - E. [(((h((), S AT.Mey n
m 72
sup ||l]| M.
Lastly, also need to deal with the truncation get I generalization bound. That is to say, need to bound
sup Eg [(ha(x) -
cell & M,
Lemma 25. We will have with probability at least 1 - O(1)d_
sup Eg [(ho(x)) - (mo))))1)o(( S (1 + m M!!
[e]] Me
Proof. By Cauchy inequality, we have
[E, [(ho(()) h(=r))²) I [ 1)² F. (()h()))(==))]]] ((ho(x) h(x)) > 7)
< (E, [1)(x))] + E. [1)x)*)) ((ho(z) - h(x)) 77)
(21)
37

---

Published as a conference paper at ICLR 2024
Recall that E2 h(x)** O(1). In addition, have
m
Ey [xx)
Er
Go Go (ax(())
& m
E (c:p((x))+ + 6:14
i=1
S m² Md
O(1) + _ Σ by S m² M' Me
if under the high probability event 1 E i== b S 1. Furthermore, we have
1 Eg (ho(x) h(z))*) ~ 74 (1 + m=
IP (1ho(x) - h(x)) > Plugging this back, we will have with probability at least 1- 0
Eg ((hg(x) ( (1o))))(ao( 2 72 (1 + m Mc)
sup
|ell M.
We now combine everything together. Let us choose d=a and > + and recall m = d= In that case, ||ě]| = ((log dy//222k(++))dd-) Therefore, when d larger than some constant that is only depending on r.p. a, we are allowed set Mc (log d) d) for some large O(1). In that case, we have
ll 1ll 111200>> & (log (2ggd)))((ttt))) d
= 4 d ((1
20 k/2 32/2 (log d) (
We will pick up our truncation level / In that case, for any E (0, 1), will have
llkg - h (1 O(log ( (dd-a)
TECHNICAL BACKGROUND
E. HERMITE POLYNOMIALS
Definition 3 (ID Hermite polynomials). The k-th normalized probabilist's Hermite polynomial, h2 : R R, is the degree k polynomial defined
(-1) (x)
h2(2)=
de
(22)
V
k!!
(x(
where (z(x)) exp(-22/2)/ is the density of the standard Gaussian.
The first such Hermite polynomials are
2² h2(2)
23 -3-32 32
(2) (2) 2, h2(22)
V2
V6
Denote (O, 1) to the standard Gaussian ID. key fact is that the normalized Hermite
polynomials form an orthonormal basis of (33); that is Ex-eshy (2) (x)) 5nn-
Given f E (9), denote by f(z) Ex be the Hermite expansion of f where
f = = [[())Aa(z)] V2 - / ((=))u(>)) &
the Hermite coefficient f. The following lemma will be useful, which can found Propo- sition 11.31 of O' O'Donelll (2014).
38

---

Published as a conference paper at ICLR 2024
Lemma 26. Given f.g (3),, we have for any u, gd-1 that
Exny [f(u' x)g((' T x)] = > fueu(uTuje
K=0
The multidimensional analog of the Hermite polynomials is Hermite tensors: Definition 4 (Hermite tensors). The k-th Hermite tensor in dimension d, Hes Rd (Re)O is defined as
(-1) (z)
He*(z)
Vk! (x)
where (1((x) -
the density the d-dimensional standard Gaussian.
The Hermite tensors form an orthonormal basis of (,)()): that is, for any f E {T)( one can write the Hermite expansion
f(x) Σ Heul)) k30
where C&(f) := Eeny (x)
We define the Hermite projection operator (Pa/)(x) == (C&(), (x)) Intuitively speaking. the operator P3 extracts out the degree part of a function when the input distribution standard Gaussian. Furthermore, denote P&* == Losigs P, and Pck operator onto the span Hermite polynomials with degree no more than k. and degree less than k. is clear that ||P || & [| for any f E (7))). This can be shown by simple Hermite expansion for f.
Losien Pj as the projection
The next lemma can be shown by direct verification.
Lemma 27. We have
1 Vk!
Hea(x)
E. Eany [(x + iz) iz)
Lemma 28. If |u]| 1. we have
hyuuaa x) (z) u
Proof.
1
(He)(z), @k)
Vk! (Exm [x + +1z) 12) un0
E
Eeny
[[u T x + i(u 2))²}
Ad!
13
BawD [(u] x + iz)* ho(u Tx).
h!!
E.2 GAUSSIAN HYPERCONTRACTIVVTT
By Holder's inequality, we have || X L. < X I for any random variable and any & q. The reverse inequality does not hold general, even up constant. However, for some measures like Gaussian, the reverse inequality will hold for some sufficiently nice functions like polynomials. The following lemma comes from Lemma 20 in Mei al. (2021).
Lemma 29. For any E N and (3) be degree é polynomial on IR where / is the standard Gaussian distribution, for any q > 2, we have
(E.- 2) < (g 1)a E.-. [f(2)
The next Lemma is also from Mei al. (2021) and is designed for uniform distribution on the sphere in d dimension.
39

---

Published as a conference paper at ICLR 2024
Lemma 30. For any l N and f (2) to be degree é polynomial, for any 2. have
(Ba-unir(ss-)) [)(=))) 2/9 < (q 1) (n--ni(((--))) [](=)*]
For the case where the imput distribution is standard Gaussian in d dimension, we shall use the next Lemma from Theorem 4.3, Prato & Tubaro (2007).
Lemma 31. For any i E N and f E (7)) to be degree polynomial, for any e 2, we have
Bxony [f())* & (I,((1) (Eenn [f(z)
where we use (1,e(1) to denote some universal constant that only depends on l.
E.3 POLYNOMIAL CONCENTRATION
In this subsection, we will introduce several Lemmas to control the deviation of random variables which polynomially depend on some Gaussian random variables. We will use slightly modified version of Lemma 30 from Damian et al. (2022).
Let g be polynomial of degree p and x (O,, Then there exists an absolute positive constant Cp depending only on p such that for any 8 > 0,
Lemma 32.
P ([lo(x) B[g(z)]] > 8v [ar((((z)) < 2xx (-Cp (62,
Consider the case that (#1,..., In) and 9(x) E, 2, ridd (o), R and _ddxn, Plug them into the above Lemma, and we get the following corollary.
Corollary 3. Let g be polynomial of degree p and xa NO, Ia). E [n]. Then there exists absolute positive constant Cp depending only on such that for any 0.
Σ Elg(z)]]
Vn /Va(g(())
< 2exp (-Cp min (62.
E.4 UNIFORM GENERALIZATION BOUNDS
Definition 5 (Rademacher complexity). The empirical Rademacher complexity of a function class F on finite samples is defined as
Rad (F) Eg
sup
Σ (x,)
(23)
fEF
where E1,E2.....5 are i.i.d Rademacher random variables: P(C: = 1) IP((: =-1) Let (F), = []) be the population Rademacher complexity.
Then we recall the uniform law of large number via Rademacher complexity, which can be found in Wainwright (2019, Theorem 4. 4.10).
Lemma 33. Assume that I ranges in [0, for all f E F. For any 2 1, for any ó E (0, w.p. at have
least 1 8 over the choice of the i.i.d. training S (X1,...; Xn),
1 ((x)) < Rad, + RV
1og(4/5)
sup
(24)
JEF
1=1
Then we recall the contraction Lemma in Vershynin (2018, Exercise 6.7.7) to compute Rademacher complexity.
Lemma 34 (Contraction Lemma). Let (p R R with i 1...., be B-Lispchitz continuous. Then,
1 Eg sup FEF 1=1
Σ
Eups o (xx) < (R),..F)
n
40

---

Published as a conference paper at ICLR 2024
Next, estimate the Rademacher complexity for random feature models. Denote Vz) u!
with i..dd sampled from the uniform distribution on the unit sphere, and i..d. ((0, 1) generated. Here a-)) is a le degree polynomial with ( coeffi-
Vauis V2 -i= Ugd
ztst
v2
cients. Denote our kernel function class
(uaav : || M..
Then we have the following lemma for the Rademacher complexity of G.
Lemma 35. With probability at least 1 2c 2e- 2e-0(1)mm .e have the following estimation for the Rademacher complexity of our function class G
Rad (9) S Muj
Proof.
1 n
Vii + V2
Radn(() = En,e
sup
-
E.u Tuuo a
DeEG TI i=1
1 Ex,
n
Vxa + 8
sup T 1oEg
S10
71
V2
1=1
Ma
Vxx + $ ( V2
Ez,ş
Σ
S1o
i=1
M
Vx, +
(25)
E E
Σ E,
/2
Ma
v. + $3 V2
E
Varg
Σ
E2o
1=1
M
E2,
1/ x + 8j
S Ma Muvmi
m 2
5m
(1 s²*)
V2
By Corollary 3, we can concentrate
m 2==1
(1 + ) and get
m
22 2
m
j=1
with probability at least - 2c -21mmmm 122
Plug that in and we get our final bound.
E.5
CONVEX OPTIMIZATION
Denote fxx as a C function defined in Rd. Assume that
There exists m such that f(r
2 m
||-|| convex.
[V)(x) f()il
< L || w||-
The following result standard and can found most convex optimization textbooks like Boyd & Vandenberghe (2004).
Lemma 36. There exists unique such that fxx)) inf ff fx). And if we start the point o and do gradient descent with learning rate n, if 1) & m+L. then will get
|| x 212 112 < k ||000 x²
where 2mL 11)
41

---

Published as a conference paper at ICLR 2024
E.6 UNIVARIATE APPROXIMATION
In this subsection, we use (2) to denote (zLU(z) and set A T 1.
Lemma 37. Let ~ Unif({-1,1)) and let have density ps(t). Then there exists v(a, b) supported on (-1,1) [A, 2A] such that for any |a|| A,
Ba)oo((ab)o(aa + b]]=1 and
sup a.b
(o(a 6) <
S (fo(t)dt
Proof. Let b(a,b) = cloE|A,24] where JA (p)dd
Then for |ar| & A.
24
Ba, b)o((x + +b] =
[)(xx + o + ())uu(t)dt
24
t(us(t)dt
1
Lemma 38. Let ~ Unif([-1,1}) and let have density ps(t). Then there exists v(a, b) supported on (-1,1) [A, 2A] such that for any |a|| A,
Ra,//u(a, b)x b)] = x
and sup a.b
(u(a,
$ SA (s(l)bb
Proof. Let b) = calbE[A.22] where
Then for |x| & A,
S2a ps(t)ds"
2A
[a,,/u(a, b)o + +b]] =
lo(x + o0-x + ())m(())
A
24
c.
Hb(t)dt
x
Lemma 39. Let a ~ Unif((-1,1)) and let b have density (s((). Let f IR - R be any C
function. Then there exists v(a, supported on (-1,1} 1} [0, 2A such that for any |o| A.
Ea,[V(a, ()x + b)] = f(x)
and
(x) (133
sup b(a a.b
sup
[2A (,(t) + infze(o.A] (t)
[
Proof. First consider v(a, b) Al (-bb. Then when > 0 we have the following equa- tion by integration by parts:
malt)
E (v(a (ax b)]
= ("((=))((==+++00((((0><+<<
= +
42

---

Published as a conference paper at ICLR 2024
where C1 = In addition when x 0,
f(A) and C2 f'(0) - f(())
En, Ea,blu(a, ()x + 6)]
= -10
-
+(A)+ f()) xf)
= C22
so this equality is true for all æ. We can use the previous two lemmas to subtract the C1 + C2r term. That is say, we can
1EAA,221 + 1 v(a b) = -1 Sa tuo(t)dt bE[A,22] C2 [a (s(t)
Mb(() 1oEjo. (-
in order to have Ba,s(uuaa b)o(ax + b)] = (xx) for any |oo| < A. In this case, we have
sup (ula 00 a
sup
(z)
22A us(t)dt
1
+ infeeto.A] (o)
0,1-
Remark 7. When f is a polynomial and (at(t) has a heavy tail, sup |m(a, b) will only depend on A polynomially, More concretely, consider the case f(2) = Eosigr G2 where sup, /cal O(1). In this case, we have
((0) -(A9) O(A))
sup 1El-A.1),=0011,
Furthermore, since have assumed (t) 2 (1d) 11 we have
infrejo.A) mb(() O(A)) and sup (u(a; b) O(APt9) ab
12 (o)d)
43


Reviews:
[{"review_id": "ahoUtnvaaD", "reviewer": "Reviewer_oysL", "summary": "In this paper, the authors consider learning target functions of the form $g \\circ p$ where $p$ is a degree $k$ polynomial and $g$ is a degree $q$ polynomial. They consider a specific 3 layer neural network with a skip connection and a bottleneck layer of size 1. They show, using layer-wise training that $O(d^k)$ samples are sufficient to train the second layer weights and $O(1)$ samples are sufficient to train the last layer weights. In particular, the total sample complexity $O(d^{k})$ is much smaller than $O(d^{kp})$ required for inner-product kernels. Note also that information theoretically, $\\Omega(d^k)$ is information theoretically minimal to learn degree-$k$ polynomials, however they only consider a particular subset of polynomials that has much smaller functional space dimension (except for $k=2$).", "questions": "- Why not directly consider a general function $g$ instead of a polynomial? The random feature analysis shouldn’t change much in that case (it simply requires to show a $\\eps$-approximate certificate using uni-dimensional random features).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The paper considers the problem of understanding the benefits of feature learning for multi-layer neural networks. The one layer case has attracted a lot of attention, while multi-layer is so far way less understood.\n- They are able to show a large separation with kernel methods. The reason is that the first layer is able to extract a good representation of the data (the polynomial $p(x)$) from only seeing samples $g \\circ p (x)$.\n- This class of target functions naturally generalizes the single index model ($k=1$), which is the natural class of functions for one-hidden layer neural networks (for one-hidden layer , the network has direct access to linear functions of the data).\n- The paper is easy to follow and the assumptions and proof techniques are accurately presented and discussed.", "weaknesses": "- The architecture and algorithm are chosen specifically to succeed for this specific class of target functions (composition of degree-$k$ multivariate polynomials with univariate functions). Several previous works have considered such layer-wise training on non-regular architectures for specific hierarchical classes of functions, including [Allen-Zhu,Li,2020] and [The staircase property, Abbe, Boix, Brennan, Bresler, Nagaraj, 2020]. It is unclear how this paper contributes in terms of novel ideas in that direction.\n- Overall, the architecture choice and algorithm makes the analysis straightforward. It reduces the problem to sequentially fitting two linear random feature models, which is by now quite well understood. See for example [Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration, Mei et al., 2022] which gives the sample size and network width $\\Theta(d^k)$ to learn any degree-$k$ polynomials (for spherical data). The main technical innovation is Lemma 2, which relies on the specific construction of assumption 4. \n- For these reasons, I wonder how much this analysis can be generalized to other settings. Especially, in more realistic settings, we expect a non-linear evolution of the parameters, which this paper avoids using their specific architecture. In contrast, while [Nichani et al., 23] only considers one step size (so ultimately also a linear model step), they use a more standard architecture (but get a worse sample complexity dependency).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "qhYBTonOfU", "reviewer": "Reviewer_znhr", "summary": "The authors examine the problem of learning hierarchical polynomials of Gaussian data with three-layer networks trained using layerwise gradient descent. They compare the performance of three-layer networks to kernel approaches and two-layer networks, demonstrating a clear sample-complexity improvement without requiring any low-rank structure, as is typically done with sparse multi-index targets learned efficiently by two-layer neural networks. The theoretical analysis is reinforced by extensive discussion.", "questions": "- It would help the unexperienced reader to include a small comment on [Ben Arous et al. 2021] in the related works section to give a complete overview on learning single-index models.\n\n- In remark 2 is written \"can be extended to any $L= \\omega_d(1)$\". What is the original condition on L?\n\n- Please state after Theorem 1 or Corollary 1 that there is a clear sample-complexity improvement for three-layer networks over two-layer networks. At the moment, this is done only with respect to kernel methods.\n\n- Could the author comment with references characterizing the failure of two-layer networks in learning these high-rank target functions? Are there any provable guarantees?  \n\n- Before Lemma 1 mention that $n = O(d^k)$ is used, and repeat references that prove for what reason you fit the best degree k polynomial.\n\n- Could the author comment on the requirements for the gradient step size in the first phase? It is mentioned that [Nichani et al. 2023] considered one large gradient step in this phase. How does this relate to your learning rate requirement in Thm 1?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is pleasant to read and the mathematical results are supported with extensive discussion.", "weaknesses": "The main weakness of the paper is the close relationship with previous works. Although a fair comparison is given, the works by [Allen-Zhu & Li, 2019;2020] and [Nichani et al. 2023] contain many of the key ideas in the manuscript.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "8hU3iikNCP", "reviewer": "Reviewer_1ASh", "summary": "This paper studies the problem of learning hierarchical polynomials with neural networks. In particular, it focuses on the problem of learning functions of the form $h(x) = g(p(x))$, where $p : \\mathbb{R}^q \\to \\mathbb{R}$ is a degree $k$ polynomial and $g : \\mathbb R \\to \\mathbb R$ is a degree $q$ polynomial (the special case $k=1$ recovers the well-studied family of single-index models).\n\nThe paper's main result is that for a subclass of degree $k$ polynomials $p$ and standard Gaussian marginals, a 3-layer NN trained via layerwise GD on the $L_2$ loss learns the target hierarchical polynomial (realizable setting) with roughly $d^k$ samples and runtime that is polynomial in the parameters of the problem.\n\nThe main conceptual message is that this sample complexity guarantee improves over kernel-based approaches which require roughly $d^{kq}$ samples and essentially cannot make use of the special hierarchical structure of the target functions.", "questions": "(1) I think a discussion on Assumption 4 (as I mention in the above section) is an important aspect that should be expanded in the current draft.\n\n(2) As an additional comment, I think that Theorem 1 could be written more formally (e.g., mention that Assumptions 1-6 hold, mention the runtime of the training process, etc.). \n\n(3) The current result heavily relies on Gaussian marginals. Do you think similar results could be established for other continuous measures? Or discrete ones (in the Boolean domain, any function is essentially a polynomial in the Fourier basis)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper's main result is that for a subclass of degree $k$ polynomials $p$ and standard Gaussian marginals, a 3-layer NN trained via layerwise GD on the $L_2$ loss learns the target hierarchical polynomial (realizable setting) with roughly $d^k$ samples and runtime that is polynomial in the parameters of the problem.\n\nI think that the result is interesting and fits well with the ICLR community. In general, the paper is easy to read: the assumptions are presented in a clear manner, comparison with prior work is well-established and the contribution seems important in the field.\n\nAt a technical level, the authors essentially show that (i) during the first training phase, the NN implements kernel regression and essentially learns the underlying polynomial $p$ and (ii) during the second stage, the NN recovers the link function $g$. The main technical tool is that using the special structure assumed for the polynomials $p$ (Assumption 4), the paper provides an approximate version of Stein's Lemma (see Lemma 2), which can be used to show Item (i) from the above technical results. I think that the paper's technical contribution is sufficient for acceptance, since it extends existing ideas and provides new insights in the area.\n\nIn general, I enjoyed reading this paper and I vote for acceptance.", "weaknesses": "I believe that it would be beneficial if the authors mentioned families of polynomials not captured by Assumption 4. This would make more clear how strong and restrictive this assumption is. This assumption highly simplifies the analysis and, hence, it would be nice if the authors could further discuss on this assumption (I see why the families of Remark 3 satisfy this condition, but I think a further discussion on how this assumption simplifies the analysis would be helpful).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "asi8SVQ3iI", "reviewer": "Reviewer_RCBe", "summary": "This paper studies the three-layer neural network for feature learning trained by two-stage layer-by-layer gradient descent.\nIn the first stage, the network training is equivalent to a random feature model $u \\sigma_1(Vx)$ where $V$ is fixed.\nIn the second stage, it’s also a random feature model. But it can be regarded as a composition of kernel methods, and achieves improved sample complexity than a single kernel methods.", "questions": "- How does the first inequality in Eq. (13) hold?\n\nOverall, the model is a composition of kernel methods (or random features). In my view, understanding composition of kernel methods than a single kernel is more accurate and has more significance than so-called “three-layer” feature learning. I'm willing to increase my score if the significant revision on the motivation and story is done.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I appreciate the authors’ discussion in Section 5. It addresses my several concerns.\n\n- In the algorithm, it uses multiple steps in the first stage to learn the feature $p(x)$ and obtain better improved results when compared to (Nichani et al. 2023) on the sample complexity.\n- I like the discussion on Sec. 5.2 and the question is quite interesting: why the compositional kernel methods can be efficiently learned than that with a single kernel? Intuitively, this question can be answered in some points but is difficult from the theoretical side.", "weaknesses": "- The model and problem setting are a bit over-claim. Though it’s a three-layer neural network, only $u$ and $c$ are trained. This is actually a two-layer neural networks when the input is not the original data but after a fixed feature mapping. More importantly, the used layer-by-layer training scheme makes such three-layer neural network degenerate to a composition of two kernel methods (random features). With weight decay, the loss function under two stages is strongly convex and smooth. Both problem setting and algorithm are totally far away from true three-layer neural networks or even two-layer neural networks.\n\n- The results cannot handle the over-parameterized regime, i.e., the results require the width $m_1 << n$. The derived results set $m_1$ and $n_1$ depending on $d$, and $m_1$ and $n$ should have different orders of $d$. Otherwise the bound will become vacuous. This is because there is one $O(\\sqrt{m/n})$ term in the convergence rate, deriving by Rademacher complexity for the random features model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)