Learning Hierarchical Polynomials of Multiple Nonlinear Features
Authors: Hengyu Fu
Zihao Wang
Eshaan Nichani
Jason D. Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Deep Learning Theory
Feature Learning
Three-Layer Neural Network
Gradient Descent
Sample Complexity
Primary area: learning theory
TL;DR: We investigate the sample complexity of learning hierarchical polynomials of multiple nonlinear features using three-layer neural networks.

Abstract: In deep learning theory, a critical question is to understand how neural networks learn hierarchical features. In this work, we study the learning of hierarchical polynomials of multiple nonlinear features using three-layer neural networks. We examine a broad class of functions of the form $f^{\star}=g^{\star}\circ \mathbf{p}$, where $\mathbf{p}:\mathbb{R}^{d} \rightarrow \mathbb{R}^{r}$ represents multiple quadratic features with $r \ll d$ and $g^{\star}:\mathbb{R}^{r}\rightarrow \mathbb{R}$ is a polynomial of degree $p$. This can be viewed as a nonlinear generalization of the multi-index model, and also an expansion upon previous work on nonlinear feature learning that focused only on a single feature (i.e. $r = 1$). 
    Our primary contribution shows that a three-layer neural network trained via layerwise gradient descent suffices for  
 - complete recovery of the space spanned by the nonlinear features 
- efficient learning of the target function $f^{\star}=g^{\star}\circ \mathbf{p}$ or transfer learning of $f=g\circ  \mathbf{p}$ with a different link function
   
within $\widetilde{\mathcal{O}}(d^4)$ samples and polynomial time.
For such hierarchical targets, our result substantially improves the sample complexity ${\Theta}(d^{2p})$ of the kernel methods, demonstrating the power of efficient feature learning. It is important to highlight that our results leverage novel techniques and thus manage to go beyond all prior settings such as single-index and multi-index models as well as models depending just on one nonlinear feature, contributing to a more comprehensive understanding of feature learning in deep learning.

