Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks
Authors: Eshaan Nichani
Alex Damian
Jason D. Lee
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Deep Learning Theory
Feature Learning
Three-Layer Neural Network
Depth Separation
Gradient Descent
Representation Learning
TL;DR: We theoretically analyze the features learned by gradient descent on three-layer neural networks, and show that there exist classes of functions efficiently learnable by three-layer networks but not by two-layer networks.

Abstract: One of the central questions in the theory of deep learning is to understand how neural networks learn hierarchical features. The ability of deep networks to extract salient features is crucial to both their outstanding generalization ability and the modern deep learning paradigm of pretraining and finetuneing. However, this feature learning process remains poorly understood from a theoretical perspective, with existing analyses largely restricted to two-layer networks. In this work we show that three-layer neural networks have provably richer feature learning capabilities than two-layer networks. We analyze the features learned by a three-layer network trained with layer-wise gradient descent, and present a general purpose theorem which upper bounds the sample complexity and width needed to achieve low test error when the target has specific hierarchical structure. We instantiate our framework in specific statistical learning settings -- single-index models and functions of quadratic features -- and show that in the latter setting three-layer networks obtain a sample complexity improvement over all existing guarantees for two-layer networks. Crucially, this sample complexity improvement relies on the ability of three-layer networks to efficiently learn *nonlinear* features. We then establish a concrete optimization-based depth separation by constructing a function which is efficiently learnable via gradient descent on a three-layer network, yet cannot be learned efficiently by a two-layer network. Our work makes progress towards understanding the provable benefit of three-layer neural networks over two-layer networks in the feature learning regime.

Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks
Eshaan Nichani Princeton University eshni sshnich@princeton.com
Alex Damian Princeton University ad27@princeton. edu
Jason D. Lee Princeton University
jasonlee@princeton.coo edu
Abstract
One of the central questions in the theory of deep learning is to understand how neural networks learn hierarchical features. The ability of deep networks to extract salient features is crucial to both their outstanding generalization ability and the modern deep learning paradigm of pretraining and finetuneing. However, this feature learning process remains poorly understood from a theoretical perspective, with existing analyses largely restricted two-layer networks. In this work we show that three-layer neural networks have provably richer feature learning capa- bilities than two-layer networks. We analyze the features learned by three-layer network trained with layer-wise gradient descent, and present a general purpose theorem which upper bounds the sample complexity and width nceded to achieve low test eгror when the target has specific hierarchical structure. We instantiate our framework specific statistical learning settings single-index models and functions of quadratic features and show that in the latter setting three-layer networks obtain a sample complexity improvement over all existing guarantees for two-layer networks. Crucially, this sample complexity improvement relies on the ability of three-layer networks to efficiently learn nonlinear features. We then establish concrete optimization-baser depth separation by constructing function which is efficiently learnable via gradient descent on three-layer network, yet cannot be learned efficiently by a two-layer network. Our work makes progress towards understanding the provable benefit three-layer neural networks over two-layer networks in the feature learning regime.
1 Introduction
The success of modern deep learning can largely be attributed to the ability of deep neural networks to decompose the target function into hierarchy of learned features. This feature learning process enables both improved accuracy [29] and transfer learning (11. Despite its importance, we still have rudimentary theoretical understanding of the feature learning process. Fundamental questions include understanding what features are Iearned, how they are learned, and how they affect generalization.
From a theoretical viewpoint, a fascinating question is understand how depth can leveraged to learn more salient features and as consequence richer class of hierarchical functions. The base case for this question is understand which features (and function classes) can be learned efficiently by three-layer neural networks, but not two-layer networks. Recent work on feature learning has shown that two-layer neural networks learn features which are linear functions of the input (see Section 1.2 for further discussion). It is thus a natural question to understand if three-layer networks can learn nonlinear features, and how this can leveraged to obtain sample complexity improvement. Initial learning guarantees for three-layer networks []6] 5 4), however, consider simplified model and function classes and do not discern specifically what the learned features are or whether a sample complexity improvement can be obtained over shallower networks or kernel methods.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

