Towards Understanding Hierarchical Learning: Benefits of Neural Representations

Introduction

Deep neural networks have been empirically observed to be more powerful than their shallow counterparts on a variety of machine learning tasks . For example, on the ImageNet classification task, a $152$-layer residual network can achieve $8$%-$10$% better top-$1$ accuracy than a shallower $18$-layer ResNet . A widely held belief on why depth helps is that deep neural networks are able to perform efficient hierarchical learning, in which the layers learn representations that are increasingly useful for the present task. Such a hierarchical learning ability has been further leveraged in transfer learning. For example, and show that by combining with additional task-specific layers, the bottom layers of pre-trained neural networks for image classification and language modeling can be naturally transferred to other related tasks and achieve significantly improved performance.

Despite significant empirical evidence, we are in the lack of practical theory for understanding the hierarchical learning abilities of deep neural networks. Classical approximation theory has established a line of “depth separation” results which show that deep networks are able to approximate certain functions with much fewer parameters than shallow networks . These work often manipulates the network parameters in potentially pathological ways, and it is unclear whether the resulting networks can be efficiently found through gradient-based optimization. A more recent line of work shows that overparametrized deep networks can be provably optimized and generalize as well as the so-called Neural Tangent Kernels (NTKs) . However, these results do not take the hierarchical structure of the neural networks into account, and cannot justify any advantage of deep architectures. More recently, show that some NTK models of deep networks are actually degenerate, and their generalization performance are no better than those associated with shallow networks.

In this paper, we provide a new persepctive for understanding hierarchical learning through studying intermediate neural representations—that is, feeding fixed, randomly initialized neural networks as a representation function (feature map) into another trainable model. The prototypical model we consider is a wide two-layer neural network taking a representation function $\mathbf{h}$ as the input, that is,

To demonstrate the importance of the representation function $\mathbf{h}$, we investigate the sample complexity for learning certain target functions using model (1). This is a fine-grained measure of the power of $\mathbf{h}$ compared with other notions such as approximation ability. Indeed, we expect $f_{\mathbf{W}}$ to be able to approximate any “regular” (e.g. Lipschitz) function of $\mathbf{x}$, whenever we use a non-degenerate $\mathbf{h}$ and a sufficiently large width $m$. However, different choices of $\mathbf{h}$ can result in different ways (for the trainable two-layer network) to approximate the same target function, thereby leading to different sample complexity guarantees. We will specifically focus on understanding when learning with the neural representation $\mathbf{h}(\mathbf{x})=\sigma(\mathbf{V}\mathbf{x}+\mathbf{b})$ is more sample efficient than learning with the raw input $\mathbf{h}(\mathbf{x})=\mathbf{x}$, which is a sensible baseline for capturing the benefits of representations.

As the optimization and generalization properties of a general two-layer network can be rather elusive, we consider more optimization aware versions of the prototype (1)—we replace the trainable two-layer network in $f_{\mathbf{W}}$ by tractable alternatives such as its linearized model (also known as “lazy training” in ) or quadratic Taylor model :

When $\mathbf{h}$ is the raw input (NTK-Raw, Quad-Raw), these are models with concrete convergence and generalization guarantees, and can approximate the training of the full two-layer network in appropriate infinite-width limits (e.g. ). However, for learning with other representation functions, these models are less understood. The goal of this paper is to provide a quantitative understanding of these models, in particular when $\mathbf{h}$ is a one-hidden-layer neural network (NTK-Neural, Quad-Neural), in terms of their convergence, generalization, and sample complexities of learning.

The contributions of this paper are summarized as follows:

We show that the Quad-$\mathbf{h}$ model has a benign optimization landscape, and prove generalization error bounds with a precise dependence on the norm of the features and weight matrices, as well as the conditioning of the empirical covariance matrix of the features (Section 3).

We study sample complexities of learning when the representation is chosen as a one-hidden-layer neural network (Quad-Neural model, Section 4). For achieving a small excess risk against a low-rank degree-$p$ polynomial, we show that the Quad-Neural model requires $\widetilde{O}(d^{\lceil p/2\rceil})$ samples. When $p$ is large, this is significantly better than the best known $\widetilde{O}(d^{p-1})$ upper bound for the Quad-Raw model, demonstrating the benefits of neural representations.

When the trainable network is instead a linearized model (or an NTK), we present a lower bound showing that neural representations are provably not beneficial: in a certain infinite-width limit, the NTK-Neural model requires at least $\Omega(d^{p})$ samples for learning a degree-$p$ polynomial (Section 5). Since $O(d^{p})$ samples also suffice for learning with the NTK-Raw model, this shows that neural representations are not beneficial when fed into a linearized neural network.

We present the problem setup and algorithms in Section 2, review related work in Section 6, and provide conclusions as well as acknowledgments in Section 7.

Notations

Preliminaries

Given dataset $S_{n}$, we define the empirical risk of a predictor $f$ as

Model, regularization, and representation

We consider the case where $f$ is either the linearized or the quadratic Taylor model of a wide two-layer network that takes a fixed representation function as the input:

For the Quad-$\mathbf{h}$ model, we add a regularizer to the risk so as to encourage $\mathbf{W}$ to have low norm. We use the regularizer $\left\lVert\mathbf{W}\right\rVert_{2,4}^{4}=\sum_{r=1}^{m}\left\lVert\mathbf{w}_{r}\right\rVert_{2}^{4}$, and consider minimizing the regularized empirical risk

Connection to a three-layer model

It is worth noticing that when $\mathbf{h}$ is indeed a neural network, say $\mathbf{h}(\mathbf{x})=\sigma(\mathbf{V}\mathbf{x})$ (omitting bias for simplicity), our NTK-$\mathbf{h}$ and Quad-$\mathbf{h}$ models are closely related to the Taylor expansion of a three-layer network

Indeed, the {NTK-$\mathbf{h}$, Quad-$\mathbf{h}$} models correspond to the {linear, quadratic} Taylor expansion of the above network over $\mathbf{W}$, and is thus a part of the full Taylor expansion of the three-layer network. By studying these Taylor models, we gain understandings about how deep networks use its intermediate representation functions, which is lacking in existing work on Taylorized models.

Quadratic model with representations

We begin by studying the (non-convex) optimization landscape as well as the generalization properties of the model (Quad-$\mathbf{h}$), providing insights on what can be a good representation $\mathbf{h}$ for such a model.

When $\mathbf{h}(\mathbf{x})=\mathbf{x}$ is the raw input, model (Quad-$\mathbf{h}$) becomes