Published as a conference paper at ICLR 2025
LEARNING HIERARCHICAL POLYNOMIALS OF MULTI- PLE NONLINEAR FEATURES
Hengyu Fu
Zihao Wang
Eshaan Nichani
Jason D. Lee
ABSTRACT
In deep learning theory, critical question is to understand how neural networks learn hierarchical features. In this work, we study the learning of hierarchical polynomials of multiple nonlinear features using three-layer neural networks. We examine a broad class of functions of the form I = g" p. where p Rd + IR" represents multiple quadratic features with polynomial of degree p. This can be viewed as a nonlinear generalization of the multi-index model (Damian et al., 2022), and also an expansion upon previous work that focused only on a single nonlinear feature, i.c. r = (Nichani et al., 2023; Wang al., 2023).
and g" R IR is
Our primary contribution shows that a three-layer neural network trained via lay- erwise gradient descent suffices for
complete recovery of the space spanned by the nonlinear features
efficient learning of the target function f* g" p transfer learning of f=go p with different link function
within O(d') samples and polynomial time. For such hierarchical targets, our result substantially improves the sample complexity (dPP) of the kernel methods, demonstrating the power of efficient feature learning. It is important to highlight that our results leverage novel techniques and thus manage to go beyond all prior settings such single-index and multi-index models well models depending just on one nonlinear feature, contributing to a more comprehensive understanding of feature learning in deep learning.
INTRODUCTION
Deep neural networks have achieved remarkable empirical success across numerous domains of artificial intelligence (Krizhevsky et al., 2012; He al., 2016). This success can be largely attributed to their ability to extract latent features from real-world data and decompose complex targets into hierarchical representations, which improves test accuracy (He et al., 2016) and allows efficient transfer learning (Devlin, 2018). These feature learning capabilities are widely regarded as a core strength of neural networks over non-adaptive approaches such as kernel methods (Wei et al., 2020; Bai and Lee, 2020).
Despite these empirical achievements, the feature learning capabilities of neural networks are less well understood from theoretical point of view. Previous work on feature learning has shown that two-layer neural networks can learn multiple linear features of the input (Damian et al., 2022), that is, multi-index models. However, the two-layer architecture inherently limits the network's ability to represent and learn nonlinear features (Daniely, 2017). Given that many real-world scenarios involve diverse and nonlinear features, recent studies have shifted focus to investigating the learning of nonlinear features using deeper neural networks. Safran and Lee (2022); Ren et al. (2023); Nichani al. (2023); Wang et (2023) have demonstrated that three-layer networks, when trained via gradient descent, can efficiently learn hierarchical targets of the form h 9 p. where represents certain types of features such as the norm |x] or quadratic form x T Ax. However, these studies are limited to relatively simple hierarchical functions and mainly focus targets of a single feature. remains unclear whether neural networks can efficiently learn wider range of
"Peking University. Email: fhy2021estu.po www.auu.kkkkuc.uc.omoooom. Stanford University. Email: zihaow8stanfos edu Princetton University. Email: {eshni ch, jasonaae))@ppppp.. iprinttnnntt

---

Published as a conference paper at ICLR 2025
hierarchical functions, particularly those that depend multiple nonlinear features. This leads us to the following central question:
Can neural networks adaptively identfy multiple nonlinear features from the hierarchical targets by gradient descent, thereby allowing an efficient learning for such targets?
1.1 MAIN CONTRIBUTIONS
In this paper, we provide strong theoretical evidence that three-layer neural networks have the ability to learn multiple hidden nonlinear features. Specifically, we study the problem learning any hierarchical polynomial with multiple quadratic features using three-layer network trained via layer-wise gradient descent. Our main contributions are summarized as follows:
A Novel Analytic Framework for Multi-Nonlinear Feature Learning. We demonstrate that when the target function belongs to a broad class of the form f = g" o P. where p Rd + I' represents guadratic (nonlincar) features and g" is link function, the first step of gradient descent efficiently learns and recovers the space spanned by these nonlinear features p within only (dd)) samples. We remark that our proof techniques are also applicable to general nonlinear features. The core technical novelty is that we develop novel and general universality argument (Lemma 1) that bridges multi nonlinear feature models to multi-index models, which allows for an accurate reconstruction of the features through a simple linear transformation on the learned representations with small approximation error (Proposition I)
Improved Sample Complexity and Efficient Transfer Learning. Leveraging the learned features in the first GD step, we prove that when the link function g a polynomial of degree p. the gradient descent on the outer layer can achieves a vanishing generalization error with small outer width and at most O(rO())) additional training samples, removing the dependence on d (Theorem 1). This significantly improves upon the sample complexity of kernel methods, which require (ddP) samples. Moreover, our analysis enables efficient transfer learning for any other target function of the form f = g with different link function g. which also only requires (p))))) additional samples.
1.2 RELATED WORKS
Kernel Methods. Earlier research links the behavior of gradient descent (GD) the entire net- work to its linear approximation near the initialization. In this scenario, neural networks act as kernels, known as the Neural Tangent Kernel (NTK). This connection bridges neural network anal- ysis with established kernel theory and offers initial learning guarantees for neural networks (Jacot al., 2018: Soltanolkotabi
al., 2018; Du al., 2018: Chizat et al., 2019; Arora et al., 2019). However, kernel theory fails to explain the superior empirical achievements of neural networks over kernel methods (Arora et al., 2019; Lee al., 2020: E et al., 2020). Networks in the kernel regime fail to learn features (Yang and Hu, 2021), not adaptable to hierarchical structures of real world targets. Ghorbani et al. (2021) proves that for uniformly distributed data on the sphere, the NTK method requires (((d)) samples to learn any polynomials of degree k in dimensions, which is impractical when i is large. Thus, central question is how neural networks can detect and capture the underlying hierarchies in the target functions, which allows for a better generalization behavior versus kernel methods.
Learning Linear Features. Recent studies have demonstrated neural networks' capability to learn hierarchical functions of linear features more efficiently than kernel methods. Specifically, Bietti et al. (2022); Ba et al. (2022) establish the efficient learning of single-index models, i.e., (x) = ((( x)). Furthermore, recent works Damian et al. (2022); Abbe et al. (2023); Dandi et al. (2023a); Bietti et al. (2023) further demonstrate that for isotropic data, two-layer or three-layer neural networks can effectively learn multi-indes models of the form f (x) g( (x). These studies adopt certain modified training algorithms, such as layer-wise training. With sufficient feature learn- ing, these networks can learn low-rank polynomials with a benign sample complexity of (d)(1)), which does not scale with the degree of the polynomial g. Empirically, fully connected networks trained via gradient descent on image classification tasks also capture low-rank features (Lee et al., 2007; Radhakrishnan et al., 2022). More recently, the learning of indel--indx and multi-index models is analyzed with more advanced algorithm framework or specified data structure. Mousavi- Hosseini al. (2024) considers learning general multi-index models with two- neural networks
2

---

Published as a conference paper at ICLR 2025
through mean mean-field Langevin dynamics, Dandi et al. (2024b); Lee al. (2024) goes beyond the traditional Correlational Statistical Query (CSQ) setting and consider algorithms that reuse samples for feature learning. MoousaviHHossein et al. (2023); Ba et al. (2023); Wang et al. (2024) considers learning linear features with structured data (such as data with a spiked covariance) rather than the commonly considered isotropic one. Cui et al. (2024); Dandi et al. (2024a) study the spectral struc- ture revealed in the learned features with one huge gradient step through spiked random feature model understand the mechanism of feature learning neural networks.
Learning Nonlinear Features. Previous studies indicate that neural networks can effectively learn specific hierarchies of nonlinear features. Safran and Lee (2022) shows that GD can efficiently learn functions such as 1 ||x| > with a three-layer network. Ren et al. (2023) demonstrates that ReLU |xx]) can be learned by a multi-layer mean-field network. Moniri et al. (2024) studies the nonlinear feature learning capabilities of two-layer neural networks with one step of gradient descent. Allen- Zhu and Li (2019; 2020) explore learning target functions of the form p + og o with p being the underlying feature through a three-layer residual network, though they either need ou(1) or cannot reach vanishing error. More recent works have addressed broader class of nonlinear features compared with the previous research and demonstrate that three-layer neural networks can learn these hidden features efficiently. Specifically, Nichani et al. (2023) demonstrates that a three-layer network trained with layer-wise GD algorithm effectively learns gop for a quadratic feature p(x) =
XT Ax with an improved sample complexity of ((d'). Wang et al. (2023) further demonstrates that such a network can in fact efficiently learn g o p for p within a broad subclass of degree polynomials and optimizes the sample complexity to O(d)). However, all of these studies focus on single nonlinear feature, limiting their applicability to scenarios involving multiple features. Our work addresses this gap by establishing the first theoretical guarantee for efficiently learning hierarchical polynomials of multiple nonlinear features, which significantly broadens the learnable function class and advances towards a better understanding of feature learning.
PRELIMINARIES
2.1 NOTATIONS
We use bold letters to denote vectors and matrices. For vector v, we denote its Euclidean norm by || For a matrix A, we denote its operator and Frobenius norm as || and ||A]| respectively. For any positive integer n,, we denote [n]] {1.2,... .n.. Moreover, for any indexes i and j. we denote dip- = if j j and 0 otherwise. We use O, θ and f2 to hide absolute constants. In addition,
we denote f S when there exists some positive absolute constant C' with f < Cg. We use Õ θ and Zi to ignore logarithmic terms. For function f X R and distribution on X, denote || [l = (Ex~o (xx) L/P When the domain is clear from context, we write I| for simplicity. Finally, we write Eg as the shorthand for Expr sometimes.
2.2 PROBLEM SETUP
Data distribution Our aim is to learn the target function f X' R, with X S Rd being the input space. Throughout the paper, we assume X' = gd-1(v that is, the sphere with radius vd in d dimensions. Also, we consider the data distribution to be the uniform distribution on the sphere, i.e., x ~ (X'), and we draw two independent datasets D1, D2, each with nı and n2 i...d samples, respectively. Thus, we draw n, + n2 samples in total.
Target function For the target function ft Rdd + R. we assume they are hierarchical functions of r quadratic features
x This structure represents a broad class of functions where p(x)= = [x A1x, x represents r quadratic features, and g" IR' - R link function. Here we consider the case < d. To simplify our analysis while maintaining generality, we make the following assumptions:
(x() (p)=)((p((x)) g (x AIX, X A,x)
Assumption (Orthogonal quadratic features). For any i, E [r], we suppose
Ex
[x
T A = O, Ex [[x] T (xx) A,x)] = dig and |Aillop S
K1 Vd
Here we assume K1
3

---

Published as a conference paper at ICLR 2025
The first assumption is equivalent to tr(),) 0 for any E [r]. For A, such that tr(A,) # 0, could simply subtract the mean of the feature to A = A, - (tr(A.)/d) so
x xx (A, (tr(A.)/d) I)x A, tr(A.).
=
The second assumption on the feature orthonormality can be attained via linear transformation on the features, preserving the overall function class. The third assumption on the operator norm bound ensures that the features are balanced, which is common in the non-linear feature learning literature (Nichani al., 2023; Wang al., 2023). Moreover, we note that when the entries of A, are sampled i..d., the assumption is satisfied with high probability by standard random matrix arguments.
Assumption (Well-conditioned link function). For the link function g², we assume g" is a degree- p polynomial with Ez (z)) (1), where (0, and p E N is constant. Moreover, R is well-conditioned, i.e., there exists
we assume the expected Hessian H EEz (z)
constant Cи such that Amin (H)
Cy.
This assumption ensures the link function adequately emphasizes all features, preventing degen- eracy to a power-dimensions subspace. The accondmooent condition is achievable through simple normalization.
Assumption 3 (Prepocessed target function). For the entire target function f*, we assume Po(f*) [x]] O and I( S k2/ Here Pk is the projection onto the function space of degree k spherical harmonics on the sphere (V-1(vd), and K2 satisfies K2 = poly(r; log d).
We will give a rigorous definition of Ps in Section 2.3.1. This assumption is analogous to pre- processing procedure conducted in Damian et al. (2022), which subtracts out the mean and linear component of the features from the target. The zero-mean condition ensures the network focuses on learning the function's 's variability rather than a constant offset. While Nichani et al. (2023); Wang al. (2023) assume the link function & has non-zero linear component, we rather assume g has nearly zero linear component, which prevents the target function from being dominated by a sin- gle linear combination of the quadratic features and keeps the learned representation space from collapsing to the one-dimenional space of that certain linear combination. This is an essential dif- ference between single-feature and multi-feature learning, because our assumptions ensure that the network genuinely learns to represent and distinguish all r features rather than conflate them, while assumptions in Nichani et al. (2023); Wang al. (2023) represent degenerate case that neural network may only learn the dominant linear combination of the features. We provide examples and counterexamples as follows.
Remark 1. These assumptions accommodate wide range of target functions. For instance, /F satisfies Assumption 3 with R22 S vrkı for any {as) )x(r) under Assumption 1. Moreover, for diagonal A₂ with A diagg(s), where a1, a2, ... a, are thogonal zero-sum vectors with entries ag E [=c/vd), we can achieve X2 0. Here c ((1) is normalizing constant. Notably, linear combinations of features like s(x) J Ek= (x Axx)
f*(x) = /F E (xTAxx)
violate our assumptions, since it represents a degenerate case with [P)()) = 11 O(1).
Three-layer neural network We adopt standard three-layer neural network for learning the target functions. Let m. m2 be the two hidden layer widths, and σi, 72 be two activation functions. Our learner a three-layer neural network parameterized by (a. W, b, V), where a E R W E m Xm22, b E Rm1,, and V E R22x The network fxx 0) is defined as
f(x:0)
1 Σ mi
1 my
a a...((ww,
(vx)) + b.) (,
1
0 (we, h (0 (x()) + b.). Σ
mi
Here, w, E Rm is the i-th row of W and (1)) (x) := 02( Vx) Ie is the random feature embedding lying in the innermost layer. We initialize each row of V to be drawn uniformly on the sphere of radius Vd, i.e., v ~ Unii((ed-((V)) initialization so that f f(x; 0(0)) 0 (Chizat et al., 2019). Explicitly, we assume that m, is an even number and for any j E [2,/ we initialize the paramters as
For a, b and W, we use a symmetric
(0) = - -0 -mmi
(0)
Unif({-11 -1 13), w!
(0)
w Wmi ((00 and 600) = (10) mj = 0.
4

---

Published as a conference paper at ICLR 2025
Here > 0 is hyperparameter to control the magnitude of the initial neurons. Different from Nichani et al. (2023) where the weights w, are initialized at zeros, we require a random initializa- tion, which enables the learned weights to capture the multiple features in all directions instead of converging to specific direction like the previous results for learning a single feature.
For the activation functions da and 02, we have the following assumptions: Assumption 4 (Activation Function). We take the outer activation function and the inner acti- vation function d2 as
2 2 - 1. (t)
It| and 02(() = - Q.i,t)
(2)
|t|
where Q:(t) is the i-th degree Gegenbauer polynomial in the d-dimensional space. Moreover, assume there exist constants Co. such that [a-(()) I Co for [t] S d, and Ex [x](x S d**Ck for k = 4. We assume 02 (()), and C2. C₁ and {eibez all constants independent of n, d, mi and m2-
We remark the outer activation 0i is a slightly modified version of the absolute value function [t], smoothed around the origin. The assumptions on d2 are based on the Gegenbauer expansion, often considered in the spherical analysis (introduced in Section 2.3.2). Compared to standard inner activations, we remove the constant term (Q)(() = 1) and the linear term (L) = 1/d) to focus on learning nonlinear features without low-order interference. Importantly, these assumptions on activation functions maintain significant generality. The assumptions are satisfied by many common activation functions with appropriate scaling. The core assumption in the Gegenbauer expansion is the non-zero component of Q2, i..., C2 ((), which rely on for a subspace recovery of the r quadratic features while other assumptions are made to simplify analysis since other components in inner activation will lead to useless noises or biases in the weights after training. Moreover, if we consider higher degree nonlinear features such as degree polynomials, we expect that d2 has sufficient emphasis Q, for efficient feature learning.
magnitude and moments
P-d is an example of the inner activation function. Remark 2. a2 Q2(t) a(d-13
Training Algorithm Following Nichani et al. (2023), our network is trained via layer-wise gra- dient descent with sample splitting. Throughout the training process, we freeze the innermost layer weights V. In the first stage, the second layer weights W are trained for step with specified learning rate mh and weight decay A₁. In the second stage, we reinitialize the bias b and train the outer layer weights for T - steps.
Transfer Learning We remark that our algorithm allows transfer learning of a different target function f that shares the same features of the original target:
(x) + f(x) = 9 (x AAx, Agx,
A,
(transferred target)
In this case. we switch the target function from f (p) to f = g(p) in the second training stage. For the loss function, we use the standard squared loss:
1 EXED, -(xx))) (original),
(0) = Σ (f(x;)) - (xx)) (() = no
I ExED, (f(x; f(x))² (transferred).
XED
n2
This layer-wise training approach, combined with the ability perform transfer learning, provides powerful framework for learning and adapting to hierarchical functions with hidden features (Kulka- mi and Karande, 2017: Damian et al., 2022; Nichani et al., 2023). The pseudocode for the entire training procedure is presented in Algorithm 1.
2.3 TECUNICAL BACKGROUND: ANALYSIS OVER THE SPIERE
We briefly introduce spherical harmonics and Gegenbauer polynomials, which forms the foundation of our analysis over the sphere gd-1 gd-1(VV). For more details, see Appendix B.5.
2.3.1 SPHERICAL HARMONICS
Let T be the uniform distribution on dd ((Vd). Consider functions in L ((dd-((Vd) (Vd), Ta-1). with scalar product and norm denoted as (,, (,)) l2 and I1 For f. E Zzo. let Va, be the linear space of
5

---

end Published as a conference paper at ICLR 2025
Algorithm Layer-wise training algorithm
Input: Learning rates 71. 72, weight decay Å1, A2, parameter e, number of steps '''
initialize a, b, W and V.
train W on dataset D
W()) W mn (W + W())
end re-initialize
6(0) -ni,((-3,3]), i [m]]
(11) v() (0)) v())
( - (a w(1), ba), v
end
train a on dataset D2
for
( + a (t-1) 1) 7/2 V (o(tt1)) + A2al1-1)]
10
end
13 return Prediction function f(:;(T)): x 1 (a (W +b())) mi
homogeneous harmonic polynomials of degree & restricted on (v-) The set (Vaehes0 forms an orthogonal basis of the L space, with dimension dim(Va.) (dd)) For each l E Zzo, the spherical harmonics ) SE(B(ay form an orthonormal basis of Va,e. Moreover, we denote by Pk
the orthogonal projections to Va which can be written
B(d.k)
Pa(f(x) Σ
(s, (xYY))))Y...xx))
We also define Pst = Ek PR. Pst = I - Pet = P and AadPo P2r = P
2.3.2 GEGENBAUER POLYNOMIALS
Corresponding to the degree l spherical harmonics in the d-dimension space, the é-th Gegenbauer polynomial Qe [-d, d] + R is a polynomial of degree L. The set {Qr)roo forms an orthogonal ba- sis on (''[dd, d). Fa-1), where Fd-1 is the distribution of vd(x,, e1) when Td-1- In particular, these polynomials are normalized so that Qe(d) 1. We present the explicit forms of Gegenbauer polynomials of degree more than 2:
L 2o(1) = Qi(1) 7. and Q2(1) = d(d -
2
Gegenbauer polynomials are directly related to spherical harmonics, leading to number of elegant properties. We provide further details on these properties in Appendix B.5.
MAIN RESULTS
The following is main theorem, which bounds the population absolute loss of Algorithm 1:
Theorem Suppose i (dda) and m2 (dd)). Let Ô the output of Algorithm after T = poly(n n2, my, m2, d) steps. Then, there exists of hyper-parameterss (c, mı T2, d1, X2) such that, with high probability over the initialization of parameters and draws of D1, Dz, we have
2p rP
d2
7-24
Ex
+
min min(mm; 12)
m2
d1/6
Complexity of g
Feature Learning Error
Moreover, for any other degree p polynomial g R - R with || < function fo = g" op by g in the second training stage, we can achieve the same result for learning the target function.
by substituting the target

---

Published as a conference paper at ICLR 2025
The full proof is provided in Appendix E.I. To interpret the results, provide the following discussion of Theorem 1.
Feature learning error This terms quantifies the requirements on the strst-stage sample complex-
ity and the inner width to sufficiently capture the non-linear features. Given d > т, if the width m2 = (tt and the sample size m4 = 2(d* + dfrpt)), we can fully capture the underlying feature information and approximate any degree p polynomials of the features. We will demon- strate how Algorithm learns these features through the learned representations Proposition and express hierarchical polynomials in Proposition 2.
Complexity of g" This term is the second-stage sample (and width) complexity given that the r features have been fully captured in the first stage. Moreover, for sufficiently preprocessed target function, i.e., K2 = (1)) we achieve the standard results of (rrP) complexity in learning degree- polynomial in the r-dimensional space in the kernel regime.
Transfer learning Leveraging the two-stage structure of training, we learn different target function in the second stage that shares the same features with the original target. This also sup- ports the fact that we have fully captured the information of the r nonlinear features in the first stage, making it possible for the efficient learning with a different polynomial head g. Moreover, by viewing the first stage as pre-training process with se(d' + ddr samples, only additional Or""xg") R22 2n)) samples are required learn any degree p polynomial of the features, which gets rid of the polynomial dependence on the ambient dimension of d.
Comparison with previous works Compared with the sample complexity of .id + dr?) in Damian al. (2022) for learning multi-index models, have similar polynomial dependence on and the dependence on d increases from dl to d because of the increased complexity of quadratic features rather than linear ones. Moreover, our approach significantly improves upon the ((dP)) sample complexity required by kernel methods to learn degree p polynomials of quadratic features (i.e., degree 2p polynomials of the input). Crucially, our polynomial dependence on d in the overall sample complexity is independent of the degree of the link function g.
Near optimality of the sample complexity We remark that our sample complexity of O(d*) is nearly optimal with respect to d for all algorithms that use one step of gradient descent for feature learning. Our assumptions on the target functions imply that the leap index! of our target functions are basically 4 (more specifically, the second order information of o p. where p are quadratic features), and we also utilize Pf( for recovering the subspace of the r quadratic features, which will be discussed in details in Section 4. Dandi et al. (2023b) indicates that 2(d²) samples are required for an efficient learning of terms in Pf(f*), which substantiates the near optimality of our result.
PROOF ROADMAP OF THEOREM
The proof of Theorem unfolds in two training stages. First, by a novel universality argument (Lemma 1), we show that after the first training stage, with sufficient training samples, the net- work learns to fully extract out the hidden features p (Proposition 1). Next, we show that during the second stage, the network is capable of expressing the link function with mild outer width my (Proposition 2). We conclude the proof through standard Rademacher complexity analysis to quantify the generalization error of the second-stage model (detailed in Appendix E.I).
4.1 STAGE 1: LEARNING THE FEATURES
We provide a brief analysis on the learned representations after the first training stage. Denote wj (0) (0) I(O,Img)) According to Algorithm 1, by setting e sufficiently small, after one-step gradient descent on W, we know for each j E [mi]]
(0)
V (0) L(B = 71 mi
1
Σ
/ (0))))))(()))) (low,, (lw,,,(((xx)) (x()))
w.
71
XED:
2em
(0)
Σ
0j
mi
TI
xED:
'The leap index of target function f is the first integer l that Pef * 0. Our assumptions imply diminishing Pca(f") and non-degenerate Pz(f") ) as + oo.
7

---

Published as a conference paper at ICLR 2025
By taking 1/1 201z 17 for 7) 0 be chosen later and A1 1/1 have
(0)
w mh (()(0) + / na
(x,
xED;
Then for any second-stage training sample x' E D2. the inner-layer representation becomes
(0)
= na,
( .0 02(')))))))
(xx,)
(x; (x:)
m2
71
xED
naj
(0)
W
(x,), (x))
h(7(x)
Our main contribution in this part is that the first-step trained presentations representations h (x) approximately spans the space of the target features (x A x A2x,.. x A,)). Thus, the target features p(x) can be reconstructed through linear transformation from the learned represen- tations h ( which is formalized in the following proposition.
Proposition 1 (Reconstruct the feature). Suppose m2, Mdd'). With high probability jointly on V, Dj and 22, there exists matrix B Rexma such that for any x E D2, we have
(x) - -(x) =
d
dr
r
m2
n2
d1/G
The proof is provided in Appendix C.3. We summarize the main idea of the proof as follows.
Universality of features The foundation of the proof lies in the universality result that the joint distribution the multiple features p is approximately multivariate standard Gaussian:
(x A,x, A2x x A,,) (O,,, d>
It is worth mentioning that we provide general universality theory that quantifies the difference between the distribution of any r-dimensional function (not limited in quadratic forms) and the r-dimensional Gaussian distribution, which is presented in Lemma 1.
Lemma 1 (Universality of vector-valued functions). Suppose X ~ (N La) is an d-dimensional standard Gaussian variable. If a function p Rd - R" satisfies Ex [p(X)] 0,- and Cor(p(X), p(X)) = I, then we have
E [IVIAAA
1/4
W((Law(p(X)),NWOIII))) &
V
X
V
Here =
[x)] and W, denotes the Wasserstein-1 distance.
The proof is provided in Appendix B.2. This lemma extends the previous universality results of univariate Gaussian approximation theory (Chatterjee, 2007) to the multivariate version and could be of independent interest for the field of high dimensional probability theory. As a corollary, when we take p to be quadratic features satisfying Assumption 1, we ensure the W, distance is bounded by ( /d) (see Lemma 16 in the appendix for the formal statement). This approximation error finally contributes to third term in the emror bound of Proposition (Equation (3)).
Utilizing the second-order information of g" Lemma 1 establishes crucial link between model and the multi-indes model studied by Damian et al. (2022). This connection allows us to sim- plify the analysis non-linear features and utilize the second-order information of the link function g" to fully recover the feature space. In the context of multi-index models where f*(() g" (p(x)) with pxx) Ux. has been shown that for a prepossessed target with non-degenerate expected Hessian H E V2 (z)]. the learned representations, dominated by the degree 2 component of
8

---

Published as a conference paper at ICLR 2025
f" which takes form Ex (x)x** U HU, are spanned by {u, @ ujttje(r)- Extending this to our setting with quadratic features and applying the universality argument from Lemma 1, we demonstrate that the degree component of our f*, namely Ex
(2. approximately spanned by the quantities (A, @ Aj AjKije(r). which is formalized in Proposition in Appendix C.I. Here (x) represents the tensorized quadratic spherical harmonics. Under Assumption 3, it turns that after the first step of GD (Stage 1 of Algorithm 1), the learned representations are dominated by this degree 4 component (Proposition 4 in Appendix C.2). This domination enables efficient re- covery of the "span" of the hidden features p. For a visual representation of our proof strategy, we also present our main idea of the proof in Figure 1. Remarkably, we find that the reconstruction matrix admits a surprisingly simple form B* x H H (p(v)( p(v2).... . (Vvma)) ). We provide empirical support for the effectiveness of this reconstruction through experiments in Section A.
Proof Idea of Proposition
(xx) 12)
a Pi(
fr Po(f*) + PP(F")
+
a
R (universality)
g" ())
bo + (b1.p)
E2
(2)2 02
p pp)
+
(Stage
(xx)
Low order terms
(f)) 62
(v @ Y2(x))
High order terms
(Controlled by Assupption 3)
N (universality)
Linearly reconstruct p(x') through
(En (z) p(v,) & p(x') = p(v;) Hp(x') , E [ma]
Figure 1: The proof idea of Proposition 1. Block characterizes the constant and linear terms of g", which is approximately equivalent to the low-order terms (²*()*) by our universality theory and results into biases in the learned weights (x') after Stage 1. This bias is vanishing with oo by our assumptions on Pf(f*) and P2(f*). Block 2 describes the second-order information of 9" (approximately Pa (f))), which is of the greatest importance and captured by the quadratic component c2Q2(-) in the inner activation (-)-) and converted into quantities spanned by the quadratic features Block 3 represents the remaining terms of f", which leads to high-order nuisance in the learned weights, but still dominated by the second term due to Assumption when d is large, which enables to utilize the terms in blue (resulted from Block 2) to reconstruct the features efficiently.
4.2 STAGE 2: LEARNING THE LINK FUNCTION
By the deduction above, after the first training stage, the model becomes random-feature model (Rahimi and Recht, 2007):
f(x':0)
mL Σ ajo nna)
1
(0) (w,. (x))) + 66)).
(4)
Here == (a, ( (1), V), with a [01, [2,,02.-.. ..,............................. 4ml E R being the trainable parameters
in the second stage. Leveraging the construction in Proposition 1, we can construct corresponding weight vector a in the outer layer to express the polynomial g QB (B) = g(p(x)).
Proposition 2 (Expressivity of the second-stage model). Suppose g is degree p polynomial with ligilza S 1. Then there exists learning rate l such that, with high probability over D,, D2, W and V, there exists a* E Ihm such that the parameter 0* = =(a", W(1) b(1), V) achieves small empirical loss:
1 Σ
<11a"11 mf
d
dpp+1
2244
(f(x;0*) -(p(x)))²
+ m2
n2 xEDz
n1
d1/3
Here a" satisfies ||a ||a11 Ō (mmmp
The proof is provided in Appendix D.I. We provide following discussions.

---

Published as a conference paper at ICLR 2025
Error propagation To explain the increased polynomial dependence on r, remark that the approximation eπror in Proposition gets multiplied by the averaged Lipschitz smoothness of the link function g. which is upper bounded by O(r"=') This product is then squared due to the use of squared loss in Proposition 2.
Reduced complexity of a Moreover, we remark that the complexity of a, i.e., ||a|| gets rid of the polynomial dependence on d, which is greatly reduced compared with a naive random- feature model that requires ||al|2 (mm, This directly saves the second-stage sample complexity and the outer width mi, since n1, m2 = (mm ||a" (12) is required for efficient approximation and generalization (Ghorbani et al., 2021). We also examine this reduced dependency by comparing our model with naive random feature model in learning hierarchical target functions in Section A.
Arbitrariness of g Thanks the two-stage architecture and the sufficient learning of the features, the choice on the link function can be an arbitrary degree p polynomial, not limited to the truth target 9". This allows us to conduct transfer learning tasks Stage of Algorithm 1. Finally, by standard Rademacher complexity analysis on the random feature model presented in Appendix 1, we conclude our proof.
5 CONCLUSIONS AND DISCUSSIONS
Comparison with Nichani et al. (2023); Wang et al. (2023) As discussed under assumptions and the initialization of our neural networks, work differs significantly in the targets of inter- ests, the parametrization of neural networks, the mathematical strategies, and the intuitions behind the results. Our assumptions ensure nearly zero linear component and a non-degenerate second order term of the link function g which significantly contrasts the assumptions posed in Nichani al. (2023); Wang et al. (2023) that emphasize the linear component. Our random initialization (rather than a deterministic initialization used in the aforementioned two works) in the weights of the three-layer neural networks allows the learned weights to capture multiple features in all direc- tions simultaneously after training rather than converge to a single direction. We develop novel universality result to relate multiple nonlinear features to multivariate Gaussian, while these two works adopt existing result of the approximate Stein's 's lemma which only applies to single nonlinear feature. Most importantly, subspace recovery is completely different from and also significantly harder than single feature recovery considered in Nichani et al. (2023); Wang et al. (2023).
Conclusions In this work, we have shown the provable capabilities of three-layer networks in ef- ficiently learning targets of multiple quadratic features. Leveraging a novel universality result, we have shown that one gradient step suffices for a full recovery of the subspace spanned by multiple quadratic features. In addition, leveraging the learned features, we have demonstrated the transfer learning capabilities this three-layer neural network with a constant polynomial sample complex- ity guarantee. To the best of our knowledge, this is the first theoretical result of efficiently learning such a board target function class of multiple nonlinear features with neural networks. We have made a great improvement the sample complexity by highlighting feature learning compared to kernel methods.
Future works First, it may be possible that the sample complexity bound of Ö(d*) could be improved to the information-theerrtic optimal sample complexity O(d) in learning general hier- archical polynomials of quadratic features. We think that this result may be achieved when consider more advanced algorithms that utilize the samples more thoroughly such as using multiple steps of GD, which could be a great future extension of our work. Moreover, our methodology is not inherently limited to quadratic features. The principles shown in Figure 1 and techniques devel- oped here give a foundation for understanding the learning of even more complex function classes. Another natural future direction of our work is to understand whether and when our results can be generalized learning multiple high-degree features.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors would like to thank Alex Damian and Yunwei Ren for the many insightful discussions. JDL acknowledges support of Open Philanthropy, NSF IIS 2107304, NSF CCF 2212262, NSF CA- REER Award 2 2144994, and NSF CCF 2019844.
REFERENCES
Emmanuel Abbe, Enric Boix Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: laap complexity and saddle-to-saddie dynamics. In The Thirty Sixth Annual Conference on Learn- ing Theory, pages 2552-2623. PMLR, 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Ad- vances in Neural Information Processing Systems, 32, 2019.
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprini vrXiv:2001.044133 2020.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019.
Jimmy Ba. Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Jimmy Ba. Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dimensional structure: spiked random matrix perspective. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36, pages 17420-17449. Curran Associates, Inc., neurips co wwcpppapr_fiiie sareeer_filesseeeeeee willssppperr/2222 Paper-------- Conference .dff.
2023. URL
Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks, 2020.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35:9768-9783 2022.
Alberto Bietti, Joan Bruna, and Loucas Pillaud- Vivien. learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310. 197933 2023.
Souray Chatterjee. Fluctuations of eigenvalues and second order poincaré inequalities, 2007. URL
https: htttps://rxxiv.or///e alss 0705 1224.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M. Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning two-layer networks after one gradient-step. 2024. URL aittps://arxiv. /aaxiviiiiv.. org/ abs/2402. 04980.
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.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two- layer neural networks, one (giant) step at a time. arXiv preprint arXiv:2305.182270 2023a.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time, 2023b. URL https: /arxivv org/ abs/ 2305. 18270.
11

---

Published as a conference paper at ICLR 2025
Yatin Dandi, Luca Pesce, Hugo Cui, Florent Krzakala, Yue M. Lu, and Bruno Loureiro. A ran- dom matrix theory perspective on the spectrum of learned features and asymptotic generalization capabilities, 2024a. URL https:://////. httpppp://axii....oo
wrxiv.org//.........ooo...o..o.. abs/ //t.//11..1.. 18938.
Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krza- kala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents, 2024b. URL 03220.
Amit Daniely. Depth separation for neural networks, 2017. URL Utttts:://aaxxi 1702. 08489.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. 2018.
arXiv preprint
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterizen neural networks. arXiv preprint arXiv:1810.02054, 2018.
Weinan E, Chao Ma, and Lei Wu. A comparative analysis of optimization and generalization prop- erties of two-layer neural network and random feature models under gradient descent dynamics. Science China Muthematicc, 63(7); 1235-1222255258 jan 2020. doi: 10.1007/111425-01-11238 URL
https: //tii//ooiooooo . .rgg www11100022114410 01 9- 1.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029 1054, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Tom H. Koornwinder. Dual Addition Formulas Associated with Dual Product Formulas, page 373-392. WORLD SCIENTIFIC, January 2018. ISBN 9789813228887. 9789813228887 ,0199,
doi: 10.1142/
URL http: //tp::/dxx... doi. www/o0.gg.....oo
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet classification with deep convo- lutional neural networks. In Advances in Neural Information Processing Systems, 2012.
Mandar Kulkarni and Shirish Karande. wayer- training of deep networks using kernel similarity, 2017. URL https: : / arxiv. org attaabbb////..////// 07115.
Honglak Lee, Chaitanya Ekanadham, and Andrew Ng. Sparse deep belief net model for visual area v2. volume Vol 20, 01 2007.
Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. Advances in Neural Information Processing Systems, 33:15156-15172, 2020.
Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learss low-dimensional polynomials with sgd near the information thhooreticcccc limit, 2024. URL 2406 01581.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pages 3351-3418. PMLR, 2021.
Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non-linear feature learning with one gradient step in two-layer neural networks, 2024. URL https: /ttpps://araae //axxiv. org/ htg//bb/// 2310.30022 07891.
Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A. Erdogdu. Gradient-basse feature learning under structured data, 2023. URL https:://axttt org/ abs 2309. 03843.
12

---

Published as a conference paper at ICLR 2025
Alireza Mousavi-Hosseing Denny Wu, and Murat A. Erdogdu. Learning multi-index models with neural networks via mean- saun-field langevin dynamics, 2024. URL httht:://// /ttptpps://riii...// aataaaiivvioooo abs/
Eshaan Nichani, Alex Damian, and Jason Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. arXiv preprint arXiv:2305.06600 06986, 2023.
Giuseppe Da Prato and Luciano Tubaro. Wick powers in stochastic pdes: an introduction. 2007. URL https: api. wwwant wrg/Corpus ID: 55:595221115 55493217.
Adiyyanarayanan Radhakrishnan, Daniel Beaglchole, Parthe Pandit, and Mikhail Belkin. Feature learning in neural networks and kernel machines that recursively learn features. arXiv preprint arXiv:2212. 13881, 2022.
Ali Rahimi and Benjamin Recht. Random features for large-scale kemel machines. In Platt, D. Koller, Y. Singer, and S. Roweis, editors. Advances in Neural In- formation Processing Systems, volume 20. Curran Associates, Inc., 2007. URL https: http///peeoennnneeeeeeeeeeeeeeee neurips.... htcc/paper_fiiii wwwee/illss/peeeeeee httppper_iiessppppe//00///////// 155000ff0oabcc5552fffffeee11
pdf..
Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. arXiv preprint FrXiv:2304.0106. 2023.
Nathan Ross. Fundamentals of stein's method. 2011.
Itay Safran and Jason Lee. Optimization-based separations for neural networks. In Conference on Learning Theory, pages 3-64. PMLR. 2022.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions Information Theory, 65(2):74 2018.
Ramon van Handel. Probability in high dimensions. 2016. URL https::tt////tt. //wb:::///bbb.. math.
princeton. ceduu-rvan/APCSS... pdf.
Zhichao Wang, Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propa- gation in deep neural networks, 2024. URL https: www.i.owwwwi..oaooooooooooooeooo abs /2402 10127.
Zihao Wang, Eshaan Nichani, and Jason D. Lee. Learning hierarchical polynomials with three-layer neural networks, 2023.
Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel, 2020.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pages 11727-11733 PMLR, 2021.
13

---

Published as a conference paper at ICLR 2025
Appendix
A Numerical Experiments
15
Techinical Background
17
B.1 Asymptotic Notation
17
B.2 Multivariate Gaussian Approximation
17
B.3
Hypercontractivity of Polynomials
18
B.4 Moments and Factorization of Polynomials
19
B.5 Spherical Harmonics and Gegenbauer Polynomials
20
Approximation Theory of the Inner Layer
C.I Asymptotic Analysis of the Learned Feature
21
C.2 Boundedness of the learned feature
30
C,3 Proof of Proposition 1.
36
C,4 Proof of Other Supporting Lemmas
D Approximation Theory of the Outer Layer
43
D.I Proof of Proposition 2
43
D.2 Omitted Proofs in Appendix D.I
45
Generalization Theory
50
E.1 Formal Proof of Theorem
50
E,2 Omitted Proofs Appendix E.1
54
14

---

Published as a conference paper at ICLR 2025
A NUMERICAL EXPERIMENTS
We empirically verify Theorem and Proposition 1. We consider learning functions with r = 3 quadratic features. Regarding the target function, choose the target functions be of the form
(x) = fd,p(x) [[fap(x)]
with fd,p(x) Σ (x AIx) PEN.
(5)
/Vrr(fapxx)
For the underlying features, take p(x) [x' A Aax] with A₁ diag (c. an),
and c > 0 is a normalizing constant. To ensure the orthogonality of the features and tr(As) = 0,
choose the ambient dimension d to be divisible by 4 and take aş to be
aj = Vec (11, 11,1,,1,110 11),,
Vec ([1, 11-11,1
-1,, 1, - 1]), and ag = Vec ([1. -, -1,111
Here a vector of ones in d/4 dimensions, and c = 2
to ensure that Exx [(x Asx)²
for each k- 3.
For the network architecture, we choose 01 as per (2) and 02 = Q2. with network sizes set to mi 10000 and m2 20000. We compare our proposed model (4) (given by Algorithm I) against the naive random-feature model defined as
RAP (x':0)
m1 Σ ajoi (na)) (w,, h (0) (x))) + 1)
where a is the only trainable parameter throughout the training process. Our experiments involve learning fåp with p = 4 and d (8, 16, 32). To examine our model's transfer learning capabilities, we also train the model on an initial target function fū,2 with d = 16 and пı = 2 the first stage, then transfer to targets fä with p = 4, 6, For each task, we explore a range of sample sizes from 2 to 26 The results of these experiments are presented in Figure 2.
Improved sample complexity and Polynomial dependence on d The left panel of Figure demonstrates that model outperforms the naive random-feature model across all dimensions. As the dimension d increases, both models show larger test errors, but our model exhibits less sen- sitivity d. This aligns with our theoretical analysis in Theorem I that the sample complexity of kernel methods should be Mdd#==) times greater than that of our model. Moreover, we redraw Figure 2 by plotting the test error against log, n. As shown in Figure 3, the loss curves for our model (Algorithm 1) align closely for different values of d, indicating that it achieves low error rates with only (dd)) samples. In stark contrast, the naive random feature model exhibits significant separation between curves for different values, requiring more than Õ(d*) samples to achieve comparable error rates. This graphical evidence powerfully demonstrates how our approach elimi- nates the dependence on dimension (ddP) presented kernel methods, resulting in substantially improved sample complexity in high-dimensional settings.
Efficient transfer learning The right panel of Figure 2 showcases our algorithm's 's strong transfer learning capabilities. our algorithm successfully learns all three transferred target functions with benign second-stage sample complexity. Notably, as the degree p increases, the test error grows no faster than r", which significantly slower than d²r This supports our theoretical result that the second-stage sample complexity depends on the number of features r rather than the ambient dimension d, underscoring our model's strong transfer learning capabilities.
Accurate reconstruction of quadratic features further demonstrate our model's 's feature learn- ing capabilities, we extract the leamed features (1) after the first training stage of Algorithm 1, us- ing fie.2 as the target. We then reconstruct these features using a linear transformation B* E RRr as described in Proposition 1. We examine how reconstruction accuracy changes with first-stage sample sizes. Figure shows the correlation between true and reconstructed features for each sam- ple size. As nı increases, all features are better approximated simultaneously. Notably, dt samples prove sufficient to reconstruct the features with high accuracy, supporting our model's 's effective fea- ture learning ability.
15

---

Published as a conference paper at ICLR 2025
Alsveihun
Randan Rotzoo Madel
Trander Tearsing
13
0.1
0.1
0.3
(a) Comaprison between Algorithm and random-feature model (b) Performance transfer learning
Figure 2: For the left panel, Algorithm uses two equally sized datasets, while the random feature model uses the full dataset. For the right panel, we conduct transfer learning with n2 2 pretrain- ing samples and plot the dependence on n2g The figure reports the mean and normalized standard error of the test error using 10,000 fresh samples, based independent experimental instances.
Alporithm
Random Frataure Model
12
1.2
00
0.5
0.2
0.2
1ogg
logg 12
Figure 3: Test error of Algorithm and the naive random feature models with x-axis being the relative sample complexity (loga n). We plot the test error of independent instances for each E {8, 16, 32).
Correlation belween the true and the reconstemcled leatures
Figure The linear correlation between the three true features and their corresponding recon- structed features for varying first-stage sample sizes The reconstructed features are standardized to match the variance of the true features. For 1,2, the i-th seatter plot represents 10, 000 test sample points of (B (1) (x)),, x Ax)) for nı E {d? d3 d where 16.
16

---

Published as a conference paper at ICLR 2025
B TECHINICAL BACKGROUND
B.1 ASYMPTOTIC NOTATION
Throughout the proof we will let C be fixed but sufficiently large constant.
Definition (high probability events). Let < = (log((dn)2mmmmm)) We say that an event happens with high probability it happens with probability at least - poly (d, n1, n22 my, m2)e
Example 1. If ~ (0, then |x| S 22 with high probability.
Note that high probability events are closed under union bounds over sets of size poly(d, n1, f2, M1, M2), such as D1, D2 and (wylsemmi)- We will also assume throughout the paper that S C-Id. B.2
MULTIVARIATE GAUSSIAN APPROXIMATION
In this section, we assume that X (O, Ia) and aim to establish an upper bound of Wasserstein distance between the distribution of p(X) and the standard r-dimensional Gaussian distribution, ie., Lemma I.
To prove Lemma 1, we introduce Stein's 'ss method (Ross, 201 for multivariate Gaussian approxima- tion. We will use the following additional notations.
Gf(x) == So² Ez)N((.)) [5(e-'x VI 1i-ee-2 2 [(Z)] d denotes the potential op- erator of f. (p) = VPI, Vp2..., ... Vp.] T E Rexa denotes the Jacobian matrix of
Now we state the supporting lemmas to prove Lemma 1.
Lemma (Corollary 9.12 in van Handel (2016)). For any probability measure μ in IR,, we have
Wi (i S
By 1
sup
Y)
Lemma (Lemma 9.21 in van Handel (2016)). Suppose X (X1, Xa) ~ NN. I.) is an d-dimensional standard Gaussian variable. Then for any functions g Rd R and h E R. we have
Cov(g(X), h(X)) = Ex (VV(()), Vgh(X))]
With the lemmas above, we begin our proof of Lemma 1.
Proof of Lemma 1. By invoking Lemma 2 with f = Law(p) and Y = p(),, for any g : R
R with /gll < 1 and || S V / we aim to bound
0
0g
Ex (g(p(X)) (Vg(p(X)), p(X)
Ex
Px)
Oyh xxp(X)
Oy: y=p(X)
Since for any i [r], E [X)]
= O, we have
0g
0g
E
pi
X
= Cov
Pi(X), Oy:
Dy (x=)
y=p(X))
dg
V )
E
V x Ty: y=p(X)
82g g
(x);(X), Vx
dy.ūyy p(x())
82g
(VD, Vx9P:(X())
Oy:oy; (x=p(X)
17

---

Published as a conference paper at ICLR 2025
where the second equality follows from Lemma 3 and we obtain the third equality by the chain rule. Thus, we have
E 3
For a special case, for any i,j E [r]. we take g(y) - yiyl in (7), obtaining that
E [2pp( X pp: (X)) 0, {
BKV X Vx9p:(X)M)
E [2pt [X] - 11 = 1-1iz i=
Thus, EE [I, (G Opxr Since for any i,j [r]. We can therefore estimate
| S L we have
|| - V
2
Wi W((LN(D(()))N(OO.L.))
Σ
Elj8, (X), VxID:
π i.j(
2
VVp; Vx92(
( (X))1 1/2
Σ
(E
where we invoke Poincaré inequality in the last inequality. For any i, E [r]] we have
E
V x (V xp, (X), VxGpx()))]] (xx,))))))
E
Vxp, (X
+
S F
V p
2
Vx
(X)
||'' p3 19]
1/2
1/2
1/2
| [1]]
1/2
S 2
E
||V
+2E
V
p
[1V'p, p. 11]
1/2
1/2
1/2
[| p: [1]
1/2
2E
E
|
2
E
The last inequality follows from the inequality in Page 308 in van Handel (2016). By adding up all the terms along and j. we have
WL(aww(PC (p I,)),,
Σ
2
| p,
[1*]
1/2 IE
11
1/2
+ 2E [|Vp)11"' 1/2
E [[ve's11") 1/2
[jE(r)
Σ ()E((
E
11
1/4
[IVP:11]
1/4
+E
[IVR>1"] // E | )
[iVaall]]] 1/4
E
We complete our proof.
B.3 HYPERCONTRACTIVVTY OF POLYNOMIALS
The following Lemma is cited from Mei et al. (2021) and is designed for uniform distribution on the sphere in d dimension.
Lemma 4, For any E N and f (sd-)) to be a degree polynomial, for any q > 2, we have
( (Rzunitiis I(V)) [ 2/9 < (g 1) B Uuu =(Va)) [f(2]]]
We remark that the results above also multiplicative.
18

---

Published as a conference paper at ICLR 2025
Lemma 5. For any E N and f E ((8 to be degree & polynomial in the components of each 21. Z...... Zy. for any z 2, we have
2/9
E 2 (st 1(va))* ((x))1)
(q 1) 3 E Uniu (s-))))))) [
Here Vec ([Z,, Zo....., zt)).
For the case where the input distribution is standard Gaussian in d dimension (denoted as ?), we have the next Lemma from Theorem 4.3, Prato and Tubaro (2007).
Lemma 6. For any E N and (2) be a degree & polynomial, for any q 2 2, we have
Eg [f(z]] s (2,e(1) (Eeny [f(z)
where we use (1)) denote some universal constant that only depends on q, l.
Moreover, introduce lemmas to control the deviation of random variables which polynomially depend on some Gaussian random variables. We will use a slightly modified version of Lemma 30 from Damian al. (2022).
Lemma 7. Let 9 be a polynomial of degree and positive constant Cp depending only on p such that for any 6 1,
N(O.Ia). Then there exists absolute
PP P((x(x) - [xg( >> 6 ovVar(((x))
s 2exp (-Cr02/0)
We also have the spherical version of Lemma 7.
Lemma 8. Let 9 be a polynomial of degree p and x ~ Sd- (Vd). Then there exists an absolute positive constant Cp depending only on p such that for any 8 1,
[lo(x)) B[g(x)]] > o. VVr(g(x))
S 2eep (-C,62/0))
IP
Thus, for a degree-p polynomial g, we have g(x) S 2/2 I with high probability. B.4 MOMENTS AND FACTORIZATION OF POLYNOMIALS
In this section, we present formulae for calculating moments of Gaussian or spherical variables, cited from Damian et al. (2022).
Lemma (Expectations of Gaussian tensors). For w NN,,,, Ia) and k E N, we have
Ew 22k (2k 1)!!Sym(I
Here Sym(T) is the symmetrization of k-tensor T E (IRd) across all k axes.
Leveraging this calculation, we can factorize any polynomial g into inner products between high- order tensors and bound the Frobenius norm of the tensors.
Lemma 10. (Lemma 2/ in Damian et al. (2022)) Given Let g : R - R be degree-p polynomial. Then there exists To. T TI...., Tp such that
((2) = Σ 200)
with | I 1 77, A 0, 1.1....P
k=
Here || = E.-N [z](
As corollary, we then have Vg(z) Eb 1T x2 and
P
[Vg(z)] < h | ||2||
S ||
Σ
kr
|| k-
(8)
k=1
For spherical variable Uni also compute its moments.
19

---

Published as a conference paper at ICLR 2025
Lemma 11 (Expectations of Spherical tensors). For Unif(S
Ez [z 12 d Bw-NNoa.Ld) (wStk] (vd)) and N, we have
(v- [b2k]
where X(d) represents the chi-distribution with the degree of freedom being d, and its moments can be computed as
(e [02K]] II (d 2j) = O(d*).
As an example, the moments of spherical quadratic forms x T Ax can be computed explicitly as
d
Ex [ Ax] tr(A), and Ex [(x Ax) (1]xxxxxx] Bx)] d 2 (tr(A)tr(B) 2(A, B)).
Thus, satisfy Assumption require tr(Aa) 0, ||AlIE v (d and (Ak, Ar) O for any k E E [r]-
B.5 SPHERICAL HARMONICS AND GEGENBAUER POLYNOMIALS
We introduce some facts of spherical harmonics and Gegenbauer polynomials, with the first four properties from Ghorbani et al. (2021) and the last one from Koornwinder (2018).
1. F, Sd (Vd).
x))( =
(9)
2. For sd-1
(Q,((x,)) )). Qu B(d, k:) 6g Qu ((x, y)).
(10)
Here B(d, k) denotes the dimension of subspace of degree k spherical harmonics
2k + d k + d 3 )) O(d*).
B(d, k) dim(Vax)
k
1: --1
3. For e
B(dk)
Qx((x,y)) = k(d,k)
Σ
(x)
(11)
1=1
4. For any No1.
t
k d
Qx(()= 2k Qt) +
2k + d - Q 1t+
(t).
(12)
d
5. For any i,j E N.
min(i.1)
i
Q100((1))
Σ
Ditj
k
A1 (1Qu4-22(())
(13)
Here, have
2(i + -(itj220 2k) + d - (dd-2)/2))
((d-))2), ((d-2)/2),((-=2)))))
((d 2)/2); x ((d 2)
(t,)
itj-2k
2
(d,2), (d) (d/2)
We note that (=)x = 2(2 +1 ... (2 k - 1) = [(= (2) is the Pochhammer symbol. Given
any and have dt d²b((.)) (i.j)
-k + when d + co. We derive quantitative bound on the scale of i4j-2k in Lemma 12.
(k.f)
20

---

Published as a conference paper at ICLR 2025
Lemma 12. For any i.j 3 k 0. denote
((d-2)/2) ((d (1d22//)((dd2222)2__-1 x)
-2/2)) 2) 2)
(..) 443-2k
(d (d-2), 2), (d 2), (1/2)
Then, when d 3 4, it holds that (3-1) 4+j-2k I [d-
Proof of Lemma 12. Note that when d 3 4, i.e., d 2 > d/2,
(6+1.))
(++1)+--25
2 i
k)
(monotone decreasing with j)
(4.1)
(1-2+i)
-4 2k
2-2 +i () 2
(d - 2 i) +i)
< 1.
Thus, we have (6+1J)
(1.i) 22k- Similarly, we have (6++1) (+0+1)-26 S C---2k Consequently,
(id)
S
((41)+j-2k
for any i,j > k, we have
(1,))
(k.k)
445-2k S
((d ( (d/22
2)
((d-2)//2)
(d (dd(d/2)
(d - 2)/2
(d 2) + K 1)
< (d 2)
The proof is complete.
C APPROXIMATION THEORY OF THE INNER LAYER
Since we focus on the first training stage throughout this section, we denote n = n, for notation simplicity when the context is clear, and let the training set be D {x1, Xx.,,
C.I ASYMPTOTIC ANALYSIS OF THE LEARNED FEATURE
In this subsection, analyse the learned feature h1)) (x') in the asymptotic way, i.e., M2, Note that we can rewrite the learned feature as
h (x))
f (xx) (x:), (x)))
11222
1
(xxx) (x, (")
where the initial kernel k:10 xx x') is defined
my (x,x') x')
(a2 (V2((V)) x Ey [x xv [x))]-
m2
In this case, have for any [m22).
(x)), = 1 K
(0) (x, x') (2z x'v x:)
m2. Ex (xxx)
Here the infinite-innerwwitth kernel (0)) is defined
K [(0) (x)))xxx) x') = Ey [az(v (Vxx)VTx()
0x(xxx))
B(d,i)
21

---

Published as a conference paper at ICLR 2025
Recall that (2(t) A-d
- d(d-1)- Q2(v)xx= - -I, - II/(d(d -00-1] 1)). Let's focus on the
contribution of the quadratic term Q2 in K (x,'', and (v x), which is
Rx (1)=)))
B(d,2)
2)d Ex (x)x I.v,V] (xx 1,,'' -1)]
16 (Ex (xx((xx - ). (V,v, - 1 - -1)).
The following proposition provides an approximation of the tensor Ex (x)(xxx -I 1)) lays the foundation of our feature reconstruction theory.
-1102) which
Proposition 3. Consider two linear operators T and T that map 1aaxd to Indxd and satisfy
T(W) = Ex
(x) (W, xx T - -I]] and
(14)
T (' Σ 1
(W, AA) HR
(15)
A
k=1
for any W Rdxd where H is the expected Hessian matrix H E.-A(O, [V]g(z)]]. Then for any IRdxd we have
[[((W) (TW))T((WW)) <d-1/6 log" d. [WiF
Here L= Ōrr is the Lipschitz constant of g" that holds with high probability.
The proof is provided in Appendix C.1.1. This proposition shows that, when H is well-conditioned and d >> r, T can fully recover the space spanned by A1. Ag...., Ar. which enables us to recon- struct the features efficiently. Specifically, when taking W Σド= "JR,JA, Aj for any & E [r].
we have T(W R T* (Wx) Ak.
Now we consider the construction of B If we set B my (po(vi),..... Po(vmz) ) for some
vector- valued function Po Rdd + R" and denote Po(v) .v.......PP...... we directly have for any ls E [r].
B (x))
1
Bx (x, x)po,
m2
E, [Ex (x) 11822 po. @ (x'x'T -]))]-
d6
d6 (T" [Ev T - -D))), TT 1).
Thus, it suffices to solve which is equivalent solving
T TE (vo. -
-D]] x Ar. h = =12,.....
Ey (v -D] -I] We H Hk
Since have Ey [(v Avv)(vvT- I)] x Ag. we can explicitly construct po,
(v) x V Ayv. ie., po(v) x '-v).
1=1
Thus, with well conditioned H. we can fully reconstruct the features.
22

---

Published as a conference paper at ICLR 2025
C.1.1 PROOF OF PROPOSITION 3
To prove Proposition suffices to prove that the upproximation error
R RWW V) = (T(W) - T (W), V) S d_ log" holds for any test matrix V with || = We rely on the following three lemmas. Lemma 13 (Bound R(A;, A,)). For any i,j E [rr] have
Lr log² d vd
Ex(g"(x Ai .... Axx) (A,,, - -I
(Vg(z)
5
13
Here I, = Ro is the Lipschitz constant of g" that holds with high probability.
Following the proof above, we have the following more general lemma.
Lemma 14 (Bound R( R(A;, B)). For any matrix B Rddd satisfying E |x Bx
0.
x Bx 2 and (B, A:) = 0 ffr any = 2, ........... V, we huve
E [xx AIx,.. ..... x A,, IJ []]] 5 dd Lr2 log" d.
Lemma 15 (Bound R(B1, B2)). For any two matrices B,, B2 E Rddd satisfying [x B,x] = 0,
E (x B,x) 1 and (B,, = for any j = 1,2 and i = 1, .......77 we have
E [gg Aix..... (BI, xx - (), xx T - 11] 2 d Lr Nn log d.
Here L L a Lipschitz constant satisfying (gx))(x))l2 S L with high probability
The proof of the three lemmas is provided in Appendix C.1.2. With the lemmas above, we begin our proof of Proposition
Proof of Proposition 3. Given any W E pdxd we assume | |F 1 without loss of generality. Let's decompose W as
Ar Art1
+ Art2B, where (B, As) = (B, Ia) = 0, E [r].
Vd
Here the coefficients 1rt2 satisfy Er+ X S so KE Au S Vr. Since TBB) =((I)- T (I) Oaxd. we have
((W) T (W)lI
(T - (Au)
Vdl (1(a -
Ia
Ia Vdl
+ + [T( - (B)lll
Aul (TlI)((AN) - T (Ax) IF + - | [
Since both T(Aa) and (AA) are traceless, by Lemma 13 and 14, we have for any = 1,2, ....... ,
(T(A)) - (Ax))l = max
[+ (T(A.)
- (As),, V)
S Vr d 1/r log² d d 1/ k1 log d
- 1-4 1/4 Lr x1 log2
This is because we can decompose V = -E A+ + B with (B', Ag) = 0 and apply the two Jemmas to obtain the results above. Similarly, by Lemma 14 and 15, we have
[T
(T(B).
S VF d 1/4 Lr Ki log² d + d 1/6 2 Kj log" d
1/4
I d 1/ log d.
23

---

Published as a conference paper at ICLR 2025
Thus, have
( - S
Aue d 1/4 1/4Tn/2 log² d + 12 1/6 2 log² d
k=1
- d 1/6 Lr 2 log d.
Here invoke ExE Dall S VF= (1) in the last inequality, The proof complete.
C.1.2 OMITTED PROOFS APPENDIX C.I.I
The following lemmas lay the foundation for our approximation process.
Lemma 16. Suppose Assumption holds. Then the Wasserstein-1 distance between the distribution of (x , x3 x) and standard Gaussian (0,, I,) can be bounded by
y²K1
Wi((x AIx, x N I,)) S
(16)
Vd
Moreover, for any orthogonal unit vectors Un u2.... u, E Rd, we have similar bound of
(r + s)²
W1 A1 X Ax,, x, u (), S
(17)
vd
Proof of Lemma 16. For fixed matrix A,, define the function fi(z) dx A, and let z |
Observe that when 2 ~ (0, 1), we have x ~ Unif(Set (Vd)). Therefore [f((z)kett) is equal in distribution to [x Ax]ee()) We have for any E [r].
A,z
A,z
Vf((z) = 2d
||z||
||
and
A;
2A,2z ,z
2zz A,
z T
A,Z
z T
A, A,z
(2 2d
||z|
||2]|
||
|z] |
|z]]
Thus, we have
||A
1z T A
Vd
|x A,x |||||
ff(22) $ 2d 2d
+
||
|
|| $
||z||
and
|+>i))) S
Aillop
||z]|
Since ||z||² is distributed as a chi-squared random variable with d degrees of freedom, and thus
24
IT - 2j)
Therefore, we have
V
fz)))
4
1/4
S
A
|| ||2|1-8] 1/4
I ||
Then, using the fact that x and ||z|| are independent,
E
$ VdE
|z
1/4
E
[|AAX|1
+E
||2||
4] 1/4
E (xx 11 21.
24

---

Published as a conference paper at ICLR 2025
Thus by Lemma I we have
Wi((x' A Ax,,,.,,,AA,,, A(),, Wi(fi(z), ... (.)).
x [13ン)())111 1/0
Σ
[1
K1
vd
Now let's focus on the function gj gj(z) = vd u
holds that
(gz(z)= = vd
u
u
2
|
|Z|
and
(2);(2) = Vd
u
+ zu;
3 u
z zz
| 1
||z/|
Thus, we have
2vd
5vd
[2)
v
and
[z)
-
||||
|z 12
which directly gives rise to
E (2))
1/4
S1 and
|
(z)
2/4
x vd
op
Again by Lemma 1, we have
(1 AIX,
x
A u2 x.. x), Inte)
(z)..
(z) ,z)))))))...........).)).))))))) NV Ir+s))
E
V.f((z)]
Σ
||
(z)
1/4
E | [
(r 2
Vd
The proof is complete.
With the lemma above, we begin our proof of Lemma 13.
Proof of Lemma 13. For E Rd, define H (z) = g(z)zzz. Then by Stein's Lemma, we have
1. [z)] = Ex~ [V2 gz))),)),, + di, (Z-N(O....) [g(z]]. Moreover, let R > 0 be truncation radius and we define H (z) H(clip(z, It).. Here the clipping function is defined as
clip(z, Rt; max(min(2;, Rt), -R), i = 1=11,2 ......
By (8), we know g(clip(z, R)) is ORt " " )-Lipschitz continuous, so H has a Lipschitz constant of O DRR8 = Thus, by Lemma 16, have
Rr 2 2
Ex [g(clip(x] Aix, x x), R)) E EZ~N1O.,1 (g(clip(z, R))) S
vd
13r 2-
E T A1x,..., X A,x)] EZ-N(O... [T( x
Vd
25

---

Published as a conference paper at ICLR 2025
Since [x[(x' ] 1 for any &s [r]. by Lemma 8, choosing R C' log for appropriate
constant C' can ensure that
1x Aix x A,) - 11 (x Aix, x AA,x)] S d'
1
Ex [g(clip(x' T Aix,.... : x), R) - g(x Aix..... x [A,x)] $ d'
1
E2 V(=...) H())) [x))] 1 d'
1
(L (o(clip(z, Rt) - [(()] - d
Altogether, have
Lrnk2 log d
E [x AIX. X XA,, 1) - -o
( g()x))
vd
Here L Cr for some constant Cg > 0 is the Lipschitz constant of g that holds with high probability. The proof is complete.
Following the above proof and replacing A; and A, by any other traceless matrices BN and Bg that are orthogonal to all As. we directly have the following corollary:
Corollary 1. For any two matrices B1, B2 E Rdxd satisfying E [x B, = 0 and (B,, A.) =
for any i 1,2,..... ... and 1,2, have
E g (x A,x, x (A,, xx - (I) xx
[]]]
[By op [
S
TK1
|ByBl Lr log d,
Vd
for any j = 1,2, 2, and
E [g" (x
A1x, ..... TT (B,, () xx 1>
B1
|B2llop
TK
op [B
|BilF [il:I|Baa |B2llF Lr log² d.
|B
Also, by (17) we know for any unit vector u E Rd, (x' A1x, x Agx.... .... ,x A, xx) is ap- proximately Gaussian when d is sufficiently large, which gives rise to the following lemma by the same deduction.
Corollary 2. For any E [r] and unit vector u1 Ug Rd and matrix satisfying the same requirements in Lemma 14, we have
2 log² d vd
E g (x (xT AIX, T x) xx -I) ((u3x)) -I)
1)}1 x
B lo |
E [9 (x A1x X
(I)((xx))-111] 1)]] /
+
| Lr log d
Vd
for any j = 1, 2, and we further have
2 log d Vd
[x A1x
(u] 1)((ux))) x)² -1)]
With the lemmas above, derive stronger version of Corollary i.e., Lemma 14, in which the error gets rid of the dependence on ||Bllpp'
Proof of Lemma 14. Let > 1/ be threshold to be determined later. Decompose as follows:
B
u. Σ M (uzuy 1))
Σ
Ague Asu; A B.
A
/=1
26

---

Published as a conference paper at ICLR 2025
where {u, 1/=1
are orthogonal unit vectors and
B E
+I d Σ
E Nuu Ag-
Axlle
10g127
k=1
1aylar
By construction, we have
tr( =
Σ
d Σ
Dyy/T
jEld
Moreover, for any 1 E [r], we have
Ab) E Yuu, + Σ
Yuu (Ax, B) = 0.
19/25
Therefore by Lemma 1, we have
g (x
A ...., x (A,,, T (B, 11
TK1 Vd
op B
B
Lr logo d.
(18)
F
Since Σ 2
|
(d + (2d( = O(1, there are at most O(r indices j satisfying
My > T, which gives rise to
E kyd
INS 7-2. Σ
ST-
Dayle
Thus, we can bound the Frobenius norm of B by
B
2
Σ
+
Σ
11
+
Σ
E
A
A L
Ay2T
1=1
2
=
B F
Σ
23 +
Σ
+
Σ
2nu7 Anus)
Az2T
2
< Σ
23 +
(
Σ
A
Σ Aj
k=1
rR?
S 1
+
d7
Thus, we have
B
S
1+ +
Frn ddt
and
p
)
T+
Σ
A
Σ
Aoue Aug
k=1
FK1
dr
Thus, plugging the norm bounds into (18), we obtain that
/3/2
g" [x] A A,
Bx 11) 2
TK1
+
+ r
PR1 dr
Lr Lrlog* d.
vdl
dr
Next, applying Corollary 2 with u = U1 ....... us we have
Lr ki Lr² log² Vd
E (
Aix...., .,,, ..X,...,TAAAxX Xx - I) (u x/d, -1)]]
Vj E (4]
27

---

Published as a conference paper at ICLR 2025
Thus, have
( /Es, uujuT - 211) xx -1)
xxx
E
g g'
..... X ATA .Ix,..,,
-1)
Lr R1 log2d
Σ Ma
Vd
Lr² log d vdr
Besides, by Lemma 13, we have
E Yy Asu, XR, A -1)]
g g" 1xxxxxXXxxXX Aix,.. ..... x A,x) XA,, 1)
Σ A& l
Daylorr
Lr² log d
Σ You Anu
( (Vg(())) +
vd
Nyn Azu
k=1
Dyglo
LrKi
vdr
Altogether, we have
(x T Aix,..., AA,, - -
(I) xx - -1]]
TR1 7K1 + 23/222 +
Lr d L log² d Lrsi dd
Vd dr
dr
vdt
< d 1/4 Lr² Lr²rx Ki log² d.
where we set t = kid 1/4 The proof complete.
Following the proof above, we complete the proof of Lemma 15.
Proof of Lemma 15. Similar to the proof of Lemma 14, we decompose BJ and B₂ as follows:
B
As u.
1=1
= Σ Xij Uu:uus
41)
Σ A
Σ
AujuTAAuus A + B,
1=1
where {ui. 1=1 are orthogonal unit vectors for i = 1, 2, respectively, and
B Σ u +
Σ A
Σ u. Astis Ap-
|
1=1
Then following the proof of Lemma 14, know for any 1,2 and & = 1, .......,
VrKj
tr B = [B, = 0,
B, F 11+
B
ST+
Σ
1 -1
Vdr
dt
Let's denote the bi-linear operator I(-..) Iddd Rdxd E being
I(A, B) E (x) - -I) -1)
(19)
28

---

Published as a conference paper at ICLR 2025
By Corollary I and the proof of Lemma 14, we have
I(BI,
x
PR1 Vd
1
r
+
(++
)
(1+ (+v)rr)
Vdr )) Lr log d.
yr
(20)
dr2
dr
r
Big
Σ
A (u (u-i5
1.3 1/d)
T
12-
7// xj
-T-
TK1
+ T +
TK1
Lr log
(21)
dr
dr
r
B
Σ الووول u
inj As Ak
Aulli
p0//2?
ST-1
rH61 +
TK1 dr + T +
Lr log d.
(22)
Vd
dr
Here -i means 2 when i = and 1 when i = 2. Moreover, by Lemma 13, we have
Σ
A A
Σ
Asuzs An
A&ll
A Il
Aruld )
)
Σ
Σ A2,
م ) Access)
(1+
Lr log² d Vd
y k
(23)
dr2
and
Σ Aas (نiس 1/d)) )
r
Σ uT. Asuss
Au
A FF |Ac,l
u Actis) )
Lo log d
Σ
Pe-ial
Vd
1A isl2
Lr log" d
rki
(24)
Vdr2
Vd
Finally, we have
r
Σ
dis (uisuls - I/d),
Σ
des (nnneay I/d)
L log²
Pesll
Σ pewl
Vd
Maz
Lv ki log² d Vd
(25)
29

---

Published as a conference paper at ICLR 2025
Summing (20) to (25) altogether, we have
[2)] X
+
rK
+
(++
TK1
( 1+ vdr ))
vr
Lr log² d
Vd
dr
dr
+
TK1
+ +
MR1
Lr log" d
Vd
dr
dr
L log² d ddr
x²
TK1
2 log d
1
L log² d vd
+
Vd
dr
dd72
Vd
72
S d 1/6 Lrr2k log" d, where we take T kid 1/6, The proof is complete.
C.2 BOUNDEDNESS THE LEARNED FEATURE
In this section, we aim upper bound the magnitude of the learned feature (w, h (x))). Since focus on the first training stage throughout this section, we denote n = n1 for notation simplicity when the context is clear, and let the training set be Di X2.. We have the following proposition:
Proposition 4. Suppose m2 d'C' and Ci for some sufficiently large C. With high probability jointly on V and the training dataset Dr, and with probability at least 1 4n cxp(-12 on for any x' D2, we have
+22 /mz
Vmzi log (m2n2)
(w, (2 (x'))
(1P32(5*)) vdlpp((")ll.>)
d3
di d'vn
d6
As corollary, when m2 Z d 2pt2
2 dt 22,2p and |P S
Va
we have for any x' E D2,
K2c
(w, h (1) (x')))
(26)
d6
Thus, by taking the learning rate 17 = Cm2 x for an appropriate constant C 0, we can
ensure that (n, h ( (x)) I 1 with high probability.
Proof of Proposition 4. Note that
1
m2
(w, (w,
wAf (x, K (xi, x 002 (v] xs)
m22
m2n
mg
(xx) x m2 (x,, x') x2) [v] (x)x)).
m2
We do decomposition follows
m2
(x K (0) (x,, (x,,, x20,022
m2
(xx: K (0) (x,, K (xx;,)
w )/joo))))))
m2
j=1
A1
- Ex (xx))x ( x
11
A2
m2
Wi Ex
(x,
112
Aa We consider derive upper bound A1, A2 and Ag, respectively.
30

---

Published as a conference paper at ICLR 2025
Lemma 17 (Bound A1). Suppose m2 C With high probability jointly V and the training
dataset D1. and with probability at least 1 2n (xp((1(/2) on w, for any x' E D2. we have
|Ai| S
m2d3
Lemma 18 (Bound A2). Suppose m2 2 C' and 3 Ci for some sufficientll large C. With high probability on the training dataset D1, for any x' E D2, we have
12
d /mgn
Lemma 19 (Bound Az). Suppose m2 2 C² for some sufficiently large C. With high probability jointly on V and the training dataset D1, and with probability least 222 exp(-(=/2) on w, we have
e log (m2n2))
(IPsa())llaa + +
Aal S
Vmzdd
Similarly, for single point x', with high probability on and D1, with probability 1-2cxp( on w, we have
A S 3 (menz)
(1P22(0)l1> + vālIPa(NIlz))
mmzd
The proof of the three lemmas are provided in Appendix C.2.1. Combining the results in the three lemmas above directly concludes our proof.
C.2.1
OMITTED PROOFS FOR PROPOSITION
Proof of Lemma 17. We can rewrite A as
|A
1 n
(xx)
Ki0 (x,, x)) K( (xx,x,
1 112 W w (v) x:)
m2
j=1
m2 Since /mnz j= wy0z vV x) N( Exe 02 (2( we know given any x and V, with
probability at least 1 - 2 2exp(-2 -12 22) on w, we have
1 m2
w. w3o2 (v, x)
22
V x)
m2
mm
m2
j=1
Moreover, by (34) in the proof of Lemma 24, know for any and 0, have
(2 2 - -Ev, [o2(v]] x(]]
Pr
>
j=1
Pr
K
(x,
K (x, xx)
>
t
m2
-t/2
< 2exp
C. 24
+
C2 3
Altogether, when m2 > d'CA,, by taking t= C 2/2 /d' for sufficiently large and union bounding over the dataset Dj, we can ensure that with probability at least 1 - exp((1) on V and at least - 2n exp(-12/2) 2) on w, i.c., high probability on w, V, we have
1
me
Ct
C2 d
iCa /mzdd
Wjoz (v)
Vx Dj.
(27)
m2
/mz
V
Vmzd²
31

---

Published as a conference paper at ICLR 2025
Here C3 is constant. We denote this joint event by E1. On the other hand, by Lemma 24, with high probability on V, we have for any X, E D and x' E D2.
K my x, x') K(O) (x,, x))
(28)
/mgdd
We denote this event by Eg. Last, we truncate the range of the target function f. Denoting the truncation radius as (Cn)² for a sufficient large constant C' and 1j = (og(dmmmnn1pg
Pr [1) R] 2e (this could be guaranteed by Lemma 4). Given iid. samples
X1.X2 X g (Vd), we have
Pr Pr(l RR [m]] > 1 2ne 22
(29)
Thus, with high probability on the dataset D, we have [s(x) iP for any x E DD. We denote this event by Fg. Thus, combining (27) (32) and the truncation radius of f, we directly have
ICs
Czip+2
| S P.
/mzd2
/m2d
m2d3
with high probability (under events E,, Eg and Es). The proof is complete.
Proof of Lemma 18, We rewrite A2 as
1 wjoz (v)o2(vx)) x)
A2
Ex
(xx))
(x,,
Wjo2 w (v]xx)) x:) -x(())( (x, x')
1=1
Denote =xx
ma E wjr2( [v] x). By the proof of bounding Ai we could choose the truncation radius as R = (Cn)² such that [x)) S R for all x E with high probability (1 - 2ne 2n)- the dataset D. Now we denote truncated version of Y by
= f()) (x, x))
((x) (x)) R R)K (x,x'
x)1
1 m2
Σ
wy
LCs v /mzdd
m2
Here C3 is constant defined in (27), Now, we decompose the concentration error
Y(x)) - Ex [x)) = Σ (Y(a)) v((x))
(((x;) - Er,
7(x;))
1=1
1=1
La
61
F(xi)) - Ex; (x:)))
+
Exg
71
We know with probability at least - 2ne
D. Co 0.
Bounding L.. We attempt to use Bernstein's 's type bound. First we derive uniform upper bound of Y yx). By the definition, we have
7(x) R
kTt)
(x,)))
1 m2
Σ
1C3 mmzd
m2
j=1
C2 iC's
R
& /m2d
RC
d3 \m²
32

---

Published as a conference paper at ICLR 2025
Then, we bound the second moments of ((x) - Exx [xx)
which is
Var [(x)
- Ex
1/2 (x)
27
S 1² Ki Ex
K (x,xxxxx')
1
0j02
x)
V
1Ca M2zd
m2 j=1
222k R1 Ex
(x100(x,,,)
21
5C'3
mm2d
of
CC's
72 N1
k33
/m2dd
Car
M2d
Here C4 is a constant. Thus, by Bernstein's inequality, we have
R
Pr
L1
< 2 exp
da
m2
C4 C.Ca
Thus, when n >/ Cd2dd, by taking t = 1² and R = (On)² & eP, with high probability on D, V and w, we have
16 <
d4 /mizn
Bounding Cz.
suffices to bound
Ex
(x)
-1x [x)(
< 18x (x)) k10) (II, x²)
1 m2
f*(x) > R
-
1 m2
woo2 (v] x)
LCg
112 1=1
m2 j=l
Ex ((('()))3|$P))
((x) > R or
1 m2
woo2 (v) x)
1Cs
Ex
x()] [x)
$ 7Css Vmzdd
m2 j=1
/m2dd
(exp(-n1) + exp(-1)))ck
d3 mmz
Taking 1) > 2log n 2 log d log (C3) and i > Cn, we ensure that L2 S 1/(d" mm2n). Altogether, with high probability (event Eg) on D, we have
2.p+2
AA
(xx) - Ex [x))
<
da Mizhi
The proof is complete.
Proof of Lemma 19.
We remember that
1 ma
m2 wy
Ag
my
Ex
fx)) (x, (x)))
Σ
wyh((,,.') x'),
m2
m2 j=1
where h(v,x') = Ex (xxx))) (v (vTx)) x)). To bound h(v.x') uniformly, we have the
(x,x
following lemma:
Lemma 20. With high probability on V and the datasets D2 and D2, we have for any j E [mz] and x' E D2.
22 2 (m2nz)
1h(v,x))
(IPS=(f))ler (V(pp(())l)
d6
33

---

Published as a conference paper at ICLR 2025
The proof of Lemma 20 is deferred to the end of this section. Thus, condition on the event above, by invoking the upper bound of Gaussian tail and uniformly bounding over x' E D22 we have with probability 1 - 2n exp( /2) on w, for any x' D2, we have
mg
e log (mzuz) vmzd
(IPse(0lle> ValIP(()llza)
As S
h2
Vmz
m2
Also, for single point x', with probability - 2exp(-((/2) on w, we have
1 log looog (m2u2)
(IP>2(0)ll>> vaIIa((IIL))
As x
/mzd
The proof is complete.
Proof of Lemma 20. Recall that the activation function d2 admits a Gegenbauer expansion
(t QQi(t).
Let's fix x' and v. Note that we can decompose h(v, x') as
x(v,x))
Ex
f (x)) (x, (x.x')02( x)
(x(x) Σ 2Q(( x')
E
Ex
i(d,i)
j=2
Ex (x)
Bid, 1)a (Yi(x), Y.(x)) Cj
(Y, Y, (v))
B(d.j)
1= 1=2
Goj B(d, (dd.))(dd.j (Y ((xx)&Y(V))).) E,,)),,, (x)))
ctcy B B(d,j)
he (v.
By the definition of (v, we have
[12 (v,)))
(30)
=Evx'
(Xx(x') @ E (())))
(Ex (J'('')Y,, 2
X,((x))
-B(,,))B(d, Exxx
[xx)[(x)( (x (x
min(
B(d.) Exxx
f(x)f(x))
Σ
6 24 Qits xx)
min(b.))
(hd)
((d.) Exx
f(x)f(x))
Σ
i4j-2%
(Y (x), Y 2x+)-)((x'))
B(d,i+j 2k)
min(i.))
(1.))
((d,
- 2k
Ex (
B(d,i+j-21 2
min((
(iJ)
B(dd,))B(d.))
||| (f)l
(31)
B(d,i+ÿ-26 2k)
34

---

Published as a conference paper at ICLR 2025
Since hi (v.x') is a degree i polynomial of x' and degree polynomial of v, by Lemma 5, have for any q > 2,
2/4 s (q 1) [V x') xx [ x') x')
Let 5 = (2ei log (m2n22) (6+))/2 for some > 1, taking q = 1 - 1 (2(((+))) and Markov inequality, we have
Ev,x' [lhi xx)
Pr
[V, > /. Ev (v,,
8 Lvxx' [ (V
S (q - 1) ((tt)///8-1
i i+j
2ec log (m2n2)
exp
1+
e exp(-(i++ (1/2). = (m22m))-((tt))
Thus, with probability at least ((m222)1 (++5)
(V,(V;;,') (2ei log (mzn2)) (++))/2
Ev [V, (v,)))
[(++)/2]
B(d, (dd,)))(( (PP+)-2x(f)11
(2ei {(ee log (mmnz)) (1+5)/2
$
B(d,i+j - 2k)
In the second inequality we invoke (31). Summing over and j gives rise to
h(v,x')
(v.x')
B B) dd.j)
lcyl (2er log (m2n2a))(e (6+))/2
((+))/2]
(dd
(8.5) B(d,i - -2k)
[[(P+)-2*(
B(d,
k=0
lcyl (Zel log (m2n2)) (1+5)/2
B(d,i) (dd
1 jm
2
ccyl (2et log (mzn2)) (1+5)/2 ((+1)/2)
B(d, ((4,)))( (PP4)-2x())1l)
Bid, i)2 (d, j/3/2
B(d,i+j-2h) 2k)
2=
log(mgn2) d5/2
Σ E=
Do
1
Σ
(15) (2ex log (m2n2)) (1+3)/2 (1+3)//
[|(P(()))
B(d,))
i(d,i)
itj-t
In the second inequality invoke Cauchy inequality. Then by plugging the bound 6(63) in Lemma 12, we have
lo(vxx))
(log(m272)
A(2tdd-2)
(.3)max(k:22)
a ley (2er log (m) (6+5)/2
8
Σ
Al (P((()))
d3/2
B(d,e(dd (2)
B(d) 2)x k
E=0
i4j-t=2k
log (2ggm))2)
(log(m2n1))
clog(mgm2)
[P22(f)ll3) )
|P>2(f))le
d5/2
d7/2
d3
2 log" (mon2)
ia loog (m2n22)
||>2(f))lz +
||
d
d31/2
35

---

Published as a conference paper at ICLR 2025
The probability of this event is at least
(m2n2) ((+t)) exp((i + 5)/2) =1 m2nz( (-(m2n2) --11/2)) 1/2)2 (m2n2) )e
j=2
which is high probability event when uniformly bounding over x' E D2 and v V1, V2 V..... Vmg- The proof is complete.
C.3 PROOF OF PROPOSITION
C.3.1 THE FORMAL STATEMENT OF PROPOSITION I AND THE COROLLARY
Let's consider a formal version of Proposition We remind the readers that throughout this section we denote n T2 for notation simplicity, since we only focus on the first training stage.
Proposition 5 (Reconstruct the feature). Suppose m2, n. d Cd for some sufficiently large C. With high probability jointly V and the training datasets DD and D2, there exists a matrix B
Rrxma satisfying || S
de
[1 such that for any x' E D2, have ma
Amin
Vr (xh(((x) p(x)) ~ Amin (H)
(P+2d
ud
pt3/2d
L log
m2
V m22
d1/6
Here L L xx is Lipschitz constant satisfying Vg (p(x))l S L with high probability.
With the proposition above, we directly have the following result.
Corollary 3. Under the same assumption Proposition with high probability, we have
nn/2
2
p43/2d
L Ki
log" d
g (B*h(((x)) - < 1 Amiu (H)
sup xED,
m2
m2
Vn
d1/6
We provide the main proof of Proposition in Appendix C.3.2, and defer the proof of Corollary and other supporting lemmas to Appendix C.3.3.
C.3.2 PROOF OF PROPOSITION
Proof. Denote the target features by p(v) AIv, v A.v] R" for any E Rd, and we further let P = [p(v1).p(va) p(v2), pVVVm.)] following decomposition
(((Y))) E Rmzxr, Then for any x' E D2, we have the
P T (x))
(*xx)) (x,,
(x,, (V) x.lxxx)px x()))vv)
m2
m2n
1=1
(x:)
(x,, x') - (x;' (223))))22vv] 02(vV (2Vvxx))
m2
D1,1
f (x),,
Σ oz (v)vx:))((v)
B(d, p(xi)
m2 j=l
1=1
Dia2
(xxx)
(x,, Ex (xx
(x,x') x)((x))(
n nB(d2
D2
Ex (xxx))
B(d,2)
D₂
36

---

Published as a conference paper at ICLR 2025
We will derive upper bound the concentration enror terms D1,1, D1.2 and D2, respectively. Moreover, leveraging the asymptotic analysis in Appendix C.I, we show that D3 s d Hp(x)(x)) with high probability,
Lemma 21 (Bound DI.1 and 11,2). Under the same assumptions in Proposition 5, with high prob- ability on V, D1 and 222 we have
9C//4 p+2
9 O2
11, I
and
|D122 S
mad
/m2d3
Lemma 22 (Bound D2). Under the same assumptions in Proposition 5, with high probability Dj and 22, we have
P4312
D2lloo S
And
Lemma 23 (Compute Dz). Under the same assumptions in Proposition 5, with high probability for any x' E D2, we have
c D 2(d,2) - Hp(x')
(Lr3 log d
S
d6+1/6
We defer the detailed proof of the three lemmas to Appendix C.3.3. Combining all the results above and choosing
B B(d,2) HT
2
have with high probability V, Dj and Dz.
B'h (x) - p(x) -xxx)
Bdd, ((d.2)ddl-3 - H (1,1+ +12 D2 D3 B(d, 2)d(d [d(dd-1) - 1) Hp(x')
&
c
VF
|ID..ll.o + +|D2120 D22loe
Amin Ami
&
+ D3
Hp(x')
Bid, ((d,2)d(d-1)
VF
F 2d
uB
pt3/2d
2 log d1/6
Amin H)
1112
m12
To bound [B*] lopp note that
B
11
B
d2
PPP
T
m min
2
H
op
d2
m2
p(v, p(v))p(v;)
22 min
H
m2
op
Moreover, for any j E [m22), we have
|P(v))D(v,)TI 1lop
1ID(vv))1
= (v] S Srd.
1=1
37

---

Published as a conference paper at ICLR 2025
and we have
Ey (olv))(())>1]1)
Ey Σ (v AAv)
d | [p(v)p(v)']Illp
rd..
The second inequality holds because p(v) is positive semi-definite. By Matrix Bernstein Inequality, we have
rd 2
Pr
Σ p(V,
p(v))
3 1 +
Vrde
S exp
27n2
m2
rdd
rd
y /Tde
lop
mg
V mg
2
?
exp
1
fide
3 m2
Thus, when m2 d4, know with high probability V,
Vrde mm2
p(v,)p(v,)
S
S 1.
m2
y=1
op
Thus, we have B llap S
The proof is complete.
Amin (H)
m2
C.3.3 OMITTED PROOFS IN APPENDICES C.3.1 AND C.3.2
Proof of Lemma 21. Let's first bound D D1,1- We can rewrite D1,1
D1.1 = 1 - (x,)
1
K (x, (-(0) (x,, x')
o2(vvxx.)pt x(),)
m2
1=1
By Lemma 25, for any le [r], we have with high probability on V
Cud
v AAV))o2(v]xx) B(d,2) 2)
x A
m2 j=1
m2z
Thus, by enumerating Ag over (A1, A2,..., Ar), we have with high probability on V,
m2
9id-
Σ
P(v xx)
B(d, 2) p(x,)
C2
m2
m2
On the other hand, by Lemma 24, with high probability V, we have for any x; D1,x' D2,
( (xx,x)) K (00) (xx,x')
(32)
/m2d²
Moreover, under the event Es (defined in (29)), with high probability on the dataset D,, we have [s(x) & iP for any Dj. Thus, altogther we have
9rd 9rd-1C2/4
9C2/4 pt2
Dialll Sp. /m2d
/mz
m2d
38

---

Published as a conference paper at ICLR 2025
with high probability. To bound 21,2, from the proof above, we know with high probability,
m2
9rd C
Σ
p(v,))o(vvxxx)
B(d, P(x,)
C2
S
m2
j=1
Moreover, for any Dj and x' 22,
C2 (x, Ev [02(v' (x)ox(''Vv xx) S v Ev (02(v [x)] (xT)\
(33)
d²
Thus, we can bound 1,,2 with high probability by
C2
9id C//4 9/C//422
| I
d²
Vmz
vm2d3
The proof is complete.
Proof of Lemma 22. Thus, let's focus on the concentration of a single element
x')x A k 1, TT Similar to the proof of Lemma 17, we denote a truncated version of Yv by
(x(x) = f" (x) (x,
(x) f (f) (x) I R) R) (x,x') x')xx A k ...,,....? r. Here, R = (Cn) for some large constant C. Now, we decompose the concentration error as
1
(x,x) - Ex [Vx(x)]
(xx)) (x(xi)
+
(x(xx)) - Ez. (x,)))
7 1=1
7
Eo
1
Ex,
[F((x))
-Ex, (xx())))
11
L2
By (29), we know with probability at least
2ne 2n, Co O.
Bounding Ct.
First we derive a uniform upper bound of Ya (x), which is
[xx(x)] = R (0) (x, x') x')x T A,x
$ R k-(0) (x, x |x AAx
$ R C2 d |
d²
RC2 |ARkllo
d
Then, we bound the second moments of Fx(x) - Ex [T((x)] Again by Lemma 8, we know that
there exists sufficient large constant C > 0 s.t. Pr |x Axx > Col I 2exp(-().. By taking
i > 2 log d, we have
Var (x)
< Ex
Y2 (xx
Ex
72 (x))
Axx|| S Ci)]
+ +Ex (x) {1x A > cij]
(x11(((x)))
12
I C2p2ni
Exx
Ex
Ex
[1 f f A c(J
dc
2T C2 exp(-1)
Σ
5
+
dl
2
C'r2k12
d6
39

---

Published as a conference paper at ICLR 2025
Here C' is sufficiently large constant independent of d. We invoke (33) in the second inequality. Thus, by Bernstein's 's inequality, we have
C
Rv
C n
Pr
C1
< exp
2nd
d
C'r nd
RC2 A| lop
F2 mdl
3nd
= =xexp
Cs Ax lopp
44
3C"
Thus, when n 2 C3 AA dt, we have with high probability on the training dataset D1,
Ru
C²
Ri
3/2/2
1
S
d
v Vnd
Bounding L2. suffices to bound
Ex (x)]
Ex [V((x)]
- Ex (x)
R} k (x, AaX)
Pr [x) nJ Ex (x, }
s Ex (f)'
Ex 1(x T Axx)*1
C2 (8 1)
s 1 exp(-nn/2)
7C2 exp(-n(2))
d²
Here invoke (33) and Lemma 8 in the last inequality. By taking 1) 2 2 logn 8 log d can ensure that with high probability, we have
3/2/2
Σ (x,) Ex (x)(
S |
122l x
Vnds
Thus, by taking over [r]. we have with high probability over the training set D1, we have
122
R/3/2
R13/2
P+3/2
D2llo S
x
S
B(d,2) Vnd
vnds
/nds
The proof is complete.
Proof of Lemma 23. Note that for any i E [r]. we have
Ex
f (x)
(xx, A4x
QQi(x
Ex
(x)
Σ
- AAx
B(d,i)
1=2
06
of
Σ
Ex (x))) Xx)x Aax]
i)
==
[Fx A x1], -1) -1)
B(d, 2)d(d
&
Ex (xXx AAX]
i(d,i)
2
(T(A)))X'' -I)
Ex xxxxxx
B(d,2)d(d-1)
B(d,i)
3=3
40

---

Published as a conference paper at ICLR 2025
Here T is the linear operator defined in (14). Recall by Proposition 3, have
T(AA) Σ H
5r21/6 log2d. d.
Let's denote R T(A:) - Ei H Hr x'x,''x' 1) is quadratic function of x', and Ex (RR,'''
that |R&ll x d d log d. Since
- 1)2]
2d
R
By
2+2
Lemma 8, there exists constant C' > 0 such that
Pr Pr| Xx - 1) Cu/Ex
Bx (R x 1)21
S 2 2exp(-1).
Thus, by enumerating [r) and x' E D2, we obtain that with high probability (1 - nr exp(- on D2, for any & E [r], we have
LLr k2 log² d
('(AA),x'' I1 Σ Hx,jx AxX
x
d/6
Moreover, we have for any x'
Ex (x) A&X
B(d,i)
1=3
00 df
Ex
[xx Axx]
B(d,i)
of
Vxx Ex Xx (x) Axx)]
i(d,i)
v Bx (F Agx)a].
B(d, 1/3/2
1=3
Again by Lemma we know that there exists a sufficient large constant C
0
Pr ||x A&x
> Ci} < 2exp(-t). By taking & > (2p + 2) log d, we have
Ex [x) (xx)]] = (x) 1xx)21 A&X S sC.
+ Ex [f 11xx A > CtH]
$ C²2 + 20 exp(-1) cxpp())
- C2.
Altogether, with high probability on D2, for any K E [rr], we have
o
Ex (x)) (x, (x,,'x A
Σ /=1
Hx Aix'
B(d, 2) 1)
u log² d Bdd 23d
Ciei
i(d,i)3/2'
Thus, by paralleling the entries together, we have with high probability on D2
& Ds B(d,2)'d(d-) 1)
2 Ri log² d
Cid²
Hp(x')
((d,2)2d(/6 11
B(d, 2) B(d i)3/2
00
2 K1 log 2 d
6+16
The proof is complete.
41

---

Published as a conference paper at ICLR 2025
Proof of Lemma 3. By the mean value theorem, we have
g(B'h((x)) (x)) - s(((x)) x
Vo(AAB"i"(x)++ + (1- (01)))(()){[1 B'h(2(x) - p(x)
sup xe
Recall by (8), we have [Vg(z)|2 x Illll Ex=1
| Note that with high probability,
supxED, |Ip(x)] Ö(vr). Therefore
V (x) + (1- ()-))((x))]|| & aa T =
sup sup xED; A=
Altogether, by Proposition 5,
(B' (1) (x)) - g(p(x))
sup xED=
< ||ا||
Bh p(x)|1
pp/2
2
m43/2d
LL m log
||ا|ا
+
Amin
(H)
m22
mm2
vn
d1/6
The proof is complete.
C.4 PROOF OF OTHER SUPPORTING LEMMAS
We first present the concentration of the initial kernel King (x, x').
Lemma 24. Let (0) (x,x') 1 my (o2 (x), 02 (Vx')) be the initial kernel with inner width being
M2, and (*(0) (x,x') Bu- (va) [ x)o2(vTx')) be the infinite-width kernel. Then there exists constant C when m2 2 Cd... with high probability probability w, V and the training dataset DD for any x E DD and x' E D2, we have
k!0) (x, x') K(0) (x, x²)
smzzd
Proof of Lemma 24. By Assumption 4, for any x,x E D and v E (vd) we have
(v xx)2 x') $ C²
and
Ev 02((V (v x) 02 (v xx')
sv Be (22(v xx)* Ev (02(vTx)))
C4
5
d1
Thus, by Bernstein inequality, we have
Pr 100) (x, (x,x') x') - (0) (x,0(x,xx) >
< 2 exp
2mng
m2
C4
C2
m2dd
3mz
m2
- -1/2
exp
(34)
C4
+
C2
d4
3
By enumerating x, x' over D, have
11/2
Pr max X
K:10) (x, x') xx(0 mg
>
V
S
n² exp
C
ma
Thus, when ma > d', we can take f. = 22 /dd to bound the probability by poly(d, n. m2)e which concludes our proof.
Then we present the concentration of the reconstructed features.
42

---

Published as a conference paper at ICLR 2025
Lemma 25. Suppose ma 2 C2 spherical harmonic, with high probability V, for any D, have
| Given any A such that v Av is quadratic
m2 v
9ed 9rd C
Av,)o2(V,
C2 B(d,2) 2)
Ax
Vm2
Proof of Lemma 25. Given any fixed x E D and A such that v Av is a quadratic spherical har- monic, we have
Ey [(v AvV x)] 1 V By l(v E. [x)(v
1 (4 - 1)2²2 E [(v; Av,) } d 2
-
and
I(v) Av)o2(v x)1 S d ||| C. = dCa ll
Since By [[vTAvv]]] Av) Bldz Ax, by Bernstein Inequality, have
9dd C//
Pr
vv Av. Avi)oo(v
62
x T Ax
172
Bdd, 2)
m2
i=1
8ld-201/20
< exxp
2mng
81d-20!/2
9ed- C//4
d(Co ||A
m2
3mz
2/2
exp
C All ez +
27C 74
Vm2
Thus, when m2 > C3CT1/2
1/2
|| by enumerating x E D, we obtain that with high probability
on V, for any x E D, we have
9ud-
[v] Av
C2
x Ax S
m2
B(d,2)
m22
1=1
The proof is complete.
D APPROXIMATION THEORY OF THE OUTER LAYER
D.1 PROOF OF PROPOSITION 2
Since we mainly focus on the first training stage throughout this section, we may sometimes denote n My for notation simplicity, and let the training set be Dj [x1, X formal version of Proposition 2.
.......,,, xx) Let's consider
Proposition Suppose g is a degree p polynomial. By setting = C5-5x5 m2 1/22
m= do for some
constant C > O, with high probability over Dz. D2, mm and V, there exists a* R mi such that the parameter f+ (a", W(1), ((1), V) gives rise
Cz(0'):
Σ (f(x;0') s(p(x)))²
112 xED,
II log d 2 d1/6
rp
252
id3
p+3/2d
I 11 Ami (H)
+
m2
M2
Vn
++1 2
k
B
k
op
T
p-k
m1
k=
43

---

Published as a conference paper at ICLR 2025
Here a' satisfies
|a m1
P Σ 0==0
2 r
B* 1lop
52(x2 rP).
To prove the proposition, let's introduce the infinite-Quter-widthh model transition term between the finite-outer- width model and the target function. We define the infinite-outer-wide model as
fo ma (x; B.,., [u(a, b, w)or (an(ww, h (x)) + +b]]
where (x') = - f (xx) King (x,, 02 (V (VTx:). xx).
We can decompose the L loss of the truth model [xx 0)
L(0*) = Σ
(
xED2
1 Σ
(f(x; 0) - fo (x('))))))))) + foo.ms(x') - + (BBh(())
s(p(x)))
x
(f(x;0) - foomms (x)))
xED;
L1
Σ
xED2
((Bh - =(p(x)))
-(p)))
xED=
We have bounded Ti in Corollary 3. We state Lemmas 26 and 27 as follows bound L, and T.2, respectively.
Lemma 26 (Bound L2). Given B E Rrxmz and setting the learning rate 06 m2
1,, 1/2 11/2,6
for a constant 0, there exists : {+1} x R m R. such that
||0||2>
1 L2
Σ n * IB* 10p
k=0
and, with high probability over Dr, D2 and V, the infinitewidtth network satisfies
(fx,m (x:1) ((Bhhh)) (x)))² So
iningmim
xEDr
Lemma (Bound L1). Given the function v {+1) x Rm Lemma 26. With high probability over D1, D2, (w and V, it holds that for any D2.
cla b, w.)oi(na, h (x)) + b.) foc.m (x; S
with
m1
mi
Σ u(a,, by w:)? S 2
The proof Lemmas 26 and 27 is provided in Appendix D.2. Now we begin our proof of Proposi- tion 6.
44

---

Published as a conference paper at ICLR 2025
Proof of Proposition 6. By Corollary 3, Lemma 27 and Lemma 26, by defining the vector a* E R by a v(a; (0) b(t) ( and letting l* (a", W b(1) V), we have with high probability that
(2) E L1 + 12 + L
2
2
ud
L log2 d'
2 11
X min
HH)
11
m2
vn
d2/6
2
p+1 llal122
Σ
B
lop
r
m1
R=
1
2
2
119
2
pp
2
ud
2pt3/2d
L
d d1/6
log
1.2
X min
H
m2
m22
Vn
2
++1 |1911222
+
-k
|B
lop
r
m1
k=
Here a* satisfies
my
|aa
uaa,, br, wi)
2 m1 |01|22
p "
2
smp&
19 l2
E0 |+1
2
The proof is complete.
D.2
OMITTED PROOFS IN APPENDIX D.1
D.2.1 RANDOM FEATURE CONSTRUCTION OF UNIVARIATE POLYNOMIALS
In this section, before proving Lemmas 26 and 27, we first construct univariate polynomials using the outer activation function σ1 and the random features a and b progressively.
Lemma 28. There exists vola, b), supported on {+1} x [2, 3). such that for any |z| < 1
Ea,(Looaa, ()z + +b] =1, sup (a 6)1 S
Proof. Let vo(a, b) = 12. 12=1 ()
10elz.3 ((b)
Then, since z + b > 1,
3
5.
Ea,luvo(a, b)o(az= b)] = 6
(b
= )o + b)db
2
2
3
5.
= 6
(b
2
(2z 2b 1)db
3
5.
3
5
= 6
/
(b
db + 6
/ (b-3 2) - 1))db
2
1.
The proof is complete.
Lemma 29. There exists Un (a, b), supported on {+1} x [2, 3). such that for any |2| S
Eg ()z + b]] = z, sup (o 6)] S
a
45

---

Published as a conference paper at ICLR 2025
Proof. Let b) 11==1-22200 61)
1ne ju(b)
Then, since + b > 1,
{
n
(a, + +b]]
(-24b + 61)o + bjdb
2
(-2Ab + 62) + 26 1)db
2
(-24b 61) + 2 /." (-24b + 61) (2)4b+6)((2b.)) 1)db [
1
2
The proof is complete.
Lemma 30. There exists V2 b), supported on {+1} x 1-2, 3), such that for any |2| S
ba)l + b)] = 2², sup (u(a, 6)1 S 1.
o.b
Proof. First, see that
S
2+z
o(2 olz
+ b)db
ob))d
24
[]] bddb +
2+ f (2b (2b- 1)db
(-2b 1)db +
2+z
-1
2 [b² b]2+z
[-b² b]]
2
(2 2)² +(z - +)++
(2 + 2)² -(x 22
14
2 + 3
Let (2 b) = 1amm
laer-22 2.21 ((
(v b) Then
1 / a(2 b)db
7
(al +b]] = bb]=
_
7 7
= +
3
2.
The proof is complete.
Lemma 31. Let v(b) - _kk(k 1)(k: 2)(1 - k))- lec/0.11 Then m(b)
Eo(vee + b)] = 1,>0 k(k - 1)
2
k2 1.
2
Proof. Plugging in un (b) and applying integration by parts yields
/'"
1 k(k l) - 2)(1 6) - + b)db
Es Esluu(bo((2 + +b]
((z
12k(k 1)(1 - 30 6)1d /'' 2 ki( 1)(1 - 3)1-220(== b)db
1
-k(k - 1)0(=) + 5 k(1 b) ,)² d' (= + b]]b f' 1 k(1 "g"--d"l + b)db
k - 1)o(z) - kd kz)
1
/' k(1 - b)*-11 1 |2+ s1 db
46

---

Published as a conference paper at ICLR 2025
When 1 3 >> O, we have
- -f k[1 b)) 1 db
1-=
k(1 - b)b-idb
db [(1 b)m1b- 2
When -1 2 0, we have
L. k(1 - -)) 124b//10b 1// 1:1 b) db -1
Since E [-,,1], we have that o(2) = 2² and (z) 22. Therefore for E [-1,1]
Ea(ue(b)o(2 + +b] = 2² -1200 k(k - 1) 22 k2 1.
2
The proof is complete.
Lemma 32. There exists un (a, b), supported on {+1} x [-2,3), such that for any [2| S
I Pa.l(vale ()z + b)] = sup (x((,,)) S poly(k).
a.b
Proof. We focus on h: > 3. We have that
k(k k(k-1)
Es(ux(bjo(2 + 6)] = =² . 1,>00
kz
2
k(k
11)
[=lxx(0)0[=++b) = 2- 1:20
+ kz
2
Therefore if even
Bs(r(b)o(= (()) +1) = - k:lk - 1)
Let un(a, b) 2vx(b) k<h 1008811)220000 1)02(a.b) + Then
(a + b)] F. Ea(rx(b)a(= - b) + (b)o(( - b)] +(hh 1)=² + 2 2b.
If l: is odd,
Es Es(u(b)o(z +
+ b]] 2k:
Let va(, b) = (b) + 2k (a,b). b). Then
Eu, (a, (az + 6)] (b(ox (z))) cx(b)o(2= 6]] + 2kz 2h.
The proof is complete.
D.2.2 PROOF SUPPORTING LEMMAS APPENDIX D.
Proof of Lemma 26. Let's consider a general version of Lemma 26.
Lemma 33. Let 9 ER E. be a degree p polynomial, and let B E R Given a set of
vectors D = [Z1, Z2.... 2n) C mm that satisfies 11 I for any z E D with probability at least 2(n1 + (2)
over w: {+1} x R x R + R so that for all z E D.
N0 Img) (uniformly over D). Then, there exists
E Ba (o(a, b, w (w, 2) + b)] = =B( + ( and
P Σn k=0
|el| $ 1lallz
B
10p
m2 ",, -1/2 de for a con- stant C' > ensure (, (2) (x')) < 1 for any x' E D2 with high probability on V, Dj, and probability at least 2(n + n2) exp( 22 /2) on w. Thus, taking D = [h(x))xeD, concludes our proof.
Thus, according to Proposition 4, we could set the learning rate 7) = Ct Cr-5k,
47

---

Published as a conference paper at ICLR 2025
To prove Lemma 33, first decompose 9 into of polynomials of different degrees and con- struct a function e to express these polynomials accordingly.
Lemma 34. Given z E R Ler B* Rrxma and Tk E (R") such that
0k Then, there exists og IR - R
Ew (w((, 2 Tx =Tx ((B*z)
Here Uk satisfies
llullla> Sn [ leo |Tell and sup suplos(w)) (w)) i m2 k/2
|B 1 | (35)
Proof of Lemma 34. suffices solve
E Ew(t(w)w@k 8k1 = = "1= @k (Tx);
where B (Tx) (Rm2)0t 0o This is achieved setting
v(w) == 1 -k Vec(w ok Mat(E[[w02])) -211)-1 Vec B* (Tx)).
Then,
||ul 17 -2 Vee(B 8k (T&))² (2 Vec 34 (Tx)).
Since
Mat(E( w 82k) z k!ll Sym (Bm2))
we have
|ull222 I n -2 B 8k (Tk) (Tx)) S vj 24 | 12 | 1|
2h
Finally,
sup [(w] 17 * sup
Vec(w Mat(E( Vec(B 26 (Tx))
/| Ir
BB (T&)
S
k/2 m2
B | |
The proof is complete.
Then we begin our proof of Lemma 33.
Proof of Lemma 33. We can write
g(z)
(Rr, 20kg
km
By Lemma 10, we have |RalIP Z
|1
Define 2y (a, b) to be the function so that Ea, in (a, b)o (az + b)] 2h, and let vx(w) be the function where Ew (w())(((ww 2) = (Te, (B²2) (8)) Next, define
P v(a, b, w) Σ ba(a,b)ux(w)
Here un (a, 6) defined in Lemma 34. Then we have that
|ell r.
Elen (w) < ||lll
"
k B
op
p
0
k=
48

---

Published as a conference paper at ICLR 2025
Note that || O(poly(mg,d)) and l0r (nnana z) + +b] na(w, z) + b) has polynomial growth. Since we have taken C (og(dd)n)mmmm)) for some sufficiently large C. 0, we know by Cauchy inequality,
Ba ( b)1(m(w 1II1 S
dm1ngmmmng
Thus, we then have that
Ba [ b, (na 2) +
B.,, (na 2) + +b- 1 (nw.
+ b) 1 nnwalllsl + (diinammm
Ew [ux (w) 1/1/w.a//s1 +o dni
Ew [vx (w)
k=
1
KR (B*2) 8k)
ddingmmmm
k=0
(B' 2)
dn dnin2mming
The proof is complete.
Proof of Lemma 27. Fix x E D2- For notation simplicity, we denote f (x) = foo (x; v). Con- sider truncation radius R > 0 be chosen later and let Es be the w such that
sup |u(a, b, w)l S R and n/w, hx)) 1.
0,5
By the construction of v(a b, w) in the proof of Lemma 33, know can be seen as a degree- p polynomial of w. Thus, by Lemma 7, by taking R C.P/2 | for some sufficiently large
C' > 0, we can ensure that
Pr[sup joaa b, w)l < R exp(-1). a.b
Moreover, by Proposition 4, conditional on a high probability event on V, D1 and D2, by taking 55d we have Pr|n(w, (x)) < the random variables
> 1 1-4 4- exp( /2) 2) for a single x. Now consider
2, 1(w; Jo bi, wi)ai (na;(w:; h (x)) + b,), i= 1,2,..., m..
We directly have that |Z,| S 2/2 |elllaa and with high probability.
my w:)² S uP || |2
Σ
olaa,
Therefore by Hoeffding incquality, with probability at least 2 exp((1)), we have
W. E, u(a,, by ( (x))) by)
m
r+41 |101
ElLwER, v(a,
+b)] S
m1
49

---

Published as a conference paper at ICLR 2025
Similar to the proof of Lemma 33, note that both v(a, b, w) and lo (na(w, z) + b)] has polynomial growth. Since we have taken L = log(dningmmmgg for some sufficiently large C' > O, we know by Cauchy inequality,
E EEEE e(a, b, w w (na(w. (x)) + b]] - fz² (x)
Bllwer, vla, b, (na h (x)) + b)]
2 ((w & (B)eaa b, (n ( (xx)
1)
2 exp(-CC log me
mi
Finally, union bounding over x E D2, we see that
m2
10+1 |
u(ai, w:)o (ma: h (x)) + +b)) (x) x
sup xED2
m1
m1
m1
The proof is complete.
E GENERALIZATION THEORY
E. FORMAL PROOF OF THEOREM
The proof is divided into two parts. The first part of proof formalizes the proof we present in Section
4. The second part presents the generalization theory after construct a* that gives small L error by Proposition 2, with the formal version presented Proposition
E.I.
PART1: ANALYSIS BEFORE FEATURE RECONSTRUCTION
Denote Wj - -w (0)
-O ((0,) Img) Note that for any DD and E [mi], have
(w) (0) (xx)) = (ew h (x)) ~ N 0, | [x00))))] (x)
Since (x) of S C² By setting - = Co we know
(0) (0) (x)) S with probability at least 2 exx((1)). Thus, uniformly bounding over x Di and [m1]. we know with high probability over W, we have
(w)) ;
(w/")) (0) (x)) < 1 for any x E D1, E [m]].
Then, according Algorithm 1, after one-step gradient descent on W, we know with high probability, for each j E [m]).
(0)
V w/ 10) ((0(0)) 7
Σ ((x)h( (lew,, (0) (x)()))
-1] m1
711
xEDi
2em
(0)
f (x)) wj.
which is a linear transformation on w,. By taking 77 2ema
n for some 17 > O to be chosen later
and A 72 we have
4 w C(0(0) + 1
w
W
V (0 L(((0)
=
(U) 1
na
m2
11 1=1
50

---

Published as a conference paper at ICLR 2025
Then for any second-stage stage training sample x' D2, the inner-layer neuron becomes
(0)
na.
(x)h wy. (2,
m1
71
(0)
K xx xx)
naj
Wj+
naj
(0)
(w,, h (x')).
Thus, after the first training stage and reinitialization on = b(t), the model becomes the following madoom-featur model in the second stage:
f(x':0)
mi Σ ajoi naj (0) (w,, h (x)) 6
1
By Proposition we know there exists at R such that with high probability over Dj, D2. {w:) and V, by taking the parameter fr = (a", W()), b(1), V), holds that
r.P
p+3/2d
L log
(2( S |ل
Amiu
H
m2
V m2
d1/6
+
12
p
π
1
Tp r
p-k
m
k=
Here a* satisfies
|a 2
Σn n B 1.op v
my
The first part of the proof is complete.
E.1.2 PART2: GENERALIZATION THEORY
Denote the population absolute loss as (i(fg) = Ex ll - g(x)]]. Moreover, we consider truncated loss function as
(2) min min) and (f,( Ex (f((((x)) -x()))
where т > 0 is the truncation radius. Moreover, we denote the empirical truncated absolute loss as
(r,((,
1 Σ
- -(((x)-gxx))
T12 xED2
Suppose Algorithm gives rise to a set of parameters θ (á, W b()), V), and we have con- structed 0" = (a*, W b(1) V) that leads to small empirical loss, we decompose the population
absolute loss as
C((((()),f) = (f(((())) (1. - - (f;:)))
Li(S(;0), 1r) (1.-(((;0))))
La
Here with little abuse of notation, consider f 9 (p) for learning the original target function and denote f* -==00 (p) with being any degree polynomial for the transfer learning setting. Next, bound L,, La and L3 respectively.
Bound L, With a little abuse of notation, we denote (z( = L2(0) for = (a, w(), b()), V) since only optimize a in the second stage. By Proposition 6, we know with high probability, the
51

---

Published as a conference paper at ICLR 2025
empirical L loss of 8* is bounded by
(²(aa)
1
Σ
(f(x;0)) - (x))
n2
xED2
rp
ud
233/22 d
L
11 log² d2/6
S |191
+
1 min
(H)
1112
1
V
1+1 l
+
7
up
r
m1
K=
Here a* satisfies
|
2/2 x
Σ
"1 -k
|
llop
p-kk
my
In the second training stage, let's set the weight decay in the second training stage as
2
p+2
ud
tt3/2
L log d
A = A= = ||a
||a1l222
+
12 min
(H)
m2
mmz
Vna
d2/6
p
Σ
B
mi
k==
so that the empirical 2 loss is directly bounded by
Ea(a") =
Σ (f(x;0)) (x)) s A || xED;
We further consider the regularized second-stage training loss to be
Σ (f(x: (a, W b(), v)) - (x)) 2 )
(a)
2 |la|
n22
xED2
Note that this loss is strongly convex, so it has a global minimum a (00) = argmin (a)) Thus, we have
(2)(((()) I (a') $ A ||a 1|
Since 22,2 (a) A- strongly convex, and we can write f(x; (a, W b(U) V)) = T (x)). where (xx) = Vec (mit o
(0) (w,, (x)))))))) (x)) + +bf")). Therefore, by Lemma 4 and our choice of n) to ensure v/,w, (x))) S with high probability, we know with high probability,
(na!")
Amax
V2L2A <
Σ 11x()) mi 122 xED2
Thus, (a) is X + 0 m )- smooth. By choosing the second-stage learning rate = M(m1), after T = Ö(>-1) -oolo(y (d, n, mi 12, (19l)²) steps, we can reach an iterate à = (7) so that (2( S (2(at) and ||ál|2 S ||a*||2-
Denoting Ó = (å, (/1)), V), it holds that
Σ - [xxx s Cc)( I VL2(ar)
xED=
Thus, we have
L = (==;((((:))) < V
- (x)))
122 xED
pp/2
2
pt3/2d
eLr ki log² d d1/6
||gllza
+
+
Amin
H
m2
1m222
Vn
p+1
p-kk
17
B
lop
mi
k=
52

---

Published as a conference paper at ICLR 2025
Here à satisfies
|al m1
|
Σ
k
B
oop
mi
k=
We assume |aal < my Ba where B, satisfies
2
P Σ k=
p=z op
r
Bound L2 To bound T2, we rely standard Rademacher complixity analysis. The following lemma provides an upper bound on the Rademacher complixity of the random feature model.
Lemma 35. Let F = US = (a, W b(), 1|ll| 5
Vmn Recall the empirical
Rademacher complexity of F as
(R( E
sup
Σ oxf((x))
JEF n2
Here the dataset (x1, (2,,x,,....
Xna] D22 Then with high probability, have
B. 2122
R. (F) S
The proof is provided in Appendix E.2. Since the e, is 1-Lipschitz, by standard Rademacher com- plexity analysis, we have that with high probability that
L2 = Exer (x(x;)) ((x;0)- ((x) - 1 Σ l (f(x:)) (x)) 112 xE
RR. + +
112
B2
T
T12
n12
Bound L.g Finally, we relate the truncated loss &, to the T. population loss.
Lemma 36. By letting 1(max(1P, Ba)), with high probability over ê, we have
La Ex -)(x0) [x) Ex ( so B1nimm
Here recall that n2- The proof is provided in Appendix E.2.
53

---

Published as a conference paper at ICLR 2025
Put the loss together By invoking the upper bound of L1, L2 and Lg and plugging the values of T. m7
|BB"llop Amin (H). Ba. L and | we have
Ex [(x:0 M(x)]
L1 + L2 1 LB
B²
1
(1((
-
+T
+
+o
nnngmimad
n12
122
xED2
rp/2
1222
ud
(pt3/2d
LLr
log" d
11
12
Amin
H
m2
m2
d1/6
iP+1
|1g1|22
p
B2 +
|B
r
V
+ T
m1
112
112
A=U
rp/2
1d3
p+3/2d
1
log²
d
Amin (H)
m2
d//6
6p 6ptlyp/2 Riz
(r-14 V Xmin (H))
m1
2p 1/4 1p 1pptlyp/22 2
(H))
min
n22
TP
d6
drr+1
p K1
+
+
min m1)
12
d1/6
The proof is complete.
E.2 OMITTED PROOFS IN APPENDIX E.
Proof of Lemma 35. Given 0 (a, W b(1), V). since we can write
fx(x) = a (x),, where (x) Vec (mi'or 01 (na!" (0) (w,, h (x)) + 6.")).
By Proposition and our choice of 1 to ensure (n/w,, (x)) < with high probability for any D2, we obtain that for any E [mi] and E D2.
Iaa (0) (w:, (x)) + (000) < a (0) n(w,, h (x)) + (.90 1
Thus, [[ 2 S m. by the standard linear Rademacher bound, with high probability, the em- pirical Rademacher complexity is upper bounded by
mm
B. TT2
Rr S
[x))) <
122
xED=
The proof is complete.
Proof of Lemma 36. We can bound the difference between &, and L1 loss by
[x: 0) (xx) - Ex [e (f(x; - - -) (xx) ]
Ex
Ex f( (xx; 0) - fr (x) 1 [(xx 8) - f*t (x) > r}]
Ex
| ((x:)))-00 /000
Pr [(x:0) (xx) 27]
(1(x; ((x)) Pr [(x;0) > 7/2 (x r/2]
Ex
(36)
54

---

Published as a conference paper at ICLR 2025
Recall that we can write
0x: = â (x), where (x()) Vec (mi'or 01 (na,
(0) (w,, (w,
(x)) + b)
By following the proof of Lemma 35 and applying Proposition for one single sample point x (instead of the whole set D2), we know with high probability over V, and Dj (we denote this event by E1), we have for any E [m1].
nww, (x)) S
holds with high probability x. Also, since for any i E [ma]] w; ( mg Amn), know |wall2 S /mzl for any i E [m1] with high probability. We denote this joint event on W1, W2.... Wm, by Ez. Thus, conditional on events E₁ and E2, have
||ll Vmn
f(x:0) 03 E
with high probability on x.
We denote this conditional event by Eg,1. Moreover, since Ct m2 do we have
|
mi
[(x: 8)
7
1|5ll2
K mg X2, x ) (0) (xx)
+3
m
T12
2
m27
||l|
C
vmzCa +3
Vm2dd ||ál|2
holds for any x. Moreover, since f*)x a degree-2p polynomial of x, we know by Lemma
8, with probability least 1 exp(-i), have | I Cyp for sufficiently large Cj > 0. Besides, have Ex [f'(x)?] E 1. Altogether, conditional on E₁ and Ez, by choosing
C" max(i?, m 1/2 [|áll2) = [(mmx((1 Ba) for some sufficiently large C, we have
Ex [xx; 8)2 + (x))) [[ Pr [(x; 0) > 7/2] + Pr Pr > 7/2]
2 m2d 12 |all + 1) (Pr[x & Er,I + Pr[x @ Em.2))
$
Pmimmgning
The last inequality holds because of the definition of high probability events and the choice of with C' Jog(dm/m2n) 12) for sufficiently large C. Plugging the result into (36) concludes our proof.
55


Reviews:
[{"review_id": "ovVkxBERK1", "reviewer": "Reviewer_jyeu", "summary": "The authors study layerwise training on a three layer neural network trained on a more general target function with a more involved \"hierarchical feature\" structure. They show sample complexity improvements over the classic kernel methods.", "questions": "please see the weakness section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- Feature learning is one of the most fundamental problems in deep learning theory.\n\n- The proof seem to be correct. Nice proof sketches are provided in the main text.", "weaknesses": "- I don't find the results of this paper particularly interesting. I would have easily guessed that these results should be obtainable for the case where r \\neq 1, after looking at Nichani et al., 2023; Wang et al., 2023. See also my other comments below:\n\n- The results of the paper are not tight; i.e., the sample complexity of O(d^4). Essentially, one would expect a sample complexity of $O(d^2)$, because there are $O(d^2)$ parameters to be learned in this setting. In particular, a more interesting study will consider $n \\asymp d^2$ and provide a precise characterization of the limiting test error. In the single index case, this is similar to the benefit show by Ba et al. 2022 (and others) over the initial results of Damian et al 2022.\n\n- In assumption 2, it is assumed that the Hessian has eigenvalues that are bounded below by a certain quantity. Can this be translated into a condition on the target function in terms of some quantity like information-exponent or leap-index? Right now, the assumption seems very artificial.\n\n- What do the authors mean by \" Moreover, when the entries of Ai are sampled i.i.d., the assumption is satisfied with high probability by Wigner’s semicircle law\" in line 157?\n\n- In the layer-wise training, the inner most layer is not trained and is kept at initialization; i.e., random features.\n\n- Assumption 4 on the activation function is completely unrealistic and is only an artifact of the proof.\n\n- What is the role of the update on the bias. Some prior work in the area consider neural networks without bias (e.g., the papers I list below). Can the authors comment on this please?\n\n- The discussion of prior work on \"Learning Linear Features\" is very limited. There are many more results on this topic, studying feature in this case with detail. The authors should consider citing and discussing these papers (alongside any other paper on the topic that I might have missed here).\n\n [1] Dandi et al., The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents\n\n [2] Moniri et al., A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks\n\n [3] Cui et al., Asymptotics of feature learning in two-layer networks after one gradient-step\n\n [4] Dandi et al., A Random Matrix Theory Perspective on the Spectrum of Learned Features and Asymptotic Generalization Capabilities\n\n [5] Wang et al, Nonlinear spiked covariance matrices and signal propagation in deep neural networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZXxPwxBzlK", "reviewer": "Reviewer_bUVE", "summary": "The paper studies hierarchical feature learning in fully-connected neural networks when trained with gradient descent. In particular, the authors consider a nonlinear (quadratic) multi-index model with $f^*(x) = g^*(x^\\top A_1 x, \\ldots, x^\\top A_r x)$ for $x \\in \\mathcal{S}^{d-1}$ and $r\\ll d$. Under approximate orthogonality, well-balancedness conditions of the quadratic features, and additional well-conditioned assumptions on the link function, the authors show that a layer-wise gradient descent method on a three-layer neural network (first layer is never trained, and activations are carefully chosen) can learn $f^*$ with roughly $\\tilde{O}(d^4 + d^2r^{p+1})$ (ignoring $\\kappa$) to irreducible error $O(r^{p+2}/d^{1/6})$, and approximately recover the $r$-dimensional subspace to allow for transfer learning to a different link function.", "questions": "- Why did the authors consider only uniform over the sphere compared to Gaussian, which is more commonly studied in this literature? Do the results extend to more general distributions?\n- There is no lower bound on the optimality of the dependence on $d$ and $\\kappa$ in the bound. It should be possible to get $d^2$ potentially. What dependence on $\\kappa$ is optimal?\n- Figure 1 is not very easy to understand. Might be helpful to add more explanation to the caption and label the different blocks/terms. It took me a while to parse it and I did not find it very helpful to understand the proof. \n- The irreducible error term of $O(r^{p+2}/d^{1/6})$ is not discussed. This is much larger than the $\\frac{\\kappa^{1/3}}{d^{1/6}}$ term in Niching et al. (2022) and requires $r$ to be much smaller than $d$. Is this necessary?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper studies an interesting model of feature learning, generalizing prior work that consider either linear multi-index settings or non-linear single-index models.\n- The authors show that gradient descent recovers the subspace which allows them to give transfer learning results for any link function.\n- To prove the results, the authors prove a new universality theorem for multiple-dimensions, generalizing previous known results. Though the proof seems to follow from existing tools in high dimensional probability, it seems novel.", "weaknesses": "- The paper makes several assumptions and tweaks that seem to be done to make the proofs easier: layer-wise training, removing lower order terms in activation functions, taking almost orthonormal equally important features, etc. Most of these are standard is deep learning theory, but the activation assumption seems more artificial, and tailored to make the analysis easier. Some assumptions (Assumptions 2 and 3) can potentially be relaxed if the goal is only getting low loss and not subspace recovery.\n- Compared to prior work (Nichani et al. 2022) that was for general hierarchal functions, this paper is narrow in its scope and relies heavily on tools already available in this prior work, while requiring additional assumptions. I do understand that the prior work did not look at multiple quadratic features, but limiting to quadratic seems like a limitation of the techniques.\n- The experiments are limited to following the theoretical assumptions precisely. This is not very insightful. If the experiments are to remain synthetic, then it would be worthwhile seeing how the scaling with different parameters in the bounds shows up in experiments. Also, how much can you relax the assumptions, and still get the same empirical results. It would also be worthwhile to come up with a real-world dataset where this quadratic model makes sense, and gains of three vs two layers can be shown.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WXnyl8y4iZ", "reviewer": "Reviewer_okDB", "summary": "In this work, the authors studies theoretically the setting of using three-layer neural networks to learn a type of \"multi-index quadratic models\" under a uniform data distribution on the sphere. With suitable activation functions and layer-wise training, the model is proved to be able to learn both the output function and the quadratic features with a sample complexity that scales as $d^4$.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The type of model studied in this paper generalizes the one considered by Nichani et al. (2023) by allowing multiple quadratic features instead of just one in the first layer. To prove the learning guarantee, the authors introduced a nice technical result on the distribution of multiple nonlinear functions of Gaussian inputs, which is interesting mathematically and novel to my knowledge. The theoretical results appear to be sound though they have not been thoroughly checked.", "weaknesses": "**Beyond the mathematical and technical**: It could be argued that the series of assumptions on the layer-wise learning algorithm, the activation functions, the data distribution, etc. are quite specific and restrictive, and the contribution of this work (relative to e.g. Nichani et al., 2023 and Wang et al., 2023, which share much of a common analytical framework) resides more in the improved mathematical results than advancements in new learning theory insights.\n\n**Some confusion regarding the proof of Theorem 1**: The authors' description of the second stage in the proof of Theorem 1 gives the impression that it breaks down into two parts: 1) showing the existence of an $a^*$ achieving small empirical loss, and 2) using Rademacher complexity analysis to bound the generalization error of $a^*$. But they alone wouldn't be sufficient since we know only the existence of the $a^*$ but not its learnability by the proposed algorithm. And indeed, the actual proof steps taken by the authors is different (and correctly so), where they use the strong convexity of the loss function to show that an alternative solution $\\hat{a}$ can be learned by gradient descent and achieves low empirical loss as well as generalization error. Even though $a^*$ still makes its appearance in the proof argument in Section D.1.2, I suspect that it is not truly necessary and only complicates the proof. If this assessment is correct, it would imply that Proposition 2 is redundant for the proof of Theorem 1, which should be commented on in order to avoid confusion.\n\nMinor typos:\n1) In Section 1.1 on Page 2, \"outer layer can achieves\" --> \"outer layer can achieve\".\n2) In Section 4.1 on Page 7, \"trained presentations representations\" --> \"trained representations\".\n3) The links referencing Proposition 6 are misplaced.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2FjKT3LSnZ", "reviewer": "Reviewer_BXau", "summary": "The authors prove that three-layer networks can learn nonlinear, multi-index targets more efficiently than kernel methods. Quite a bit of literature already exists on 2- and 3-layer networks learning single-index targets of linear and non-linear features, and also multi-index targets of linear features. To the best of my knowledge this is the first set of results for learning multiple, non-linear features. \n\nThe analysis is valid of an algorithm that does a single step on the first layer weights, W, and then completely trains the rest of the network. The authors use this two-stage training methodology to also prove a transfer learning result, where the link function is changed. \n\nFinally the authors provide some empirical evidence that supports their theory. In particular, their approach appears to improve sample complexity over a random feature model. They also support their findings regarding efficient transfer learning and the reconstruction of quadratic features.", "questions": "1. The “balanced features” part of Assumption 1, where the operator norm of Ai satisfies a particular bound. This bound essentially tells us that each feature cannot depend on a sparse subset of inputs. Can you discuss the implications of that assumption and pertinence to real distributions? \n2. Similarly, can you discuss the implications of Assumption 4, especially the Gegenbauer expansion? Activation functions used in practice are not even, and many of them are actually odd.\n3. Do you expect that the d^4 dependence on dimension is necessary for quadratic features?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a strong submission in most respects: relevance, rigor, and clarity. I have minor comments and questions below.", "weaknesses": "1. Some of the assumptions might be too restrictive for the results to be relevant for practice (eg parts of Assumption 1, and Assumption 4). \n2. Algorithm 1 does but a single update on W. Then parameter a is trained on multiple steps. While this style of analysis is common in the area, there already exist articles that include analysis of many-step gradient descent for all parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zn21TBp7Ab", "reviewer": "Reviewer_Xpr2", "summary": "This work studies the problem of learning functions $f^* : \\mathbb{R}^d \\to \\mathbb{R}$ of the form $f^* = g^* \\circ p$ where $p : \\mathbb{R}^d \\to \\mathbb{R}^r$ returns $r$ quadratic features of the input, and $g^* : \\mathbb{R}^r \\to \\mathbb{R}$ is a polynomial of degree $p$. By analyzing a gradient-based algorithm for training a three-layer neural network with specific choices of activation, the authors prove that a sample complexity of $\\tilde{O}(d^4)$ is sufficient for learning this function, as opposed to the $\\Omega(d^{2p})$ sample complexity lower bound for learning with rotationally invariant kernels.", "questions": "* What is the main technical challenge in directly generalizing the result of Nichani et al., 2023 to the setting of having multiple non-linear features? Specifically, if $f^* = g \\circ p$ where $p = (p_1,...,p_r)$ and each $p_j$ is a degree $k$ polynomial of $x$, and $g$ is a degree $q$ polynomial, is it possible to argue that in the regime of constant $r$, the sample complexity of learning $f^*$ with three-layer neural networks is $d^k$, as opposed to $d^{kq}$ of kernel methods?\n\n* On the discussion of learning linear features, a relevant work may be [1] that considers learning general multi-index models with two-layer neural networks via a noisy gradient descent algorithm. For single-index models, while the authors have included works in the CSQ framework, there are more recent advances on gradient-based learning with two-layer neural networks that go beyond the CSQ framework [2, 3], or consider structured rather than isotropic input covariance [4, 5].\n\n---\n[1] A. Mousavi-Hosseini et al. \"Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics\". arXiv 2024.\n\n[2] Y. Dandi et al. \"The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents\". ICML 2024.\n\n[3] J. D. Lee et al. \"Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit\". arXiv 2024.\n\n[4] A. Mousavi-Hosseini et al. \"Gradient-Based Feature Learning under Structured Data\". NeurIPS 2023.\n\n[5] J. Ba et al. \"Learning in the presence of low-dimensional structure: a spiked random matrix perspective\". NeurIPS 2023", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Presenting a guarantee for learning a polynomial of multiple quadratic features using three-layer neural networks is a new direction in the literature which extends beyond the guarantees for learning multi-index models with two-layer neural networks and learning polynomials of a single non-linear feature which existed in prior works. More broadly, this work contributes to finding settings in which networks deeper than two layers are natural for learning certain function classes, which is an area of particular interest for the community.", "weaknesses": "* The feature learning regime considered in this paper is limited in the sense that the first layer weights are fixed at random and not trained. Training these weights could lead to learning more complex hierarchies, such as hierarchies of multi-index models, although admittedly challenging for a theoretical analysis.\n\n* The choice of activation for the second layer is rather specific, and choosing a different activation for each layer seems a bit odd when compared with standard three-layer neural networks. However, I believe there might be a way to choose the same activation as my impression is that $\\sigma_1$ only has a constant part and no $Q_1$ part, thus one only needs to introduce a fixed bias to the output of the network. If this interpretation is correct this might be worth mentioning in the main text.\n\n* While the statistical model is naturally suitable for three-layer networks, I could not find any lower bounds in this submission which explain that learning this function with a two-layer neural network is provably less efficient (in terms of samples or computation). I am wondering if the authors could use the results of prior works to show a statistical or computational lower bound holding in this setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)