Title: Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep Learning Theory

Feature Learning

Three-Layer Neural Network

Depth Separation

Gradient Descent

Representation Learning
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

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.

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 a hierarchy of learned features. This feature learning process enables both improved accuracy [29] and transfer learning [21]. Despite its importance, we still have a rudimentary theoretical understanding of the feature learning process. Fundamental questions include understanding what features are learned, how they are learned, and how they affect generalization.

From a theoretical viewpoint, a fascinating question is to understand how depth can be leveraged to learn more salient features and as a consequence a richer class of hierarchical functions. The base case for this question is to 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 be leveraged to obtain a sample complexity improvement. Initial learning guarantees for three-layer networks [16, 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.

On the other hand, the standard approach in deep learning theory to understand the benefit of depth has been to establish 'depth separations' [51], i.e. 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, 44] 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 a 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 a three-layer network trained with a layer-wise variant of gradient descent (Algorithm 1). Our main contributions are as follows.

Theorem 1 is a general purpose sample complexity guarantee for Algorithm 1 to learn an arbitrary target function f ∗ . We first show that Algorithm 1 learns a feature roughly corresponding to a low-frequency component of the target function f ∗ with respect to the random feature kernel K 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 1D 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 1 in two statistical learning settings which satisfy such hierarchical structure. As a warmup, we show that Algorithm 1 learns single-index models (i.e f ∗ ( x ) = g ∗ ( w · x ) ) in d 2 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 g ∗ . We next show that Algorithm 1 learns the target f ∗ ( x ) = g ∗ ( x T Ax ) , where g ∗ is either Lipschitz or a degree p = O (1) polynomial, up to o d (1) error with d 4 samples. This improves on all existing guarantees for learning with two-layer networks or via NTK-based approaches, which all require sample complexity d Ω( p ) . A key technical step is to show that for the target f ∗ ( x ) = g ∗ ( x T Ax ) , the learned feature is approximately x T Ax . This argument relies on the universality principle in high-dimensional probability, and may be of independent interest.

We conclude by establishing an explicit optimization-based depth separation: We show that the target function f ∗ ( x ) = ReLU( x T Ax ) for appropriately chosen A can be learned by Algorithm 1 up to o d (1) error in d 4 samples, whereas any two layer network needs either superpolynomial width or weight norm in order to approximate f ∗ up to comparable accuracy. This implies that such an f ∗ is not efficiently learnable via two-layer networks.

The above separation hinges on the ability of three-layer networks to learn the nonlinear feature x T Ax and leverage this feature learning to obtain an improved sample complexity. Altogether, our work presents a 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 [31, 50, 23, 17]. The NTK theory shows global convergence by coupling to a kernel regression problem and generalization via the application of kernel generalization bounds [7, 15, 5]. The NTK can be characterized explicitly for certain data distributions [27, 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

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, [27] shows that kernels cannot adapt to low-dimensional structure and require d k samples to learn any degree k polynomials in d 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, 3, 25, 26, 20, 54, 33, 35]. Crucially, our analysis is not in the kernel regime, and shows an improvement of three-layer networks over two-layer networks in the feature-learning regime.

Feature Learning. Recent work has studied the provable feature learning capabilities of two-layer neural networks. [9, 1, 2, 8, 13, 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 f ∗ ( x ) = g ( Ux ) for U ∈ R r × d ). Here, x ↦→ Ux is the 'linear feature.' Such target functions include low-rank polynomials [18, 2] and single-index models [8, 13] for Gaussian covariates, as well as sparse boolean functions [1] such as the k -sparse parity problem [10] for covariates uniform on the hypercube. [43] draws connections from the mechanisms in these works to feature learning in standard image classification settings. The above approaches rely on layerwise training procedures, and our Algorithm 1 is an adaptation of the algorithm in [18].

Another approach uses the quadratic Taylor expansion of the network to learn classes of polynomials [9, 41] This approach can be extended to three-layer networks. [16] replace the outermost layer with its quadratic approximation, and by viewing z p as the hierarchical function ( z p/ 2 ) 2 show that their three-layer network can learn low rank, degree p polynomials in d p/ 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 T Ax ) for degree p polynomial g ∗ yields a sample complexity of d p +1 . However, [16, 5] are proved via opaque landscape analyses, do not concretely identify the learned features, and rely on nonstandard algorithmic modifications. Our Theorem 1 directly identifies the learned features, and when applied to the quadratic feature setting in Section 4.2 obtains an improved sample complexity guarantee independent of the degree of g ∗ .

Depth Separations. [51] constructs a function which can be approximated by a 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 a shallow (depth 3 network) is a necessary condition for learnability via a deeper network.

These issues have motivated the development of optimization-based, or algorithmic, depth separations, which construct functions which are learnable by a three-layer network but not by two-layer networks. [45] shows that certain ball indicator functions 1 ( ∥ x ∥ ≥ λ ) are not approximatable by two-layer networks, yet are learnable via GD on a special variant of a 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 ≳ d 36 ) is weak. [44] shows that a multi-layer mean-field network with a 1D bottleneck layer can learn the target ReLU(1 - ∥ x ∥ ) , which [47] previously showed was inaproximatable 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 poly( d ) . Our Section 4.2 shows that three-layer networks can learn a larger class of features ( x T Ax versus ∥ x ∥ ) and functions on top of these features (any Lipschitz q versus ReLU ), with explicit dependence on the width and sample complexity needed ( n, m 1 , m 2 = ˜ O ( d 4 ) ).

2 Preliminaries

2.1 Problem Setup

Data distribution. Our aim is to learn the target function f ∗ : X d → R ,with X d ⊂ R d the space of covariates. We let ν be some distribution on X d , and draw two independent datasets D 1 , D 2 , each with n samples, so that each x ∈ D 1 or x ∈ D 2 is sampled i.i.d as x ∼ ν . Without loss of generality, we normalize so E x ∼ ν [ f ∗ ( x ) 2 ] ≤ 1 . We make the following assumptions on ν :

Algorithm 1 Layer-wise training algorithm

Input: Initialization θ (0) ; learning rates η 1 , η 2 ; weight decay λ ; time T {Stage 1: Train W } W (1) ← W (0) - η 1 ∇ W L 1 ( θ (0) ) a (1) ← a (0) θ (1) ← ( a (1) , W (1) , b (0) , V (0) ) {Stage 2: Train a } for t = 2 , · · · , T do a ( t ) ← a ( t - 1) - η 2 [ ∇ a L 2 ( θ ( t - 1) ) + λa ( t - 1) ] θ ( t ) ← ( a ( t ) , W (1) , b (0) , V (0) ) end for ˆ θ ← θ ( T ) Output: ˆ θ

Definition 1 (Sub-Gaussian Vector) . A mean-zero random vector X ∈ R d is γ -subGaussian if, for all unit vectors v ∈ R d , E [exp( λX · v )] ≤ exp ( γ 2 λ 2 ) for all λ ∈ R .

Assumption 1. E x ∼ ν [ x ] = 0 and ν is C γ -subGaussian for some constant C γ .

Assumption 2. f ∗ has polynomially growing moments, i.e there exist constants ( C f , ℓ ) such that E x ∼ ν [ f ∗ ( x ) q ] 1 /q ≤ C f q ℓ for all q ≥ 1 .

We note that Assumption 2 is satisfied by a number of common distributions and functions, and we will verify that Assumption 2 holds for each example in Section 4.

Three-layer neural network. Let m 1 , m 2 be the two hidden layer widths, and σ 1 , σ 2 be two activation functions. Our learner is a three-layer neural network parameterized by θ = ( a, W, b, V ) , where a ∈ R m 1 , W ∈ R m 1 × m 2 , b ∈ R m 1 , and V ∈ R m 2 × d . The network f ( x ; θ ) is defined as:

Here, w i ∈ R m 2 is the i th row of W , and h (0) ( x ) := σ 2 ( V x ) ∈ R m 2 is the random feature embedding arising from the innermost layer. The parameter vector θ (0) := ( a (0) , W (0) , b (0) , V (0) ) is initialized with a (0) i ∼ iid Unif( {± 1 } ) , W (0) = 0 , the biases b (0) i ∼ iid N (0 , 1) , and the rows v (0) i of V (0) drawn v i ∼ iid τ , where τ is the uniform measure on S d -1 (1) , the d -dimensional unit sphere. We make the following assumption on the activations, and note that the polynomial growth assumption on σ 2 is satisfied by all activations used in practice.

Assumption 3. σ 1 is the ReLU activation, i.e σ 1 ( z ) = max( z, 0) , and σ 2 has polynomial growth, i.e | σ 2 ( x ) | ≤ C σ (1 + | x | ) α σ for some constants C σ , α σ > 0 .