which is the quadratic Taylor model of a wide two-layer neural network. This model is analyzed by Bai and Lee who show that (1) the (regularized) risk $\widehat{\mathcal{R}}_{\lambda}(f_{\mathbf{W}})$ enjoys a nice optimization landscape despite being non-convex, and (2) the generalization gap of the model $f^{Q}_{\mathbf{W}}$ is controlled by $\left\lVert\mathbf{W}\right\rVert_{2,4}$ as well as $\|\frac{1}{n}\sum_{i\in[n]}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}\|_{\rm op}$. Building on these results, show that learning low-rank polynomials with (Quad-Raw) achieves a better sample complexity than with the NTK. Besides the theoretical investigation, empirically show that (Quad-Raw) model also approximates the training trajectories of standard neural networks better than the linearized model.

General case

We analyze optimization landscape and establish generalization guarantees when $\mathbf{h}$ is a general representation function, extending the results in . We make the following assumption:

(Optimization) Given any $\epsilon>0$, $\tau=\Theta(1)$, and some radius $B_{w,\star}>0$, suppose the width $m\geq\widetilde{O}(B_{h}^{4}B_{w,\star}^{4}\epsilon^{-1})$ and we choose a proper regularization coefficient $\lambda>0$. Then any second-order stationary point $\mathbf{W}$ is a second-order stationary point (SOSP) of a twice-differentiable loss $L(\mathbf{W})$ if $\nabla L(\mathbf{W})={\bm{0}}$ and $\nabla^{2}L(\mathbf{W})\succeq{\bm{0}}$. (SOSP) $\widehat{\mathbf{W}}$ of the regularized risk $\widehat{\mathcal{R}}_{\lambda}(f^{Q}_{\mathbf{W}})$ satisfies $\|\widehat{\mathbf{W}}\|_{2,4}\leq O(B_{w,\star})$, and achieves

(Generalization) For any radius $B_{w}>0$, we have with high probability (over $(\mathbf{a},\mathbf{W}_{0})$) that

Efficient optimization; role of feature isotropicity

Theorem 1 has two main implications: (1) With a sufficiently large width, any SOSP of the regularized risk $\widehat{\mathcal{R}}(f^{Q}_{\mathbf{W}})$ achieves risk close to the optimum in a certain norm ball, and has controlled norm itself. Therefore, escaping-saddle type algorithms such as noisy SGD that can efficiently find SOSPs can also efficiently find these near global minima. (2) The generalization gap is controlled by $M_{h,{\rm op}}$, which involves the operator norm of $\frac{1}{n}\sum_{i=1}^{n}\mathbf{h}(\mathbf{x}_{i})\mathbf{h}(\mathbf{x}_{i})^{\top}$. It is thus beneficial if our representation $\mathbf{h}(\mathbf{x})$ is (approximately) isotropic, so that $M_{h,{\rm op}}\asymp O(1/\sqrt{D})$, which is much lower than its naive upper bound 1. This will be a key insight for designing our neural representations in Section 4. The proof of Theorem 1 can be found in Appendix A.

Learning with neural representations

We now develop theories for learning with neural representations, where we choose $\mathbf{h}$ to be a wide one-hidden-layer neural network.

We consider a fixed, randomly initialized one-hidden-layer neural network:

where $\mathbf{v}_{i}\stackrel{{\scriptstyle\rm iid}}{{\sim}}{\sf N}(0,\mathbf{I}_{d})$ and $b_{i}\stackrel{{\scriptstyle\rm iid}}{{\sim}}{\sf N}(0,1)$ are the weights. Throughout this section we will use the indicator activation $\sigma(t)=\mathds{1}\left\{t\geq 0\right\}$. We will also choose $\phi(t)={\rm relu}(t)^{2}/2$ so that $\phi^{\prime\prime}(t)=\mathds{1}\left\{t\geq 0\right\}$ as well.We can use a non-smooth $\sigma$ since $(\mathbf{V},\mathbf{b})$ are not trained. Our results can be extended to the situation where $\sigma$ or $\phi^{\prime\prime}$ is the relu activation as well.

We define the representation function $\mathbf{h}(\mathbf{x})$ as the whitened version of $\mathbf{g}(\mathbf{x})$:

We summarize our overall learning algorithm (with the neural representation) in Algorithm 1.

2 Learning low-rank polynomials with neural representations

We now study the sample complexity of Algorithm 1 to achieve low excess test risk compared with the best low-rank degree-$p$ polynomial, that is, sum of polynomials of the form $(\bm{\beta}^{\top}\mathbf{x})^{p}$. This setting has been considered in a variety of prior work on learning polynomials as well as analyses of wide neural networks .

We need the following additional assumption on the random features.

where $q(z)$ denotes a polynomial in $z$ and its degree is denoted as ${\rm deg}(q)$. For general distributions of $\mathbf{x}$, we show Assumption 2 still holds under certain moment conditions on the distribution of $\mathbf{x}$ (see the formal statement and proof of both results in Appendix B).

We focus on low-rank polynomials of the form

We state our main result for the Quad-Neural model to achieve low excess risk over such functions.

Suppose Assumption 2 holds, and there exists some $f_{\star}$ of the form (5) that achieves low risk: $\mathcal{R}(f_{\star})\leq\mathsf{OPT}$. Then for any $\epsilon,\delta\in(0,1)$ and $\tau=\Theta(1)$, choosing

$n_{0}=\widetilde{O}(D\delta^{-2})$, and a proper $\lambda>0$, Algorithm 1 achieves the following guarantee: with probability at least $1-\delta$ over the randomness of data and initialization, any second-order stationary point $\widehat{\mathbf{W}}$ of $\widehat{\mathcal{R}}_{\lambda}(f^{Q}_{\mathbf{W}})$ satisfies

In particular, for any $\epsilon>0$, we can achieve $\mathcal{R}(f^{Q}_{\widehat{\mathbf{W}}})\leq(1+\tau)\mathsf{OPT}+2\epsilon$ with sample complexity

According to Theorem 2, Quad-Neural can learn polynomials of any degree by doing the following: (1) Choose a sufficiently large $D$, so that the neural representations are expressive enough; (2) Choose a large width $m$ in the quadratic model so as to enable a nice optimization landscape, where such $m$ only appears logarithmically in generalization error (Theorem 1).

Improved dimension dependence over Quad-Raw and NTK-Raw

In comparison, the sample complexity for learning with the Quad-Raw (quadratic neural network with the raw input) is

(see, e.g. [8, Thm 7]). Therefore, Theorem 2 shows that neural representations can significantly improve the sample complexity over the raw input, when fed into a quadratic Taylor model.

Overview of techniques