On the other hand, the standard approach in deep learning theory to understand the benefit of depth has been establish "depth separations" [ST]. i.c. functions that cannot be efficiently approximated by shallow networks, but can be via deeper networks. However, depth separations are solely concerned with the representational capability of neural networks, and ignore the optimization and generalization aspects. In fact, depth separation functions such as [51] are often not learnable via gradient descent [36]. To reconcile this, recent papers [45] [4] have established optimization based depth separation results, which are functions which cannot be efficiently learned using gradient descent on a shallow network but can be learned with a deeper network. We thus aim to answer the following question:
What features are learned by gradient descent on three-layer neural network, and can these features be leveraged to obtain a provable sample complexity guarantee?
1.1 Our contributions
We provide theoretical evidence that three-layer neural networks have provably richer feature learning capabilities than their two-layer counterparts. We specifically study the features learned by three-layer network trained with layer-wise variant of gradient descent (Algorithm) I Our main contributions are as follows.
Theorem 1 is general purpose sample complexity guarantee for Algorithm 1 learn an arbitrary target function f". We first show that Algorithm 1 learns a feature roughly corresponding to low-frequency component of the target function f with respect to the random feature kernel IK induced by the first layer. We then derive an upper bound on the population loss in terms of the learned feature. As a consequence, we show that if f" possesses a hierarchical structure where it can be written as a ID function of the learned feature (detailed in Section 3. then the sample complexity for learning f* is equal to the sample complexity of learning the feature. This demonstrates that three-layer networks indeed perform hierarchical learning.
We next instantiate Theorem
in two statistical learning settings which satisfy such hierar- chical structure. As a warmup, we show that Algorithm [ leerns single-index models (i.c f (x) = 9 (w xr)) in dl samples, which is comparable to existing guarantees for two-layer networks and crucially has d-dependence not scaling with the degree of the link function
(xTAx), Ax), where g" is either Lipschitz or a degree p = O(1) polynomial, up to (1 error with dt samples. This improves on all existing guarantees for learning with two-layer networks or via NTK-based approaches, which all require sample complexity dd(p) A key technical step is to show that for the target (x( = g (x" Ax). the learned feature is approximately x" Az. This argument relies on the universality principle in high-dimensional probability, and may be of independent interest.
g". We next show that Algorithm ] learns the target (x) g" g²(x"
We conclude by establishing an explicit optimization-base depth separation: We show that the target function (x) = ReLU(x" (x"AA) for appropriately chosen A can be learned by Algorithm I up (1(1) error in d' samples, whereas any two layer network needs either superpolynomin width or weight norm in order to approximate f up to comparable accuracy. This implies that such an fr is not efficiently learnable via two-layer networks.
The above separation hinges on the ability of laree-layer networks to learn the nonlinear feature x' Ax and leverage this feature learning to obtain an improved sample complexity. Altogether, our work presents general framework demonstrating the capability of three-layer networks to learn nonlinear features, and makes progress towards a rigorous understanding of feature learning, optimization- based depth separations, and the role of depth in deep learning more generally.
1.2 Related Work
Neural Networks and Kernel Methods. Early guarantees for neural networks relied on the Neural Tangent Kernel (NTK) theory B1 50
17 The NTK theory shows global convergence by coupling to a kernel regression problem and generalization via the application of kernel generalization
23
bounds 7 15 5 The NTK can be characterized explicitly for certain data distributions [2 39 38
which allows for tight sample complexity and width analyses. This connection to kernel methods has also been used to study the role of depth, by analyzing the signal propagation and evolution of
2

---

the NTK in MLPs 42 48 28, convolutional networks [6 55 56 40), and residual networks 30
However, the NTK theory is insufficient as neural networks outperform their NTK in practice 6 32 In fact, ( shows that kernels cannot adapt to low-dimensional structure and require dh samples to learn any degree k polynomials in dimensions. Ultimately, the NTK theory fails to explain generalization or the role of depth in practical networks not in the kernel regime. A key distinction is that networks in the kernel regime cannot learn features [57]. A recent goal has thus been to understand the feature learning mechanism and how this leads to sample complexity improvements 53 22 58
B 25 26 20 33 35). Crucially, our analysis not in the kernel regime, and shows
an improvement of three- layer networks over two-layer networks in the Seature-learning regime.
Feature Learning. Recent work has studied the provable feature learning capabilities of two-layer neural networks. 99 2 8 1 10] show that for isotropic data distributions, two-layer networks learn linear features of the data, and thus efficiently learn functions of low-dimensional projections of the input (i.e targets of the form (x) (U2) for U E R Here, I Ur is the "linear feature." Such target functions include low-rank polynomials 18 and single-index models [ 13 for Gaussian covariates, well sparse boolean functions
1 such the k-sparse parity problem 100 for covariates uniform on the hypercube. [43 draws connections from the mechanisms in these works feature learning in standard image classification settings. The above approaches rely on layerwise training procedures, and our Algorithm Dis is an adaptation of the algorithm in [88.
Another approach uses the quadratic Taylor expansion of the network to learn classes of polynomials [ This approach can be extended to three-layer networks. [16] replace the outermost layer with its quadratic approximation, and by viewing 2² as the hierarchical function (=#/2)2 show that their three-layer network can learn low rank, degree p polynomials in d'/2 samples. [5] similarly uses a quadratic approximation to improperly learn a class of three-layer networks via sign-randomized GD. An instantiation of their upper bound to the target g (x* Ax) for degree p polynomial, g" yields sample complexity of dd+11 However, [16 5 are proved via opaque landscape analyses, do not concretely identify the learned features, and rely on nonstandard algorithmic modifications. Our Theorem directly identifies the learned features, and when applied to the quadratic feature setting in Section obtains an improved sample complexity guarantee independent of the degree of g".
Depth Separations. 5 constructs function which can be approximated by poly-width network with large depth, but not with smaller depth. 24 is the first depth separation between depth 2 and 3 networks, with later works [46] 19 47 constructing additional such examples. However, such functions are often not learnable via three-layer networks [34]. [36] shows that approximatability by shallow (depth 3 network) is a necessary condition for learnability via a deeper network.
These issues have motivated the development of optimization-bassd, or algorithmic, depth separations, which construct functions which are learnable by three-layer network but not by two-layer networks. 45 shows that certain ball indicator functions 1 [[|x=|| 2 A) are not approximatable by two-layer networks, yet are learnable via GD special variant of three-layer network with second layer width equal to 1. However, their network architecture is tailored for learning the ball indicator, and the explicit polynomial sample complexity (n Z d66) is weak. [44] shows that a multi-layer mean-field network with a 1D bottleneck layer can learn the target RcLU(1 - ||=])), which 47) previously showed was inaproxinatable via two-layer networks. However, their analysis relies on the rotational invariance of the target function, and it is difficult to read off explicit sample complexity and width guarantees beyond being (oly Our Section 4.2 shows that three- layer networks can learn larger class of features (xx" Ax versus [|x=]) and functions on top of these features (any Lipschitz versus ReLU), with explicit dependence on the width and sample complexity needed (n, mi, m2 = Odd)).
2 Preliminaries
2.1 Problem Setup
Data distribution. Our aim is to learn the target function f : X'a - R,,tthh Xa C Rd the space of covariates. We let v be some distribution on X'a and draw two independent datasets Dj, Dz, each with n samples, so that each x E D2 or x E D2 is sampled i.i.d as x ~ v. Without loss of generality, we normalize so Egan [f (x)) < 1. We make the following assumptions on D:
3

---

Algorithm Layer-wise training algorithm
Input: Initialization (();; learning rates 7/, 122 weight decay x; time T
(Stage I: Train W
W() -WOO ?n V T.1 (0(0)
al)) < -0
((1) - (al)) bWW, B(0)) )(0)
(Stage 2: Train a)
for I 2, ... T do
aL (1-11) 122 V. L (0(t-1) \a(t-11)
@10) + (a W(1), bl0), )(0))
end for
< 8(T)
Output: 0
Definition (Sub-Gaussian Vector). A mean-zero random vector X E R is y-subGaussian if. for all unit vectors E R Elexp(AXX v)] S xxp(>>>²) for all X E R.
Assumption 1. E [a) and C.-subGGussian for some constant Cn--
Assumption 2. f" has polynomially growing moments, i.e there exist constants (C₁,e) such that Egwr [f s Cyy for all 1
We note that Assumption 2 is satisfied by number of common distributions and functions, and we will verify that Assumption holds for each example in Section
Three-layer neural network. Let m., my be the two hidden layer widths, and σ1, be two activation functions. Our learner is a three-layer neural network parameterized by (a, W, V). where mm W mmmixm, E IRm,, and V Rmxxd The network f(;; 0) is defined as:
fx:00) = 1 a 01 (Woz(Vx) - b) Σ 0000 ((w0, h 100 + +b.)
Here, E A is the ith row of W, and (x) = (Vz) E R is the random feature embedding arising from the innermost layer. The parameter vector ((0) == (a) W(O), b(0), v()) the biases 6(0) ~idd AO 1). and the rows 4/ of (0) drawn v; riid where is the uniform measure (1), the d'dimensional unit sphere. We make the following assumption on the activations, and note that the polynomial growth assumption
initialized with a!") -iid Un)1( (0))
is satisfied by all activations used in practice.
Assumption 3. ai is the ReLU activation, i. a (z) = max 0), and a2 has polynomial growth, ie [02(1)) < (1 + (xr))a> for some constants Ca: Oar
Training Algorithm. Let (f) denote the empirical loss on dataset De: that is for i = 1, 2: == : 1 ΣxeD, (f(z; 0) - -x))) Our network is trained via layer-wise gradient descent with sample splitting. Throughout training, the first layer weights V and second layer bias b are held constant. First, the second layer weights W are trained for t = timesteps. Next, the outer layer weights a are trained for f T timesteps. This two stage training process common in prior works analyzing gradient descent on two-layer networks [8] 8 I 10], and as we see in Section 5 is already sufficient to establish a separation between two and three-layer networks. Pseudocode for the training procedure is presented in Algorithm [
2.2 Technical definitions
The activation d2 admits random feature kernel K : X' Xa - E and corresponding integral operator KK : L v) + (X),
Definition (Kernel objects). 02 admits the random feature kernel
K(x, x²) Eveor Toz(x v)02(z' v)]
(2)

---

and corresponding integral operator
(TKf) := : (xx, (x)))
(3)
We make the following assumption on IK, which we verify for the examples in Section Assumption 4. f has polynomially bounded moments, i.e there exist constants CK X such that, for all 1 < q & d, [[ I [Kg
We also require the definition the Sobolev space: Definition 3. Let ([-1,,,
1]] be the Sobolev space of twice continuously differentiable functions for for 1,2.
q: [[1, + R equipped with the norm ||all, == mäX,sk max (x)
2.3 Notation
We use big O notation (i.e O, O, (2) to ignore absolute constants (Co, Cy, etc.) that do not depend on d, n, m1, m2- We further write S ba if aa O(ba), and o(ba) if lima-soc 0a/ba = 0. Additionally, we use Ô notation to ignore terms that depend logarithmically on dumim For f: X2 define II [I Il (=()) and [|9l)e)(xx.>)) (|alle=(>) are defined analogously for functions 9 (1-1(1) - R. When the domain is clear from context, write || k. |[g|] z.- We let (P v) be the space of f with finite ||| (X,,e) Finally, we write E, and E, as shorthand for Exwy and Eper respectively.
(f)=) To simplify notation we also call this quantity
3 Main Result
The following our main theorem which upper bounds the population loss of Algorithm
Theorem 1. Select E ([,,((-,,,),, Let a = miT, TJ. and assume m1. m2 = f There exist 7j, X, 122 such that after T = poly(n, mi, m2 d, l1 timesteps, with high probability over the initialization and datasets the output of Algorithm I satisfies the population L loss bound
Egg (I(=: ((()))
|gl. min (2, m. m22
Kf
l1
11a12,00
+
(7-0 -'
(4)
accuraccy of feature learning
sample complexity feature learning
complexity of
The full proof of this theorem is in Appendix D The population risk upper bound 4 has three terms:
1. The first term quantifies the extent to which feature learning is useful for learning the target f*, and depends on how close f* is having hierarchical structure. Coneretely, if there exists q R + R such that the compositional function o 77 KK is close to the target f* then this first term is small. In Section 4 we show that this true for certain hierarchical functions. In particular, say that f" satisfies the hierarchical structure I" = g" h" If the quantity f nearly proportional to the true feature h. then this first term negligible. As such, we refer to the quantity Kf as the learned feature.
2. The second term is the sample (and width) complexity of learning the feature Kf*. It is useful to compare the |K term to the standard kernel generalization bound, which requires 2 (J*, za. Unlike in the kernel bound, the feature learning term in (4) does not require inverting the kernel K as it only requires lower bound on | f This difference can be best understood by considering the alignment of f with the eigenfunctions of IK. Say that f* has nontrivial alignment with eigenfunctions of K. with both small (Amin) and large (Amar) eigenvalues. Kernel methods require ?(>mi))) samples, which blows up when Amnn small; the sample complexity kernel methods depends on the high frequency components of f". On the other hand, the guarantee in Theorem [ scales with
5

---

| O()maz), which can be much smaller. In other words, the sample complexity feature learning scales with the low-frequency components of f*. The feature learning process can thus be viewed as extracting the low-frequency components of the target.
3. The last two terms measure the complexity of learning the univariate function q- In the examples in Section 4 the effect of these terms is benign.
Altogether, if f* satisfies the hierarchical structure that its high-frequency components can be inferred from the low- low-frequency ones, then a good for Theorem exists and the dominant term in 4 the sample complexity of feature learning term, which only depends the low-frequee wwwfrequency components. This is not the case for kernel methods, as small eigenvalues blow up the sample complexity. As we show in Section 4 this ability ignore the small eigenvalue components of f during the feature learning process is critical for achieving good sample complexity in many problems.
3.1 Proof Sketch
At initialization, f(x; 0(0)) xo O. The first step of GD on the population loss for neuron wy is thus
( - V w.
Ex
(5)
Em
(6)
1 1 y 200
E. r (w)
(7)
12
Therefore the network f(a1; after the first step of GD given by
1 mi
f(11:0())
Σ
ajoz
(w) (1) h (000 b,)
(8)
11
j=1
m1
(0)
7 B Ez
aj01
+
1000 20
(9)
We first notice that this network now implements a ID function of the quantity
1 E. (+))))))))))))))))
o(z') = b
(10)
m2
Specifically, the network can be rewritten as
f(z:; f())
mı Σ ajor 6(x')(()))))))))) (af) b;) 11,00 100020
(11)
Since f implements hicrarchical function of the quantity d(z), we term the learned feature.
The second stage of Algorithm is equivalent to random feature regression. We first use results on ReLU random features to show that any E W2 ([-1, 1]) can be approximated on [-1, as
<(2) ~ 1 و aza (a, (0) + by) for some ||a*|| $ ll (Lemma Next, we use the standard kernel Rademacher bound to show that the excess risk scales with the smoothness of q (Lemma
Hence efficiently learn functions of the form q Φ.
suffices to compute this learned feature é. For my large, we observe that
7 m2
6(x')
E. If v)a(a' v)l)
(x)
(12)
The learned feature is thus approximately 7 Kf Choosing 7j so that |o(x') S 1, we see that Algorithm learns functions of the form ø (7 Kf*). Finally, we translate the above analysis the finite sample gradient via standard concentration tools. Since the empirical estimate to K concentrates at a 1/ / rate, n 2 || samples are needed to obtain a constant factor approximation (Lemma

---

Examples
We next instantiate Theorem in two specific statistical learning settings which satisfy the hierarchical As a warmup, we show that three-layer networks efficiently learn
prescription detailed in Section 3 single index models. Our second example shows how three-layer networks can obtain sample complexity improvement over existing guarantees for two-layer networks.
4.1 Warmup: single index models
Let f* = --"""" (w" z). for unknown direction w" E Fd and unknown link function g² R - R, and (N, I). Prior work [18 13 shows that two-layer neural networks learn
take X2 = Rd with such functions with an improved sample complexity over kernel methods. Let (2) = so that the network is of the form fxx; 0) = I a² a (WVx). We can verify that Assumptions to are satisfied, and thus applying Theorem [ in this setting yields the following:
Theorem Let (x)x g'(w" x), where ||00 |12 = 1. Assume that g", (ggY and (g")" are polynomially bounded and that E (=))
O. Then with high probability Algorithm 7
satisfies the population loss bound
= Ex (F(=: (f(=:0)- ((x) =
- (minnn( mi, m2)
d²
1 V
(13)
min
Given widths m., my = ( n = ((f) samples suffice to learn f³, which matches exist- ing guarantees for two-layer neural networks [3
18 We remark that prior work on learning
single-index models under assumptions on the link function such as monotonicity or the condition However, our sample complexity improves on that of kernel methods, which require d' samples when g" is a degree polynomial.
Ez~N(o, ((g (=))] & require d samples 49 37
Theorem 2 is proved in Appendix E. brief sketch is as follows. Since (2(2) z, the kernel is K K(x, x²)
(x) = Bx o' fff( 1 [x (x -)]
By an application of Stein's Lemma, the learned feature
E'
(14)
d'
Since (x) = (w x). (x) = w (n")(ww v). and thus
(z)) = 1 E0- (z))[u" x x = r.
(15)
The learned feature is proportional to the true feature, an appropriate choice of 17 and choosing q ==" 9" in Theorem implies that ||K || d² samples are needed to learn f*
4.2 Functions of quadratic features
The next example shows how three-layer networks can learn nonlinear features, and thus obtain sample complexity improvement over two-layer networks.
Let X₂ = 5d-1dd1-i (vdd, the sphere of radius vd, and v the uniform measure on Xa. The integral opera- tor K. has been well studied [39 27 (8), and its eigenfunctions correspond to the spherical harmonics. Preliminaries on spherical harmonics and this eigendecomposiitin are given in Appendix F
Consider the target (x(x) 9 (xx Ar), where A E Rdxd is symmetric matrix and g" I - BR-Rii an unknown link function. In contrast to a single-index model, the feature x Ax we aim to learn is quadratic function of Since can write x" Ar x" (A Tr(A) 2))= + Tr(A),
without
loss of generality assume Tr(A) = 0. We also select the normalization ||Al] ensures that Eg ((x" Ax = 1. We first make the following assumptions on the target function.:
d+2
2d = 0(1: this
Assumption E., [f (x))
E=~ (()) @(1), g" 1-Lipschitz and (g")" has
polynomial growth.
The first assumption can be achieved via a preprocessing step which subtracts the mean of f*, the second is a nondegeneracy condition, and the last two assume the target is sufficiently smooth.
We next require the eigenvalues of A to satisfy an incoherence condition:
7

---

Assumption 6. Define K = || Vd. Then = p(Vd/polylog(d))
Note that I Vd. If has rank ( and condition number ((1), then O()). Furthermore, when the entries of A are sampled i.idd バ (1) with high probability by Wigner's semicircle law.
Finally, we make the following nondegeneracy assumption on the Gegenbauer decomposition of 02 (defined in Appendix (F.. We show that (x2) Odd-2) and later argue that the following assumption is mild and indeed satisfied by standard activations such as d2 = ReLU.
Assumption 7. Let (0( be the 2nd Gegenbauer coefficient of 02. Then (22) = (d)
We can verify Assumptions I to hold for this setting, and thus applying Theorem following:
| yields the
Theorem 3. Under Assumptions 5
with high probability Algorithm 7 satisfies the population
loss bound
Ez (f(z:0)-
P min(n, m1, m2) + Vn
d²
1/3
S
+
Vaa
(16)
We thus require sample size n ((d*) and widths my, m2 m ((d")) to obtain 0a(1) test loss.
Proof Sketch. The integral operator K. has eigenspaces corresponding to spherical harmonics of degree k. In particular, [27 shows that, in L
Kf
P&f"
(17)
A:0
where Ps is the orthogonal projection onto the subspace of degree k spherical harmonics, and the ca are constants satisfying ex = (dd=). Since f an even function and EE, [f,[(((x)) = 0, truncating
this expansion at l yields
KS" 2(d P2f" + OOd"))
(18)
thus suffices to compute Pef To do so, we draw a connection to the universality phenomenon in high dimensional probability. Consider two features x² Az and x Br with (A, B) = O. We show that, when d is large, the distribution of x" Ax approaches that of the standard Gaussian, while x" Br approaches mixture of x² and Gaussian random variables independent of xT Az. As such, show
Ex [g" (x" (xT Ax) Bx] E, [9 (xT Ax)] Ex [x Bx
(19)
E2 [x" An)x Ar] Pv (. [9 (=)=== E =
((g (()))
(20)
The second expression can be viewed as an approximate version of Stein's lemma, which was applied in Section 4.1 to compute the learned feature. Altogether, our key technical result (stated formally in
Lemma 20 is that for f" ====" (x" Ax). the projection Pef satisfies
(P2f*) (o..wo..nll xJ Ar + oa()
(21)
The learned feature is thus ( (IT Ar + od(()) plugging this into Theorem yields the d' sample complexity.
The full proof of Theorem 3 is deferred to Appendix E.2 In Appendix E.3
we show that when 9" is degree = (1) polynomial, Algorithm learns I in (dd)) samples with an improved error floor.
Existing guarantees for two-layer networks cannot efficiently learn functions of the form f"(x) = (((x)=='fe (x" (x'"Ax)) Ar) for arbitrary Lipschitz g" In fact, in Section 5 we provide an explicit lower bound against two-layer networks efficiently learning a subclass of these functions. When g" is a degree p polynomial, networks in the kernel regime require dPp >> d" samples to learn f (27) Improved guarantees for two-layer networks learn degree p' polynomials in rank r & d projection of the input [18]. However, g" (aTTAz) Az) cannot be written in this form for some r < d. and thus existing guarantees do not apply. We conjecture that two-layer networks require ((()) samples when g" a degree polynomial.
Comparison to two-layer networks.
pe' samples when the target only depends
Altogether, the ability of taree-layer networks to efficiently learn the class of functions [x(x)= g" (at Ar) hinges on their ability to extract the correct nonlinear feature. Empirical validation of the above examples is given in Appendix A

---

An Optimization- Depth Separation
We complement the learning guarantee Section 4.2 with lower bound showing that there exist functions in this class that cannot be approximated by polynomial size two-layer network.
The class of candidate two-layer networks is as follows. For parameter vector 0 (a, W, b2), where E T W E Rmxd bi E I ba E IR, define the associated two-layer network as
No(z) == a oWW> + b1) + = Σ
a,o (w) x + b1,a) + +
(22)
0o |bb oc, l denote the maximum parameter value. We make the following assumption o, which holds for all commonly-usee activations.
Let |10] == max( |all W
Assumption 8. There exist constants Co, such that lo(2)) S C1 (1 + I21)
Our main theorem establishing the separation is the following.
Theorem Let d be a suffiently large even integer. Consider the target function f*()
(e Ax) Co. where A J2U
(Larz
14/2
O
Ia/2
0
U² for some orthogonal matrix U and co-
Bgw (ReLU(x) Ax)] Under Assumption there exist constants C1 C2, C3, Cz depending only on Czd any two layer neural network No(x) of width m and pop- ulation 1.2 error bound [[oo- - 1r1 S e must satisfy max(m, |10llcc) 2 C1 exp ( (C2e- 11/2 log(de))-
(Co,0o) such that for any 03
However, Algorithm 1 outputs a predictor satisfying the population L loss bound
Ex (((=:0)- 20
dt
V
d 1/6 d/6/
(23)
n
after T = poly(d, my, m2, n)) timesteps.
The lower bound follows from a modification of the argument in [9 along with an explicit de- composition of the ReLU function into spherical harmonics. We remark that the separation applies for any link function g" whose Gegenbauer coefficients decay sufficiently slow. The upper bound follows from an application of Theorem Tto to smoothed version of ReLU, since ReLU is not in 02 1, 1)). The full proof of the theorem is given in Appendix G
Remarks. In order for two-layer network to achieve test loss matching the d' error floor, either the width or maximum weight norm of the network must be exp(I(dd)) for some constant 5; this lower bound is superpolynomial in d. As a consequence, gradient descent on poly-width two-layer neural network with stable step size must run superpolynomially many iterations in order for some weight to grow this large and thus converge to a solution with low test error. Therefore f* not learnable via gradient descent in polynomial time. This reduction from a weight norm lower bound to runtime lower bound is made precise in [45].
We next describe a specific example of such an f*. Let S be a d/2-dimensional subspace of Rd. and let A Ps - 2 where Ps, Ps are projections onto the subspace and and its orthogonal complement respectively. Then, (x) ReLU 2d ||Psx|] dà
44 established
an optimization-basee based separation for the target ReLU(1 ||=-]),, under different distribution However, their analysis relies heavily on the rotational symmetry of the target, and they posed the question of learning ReLU(1 Ps for some subspace S. Our separation applies similar target function, and crucially does not rely on this rotational invariance.
6 Discussion
In this work we showed that three-layer networks can both learn nonlinear features and leverage these features to obtain provable sample complexity improvement over two-layer networks. There are number of interesting directions for future work. First, can our framework be used to learn hierarchical functions of larger class of features beyond quadratic functions? Next, since Theorem D is general

---

purpose and makes minimal distributional assumptions, it would be interesting to understand if it can be upplied to standard empirical datasets such as CIFAR-10, and what the learned feature f and hierarchical learning correspond to in this setting. Finally, our analysis studies the nonlinear feature learning that arises from a neural representation (22)) where V fixed initialization. This alone was enough to establish separation between two and three-layer networks. A fascinating question to understand the additional features that can be learned when both V and W are jointly trained. Such an analysis, however, is incredibly challenging in feature-learning regimes.
Acknowledgements
EN acknowledges support from National Defense Science & Engineering Graduate Fellowship. AD acknowledges support from a NSF Graduate Research Fellowship. EN, AD, and JDL acknowledge support of the ARO under MURI Award 191
the Sloan Research Fellowship. NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994. The authors would like to thank Itay Safran for helpful discussions.
References
[1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Proceedings of Thirty Fifth Conference on Learning Theory, pages 4782- 4887, 2022.
[2] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning neural net- works: leap complexity and saddle-to-saddle dynamies. In The Thirty Sixth Annual Conference on Learning Theory, pages 2552-2623. PMLR, 2023.
[3] Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
[4] Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning, arXiv preprint TrXiv:2001.044133 2020.
[5] Zeyuan Allen-Zi 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.
[6] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[7] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analy- sis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning (ICML), 2019.
[8] 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-37946, 2022.
[9] Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher- approximation of wide neural networks. In International Conference on Learning Representations, 2020.
[10] Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[11] William Beckner. Sobolev inequalities, the poisson semigroup, and analysis on the sphere sn. Proceedings of the National Academy of Sciences of the United States of America, 89 11: 4816-9, 1992.
10

---

[12] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from sighdimensional inference. The Journal of Machine Learning Research, ((1))4788-4833 2021.
[13] 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.
[14] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye "aandemman-Mlln and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githut com/ google/jax
[15] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, 2019.
[16] Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Ad- vances in Neural Information Processing Systems, 33:22134-22144 2020.
[17] Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In Advances in Newral Information Processing Systems (NeurIPS), 2019.
[18] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
[19] Amit Daniely. Depth separation for neural networks. In Conference on Learning Theory, pages 690-696, PMLR, 2017.
[20] Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33:20356-20365, 2020.
[21] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidircctional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171 1171-4 Minneapolis, Minnesota, June 2019, Association for Computational Linguisties. doi: 10. 1/SS3//1/N11-1423 URL
org/N19-- 1423
[22] Simon S. Du and Jason D. Lee. On the power of over-parametriizatoo in neural networks with quadratic activation. Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 1329-1338, 2018.
[23] Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations (ICLR), 2019.
[24] Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Conference on learning theory, pages 907-940. PMLR, 2016.
[25] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz. and Andrea Montanari. Limitations of lazy training of two-layers neural network. In Advances in Neural Information Processing Systems, 2019.
[26] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems, 2020.
[27] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49: 19:10991105, 2021.
[28] Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. Exact convergence rates of the neural tangent kernel in the large depth limit. arXiv preprint arXiv:1905.13654 2019.
11

---

[29] Kaiming He. X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). pages 770-778, 2015.
[30] Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual net- works generalize better than deep feedforward networks? - a neural tangent kernel per- spective. In H. Larochelle, M. Ranzato. R. Hadsell, M.F. Balcan, and H. Lin, editors, Ad- vances in Neural Information Processing Systems, volume 33, pages 2698-2709 Curran Associates, Inc.. 2020. URL httts::/ http://iiig..aaopooooocop//////a pcppppp////p////////lee/// [c3$5500f2bcc222249467777-1---23
newrins pdf
[31] Arthur Jacot, Franck Gabriel, and Clément Hongler. generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Neural tangent kernel: Convergence and
[32] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. In Advances in Neural Information Processing Systems, 2020.
[33] Yuanzhi Li, Tengyu Ma, and Hongyang R. Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Proceedings of Thirty Third Conference on Learning Theory, Proceedings of Machine Learning Research, pages
2020.
[34] Eran Malach and Shai Shalev-Shwand good? In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
Is deeper better only when shallow is
http://wwwedeanaaeee.e.....eeeee pdf
[35] Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 7379-
7389, 2021.
[36] Eran Malach, Gilad Yehudai, Shai Shalev-Schwartz and Ohad Shamir. The connection between approximation, depth separation and learnability in neural networks. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory, volume PMLR, 15-19 Aug 2021. html
134 of Proceedings of Machine Learning Research, pages mlr, press/ pepxess/v134//e
URL htttps://proceeddiiiiii wtt.:///........../eeeeeemm
[37] Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 46(6A):2747-2774, 2018.
[38] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Generalization error of random features and kernel methods: hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis, Special Issue on Harmonic Analysis and Machine Learning. 59:3- 84, 2022.
[39] Andrea Montanari and Yiqiao Zhong, The interpolation phase transition in neural networks: Memorization and generalization under lazy training. arXiv preprint arXiv:2007.12826, 2020.
[40] Eshaan Nichani, Adityanarayanan Radhakrishnan, and Caroline Uhler. Increasing depth leads to u-shaped test risk in over-parameterized convolutional networks. arXiv preprint arXiv:2010.096100 2020.
[41] Eshaan Nichani, Yu Bai, and Jason D Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. Advances in Neural Information Processing Systems, 35:14568-1458 2022.
[42] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl- Sohl-Dickstein, and Surya Ganguli. Expo- nential expressivity deep neural networks through transient chaos. In D. Lee, M. Sugiyama, U. Luxburg. I. Guyon, and R. Garnett, editors, Advunces in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL http::tttt://prcedddinge neurips cc
Paper. pdf
12

---

[43] Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Feature learning in neural networks and kernel machines that recursively learn features. arXiv preprint arXiv:2212. 13881, 2022.
[44] Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. arXiv preprint arXiv:2304. 30..01063, 2023.
[45] Itay Safran and Jason Lee. Optimization-baseed separations for neural networks. In Conference on Learning Theory, pages 3-64. PMLR, 2022.
[46] Itay Safran and Ohad Shamir. Depth- width tradeoff's in approximating natural functions with neural networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2979-2987. PMLR, 06-11 Aug 2017. URL httpps://proceedingg mlr press/770saafran1 html
[47] Itay Safran, Ronen Eldan, and Ohad Shamir. Depth separations in neural networks: What is actually being separated? In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Secone Conference Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 2664-2666. PMLR, 25-28 Jun 2019. URL htttps://proceeding: mlr
press/v99/ 0atss//99//afrfa199... wttp://faaanna.aaaammaa.mmm html
[48] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl- Deep infor- mation propagation. In International Conference on Learning Representations, 2017. URL
[49] Mahdi Soltanolkotabi. Learning relus via gradient descent. In Advances in Neural Information Processing Systems (NeurIPSS) 2017.
[50] Mahdi Soltanolkotabi, Adel Javanmard, and Jason D. Lee. Theoretical insights into the op- timization landscape of werrparameterizeed shallow neural networks. IEEE Transactions on Information Theory, 65:742-769, 2018.
[51] Matus Telgarsky. Benefits of depth in neural networks, In Conference on learning theory, pages 1517-1539. PMLR, 2016.
[52] Ramon van Handel. Probability in high dimension, 2016. URL http://web.man princeton wwwpervaa///////eeeeeaeaaaaaaaae "rvan/ http:/PC555..ppp pdf
[53] Colin Wei, Jason Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. arXiv preprint arXiv:1810.053699, 2018.
[54] Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, Proceedings of Machine Learning Research, pages 3635-3673. 2020.
[55] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pen- nington. Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks. In Jennifer Dy and Andreas Krause, editors, Pro- ceedings of the 35th International Conference on Machine Learning. volume 80 of Pro- ceedings of Machine Learning Research, pages 5393-5402. PMLR, 10-15 Jul 2018. URL mlr wwwps/e80/eeeeppp http///80xxxx11110. xiao18a..... html
[56] Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generalization in deep neural networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 10462-10472. PMIR, 13-18 Jul 2020. URI. https: //proceedings.no mlr. www.srresssseeee ptress/111/xxaa2201 html
[57] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-wice neural networks. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[58] Gilad Yehudai and Ohad Shamir, On the power and limitations of random features for under- standing neural networks. In Advances in Neural Information Processing Systems, 2019.
13

---

Limitations
One limitation of our work is that the innermost layer weights V are fixed throughout training, and our training procedure layerwise. We remark that this is already enough to establish an algorithmic separation between two and three-layer networks (Section D and that prior works on two-layer networks also rely on layerwise training procedures. is very interesting direction of future work to understand if training V can lead to learning a larger class of functions. Another limitation is that our examples in Section 4 rely on the data distribution v being either the standard Gaussian or uniform on the sphere. This allows us to carefully characterize the integral operator K, and we remark that other works on two-layer networks make similar distributional assumptions (such as uniform on the hypercube). However, it is very important direction of future work to extend these lines of work to more general data distributions.
A Empirical Validation
Single Index Setting
Quadratic Feature (Random Entries)
Test Loss
Correlation of Features
Teat Loss
Correlation of Features
Tost Lo0s
2014 LO00
Tex
if
a Loumber of uaiiiin samples)
a coumee of traning sumples)
ntumber waining sumples)
essumber of miniig sumplest
Quadratic Feature (Random Projection)
Correlation Between Learned and Predicted Features
Test Loss
Correlation of Features
Single Index, of =
Random Projection, d 16
and
992 Pucatoo Fomie
9U1 (Lomned Fontre)
Tea Lave
triingg samplos)
training crmples)
A000 9
We ran Algorithm D on both the single index and quadratic feature settings described in Section 4 Each trial was run with 5 random seeds. The solid lines represent the medians and the shaded areas represent the min and max values. For every trial we recorded both the test loss on a test set of size 25 and the linear correlation between the learned feature map 6(r) and the true intermediate feature h* (x) where h" (x) I for the single index setting and h h²(x) = IJ Ar for the quadratic feature setting. Our results show that the test loss goes to 0 as the linear correlation between the learned feature map & and the true intermediate feature h² approaches 1.
Figure 1:
We empirically verify our conclusions in the single index setting of Section 4.1 and the quadratic feature setting of Section 4.2
Single Index Setting We learn the target function g" (w* x:) using Algorithm I where w' Sd-1 is drawn randomly and (x) sigmoid(x) (o-' # 0. As in Theorem 2, we choose the initial activation 02(2) = 2. We opti-
I+e- which satisfies the condition
mize the hyperparameters M/1 using grid search over a holdout validation set of size 215 and report the final error over a test set of size 215
Quadratic Feature Setting We learn the target function 9" (x" Ax) using Algorithm where g" (x) = Rz)U(x). We ran our experiments with two different choices of A::
is symmetric with random entries, i.. Aij ~ N 1) and A₁; = Ais for i j. A is a random projection, i.c. A = IIs where S is random d/2 dimensional subspace.
14

---

Both choices of A were then normalized so that tr A = 0 and |||] P by subtracting the trace and dividing by the Frobenius norm. We chose initial activation (2 (z) We note that in both examples, holdout validation set of size 215 and report the final error over a test set of size 215
(1)). As above, we optimize the hyperparameters 7/1 À using grid search over
To focus on the sample complexity and avoid width-related bottlenecks, we directly simulate the infinite width limit (mz - oc) of Algorithm by computing the kernel KK in closed form. Finally, we run each trial with 5 random seeds and report the min, median, and max values in Figure
Comparison Between Two and Three-Layer Networks We also show that the sample complexity
separation between two and three layer networks persists in standard training settings. In Figure 2 we train both a two and three-layer neural network on the target (x) ReLU(x" Ax). where A is symmetric with random entries, as described above. Both networks are initialized using the μP parameterization (57] and are trained using SGD with momentum on all layers simultaneously. The input dimension is d 32, and the widths are chosen to be 256 for the three-layer network and 2048 for the two-layer network, so that the parameter counts are approximately equal. Aigureploo the average test loss over 5 random seeds against sample size; here, we see that the three-layer network has a better sample complexity than the two-layer network.
[(x) g(x = 32
6.7
2ayyr Deonds
1aysr screcs
66-
65
Test Loss
64-
03-
022
61-
68-
10'
18'
Samples
Figure 2: Three-layer neural networks learn the target fx( = ReLU(xTAA Az) with better sample complexity than two-layer networks.
Experimental Details. Our experiments were written in JAX 14 and were run on a single NVIDIA RTX A6000 GPU.
B
Notation
B.1 Asymptotic Notation
Throughout the proof we will let be a fixed but sufficiently large constant.
Definition (high probability events). Let e C log (d)) We say that event happens with high probability if happens with probability at least 1 - poly(d, n, mi, m2)e
Example 5. If (0. 1) then |z| S /2i with high probability.
Note that high probability events are closed under union bounds over sets of size poly(d, n, my m2). We will also assume throughout that S C-1d.
B.2 Tensor Notation
For a k-tensor E (R2)3k, let Sym be the symmetrization of T across all & axes,
Sym(T) in
1
Σ T Foayrraso
k!
oeSt
Next, given tensors A E (Rd)6a B E (R))8), we let A 2 B E (1Rd)aat be their tensor product.
15

---

Definition 5 (Symmetric tensor product). Given tensors A E (RR))0, B (Rd)&), we define their symmetric tensor product A E (TR)) as
A == Sym (A B).
We note that & satisfies associativity.
Definition (Tensor contraction). Given a symmetric k-tensor (Rd)*k and an a-tensor A (IRd) 80, where k 2 a, we define the lensor contraction T(A) be the k lensor given by
T(A) (a)ddd
Σ
Sa be- Asamnda
When A is also a k-tensor, then (T, T(A) TR denotes their Euclidean inner product. further define ||T|| = (T.
C Univariate Approximation
Throughout this section, let μ(b) == exp(-bP/2)
denote the PDF of a standard Gaussian.
1/2m
Lemma Let Unif({-1, and let N(0,1) 1). Then there exists v(a, b) supported on {-1,1) x [0,2] such that for any |ar] S
Eu,blu(a, b)o b)] = and sup
(e(a b)] 1.
a.b
Proof. Let v(a,b) = claE(1.2] where c= [(1)- Then for |0|| 1
[a, b)x b]] {
2
fo(x + +b) +o(-x b)]u(b)(b
= c / bb(bb
1.
Lemma 2. Let ~ Unif((-1,1) and let ~ NO,, 1).
Then there exists v(a, b) supported on
{-1, x [0,2] such that for any |x] S
Ea,//((a, ()x + b)] = x and sup
(u $)155 1.
a.
Proof. Let v(a, b) calse(112] where
To n(b)db Then for |x| < 1.
Ea,s/o(a, b)r +b] +b]]=c = {" 1o(x + b) o(-r + b)]u(b)dbb
cr
pbb)db
Lemma 3. Let f. 1-1,1] R. be any twice differentiable function. Then there exists v(a, 6)
supported on {-1,1} [0, 2] such that for any |x] < 1,
En,o/v(a, b)o(ax 6)] f(x) and sup (((a,))]= X sup -
(f()(()))
a.b
(0
16

---

Proof. First consider v(a.b) = 2 Then when & 2 0 we have by integration by parts:
+b]
I [1"
L. (((b))bb f( (((b)db
-0'00x= -0000=
[(x) + +2 02
where C1 = -
and c2 - (-(1). In addition
when x < O.
Easl ()o + b]]
! b) (00
=f((0)= (((0)=-(((1(++1]])) f(((==+))))))) [ (((b)) [ f f((jdbb
f(x) + c
so this equality is true for all By Lemmas and we can subtract out the constant and linear terms so there exists v(a, b) such that
B (va (ax + b)] f(x).
In addition, using that μ(b)= 2 for E [0, 1] gives that this v(a, b) satisfies
sup (e b)] $ |cal + le2l max (x)) sup
(x))
a.b
wE(-1.1)
zE -E1-1,11
kE (0,1,2)
D Proofs for Section 3
The following a formal restatement of Theorem I
Theorem 6. Select q W2.000 ([-1,)) Let th min. n. and 2 ||K 2. mm 2
||Kf* m2 Z |K ( There exists choice of ग O( x poly(d, n, mı, m2; ol and X, 12 such that with high probubility the outpul of Algorithm I satisfies the population I loss bound
T
Ex (I(=:))) - -1())))))()) s olali, | (
20ir +2x+1
b
n
m2
m1
1aa12.
1 1,000
14+1
+ l1q ( - Kfr)
f=11 +
+
m²
Proof of Theorem 6 The gradient with respect to wy is
Tn (0(0))
1 Σ n
((re) Iwy 0(0)
(1(00
1
Σ
- ((x))
10) (x)
n :=1
mi
(0)
11,000
1
(f(I;;(0)) - f (xx))
(1))
m1
iml
17

---

Therefore
wj 111 Va, w, (0(0))
(1)
n
77 (0)
Σ
(ax)) [ (10s; 11 (xr))
1 0000000
m² j
n
One then has
f(z; (a, W(), W()
b(1) v(0)))
1 100
T Σ
(((x)) (z)) + +b,)
(x))
m2
1=1
o (ag
(0) -T(())))) + bz) 100000
m1
where Φ(z) = m EE= (f:) (z,; (()) (1( (x,), (x))
The second stage of Algorithm [ is equivalent to random feature regression. The next lemma shows that there exists a" with small norm that acheives low empirical loss the dataset D2.
Lemma 4. There exists a" with ||a" lel S 11 vmi such that 8" = =aa", (a",W()),b (((), v(0))
satisfies
(20+2a0+2xt]
22,
L2(0") S 1111l1,001|
4
m2
m1
2 +
1 + 11 (1- kff-Kf()-- - f
n
The proof of Lemma is Kf'', and then invoke the ReLU random feature expressivity results from Appendix C
is deferred to Appendix D.I
We first show that @ is approximately proportional
Next, set
(2ett2aot2xxt)
20++2x+1
= 112
[9l1i, 1B
+
12
m2
m2
26+1 +
111l2,00'
+ 11a (
n
mi
so that (2(0)) S ||a 11
Define the regularized loss to be i(a) La((a, WW,, b00), v(O)) arg min, M()), and (!)) be the predictor from running gradient descent for steps initialized at (0)) We first note that
1all Let a
)(al))) $la')\
Next, we remark that i(a) is A-strongly convex. f(; (a, W(),, b(0), v(o)) at (x)( where ('xx) Vec
Additionally, we can write o (a 9 mo(x) + +b)0) 1 100000 In
mi
Appendix D.2 we show supxED, [|x((=]] & Therefore
Amaz V
2
Σ DE
||(())11
x m1
so L is O(2) + smooth. Choosing a learning rate 7/2 = O(1)), after T = 0 -(x2m,) =
Xmny
poly(d, n, mi, m2, [ steps reach iterate à (T)) 0 (ā, (()), b(0), v(1))) satisfying
L2(0) & L2(0") and ||اا|| & |1 112-
18

---

For T > 0, define the truncated loss + by (z) = min(=?, (2)) We have that (+(z) S =², and thus
Σ 4 (1(1;0) - ((()) - L2(
T D5
Consider the function class
F(B,) = {f(: (a, W() (10,,,)) v(o)): ||a||2 S Ba)}
The following lemma bounds the empirical Rademacher complexity of this function class
Lemma 5, Given a dataset DD recall that the empirical Rademacher complexity of F is defined as
1 sup - |EF T2
Σ aaf((z))
RD(F) = BoE(+1p
Then with high probability
B² nm1
RD. (F(B)(())) S
Since & loss is 2r-Lipschitz, the above lemma with B, = O(1aa1122) - 0111112,.0 V
along
with the standard empirical Rademacher complexity bound yields
Eyye (f(1;; - (x))
- n
Σ e
- -(()) + RD + x2 V
xED2
(22t2200t2xt)
220+2x+1
Iall
|K |f
m2
m1
2641
1u11200
+ 119 (T Kft) - f*1 +
n
m2
+ +7V 11011
V=
Finally, relate the &, population loss the @ population loss.
Lemma 6. Let O(max(i', (l Then with high probability over 8,
11112,.0
Ex -)))))))(()))))) < E,e, Exe, (1(1;)))-1000
m1
Plugging in r and applying Lemma66yyelle
Ex (f(x:)0) ((()) x I ( (22tt2ast2xx))
2a +2x+1
1
m2
m2
11112.00
114
12,
1 + 4e+1
+ 11 Ks*)
+ V
m1
n
as desired.
D.I Proof of Lemma
We require three auxiliary lemmas, all of whose proofs are deferred to Appendix D.4 The first lemma bounds the error between the population learned feature K, and the finite sample learned feature over the dataset Da.
19

---

Lemma 7. With high probability
/
p2t22rt1
12,+1
sup (d(x) - (K)) rED,
m2
m1
The second lemma shows that with appropriate choice of 17, the quantity no(x) is small.
Lemma 8. Let 2 |K 2 22222.+1 mm 2 |K and m2 2 |K 1 1, 20
There exists 7j = [f (2) such that with high probability, supsEd, [T(()) I 1 and SuPzED, [7] (x)(
The third lemma expresses the compositional function q o TJ via an infinite width network.
Lemma 9. Assume that n, mi, m2 = 1([KKff(173) There exists v {+1} x R - R such that &u (o(a,b) b]] S |1al|,,00 and, with high probability over D22, the infinite width network
f0 (x)) == ()1 + b)
satisfies
f = (()==)(7((z))
for all x. E D2.
Proof of Lemma 4 Let v be the infinite width construction defined Lemma 9 Define a* E R
to be the vector with a² v(a,)) (f.)). We can decompose
L2(0r) = Σ
-) (x))
ZED2
(f(z;0"- (z)) +
(f) a(T)(z)))
xED₂
FED2
1
(g(To(z)) - {(T((())))))
(IT((())) - fz2
zED2
Take mi 2 The first term is the error between the infinite width network (08 and the finite width network f(x;; 0*). This error can be controlled via standard concentration arguments: by
Corollary we have that with high probability ||a*|| S || and by Lemma 17 we have
ymi
1
Σ RED
m
Next, by Lemma 9 we get that the second term is zero with high probability. We next turn to the third term. Since is [|q 1, - -Lipschitz on [-1, 1], and supzeD, [T0(x)] I s,supee. 1 (x)) < by Lemma g we can apply Lemma to get
sup (x(70(z)) ( ((7(K))))))) < |19l),,00 sup M0(x) 7(Kf"))))
xED,
=ED;
12001
11al2,.00 77
V
7
111
mi
(12et2o0t1
200-1
| [K |11a11,,0
n
m2
mi
and thus
( (,20420, +2x++
+2x+1
Σ
- -(((())))(((
<
||
TV
m2
m
20

---

Finally, we must relate the empirical error between o fj- Kf" and f to the population error. This can be done via standard concentration arguments: in Lemma 19 we show
Σ ( l - + l1 + 22+1 2ED2
Altogether,
(22420
200
|/ l1a +
2241
11911,.0
L2(00) S [olli,oo|I((1
_
m2
mi
n
mi
+ 1q (T((K))))
-
< |K
[2
+2x+1
1
+
26+1
m2
m
+ 1 (T(K))))
f'1 L2
since
|K]] op $ and thus |k
22 2
D.2 Proof of Lemma
Proof. The standard linear Rademacher bound states that for functions of the form ++ w" (/)) with || S Bu. the empirical Rademacher complexity is upper bounded by
Bw
|x(())11
RED
We have that f(x;0) = a (xr),, where
- (a, (0) 7)(x)) + +b)))) -1000
((x) = Vec
m1
By Lemma with high probability, sup [T(((a) S Thus for x D2
1b,1) mi
m
with high probability by Lemma 10
Altogether, we have
Ba n 1/m1
B2 nm
Ro, (F(B.)
D.3
Proof of Lemma 6
Proof. We can bound
[f(=:0) -(())) -kx ( -((11)))((()))
Ex
(((()))))(((())))))))))))(())) -M(s)) - -110))))
En
(((=)))((m)) 1
E.
1/2 (r(=:)) - -5()) '1
(P(111=50) (mx)) 25)
1/2
E,
[(:;0)]111 1/2
+ +B, (1(0))112 (0(1(1:00) 22) + +P(5(()))
E2
21

---

Next, we bound f(; 8):
f(z;; 0] =
aoo (ag)) -nδ(() +b, 100000000
m1
m1
Lä [0 +bby)
m1
||áll2 vm
In)(a))
+ m1 |اا|اااblا
By Lemma 10. with high probability over the initialization we have [6]2 S v mi. Next, by Lemma 4
with high probability over the initialization and dataset we have |lāll2 2 a 11 S 19112,00 mmj. Thus with high probability we have
[(x: 0) 2 10 (mo(x)) +
uniformly over æ.
We naively bound
[o(x)| - _ 1 (x))
1
(x))]] (1)) (az)
m2
By Lemma and Lemma with high probability we have [[ 010)) fxx)) S cl for all i. With high probability, we also have [1:00 (xx)] E Vmg. Additionally, we have
2 02 ( 12)2
h
(w)]
Σ
j=1
S m2
02 11)1
m2
x m2
|x 2y )on)
Since æ is C. subGaussian and
||a||| = 1, we have
Ep
10) (
S m EE x 2i t00
S m
Altogether,
[x(())) $ S1 1e
and thus
E2 [(c=:)) 0) 11 S 111123 2.2 < $11a112000
By Assumption 2 Ex [(*())) 11
We choose O (max(1) l1 By Lemma 7nd Lemma (/ & are both high probability events. Therefore by choosing C' sufficiently large, we have the bound
o) > +} and ([5*(t)) 7}
[ 27) P(I(((())) +)] S 222
Altogether, this gives
(1(o:)) - E [er ((a)))
11912.000 |K 11 2
1
mf
m
m
22

---

D.4 Auxiliary Lemmas
Proof of Lemmal
7
We can decompose
16(x) - (x))
1
1
(1(m;:010)) - (xxs) (x,) - (x)
m2
i=1
f(x;; (111
M2n
(x.), - ( (z))) (0)0
(2:)
m2
1
(x))
(100 ((KK))))))(()))))
+
m/
By Lemma
and Lemma
we can upper bound the first term by -
For the second term, by Lemma 1 and Lemma 13 we have for all æ E D2
2
(xx))) - (1<)))
[1)
2
1 m
m2 mzi 20+224224+
2
m2
m2
11
12
The third term can be bounded as
m2
20, +1
Σ
E. Lf 0z)02(x- Ex v)02(x [i]
$
m2
12
Proof of Lemma 8
Conditioning on the event where Corollary I holds, and with the choice πe
K e-
|K
l -x, we get that with high probability
1
sup (77 Kf)) S 2
mED
Therefore
sup M)(()) Cx ex | rED:
11 x sup [d(x)
122220001
2aa
$ C E -- | 1 sup (Kf*)()))+
n
m2
m1
12222200
20
$ +0
|IK]
n2
m2
m
1.
Proof of Lemma 9 Conditioning on the event where Lemma X holld and applying Lemma that there exists such that
we get
E
b)] ((T)((=))
for all x E Dz. Since v(a, b) for < O, we have that v(a, b) v(a, b) Thus the desired
claim is truc.
D
23

---

D.5 Concentration
Lemma 10. Let my 2 L. With high probability,
(0)
1.
m
Proof. By Bernstein, we have
1
m2
(0)
2
2
11.
m
my
1=
Lemma 11. With high probability, (xx;010)) S v in
mi for all
Proof. We have
m1
f(1; ((0
a, (b.)(b)).
Since O2 ~ Unif([=1}) and b, ~ (0, 1), the quantities a,o (b.) are 1-subGaussian, and thus by Hoeffding with high probability we have
f(1; S m1
Lemma 12. With high probability
sup E. [x
u))] 1 and
Ez (002(x100)) S 1.
sup jE(m2]
zED,
Proof. First, have that (2 w) < 1 (a v) la Therefore
[o(o2(1 e $ 1 + E, (x v)
S 11 + || d 20. 220,
Next, with high probability we have sup we can thus bound
|| 2 - -|1=1| S uC2 S Since E||=|| S dr Z
sup E, [02 D2
S ())
das
sup ||0||
Te
xED,
d 201 (d + 1) 12200
x 1.
The proof for the other inequality is identical.
Lemma 13. Let m2 2 21, With high probability, sup (x)] i /mz
Proof. Since x is C.,, subGaussian and v is 1/v subGaussian, with high probability we have x u O(C,.) = O()).
Union bounding over a E D2,j [ma], with high probability we have sup pp-0yl S O()). Therefore
sup 1226 uz)1 S O(1²))
2,0;
24

---

Next, see that
1
m2 h 100
(x))
02 (r - 0z)²
m2
m2
j=1
Pick truncation radius = O1*)). By Hoeffding's inequality and union bound over x E D1 U we have that with high probability
m2
22+t1
2 E.y [ 10 u) 2
sup zED,
JD,
m2
j=1
M2
Observe that
Eg (o000) u)
E. [o(o []]]] E. R00)) v)?]
S (o2(( w)² R)B[o2(x u) 1/2
$ d'
Therefore with high probability
1
h
(0)) (x)))) (x)
E. [o[o v)
[20
m2
m2
d'
by Lemma 12 Lemma
also tells us that E.g [az(x r) 0(1). Altogether, for m2 2 200
we have
(z)) <
m2
and hence || (x))) S /m2, as desired.
Lemma 14. Paw (lf* Cye(") Se and P Pa ([Kf))) 2 Ckei |Kf S
Proof. By Markov's ' inequality, we have
MM
P((f*()) > = (x) > 19) =
t
t
Choose Cyec" We select g アー which is at least for C' in the definition of sufficiently large. Plugging in, we get
Gjee Cheat = =e -do Set
P(()(<)) Cseet) <
since le!/t-1 21
An analogous derivation for the function
K Tf yields the second bound
Corollary 1. With high probability, sup.eD, [f*()] < Cye of and supzepa (x)] x Cxe- x*||K
Proof. Union bounding the previous lemma over E D2 yields the desired result.
Lemma 15. With high probability.
8
(x)) E. (F(]))
2
72
2 n
25

---

Proof. Consider the quantity
1 (2))2222; - En ( (m)))
With high probability, supreDaje(my (x uj)l S tta Pick truncation radius Olitar). ). By Hoeffding, we have with high probability.
264
(x,)o222 15) (as)))))))) B [ uy) ((s))
Furthermore, note that
E. (z)o2( - 2,) (o)a -., - -0)]] ([ (z)0 2,)] -2)]]
P(f U3) R) 02 uz)a]
Thus
" Σ f (x:))22(2 0,) Ex [Ma)) (
$ 2e+2as +1
By union bound, the above holds with high probability for all j E [m2)]. and thus
$ m2 2
(x,)) E. [()))]
Lemma 16. With high probability.
m2
Ew (1(')0a) uz)02(x uz)l - (x)))(( (o0000)200 v)02(x000)) S 22a
sup xet
m22
112
Proof. Fix D2. Consider the random variables Z(v) = Eg (x')==xz22(( u)02(1 v)]. With high probability we have [02(2 u2)l S to and thus
I(,z)] = E2 (f" 0z)02(1 u,)ll 2 (h E. [f')' o,) L0. For all j E [m2]] Choosing truncation radius R t" with Hoeffding have with high probability that
200
E. (f 25)lo2(x 1,)1 z(n)ER - Er [f (z')o v)o2(1 w) (())))))))
m2
m2
Next, we have that
[ftrv v)02(( - 4) - Er (fl" (oo((a)))))
= Ea (f*(x') v)02(# u) )|
((Z(t) R) Er [f (x') ozl v) 02 w) 1/2
24 m2
Conditioning the high probability event that supje(ma) [Z(r;)] S R, we have with high probability
that
20,
E. (x')02z (x))oz(x. (.)(x--)))))) vz) E. Er (f 2)02(1 (3)2]()-))] S
m2
m2
1=11
A union bound over D2 yields the desired result.
26

---

Lemma 17. With high probability
mi
|ell
sup D)
v(a,, (()0(( + b,)16.>0 (x)) x
m
m1
Proof. Condition on the high probability event supzED, [7)((2) S 1. Next, note that whenever b > 2 that v(a,b) = 0. Therefore we can bound
bo)a ( (a7)(2) 1o, - 2||e||
Therefore by Hoeffiding's incquality we have that
|ol| ' m1
v(as, (a=)(()) + b;)16, (x) S
The desired result follows via Union bound over D2.
Lemma 18. With high probability.
|
Σ u(a,,,,)2 -
1101
S
m1
Proof. Note that
v(a b,)2 - 111113
Thus by Hoeffding's ' inequality we have that
mi Σ
|llll m2
b,),,
1101 2
m1
Corollary 2. Let mi Z Then with high probability
v(ac b,)2 S |1 m1
Proof. By the previous lemma, we have that
|1ll mi
b,), - 111112. $
S |101/3-
m
Thus
mi Σ
o(a; b,) S |1uliam + 110112 mi $ ||/| mi-
Lemma 19. With high probability.
Σ (()(K)))))){(((()))) - s o - 11a1
20+1
27

---

Proof. Let S be the set of x so that (T(Kf))(2) & 1 and [f*()) & t.. Consider the random
variables (I(T((K))())) f())² 1xES. We have that
(((((K/))))))((( 1EEs x sup xEI-1,11
lq(2)| + t
and
(((((((()))))) lees] =s z
|a(2) +220 - 1111 (T(({*)) fr 12
sup
=c|-1,1]
Therefore by Berstein's 's inequality we have that
Σ zEDs
(=)) - (x)) 1 - 11(9 (7(Kf))) - f 1
T1
(supze(-1.)|((<))1 11(a o (T(Kf*)) - ff f*) supzel-1.11 lq(z)) 20
S ll (π(K*))) (") suPse[-1.1) (x)) + ,2
Conditioning on the high probability event that E S for all x E D2, we get that
suP=e[-1.11 lq(2)) + 2
(((T(K)))))) - (())) & l1
- - -
2ED2
11g (7kf')) 1 suP:e[-1.1] la(=)1 + 26+1
n
E
Proofs for Section
E.1 Single Index Model
Proof of Theorem It is easy to see that Assumptions T and 3 are satisfied. By assumption g" polynomially bounded, i.e there exist constants Cg. 2g such that
g²(2) < Cg(1 + +=1)) C2200 - -1 +1++122 -12
Therefore
| 11g (R. S C,2200 + Eg~N(0.1)
(0.1) [| (1//00012//) s C2 100
Thus Assumption satisfied with Cj 20 av 12 and l a2
Next, we see that
K(x,x') [x u)( . w)] x
d
Therefore
(KK)(() E, (f x/d 2 Ear (V xz.
Furthermore, we have
E /1 R. (w g'(w" x')l = w w'
Altogether, letting ci E.- (g' (2)] R(1), have (Kf))(x) -/c ww Assumption 4 is thus satisfied with 1/2.
28

---

Next, see that |Kf 1| c1/d. We select the test function q to be g(z) 1 d/c1 2), so that
((T(K()(()) = g(x z) (x)
Since rj x)), we see that
sup |q(2)] sup
sEl-e(ex),0(x [9(2) = poly(()
=E(-1,1|
sup |a((=) 17 d/c1
sup xE(-A(xx,,((x))
[g'(2] = poly(1)
sE[-1.1]
sup [=)] गÿ 2d / =E[-1,1]
sup
1g"(2)] = poly(1)
=E| =E1-0(x),0(xx)
Therefore bound the population loss as
- (()))
d poly(1)
E
min(n, m1, m2)
vn
E.2 Quadratic Feature
Throughout this section, call x" Ar degree 2 spherical harmonic if A symmetric, IE 1 Ax and E [xx Ax)²] 1. Then, we have that Tr(A) = 0, and also
d+2
E[x 21 (AA 3x2 (2002) 2x2l1 = ||
O11)
V
2X2
2d
See Appendix F for technical background on spherical harmonics.
Our goal is to prove the following key lemma, which states that the projection of f onto degree spherical harmonics is approximately x" Ar.
Lemma 20, Let q be a L-Lisschitz function functioon with 14(0)] < L, and let the target f be of the form (z) qxx Az), where x. Az is spherical harmonic. Let 0₁ = =E=-N(0,) (2)]. Then
P2f - -xxx Ar L3 $ Lx 1/12 log d.
We defer the proof of this Lemma to Appendix E.2.1
As consequence, the learned feature K approximately proportional to xf Ax. Lemma 21. Recall e E== Then
|| - (2(())11" Arll S 2/6 log
Proof. Since [x)* 0, Pof 0. Next, since f is an even function, P&f 0 for & odd. Thus
| XB(0)) $04.
Additionally, by Lemma 20 we have that
Pas" axx Arl $ T2 ai All 1/12 log d.
Since 23(c) = O(d 2), we have
Kf - 25(0)a1=" Azll S 2x21/6-2-1/12 log d.
Corollary 3. Assume o(Vd). Then
Ax [DK||KKf* 1 < 1/6 log
29

---

Proof.
-² Az - |Kf Kf
12a
|||]/| || 12" Ax | 1|2 - 1610
||K ||اا Kf - (2( 10c1=AAx/ll
Axl11
|<< - 22(0)
X Lx 1021//d2221100 log /
Lx logd.
Proof of Theorem 3
By our choice of v, we see that Assumption T is satisfied. We next verify
Assumption Since f* is I-Lipschitz, we can bound (z)( S Ig" (0) [a], and thus
Ex [9" Ar)
(0)) Ex ||- |a=Ar| Ar
1/9
19 (0)) +
(1 g (0))g
where we used Lemma 35 Thus Assumption 2 holds with l 1.
Finally, we have
Kf = Σ Paff
A22
By Lemma 21 we have ||K]/" L2 2 (o))e for d larger than some absolute constant. Next, by
Lemma 35 we have for any g - d2
||K S
/
622
(Vg/d)*
q
d
/qd
S 2qd 2
Therefore
|Kf S c @ |f S <f
since xā(02) = ()d Thus Assumption holds with f= 1.
Next, observe that |T/f"|1> E d=2. We select the test function to be g(2) 0" (K 2).
We see that
(((({))())) = 9" (IKM
and thus
9 (1. Ar) 9" (IT]I
2 1 Ax |K | 13
2 2/ 1/6 d 1/12 log d,
30

---

where the first inequality follows from Lipschitzness of g", and the second inequality is Corollary Furthermore since f xx), we get that 77 |K @ and thus
sup 1g(=) = sup
1 poly(()
=E|-1.1]
=E(-6((2),8(23)
sup (x)) ... |Kk sup
(()))((2) poly(t)
=E|-1.11
=E(-6(cx),((x))
sup d"())
(
sup xE (((-0())))()))
(g)) poly(1)
xE[-1,1]
Therefore by Theorem 6 we can bound the population loss as d' poly(() min(n, m1, m2) vn
Eg f(x: (*) &
1
+ 13133-11101)
E.2.1 Proof Lemma 20
The high level sketch of the proof of Lemma 20 is as follows. Consider a second spherical harmonic x² Br satisfying [[(=""Aa Az) (x)] 0 (a simple computation shows that this is equivalent to Tr(AB) = 0). We appeal to a key result in universality to show that the large limit, the distribution of x. Arr converges standard Gaussian; additionally, x!' Br converges independent mean- zero random variable. As consequence, we show that
Elg(#" Ax)x² Ax] x [7)) = BzwN((.))ld(
c1.
and
Elg(=" Artx. Br) Ru Elg(x" Ax)] B[x" Br] = 0.
From this, immediately follows that Paf* Rs cxxx Ax.
The key universality theorem is the following.
Definition For two probability measures p,u, the Wasserstein 1-distance between μ and is defined as
Wi(p,t)
sup
-Es
1 S1
where Ilfl = sup [f(x)-f()) is the Lipschitz norm of f.
1--vil
Lemma 22 (52) 9.20). ] Let z No La) be standard Gaussian vector, and let f: d R satisfy [=)( = 0, [[([()))}] Then
W1 (LLw 1)) [IV)(>)] 1/4 E [[v)
1/4
where W1 is the Wasserstein / Iddistance.
We next apply this lemma to show that the quantities x. Ax + xx Bu and x. Ax. + x u are approxi- mately Gaussian, given appropriate operator norm bounds on A, B.
Lemma 23. Let xT Ax and T Bx be orthogonal spherical harmonics. Then, for constants C1, with e + + 1, we have that the random variable Y = cxx² Ax + o2x Be satisfies
Wi(Y,N(0,,) & ||Allop + B|pp
Proof. Define the function f(=)
c FAg 1=
+ cnd cnd TB and let 2yd Observe that when T=
(0, have x Uni Therefore f(z) equal in distribution to Y. Define
[=)) = d d= |||
We compute
Az
V=(
2d
|
||
31

---

and
A
2Az
222 A
TT Az ||
x" A2222
(2 = 2d
+ 4
||=1|²
||=||
||=||
||=||
Thus
||Az||
1²" Az
vd
12- Ax
[V
< 2d
||
||=|"
||
|||||
||
and
v2
S
||A
||
||=|] distributed chi-squared random variable with degrees of freedom, and thus
E
||
2k
1
IT (d 2j)
Therefore
E
| (+)
1/4 X
d||A|| op
B
||| [1]]
1/4
S ||
and, using the fact that and ||=||| are independent,
(+))))))]]] 1/4
I vdB ||||| 1
E
|
1/4
+
E
|||||
9 1/4
[[( Ax) 1/4 S1.
As consequence, we have
E [IV 1/4 $ and E
1/4
x Al ||B
Thus by Lemma 22 we have
Wi(Y,N(O,1) =w(r(), 2 | + B
Lemma 24. Let xx Ax be spherical harmonic, and |ud] = 1. Then, for constants C1, C2 with
c7 c have that the random variable Y C1æ² Ax czx u satisfies
WI(Y,N(O,11) Z |
where W1 is the 1-Wasserstein distance.
Proof. Define fz(2)
Vd=" [xll
We have
v2((2) vāl
||
and
uzt + 2ut
rTu us
2T uzz
f2(x Vdl
I
||
|| 3
|| 5
Thus
[Vf((2)] - Vd
Vd
and
V2 fe())())
x
||>||
|
So
[+))"))
E [[IVF((>)1*
E
S and
We finish using the same argument above.
32

---

[Bllop are small, Az,Ax, x² Bxx) is close in distribution to the standard Gaussian in 2-dimensions. As a consequence, Elg(x Ax)x² Bx] ~ [B]g(2))] [=[=]] where 21, 22 are i.d Gaussians. This intuition is made formal in the following lemma.
Lemma 23 implies that, when |
Lemma 25. Let x² Ax, Ba be two orthogonal spherical harmonics. Then
El9(x" Ar)x Bu]] - L
< (VIIA)) log d + | log d)
Proof.
Define the function F(t) = So (())ds. Then (t) = qtt), so by a Taylor expansion we get
that
[F + ey) - F(x) - eyg(aa) S ²²
Therefore
[B[((= Ar)a" Bu] S E<LE((> Bx)² + I [F Ar + ex Bu)] - B[F [c))]]
Pick truncation radius R, and define the function F(z) = F Fmmax( R, min(R, 2)). F has Lipschitz
constant [u 1=(2), and thus since W1 W (x"' Ar, 1)) - || we have
EF("" Aa:) P()) S sup |ę(=)| ||Allop
=E|-R,8
Next, we have
E F("" Az) - F(="Ax)) Az)]
S 1|2 F( Ax) s P( Ar| > R) B[F B[F(x"AA))
Likewise,
E. F(2) F(=)] $ [[2] R) (
Since y is L--Liscchiz, we can bound [[=) < |=]1((0) + = and thus
E [[F(=)* - L and B[F Ax)²] S L.
The standard Gaussian tail bound yields P([2] > R) S exp( (C for appropriate constant C1 and polynomial concentration yields PP( 2=" Ax > R. S exp(-C22 C2 for appropriate constant Cz. Thus choosing R C3 log d for appropriate constant C3. we get that
E Fx' (x T Az) F(x" T Ax) + E2 F(2) - F(=)]1 2 d'
Altogether, since 1g(=)) S [q(0)] + L|2l. we get that
|F Ax) - Es.Nico,)) F(z) S
By an identical calculation, we have that for < 1,
EF(x" Ax + xx Bx) - F...o..)(av1++....... $ (1A41oop + =([BBllpp) L Llog d
Altogether, we get that
[F= Ar + cx Bx)] E[F(x*AA))
< (IAlllop <||B|l.p) L logd (=-)) -e.vrroo///////1++))
Via simple calculation, sees that
F(aV1+ - =()) [c(0)]= (110))ll1(11111111
L 2² - L|z|e² + L=262
Therefore
[F Ax + ex B&)) - F( Az)] S ((111llpp + <||B1l.p)
Ax
log + +
so
E g( Az Az)x Bx] s L(e (1All logd + | |BBp logd).
Setting e = VIl op log d yields
B [lq((= Ar) Bx| S |1Al log d B| log d).
as desired.
33

---

Similarly, we use the consequence of Lemma 24 that (x" Axx u² u²x) is close in distribution to 2d
standard Gaussian, and show that Elq(2* Ax) (n"x)² - 1)] ~ (1 -1] = 0.
Lemma 26.
B[g(=" Ax) - X || 1/3 2/3 d.
Proof. Let G(t) S/ F(s)ds. Then (t) = F(t). (t) g().. , so Taylor expansion yields
G( + ( = G(x) -00
G(z cy) G(x) ++
Thus
()) = 2 [G +z) G(x - - -
Therefore
[[[((" (x) x)) [x)))]] 2 eT. + 6 2 Ax eu x)) + G(x" Ax - eu x)) 2G( Ar)]l
For truncation radius R. define G(z) = G(max( R, min(R, 2))). We get that has Lipschitz constant slp [F()) x LRR Therefore
EG (z)) Eo,, X LR AA
and by similar argument in the previous lemma, setting R Cj log d yields
E G Az) G(x² Ar)] [B,((2) G(=)] I
Altogether,
EG(xT (G))TAx)= -E,((z)] 2 || log d.
By an identical calculation,
BG(x Ax + cu x) E.G((vI + 2) & ||Allop Llog d.
Additionally, letting be independent standard Gaussians,
E [20 - -2G(z) R [G(= + ew) + G(z ew) 2G(=))
E.,w (q<=)/) O(eL)
E.[u(=)] + O(eL).
Altogether,
B[g((" Ar)("² Ex(g(>)]| & eL + 21A111 L log²
||| 1/3 1/3 d,
where we set e = || 1/3 2/3
2/3 d.
Lemma 25 shows that when |B|| op < 1, Elg(2" "'l)(("'A)='" Bx] However, we need show this is true for all spherical harmonics, even those with | O1). accomplish this, we decompose B into the sum of low rank component and small operator norm component. We use Lemma 25 to bound the small operator norm component, and Lemma 26 to bound the low rank component. Optimizing over the rank threshold yields the following desired result:
Lemma 27. Let A, be orthogonal spherical harmonics. Then
E [g( Ax)xx" Bx] < /| log d.
34

---

Proof. Let r > || be a threshold to be determined later. Decompose B as follows:
B =
d
T
Σ da
-
1
Σ |A:
of Au: A ++1 B,
|
14,2
where
B
Σ Aauru T +
Σ
d +
Σ
Auce A.
d
||
14:127
By construction, we have,
Tr TT ( = Σ
+
>==
A
iE(d
and
(B, = Σ Anu Au
Au +
Σ
douf Au; B) = O.
D
Therefore by Lemma 25
B aataa Ar)x T Bu]] S L VVAAllop log [ F
log d.
There are at most O indices i satisfying Aull > T, and thus
Σ PrST
M Σ 1 ST-
We thus compute that
2
B
2
Σ
X
+ d
Σ
d
Σ
deuf Au;
F
Aist
x
Σ
x + |A/l
Σ DENT
A
1
2 ||
S 1.
and
ST+ +
+ 1|Allop +11|Allopp
Σ Aoug Aui |A:120
S r + ||A1| 7
Next, since
[[[a(=" Ax)] - [x[((())]
S /| Lemma 26 yields
q(0" Ax). Σ de (muuf" - -
|al [a(a" Az) wt) - -E([-E((de E[[(a" Ax)
A |
Pal |||
1/3 1/g2/3 d op
| 1/3 log
2/3 d.
35

---

Finally,
B q(a Ax)
1
E
Aae Au; xf Ax
Z Σ
Aus x L11
||Al
Dalls
|adls
Altogether,
|[[((" Az)a² Ba] S log d (11411022 + + |1 +
|Allop + 114llpp=71
S T1l log d.
where we set T = 11
Finally, we use the fact that x.' Ar is approximately Gaussian to show that Elg(x" A Ax) Fo ct
Lemma 28. Let x² Ax be a spherical harmonic. Then
|Ela(" Ax)x Ar) -c x L|| log
Proof. Define H2) = (())2. For truncation radius R. define H2)) = (maxx(-aaa) - R, min (R, z))).
For v, [-R, R] we can bound
[H(x) - H(v) = 11(x) - (v)v
< (=]1((()) q(v) + la(v)]l<> vll
2 RLIl - ull-
Thus 77 has Lipschitz constant O(RL). Since Wi (2 Ax, 1)) S | we have Ax) Ax)
1 (T(=) & RI||
Furthermore, choosing R = C log d for appropriate constant C, we have that
H, Az) H Az)] P( R) EH Ar) 2
L [x] H1 F(1) R) {[ d.
Altogether,
Ax)) - [{))]] 2 L| logd. d.
Substituting Hx (x) gx Ar)x" Ar and E [=)] = B.(a(z)=] = E=]g(=)]] Ci
yields the desired bound.
We are now set to prove Lemma 20
Proof of Lemma 20 Let Paf x T22. Then Paf a Ar
12 & [T -cAll
Write T2 a.A A where (A- 0. We first have that
(x)* A x] = E x T TT x- A (x, = 2x2l|
Also, by Lemma
we have
E [ (x)x) A
x R [q(x"
Am) T A x
||A F. T.I1-111// 2/6 log d.
Therefore || -
|||
1/6 op
log d. Next, see that
E [ (r)xT" Auj = 20 A = a.
so by Lemma 28 we have |01 al S L|| op
log d. Altogether,
T2 - - 5 |clAc111A11 +|I || z L||Al| 1/6 log d // 1/1/2111g/0 d.
36

---

E.3 Improved Error Floor for Polynomials
When q is a polynomial of degree p O(1), we can improve the exponent of d the error floor.
Theorem 7. Assume that q is a degree p polynomial, where p
O Under Assumption 5
Assumption g and Assumption 7 bound
with high probability Algorithm
7
satisfies the population loss
Ex (
d=
+
min(n, mi m2)
d
The high level strategy to prove Theorem 7 is similar to that for Theorem 3 as we aim to show K is approximately proportional to z" Az. Rather to passing to universality in Lemma 20 however, we use an algebraic argument to estimate P2f
The key algebraic lemma is the following:
Lemma 29. Let Tr(A) 0. Then
A A (Toh 1) =
Ede A A (8B-)((
where the constants dy defined by
dr.s := (2k: - 25 - 1)!!(k 1)!
2-
(2k- 1)!!(k - s)!
and we denote (-1)!!
Proof. The proof proceeds via a counting argument. We first have that
A (100-1-1)
Σ
(A
Consider any permutation a E S2k. We can map this permutation to the graph G(o) on & vertices and k- 1 edges follows: for E k - 1], then we draw an edge e(m) between a and In the resulting graph G(o)) each node has degree at most and hence there are either two vertices with degree 1 one vertex with degree O. For vertex v. let ei (v), (2(v) E [k - be the two edges v is incident to if u has degree 2, and otherwise ei (v) be the only edge is incident For shorthand, let (f,,..., ( (an, ( 01,,... 0k-1, Ork- ().
(2m- (2a- 2a] and (nm) {2b- 2b).
there are two vertices (u1, u2) with degree 1, we have that
(A kk) a ( Aiooe Ajatziugg)
II
oigi M2
Let u1, W2 be connected to eachother via path of s total vertices, and let P be the ordered set of vertices this path. Via the matrix multiplication formula, one sees that
Σ
bect) Secey (A
(A*)
Σ
II
A (ne
WgP
where the sum is over the k a's that are still remaining in (e)(v): u P} Likewise, if there is one vertex un with degree 0, we have
(A @k)
tennooeg Ai, II
Aagto)
and thus, since P = {ur}
Σ
(A
(011) A
Σ
II
Aaey (eyl)))
(011....00-1) ugP
37

---

Altogether, we have that
2 & (8h(118-1)
Σ oESak
A Σ 2 P
(2h))
where 8, P are defined based on the graph (o)). Consider a graph with fixed path P, and let Sp be the set of permutations which give rise to the path P. We have that
A 1
Σ 1 E
22)
SESp
uEP
There (k: 1)--1)... 1)... s choices for the edges use in the path, and each vertex there are two choices for which edge should correspond to 2v or 2u + 1. Additionally, there are 2 ways to orient each edge. Furthermore, there are
ways to choose the ordering of the path. Altogether,
[k==]]
there are 22--1 (k-1)) R! (K-2)] [k-
ways to construct path of length We can thus write
1 (2k) Σ 2- (Ac
1) k!
A A (
Σ
II Aa. Aa. DEP
(k s)! s)=
where this latter sum is over all permutations where the mapping corresponding to vertices not on the path have not been decided, along with the over the unused a's s. Reindexing, this latter sum is (letting (f1,...,izz (2k-2s) = (01,01.. O.k-s:
(tk=s))
Σ
Σ
II
A An = (2k - 2))1A9)-0((10--)))
2(10k--))
22 (an a E(ddtt-
Altogether, we obtain
(2k 25)! (k - 1)! k!
A (100-1))
220-1
A A (100-)
(2k)) (k- 8)! (k 8)!
k!! (2k 2s)! (k: 1)!
A A
(2k) (l: 8)! (I: - =s]
(2k - 2s 1)1! (k 1)!
2
A A A
(2k: !!!!"*
(k s)!
Σ da.s A A s (108--0)
as desired.
Definition 8. Define the operator : Rddd - Rddd by T(M) = M Tr(M) 1
Lemma 30. Let P2f (x) = x² Tzx. Then [T - E. (g)) (x" Ar)] All S Va
Proof. Throughout, we treat p 01) and thus functions of independent of d as (l) quantities. Let be of the form g" (z) -E 26 We then have
( = Σ
Therefore P2f =x""" TT,,, where
To (2k 1)!!k
XK+1 (A (10-11)))))
X2
0=
Applying Lemma 29 one has
T ΣT(A)) Σ
(2k - 1)!lk
Xk+1 1
dk. 1000--000 a (100(k-))).
X2
==0
kma
38

---

Define
P
B. Σ
ax (x (2k:
1)llk Xk+1
dx. 0h ((
X2
k=
We first see that
B1 Σ (2k 3)1!
kag Xb+1 A 88 (70-1(70((-11)
X2
Next, see that
Xk+1
(d 2k) + 2k 2) Xk-1 XX-1 =X-1
2(d+2)
X2
Thus
(2k 3)!! kaxXk 1 ( + 0(1/d) Σ 22k 3)1! k/aa (/0x14@)-1(1S(t-1))
81
k=1
A 8ok 8k (2K- O(1/d)
k=1
E. ((= Ax)] + 0(1/d).
since
A cb (1 (k))
S Eg (x Ax)*1 (x"Ax)*) S Ez (x 4z)
where Tr(A) = 0 implies E, (x Ax = O(1) and we invoke spherical hypercontractivity
(Lemma 35 Similarly, [ = 0(1) and thus
T2 F. [(0))(( Ax) All - Σ IT(A 2 Vd
where we use the inequality
[T(xX] S IIIII +|Tr(XII-
S 2|| ||
Vd
along with
|| S $| S
x
Vd
for 2 2.
Lemma 31. Let c2 = E2 (g"(("")((''' Ax) Then |k - ()) Axll L2 S 5/2
Proof. Since [[]((r)) 0, Pof = 0. Next, since f is an even function, Paf = 0 for k odd. Thus
| - 1lo
Additionally, by Lemma 30 we have that
|Pf q Ar 12 S Th c1 Alp z Vd
Since \²(0)= ( (d)2)), we have
|K - >3(0))c (.ll))cTTAallpp S kd 5/2 5/2
Corollary 4. Assume i = ol o(vd). Then
| |r Ar - 1KK k\va
39

---

Proof.
x² Ax
|K | 12
|Kg 11 12=1 Ar Ax - Kf
||K/ 11 |&
(0) -x | 1c1l
x kd kd2
|
S s\/Vd.
The proof of Theorem Theorem
7 follows directly from Corollary in an identical manner to the proof of
Preliminaries on Spherical Harmonics
In this section restrict to Xa- (v-)(Vd),, the sphere of radius vd, and the uniform distribution on Ng.
The moments of are given by the following 18
Lemma 32. Let Then
82k
X (2k 1)!!1&
where
X
-(1).
d + 2j
f=
For integer 2 0, let Va,e be the space homogeneous harmonic polynomials on gd of degree restricted Xz. One has that Va,e form an orthogonal decomposition of (2²vv) (7), ie
(2(v) = ①
Va.e
t=
Homogeneous polynomials of degree l can be written as T( for an e-tensor T E (RdySt, The following lemma characterizes Va.i:
Lemma 33. Ttx Va and only ifT())
Proof. By definition, degree homogeneous polynomial p(x) E Va, if and only if Ap(x) 0 for all Sdd1 Note that V²p(x) = ((I 1) so this is satisfied if and only if
0 = tr ((2 8(- & I (T(T), (1-2))
As this must hold for all æ, this holds if and only if T([)
From the above characterization, see that dim( Va,)) [3d, k). where
B(d, k) 2k + d - 2 k + - (k: + d-dd3300 -3))((2k + -dd 2) (1 + oa(1)) b kt(d 2)!
3: 1-
Define Ps: (1) - (v) be the orthogonal projection onto Va,e. The action of Po, Pı, Pa on homogeneous polynomial is given by the following lemma:
40

---

Lemma 34. Let E (R2) be symmetric 2k tensor, and let p(x) = T(x be a polynomial.
Then:
Pop Xx(2k- (T
Pip- = 0
1:22 1)!!x++1
-((sh-1)- -T(Toh) 2) xxx"
Pap-
X2
Proof. First,
Pop E[T Xx 1
Next, since p is even, Pip = O. Next, let Pap = T2T For symmetric B so that Tr(B) = 0, we have that
B [[ xx T Bx] Tyxx" Bz).
The LHS is
B(T(r x Bx = (2k + 1) (TBB) @K-11
2k
(2k + 1)!!
2k + (T (T(T@K-1), B)
2k (2k - (T -T(IOK) _ B),
I
where the last step is true since Tr(B) = 0. The RHS is
E r T Tazx' T Br] 31!x2 (1
-2x2(T2, B).
Since these two quantities must be equal for all with Tr(B) = 0, and Tr(T:) = 0, we see that
k(2k 1)llx++1
(((BAA-1) -T(Tak) 4
T2
X2
as desired.
Polynomials over the sphere verify hypercontractivity:
Lemma 35 (Spherical hypercontractivity 11 38)). Let be degree polynomial. Then for q
[l S (q - 1) 11fll
F1
Gegenbauer Polynomials
For an integer d 1, let pa be the density of x e1, where x Uni((S and e a fixed unit vector. One can verify that pa is supported on [-1, 1] and given by
T(d/2) (तI()=1)
x² dr
where I'(n) the Gamma function. For convenience, we let Za
Tid/2) (
denote
the normalizing constant.
The Gegenbauer polynomials (old)
are a sequence of orthogonal polynomials with respect
kg220
to the density pa, defined as Go (x) = 1, G (x) = x, and
(1) d 4 ((d) ( (अ)
k ((d)
+ K _ (x).
(24)
d
41

---

By construction, k is polynomial of degree k. The ( are orthogonal in that
Exzuua ( (1) (x)) B(d
For function E (Maa), write its Gegenbauer decomposition as
f(z)= - Σ B(d, GK (12((a)) G/ (z),
where convergence is in L (Ma). For an integer k, we define the operator PE (2((s)) + (paa)) be the projection onto the degree Gegenbauer polynomial,
p f = (J,, G (PHH
We also define the operators pld) & p(()) and 2k = Li=k p!d) Sk
=
Recall that = Unit(S Let jua be the density of et, where ~ v. For function E (Ta)) we define its Gegenbauer coefficients as
(x(a) = Ea G (d 1/2 ei = (o(d" e1. (
By Cauchy, we get that |^e())) S || B(d, k:) -1/2 0(1 (a
A key property of the Gegebauer coefficients is that they allow us express the kernel operator K in closed form 27 39]
Lemma 36. For a function g E LVV,,, the operator K. acts as
1kg Σ Pk(
E20
One key fact about Gegenbauer polynomials is the following derivative formula: Lemma 37 (Derivative Formula).
d
G(d)
k(A: + 2) (I).
((d+2)
dx
d-1
k-1
Furthermore, the following is a corollary of eq. Corollary 5.
24
(2k 1)!!
(d) (0)
(-1)*
1 1 (d + 2j - 1)
G Proofs for Section 5
The proof of Theorem relies on the following lemma, which gives the Gegenbauer decomposition of the ReL.U function:
Lemma 38 (ReLU Gegenbauer). Lei ReLU(x) = max(x, 0). Then
ReLU(x)
1
(d) (x)
G1 (d)
(z)
[ d-
[d 11)
(2k: 3)!!
+ (11)
B(d, (4,2kk)G (d)) (xr).
3
d-1 2
+ 2j 1)
k1
As a consequence,
|| pld ReLU(x))
Σ
(2k - 3) 2k)
12(Ua)
3 d-1 1- (d + 2j - 1)2
k2m
42

---

The proof of this lemma is deferred to Appendix G.1
We also require key result from product function.
19] which lower bounds the approximation emror of an inner
Definition 9. ʃ (1) (1)
an inner product function if f(x, (((2, 2f) for
some Φ: [-1,1] + R.
Definition 10. g: 8d1 (1) x 1d-1(1) + R is a separable function if gf, = (b( a), (u', ,²)) for some v, v E (1) and v [-11, 1|2 R.
Let Da be the uniform distribution over (1) x 51-1(1). We note that if (Ir, Vd then (r, x²) ~ pd. For an inner product function f. we thus have ( = | We overload
notation and let (p(2))
= -11%%-) (((a))))
12((2))
Lemma 39. 19 Theorem 3] Let f be an inner product function and . gr be separable functions. Then, for any k 2
(IBAflleetua
2 (190ll
9
[|PefllLP(aa)
(P()))
We now can prove Theorem
Proof of Theorem We begin with the lower bound. Let x
J1 x22
where 21, 32 E Rd/2
Assume that there exists some 0 such that || Nollza(2) S.. Then
2² ž E. ( (N(=)))(()))
= Eyaa
F Rnns (WFF) (f No(x))
where r is the random variable defined as = | I| and μ is the associated measure. The equality comes from the fact that conditioned r, 21 and xz independent and distributed uniformly the spheres of radii /F and vd respectively. We see that EEr d/2, and thus
P(rr d/2 d/4) S exp(-S((d)
Let 5 inf (nf/43//43//44 B2, 2 - (VF)
(f) (z)) We get the bound 8. P(r E [d/4.3d/] 3/44))
and thus 6
Therefore there exists E [d/4,3d/4] E B ((FF) -[V (x) - No(x)) < 2/22.
3d/4] such that
Next, that when
| 2 =r, ||12| =d have that
2 (21, = 2/ r(d r(d-r) (T,, (2)),
xT Ax
Vd
where 32 84/2-1 i.i.d. Defining ((2) ReLU (2 (2/rrd Co. thus have
7( T2))
R=. Ax) co fx))
Tj T2
Furthermore, defining I
choosing the parameter vector Ø = (a, W. by, b2), where W
WU Vr-I
yields a network so that N₂(z) = (z()). Therefore we get that the new
vd-r.11
network Ng satisfies
E
(2)( - -N(z))) 2/2
43

---

where T1, T2 are drawn i.d over Unif(554/-1(11)
We aim invoke Lemma 39 We note that #2) 2a/2. and that q is inner product function. Define 9: (T) = 0:a [ b1,)) We see that gi is separable function, and also that
m Ng(x)= Σ (x( + b2-
Hence N3) is the sum of m + separable functions. We can bound the single function as
19:(T) - |a:|Ca
b1,1)
C1 + Vd + B) $ B(1 + Bd + B)dr
(Bdd/2jart)
since || S max(VVF rd)]) S Bdj2. Therefore by Lemma 39
(IE2
2(m + 1) B)
Eg, (x) N 12
/Ba/2.t
By Lemma 38 we have that
(2k 3)!! B(d2,,
B( 2
4/2-1 )a k (d/2 2j - 1)2
(d/2
22m
Simplifying, we have that
(2k - B(d/2 2k)
(2k 3)!!²
(d/2 + 2k 3)1 2)
- (d/2+ 2- 1)2
(2k)
(d/2 2| 2 (d/2+ - 1)2
(2k 3)!!
(d/2 + 4k 2)
k-2
d/2
II
(2k)
(d/2 + 2k 2k- 1)
+ 2k - 3)
d/2 2j
(2k 3)!!²
(d/2 + 4h: - 22)
(2k)
(d/2 + 2k - (d/2 + 2k- 3)
1
(d/2 + 4k 2)
2k(2k 1)(2k -2) 2) (d/2 + 2k (d/2 2k 3)
By Gautschi's 's inequality, we can bound
1 d d-2 2
v
d
-1)
-1/2
3
S o
4/4
< 4d 1//2
2
I(()
Therefore
(d/2 + 4k - 2)d
|| (eLL
2
2k(2k: 1)(2k 2)
(d/2 + 2k: (d2 + 2k 2
(d/2 + 4k 2)d
128
(d/2 2 (d 2k 3)
3
for k S d/4. Altogether,
4/4
1
PP ReLU( (27
2
128d
k
512m2
Am
44

---

for m S d/8. Since ((2) = ReLUU22 - -co =2 /
r(d-r) RcLU(z) Co. we have that
P2 [P2amellea( 2 r(d -r)
P2 RcLU(x)]]
d
v
512m2d
1
16m
We thus have, for any integer < d/8,
2/2 2 Ex, (7((T1,T2)) N²(z))²
2(m + B)
(IB2207lle>
2 12222
BB222.2
Choose S 512gz; we then must have
2(m 1)(B8/22)ott
BBa/2,28
32k
(m + 1)(Bd3/2)oa
a
2 B(d/2,26)2/2 2 2
64k
64k (2h))
C exp
k: -lood-logg llg 2
log(2))) log
Ci exp
(k logd
log k. k klo((2k))
k log 2)
loo d
C1 exp
log k 2k log
C2k log #
C1
exp
for any l S Ozd. Selecting k = LV 512c yields
max(m, B) Ci exp C2e 1/2 log(de)) d 3/2 2 O exp 11/21/g((de)
for less than a universal constant Cg.
We next show the upper bound,
is easy to see that Assumptions and Jare are satisfied. Next, since the verification of Assumptions and 4 only required Lipschitzness, those assumptions are satisfied as well with E, 1. Finally, we have
ES [f*(2))] < E. [c Az)]
1 Ex ((
Ar)
d
<1
2
2
Next, observe that ||Kf* S ? Define A
d+2 A. This scaling ensures |
Then, we can write (²(x) = 9" (x (x)) for g¹ g2)
ReLU(z) co. For 0, define the d+2
2d
smoothed ReLU (=) as
2 25 -E -6 -6505e
ReLU,(2)
x e)²
x z26
45

---

One sees that ReLU, is twice differentiable with |ReLUU.l1. and |ReeLell22.0 = 2c We select the test function q be g(2)) /
=) co. We see that ((7(K/()<<) = ReLU, (1K/))))
and thus
If (T((F (x))))
Vd+2/ ReLU(x"A')) Az) RcLU, (x)
2d
< |ReT.U<" Ax) ReI.U (x) 12 ReT.U. (2' Aa) ReLU, (IRMIE((M)))))
[ReLU(="A)) Az) ReLU, Az) | |K |
ReLU(<"Az) - ReLU. (a'uul + Ld11/12 log d,
where the first inequality follows from Lipschitzness and the second inequality is Corollary 3 using
There exists a constant upper bound for the density of x² Axx and thus we can upper bound
Ax) ReL. Ax) - L = de S é e. Furthermore since 1) = @(| - we get that r/
1 ||K O(x)), and thus
sup |g(=)] =E[ xE(-6(xx),((x))
sup
=
=E(-1.1)
sup 1d(()) = .. | =E|-1.1]
sup =EI- =E(-0((x),((x1]
(ReI.U.)(z)) = poly(1)
sup
1q"(z) = (Kf
(ReLU.)) = poly(1)e-
=E[-1,1]
0(((
Therefore by Theorem 6 we can bound the population loss as
Er, (((=:0)- 2o
d
+ + + 1-
(n, m1, m2)
mi
n
Choosing = d 1/4 yields the desired result. As for the sample complexity, have 1|q/12,00 0(--))= = (dd and so the runtime is poly(d, mi, n2,
G.1 Proof of Lemma
38
Proof of Lemma
38
For any integer k, we define the quantities 1 B2441
(d) 2G (x)dpa(<)
A
['
G(d) (
We also let Zμ
T(d/2)
be the normalization constant.
/ 1
46

---

Integration by parts yields
1/4 Zd / Gld (1))=(1 -x?) 2 dr
d-3
-da G (x)
d (1 x2 - x
1
Za
2k(2k + d - 2)
/
- (I) (1) =²)²²² dr
0 d
d-1
Zd (d) (0) +
Za 2k(2k + d - 2) ld+2)
d 24
Zo42
(d 1)²
B2x 2k-1
From Corollary 5 we have
(d) (0) 17 -0 (d + 2j - 1)
(2k: - 1)!!
(-1)*²
Thus
A Zd
(d 1)m (d 2j - 1) (-1)* + 2h:(2 dd-2 2) Zd Bld+2)
(2k- 1)!!
(25)
(d (d-1)² 1)2 Za+2
The recurrence formula yields
[
B2k+1
(d)
Gld) 26+1 (x)ddua(z)
(26)
2k +
4k 2
2
GG
x
2k
(d)
(x)
2k - d (x) duala)
2k
(27)
4k + d - 2
2k 2
2k
(d)
2k +
2k:
2
B (d)
B2k-1
(28)
I claim that
Zz
h:
(2k 1)!!
A
k: 3> and B24+1 (-1)*7. Zd
(-1) k+1 28
"T (d-2j-1)
(2k (2k-3)1 3)!!
IT 1= (d 2j 1)
We proceed by induction on k. For the base cases, first have
1/ {( adpa(x))
7 =3²)²=² 2-3
J. Zaa
du
Za
d
where we the substitution
(1 x2)
+ Next,
B /' zdpa(x)
Za
Al d-1'
Next, eq. 25 gives
1/ Zea [d- - 1)2
-1
2d
Za
+
(d- 1)2 d+1
Za
(d. 1)(d+1)
Finally, eq. 266 gives
00000
d+2
A
(d)
3
2
B(d
d
Za
d
2
d
d(d
+ 1)
Zu
(d
1)
47

---

Therefore the base case is proven for h 0,
Now, assume that the claim is true for some for all d. We first have
1 Za (d 1)l1 (d + 2j
(2k: + 1)!!
(-1)
g++ + (2k + ()h + d)
Zd B(o+2)
26+2
(d 1)2
Za+2
2k+1
(2k: 1)!!
(2k + 2) d)
(2k: 1)!!
=ld-
(-1) + Za
T (d 2j + 1) (-1)
1=0 (d + 2j 1) (d- --
(d 1)2
(2k - 1)
(d + 2k: + 1) + 1)
(2k: + (2k + d)
(-1) Za
11 [d + 2j 1)
d
d d-1
(2k 1)!!
-d+
= (-11)*+1 Za
1=0 (d 2j
d -
(2k: 1)!!
= (-1))++ 12 Za
;= (d 2j
Next, we have
(1d)
4k + d ++2
(d)
2k +
b(d)
2k +d
2242
2k +d 2k+1
4k + d +
(2k - 1)!!
2k +
(2k 1)!!
=-11)* Zd
2k +
1 (d + 2j 1)
2h + d oo (d + 2j
(2h: 1)!!
(4k + d
- (2k 2)(d + 21 +
(-11) Za
[1t+0 (d 2j 1)
2k +
(2k: 1)!!
(2k + 1)(2k d)
=(-1)* Za
IT,=0 (d 2j 1)
2k + d
(2k: + 1)!!
= (-1) Zd
1 (d + 2j - 1)
Therefore by induction the claim holds for all k; d.
The Gegenbauer expansion of ReLU is given by
ReLU(x) = Σ (ReLU, ( (es, (z).
Note that ReLU(x)= (xx + (z])). Since |x]| is even, the only nonzero odd Gegenbauer coefficient is for (:(d) In this case,
(ReI.U, (OITU.G(())aa 2 (exunlls? 2d
Also, B(d, 1) = d. Next, we see that
(ReLU, Gl LA(HE) = S' R (2)dpa(()= /' ( (x)dua(t) = =A(d)
Plugging in our derivation for 1 gives the desired result.
48


Reviews:
[{"review_id": "vFZ4uiKKKd", "reviewer": "Reviewer_oYFv", "summary": "This paper analyzed the features learned by a three-layer network trained with layer-wise gradient descent as existing analyses are largely restricted to two-layer networks. It presented a general purpose theorem that upper bounds the sample complexity and width needed to achieve low test error when the target has a certain hierarchical structure.", "questions": "Do the authors think this can be applied to convolutional networks and networks with more depth than three layers?", "limitations": "No limitations are discussed in the paper. It would have been nicer if the authors discussed the generalizability of the work to a deeper network.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-\tThis paper has shown that a three-layer network can learn nonlinear features, although existing work only has shown that two-layer neural networks learn only linear functions of input.\n-\tIt’s good that the paper has an example section (Sec. 4.)\n", "weaknesses": "-\tThe “Hierarchical function” that this paper deals with is not early enough defined or introduced. It will make the problem that the paper solved more clear if that point is made clear in an earlier part of the paper.\n-\tThis work still does not answer the same questions for convolutional networks.\n-\tThis paper certainly showed that three-layer networks have provably richer feature learning capabilities than two-layer networks. Nonetheless, it is not guaranteed that the findings and proof can still be applied to any general “depth” as it has shown for the change from the two-layer to three-layer (although it might be high-likely.)\n-\tIn other words, this issue is about generalizability on depth scalability.\n-\tThis paper lacks discussion of its own limitations.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and thoughtful review, and address specific comments below.\n\n> “It will make the problem that the paper solved more clear if that point is made clear in an earlier part of the paper.”\n\n- Thank you for this feedback; we will add more details about what specifically we mean by a hierarchical function to the introduction section in a future revision of our paper.\n\n> “This work still does not answer the same questions for convolutional networks.”\n\n- While understanding the optimization and generalization properties of convolutional networks (or other modern architectures such as transformers) is indeed an interesting question, it is beyond the scope of our current work. We remark that showing an end-to-end learning guarantee for a two-layer convolutional network beyond the kernel regime is still an open question. Prior guarantees for feature learning in fully connected networks have largely been restricted to two-layer networks, and our work makes progress on extending these results to three-layer networks.\n\n> “Nonetheless, it is not guaranteed that the findings and proof can still be applied to any general “depth””\n\n- We agree that it is also an interesting question to understand what kinds of hierarchical functions can be learned by networks of depth >3. We anticipate that a similar layerwise training procedure and analysis could allow us to show learnability of a class of hierarchical functions with deeper networks, but we leave investigation of this to future work.\n\n> “This paper lacks discussion of its own limitations.”\n\n- We included a limitation section within the supplementary material, but we can certainly move this discussion of limitations to the main text in a future revision of our paper. "}, {"review_id": "HbBdEjw8yx", "reviewer": "Reviewer_8d6Z", "summary": "In this work the authors show that there is a three layer neural network setup with better provable learning guarantees than the current best bound for two layer setups. The setup involves a randomly initialized layer with frozen weights, which feeds into a two layer network where first the hidden layer weights are trained, and then the output weights are trained. In the case of functions of quadratic features, this three layer (but two learnable layers) network is able to use feature learning to re-weight the initial random feature kernel to improve training efficiency.", "questions": "What evidence is there that the difference between the 2 and 3 layer bounds will persist even as theoretical techniques/the bounds themselves are improved?\n\nHow do the 3 layer feature learning bounds differ from NTK learning of a similar 3 layer network? (This discussion would be useful to add to the main text to emphasize the benefits of feature learning over depth alone.)\n\nCan the method be used to show anything about the usefulness of running gradient descent over a longer period of time? How does this trade off with the sample complexity?\n\nIs there any intuition about good choices of the function q in the bounds, in a more general setting?", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The model used in the paper is a simple and sensible extension of the two layer network. The learning algorithm is also quite reasonable, both in its connection to previous work as well as to the practical end of deep learning.\n\nThe arguments seem correct (though I did not fully validate the detailed proofs in the appendices), and the intuitive explanation relating feature learning and eigenspace weighting helps shed light on the importance (and perhaps mechanisms of) feature learning in deep networks. The bounds in the 3 layer case are a marked improvement over the kernel learning and 2 hidden layer case bounds.", "weaknesses": "There are two concerns with this paper. The first is a question of the tightness of the bounds. While the bounds show great improvement over the quoted bounds in the two hidden layer case, it is not clear how tight these bounds are in the various scenarios discussed.\n\nRelatedly, there is a big question as to whether or not the bounds are useful for understanding the success of deep learning systems even on simple problems (e.g. FCN on MNIST). In particular, the proof sketch seems to suggest that just a single step of GD in the middle layer is enough to induce massive improvements in the sample complexity; however even in simple settings, it seems that it is helpful to both learn in multiple layers, as well as over multiple timepoints. I note that this is a general weakness of similar sample complexity analyses and not of this particular paper, and the authors do mention this in the discussion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and thoughtful review, and address specific concerns of yours below.\n\n> “While the bounds show great improvement over the quoted bounds in the two hidden layer case, it is not clear how tight these bounds are in the various scenarios discussed.”\n- For the quadratic feature setting (Section 4.2), the information theoretic lower bound for the sample complexity is $d^2$, as there are $O(d^2)$ free parameters in the matrix $A$. Our theorem 3 shows that the three-layer network can learn the target with $d^4$ samples, which is within a polynomial factor of this optimal sample complexity (but crucially still better than kernel methods or two-layer networks). Similarly, the information theoretic lower bound for single-index models is $O(d)$, while our algorithm obtains an $O(d^2)$ sample complexity. Improving the sample complexity to the information theoretic threshold is an interesting direction of future work.\n\n> “...there is a big question as to whether or not the bounds are useful for understanding the success of deep learning systems even on simple problems (e.g. FCN on MNIST)....it seems that it is helpful to both learn in multiple layers, as well as over multiple timepoints.”\n- We certainly agree that there is more work to be done in understanding the effectiveness of deep learning on realistic datasets. In more realistic settings, it is plausible that training all layers jointly and for multiple timesteps leads to improved results over the layerwise training procedure. We view our work as elucidating the capabilities of three-layer networks, in particular in comparison to two-layer networks. Specifically, our results provide a setting in which three-layer networks can efficiently learn a class of functions that cannot be efficiently learned using two-layer networks.\n- Furthermore, the recent work [43] shows that shallow MLPs trained via standard techniques on image datasets do learn features corresponding to the linear feature obtained via a single step of GD for two-layer networks. An interesting direction of future study would be to understand whether this phenomenon also holds for three-layer networks.\n\n> “What evidence is there that the difference between the 2 and 3 layer bounds will persist even as theoretical techniques/the bounds themselves are improved?”\n- We remark that our lower bound in Section 5 shows that no polynomially sized 2-layer network can express $f^*$ below an error threshold. This implies that no algorithm can learn $f^*$ with polynomially many samples / in poly time. Thus even if more refined bounds were proven for 2 layer neural networks, $f^*$ would still not be learnable and the separation would still exist.\n\n> “How do the 3 layer feature learning bounds differ from NTK learning”\n- The NTK is a kernel method, and thus the discussion in Section 3, point 2 also applies to the NTK. Since the NTK for any depth is a rotationally invariant kernel, the lower bound from [27] applies. Therefore, in the quadratic feature example in Section 4.2, $d^{2p}$ samples are needed to learn $g^*(x^TAx)$ for a degree $p$ polynomial $g^*$ via the NTK. We mention this point in line 305, but we will update the exposition to make clear that this lower bound applies to the NTK for a depth 3 network as well.\n\n> “Can the method be used to show anything about the usefulness of running gradient descent over a longer period of time?”\n- Understanding the sample complexity benefit of training $W$ for longer is indeed an interesting direction of future work. While initial works studying feature learning in two-layer neural networks such as [8, 18] relied on the single-step training procedure, later works [2] studied training $W$ for multiple steps and were able to improve the dimension dependence in the sample complexity. For three layer neural networks, it is possible that a more refined analysis of multi-step training could improve the sample complexity even further. However, analyzing the training dynamics of $W$ for multiple steps introduces a number of new technical challenges which will likely require developing new techniques. As such, we defer this to future work. \n\n> “Is there any intuition about good choices of the function q in the bounds”\n- For the general theorem, if $f^*$ possesses the hierarchical structure that $f^* = g^* \\circ h^*$ where $h^*$ is the learned feature, then one should choose $q = g^*$ in the main theorem. In the examples in Section 4, we see that a good choice is $q = g^*$; however, we don’t have a clean description of the optimal $q$ in the general setting.\n"}, {"review_id": "jmLRzRwHJR", "reviewer": "Reviewer_aPJJ", "summary": "The paper theoretically studies the feature learning in three layer neural networks. For the analysis, it considers layer-wise GD; more precisely, the first layer is not trained, the second layer is trained for one step, and then there is the training for last layer. Particularly, they show that three layer neural networks achieve a better bound for learning functions of the form $g(x^TAx)$ comparing to the known bounds for two-layer neural networks. They also show an optimization-based separation showing that three-layer NNs can learn functions that two-layer NNs cannot learn with polynomial width and time (assuming a non-increasing learning rate). ", "questions": "- Can the analysis be done in the SGD setting as well?\n\nHere are also a few typos:\n- line 246: $\\mathbb{E}[(g^{*})'] = 0$ instead of $\\mathbb{E}[g'] = 0$\n- line 291: inconsistency regarding $P_2 f^*$ (e.g., with equation 16) (I think also the same problem appeared in the appendix.)\n- line 298: is it the intended equation?", "limitations": "There is no negative societal impact. The limitations of theory have been discussed in the work and are (unfortunately) quite common in the current literature of deep learning theory. ", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper considers the optimization of three-layer NNs with a layer-wise training and provides a fairly general approach for studying their feature learning capabilities (in the paper, the features have also been characterized for the examples considered).\n- For functions of the type $g(x^TAx)$, a sample complexity upper bound for three-layer NNs is proved which is more efficient than the known upper bounds for two-layer NNs. \n- There is an optimization-based separation result provided between two-layer and three-layer NNs. \n- Paper is generally very well written.\n", "weaknesses": "- Generally the limitations of the work have been discussed in the paper. These limitations are usually quite common in the deep learning theory literature. For example, the training is layer-wise, and the first layer is not trained at all and second layer is only trained for one iteration. In this sense, the work is similar to theoretical analysis done on two-layer NNs in many recent works. \n- The simulations have been implemented for the settings as the theorems. However, it would have been interesting to also run experiments with the more common setting (e.g., training all parameters together) for both two-layer and three-layer NNs and compare the results with the results of the model in the theoretical settings (e.g., with layer-wise training). This would potentially show if the assumptions made in theory are too restrictive.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and thoughtful review, and address specific comments below:\n\n> “However, it would have been interesting to also run experiments with the more common setting (e.g., training all parameters together) for both two-layer and three-layer NNs…”\n- We agree that this is an interesting experiment to run, and will add it to a future revision of our paper.\n\n> “Can the analysis be done in the SGD setting as well?”\n- For the first stage we do need a large batch size, but the second stage can be done in the online SGD setting. Since the loss is convex in $a$, the analysis for this would be a standard application of online convex optimization results.\n\nThank you for pointing out the typos, we will update these in the next revision of our paper. Line 298 should read $\\mathbb{K}f^*(x) = \\Theta(d^{-2})\\cdot(x^TAx + o_d(1))$.\n"}, {"review_id": "vmGAZLtRh8", "reviewer": "Reviewer_3iYS", "summary": "This paper studies the feature learning capability of a 3-layer neural\nnetwork, where the bottom layer is random and fixed, the middle layer is\ntrained for only one step from zero, and the upper layer is trained in\nthe rest of the gradient descent steps. The paper characterizes the\nricher feature learning capability of this 3-layer network. That is, the\none-step training of the middle-layer weights essentially maps the\noriginal random features to another space, which enlarges the type of\nfunctions that can be learned. Two special cases are then considered:\n(i) the single index model (similar to [18]); and (ii) more\nsignificantly, the quadratic feature model (more complex than [18]). The\nresults demonstrate a significant reduction in the sample complexity\ncompared to kernel methods. Further, the latter class of functions in\n(ii) cannot be approximated by 2-layer networks, and thus a\ndepth-separation between 2-layer and 3-layer networks is established.", "questions": "1. Can the authors comment on other possible features that this\napproach can learn effectively (beyond quadratic features)? \n\n2. The one-step training of the middel-layer weights seems to be a\ncrucial element for learning a specific type of features. At the same\ntime, one-step training could be limiting. Can the authors comment on\npossible generalizations for multi-step training on the middle-layer? \n\nPost rebuttal phase:\n\nThe reviewer wishes to thank the authors for their response and preliminary thoughts on these questions.", "limitations": "The discussion of limitations is adequate.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The results are novel and solid. The precisely characterization of\nthe feature learning capability of this particular 3-layer network\n(Theorem 1), and the ability to learn functions of quadratic features,\nare both significant contributions to the field. \n\n2. The paper is also very well-written and presents the main intuitions\nquite clearly.", "weaknesses": "Nothing particular.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and thoughtful review, and address specific comments below.\n\n> “Can the authors comment on other possible features that this approach can learn effectively”\n- We hypothesize that our results here can be extended to learn arbitrary polynomial features, and thus allow us to learn hierarchical functions $f^* = g^* \\circ h^*$ where $h^*$ is a polynomial, with sample complexity scaling with the degree of $h^*$. However, we don’t currently have any concrete results in this setting and hence defer this to future work.\n\n> “Can the authors comment on possible generalizations for multi-step training on the middle-layer?”\n- While initial works studying feature learning in two-layer neural networks such as [8, 18] relied on the single-step training procedure, later works [2] studied training $W$ for multiple steps and were able to improve the dimension dependence in the sample complexity. For three layer neural networks, it is possible that a more refined analysis of multi-step training could improve the sample complexity. However, analyzing the training dynamics of $W$ for multiple steps introduces a number of new technical challenges which will likely require developing new techniques. As such, we defer this to future work. \n"}]

Author rebuttal:

Meta-review: This paper analyzes feature learning in three-layer neural networks and proves that there are functions that are efficiently learnable with three layers but not two layers.

The reviewers acknowledge that these results improve the state-of-the-art in the theoretical understanding of feature learning, and generally  found the paper to be well written and clear. Some of the weaknesses brought forward in the discussion include limitations of layer-wise analysis, restriction to simple architectures, and the tightness of the bounds. These concerns were adequately addressed by the authors and overall the consensus is that this is an interesting paper that should be accepted to NeurIPS.
Decision: Accept (spotlight)