Training Algorithm. Let L i ( θ ) denote the empirical loss on dataset D i ; that is for i = 1 , 2 : L i ( θ ) := 1 n ∑ x ∈D i ( f ( x ; θ ) -f ∗ ( x )) 2 . 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 = 1 timesteps. Next, the outer layer weights a are trained for t = T -1 timesteps. This two stage training process is common in prior works analyzing gradient descent on two-layer networks [18, 8, 1, 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 1.

2.2 Technical definitions

The activation σ 2 admits a random feature kernel K : X d ×X d → R and corresponding integral operator K : L 2 ( X d , ν ) → L 2 ( X d , ν ) :

Definition 2 (Kernel objects) . σ 2 admits the random feature kernel

and corresponding integral operator

We make the following assumption on K , which we verify for the examples in Section 4:

Assumption 4. K f ∗ has polynomially bounded moments, i.e there exist constants C K , χ such that, for all 1 ≤ q ≤ d , ∥ K f ∗ ∥ L q ( ν ) ≤ C K q χ ∥ K f ∗ ∥ L 2 ( ν ) .

We also require the definition of the Sobolev space:

Definition 3. Let W 2 , ∞ ([ -1 , 1]) be the Sobolev space of twice continuously differentiable functions q : [ -1 , 1] → R equipped with the norm ∥ q ∥ k, ∞ := max s ≤ k max x ∈ [ -1 , 1] ∣ ∣ q ( s ) ( x ) ∣ ∣ for k = 1 , 2 .

2.3 Notation

We use big O notation (i.e O, Θ , Ω ) to ignore absolute constants ( C σ , C f , etc.) that do not depend on d, n, m 1 , m 2 . We further write a d ≲ b d if a d = O ( b d ) , and a d = o ( b d ) if lim d →∞ a d /b d = 0 . Additionally, we use ˜ O notation to ignore terms that depend logarithmically on dnm 1 m 2 . For f : X d → R , define ∥ f ∥ L p ( X d ,ν ) = ( E x ∼ ν [ f ( x ) p ]) 1 /p . To simplify notation we also call this quantity ∥ f ∥ L p ( ν ) , and ∥ g ∥ L p ( X d ,ν ) , ∥ g ∥ L p ( τ ) are defined analogously for functions g : S d -1 (1) → R . When the domain is clear from context, we write ∥ f ∥ L p , ∥ g ∥ L p . We let L p ( X d , ν ) be the space of f with finite ∥ f ∥ L p ( X d ,ν ) . Finally, we write E x and E v as shorthand for E x ∼ ν and E v ∼ τ respectively.

3 Main Result

The following is our main theorem which upper bounds the population loss of Algorithm 1:

Theorem 1. Select q ∈ W 2 , ∞ ([ -1 , 1]) . Let η 1 = m 1 m 2 η , and assume n, m 1 , m 2 = ˜ Ω( ∥ K f ∗ ∥ -2 L 2 ) There exist η, λ, η 2 such that after T = poly( n, m 1 , m 2 , d, ∥ q ∥ 2 , ∞ ) timesteps, with high probability over the initialization and datasets the output ˆ θ of Algorithm 1 satisfies the population L 2 loss bound

The full proof of this theorem is in Appendix D. The population risk upper bound (4) has three terms:

The first term quantifies the extent to which feature learning is useful for learning the target f ∗ , and depends on how close f ∗ is to having hierarchical structure. Concretely, if there exists q : R → R such that the compositional function q ◦ η · K f ∗ is close to the target f ∗ , then this first term is small. In Section 4, we show that this is true for certain hierarchical functions. In particular, say that f ∗ satisfies the hierarchical structure f ∗ = g ∗ ◦ h ∗ . If the quantity K f ∗ is nearly proportional to the true feature h ∗ , then this first term is negligible. As such, we refer to the quantity K f ∗ as the learned feature .

The second term is the sample (and width) complexity of learning the feature K f ∗ . It is useful to compare the ∥ K f ∗ ∥ -2 L 2 term to the standard kernel generalization bound, which requires n ≳ 〈 f ⋆ , K -1 f ⋆ 〉 L 2 . Unlike in the kernel bound, the feature learning term in (4) does not require inverting the kernel K as it only requires a lower bound on ∥ K f ⋆ ∥ L 2 . This difference can be best understood by considering the alignment of f ⋆ with the eigenfunctions of K . Say that f ⋆ has nontrivial alignment with eigenfunctions of K with both small ( λ min ) and large ( λ max ) eigenvalues. Kernel methods require Ω( λ -1 min ) samples, which blows up when λ min is small; the sample complexity of kernel methods depends on the high frequency components of f ∗ . On the other hand, the guarantee in Theorem 1 scales with

∥ K f ⋆ ∥ -2 L 2 = O ( λ -2 max ) , which can be much smaller. In other words, the sample complexity of 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.

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-frequency ones, then a good q for Theorem 1 exists and the dominant term in (4) is the sample complexity of feature learning term, which only depends on the low-frequency 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 to 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 . The first step of GD on the population loss for a neuron w j is thus

Therefore the network f ( x ′ ; θ (1) ) after the first step of GD is given by

We first notice that this network now implements a 1D function of the quantity

Specifically, the network can be rewritten as

Since f implements a hierarchical function of the quantity ϕ ( x ) , we term ϕ the learned feature .

The second stage of Algorithm 1 is equivalent to random feature regression. We first use results on ReLU random features to show that any q ∈ W 2 , ∞ ([ -1 , 1]) can be approximated on [ -1 , 1] as q ( z ) ≈ 1 m 1 ∑ m 1 j =1 a ∗ j σ 1 ( a (0) j z + b j ) for some ∥ a ∗ ∥ ≲ ∥ q ∥ 2 , ∞ (Lemma 3). Next, we use the standard kernel Rademacher bound to show that the excess risk scales with the smoothness of q (Lemma 5). Hence we can efficiently learn functions of the form q ◦ ϕ .

It suffices to compute this learned feature ϕ . For m 2 large, we observe that

The learned feature is thus approximately η · K f ∗ . Choosing η so that | ϕ ( x ′ ) | ≤ 1 , we see that Algorithm 1 learns functions of the form ϕ ◦ ( η · K f ∗ ) . Finally, we translate the above analysis to the finite sample gradient via standard concentration tools. Since the empirical estimate to K f ∗ concentrates at a 1 / √ n rate, n ≳ ∥ K f ∗ ∥ -2 L 2 samples are needed to obtain a constant factor approximation (Lemma 7).

4 Examples

Wenext instantiate Theorem 1 in two specific statistical learning settings which satisfy the hierarchical prescription detailed in Section 3. As a warmup, we show that three-layer networks efficiently learn single index models. Our second example shows how three-layer networks can obtain a sample complexity improvement over existing guarantees for two-layer networks.

4.1 Warmup: single index models

Let f ∗ = g ∗ ( w ∗ · x ) , for unknown direction w ∗ ∈ R d and unknown link function g ∗ : R → R , and take X d = R d with ν = N (0 , I ) . Prior work [18, 13] shows that two-layer neural networks learn such functions with an improved sample complexity over kernel methods. Let σ 2 ( z ) = z , so that the network is of the form f ( x ; θ ) = 1 m 1 a T σ 1 ( WVx ) . We can verify that Assumptions 1 to 4 are satisfied, and thus applying Theorem 1 in this setting yields the following:

̸

Theorem 2. Let f ∗ ( x ) = g ∗ ( w ∗ · x ) , where ∥ w ∗ ∥ 2 = 1 . Assume that g ∗ , ( g ∗ ) ′ and ( g ∗ ) ′′ are polynomially bounded and that E z ∼N (0 , 1) [( g ∗ ) ′ ( z )] = 0 . Then with high probability Algorithm 1 satisfies the population loss bound

̸

Given widths m 1 , m 2 = ˜ Θ( d 2 ) , n = ˜ Θ( d 2 ) samples suffice to learn f ∗ , which matches existing guarantees for two-layer neural networks [13, 18]. We remark that prior work on learning single-index models under assumptions on the link function such as monotonicity or the condition E z ∼N (0 , 1) [( g ∗ ) ′ ( z )] = 0 require d samples [49, 37, 12]. However, our sample complexity improves on that of kernel methods, which require d p samples when g ∗ is a degree p polynomial.

Theorem 2 is proved in Appendix E.1; a brief sketch is as follows. Since σ 2 ( z ) = z , the kernel is K ( x, x ′ ) = E v [( x · v )( x ′ · v )] = x · x ′ d . By an application of Stein's Lemma, the learned feature is

Since f ∗ ( x ) = g ∗ ( w ∗ · x ) , ∇ f ∗ ( x ) = w ∗ ( g ∗ ) ′ ( w ∗ · x ) , and thus

The learned feature is proportional to the true feature, so an appropriate choice of η and choosing q = g ∗ in Theorem 1 implies that ∥ K f ∗ ∥ -2 L 2 = d 2 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 a sample complexity improvement over two-layer networks.

Let X d = S d -1 ( √ d ) , the sphere of radius √ d , and ν the uniform measure on X d . The integral operator K has been well studied [39, 27, 38], and its eigenfunctions correspond to the spherical harmonics. Preliminaries on spherical harmonics and this eigendecomposition are given in Appendix F.

Consider the target f ∗ ( x ) = g ∗ ( x T Ax ) , where A ∈ R d × d is a symmetric matrix and g ∗ : R → R is an unknown link function. In contrast to a single-index model, the feature x T Ax we aim to learn is a quadratic function of x . Since one can write x T Ax = x T ( A -Tr( A ) · I d ) x +Tr( A ) , we without loss of generality assume Tr( A ) = 0 . We also select the normalization ∥ A ∥ 2 F = d +2 2 d = Θ(1) ; this ensures that E x ∼ ν [( x T Ax ) 2 ] = 1 . We first make the following assumptions on the target function.: Assumption 5. E x [ f ∗ ( x )] = 0 , E z ∼N (0 , 1) [( g ∗ ) ′ ( z )] = Θ(1) , g ∗ is 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:

Note that κ ≤ √ d . If A has rank Θ( d ) and condition number Θ(1) , then κ = Θ(1) . Furthermore, when the entries of A are sampled i.i.d, κ = ˜ Θ(1) with high probability by Wigner's semicircle law.

Finally, we make the following nondegeneracy assumption on the Gegenbauer decomposition of σ 2 (defined in Appendix F). We show that λ 2 2 ( σ 2 ) = O ( d -2 ) , and later argue that the following assumption is mild and indeed satisfied by standard activations such as σ 2 = ReLU .

Assumption 7. Let λ 2 ( σ 2 ) be the 2nd Gegenbauer coefficient of σ 2 . Then λ 2 2 ( σ 2 ) = Θ( d -2 ) .

We can verify Assumptions 1 to 4 hold for this setting, and thus applying Theorem 1 yields the following:

Theorem 3. Under Assumptions 5 to 7, with high probability Algorithm 1 satisfies the population loss bound

We thus require sample size n = ˜ Ω( d 4 ) and widths m 1 , m 2 = ˜ Ω( d 4 ) to obtain o d (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 2 ,

where P k is the orthogonal projection onto the subspace of degree k spherical harmonics, and the c k are constants satisfying c k = O ( d -k ) . Since f ∗ is an even function and E x [ f ∗ ( x )] = 0 , truncating this expansion at k = 2 yields

It thus suffices to compute P 2 f ∗ . To do so, we draw a connection to the universality phenomenon in high dimensional probability. Consider two features x T Ax and x T Bx with ⟨ A,B ⟩ = 0 . We show that, when d is large, the distribution of x T Ax approaches that of the standard Gaussian, while x T Bx approaches a mixture of χ 2 and Gaussian random variables independent of x T Ax . As such, we show

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 ∗ = g ∗ ( x T Ax ) , the projection P 2 f ∗ satisfies

The learned feature is thus K f ∗ ( x ) = Θ( d -2 ) · ( x T Ax + o d (1)) ; plugging this into Theorem 1 yields the d 4 sample complexity.

The full proof of Theorem 3 is deferred to Appendix E.2. In Appendix E.3 we show that when g ∗ is a degree p = O (1) polynomial, Algorithm 1 learns f ∗ in ˜ O ( d 4 ) samples with an improved error floor.

Comparison to two-layer networks. Existing guarantees for two-layer networks cannot efficiently learn functions of the form f ∗ ( x ) = g ∗ ( x T Ax ) 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 d 2 p ≫ d 4 samples to learn f ∗ [27]. Improved guarantees for two-layer networks learn degree p ′ polynomials in r p ′ samples when the target only depends on a rank r ≪ d projection of the input [18]. However, g ∗ ( x T Ax ) cannot be written in this form for some r ≪ d , and thus existing guarantees do not apply. We conjecture that two-layer networks require d Ω( p ) samples when g ∗ is a degree p polynomial.

Altogether, the ability of three-layer networks to efficiently learn the class of functions f ∗ ( x ) = g ∗ ( x T Ax ) hinges on their ability to extract the correct nonlinear feature. Empirical validation of the above examples is given in Appendix A.

5 An Optimization-Based Depth Separation

We complement the learning guarantee in Section 4.2 with a lower bound showing that there exist functions in this class that cannot be approximated by a polynomial size two-layer network.

The class of candidate two-layer networks is as follows. For a parameter vector θ = ( a, W, b 1 , b 2 ) , where a ∈ R m , W ∈ R m × d , b 1 ∈ R m , b 2 ∈ R , define the associated two-layer network as

Let ∥ θ ∥ ∞ := max( ∥ a ∥ ∞ , ∥ W ∥ ∞ , ∥ b 1 ∥ ∞ , ∥ b 2 ∥ ∞ ) denote the maximum parameter value. We make the following assumption on σ , which holds for all commonly-used activations.

Assumption 8. There exist constants C σ , α σ such that | σ ( z ) | ≤ C σ (1 + | z | ) α σ .

Our main theorem establishing the separation is the following.

Theorem 4. Let d be a suffiently large even integer. Consider the target function f ∗ ( x ) = ReLU( x T Ax ) -c 0 , where A = 1 √ d U ( 0 I d/ 2 I d/ 2 0 ) U T for some orthogonal matrix U and c 0 =

E x ∼ ν [ ReLU( x T Ax ) ] . Under Assumption 8, there exist constants C 1 , C 2 , C 3 , c 3 , depending only on ( C σ , α σ ) , such that for any c 3 ≥ ϵ ≥ C 3 d -2 , any two layer neural network N θ ( x ) of width m and population L 2 error bound ∥ N θ -f ∗ ∥ 2 L 2 ≤ ϵ must satisfy max( m, ∥ θ ∥ ∞ ) ≥ C 1 exp ( C 2 ϵ -1 / 2 log( dϵ ) ) . However, Algorithm 1 outputs a predictor satisfying the population L 2 loss bound

after T = poly( d, m 1 , m 2 , n ) timesteps.

The lower bound follows from a modification of the argument in [19] along with an explicit decomposition 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 1 to a smoothed version of ReLU , since ReLU is not in W 2 , ∞ ([ -1 , 1]) . The full proof of the theorem is given in Appendix G.

Remarks. In order for a two-layer network to achieve test loss matching the d -1 / 6 error floor, either the width or maximum weight norm of the network must be exp ( Ω( d δ ) ) for some constant δ ; this lower bound is superpolynomial in d . As a consequence, gradient descent on a poly-width two-layer neural network with stable step size must run for superpolynomially many iterations in order for some weight to grow this large and thus converge to a solution with low test error. Therefore f ∗ is not learnable via gradient descent in polynomial time. This reduction from a weight norm lower bound to a 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 R d , and let A = d -1 2 P S -d -1 2 P ⊥ S , where P S , P ⊥ S are projections onto the subspace S and and its orthogonal complement respectively. Then, f ∗ ( x ) = ReLU ( 2 d -1 2 ∥ P S x ∥ 2 -d 1 2 ) . [44] established an optimization-based separation for the target ReLU(1 - ∥ x ∥ ) , under a different distribution ν . However, their analysis relies heavily on the rotational symmetry of the target, and they posed the question of learning ReLU(1 -∥ P S x ∥ ) for some subspace S . Our separation applies to a 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 a provable sample complexity improvement over two-layer networks. There are a number of interesting directions for future work. First, can our framework be used to learn hierarchical functions of a larger class of features beyond quadratic functions? Next, since Theorem 1 is general

purpose and makes minimal distributional assumptions, it would be interesting to understand if it can be applied to standard empirical datasets such as CIFAR-10, and what the learned feature K f ∗ and hierarchical learning correspond to in this setting. Finally, our analysis studies the nonlinear feature learning that arises from a neural representation σ 2 ( V x ) where V is fixed at initialization. This alone was enough to establish a separation between two and three-layer networks. A fascinating question is 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.

References

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Proceedings of Thirty Fifth Conference on Learning Theory , pages 47824887, 2022.

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory , pages 2552-2623. PMLR, 2023.

Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems , 32, 2019.

Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413 , 2020.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019.

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.

Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning (ICML) , 2019.

Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems , 35:37932-37946, 2022.

Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations , 2020.

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.

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.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research , 22(1):4788-4838, 2021.

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:97689783, 2022.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

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.

Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances in Neural Information Processing Systems , 33:22134-22145, 2020.

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Amit Daniely. Depth separation for neural networks. In Conference on Learning Theory , pages 690-696. PMLR, 2017.

Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems , 33:20356-20365, 2020.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423 .

Simon S. Du and Jason D. Lee. On the power of over-parametrization in neural networks with quadratic activation. In Proceedings of the 35th International Conference on Machine Learning , Proceedings of Machine Learning Research, pages 1329-1338, 2018.

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.

Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Conference on learning theory , pages 907-940. PMLR, 2016.

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.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems , 2020.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics , 49:1029-1054, 2021.

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.

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.

Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks? - a neural tangent kernel perspective. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 2698-2709. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/ 1c336b8080f82bcc2cd2499b4c57261d-Paper.pdf .

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

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.

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 2613-2682, 2020.

Eran Malach and Shai Shalev-Shwartz. Is deeper better only when shallow is good? In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper/2019/file/ 606555cf42a6719782a952aa33cfa2cb-Paper.pdf .

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 73797389, 2021.

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 134 of Proceedings of Machine Learning Research , pages 3265-3295. PMLR, 15-19 Aug 2021. URL https://proceedings.mlr.press/v134/malach21a.html .

Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics , 46(6A):2747-2774, 2018.

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.

Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. arXiv preprint arXiv:2007.12826 , 2020.

Eshaan Nichani, Adityanarayanan Radhakrishnan, and Caroline Uhler. Increasing depth leads to u-shaped test risk in over-parameterized convolutional networks. arXiv preprint arXiv:2010.09610 , 2020.

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-14581, 2022.

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016. URL https://proceedings.neurips. cc/paper/2016/file/148510031349642de5ca0c544f31b2ef-Paper.pdf .

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.

Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. arXiv preprint arXiv:2304.01063 , 2023.

Itay Safran and Jason Lee. Optimization-based separations for neural networks. In Conference on Learning Theory , pages 3-64. PMLR, 2022.

Itay Safran and Ohad Shamir. Depth-width tradeoffs 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 https://proceedings.mlr. press/v70/safran17a.html .

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-Second Conference on Learning Theory , volume 99 of Proceedings of Machine Learning Research , pages 2664-2666. PMLR, 25-28 Jun 2019. URL https://proceedings.mlr. press/v99/safran19a.html .

Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations , 2017. URL https://openreview.net/forum?id=H1W1UN9gg .

Mahdi Soltanolkotabi. Learning relus via gradient descent. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Mahdi Soltanolkotabi, Adel Javanmard, and Jason D. Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory , 65:742-769, 2018.

Matus Telgarsky. Benefits of depth in neural networks. In Conference on learning theory , pages 1517-1539. PMLR, 2016.

Ramon van Handel. Probability in high dimension, 2016. URL http://web.math. princeton.edu/~rvan/APC550.pdf .

Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. arXiv preprint arXiv:1810.05369 , 2018.

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.

Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. 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, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 5393-5402. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/xiao18a.html .

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. PMLR, 13-18 Jul 2020. URL https: //proceedings.mlr.press/v119/xiao20b.html .

Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Proceedings of the 38th International Conference on Machine Learning , 2021.

Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. In Advances in Neural Information Processing Systems , 2019.

Limitations

One limitation of our work is that the innermost layer weights V are fixed throughout training, and our training procedure is layerwise. We remark that this is already enough to establish an algorithmic separation between two and three-layer networks (Section 5), and that prior works on two-layer networks also rely on layerwise training procedures. It is a 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 ν 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 a very important direction of future work to extend these lines of work to more general data distributions.

A Empirical Validation

10 2

10 3

10 4

10 5

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Test Loss

Test Loss

d

=

8

d

=

16

d

=

32

d

=

64

d

=

128

Trivial Predictor

10 2

10 3

10 4

10 5

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

R

2

correlation of

f

and

h

/star

Correlation of Features

d

=

8

d

=

16

d

=

32

d

=

64

d

=

128

Single Index Setting

10 2

10 3

10 4

10 5

10 6

10 7

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Test Loss

Test Loss

d

=

8

d

=

16

d

=

32

d

=

64

Trivial Predictor

10 2

10 3

10 4

10 5

10 6

10 7

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

R

2

correlation of

f

and

h

/star

Correlation of Features

d

=

8

d

=

16

d

=

32

d

=

64

Quadratic Feature (Random Entries)

10 2

10 3

10 4

10 5

10 6

10 7

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Test Loss

Test Loss

d

=

8

d

=

16

d

=

32

d

=

64

Trivial Predictor

10 2

10 3

10 4

10 5

10 6

10 7

n

(number of training samples)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

R

2

correlation of

f

and

h

/star

Correlation of Features

d

=

8

d

=

16

d

=

32

d

=

64

Quadratic Feature (Random Projection)

-

2

0

2

4

h

/star

(

x

) =

x

b

-

3

-

2

-

1

0

1

2

3

4

f

(

x

)

(Learned Feature)

Single Index,

d

=

64

n

=

2 4

n

=

2 8

n

=

2 16

-

2

0

2

4

h

/star

(

x

) =

x T Ax

-

3

-

2

-

1

0

1

2

3

f

(

x

)

(Learned Feature)

d

=

16

n

=

2 8

n

=

2 16

n

=

2 24

Correlation Between Learned and Predicted Features

·

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 1 where w ⋆ ∈ S d -1 is drawn randomly and g ⋆ ( x ) = sigmoid( x ) = 1 1+ e -x , which satisfies the condition E z ∼N (0 , 1) [ g ′ ( z )] = 0 . As in Theorem 2, we choose the initial activation σ 2 ( z ) = z . We optimize the hyperparameters η 1 , λ using grid search over a holdout validation set of size 2 15 and report the final error over a test set of size 2 15 .

Quadratic Feature Setting We learn the target function g ⋆ ( x T Ax ) using Algorithm 1 where g ⋆ ( x ) = ReLU( x ) . We ran our experiments with two different choices of A :

A is symmetric with random entries, i.e. A ij ∼ N (0 , 1) and A ji = A ij for i ≤ j .

A is a random projection, i.e. A = Π S where S is a random d/ 2 dimensional subspace.

Both choices of A were then normalized so that tr A = 0 and ∥ A ∥ F = 1 by subtracting the trace and dividing by the Frobenius norm. We chose initial activation σ 2 ( z ) = ReLU( z ) . We note that in both examples, κ = Θ(1) . As above, we optimize the hyperparameters η 1 , λ using grid search over a holdout validation set of size 2 15 and report the final error over a test set of size 2 15 .

To focus on the sample complexity and avoid width-related bottlenecks, we directly simulate the infinite width limit ( m 2 →∞ ) of Algorithm 1 by computing the kernel K in closed form. Finally, we run each trial with 5 random seeds and report the min, median, and max values in Figure 1.

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 f ∗ ( x ) = ReLU( x T 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. Figure 2 plots 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.

10 3

10 4

10 5

10 6

Samples

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

0

.

7

Test Loss

f

∗

(

x

) =

q

(

x T Ax

)

,

d

=

32

2 layer network

3 layer network

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 C be a fixed but sufficiently large constant.

Definition 4 (high probability events) . Let ι = C log( dnm 1 m 2 ) . We say that an event happens with high probability if it happens with probability at least 1 -poly( d, n, m 1 , m 2 ) e -ι .

Note that high probability events are closed under union bounds over sets of size poly( d, n, m 1 , m 2 ) . We will also assume throughout that ι ≤ C -1 d .

B.2 Tensor Notation

For a k -tensor T ∈ ( R d ) ⊗ k , let Sym be the symmetrization of T across all k axes, i.e

Next, given tensors A ∈ ( R d ) ⊗ a , B ∈ ( R d ) ⊗ b , we let A ⊗ B ∈ ( R d ) ⊗ a + b be their tensor product.

Definition 5 (Symmetric tensor product) . Given tensors A ∈ ( R d ) ⊗ a , B ∈ ( R d ) ⊗ b , we define their symmetric tensor product A ˜ ⊗ B ∈ ( R d ) ⊗ a + b as

Definition 6 (Tensor contraction) . Given a symmetric k -tensor T ∈ ( R d ) ⊗ k and an a -tensor A ∈ ( R d ) ⊗ a , where k ≥ a , we define the tensor contraction T ( A ) to be the k -a tensor given by

We note that ˜ ⊗ satisfies associativity.

When A is also a k -tensor, then ⟨ T, A ⟩ := T ( A ) ∈ R denotes their Euclidean inner product. We further define ∥ T ∥ 2 F := ⟨ T, T ⟩ .

C Univariate Approximation

Throughout this section, let µ ( b ) := exp ( -b 2 / 2 ) √ 2 π denote the PDF of a standard Gaussian.

Lemma 1. Let a ∼ Unif( {-1 , 1 } ) and let b ∼ N (0 , 1) . Then there exists v ( a, b ) supported on {-1 , 1 } × [0 , 2] such that for any | x | ≤ 1 ,

Proof. Let v ( a, b ) = c 1 b ∈ [1 , 2] where c = 1 µ (1) -µ (2) . Then for | x | ≤ 1 ,

Lemma 2. Let a ∼ Unif( {-1 , 1 } ) and let b ∼ N (0 , 1) . Then there exists v ( a, b ) supported on {-1 , 1 } × [0 , 2] such that for any | x | ≤ 1 ,

Proof. Let v ( a, b ) = ca 1 b ∈ [1 , 2] where c = 1 ∫ 1 0 µ ( b ) db . Then for | x | ≤ 1 ,

Lemma 3. Let f : [ -1 , 1] → R be any twice differentiable function. Then there exists v ( a, b ) supported on {-1 , 1 } × [0 , 2] such that for any | x | ≤ 1 ,

Proof. First consider v ( a, b ) = 1 b ∈ [0 , 1] µ ( b ) 2 f ′′ ( -ab ) . Then when x ≥ 0 we have by integration by parts:

1 2 where c 1 = f (0) -f (1) -f ( -1) + f ′ (1) -f ′ ( -1) and c 2 = f ′ (0) -f ′ (1) -f ′ ( -1) . In addition when x < 0 ,

so this equality is true for all x . By Lemmas 1 and 2 we can subtract out the constant and linear terms so there exists v ( a, b ) such that

In addition, using that µ ( b ) ≳ 1 for b ∈ [0 , 1] gives that this v ( a, b ) satisfies

D Proofs for Section 3

The following is a formal restatement of Theorem 1.

Theorem 6. Select q ∈ W 2 , ∞ ([ -1 , 1]) . Let η 1 = m 1 m 2 η , and n ≳ ∥ K f ∗ ∥ -2 L 2 ι 2 ℓ +2 α σ +1 , m 1 ≳ ∥ K f ∗ ∥ -2 L 2 ι , m 2 ≳ ∥ K f ∗ ∥ -2 L 2 ι 2 α σ +1 . There exists a choice of η = Θ( ι -χ ∥ Kf ∗ ∥ -1 L 2 ) , T = poly ( d, n, m 1 , m 2 , ∥ q ∥ 2 , ∞ ) , and λ, η 2 such that with high probability the output ˆ θ of Algorithm 1 satisfies the population L 2 loss bound

Proof of Theorem 6. The gradient with respect to w j is

Therefore

One then has

Lemma4. There exists a ∗ with a ∗ 2 ≲ q 2 , ∞ √ m 1 such that θ ∗ = ( a ∗ , W (1) , b (0) , V (0) ) satisfies

The second stage of Algorithm 1 is equivalent to random feature regression. The next lemma shows that there exists a ∗ with small norm that acheives low empirical loss on the dataset D 2 .

The proof of Lemma 4 is deferred to Appendix D.1. We first show that ϕ is approximately proportional to K f ∗ , and then invoke the ReLU random feature expressivity results from Appendix C.

Next, set

so that L 2 ( θ ∗ ) ≲ ∥ a ∗ ∥ 2 2 λ .

Define the regularized loss to be ˆ L ( a ) := L 2 ( ( a, W (1) , b (0) , V (0) ) ) + λ 2 ∥ a ∥ 2 2 . Let a ( ∞ ) = arg min a ˆ L ( a ) , and a ( t ) be the predictor from running gradient descent for t steps initialized at a (0) . We first note that

Next, we remark that ˆ L ( a ) is λ -strongly convex. Additionally, we can write f ( x ; ( a, W (1) , b (0) , V (0) )) = a T ψ ( x ) where ψ ( x ) = Vec ( 1 m 1 σ 1 ( a (0) j ηϕ ( x ) + b (0) j ) · 1 b (0) j > 0 ) . In Appendix D.2 we show sup x ∈D 2 ∥ ψ ( x ) ∥ ≲ 1 m 1 . Therefore

so ˆ L is O ( 1 m 1 ) + λ smooth. Choosing a learning rate η 2 = Θ( m 1 ) , after T = ˜ O ( 1 λm 1 ) = poly ( d, n, m 1 , m 2 , ∥ q ∥ 2 , ∞ ) steps we reach an iterate ˆ a = a ( T ) , ˆ θ = (ˆ a, W (1) , b (0) , V (0) )) satisfying

For τ > 0 , define the truncated loss ℓ τ by ℓ τ ( z ) = min( z 2 , τ 2 ) . We have that ℓ τ ( z ) ≤ z 2 , and thus

Consider the function class

The following lemma bounds the empirical Rademacher complexity of this function class Lemma 5. Given a dataset D , recall that the empirical Rademacher complexity of F is defined as

Then with high probability

Since ℓ τ loss is 2 τ -Lipschitz, the above lemma with B a = O ( ∥ a ∗ ∥ 2 ) = O ( ∥ q ∥ 2 , ∞ √ m 1 ) along with the standard empirical Rademacher complexity bound yields

Finally, we relate the ℓ τ population loss to the ℓ 2 population loss.

Lemma 6. Let τ = Θ(max( ι ℓ , ∥ q ∥ 2 , ∞ )) . Then with high probability over ˆ θ ,

Plugging in τ and applying Lemma 6 yields

as desired.

D.1 Proof of Lemma 4

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 f ∗ and the finite sample learned feature ϕ over the dataset D 2 .

Lemma 7. With high probability

The second lemma shows that with appropriate choice of η , the quantity ηϕ ( x ) is small.

Lemma 8. Let n ≳ ∥ K f ∗ ∥ -2 L 2 ι 2 ℓ +2 α σ +1 , m 1 ≳ ∥ K f ∗ ∥ -2 L 2 ι , and m 2 ≳ ∥ K f ∗ ∥ -2 L 2 ι 2 α σ +1 . There exists η = ˜ Θ( ∥ K f ∗ ∥ -2 L 2 ) such that with high probability, sup x ∈D 2 | ηϕ ( x ) | ≤ 1 and sup x ∈D 2 | η · K f ∗ ( x ) | ≤ 1

The third lemma expresses the compositional function q ◦ ηϕ via an infinite width network.

Lemma 9. Assume that n, m 1 , m 2 = ˜ Ω( ∥ K f ∗ ∥ -2 L 2 ) . There exists v : {± 1 } × R → R such that sup a,b | v ( a, b ) | ≤ ∥ q ∥ 2 , ∞ and, with high probability over D 2 , the infinite width network

satisfies for all x ∈ D 2 .

Proof of Lemma 4. Let v be the infinite width construction defined in Lemma 9. Define a ∗ ∈ R m 1 to be the vector with a ∗ j = v ( a (0) j , b (1) j ) . We can decompose

Take m 1 ≳ ι . The first term is the error between the infinite width network f ∞ v ( x ) and the finite width network f ( x ; θ ∗ ) . This error can be controlled via standard concentration arguments: by Corollary 2 we have that with high probability ∥ a ∗ ∥ ≲ ∥ q ∥ 2 , ∞ √ m 1 , and by Lemma 17 we have

Next, by Lemma 9 we get that the second term is zero with high probability. We next turn to the third term. Since q is ∥ q ∥ 1 , ∞ -Lipschitz on [ -1 , 1] , and sup x ∈D 2 | ηϕ ( x ) | ≤ 1 , sup x ∈D 2 | ( η · K f ∗ )( x ) | ≤ 1 by Lemma 8, we can apply Lemma 7 to get

and thus

Finally, we must relate the empirical error between q ◦ η · K f ∗ and f ∗ to the population error. This can be done via standard concentration arguments: in Lemma 19, we show

Altogether,

D.2 Proof of Lemma 5

Proof. The standard linear Rademacher bound states that for functions of the form x ↦→ w T ψ ( x ) with ∥ w ∥ 2 ≲ B w , the empirical Rademacher complexity is upper bounded by

We have that f ( x ; θ ) = a T ψ ( x ) , where

By Lemma 8, with high probability, sup x ∈D 2 | ηϕ ( x ) | ≤ 1 . Thus for x ∈ D 2

with high probability by Lemma 10. Altogether, we have

D.3 Proof of Lemma 6

Proof. We can bound

Next, we bound f ( x ; ˆ θ ) :

By Lemma 10, with high probability over the initialization we have ∥ b ∥ 2 ≲ √ m 1 . Next, by Lemma 4, with high probability over the initialization and dataset we have ∥ ˆ a ∥ 2 ≲ ∥ a ∗ ∥ 2 ≲ ∥ q ∥ 2 , ∞ √ m 1 . Thus with high probability we have uniformly over x .

We naively bound

By Lemma 11 and Lemma 12, with high probability we have ∣ ∣ f ( x i ; θ (0) ) -f ∗ ( x i ) ∣ ∣ ≲ ι ℓ for all i . With high probability, we also have ∥ ∥ h (0) ( x i ) ∥ ∥ ≲ √ m 2 . Additionally, we have

Since x is C γ subGaussian and ∥ v ∥ = 1 , we have

Altogether, and thus

By Assumption 2, E x [ ( f ∗ ( x )) 4 ] ≲ 1 .

We choose τ = Θ(max( ι ℓ , ∥ q ∥ 2 , ∞ ) . By Lemma 7 and Lemma 8 { ∣ ∣ ∣ f ( x, ˆ θ ) ∣ ∣ ∣ > τ } and {| f ∗ ( x ) | > τ } are both high probability events. Therefore by choosing C sufficiently large, we have the bound

Altogether, this gives

D.4 Auxiliary Lemmas

Proof of Lemma 7. We can decompose

By Lemma 11 and Lemma 13 we can upper bound the first term by √ ι m 1 .

For the second term, by Lemma 13 and Lemma 15 we have for all x ∈ D 2

The third term can be bounded as

Proof of Lemma 8. Conditioning on the event where Corollary 1 holds, and with the choice η = 1 2 C -1 K e -1 ∥ K f ∗ ∥ -1 L 2 ι -χ , we get that with high probability

Therefore

Proof of Lemma 9. Conditioning on the event where Lemma 8 holds and applying Lemma 3, we get that there exists v such that

for all x ∈ D 2 . Since v ( a, b ) = 0 for b ≤ 0 , we have that v ( a, b ) = v ( a, b ) 1 b> 0 . Thus the desired claim is true.

D.5 Concentration

Lemma 10. Let m 1 ≳ ι . With high probability,

Lemma 11. With high probability, ∣ ∣ f ( x ; θ (0) ) ∣ ∣ ≲ √ ι m 1 . for all x .

Proof. We have

Since a i ∼ Unif( {± 1 } ) and b i ∼ N (0 , 1) , the quantities a i σ 1 ( b i ) are 1 -subGaussian, and thus by Hoeffding with high probability we have

Lemma 12. With high probability

Proof. First, we have that σ 2 ( x · v ) 4 ≲ 1 + ( x · v ) 4 α σ . Therefore

Next, with high probability we have sup x ∈D 2 ∣ ∣ ∣ ∥ x ∥ 2 -E ∥ x ∥ 2 ∣ ∣ ∣ ≤ ιC 2 γ ≲ ι . Since E ∥ x ∥ 2 ≲ dγ 2 x ≲ d , we can thus bound

The proof for the other inequality is identical.

Lemma 13. Let m 2 ≳ ι 2 α σ +1 . With high probability, sup x ∈D 1 ∪D 2 ∥ ∥ h (0) ( x ) ∥ ∥ ≲ √ m 2

Proof. Since x is C γ subGaussian and v is 1 / √ d subGaussian, with high probability we have x · v ≤ O ( C γ ι ) = O ( ι ) .

Union bounding over x ∈ D 2 , j ∈ [ m 2 ] , with high probability we have sup x ∈D 1 ∪D 2 ,j ∈ [ m 2 ] | x · v j | ≤ O ( ι ) . Therefore

Next, see that

Pick truncation radius R = O ( ι α σ ) . By Hoeffding's inequality and a union bound over x ∈ D 1 ∪D 2 , we have that with high probability

Observe that

Therefore with high probability

by Lemma 12. Lemma 12 also tells us that E v [ σ 2 ( x · v ) 2 ] = O (1) . Altogether, for m 2 ≳ ι 2 α σ +1 we have

and hence ∥ ∥ h (0) ( x ) ∥ ∥ ≲ √ m 2 , as desired.

Lemma 14. P x ∼ ν ( | f ∗ ( x ) | ≥ C f eι ℓ ) ≤ e -ι and P x ∼ ν ( | ( K f ∗ )( x ) | ≥ C K eι χ · ∥ K f ∗ ∥ L 2 ) ≤ e -ι

Proof. By Markov's inequality, we have

Choose t = C f eι ℓ . We select q = ι e 1 -1 /ℓ , which is at least 1 for C in the definition of ι sufficiently large. Plugging in, we get

since ℓe 1 /ℓ -1 ≥ 1 .

An analogous derivation for the function K f ∗ ∥ K f ∗ ∥ L 2 yields the second bound

Corollary 1. With high probability, sup x ∈D 2 | f ∗ ( x ) | ≤ C f e · ι ℓ and sup x ∈D 2 | ( K f ∗ )( x ) | ≲ C K e · ι χ ∥ K f ∗ ∥ L 2

Proof. Union bounding the previous lemma over x ∈ D 2 yields the desired result.

Lemma 15. With high probability,

Proof. Consider the quantity

With high probability, sup x ∈D 2 ,j ∈ [ m 2 ] | f ∗ ( x i ) σ 2 ( x i · v j ) | ≲ ι ℓ + α σ . Pick truncation radius R = O ( ι ℓ + α σ ) . By Hoeffding, we have with high probability.

Furthermore, note that

Thus

By a union bound, the above holds with high probability for all j ∈ [ m 2 ] , and thus

Lemma 16. With high probability,

Proof. Fix x ∈ D 2 . Consider the random variables Z ( v ) = E x ′ [ f ∗ ( x ′ ) σ 2 ( x ′ · v ) σ 2 ( x · v )] . With high probability we have | σ 2 ( x · v j ) | ≲ ι α σ and thus

For all j ∈ [ m 2 ] . Choosing truncation radius R = ι α σ , with Hoeffding we have with high probability that

Next, we have that

Conditioning on the high probability event that sup j ∈ [ m 2 ] | Z ( v j ) | ≤ R , we have with high probability that

A union bound over x ∈ D 2 yields the desired result.

Lemma 17. With high probability,

Proof. Condition on the high probability event sup x ∈D 2 | ηϕ ( x ) | ≤ 1 . Next, note that whenever b > 2 that v ( a, b ) = 0 . Therefore we can bound

Therefore by Hoeffding's inequality we have that

The desired result follows via a Union bound over x ∈ D 2 .

Lemma 18. With high probability,

Thus by Hoeffding's inequality we have that

Corollary 2. Let m 1 ≳ ι . Then with high probability

Proof. By the previous lemma, we have that

Lemma 19. With high probability,

Proof. Note that

Thus

Proof. Let S be the set of x so that | η ( K f ∗ )( x ) | ≤ 1 and | f ∗ ( x ) | ≲ ι ℓ . Consider the random variables ( q ( η ( K f ∗ )( x )) -f ∗ ( x )) 2 · 1 x ∈ S . We have that

and

Therefore by Berstein's inequality we have that

Conditioning on the high probability event that x ∈ S for all x ∈ D 2 , we get that

E Proofs for Section 4

E.1 Single Index Model

Proof of Theorem 2. It is easy to see that Assumptions 1 and 3 are satisfied. By assumption g ∗ is polynomially bounded, i.e there exist constants C g , α g such that

Therefore

Thus Assumption 2 is satisfied with C f = C g 2 α g α α g / 2 g and ℓ = α g / 2 .

Next, we see that

Therefore

Furthermore, we have

Altogether, letting c 1 = E z ∼N (0 , 1) [ g ′ ( z )] = Ω(1) , we have ( K f ∗ )( x ) = 1 d c 1 x T w ∗ . Assumption 4 is thus satisfied with χ = 1 / 2 .

Next, see that ∥ K f ∗ ∥ L 2 = c 1 /d . We select the test function q to be q ( z ) = g ( η -1 d/c 1 · z ) , so that

Since η = Θ( dι -χ ) , we see that

Therefore we can bound the population loss as

E.2 Quadratic Feature

Throughout this section, we call x T Ax a degree 2 spherical harmonic if A is symmetric, E [ x T Ax ] = 0 , and E [ ( x T Ax ) 2 ] = 1 . Then, we have that Tr( A ) = 0 , and also

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 2 spherical harmonics is approximately x T Ax .

Lemma 20. Let q be a L -Lipschitz function with | q (0) | ≤ L , and let the target f ∗ be of the form f ∗ ( x ) = q ( x T Ax ) , where x T Ax is a spherical harmonic. Let c 1 = E z ∼N (0 , 1) [ q ′ ( z )] . Then

We defer the proof of this Lemma to Appendix E.2.1.

As a consequence, the learned feature K f ∗ is approximately proportional to x T Ax .

Lemma 21. Recall c 1 = E z ∼N (0 , 1) [ q ′ ( z )] . Then

Proof. Since E [ f ∗ ( x )] = 0 , P 0 f ∗ = 0 . Next, since f ∗ is an even function, P k f ∗ = 0 for k odd. Thus

Additionally, by Lemma 20 we have that

Since λ 2 2 ( σ ) = Θ( d -2 ) , we have

Corollary 3. Assume κ = o ( √ d ) . Then

Proof.

Proof of Theorem 3. By our choice of ν , we see that Assumption 1 is satisfied. We next verify Assumption 2. Since f ∗ is 1-Lipschitz, we can bound | g ∗ ( z ) | ≤ | g ∗ (0) | + | z | , and thus

where we used Lemma 35. Thus Assumption 2 holds with ℓ = 1 . Finally, we have

By Lemma 21 we have ∥ K f ∗ ∥ L 2 ≥ 1 2 λ 2 2 ( σ ) c 1 for d larger than some absolute constant. Next, by Lemma 35 we have for any q ≤ 1 4 d 2

Therefore

since λ 2 2 ( σ 2 ) = Ω( d -2 ) . Thus Assumption 4 holds with ℓ = 1 .

Next, observe that ∥ K f ∗ ∥ L 2 ≲ d -2 . We select the test function q to be q ( z ) = g ∗ ( η -1 ∥ K f ∗ ∥ -1 L 2 · z ) . We see that

and thus

where the first inequality follows from Lipschitzness of g ∗ , and the second inequality is Corollary 3. Furthermore since η = Θ( ∥ K f ∗ ∥ -1 L 2 ι -χ ) , we get that η -1 ∥ K f ∗ ∥ -1 L 2 = Θ( ι χ ) , and thus

Therefore by Theorem 6 we can bound the population loss as

E.2.1 Proof of Lemma 20

The high level sketch of the proof of Lemma 20 is as follows. Consider a second spherical harmonic x T Bx satisfying E [( x T Ax )( x T Bx )] = 0 (a simple computation shows that this is equivalent to Tr( AB ) = 0 ). Weappeal to a key result in universality to show that in the large d limit, the distribution of x T Ax converges to a standard Gaussian; additionally, x T Bx converges to an independent meanzero random variable. As a consequence, we show that and

The key universality theorem is the following.

From this, it immediately follows that P 2 f ∗ ≈ c 1 x T Ax .

Definition 7. For two probability measures µ, ν , the Wasserstein 1-distance between µ and ν is defined as

where ∥ f ∥ Lip := sup x = y | f ( x ) -f ( y ) | ∥ x -y ∥ 2 is the Lipschitz norm of f .

̸

Lemma 22 ([52][Theorem 9.20) . ] Let z ∼ N (0 , I d ) be a standard Gaussian vector, and let f : R d → R satisfy E [ f ( z )] = 0 , E [( f ( z )) 2 ] = 1 . Then where W 1 is the Wasserstein 1-distance.

We next apply this lemma to show that the quantities x T Ax + x T Bx and x T Ax + x · u are approximately Gaussian, given appropriate operator norm bounds on A,B .

Lemma 23. Let x T Ax and x T Bx be orthogonal spherical harmonics. Then, for constants c 1 , c 2 with c 2 1 + c 2 2 = 1 , we have that the random variable Y = c 1 x T Ax + c 2 x T Bx satisfies

Proof. Define the function f ( z ) = c 1 d z T Az ∥ z ∥ 2 + c 2 d z T Bz ∥ z ∥ 2 , and let x = z √ d ∥ z ∥ . Observe that when z ∼ N (0 , I ) , we have x ∼ Unif ( S d -1 ( √ d )) . Therefore f ( z ) is equal in distribution to Y . Define f 1 ( z ) = d z T Az ∥ z ∥ 2 . We compute

and

Thus and

∥ z ∥ 2 is distributed as a chi-squared random variable with d degrees of freedom, and thus

Therefore

and, using the fact that x and ∥ z ∥ are independent,

As a consequence, we have

Lemma 24. Let x T Ax be a spherical harmonic, and ∥ u ∥ = 1 . Then, for constants c 1 , c 2 with c 2 1 + c 2 2 = 1 , we have that the random variable Y = c 1 x T Ax + c 2 x T u satisfies

where W 1 is the 1-Wasserstein distance.

Proof. Define f 2 ( z ) = √ dz T u ∥ z ∥ . We have and

Thus so

We finish using the same argument as above.

Lemma 23 implies that, when ∥ A ∥ op , ∥ B ∥ op are small, ( x T Ax,x T Bx ) is close in distribution to the standard Gaussian in 2-dimensions. As a consequence, E [ q ( x T Ax ) x T Bx ] ≈ E [ q ( z 1 )] E [ z 2 ] = 0 , where z 1 , z 2 are i.i.d Gaussians. This intuition is made formal in the following lemma.

Lemma 25. Let x T Ax,x T Bx be two orthogonal spherical harmonics. Then

Proof. Define the function F ( t ) = ∫ t 0 q ( s ) ds . Then F ′ ( t ) = q ( t ) , so by a Taylor expansion we get that

Therefore

Pick truncation radius R , and define the function F ( z ) = F (max( -R, min( R,z )) . F has Lipschitz constant sup z ∈ [ -R,R ] | q ( z ) | , and thus since W 1 ( x T Ax, N (0 , 1) ) ≲ ∥ A ∥ op , we have

Next, we have

∣ ∣ E [ F ( x T Ax ) -F ( x T Ax ) ]∣ ∣ ≤ E [ 1 | x T Ax | >R ∣ ∣ F ( x T Ax ) ∣ ∣ ] ≤ P ( ∣ ∣ x T Ax ∣ ∣ > R ) · E [ F ( x T Ax ) 2 ] 1 / 2 . Likewise,

Since q is L -Lipschitz, we can bound | F ( z ) | ≤ | z || q (0) | + 1 2 L | z | 2 , and thus

The standard Gaussian tail bound yields P ( | z | > R ) ≲ exp ( -C 1 R 2 ) for appropriate constant C 1 , and polynomial concentration yields P (∣ ∣ x T Ax ∣ ∣ > R ) ≲ exp( -C 2 R ) for appropriate constant C 2 . Thus choosing R = C 3 log d for appropriate constant C 3 , we get that

Altogether, since | q ( z ) | ≤ | q (0) | + L | z | , we get that

By an identical calculation, we have that for ϵ < 1 ,

Altogether, we get that

Via a simple calculation, one sees that

Therefore so

Setting ϵ = √ ∥ A ∥ op log d yields as desired.

Similarly, we use the consequence of Lemma 24 that ( x T Ax,u T x ) is close in distribution to a 2d standard Gaussian, and show that E [ q ( x T Ax ) ( ( u T x ) 2 -1 ) ] ≈ E [ q ( z 1 ) ( z 2 2 -1 ) ] = 0 .

Lemma 26.

Proof. Let G ( t ) = ∫ t 0 F ( s ) ds . Then G ′ ( t ) = F ( t ) , G ′′ ( t ) = q ( t ) , so a Taylor expansion yields

Thus

Therefore

For truncation radius R , define G ( z ) = G (max( -R, min( R,z ))) . We get that G has Lipschitz constant sup | z |≤ R | F ( z ) | ≲ LR 2 . Therefore

and by a similar argument in the previous lemma, setting R = C 3 log d yields

Altogether,

By an identical calculation,

Additionally, letting z, w be independent standard Gaussians,

Altogether,

where we set ϵ = ∥ A ∥ 1 / 3 op log 2 / 3 d .

Lemma 25 shows that when ∥ B ∥ op ≪ 1 , E [ q ( x T Ax ) x T Bx ] ≈ 0 . However, we need to show this is true for all spherical harmonics, even those with ∥ B ∥ op = Θ(1) . To accomplish this, we decompose B into the sum of a 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,B be orthogonal spherical harmonics. Then

Proof. Let τ > ∥ A ∥ op be a threshold to be determined later. Decompose B as follows:

where

By construction, we have,

and

There are at most O ( τ -2 ) indices i satisfying | λ i | > τ , and thus

We thus compute that

and

Next, since ∣ ∣ E [ q ( x T Ax ) ] -E z [ q ( x )] ∣ ∣ ≲ L ∥ A ∥ op , Lemma 26 yields

Finally,

Altogether,

Finally, we use the fact that x T Ax is approximately Gaussian to show that E [ q ( x T Ax ) x T Ax ] ≈ c 1 . Lemma 28. Let x T Ax be a spherical harmonic. Then

Proof. Define H ( z ) = q ( z ) z . For truncation radius R , define H ( z ) = H (max( -R, min( R,z ))) . For x, y ∈ [ -R,R ] we can bound

Thus H has Lipschitz constant O ( RL ) . Since W 1 ( x T Ax, N (0 , 1)) ≲ ∥ A ∥ op , we have

Furthermore, choosing R = C log d for appropriate constant C , we have that

Altogether,

Substituting H ( x T Ax ) = q ( x T Ax ) x T Ax and E z ∼N (0 , 1) [ H ( z )] = E z [ q ( z ) z ] = E z [ q ′ ( z )] = c 1 yields the desired bound.

We are now set to prove Lemma 20.

Proof of Lemma 20. Let P 2 f ∗ = x T T 2 x . Then ∥ ∥ P 2 f ∗ -c 1 x T Ax ∥ ∥ L 2 ≲ ∥ T 2 -c 1 A ∥ F .

Write T 2 = αA + A ⊥ , where ⟨ A ⊥ , A ⟩ = 0 . We first have that

Therefore ∥ ∥ A ⊥ ∥ ∥ F ≲ L ∥ A ∥ 1 / 6 op log d . Next, see that

so by Lemma 28 we have | c 1 -α | ≲ L ∥ A ∥ op log d . Altogether,

E.3 Improved Error Floor for Polynomials

When q is a polynomial of degree p = O (1) , we can improve the exponent of d in the error floor.

Theorem 7. Assume that q is a degree p polynomial, where p = O (1) . Under Assumption 5, Assumption 6, and Assumption 7, with high probability Algorithm 1 satisfies the population loss bound

The high level strategy to prove Theorem 7 is similar to that for Theorem 3, as we aim to show K f ∗ is approximately proportional to x T Ax . Rather to passing to universality as in Lemma 20, however, we use an algebraic argument to estimate P 2 f ∗ .

The key algebraic lemma is the following:

Lemma 29. Let Tr( A ) = 0 . Then

where the constants d k,s are defined by and we denote ( -1)!! = 1 .

Proof. The proof proceeds via a counting argument. We first have that

Consider any permutation σ ∈ S 2 k . We can map this permutation to the graph G ( σ ) on k vertices and k -1 edges as follows: for m ∈ [ k -1] , if σ -1 (2 m -1) ∈ { 2 a -1 , 2 a } and σ -1 (2 m ) ∈ { 2 b -1 , 2 b } , then we draw an edge e ( m ) between a and b . In the resulting graph G ( σ ) each node has degree at most 2 , and hence there are either two vertices with degree 1 or one vertex with degree 0. For a vertex v , let e 1 ( v ) , e 2 ( v ) ∈ [ k -1] be the two edges v is incident to if v has degree 2 , and otherwise e 1 ( v ) be the only edge v is incident to. For shorthand, let ( i 1 , . . . , i 2 k ) = ( α 1 , α 1 , . . . , α k -1 , α k -1 , i, j ) .

If there are two vertices ( u 1 , u 2 ) with degree 1, we have that

̸

Let u 1 , u 2 be connected to eachother via a path of s total vertices, and let P be the ordered set of vertices in this path. Via the matrix multiplication formula, one sees that

where the sum is over the k -s α 's that are still remaining in { e i ( v ) : v / ∈ P}

Likewise, if there is one vertex u 1 with degree 0, we have

̸

and thus, since P = { u 1 }

Altogether, we have that

where s, P are defined based on the graph G ( σ ) . Consider a graph with fixed path P , and let S P be the set of permutations which give rise to the path P . We have that

There are ( k -1) · · · ( k -s +1) choices for the k edges to use in the path, and at each vertex v there are two choices for which edge should correspond to 2 v or 2 v +1 . Additionally, there are 2 ways to orient each edge. Furthermore, there are k ! ( k -s )! ways to choose the ordering of the path. Altogether, there are 2 2 s -1 ( k -1)! ( k -s )! k ! ( k -s )! ways to construct a path of length s . We can thus write

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 sum over the unused α 's. Reindexing, this latter sum is (letting ( i 1 , . . . , i 2 k -2 s ) = ( α 1 , α 1 , . . . , α k -s , α k -s ) )

Altogether, we obtain

as desired.

Definition 8. Define the operator T : R d × d → R d × d by T ( M ) = M -Tr( M ) · I d . Lemma 30. Let P 2 f ∗ ( x ) = x T T 2 x . Then ∥ ∥ T 2 -E x [ ( g ∗ ) ′ ( x T Ax ) ] · A ∥ ∥ F ≲ κ √ d .

Proof. Throughout, we treat p = O (1) and thus functions of p independent of d as O (1) quantities. Let q be of the form g ∗ ( z ) = ∑ p k =0 α k z k . We then have

Therefore P 2 f ∗ ( x ) = x T T 2 x , where

Applying Lemma 29, one has

Define

We first see that

Next, see that

Thus

since

where Tr( A ) = 0 implies E x [ ( x T Ax ) 2 ] = O (1) and we invoke spherical hypercontractivity (Lemma 35). Similarly, | β s | = O (1) , and thus

where we use the inequality

Lemma 31. Let c 1 = E x [ ( g ∗ ) ′ ( x T Ax ) ] . Then ∥ ∥ K f ∗ -λ 2 2 ( σ ) c 1 x T Ax ∥ ∥ L 2 ≲ κd -5 / 2

along with for s ≥ 2 .

Proof. Since E [ f ∗ ( x )] = 0 , P 0 f ∗ = 0 . Next, since f ∗ is an even function, P k f ∗ = 0 for k odd. Thus

Additionally, by Lemma 30 we have that

Since λ 2 2 ( σ ) = Θ( d -2 ) , we have

Corollary 4. Assume κ = o ( √ d ) . Then

Proof.

The proof of Theorem 7 follows directly from Corollary 4 in an identical manner to the proof of Theorem 3.

F Preliminaries on Spherical Harmonics

In this section we restrict to X d = S d -1 ( √ d ) , the sphere of radius √ d , and ν the uniform distribution on X d .

The moments of ν are given by the following [18]:

Lemma 32. Let x ∼ ν . Then where

For integer ℓ ≥ 0 , let V d,ℓ be the space of homogeneous harmonic polynomials on R d of degree ℓ restricted to X d . One has that V d,ℓ form an orthogonal decomposition of L 2 ( ν ) [27], i.e

Homogeneous polynomials of degree ℓ can be written as T ( x ⊗ ℓ ) for an ℓ -tensor T ∈ ( R d ) ⊗ ℓ . The following lemma characterizes V d,ℓ :

Lemma 33. T ( x ⊗ ℓ ) ∈ V d,ℓ if and only if T ( I ) = 0 .

Proof. By definition, a degree l homogeneous polynomial p ( x ) ∈ V d,l if and only if ∆ p ( x ) = 0 for all x ∈ S d -1 . Note that ∇ 2 p ( x ) = l ( l -1) T ( x ⊗ ( l -2) ) so this is satisfied if and only if

As this must hold for all x , this holds if and only if T ( I ) = 0 .

From the above characterization, we see that dim ( V d,k ) = B ( d, k ) , where

Define P ℓ : L 2 ( ν ) → L 2 ( ν ) to be the orthogonal projection onto V d,ℓ . The action of P 0 , P 1 , P 2 on a homogeneous polynomial is given by the following lemma:

Lemma 34. Let T ∈ ( R d ) ⊗ 2 k be a symmetric 2 k tensor, and let p ( x ) = T ( x ⊗ 2 k ) be a polynomial. Then:

Proof. First, we see

The LHS is

where the last step is true since Tr( B ) = 0 . The RHS is

Since these two quantities must be equal for all B with Tr( B ) = 0 , and Tr( T 2 ) = 0 , we see that

as desired.

Polynomials over the sphere verify hypercontractivity:

Lemma 35 (Spherical hypercontractivity [11, 38]) . Let f be a degree p polynomial. Then for q ≥ 2

Next, since p is even, P 1 p = 0 . Next, let P 2 p = x T T 2 x . For symmetric B so that Tr( B ) = 0 , we have that

F.1 Gegenbauer Polynomials

For an integer d > 1 , let µ d be the density of x · e 1 , where x ∼ Unif( S d -1 (1)) and e 1 is a fixed unit vector. One can verify that µ d is supported on [ -1 , 1] and given by

where Γ( n ) is the Gamma function. For convenience, we let Z d := Γ( d/ 2) √ π Γ( d -1 2 ) = 1 β ( 1 2 , d -1 2 ) denote the normalizing constant.

The Gegenbauer polynomials ( G ( d ) k ) k ∈ Z ≥ 0 are a sequence of orthogonal polynomials with respect to the density µ d , defined as G ( d ) 0 ( x ) = 1 , G ( d ) 1 ( x ) = x , and

By construction, G ( d ) k is a polynomial of degree k . The G ( d ) k are orthogonal in that

For a function f ∈ L 2 ( µ d ) , we can write its Gegenbauer decomposition as

where convergence is in L 2 ( µ d ) . For an integer k , we define the operator P ( d ) k : L 2 ( µ d ) → L 2 ( µ d ) to be the projection onto the degree k Gegenbauer polynomial, i.e

We also define the operators P ( d ) ≤ k = ∑ k ℓ =0 P ( d ) ℓ and P ( d ) ≥ k = ∑ ∞ ℓ = k P ( d ) ℓ .

Recall that ν = Unif( S d -1 ( √ d )) . Let ˜ µ d be the density of x · e 1 , where x ∼ ν . For a function σ ∈ L 2 (˜ µ d ) , we define its Gegenbauer coefficients as

By Cauchy, we get that | λ k ( σ ) | ≤ ∥ σ ∥ L 2 (˜ µ d ) B ( d, k ) -1 / 2 = O ( ∥ σ ∥ L 2 (˜ µ d ) d -k/ 2 ) .

A key property of the Gegebauer coefficients is that they allow us to express the kernel operator K in closed form [27, 39]

Lemma 36. For a function g ∈ L 2 ( ν ) , the operator K acts as

One key fact about Gegenbauer polynomials is the following derivative formula:

Lemma 37 (Derivative Formula) .

Furthermore, the following is a corollary of eq. (24):

Corollary 5.

G Proofs for Section 5

The proof of Theorem 4 relies on the following lemma, which gives the Gegenbauer decomposition of the ReLU function:

Lemma 38 (ReLU Gegenbauer) . Let ReLU( x ) = max( x, 0) . Then

As a consequence,

The proof of this lemma is deferred to Appendix G.1.

We also require a key result from [19], which lower bounds the approximation error of an inner product function.

Definition 9. f : S d -1 (1) ×S d -1 (1) → R is an inner product function if f ( x, x ′ ) = ϕ ( ⟨ x, x ′ ⟩ ) for some ϕ : [ -1 , 1] → R .

Definition 10. g : S d -1 (1) ×S d -1 (1) → R is a separable function if g ( x, x ′ ) = ψ ( ⟨ v, x ⟩ , ⟨ v ′ , x ′ ⟩ ) for some v, v ′ ∈ S d -1 (1) and ψ : [ -1 , 1] 2 → R .

Let ˜ ν d be the uniform distribution over S d -1 (1) × S d -1 (1) . We note that if ( x, x ′ ) ∼ ˜ ν d , then ⟨ x, x ′ ⟩ ∼ µ d . For an inner product function f , we thus have ∥ f ∥ L 2 (˜ ν d ) = ∥ ϕ ∥ L 2 ( µ d ) . We overload

Lemma 39. [19, Theorem 3] Let f be an inner product function and g 1 , . . . , g r be separable functions. Then, for any k ≥ 1 ,

We now can prove Theorem 4

Proof of Theorem 4. We begin with the lower bound. Let x = U ( x 1 x 2 ) , where x 1 , x 2 ∈ R d/ 2 . Assume that there exists some θ such that ∥ f ∗ -N θ ∥ L 2 ( X ) ≤ ϵ . Then

where r is the random variable defined as r = ∥ x 1 ∥ 2 and µ is the associated measure. The equality comes from the fact that conditioned on r , x 1 and x 2 are independent and distributed uniformly on the spheres of radii √ r and √ d -r , respectively. We see that E r = d/ 2 , and thus

Let δ = inf r ∈ [ d/ 4 , 3 d/ 4] E x 1 ∼S d -1 ( √ r ) ,x 2 ∼S d -1 ( √ d -r ) ( f ∗ ( x ) -N θ ( x )) 2 . We get the bound ϵ 2 ≥ δ · P ( r ∈ [ d/ 4 , 3 d/ 4]) ,

and thus δ ≤ ϵ 2 (1 -exp( -Ω( d ))) . Therefore there exists an r ∈ [ d/ 4 , 3 d/ 4] such that

Next, see that when ∥ x 1 ∥ 2 = r, ∥ x 2 ∥ 2 = d -r , we have that

where now x 1 , x 2 ∼ S d/ 2 -1 (1) i.i.d. Defining q ( z ) = ReLU ( 2 √ r ( d -r ) d z ) -c 0 , we thus have

WU ( √ r · I 0 0 √ d -r · I ) yields a network so that N θ ( x ) = N θ ( x ) . Therefore we get that the new network N θ satisfies

Furthermore, defining x = ( x 1 x 2 ) , choosing the parameter vector θ = ( a, W, b 1 , b 2 ) , where W =

where x 1 , x 2 are drawn i.i.d over Unif( S d/ 2 -1 (1)) .

We aim to invoke Lemma 39. We note that ( x 1 , x 2 ) ∼ ˜ ν d/ 2 , and that q is an inner product function. Define g i ( x ) = a i σ ( w T i x + b 1 ,i ) . We see that g i is a separable function, and also that

Hence N θ ) is the sum of m +1 separable functions. We can bound the a single function as

since ∥ ∥ W ∥ ∥ ∞ ≤ max( √ r, √ d -r ) ∥ WU ∥ ∞ ≤ Bd/ 2 . Therefore by Lemma 39

Simplifying, we have that

By Gautschi's inequality, we can bound

Therefore

for k ≤ d/ 4 . Altogether,

for m d/ 8 . Since q ( z ) = ReLU(2 √ r ( d -r ) d z ) c 0 = 2 √ r ( d -r ) d ReLU( z ) c 0 , we have that

We thus have, for any integer k < d/ 8 ,

Choose ϵ ≤ 1 512 k 2 ; we then must have

or

for any k ≤ C 3 d . Selecting k = ⌊ √ 1 512 ϵ ⌋ yields

for ϵ less than a universal constant c 3 .

We next show the upper bound,

It is easy to see that Assumptions 1 and 3 are satisfied. Next, since the verification of Assumptions 2 and 4 only required Lipschitzness, those assumptions are satisfied as well with ℓ, χ = 1 . Finally, we have

.

Next, observe that ∥ K f ∗ ∥ L 2 ≲ d -2 . Define A = √ d +2 2 d A . This scaling ensures ∥ ∥ x T Ax ∥ ∥ L 2 = 1 Then, we can write f ∗ ( x ) = g ∗ ( x T Ax ) for g ∗ ( z ) = √ 2 d d +2 ReLU( z ) -c 0 . For ϵ > 0 , define the smoothed ReLU ReLU ϵ ( z ) as

One sees that ReLU ϵ is twice differentiable with ∥ ReLU ϵ ∥ 1 , ∞ ≤ 1 and ∥ ReLU ϵ ∥ 2 , ∞ = 1 2 ϵ We select the test function q to be q ( z ) = √ 2 d d +2 ReLU ϵ ( η -1 K f ∗ -1 L 2 z ) c 0

and thus

where the first inequality follows from Lipschitzness and the second inequality is Corollary 3, using κ = 1 .

There exists a constant upper bound for the density of x T Ax , and thus we can upper bound

Furthermore since η = Θ( ∥ K f ∗ ∥ -1 L 2 ι -χ ) , we get that η -1 ∥ K f ∗ ∥ -1 L 2 = Θ( ι χ ) , and thus

Therefore by Theorem 6 we can bound the population loss as

Choosing ϵ = d -1 / 4 yields the desired result. As for the sample complexity, we have ∥ q ∥ 2 , ∞ = ˜ O ( ϵ -1 ) = ˜ O ( d 1 / 4 ) , and so the runtime is poly( d, m 1 , m 2 , n ) .

G.1 Proof of Lemma 38

Proof of Lemma 38. For any integer k , we define the quantities A ( d ) 2 k , B ( d ) 2 k +1 as

We also let Z d = Γ( d/ 2) √ π Γ( d -1 2 ) to be the normalization constant.

∥ ∥ · -. We see that ,

Integration by parts yields

From Corollary 5 we have

Thus

The recurrence formula yields

I claim that

We proceed by induction on k . For the base cases, we first have

where we use the substitution u = (1 -x 2 ) d -3 2 . Next,

Next, eq. (25) gives

Finally, eq. (26) gives

Therefore the base case is proven for k = 0 , 1 .

Now, assume that the claim is true for some k ≥ 1 for all d . We first have

Next, we have

Therefore by induction the claim holds for all k, d .

The Gegenbauer expansion of ReLU is given by

Note that ReLU( x ) = 1 2 ( x + | x | ) . Since | x | is even, the only nonzero odd Gegenbauer coefficient is for G ( d ) 1 . In this case,

Also, B ( d, 1) = d . Next, we see that

Plugging in our derivation for A ( d ) 2 k gives the desired result.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nNothing particular.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe discussion of limitations is adequate.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\nWhat 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
-	This 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.
-	It’s good that the paper has an example section (Sec. 4.)

WEAKNESSES:
-	The “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.
-	This work still does not answer the same questions for convolutional networks.
-	This 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.)
-	In other words, this issue is about generalizability on depth scalability.
-	This paper lacks discussion of its own limitations.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-24