At a high level, the improved sample complexity achieved in Theorem 2 is due to the flexibility of the neural representation: the Quad-$\mathbf{h}$ model can express polynomials hierarchically, using weight matrices with much smaller norms than that of a shallow learner such as the Quad-Raw model. This lower norm in turn translates to a better generalization bound (according to Theorem 1) and an improved sample complexity. We sketch the main arguments here, and leave the complete proof to Appendix C.

NTK with neural representations: a lower bound

In this section, we show that neural representations may not be beneficial over raw inputs when the trainable network is a linearized neural network through presenting a sample complexity lower bound for this method in the infinite width limit.

More concretely, we consider NTK-Neural, which learns a model $f^{L}_{\mathbf{W}}$ of the form

(see e.g. for the derivation). Motivated by this, we consider kernel predictors of the form

as a proxy for (NTK-Neural), where $\left\lVert\cdot\right\rVert_{H_{\infty}}^{2}$ denotes the RKHS (Reproducing Kernel Hilbert Space) norm associated with kernel $H_{\infty}$. This set of predictors is a reliable proxy for the (NTK-Neural) method: for example, taking $\lambda\to 0_{+}$, it recovers the solution found by gradient descent (with a small stepsize) on the top layer of a wide three-layer network .

We now present a lower bound for the predictor $\widehat{f}_{\lambda}$, adapted from [27, Theorem 3].

that is, any predictor of the form (8) will not perform much better than the trivial zero predictor.

No improvement over NTK-Raw; benefits of neural representations

Related work

Approximation theory and depth separation. Extensive efforts have been made on the expressivity of neural networks and the benefits of increased depth. Two separate focuses were pursued: 1) Universal approximation theory for approximating dense function classes, e.g., Sobolev and squared integrable functions ; 2) depth separation theory demonstrating the benefits of increased depth on expressing certain structured functions, e.g., saw-tooth functions . More recently, the recent work merged the two focuses by studying unbounded-depth ReLU networks for approximating Sobolev functions. In all these work, the network parameters are constructed in potentially weird ways, and it is unclear whether such networks can be efficiently found using gradient-based optimization.

A growing body of recent work show the connection between gradient descent on the full network and the Neural Tangent Kernel (NTK) , from which one can prove concrete results about neural network training and generalization . Despite such connections, these results only show that neural networks are as powerful as shallow learners such as kernels. The gap between such shallow learners and the full neural network has been established in theory by and observed in practice . Higher-order expansions of the {network, training dynamics} such as Taylorized Training and the Neural Tangent Hierarchy have been recently proposed towards closing this gap. Finally, recent work by Allen-Zhu and Li shows that there exists a class of polynomials that can be efficiently learned by a deep network but not any “non-hiearchical” learners such as kernel methods or neural tangent kernels, thereby sheding light on how representations are learned hierarchically.

Learning low-rank polynomials in high dimension

In and , the authors propose a tensor unfolding algorithm to estimate a rank $k$ order $p$ tensor with $(d)^{p/2}k$ samples. Under Gaussian input data, propose a Grassmanian manifold optimization algorithm with spectral initialization to estimate a polynomial over $k$-dimensional subspace of variables of degree $p$ with $O_{k,p}(d\log^{d}p)$ samples, where $O_{k,p}$ suppresses unknown (super)-exponential dependence on $k$ and $p$. However, these methods explicitly use knowledge about the data distribution. Neural networks can often learn polynomials in distribution-free ways. show that wide two-layer networks that simulate an NTK require $\widetilde{O}(d^{p})$ samples to learn a degree-$p$ polynomial. show that $\Omega(d^{p})$ samples is also asymptotically necessary for any rotationally invariant kernel. show that a randomized wide two-layer network requires $\widetilde{O}(d^{p-1})$ samples instead by coupling it with the quadratic Taylor model. Our algorithm belongs to this class of distribution-free methods, but achieve an improved sample complexity when the distribution satisfies a mild condition.

Conclusion

This paper provides theoretical results on the benefits of neural representations in deep learning. We show that using a neural network as a representation function can achieve improved sample complexity over the raw input in a neural quadratic model, and also show such a gain is not present if the model is instead linearized. We believe these results provide new understandings to hiearchical learning in deep neural networks. For future work, it would be of interest to study whether deeper representation functions are even more beneficial than shallower ones, or what happens when the representation is fine-tuned together with the trainable network.

Acknowledgment

We thank the anonymous reviewers for the suggestions. We thank Song Mei for the discussions about the concentration of long-tailed covariance matrices. JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, and NSF CCF 2002272.

References

Appendix A Proofs for Section 3

We first derive the gradient and Hessian of empirical risk $\widehat{\mathcal{R}}(f^{Q}_{\mathbf{W}})$, which will be used throughout the rest of the proof. For a better presentation, we denote $\left\langle\cdot,\cdot\right\rangle$ as inner product and

We compute the gradient and Hessian of $\widehat{\mathcal{R}}(f^{Q}_{\mathbf{W}})$ along a given direction $\mathbf{W}_{\star}$.

We denote $\widehat{\mathcal{D}}$ as the empirical data distribution, and bound I and II separately.

where the last step used Cauchy-Schwarz on ${\left\{\left\lVert\mathbf{w}_{r}\right\rVert_{2}\right\}}$ and ${\left\{\left\lVert\mathbf{w}_{\star,r}\right\rVert_{2}\right\}}$, and the constant $C$ is the uniform upper bound on $\phi^{\prime\prime}$. Putting terms I and II together, we have

where $\lambda_{0}$ is a constant to be determined.

We argue that any second order stationary point $\widehat{\mathbf{W}}$ has to satisfy $\|\widehat{\mathbf{W}}\|_{2,4}=O(B_{w,\star})$. We have for any $\mathbf{W}$ that

Combining with the fact that $\left\langle\nabla_{\mathbf{W}}(\left\lVert\mathbf{W}\right\rVert_{2,4}^{4}),\mathbf{W}\right\rangle=4\left\lVert\mathbf{W}\right\rVert_{2,4}^{4}$, we have simultaneously for all $\mathbf{W}$ that

Therefore we see that any stationary point $\mathbf{W}$ has to satisfy

we get $36\lambda B_{w,\star}^{4}=2\tau M+\epsilon$. The Hessian of $\widehat{\mathcal{R}}_{\lambda}(f^{Q}_{\mathbf{W}})$ along direction $\mathbf{W}_{\star}$ is

We used the fact $12ab\leq a^{2}+36b^{2}$. For a second order-stationary point $\widehat{\mathbf{W}}$ of $\widehat{\mathcal{R}}_{\lambda}(f^{Q}_{\mathbf{W}})$, its gradient vanishes and the Hessian is possitive definite. Therefore, we have

We choose $m=\epsilon^{-1}(2\lambda_{0})^{-1/2}C^{2}B_{h}^{4}B_{w,\star}^{4}\geq\epsilon^{-1}C^{2}B_{h}^{4}\|\widehat{\mathbf{W}}\|_{2,4}^{2}\left\lVert\mathbf{W}_{\star}\right\rVert_{2,4}^{2}$ and the above inequality implies

A.2 Proof of Generalization in Theorem 1

where $\xi$ is i.i.d. Rademacher random variables. The above Rademacher complexity can be bounded using the contraction theorem [51, Chapter 5]:

where the last step used the power mean (or Cauchy-Schwarz) inequality on ${\left\{\left\lVert\mathbf{w}_{r}\right\rVert_{2}\right\}}$ and $\left\lVert\cdot\right\rVert_{*}$ denotes the matrix nuclear norm (sum of singular values). Now it only remains to bound the expected max operator norm above. We apply the matrix concentration lemma Bai and Lee [8, Lemma 8] to deduce that

Appendix B Results on feature covariance

We first present a Lemma for relating the covariance of nonlinear random features to the covariance of certain polynomial bases, adapted from [27, Proposition 2].

For any $k\geq 0$, suppose $D\leq O(d^{k+1-\delta})$ for some $\delta>0$, then we have with high probability as $d\to\infty$ that

where $(\cdot)^{\odot k}$ is the Hadamard product: $(\mathbf{A}^{\odot k})_{ij}=\mathbf{A}_{ij}^{k}$.

B.2 Lower bound on population covariance

then we have $\lambda_{\min}(\bm{\Sigma})\geq c>0$ with high probability as $d\to\infty$, where $c=c_{K}$ is a constant that depends on $K$ (and the indicator activation) but not $d$.

This falls into the setting of Lemma 1(b), applying which implies that with high probability (as $d\to\infty$) we have

as the indicator function is not a polynomial of any degree (so that its $L_{2}$ projection onto polynomials of degree $\leq K$ is not itself for any $K\geq 0$).

Decay of eigenvalue lower bounds with uniform data.

We now provide a lower bound for the quantity $\left\lVert{\sf P}_{\geq K+1}\mathds{1}\left\{\cdot\geq 0\right\}\right\rVert_{L_{2}(\gamma)}^{2}$, thereby giving a lower bound on $c_{K}$ defined in (10). Indeed, we have

is the Hermite decomposition of $\mathds{1}\left\{z\geq 0\right\}$. By , we know that

We now calculate the decay of $\widehat{\sigma}_{2i+1}^{2}$. By Stirling’s formula, we have

for some absolute constant $C>0$. This means that for all $i\geq 0$ we have $\widehat{\sigma}_{2i+1}^{2}\geq Ci^{-3/2}$ for some (other) absolute constant $C>0$, which gives

Therefore we have $c_{K}\geq\Omega(K^{-1/2})$ for all $K$. ∎

Covariance lower bounds for non-uniform data.

Here the equality denotes two random variables following the same distribution. We apply the Hermite decomposition of indicator function to decompose the covariance matrix $\bm{\Sigma}$:

where ${\mathcal{T}}_{1}$ and ${\mathcal{T}}_{2}$ are given as follows,

We denote the normalized $\mathbf{v}_{j}$ as $\widetilde{\mathbf{v}}_{j}=\mathbf{v}_{j}/\left\lVert\mathbf{v}_{j}\right\rVert_{2}$, and derive a lower bound on the singular value of $\widetilde{\bm{\Sigma}}_{K+1}$ with

Using the tensor product notation, we rewrite $\widetilde{\bm{\Sigma}}_{K+1}$ as

Moreover, using Lemma 1(a), we have $\lambda_{\min}\left(\widetilde{\mathbf{V}}^{\ast(K+1)}\left(\widetilde{\mathbf{V}}^{\ast(K+1)}\right)^{\top}\right)\geq 1/2$ as we picked $D\leq O(d^{K})$. Substituting into $\bm{\Sigma}_{K+1}$, we have

Therefore the smallest singular value of $\bm{\Sigma}$ is lower bounded by $\Omega(\widehat{\sigma}_{K+1}^{2}\kappa^{-(K+1)})$, which is a constant only depending on $K$ but not $d$. This finishes the proof. ∎

where $\widetilde{\mathbf{v}}_{i}\mathrel{\mathop{:}}=\mathbf{v}_{i}/\left\lVert\mathbf{v}_{i}\right\rVert_{2}$. We assume we have

The not-too-correlated condition: there exists some $\epsilon\in(0,1]$ such that for any $k\geq 0$, we have

For all large $d$ and $D\leq O(d^{k})$, we have

where $c_{k+1}>0$ is a constant that depends on $k$ but not $d$.

where $(i)$ applied the not-too-correlated condition. Repeating the above process for $K$ times leads to

Combining with existing lower bound $\lambda_{\min}((\widetilde{\mathbf{V}}\widetilde{\mathbf{V}}^{\top})^{\odot(K+1)})\geq 1/2$ (Lemma 1(a)), we see Assumption 2 holds with $\lambda_{K}=\Omega(\epsilon^{K+2}c_{K+1})$, a constant that depends on $K$ and independent of $d$.

We further note that the above two conditions are all satisfied by the rescaled Gaussian distributions: Choosing $h^{(i)}_{k}\equiv h_{k}$ (the $k$-th Hermite polynomial) for all $i=1,\dots,D$, the first condition holds with $\epsilon=1$ since $\bm{\Sigma}_{k,>k}={\bm{0}}$, and the second condition holds with $c_{k+1}=(\lambda_{\min}(\mathbf{S})/\lambda_{\max}(\mathbf{S}))^{k+1}$ (as shown earlier). Combining with the fact they only assume things about the moments of $\mathbf{x}$ (or $\mathbf{z}$; since $h_{k}^{(i)}$ are polynomials), we see that they are indeed moment-based assumptions that contain Gaussian distributions with arbitrary covariances, and thus can be fairly general.

B.3 Relative concentration of covariance estimator

Further, when $n\geq C\delta^{-2}\epsilon^{-2}\lambda_{\min}^{-1}B_{g}^{2}\log(n\vee D)$ we have with probability at least $1-\delta$ that

where $C>0$ is a universal constant. On the same event, we have the relative concentration

which implies that $(1-\epsilon)\bm{\Sigma}\preceq\widehat{\bm{\Sigma}}\leq(1+\epsilon)\bm{\Sigma}$.

We now prove the first statement, which builds on the following Rudelson’s inequality for controlling expected deviation of heavy-tailed sample covariance matrices:

Therefore, setting $n\geq C\epsilon^{-2}\lambda_{\min}^{-1}B_{g}^{2}\log(n\vee D)$, we get that

Appendix C Proofs for Section 4

This section devotes to the proof of Theorem 2. The proof consists of two main parts: expressivity of neural representation (Sections C.1 and C.2) and generalization property of Quad-Neural (Section C.3). Besides, Section C.5 presents that using data dependent regularizer also achieves improved sample complexity.

We denote by $H_{j}(x)$ the $j$-th probabilistic Hermite polynomial. We pick

From expectation to finite neuron approximation.

For a given $\epsilon>0$ and $\delta>0$, we choose $D=2\times 200^{2}k^{5}\left\lVert\bm{\beta}\right\rVert_{2}^{2k}/(\epsilon^{2}\delta)$ and independently generate $\mathbf{v}_{j}\sim{\sf N}(\mathbf{0},\mathbf{I}_{d})$ and $b_{j}\sim{\sf N}(0,1)$ for $j=1,\dots,D$. Then with probability at least $1-\delta$, we have

The desired bound can be obtained by Chebyshev’s inequality. We bound the second moment of the $L_{2}$ norm as

The last inequality invokes the identity $\frac{k!}{((k-1)!!)^{2}}+1\leq\frac{k!}{(k-1)!}+1\leq 2\frac{k!}{(k-1)!}=2k$. Therefore, choosing $D=2\times 200^{2}k^{5}\left\lVert\bm{\beta}\right\rVert_{2}^{2k}/(\epsilon^{2}\delta)$ gives rise to

From single polynomial to sum of polynomials.

To this end, we set $D=\sum_{s=1}^{r_{\star}}D_{s}\geq\frac{2\times 200^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}k_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{2k_{s}}}{\epsilon^{2}\delta}$ and define

The above inequality holds with probability $1-\delta$ by the union bound. We complete the proof. ∎

Lemma 7 showcases how to express a sum of polynomials by stacking neural random features for approximating individual polynomials. This technique will be extensively used in the remaining proofs.

C.2 Expressivity of Quad-𝐡𝐡\mathbf{h}

We show Quad-Neuralwith neural representation $\mathbf{h}$ can approximate any function $f$ of the form

To ease the presentation, we temporarily assume all the $p_{s}$ are even. We extend to odd-degree polynomials in 9. Recall we denote

We whiten $\mathbf{g}(\mathbf{x})$ by the estimated covariance matrix $\widehat{\bm{\Sigma}}$ to obtain $\mathbf{h}(\mathbf{x})=\widehat{\bm{\Sigma}}^{-1/2}\mathbf{g}(\mathbf{x})$. Note that $\mathbf{h}(\mathbf{x})$ is a $D$-dimensional vector. The approximation of Quad-$\mathbf{h}$ is stated in the following lemma.

For a given $f$ in the form of (12) with all $p_{s}$ even, and for small constants $\epsilon>0$ and $\delta>0$, we choose $D\geq\frac{4\times 50^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}p_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{p_{s}}}{\epsilon^{2}\delta}$, and $m\geq\frac{54r_{\star}D(1+\log\frac{8}{\delta})}{\epsilon^{2}}\log\frac{1}{\epsilon}$. Let $\mathbf{w}_{0,r}\overset{{\rm iid}}{\sim}{\sf N}(\bm{0},\mathbf{I}_{D})$ and $a_{r}\overset{{\rm iid}}{\sim}\textrm{Unif}(\{\pm 1\})$ for $r=1,\dots,m$, then there exist proper $\{\mathbf{w}_{r}^{*}\}$ such that with probability at least $1-\delta$, we have

By definition, $f$ can be written as a sum of polynomials with leading coefficients $\alpha_{s}$. We partition $m$ neurons into two parts according to the sign of $a_{r}$. We will use the positive part to express those polynomials with positive coefficient $\alpha_{s}$, and negative part to express those with negative coefficients. We first show for sufficiently large $m$, the number of positive $a_{r}$’s exceeds $\frac{1}{3}m$ with high probability. This follows from the tail bound of i.i.d. binomial random variables. By the Hoeffding’s inequality, we have

The remaining proof is built upon Lemma 7. We choose $D=\frac{50^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}p_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{p_{s}}}{\epsilon^{2}\delta}$, so that with probability at least $1-\delta$, there exists $\mathbf{a}$ with $\left\lVert\mathbf{a}^{\top}\mathbf{g}(\mathbf{x})-\sum_{s=1}^{r_{\star}}(\bm{\beta}^{\top}_{s}\mathbf{x})^{p_{s}/2}\right\rVert_{L_{2}}\leq\epsilon$. We further partition $\mathcal{I}_{1}$ into $r_{\star}$ consecutive groups of equal size $m_{0}$, i.e., $r_{\star}m_{0}=m/3$. Within a group, we aim to approximate $\alpha_{s}(\bm{\beta}_{s}^{\top}\mathbf{x})^{p_{s}}$ with $\alpha_{s}>0$ for some fixed $s\leq r_{\star}$. Accordingly, we choose $\mathbf{w}_{r}^{s,*}=2\sqrt{\alpha_{s}}(3r_{\star})^{1/4}m_{0}^{-1/4}\widehat{\bm{\Sigma}}^{1/2}[\bm{0}^{\top},\dots,\mathbf{a}_{s}^{\top},\dots,\bm{0}^{\top}]^{\top}$ for $r=1,\dots,m_{0}$. We have

With probability at least $1-2\delta$, we have

The proof of the claim is deferred to Appendix C.4. Based on the claim, we are ready to finish proving (13). By the triangle inequality, we deduce

The above upper bound holds with probability no smaller than $1-3\delta$. Taking

for a small $\epsilon<\left\lVert(\bm{\beta}_{s}^{\top}\mathbf{x})^{p_{s}/2}\right\rVert_{L_{2}}$, with probability at least $1-3\delta$, the following

holds true for the $s$-th group with $\alpha_{s}>0$. When $\alpha_{s}<0$, we simply set $\mathbf{w}_{r}^{s,\star}=\bm{0}$. As a result, in $\mathcal{I}_{1}$, we can express all the polynomial with a positive coefficient.

To express polynomials with negative coefficients, we use $\mathcal{I}_{2}$ analogously. By evenly partitioning $\mathcal{I}_{2}$ into $r_{\star}$ consecutive groups, for a fixed $s\leq r_{\star}$ and $\alpha_{s}<0$, we choose $\mathbf{w}_{r}^{s,*}=2\sqrt{|\alpha_{s}|}(3r_{\star})^{1/4}m_{0}^{-1/4}\widehat{\bm{\Sigma}}^{1/2}[\bm{0}^{\top},\dots,\mathbf{a}_{s}^{\top},\dots,\bm{0}^{\top}]^{\top}$. Using exactly the same argument in $\mathcal{I}_{1}$, with probability at least $1-3\delta$, for $\alpha_{s}<0$, we also have

The last step for proving Lemma 8 is to combine $\mathcal{I}_{1}$ and $\mathcal{I}_{2}$ together and choose the remaining weight parameters $\mathbf{w}_{r}^{*}$ identically $\bm{0}$ for $r\geq 2m/3+1$. Substituting into the Quad-$\mathbf{h}$ model, with probability at least $1-4\delta$, we deduce

The width $m$ satisfies $m=3r_{\star}m_{0}\geq\frac{54r_{\star}D(1+\log\frac{2}{\delta})}{\epsilon^{2}}\log\frac{1}{\epsilon}$. Replacing $\delta=\delta/4$ completes the proof. ∎

Expressivity with odd-degree polynomials.

Quad-$\mathbf{h}$ model can also efficiently express odd-degree polynomials. We rely on the following decomposition trick. Let $k$ be an integer. We rewrite a $(2k+1)$-degree polynomial as

Since QuadNTK can naturally implement the quadratic function, we only require that the neural representation $\mathbf{h}(\mathbf{x})$ can approximate $(\bm{\beta}^{\top}\mathbf{x})^{k+1}\pm(\bm{\beta}^{\top}\mathbf{x})^{k}$. This is true since random indicator functions can approximate $(\bm{\beta}^{\top}\mathbf{x})^{k+1}$ and $(\bm{\beta}^{\top}\mathbf{x})^{k}$ due to Lemma 5. We denote $\mathbf{a}_{1}^{\top}\mathbf{g}_{1}(\mathbf{x})\approx(\bm{\beta}^{\top}\mathbf{x})^{k+1}$ in $L_{2}$, and $\mathbf{a}_{2}^{\top}\mathbf{g}_{2}(\mathbf{x})\approx(\bm{\beta}^{\top}\mathbf{x})^{k}$ in $L_{2}$. Then by stacking $\mathbf{g}_{1}$ and $\mathbf{g}_{2}$, we have $[\mathbf{a}_{1}^{\top},\pm\mathbf{a}_{2}^{\top}][\mathbf{g}_{1}^{\top},\mathbf{g}_{2}^{\top}]^{\top}\approx(\bm{\beta}^{\top}\mathbf{x})^{k+1}\pm(\bm{\beta}^{\top}\mathbf{x})^{k}$ in $L_{2}$. Therefore, we only need to augment the dimension $D$ of the neural representation to approximate odd-degree polynomials. We concretize this argument in the following lemma.

For a given $f$ in the form of (12), and small constants $\epsilon>0$ and $\delta>0$, we choose $D\geq\frac{8\times 50^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}p_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{2\lceil p_{s}/2\rceil}}{\epsilon^{2}\delta}$, and $m\geq\frac{54r_{\star}D(1+\log\frac{8}{\delta})}{\epsilon^{2}}\log\frac{1}{\epsilon}$. Let $\mathbf{w}_{0,r}\overset{{\rm iid}}{\sim}{\sf N}(\bm{0},\mathbf{I}_{D})$ and $a_{r}\overset{{\rm iid}}{\sim}\textrm{Unif}(\{\pm 1\})$ for $r=1,\dots,m$, then there exist proper $\{\mathbf{w}_{r}^{*}\}$ such that with probability at least $1-\delta$, we have

Applying Lemma 5 once, there exists $\mathbf{a}_{s}$ such that $\left\lVert\mathbf{a}_{s}^{\top}-(\bm{\beta}_{s}^{\top}\mathbf{x})^{p_{s}/2}\right\rVert_{L_{2}}\leq\epsilon/r_{\star}$, when $p_{s}$ is even and the corresponding $D_{s}\geq\frac{4\times 50^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}p_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{2p_{s}}}{\epsilon^{2}\delta}$.

For an odd $p_{s}$, we apply the technique in Lemma 6. There exist $\mathbf{a}_{s,+}$ and $\mathbf{a}_{s,-}$ with corresponding random indicator features $\mathbf{g}_{s,+}(\mathbf{x})$ and $\mathbf{g}_{s,-}(\mathbf{x})$ such that

The corresponding neural representation dimension is $D_{s}\geq 50^{2}r_{\star}^{3}\frac{(p_{s}+1)^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{p_{s}+1}+(p_{s}-1)^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{p_{s}-1}}{\epsilon^{2}\delta}$. Combining the even and odd degrees together, we can choose

Lemma 9 now follows from Lemma 8 by merging $\mathbf{g}_{s,+},\mathbf{g}_{s,-}$ as a single feature $\mathbf{g}_{s}$, and $\mathbf{a}_{s,+},\mathbf{a}_{s,-}$ as a single weight vector $\mathbf{a}_{s}$ so that $\mathbf{w}_{r}^{s,*}$ can be chosen accordingly. Unifying the notation for even and odd degree polynomials, we have

C.3 Generalization of Quad-𝐡𝐡\mathbf{h}

For a given $\epsilon_{0}>0$, using Chebyshev’s inequality, we have

Choosing $n\geq\frac{8(1+\mathsf{OPT})^{2}+8r_{\star}^{2}}{\delta\epsilon_{0}^{2}}$, $\widehat{\mathcal{R}}(f_{\star})-\mathcal{R}(f_{\star})\leq\epsilon_{0}/2$ holds with probability at least $1-\delta$. We further invoke Lemma 8 and Chebyshev’s inequality again on $\frac{1}{n}\sum_{i=1}^{n}\left|f_{\mathbf{W}^{*}}(\mathbf{x}_{i})-f_{\star}(\mathbf{x}_{i})\right|$:

where the last inequality holds with probability at least $1-\delta$. We set $\frac{196r_{\star}^{2}\epsilon^{2}}{\epsilon_{0}^{2}}\leq\delta$, which implies $\epsilon^{2}\leq\frac{\delta\epsilon_{0}^{2}}{196r_{\star}^{2}}$. Accordingly, the number of neurons in the top layer needs to be at least

and the dimension of the neural representation is

This gives us that with probability at least $1-3\delta$ over the randomness of data and initializationTo achieve probability $1-\delta$, we replace $\delta$ with $\delta/3$, which only introduce a multiplicative constant in the size of $m$ and $D$., the empirical risk satisfies

Applying Theorem 1 part (2), for any second-order stationary point $\widehat{\mathbf{W}}$ and proper regularization parameter $\lambda$, we have

Towards establishing the generalization bound of $f^{Q}_{\widehat{\mathbf{W}}}$, we first find $B_{w,\star}$:

To bound $\left\lVert\widehat{\bm{\Sigma}}^{1/2}[\bm{0}^{\top},\dots,\mathbf{a}_{s}^{\top},\dots,\bm{0}^{\top}]^{\top}\right\rVert_{2}$, we first replace $\widehat{\bm{\Sigma}}$ with $\bm{\Sigma}$. We denote $\bm{\theta}_{s}=\bm{\Sigma}^{1/2}[\bm{0}^{\top},\dots,\mathbf{a}_{s}^{\top},\dots,\bm{0}^{\top}]^{\top}$, and observe $\bm{\theta}_{s}$ is the optimal solution to the following least square problem

where the last inequality follows from Lemma 5. This gives rise to

To switch back to $\widehat{\bm{\Sigma}}$, we invoke Lemma 3 on the concentration of $\widehat{\bm{\Sigma}}$ to $\bm{\Sigma}$. Specifically, with probability at least $1-\delta$, choosing $n_{0}\geq 4c\delta^{-2}\lambda_{\lceil p/2\rceil}^{-1}D\log D$ for some constant $c$, we have

Consequently, by denoting $\widehat{\bm{\theta}}_{s}=\widehat{\bm{\Sigma}}^{1/2}[\bm{0}^{\top},\dots,\mathbf{a}_{s}^{\top},\dots,\bm{0}^{\top}]^{\top}$, we have

Plugging into $\sum_{r=1}^{m}\left\lVert\mathbf{w}_{r}^{*}\right\rVert_{2}^{4}$, we have

Therefore, we can set $B_{w,\star}^{4}=108r_{\star}^{2}$. Note that $B_{w,\star}$ is independent of the width $m$.

Bounding $M_{h,\textrm{op}}$ and $B_{h}$.

The remaining ingredients are $M_{h,\textrm{op}}$ and $\left\lVert\mathbf{h}(\mathbf{x})\right\rVert_{2}$. Conditioned on the event $\frac{1}{2}\bm{\Sigma}\leq\widehat{\bm{\Sigma}}\leq\frac{3}{2}\bm{\Sigma}$, we know $\widehat{\bm{\Sigma}}^{-1}\leq 2\bm{\Sigma}^{-1}$. Therefore, we have

Note that he norm of $\mathbf{h}(\mathbf{x})$ is in the order of $\sqrt{D}$ according to Assumption 2.

The last inequality holds, due to Lemma 3 and $\widehat{\bm{\Sigma}}$ is obtained using independent samples. Conditioned on the same event $\frac{1}{2}\bm{\Sigma}\leq\widehat{\bm{\Sigma}}\leq\frac{3}{2}\bm{\Sigma}$, we have

Therefore, $M_{h,\textrm{op}}^{2}\leq 3B_{h}^{-2}$. Putting all the ingredients together and applying Theorem 1, by choosing

we establish for any SOSP $\widehat{\mathbf{W}}$, the generalization error bounded by:

We set the above probability upper bounded by $\delta$, which requires

We can now bound $\mathcal{R}(f^{Q}_{\widehat{W}})$ as

which holds with probability at least $1-\delta$.

Taking $\left\lVert\bm{\beta}_{s}\right\rVert_{2}=\sqrt{d}$, the sample size $n$ grows in the order of $\widetilde{O}\left(\frac{d^{\lceil p/2\rceil}}{\epsilon_{0}^{4}}\frac{\lambda_{\lceil p/2\rceil}^{-1}r_{\star}^{8}p^{5}}{\delta^{3}}\right)$. On the other hand, estimating covariance matrix $\bm{\Sigma}$ requires $n_{0}=\widetilde{O}\left(4\delta^{-2}\lambda_{\lceil p/2\rceil}^{-1}D\log D\right)$ samples, which is in the order of $\widetilde{O}\left(\frac{d^{\lceil p/2\rceil}}{\epsilon_{0}^{2}}\frac{\lambda_{\lceil p/2\rceil}^{-1}r_{\star}^{6}p^{5}}{\delta^{2}}\right)$. Adding $n_{1},n_{2}$ together, the sample complexity $n$ grows in the order of $\widetilde{O}\left(\frac{d^{\lceil p/2\rceil}}{\epsilon_{0}^{4}}\textrm{poly}(r_{\star},p,\delta^{-1})\right)$.

C.4 Proof of Claim 1

For a given $\mathbf{y}$, each $2\operatorname{\mathds{1}}\{\mathbf{w}_{0,r}^{\top}\mathbf{y}\geq 0\}$ is bounded in $$, hence it is sub-Gaussian with variance proxy$1$. Using the Hoeffding’s inequality, for every$\mathbf{y}$, we have

Taking $t=\sqrt{\frac{D\log\frac{3}{\gamma}\left(1+\frac{1}{D}\log\frac{2}{\delta}\right)}{m_{0}}}$, with probability at least $1-\delta$, we have

To bound the expectation, we observe that $(\mathbf{w}_{0,r}^{\top}\mathbf{y},\mathbf{w}_{0,r}^{\top}\bar{\mathbf{y}})$ is jointly Gaussian with zero mean and the covariance matrix

Therefore, we find the following probability

This implies with probability at least $1-\delta$,

Combining (15) and (16) together, with probability at least $1-2\delta$, we deduce

holds with probability at least $1-2\delta$. ∎

C.5 Learning in Quad-Neural with data dependent regularizer

We consider using data dependent regularizer for learning with unwhitened features $\mathbf{g}(\mathbf{x})$, which also yields improved sample complexity. The full learning algorithm is described Algorithm 2.

(Optimization) Given any $\epsilon>0$ and $\delta>0$, $\tau=\Theta(1)$, and some radius $B_{w,\star}>0$, suppose the width $m\geq\widetilde{O}(D^{2}B_{w,\star}^{4}\epsilon^{-1})$, sample size $n_{0}=\widetilde{O}(\delta^{-2}D)$, and we choose a proper regularization coefficient $\lambda>0$. Then with probability $1-\delta$ over $\widetilde{S}_{n_{0}}$, any second-order stationary point $\widehat{\mathbf{W}}$ of the regularized risk $\widehat{\mathcal{R}}^{\rm dreg}_{\lambda}(f^{Q}_{\mathbf{W}})$ satisfies $\|\widehat{\mathbf{W}}\widehat{\bm{\Sigma}}^{1/2}\|_{2,4}\leq O(B_{w,\star})$, and achieves

(Generalization) For any radius $B_{w}>0$, we have with high probability (over $(\mathbf{a},\mathbf{W}_{0},\widetilde{S}_{n_{0}})$) that

We recall the second-order directional derivative of $\widehat{\mathcal{R}}(f^{Q}_{\mathbf{W}})$ satisfies

where $\lambda_{0}$ is to be determined. We argue that any second-order stationary point $\widehat{\mathbf{W}}$ has to satisfy $\left\lVert\widehat{\mathbf{W}}\widehat{\bm{\Sigma}}^{1/2}\right\rVert_{2,4}=O(B_{w,\star})$. We already know from proof A.1 that for any $\mathbf{W}$, $\left\langle\nabla\widehat{\mathcal{R}}(f^{Q}_{\mathbf{W}}),\mathbf{W}\right\rangle\geq-2$ holds.

we have simultaneously for all $\mathbf{W}$ that

Therefore we see that any stationary point $\mathbf{W}$ has to satisfy

we get $36\lambda B_{w,\star}^{4}=2\tau M+\epsilon$. The second-order directional derivative of $\widehat{\mathcal{R}}^{\rm data}_{\lambda}(f^{Q}_{\mathbf{W}})$ along direction $\mathbf{W}_{\star}$ is upper bounded by

We used the fact $12ab\leq a^{2}+36b^{2}$. For a second order-stationary point $\widehat{\mathbf{W}}$ of $\widehat{\mathcal{R}}_{\lambda}(f^{Q}_{\mathbf{W}})$, its gradient vanishes and the Hessian is positive definite. Therefore, we have

By Assumption 2, we have $\lambda_{\min}(\bm{\Sigma})\geq\lambda_{k}$. Moreover, by Lemma 3, when $n_{0}=O\left(\delta^{-2}D\log D\right)$, with probability at least $1-\delta$, we have the following relative concentration of $\widehat{\bm{\Sigma}}$:

Combining these two ingredients together, we deduce

Exactly the same argument yields $\left\lVert\mathbf{W}_{\star}\right\rVert_{2,4}^{4}\leq 4\lambda_{k}^{-2}\left\lVert\mathbf{W}_{\star}\widehat{\bm{\Sigma}}\right\rVert_{2,4}^{4}$. Therefore, we choose $m=4\epsilon^{-1}\lambda_{k}^{-2}(2\lambda_{0})^{-1/2}C^{2}B_{g}^{4}B_{w,\star}^{4}\geq\epsilon^{-1}C^{2}B_{g}^{4}\|\widehat{\mathbf{W}}\|_{2,4}^{2}\left\lVert\mathbf{W}_{\star}\right\rVert_{2,4}^{2}$ and the above inequality implies

Plugging in the naive upper bound $\left\lVert\mathbf{g}(\mathbf{x})\right\rVert_{2}\leq\sqrt{D}$ in $B_{g}$, the proof is complete. ∎

where $\xi$ is i.i.d. Rademacher random variables. Recall that the whitened feature is $\mathbf{h}(\mathbf{x})=\widehat{\bm{\Sigma}}^{-1/2}\mathbf{g}(\mathbf{x})$. We further have

Consequently, the generalization error is still bounded by

When using Quad-$\mathbf{g}$ to learn low-rank polynomials in the form of

we derive the following sample complexity bound.

Suppose Assumption 2 holds, and there exists some $f_{\star}$ that achieves low risk: $\mathcal{R}(f_{\star})\leq\mathsf{OPT}$. Then for any $\epsilon,\delta\in(0,1)$ and $\tau=\Theta(1)$, choosing

$n_{0}=\widetilde{O}(D\delta^{-2})$, and a proper $\lambda>0$, Algorithm 2 achieves the following guarantee: with probability at least $1-\delta$ over the randomness of data and initialization, any second-order stationary point $\widehat{\mathbf{W}}$ of $\widehat{\mathcal{R}}^{\rm dreg}_{\lambda}(f^{Q}_{\mathbf{W}})$ satisfies

In particular, for any $\epsilon>0$, we can achieve $\mathcal{R}(f^{Q}_{\widehat{\mathbf{W}}})\leq(1+\tau)\mathsf{OPT}+2\epsilon$ with sample complexity

The proof reproduces that for Quad-$\mathbf{h}$ in Sections C.1, C.2, and C.3. Specifically, following the same argument in Lemma 9, we can establish the expressivity of Quad-$\mathbf{h}$, where for $r=1,\dots,m_{0}$, we only need to choose

Remember $\mathcal{I}_{1}=\{1,\dots,m/3\}$ where $a_{r}=1$ for $r\in\mathcal{I}_{1}$ and $\mathcal{I}_{2}=\{m/3+1,2m/3\}$ with $a_{r}=-1$. Compared to using whitened representation $\mathbf{h}(\mathbf{x})$, we remove the multiplicative factor $\widehat{\bm{\Sigma}}^{1/2}$ in $\mathbf{w}_{r}^{*}$ (see Lemma 8). The corresponding representation dimension $D=\frac{8\times 50^{2}r_{\star}^{3}\sum_{s=1}^{r_{\star}}p_{s}^{5}\left\lVert\bm{\beta}_{s}\right\rVert_{2}^{2\lceil p_{s}/2\rceil}}{\epsilon^{2}\delta}$ and the width $m\geq\frac{54r_{\star}D(1+\log\frac{8}{\delta})}{\epsilon^{2}}\log\frac{1}{\epsilon}$ remain unchanged. Then with probability $1-\delta$, we have

The rest of the proof follows Section C.3, where we need to upper bound $M_{g,\textrm{op}}$, $B_{w,\star}$, and $B_{g}$, respectively. We use the naive upper bound on $B_{g}\leq\sqrt{D}$, since each entry of $\mathbf{g}(\mathbf{x})$ is bounded by $1$. By definition, we have

Lastly, observe $B_{w,\star}^{4}=\left\lVert\mathbf{W}^{*}\widehat{\bm{\Sigma}}^{1/2}\right\rVert_{2,4}^{4}=\sum_{r=1}^{m}\left\lVert\widehat{\bm{\Sigma}}^{1/2}\mathbf{w}_{r}^{*}\right\rVert_{2}^{4}$. An upper bound has been already derived in Section C.3, which is $108r_{\star}^{2}$. As can be seen, quantities $M_{g,\textrm{op}}$, $B_{w,\star}$, and $B_{g}$ all retain the same order as using the whitened neural representation $\mathbf{h}(\mathbf{x})$ (with possibly different absolute constants). Therefore, in order to achieve

and $n_{0}$ stays the same for the covariance estimation. This yields the same sample complexity (again with a potentially different absolute constant) as using the whitened representation $\mathbf{h}(\mathbf{x})$. ∎