Learning Positive Linear Combinations of ReLUs in Fully Polynomial Time

Learning Positive Linear Combinations of ReLUs in Fully Polynomial Time

arXiv:2610.32940v1 [cs.DS] 23 Oct 2026

Ilias Diakonikolas† UW Madison ilias@cs.wisc.edu

Jelena Diakonikolas‡ UW Madison jelena@cs.wisc.edu

Puqian Wang§ UW Madison pwang333@wisc.edu

Nikos Zarifis¶ CSAIL, MIT zarifis@mit.edu

Abstract We study the fundamental problem of PAC learning positive linear combinations of k arbitrary ReLU activations on Rd with respect to the Gaussian distribution, in L2 loss. We give the first fullypolynomial time algorithm for this learning task. Specifically, our algorithm draws poly(d, k, 1/ϵ) samples, runs in poly(d, k, 1/ϵ) time, and outputs a hypothesis that is ϵ-close to the target concept with high probability. All prior algorithms for this problem require super-polynomial time in k or 1/ϵ, or leverage strong well-conditioned assumptions on the weight matrix. On a technical level, our approach develops a novel boosting procedure of a “weak” learner that recovers the hidden weight vectors to constant error, in a list-decodable manner. To develop this “weak” learning procedure, we leverage implicit tensor computation techniques to construct a low-degree filtering procedure based on Hermite tensors of the target function. This procedure amplifies separation among (clusters of) weight vectors to learn each of them to constant error. We then argue that the same procedure can be used for improving the estimation of hidden weight vectors, using gradient descent-like updates on the Euclidean sphere. A crucial ingredient of our algorithm is a pruning approach that ensures only relevant vectors that approximate the hidden weight vectors at target accuracy are maintained by the algorithm.

∗ Authors are listed in alphabetical order. † Supported by NSF Medium Award CCF-2107079, ONR Award number N00014-25-1-2268, an H.I. Romnes Faculty Fellowship and a Guggenheim Fellowship. ‡ Supported in part by the AFOSR YIP award FA9550-24-1-0076, by the NSF MFAI Award DMS-2502282, and by the NSF CAREER Award CCF-2440563. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the U.S. Department of Defense. § Supported in part by the AFOSR YIP award FA9550-24-1-0076. ¶ Supported by MIT METEOR fellowship. Some of this research was done while the author was at UW Madison, supported by NSF Medium Award CCF-2107079.

Contents 1 Introduction 1.1 Problem Formulation and Main Result . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2 Technical Overview . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2.1 Coarse Direction Recovery . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2.2 Boosting Coarse Direction Recovery . . . . . . . . . . . . . . . . . . . . . . . . . 1.3 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.4 Concurrent Work and AI Use . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.5 Paper Organization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

1 1 2 3 5 7 9 9

2 Notation and Preliminaries

9

3 Optimization via Coarse Direction Recovery, Refinement, and Pruning 11 3.1 Learning λ-Heavy Hidden Weight Vectors . . . . . . . . . . . . . . . . . . . . . . . . . . 12 3.1.1 Boosting Hidden Vector Estimates . . . . . . . . . . . . . . . . . . . . . . . . . . 12 3.1.2 Closeness Tester . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 3.1.3 Covering and Pruning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19 3.2 Proof of Theorem 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23 4 Efficient Coarse Direction Recovery 24 4.1 Low-Degree Filtering: Revealing Hidden Directions . . . . . . . . . . . . . . . . . . . . . 24 4.2 Warm-up: Quasi-Polynomial Time CDR Algorithm . . . . . . . . . . . . . . . . . . . . . 26 4.3 Polynomial Time CDR via Implicit Tensor Computation . . . . . . . . . . . . . . . . . . 32 4.3.1 Determining the Probability of the Probing Event . . . . . . . . . . . . . . . . . 33 4.3.2 Estimating the Target Matrix Using Implicit Tensor Computation . . . . . . . . 34 4.4 Proof of the Main Result of the Section: Polynomial Complexity CDR . . . . . . . . . . 47 5 Putting Everything Together: Proof of Theorem 1.3

47

1

Introduction

Understanding the computational complexity of learning neural networks is a central question in the foundations of machine learning. Over the past decade, a substantial body of work has developed provable learning algorithms under various structural and distributional assumptions; see, e.g., [JSA15, ZLJ16, ZSJ+ 17, GLM18, GK19, CKM22, WZDD23, ZWDD25, DIKR25, DLB25]. Despite this progress, a basic question has remained open even for networks with a single hidden layer: Can positive linear combinations of ReLUs be learned in polynomial time under Gaussian inputs, without strong structural assumptions on the hidden weights? This longstanding question is explicitly highlighted, for example, by Chen, Klivans, and Meka [CKM22] (see also Chen’s thesis [Che21]): “Ideally, we would like an algorithm with sample complexity and running time that is polynomial in all the relevant parameters. Even for learning arbitrary sums of ReLUs, [. . . ], it remains a major open question to obtain a polynomial-time algorithm [. . . ].” Positive linear combinations of ReLUs are arguably the simplest extension of learning a single neuron— the setting of generalized linear models. A fully polynomial-time algorithm for this basic class is an important step towards understanding the learnability of more general networks. Moreover, the techniques developed for this problem may also be useful for learning other latent variable models. In this paper, we study the task of PAC learning positive linear combinations of at most k ReLUs on Rd under the standard Gaussian distribution. Specifically, the learner has access to independent labeled examples from an unknown function in the concept class, and its goal is to output a hypothesis that approximates the target in L2 norm. The relevant parameters are the dimension d, the number of hidden units k, and the inverse error 1/ϵ. It is well-known that polynomially many samples informationtheoretically suffice for this learning task, ignoring computational complexity. The challenge is to develop a computationally efficient learning algorithm: despite a sequence of algorithmic advances, the best previous running-time bounds were super-polynomial in at least one of these parameters. The work of [DKKZ20] gave the first learner for this problem with running time polynomial in d and 2 e 1/ϵ for every fixed k. Their algorithm uses O(d)poly(k/ϵ) samples and runs in time poly(d/ϵ)+(k/ϵ)O(k ) . The work of [DK20] improved the running time using a structural theorem on small covers for near-zero 2 sets of polynomials. Its sample and computational complexities are poly(dk/ϵ) + (k/ϵ)O(log k) , i.e., quasi-polynomial in k even for constant error. More recently, [DK25] developed an algorithm based on implicit high-order moment estimation with sample and computational complexity poly(d, k) 2poly(1/ϵ) . That is, while this achieves polynomial dependence on d and k, the dependence on 1/ϵ is exponential. In summary, neither of these results gives a fully-polynomial time algorithm. The above discussion motivates the following question: Can positive linear combinations of k ReLUs on Rd be learned under Gaussian inputs with sample and computational complexity poly(d, k, 1/ϵ)? As our main contribution, we answer this question in the affirmative. Specifically, we give a learning algorithm with sample and computational complexity poly(d, k, 1/ϵ) that works for arbitrary positive linear combinations of ReLUs. Our algorithm outputs a simple interpretable hypothesis. Moreover, the technical tools we develop in the process may be useful for learning other latent variable models.

1.1

Problem Formulation and Main Result

Before stating our main result, we define the underlying concept class and the PAC learning task. Definition 1.1 (Positive linear combinations of ReLUs). Let σ(t) := max{0, t} denote the ReLU activation. For d, k ∈ Z+ , let Cd,k denote the class of functions f : Rd → R of the form f (x) = Pk i=1 βi σ(wi · x), where βi ∈ [0, 1] and ∥wi ∥2 = 1 for all i ∈ [k]. We call the wi ’s the hidden weight vectors and the βi ’s the coefficients. By positive homogeneity of the ReLU, the norm of a nonzero hidden weight vector can be absorbed into its coefficient. The constraint βi ≤ 1 fixes the scale (without loss of generality); see Remark 1.5. For a function g : Rd → R, we write ∥g∥L2 := (Ex∼N (0,Id ) [g(x)2 ])1/2 .

1

Definition 1.2 (PAC learning Cd,k under Gaussian marginals). The learner is given accuracy and confidence parameters ϵ, δ ∈ (0, 1) and access to independent examples (x, y) ∼ D, where D is a distribution supported on Rd × R such that x ∼ N (0, Id ) and y = f (x) for an unknown f ∈ Cd,k . The goal of the learner is to output a hypothesis h : Rd → R that, with probability at least 1 − δ, satisfies ∥h − f ∥L2 ≤ ϵ∥f ∥L2 . The hypothesis need not belong to Cd,k . The learner is called proper if it outputs a hypothesis in Cd,k . Our main result is summarized in the following theorem (see Theorem 5.1 for a detailed full version). Theorem 1.3 (Main result, informal). There is an algorithm that, given d, k ∈ Z+ , ϵ ∈ (0, 1), and access to independent examples (x, f (x)), where x ∼ N (0, Id ) and f ∈ Cd,k , uses poly(d, k, 1/ϵ) samples and time and outputs a hypothesis h such that, with probability at least 2/3, ∥h − f ∥L2 ≤ ϵ∥f ∥L2 . A few remarks are in order. While our learner is not proper, its hypothesis has a simpleP form: a linear k function plus a linear combination of at most k absolute-value functions, i.e., h(x) = g·x+ i=1 αi |vi ·x|. Equivalently, using |t| = σ(t) + σ(−t), it is a linear function plus a linear combination of at most 2k ReLUs. Thus, the hypothesis is itself a small network with O(k) nonlinear units. For comparison, the improper learner from [DK20] outputs a positive linear combination of quasi-polynomially many ReLUs, while the learner from [DK25] outputs a polynomial represented through compressed Hermite coefficients. For these representations, the bounds on the time needed to even evaluate the hypothesis on a single new example are also super-polynomial in k or 1/ϵ. Remark 1.4 (Additive label noise). More generally, one may consider examples (x, y) with x ∼ N (0, Id ) and y = f (x) + ξ, where ξ is independent of x, has mean zero, and satisfies E[ξ 2 ] ≤ σξ2 . Following the convention of [DKKZ20], the objective in this setting is to output a hypothesis h satisfying
∥h − f ∥2L2 ≤ ϵ2 ∥f ∥2L2 + σξ2 with high probability. The realizable setting of Definition 1.2 corresponds to σξ = 0. Our algorithm straightforwardly extends to this setting with a poly(σξ ) overhead in its complexity. Remark 1.5 (Normalization). Our result applies to arbitrary positive linear combinations of at most k ReLUs, without bounds on the coefficients or the norms of the hidden weight vectors. By positive homogeneity, we can absorb the norm of each hidden weight vector into its coefficient. Writing b for 2 the resulting coefficient sum, we have E[f (x)] = √b2π and E[f (x)2 ] ≤ b2 . By Chebyshev’s inequality, √ the empirical mean y of a sufficiently large constant number of labels gives an estimate B = 2 2π y satisfying b ≤ B ≤ 3b with high constant probability. For B > 0, we can apply our learner to fresh examples with labels divided by B and multiply its output by B. The normalized target has unit hidden weight vectors and coefficients in [0, 1], and rescaling preserves the relative-error guarantee. If B = 0, we return the zero hypothesis. Thus, our normalization is without loss of generality and requires no prior knowledge of the target’s scale. Terminology. Throughout this paper, we will say that a distribution (x, y) ∼ D supported on Rd × R is a Gaussian sum of k ReLUs if the x-marginals Pk are distributed according to the standard normal distribution, i.e., x ∼ N (0, Id ), and if y(x) = i=1 βi σ(wi · x), where σ is the ReLU function, βi ∈ [0, 1], and ∥wi ∥2 = 1, for all i ∈ [k].

1.2

Technical Overview

To explain the main ideas behind our results, we begin with the following observations. First, lowcoefficient ReLU components in the sum have little influence √ on the predicted label: in particular, if the coefficient βi associated with component i is smaller than ϵ/k, then zeroing out all such components can contribute at most an additive ϵ to the L22 loss we aim to minimize. Thus, setting a threshold √ ϵ0 ≈ ϵ/k and focusing on approximating the ReLU components with coefficients above this threshold, comes at no loss of generality. In the remainder of the technical overview, we only consider such nontrivial components with coefficients βi ≥ ϵ0 , and refer to them as “ϵ0 -heavy” in the rest of the paper. 2

Second, we claim that it suffices to recover the target vectors up to their orientation (i.e., up to their signs). The reason is the decomposition σ(z) − z2 = |z| 2 . Specifically, if we calculate the first-moment Pk parameter g = E[yx] = 12 i=1 βi wi , we get k

y−g·x=

1X βi |wi · x| . 2 i=1

(1)

Thus, after subtracting the linear part, the remaining function is a sum of absolute values. Its terms are unchanged when wi is replaced by −wi , so the directional part of the algorithm only needs to recover one of these two orientations. We measure the error of the angle using θ± (v, wi ) = min{θ(v, wi ), π − θ(v, wi )}, rather than by the ordinary angle for this reason. The birds-eye view of our algorithm is the following. The algorithm operates in rounds. Each round t = {1, 2, . . . , T } of computation maintains the following invariant: Algorithm Invariant. There is a set of vectors Vt , |Vt | ≤ k, maintained by the algorithm such that for every ϵ0 -heavy hidden vector wi , there exists a vector vi ∈ Vt that is close to it; in particular, θ± (vi , wi ) ≤ Rt , Rt > 0. Note that this mapping of vectors need not be one-to-one: one vector v can be close to multiple wi ’s. Initially, R0 ∈ (0, 1) is a constant, and this is achieved using the Coarse Direction Recovery (CDR) subroutine described in Section 1.2.1 below. We argue that we can maintain the stated invariant with Rt+1 ≤ ρRt for a constant ρ ∈ (0, 1) and all t ≥ 0, √ so that the algorithm approximates all ϵ0 -heavy hidden weight vectors wi with error θ± (vi , wi ) ≤ ϵ/k, after T = O(logρ (k/ϵ)) rounds. The last step is then to estimate the coefficients for the absolute-sum representation (1). It is not hard to see that this part boils down to solving a linear regression problem of the form min m

α∈[0,1]

E

h

m

y−g·x−

(x,y)∼D

2 i 1X αj |vj · x| , 2 j=1

where m = |VT | ≤ k. This problem is convex and it is not hard to argue that it is efficiently solvable in terms of both sample and computational complexity. The technical crux of the approach thus lies in implementing the rounds of the above high-level algorithm, which relies upon multiple core technical components described below. 1.2.1

Coarse Direction Recovery

A key technical tool in our algorithmic approach is a procedure that, given a failure probability δCDR ∈ (0, 1) and access to i.i.d. examples (x, y(x)), where x ∼ N (0, Id ) and y(x) = f (x) for a fixed but unknown f ∈ Cd,k , outputs a polynomially-sized list in k, 1/δCDR with the following property: with probability at least 1 − δCDR , for each ϵ0 -heavy hidden weight vector wi defining the labeling function, there exists at least one list element ri such that |ri · wi | ≥ 0.99. Crucially, this procedure, to which we refer as the Coarse Direction Recovery (CDR) can be implemented using polynomially many labeled examples (in d, k, 1/ϵ0 , 1/δCDR ), and in time polynomial in d and the number of samples. CDR can be understood as a particularly “weak” 1 learner for our problem, where each ϵ0 -heavy hidden weight vector is learned to constant error, in a list-decodable manner. The latter means that the output provided by CDR is a polynomially-sized list containing useful, constant-error outputs, but without any information on how to separate useful outputs from the rest of the list. It is not hard to argue that a stronger √ learner, which would only output vectors ri with sufficiently strong correlation |ri · wi | ≥ 1 − Θ( ϵ/k), would solve our learning problem stated in Definition 1.2. Yet, with the techniques described below, such a stronger requirement would lead only to a quasi-polynomial algorithm, at best. The first insight for developing the CDR subroutine comes from the properties of Hermite tensors of positive linear combinations of ReLUs (see Section 2 for relevant definitions); namely, that Hermite moments of a linear combination of activations can be expressed as a sum of rank-one tensors, an 1 Here, “weak” is in quotes to distinguish it from the standard notion in PAC learning theory, where weak learning is understood as slightly better than random, whereas we require constant correlation of vectors.

3

observation used extensively in prior work [DKKZ20, CDG+ 23, DK25]. More specifically, for a positive linear combination of ReLUs, we have Q2p+2 :=

E (x,y)∼D

[yHe2p+2 (x)] ∝

k X

βi wi⊗2p+2 ,

(2)

i=1

where He2p+2 (x) is the Hermite tensor of order 2p + 2, for p ≥ 0. The moment formulation here is what estimation from samples; the fact that it is proportional to the sum of rank-one tensors Pk enables ⊗2p+2 is what we use to obtain the guarantee of CDR. i=1 βi wi Directly recovering the individual rank-one components of a high-order tensor is, in general, a computationally challenging problem. Instead, a useful approach is to contract the tensor along all but two indices, obtaining a matrix whose leading eigenvector can reveal one of the hidden directions. To be precise: taking the inner product of Equation (2) with a tensor u⊗2p and reshaping the result as a matrix, we get that k X Q2 (u) = Q2p+2 , u⊗2p ∝ βi |wi · u|2p wi wi⊤ . (3) i=1

The key observation is that the coefficients of the rank-one matrices in Equation (3) depend on powers of the correlations, |wi · u|2p . Consequently, even a small multiplicative advantage in the correlation with one hidden weight vector can be amplified by choosing a sufficiently large degree p. To illustrate this idea, suppose that w1 is well-separated from the remaining hidden vectors w2 , . . . , wk (e.g., maxi̸=1 |w1 · wi | ≤ 1/2). Then, we can write Q2 (u) as a sum of two matrices Q2 (u) ∝ β1 |w1 · u|2p w1 w1⊤ + Qϵ (u). If a vector u is such that the spectral norm of β1 |w1 · u|2p w1 w1⊤ is much larger than the spectral norm of Qϵ (u), then the top eigenvector of Q2 (u) must have constant Pk correlation with w1 . A sufficient condition for this to hold is to have βi |wi · u|2p ≫ j̸=i βj |wj · u|2p . Given that we can choose p to be a (sufficiently) large power, it suffices to find a u so that for all i = ̸ 1, we have |u · w1 | ≥ (1 + ξ)|u · wi | for some ξ > 0 (for simplicity of this exposition, assume ξ = 1). Then, by choosing p ≫ log(k maxi βi /β1 ), we have that the spectral norm of β1 |w1 · u|2p w1 w1⊤ is at least two times larger than the spectral norm of Qϵ (u). Hence, the top eigenvector of Q2 (u) is almost parallel to w1 . Importantly, this intuition extends beyond the well-separated case. Indeed, directions wi that are close to ±w1 , i.e., |wi · w1 | ≫ 0, can be grouped together with w1 while retaining the same properties, since their contributions align with the desired rank-one matrix. Therefore, we only need to suppress directions that are far from ±w1 , rather than isolate a single component. The hope here is that by sampling u from an appropriately selected probability distribution, the above event would happen for all wi , with some nontrivial probability. This is the main intuition for why the stated guarantee of CDR should be possible. There are at least two main technical challenges with the above argument. First, since u is chosen from a probability distribution, to ensure that we can produce a polynomially-sized list containing at least one vector that aligns well with some wi , the event ensuring separation in alignment among hidden weight vectors with u described above must happen with probability no smaller than an inverse polynomial in the problem parameters. (Ensuring that such an event is satisfied for all wi , i ∈ [k], for at least one of the drawn probing vectors u follows by the union bound.) We argue that this is indeed the case for the simple distribution where u is chosen uniformly at random from the k-dimensional unit Euclidean sphere, defined over the space containing the space spanned by the targets wi , i ∈ [k]. The same can also be argued for other related distributions, such as standard Gaussians. Second, as hinted before and more importantly, we need to argue that the matrix in (3) can be estimated efficiently from labeled examples with bounded error in the operator norm, so the top eigenvector is preserved (up to constant error). For example, when u is sampled uniformly from a √ k-dimensional sphere, its typical correlation with a fixed unit vector is of the order 1/ k. Raising this correlation to degree 2p produces a factor of order k −p , which is quasipolynomially small when p = Θ(log(k/ϵ)). Estimating the resulting matrix may therefore require quasipolynomially small accuracy, and this, as a result, leads to quasipolynomial sample complexity. We provide full technical details for this approach in Section 4.2. To obtain a fully polynomial CDR implementation, the main challenge is to ensure that the estimation task can be implemented efficiently, in terms of both samples and time. A major obstacle when sampling u uniformly from the sphere or from the standard Gaussian distribution is that the variance of the empirical estimate of Q2 = Q2p+2 , u⊗2p is large because the variance of the tensor u⊗2p is large 4

for large degree p, which leads to error amplification. To illustrate this point, consider p = 1 and b 4 = Q4 + ζH. For the decompose the estimation tensor into Q4 ∈ (Rd )⊗4 and an error tensor ζH, Q sake of discussion, suppose H = Id ⊗ Id . Then the error term after tensor contraction with the filter u⊗2 becomes ⟨H, u⊗2 ⟩ = (u⊤ u)Id . If u is selected from a real Gaussian distribution, then the second moment of the spectral norm of this error matrix is E[∥⟨H, u⊗2 ⟩∥22 ] = E[∥u∥42 ] = O(d2 ). On the other hand, if u = p + iq is a complex Gaussian vector where p, q are i.i.d. standard real Gaussians, then the second order moment becomes E[∥⟨H, u⊗2 ⟩∥22 ] = O(d). This is the main intuition for why complex vectors are more helpful for keeping the estimation error small. The advantage of complex vectors leading to lower growth of the error also persists for higher moments; this is formally established for real and complex Gaussian vectors in Lemma 4.26 and Lemma 4.27, in Section 4. The choice of Gaussian vectors in our analysis is for technical convenience. Full technical details for this polynomial sample and time implementation based on complex Gaussians are included in Section 4.3. 1.2.2

Boosting Coarse Direction Recovery

The CDR procedure described in Section 1.2.1 gives a concrete approach for constructing initial, constantcorrelation candidate vectors. However, several questions remain on how to make this procedure useful for constructing a hypothesis with arbitrarily small target error ϵ > 0. Most importantly, we need a way to reduce the error of candidate vectors and isolate useful candidate vectors from the rest of the produced list. We begin with providing the intuition for how to improve the error. Ideal (and nearly-ideal) update directions. Fix one candidate vector v ∈ Sd−1 and one target ⊥v wi ∈ Sd−1 . If we qknew wi , the component of wi orthogonal to v, then we would immediately know 2

wi , since wi = 1 − wi⊥v 2 v + wi⊥v , so one update of v would suffice. Notice that wi⊥v corresponds to the “shortest path” direction for rotating v towards wi : rotating v in the direction of wi⊥v leads it to wi , following the geodesic that connects v and wi . Of course, such an ideal update direction is not available to the algorithm. However, it is also unnecessary for contracting the distance between v and wi . If, instead, we had access to a unit vector r that correlates well with wi⊥v , in the sense that r·

wi⊥v w ∥ i⊥v ∥

≥ 0.9, then r would be a “nearly-ideal” update direction: although not moving v on the

2

geodesic connecting it to wi , it would allow for a large rotation that contracts the angle between v and wi by a constant factor; this is the primary idea underlying efficient optimization on the Euclidean sphere [WZD26]. Of course, the question remains of how to construct such update directions. Observe that the “constant correlation” requirement r·

wi⊥v

∥wi⊥v ∥2

≥ 0.9 is similar to what we had stated

[ ⊥v := for CDR. In particular, if we had access to samples from a problem instance with w i

wi⊥v ∥wi⊥v ∥

as the 2

hidden weight vectors, CDR would apply directly. We argue that this is indeed possible; the small loss in correlation (0.99 vs 0.9) is a small error slack needed for the technical argument. Methodologically, the use of moments beyond the first to construct update directions that enable contracting the distance to target vectors is related to [WZDD25]; in the present work, the use of higher moments in (2) is crucial to updating a set of k ≥ 1 vectors simultaneously (in [WZDD25], k = 1, and p = 0 in (2) suffices). [ ⊥v , in Section 3.1.1, we show that a simple To argue that CDR could likewise apply to estimating w i

rejection sampling-based procedure applied to labeled examples drawn from the input distribution leads to a new sampling distribution that itself is a Gaussian sum of ReLUs with hidden weight e i = wi /∥wi ∥2 , where wi = wi⊥v + γ(v · wi )v and γ ∈ (0, 1) is a parameter of the rejection vectors w sampling procedure that we can choose. Then, choosing γ sufficiently small (relative to the relevant e i becomes negligible, and we can accuracy/heaviness threshold), the contribution of γ(v · wi )v in w e i are also constant-alignment estimates for wi⊥v , with ensure that constant-alignment estimates for w just a slightly smaller alignment constant (this is where the loss from 0.99 to 0.9 occurs). Thus, to obtain vectors with constant-factor alignment with wi⊥v s, it suffices to apply CDR to the samples from D retained after the rejection sampling. Now, having access to a procedure that can produce near-ideal update directions, we discuss how the algorithm can maintain the invariant stated at the beginning of technical overview. The process of generating candidate updates and pruning the generated list is illustrated in Figure 1.

5

Constructing a large list of candidate updates. The process of constructing new candidate vectors is now as follows. In each round of computation t, we maintain a list Vt , which contains vectors vi that approximate targets wi . At the beginning of each round t, for each v ∈ Vt , we call the Refine procedure, which performs the following operations. We generate a polynomially-sized list L(v) that contains near-ideal update directions for all wi , i ∈ [k] (and possibly many other vectors with no specific guarantee). For each vector in r ∈ L(v), we perform candidate projected gradient descent-style updates projSd−1 (u + η1 r), projSd−1 (u − η1 r), and add both obtained vectors to L(v), where η1 is an appropriately chosen step size. For technical reasons (namely, to ensure sufficient contraction), we repeat this procedure three times. As a result, we obtain a large but still polynomially sized list Vtnew of vectors, with the following property: for each hidden weight vector wi , i ∈ [k], the list Vtnew contains at least one vector vi such that θ± (vi , wi ) ≤ cRt+1 = cρRt , ρ ∈ (0, 1), and c ∈ (0, 1) is a constant that provides slack for introducing small errors when pruning the list. That is, there are new approximate vectors vi within this large list Vtnew that ensure constant contraction cρ of distance to the target vectors wi . However, as should be clear, at this point these “useful” vectors vi are like a few needles in a large haystack. The next step thus involves identifying the useful vectors and pruning the rest. Testing and pruning. As a first step in removing irrelevant vectors from Vtnew , in Section 3.1.2, we develop a testing procedure that can be used to remove all irrelevant vectors in Vtnew ; namely, all vectors with alignment to all wi larger than a small constant 1 + τ , τ > 0, times cRt+1 . This testing procedure is based on using properties of Hermite coefficients of ReLUs, with similar intuition as in construction of CDR discussed in Section 1.2.1. In particular, it is possible to argue that for even integer p, we have the following identity: ap (v) :=

X 1 E[yhep (v · x)] = βi cosp θ± (v, wi ), σ bp i

(4)

where hep (v · x) is the normalized probabilist’s Hermite polynomial of order p (see Section 2). The first σp ) E[yhep (v · x)] tells us that ap can be estimated from samples; the second term P termp (1/b 2 i βi cos θ± (v, wi ) that for sufficiently large degree p (p ≈ log(k/ϵ)/(cRt+1 ) ), we can determine if v correlates nontrivially with any target direction wi (as measured by θ± (v, wi )) or if it can be eliminated. More specifically, (with high probability), our tester returns FALSE for all v ∈ Lt such that θ± (v, wi ) > (1 + τ )cRt+1 and TRUE for all v ∈ Lt such that θ± (v, wi ) ≤ cRt+1 . As a result, the tester can be used as a rough pruning method that removes all v ∈ Lt that do not approximate any of the target directions within angle (1 + τ )cRt+1 . Ideally, the tester alone would suffice to move on to the next round of computation. However, it is possible that one wi is approximated by many vectors vi in the list, and if all were retained in Vt+1 , the size of VT would not be guaranteed to be polynomial, as the algorithm takes T = Θ(log(k/ϵ)) rounds and each call to Refine (for each vector in the list) produces a polynomially-sized list. Thus, the final step, in Section 3.1.3, is to prune the list Vtnew down to at most k vectors (and assign it to Vt+1 ). The reason we can do this is geometric: since the previous steps ensured sufficient contraction with leeway for additional error, we make the following observation. If multiple vectors v are close to the same target wi , then they must be close to each other. The final pruning step is thus the following: pick any candidate vector v in the pruned list and add it to Vt+1 . The testing procedure ensures that this vector is within the target distance of at least one hidden weight vector. Now, iterate through the remaining vectors v in the pruned list Vtnew and add v to Vt+1 if and only if it is at least 2(1 + 2τ )cRt+1 -far from all the vectors already in the list. We claim that for each wi , there is exactly one vi ∈ Vt+1 such that θ± (vi , wi ) ≤ Rt+1 (but the same vi can satisfy θ± (vi , wi ) ≤ Rt+1 for multiple wi ), thus |Vt+1 | ≤ k. Fix one wi and consider all v such that θ± (v, wi ) ≤ (1 + τ )cRt+1 ; these are the candidates v that survive the testing prune, and among them, there is at least one vector vi∗ such that θ± (vi∗ , wi ) ≤ cRt+1 , as guaranteed by the refinement procedure. Because, by the triangle inequality, all such v are within 2(1 + τ )cRt+1 of each other, at most one can get added to Vt+1 . If exactly vi∗ is added, we are done, so suppose vi∗ is not added to Vt+1 . This can only happen if there exists v′ ∈ Vt+1 such that θ± (vi∗ , v′ ) ≤ 2(1 + 2τ )cRt+1 . Thus, again by the triangle inequality, θ± (wi , v′ ) ≤ (2(1 + 2τ ) + 1)cRt+1 . Parameters c, τ are chosen so that (2(1 + 2τ ) + 1)c ≤ 1, and thus we have that each wi can be mapped to one vi ∈ Vt+1 such that θ± (vi , wi ) ≤ Rt+1 .

6

(a) After the Refine procedure, each candidate produces a list of vectors (marked in a dotted circle), of which some are closer to the hidden weight vectors. But there are also potentially many other redundant vectors.

(b) After the Test procedure, those vectors that do not “cover” hidden weight vectors at a target contracted distance are discarded (marked in gray).

(c) There might be many candidate vectors that are close to each other, covering the same hidden weight vector. After the Prune procedure, the candidate vectors are separated (the pruned vector is also marked in gray). Yet, we guarantee that each hidden weight vector is still covered by the candidate vectors after pruning.

Figure 1: The update procedure of the CoverAndPrune(ϵ0 , δ; D) algorithm. × denotes the hidden weight, • denotes the candidate vector at step t, ◦ denotes the vectors generated after step t. Only the hidden weight vectors with not-too-small coefficients are displayed.

1.3

Related Work

Our work belongs to the literature on learning structured multi-index models, and we review related work within this context. 7

Learning positive linear combinations of ReLUs. The most directly related works are [DKKZ20, DK20, DK25], whose complexity guarantees were discussed earlier in the introduction. The algorithm of [DKKZ20] uses the degree-2 Chow parameters of the target to find a low-dimensional subspace sufficient for prediction, and then searches over candidate networks in that subspace. The work of [DK20] uses higher-degree moments to find a vector space of polynomials that nearly vanish on the hidden weight vectors. Its structural result gives a small cover for the near-zero set of these polynomials, which is then used to construct a hypothesis. The work of [DK25] estimates high-order moment tensors implicitly, without writing down their entries, and uses these estimates to approximate the target’s Hermite expansion. This generalizes the implicit moment-estimation technique of Li and Liu [LL22], developed for clustering separated mixtures of spherical Gaussians. Learning general linear combinations of ReLUs. PAC learning general linear combinations of ReLUs, with both positive and negative coefficients, has also been studied. Without structural assumptions on the hidden weight vectors, [CDG+ 23] gave the first algorithm with running time polynomial in 2 d and 1/ϵ for every fixed k. Subsequent work improved the running time to (dk/ϵ)O(k ) [CN24] and (dk/ϵ)O(k) [DK24]. These algorithms handle a broader class of functions than the one studied here, but their complexity scales exponentially in the number of hidden units.2 Learning multi-index models and deeper networks. The work of [CKM22] gives fixed-parameter learning algorithms for homogeneous ReLU networks of arbitrary depth under Gaussian inputs, with polynomial dependence on the ambient dimension but exponential dependence on network size and inverse accuracy. Recent work on multi-index models gives learning guarantees for more general classes of functions that depend only on an unknown low-dimensional subspace [DIKZ25, DIKR25, DLB25]. Similarly, the latter algorithms have polynomial dependence on the ambient dimension but exponential dependence on the other relevant parameters. Parameter learning under nondegeneracy assumptions. A number of earlier works focus on parameter learning, namely recovering the parameters of the data-generating network. Examples include tensor-based methods [JSA15] and algorithms combining initialization with local optimization [ZSJ+ 17]. Related tensor methods were developed for mixtures of generalized linear models in [SJA16]. For positive linear combinations of ReLUs under Gaussian inputs, [GLM18] gives a parameter-learning algorithm assuming a full-rank hidden-weight matrix, with complexity depending polynomially on its condition number. [BJW19, GKLW19] give parameter-recovery algorithms for vector-valued two-layer networks under rank and related nondegeneracy assumptions. The work of [ATV21] develops tensorbased algorithms for scalar-output ReLU networks with biases and general coefficients. In contrast, we focus on small prediction error rather than parameter recovery, with complexity independent of separation, rank, condition-number, or any other parameters requiring strong assumptions about the configuration of the hidden weight vectors. Computational lower bounds. For general linear combinations of ReLUs, [DKKZ20, GGJ+ 20] establish super-polynomial lower bounds in the correlational statistical-query (CSQ) model [BF02, Fel08]. This is a restriction of the Statistical Query (SQ) model [Kea93], in which the learner accesses the target through estimates of correlations with query functions; see also [BFJ+ 94, VW19]. For networks with two hidden layers, [CGKM22] establishes general SQ and cryptographic hardness even in the noiseless setting. More recently, [LZZ24] proves cryptographic hardness for one-hidden-layer ReLU networks with both positive and negative coefficients and biases, in the presence of small independent Gaussian label noise. Whether general linear combinations of ReLUs without biases can be learned in fully-polynomial time under Gaussian inputs remains open. Since our algorithm can be implemented in the CSQ model, the CSQ lower bounds above rule out an immediate extension to general coefficients. 2 Here we focus on learning from random labeled examples. If query access to the function is allowed, [CKM21] gives a learner for one-hidden-layer ReLU networks with general coefficients and biases whose query complexity and running time are polynomial in dimension, the number of hidden units, inverse accuracy, and bounds on the weights and biases.

8

1.4

Concurrent Work and AI Use

All ideas and technical arguments in this paper are the intellectual ownership of the authors, who have worked on developing them over the past 2+ years. In the process of writing the paper, OpenAI GPT models were used to improve the presentation (including generating Figures 1 and 2), audit the proofs, and simplify some of the statements. We note that a very recent paper [Zhe26] appeared on the arXiv, claiming essentially the same main result. Rather curiously, that work is quoting some of the exact same core techniques as our work, but is lacking in most of the proof details. The paper [Zhe26] further uses some of the same non-standard terminology appearing in our work (such as “refine,” “prune,” “cover,” etc.) and even some of the same universal constants chosen lazily by the authors to be smaller than needed. As a result, we were forced to rush posting of this work, before we were entirely satisfied with its presentation, but with complete technical details included.

1.5

Paper Organization

The paper is structured as follows: In Section 2, we record the notation and background required in our technical sections. Sections 3 to 5 provide a detailed description and analysis of our algorithm.

2

Notation and Preliminaries

Basic Notation. For n ∈ Z+ , we denote [n] := {1, . . . , n}. We use bold lowercase letters for vectors and bold uppercase letters for matrices and tensors. For x ∈ Rd and i ∈ [d], (x)i denotes the ith Pd coordinate of x, and ∥x∥2 := ( i=1 (x)2i )1/2 denotes the ℓ2 norm of x. We write x · y for the standard inner product of x, y ∈ Rd and θ(x, y) := arccos(x · y/(∥x∥2 ∥y∥2 )) when x ̸= 0 and y ̸= 0. We denote θ± (x, y) := min{θ(x, y), π − θ(x, y)}. For unit vectors u, v, we use u⊥v to denote the component of u orthogonal to v, i.e., u⊥v = b to denote its normalization to unit ℓ2 norm. Given v ∈ Rd , (I − vv⊤ )u. For any vector u ̸= 0, we use u ⊥ we use v to denote the linear subspace of vectors orthogonal to v, i.e., v⊥ = {u ∈ Rd : u ⊥ v}. We denote the projection from Rd to v⊥ by P⊥v = I − vv⊤ . More generally, for a linear subspace F of Rd , the projection from Rd to F is denoted by PF . Throughout the paper, Sd−1 denotes the unit sphere in Rd . We use 1{E} to denote the indicator of an event E. σ(z) := max{0, z}. We write E ≳ F for two nonnegative expressions E and F to denote that there exists some universal constant c > 0 (independent of the variables or parameters on which E and F depend) such that E ≥ cF . The expression E ≲ F is similarly defined. We use the notation g(t) ≤ poly(t) for a quantity t ≥ 1 to indicate that there exist constants c, C > 0 such that g(t) ≤ Ctc . Similarly, we use g(t) ≥ poly(t) for a quantity t < 1 to denote that there exist constants c, C > 0 such e e (the that g(t) ≥ Ctc . Finally, we use standard asymptotic notation O(·), Θ(·), Ω(·), and O(·), Ω(·) latter two suppress polylogarithmic factors in the argument). Complex Space. We use Cd to denote the d-dimensional complex vector space. For z ∈ Cd , z = x+iy where i is the imaginary unit, and z̄ = x − iy is the complex conjugate of z. We use Re(z) to denote the real part of z and Im(z) to denote its imaginary part. For a complex number z = a + ib, recall that |z|2 = a2 + b2 = z z̄. Matrices and Tensors. For a matrix M ∈ Rd×d , we use ∥M∥2 to denote its spectral/operator norm. For a symmetric matrix M, λi (M) denotes its ith largest eigenvalue and ri (M) denotes the corresponding (unit) eigenvector, i ∈ [d]. tr(M) denotes the trace of a matrix. Our analysis makes use of the following standard fact. Fact 2.1 (Wedin’s Theorem). For any pair of symmetric matrices M1 , M2 ∈ Rd×d , if λ1 (M1 ) − λ2 (M1 ) − ∥M1 − M2 ∥2 > 0, then sin(θ(r1 (M1 ), r1 (M2 ))) ≤

∥M1 − M2 ∥2 . λ1 (M1 ) − λ2 (M1 ) − ∥M1 − M2 ∥2

A complex tensor W ∈ (Cd )⊗p of dimension d and order p is a multilinear map defined by a p-index array with complex entries (W)α , where α = (i1 , . . . , ip ) with i1 , . . . , ip ∈ [d]. A tensor W is symmetric 9

if (W)α = (W)απ , where απ = (iπ(1) , . . . , iπ(p) ), for any permutation π : [p] → [p]. For complex tensors W, T of dimension d and order p, we denote by ⟨W, T⟩ their entrywise inner product: X ⟨W, T⟩ = (W)i1 ,...,ip (T)i1 ,...,ip . i1 ,...,ip

P For a complex tensor W, let ∥W∥2F = ⟨W, W⟩ = i1 ,...,ip |Wi1 ,...,ip |2 denote its Frobenius norm. Let U, V be finite-dimensional complex inner-product spaces. The Hermitian adjoint of a complex-linear map W : U → V is the unique complex-linear map W∗ : V → U satisfying ⟨W(A), B⟩ = ⟨A, W∗ (B)⟩ for any A ∈ U and B ∈ V. We overload the notation ⟨·, ·⟩ with partial tensor contraction: let W ∈ (Cd )⊗(p+m) , T ∈ (Cd )⊗p , ⟨W, T⟩ is a tensor in (Cd )⊗m defined by X (W)i1 ,...,ip ,j1 ,...,jm (T)i1 ,...,ip . (⟨W, T⟩)j1 ,...,jm = i1 ,...,ip

Let Sym(W) be the symmetrization operator: (Sym(W))i1 ,...,ip =

1 p!

X

(W)iπ(1) ,...,iπ(p) .

permutation π of [p]

Identify an order-p tensor T with the vector t = vec(T) of its entries. Its covariance is the Hermitian positive semidefinite matrix Cov(T) = E[(t − E[t])(t − E[t])∗ ], and its uncentered second moment is E[tt∗ ]. We say that its covariance is bounded above by V if Cov(T) ⪯ V I, which is equivalent to sup∥W∥F =1 E[|⟨W, T − E[T]⟩|2 ] ≤ V . For real-valued tensors, these definitions reduce to the usual real covariance and second-moment matrices after flattening. Gaussian Space and Hermite Analysis. Let N (µ, Σ) denote the d-dimensional Gaussian distribution with mean µ ∈ Rd and covariance Σ ∈ Rd×d . We will primarily use the standard normal distribution N (0, Id ), denoted simply by N . The standard Gaussian distribution on a linear subspace e ∼ N (0, IF ) and defined via x e = PF x, x ∼ N (0, Id ). The Gaussian L2 inner F of Rd is denoted by x product is ⟨f, g⟩ := Ex∼N [f (x)g(x)]. For f ∈ L2 (N ), define ∥f ∥L2 := (Ex∼N [f (x)2 ])1/2 . Fact 2.2 (Gaussian Hypercontractivity; see, e.g., [O’D14]). For any real polynomial f (z) of degree at most p and real q ≥ 2, it holds that (Ez∼N [|f (z)|q ])1/q ≤ (q − 1)p/2 (Ez∼N [f (z)2 ])1/2 . We now proceed to define Hermite polynomials and the associated orthonormal basis they define. In one dimension, define the normalized probabilists’ Hermite polynomial 1 dp hep (z) := √ (−1)p exp(z 2 /2) p exp(−z 2 /2). dz p!

(5)

The family {hep }p≥0 is an orthonormal basis for L2 (N (0, 1)). Thus, any f : R → R with f ∈ L2 (N (0, 1)) has a unique Hermite expansion X f (z) = fbp hep (z), fbp := E [f (z)hep (z)]. z∼N (0,1)

p≥0

We make use of a number of basic properties of Hermite polynomials summarized below. Fact 2.3 (Basic Hermite Facts, see, e.g., [Bog98, O’D14]). For normalized Hermite polynomials hep , as defined in (5), the following hold: √ 1. (Differentiation) (hep (z))′ = p hep−1 (z), for p ≥ 1. √ √ 2. (Degree Identity) zhep (z) = p + 1 hep+1 (z) + p hep−1 (z), for p ≥ 2. p 3. (Smoothing) Let f ∈ L2 (N (0, 1)), ρ ∈ [0, 1] and denote Tρ (f )(x) := Ez∼N (0,1) [f (ρx + 1 − ρ2 z)]. Then the smoothed function Tρ (f )(x) has Hermite expansion: X Tρ (f )(x) = ρp fbp hep (x). p≥0

10

4. (Hermite expansion in one specific direction) Let f ∈ L2 (N ) and u, v be unit vectors in Rd . Then: E x∼N (0,Id )

[f (v · x)hep (u · x)] = fbp cosp (θ(v, u)).

Throughout the paper, we denote by σ bi the ith Hermite coefficient of the ReLU function σ. These Hermite coefficients admit the following characterization. Fact 2.4 (Hermite expansion of ReLU; √ see, e.g., [GGK20, DK24]). For σ(z) := max{0, z}, we have P∞ σ(z) = p=0 σ bp hep (z), where σ b0 = 1/ 2π, σ b1 = 1/2, and σ b2p−1 = 0

for all p ≥ 2,

σ b2p = p

he2p−2 (0) 2π(2p)(2p − 1)

for all p ≥ 1.

2 −5/2 In particular, |b σ2p | = Θ(p−5/4 ) and σ b2p ) for p ≥ 1. Moreover, if we formally differentiate P= Θ(p′ √ ′ the Hermite series, we obtain σ (z) = p≥0 σ bp hep (z) with σ bp′ = σ bp+1 p + 1.

Hermite polynomials can be extended to Hermite tensors in Rd . Throughout this paper, we will use the following normalized Hermite tensor. Definition 2.5 (Normalized Probabilist’s Hermite Tensor). We define the normalized order-p Hermite tensor in d dimensions as follows: for any i1 , . . . , ip ∈ [d], 1 (Hep (x))i1 ,...,ip := √ p!

X

Y

Partitions P of [p] all size-2 sets {a, b} in partition P into sets of size 1 and 2

−(Id )ia ,ib

Y

(x)ic .

(6)

all size-1 sets {c} in partition P

We gather here important facts about Hermite tensors that are used in our analysis. Fact 2.6 (Properties of Normalized Hermite Tensor). Let Hep be the pth -order normalized Hermite tensor. The following holds: √ 1. Hep (x) = Ez∼N (0,Id ) [Re((x + iz)⊗p )/ p!], for all x ∈ Rd . 2. Let f ∈ L2 (N (0, 1)). For any v ∈ Sd−1 , Ez∼N (0,Id ) [f (v · z)He2 (z)] = Ez∼N (0,Id ) [f ′′ (v · z)]vv⊤ .

3

Optimization via Coarse Direction Recovery, Refinement, and Pruning

In this section, we provide the full technical details of our analysis, assuming access to a coarse direction recovery (CDR) procedure as an oracle; the implementation of CDR is presented in Section 4. Our oracle access to CDR is formally stated as the following assumption. Assumption 1 (Coarse Direction Recovery). Let d, k ∈ Z+ and λ, δ ∈ (0, 1). Given access to a Gaussian sum of k ReLUs distribution on Rd × R, a Coarse Direction Recovery algorithm CDR(D, λ, δ) draws NCDR (d, k, λ, δ) i.i.d. samples from D, runs in time TCDR (d, k, λ, δ), and with probability at least 1 − δ returns a list L ⊆ Sd−1 of size at most LCDR (d, k, λ, δ) such that: for every λ-heavy component i ∈ [k], there exists ri ∈ L satisfying |ri · wi | ≥ 0.99. At a high level, the CDR procedure postulated in Assumption 1 guarantees that if we had i.i.d. samples from a Gaussian sum of k ReLUs distribution D, then we could efficiently compute a not-toolarge list L of “candidate” unit vectors such that for each λ-heavy hidden weight vector wi there exists a candidate r ∈ L that is “close” to wi (as measured by the nontrivial correlation |r · wi | ≥ 0.99). The list may also contain other vectors with no specific guarantee. Given the existence of such an algorithm that runs in polynomial time (which we establish in Section 4), our main result in this section is a polynomial-time learner of an unknown sum of ReLUs, as stated in the following theorem.

11

Algorithm 1 LearnSumReLUs(d, k, ϵ; D) 1: Input: k, ϵ > 0, and access to a Gaussian sum of k ReLUs distribution D. 2: Output: a hypothesis y b such that E(x,y)∼D [(b y (x) − y)2 ] ≤ ϵ with probability at least 2/3. 3: ϵ0 ←

√

ϵ/(4k) , δ ← 1/100

▷ ϵ0 : error threshold of Parameters; δ: failure probability b ← CoverAndPrune(ϵ0 , δ; D) 4: V ▷ Returns a candidate list by repeatedly calling the CDR Subroutine. b = {v1 , . . . , vm } where m = |V|. b 5: Let V ▷ Enumerate the elements b of E(x,y)∼D [yx]. 6: Draw Θ(dk 2 /ϵ2 ) samples from D to construct an estimator g 10 2 b 7: Draw Θ(k /ϵ ) samples from D to formulate the empirical distribution D. Pm αi 2 P b ← argminα∈{αi ≥0: αi ≤k} E(x,y)∼Db [(y − g b · x − i=1 2 |vi · x|) ]. 8: α ▷ Convex Quadratic Program P b · x + i∈[m] (b 9: Return: y b(x) = g αi /2)|vi · x|.

Theorem 3.1 (Direction Covering Algorithm). Let ϵ ∈ (0, 1), k, d ∈ Z+ , and D be a Gaussian sum of k ReLUs distribution on Rd × R. Given access to a procedure CDR satisfying Assumption 1, there exists an algorithm (LearnSumReLUs(d, k, ϵ; D)) that repeatedly calls CDR with parameters λ = poly(ϵ/k) e poly(k/ϵ) + and δ = poly(ϵ/k) and attains the following performance guarantee: it draws N = O(d L3CDR (NCDR + L3CDR )poly(k/ϵ)) samples from D, runs in poly(d, N, TCDR ) time and, with probability at least 2/3, outputs a hypothesis yb(x) such that E(x,y)∼D [(y − yb(x))2 ] ≤ ϵ. The pseudocode for our main algorithm is displayed below. The body of this section is devoted to the proof of Theorem 3.1. In Section 3.1, we show that the main subroutine of our algorithm, CoverAndPrune(ϵ0 , δ; D), successfully covers all hidden weight vectors wi with non-trivial coefficients. Then in Section 3.2, we present the proof of Theorem 3.1.

3.1

Learning λ-Heavy Hidden Weight Vectors

Learning λ-heavy (i.e., nontrivial-coefficient) hidden weight vector described in this subsection is the technical core of the paper. We begin this section by first arguing that the CDR procedure can be used to construct updates to candidate vectors maintained as estimates of the hidden vectors wi . In particular, in Section 3.1.1, we argue that, given a fixed vector v ∈ Sd−1 and a simple to implement rejection sampling procedure (used for selecting samples from D), CDR procedure can be used to construct a polynomially-sized list L(v) that contains constant-correlation estimates of wi⊥v for λ-heavy ReLU components. On input v, our Refine(v, r, ϵ0 , δCDR ; D) procedure uses CDR to construct a list V of gradient descent-style updates to v such that for each λ-heavy wi , there exists at least one updated vector that contracts the distance between v and wi by a small constant factor. At the end of Refine(v, r, ϵ0 , δCDR ; D), only a small subset of updated candidate vectors in the output list V is useful to algorithm progress, that is, to improving the estimation of hidden λ-heavy vectors. The next step in the algorithm is thus to prune the list, after Refine(v, r, ϵ0 , δCDR ; D) is applied to all candidate vectors maintained by the algorithm. As a first coarse pruning step, Section 3.1.2 introduces a testing procedure that can be used to identify elements of V farther from all λ-heavy wi than a prespecified threshold. As a final step, Section 3.1.3 introduces Prune(V, r) procedure that prunes the list further so that only up to k candidate vectors are retained by the algorithm. A complete algorithm statement for learning λ-heavy hidden vectors is then provided in CoverAndPrune(ϵ0 , δ; D) and the main result of the subsection (Proposition 3.10) is proved. 3.1.1

Boosting Hidden Vector Estimates

We begin by arguing that, given any v ∈ Sd−1 , the CDR procedure can be called to obtain constantcorrelation estimates of wi⊥v , for any λ-heavy component i ∈ [k]. These estimates are then used to boost the estimation of wi , by treating estimates of wi⊥ as gradient descent directions from v. Of course, it is not immediately clear that CDR can be applied here, as Assumption 1 requires samples [ ⊥v := w⊥v /∥w⊥v ∥ as the hidden directions. from a Gaussian sum of k ReLUs distribution with w 2 i i i 12

For intuition on how CDR could be applied to estimating wi⊥v and boosting the estimate of wi , recall Pk first that a sample (x, y) drawn from our distribution D satisfies x ∼ N (0, Id ) and y = i=1 βi σ(wi · x). [ ⊥v , Fix a unit candidate v and a target w . Given v, we can express w as w = cos(θ ) v + sin(θ ) w 1

1

1

1

1

1

where θ1 = θ(w1 , v). Consider the samples (x, y) ∼ D that lie in a band of width γ around the origin, [ ⊥v namely satisfy |v · x| ≤ γ. Then, it is possible to argue that σ(w1 · x) ≈ σ(sin(θ1 )w 1 ) + O(γ). Thus, Pk [ ⊥v · x), which is another when γ → 0, for such samples we have that y ≈ i=1 βi sin(θ(wi , v))σ(w i [ ⊥v lying in the sum of ReLUs with coefficients b = β sin(θ(w , v)) and with hidden weight vectors w i

i

i

i

subspace v⊥ . Moreover, the component of x that lies in the subspace v⊥ follows the standard normal distribution (the rest is irrelevant, as y depends only on wi⊥v · x = wi⊥v · x⊥v ). This means that we can apply CDR to this new instance (obtained via conditioning) to obtain a new list of unit vectors L with [ ⊥v | ≥ 0.99. Then, updating v along the the following guarantee: there exists r ∈ L such that |r · w 1

direction r with a carefully chosen step-size, we obtain a new vector v′ such that |v′ · w1 | ≫ |v · w1 |, boosting the estimation of w1 . We now formalize this intuition. Rejection Sampling. For technical reasons, instead of using a “hard” band, i.e., conditioning on |v · x| ≤ γ for some γ > 0, we will use a soft-conditioning method based on rejection sampling. Soft-conditioning is more convenient because the marginal distribution remains a standard Gaussian while ensuring that the retained samples define an instance of the same learning problem (with different parameters). Formally, we use the following rejection sampling procedure: Fact 3.2 (see, e.g., Lemma 4.7 in [DKS18]). Given parameters v ∈ Sd−1 and γ ∈ (0, 1), the (γ, v)rejection procedure RSγ,v is defined as follows: A vector x ∈ Rd is accepted by RSγ,v with probability  −2  γ −1 2 := aγ,v (x) exp − (v · x) . 2 Let xγ denote the sample x conditioned on being accepted by RSγ,v . Then, for x ∼ N (0, Id ), we have that Prx∼N (0,Id ) [x is accepted] = γ and xγ ∼ N (0, Aγ,v ) with Aγ,v := Id − (1 − γ 2 )vv⊤ . Furthermore, for any δ > 0, in order to obtain N i.i.d. samples from xγ with probability at least 1 − δ, we need O(N log(N/δ)/γ) samples from N (0, Id ). For concreteness of discussion, we introduce the following notation and terminology for the conditional distribution of samples following the rejection sampling procedure from Fact 3.2. Definition 3.3 ((γ, v)-Rejection Sampling Distribution). Given a distribution D supported on Rd × R and samples (x, y) ∼ D, consider the rejection sampling procedure from Fact 3.2 applied to the covariates −1/2 x. Whenever x is accepted, accept the whole labeled example (x, y), and let z = Aγ,v x. The distribution of transformed accepted samples (z, y) is then called the (γ, v)-rejection sampling distribution obtained from D, and it is denoted by Dγ,v . We now argue that Dγ,v itself defines an instance of a sum of k ReLUs under Gaussian marginals. Fact 3.4 (Properties of Rejection Sampling). Fix parameters v ∈ Sd−1 and γ ∈ (0, 1). Let (x, y) ∼ D Pk be such that x ∼ N (0, Id ) and y(x) = i=1 βi σ(wi · x) with ∥wi ∥2 = 1 and βi ∈ [0, 1] for i ∈ [k]. Let −1/2 (xγ , yγ ) be the sample accepted by RSγ,v . Let z := Aγ,v xγ , where Aγ,v := Id − (1 − γ 2 )vv⊤ . Then, the following statements both hold: 1. z follows the standard normal distribution, i.e., z ∼ N (0, Id ). Pk e i · z), where wi = wi⊥v + γ(v · wi )v, w e i = wi /∥wi ∥2 , and βei = βi ∥wi ∥2 ≤ βi . 2. E[yγ | z] = i=1 βei σ(w 1/2

Proof. The acceptance decision depends only on x. Since xγ = Aγ,v z, we have E[yγ | z] =

k X

k   X 
⊤  βi σ wi⊤ A1/2 βi σ A1/2 z . γ,v z = γ,v wi

i=1

i=1

13

1/2

We can write wi = wi⊥v + (v · wi )v. Therefore, Aγ,v wi = wi⊥v + γ(v · wi )v = wi . Using positive e i · z), we get homogeneity of ReLU, namely σ(wi · z) = ∥wi ∥2 σ(w E[yγ | z] =

k X

e i · z) = βi ∥wi ∥2 σ(w

k X

e i · z). βei σ(w

i=1

i=1

Finally, βei ≤ βi because ∥wi ∥2 ≤ ∥wi ∥2 , as γ ∈ (0, 1). With Fact 3.4 on hand, we now formally argue that CDR can be applied to estimate wi⊥v . Proposition 3.5 (Estimating Update Directions). Fix a unit vector v ∈ Rd and let λ, δ ∈ (0, 1). Let D be a sum of k ReLUs under Gaussian marginals and, given γ = λ/100, let Dγ,v be the (γ, v)-rejection sampling distribution obtained from D. Then, procedure CDR(Dγ,v , λ, δ) that satisfies Assumption 1 applied to i.i.d. samples from Dγ,v returns a list L such that with probability at least 1 − δ, for any wi d ⊥v · w ⊥v with βi sin(θ± (wi , v)) ≥ λ, there exists a unit vector ri ∈ L with |rd i i | ≥ 9/10. Proof. Denote the projective angle between wi and v by θi = θ± (wi , v). From Fact 3.4, we have that E γ,v [y | x] =

(x,y)∼D

k X

e i · z) , βei σ(w

i=1

e i = wi /∥wi ∥2 , and βei = βi ∥wi ∥2 . Observe that with wi = wi⊥v + γ(v · wi )v, w [ [ [ ⊥v 2 ⊥v 2 ⊥v ei − w e i ∥22 + ∥w e i⊤ w ∥w i ∥2 = ∥w i ∥2 − 2 w i   ∥wi⊥v ∥2 =2 1− ∥wi⊥v + γ(v · wi )v∥2   γ 2 (v · wi )2 1 , ≤ =2 1− (1 + γ 2 (v · wi )2 /∥wi⊥v ∥22 )1/2 ∥wi⊥v ∥22

(7)

[ ⊥v are unit vectors and applies their definitions to e i, w where the second equality uses that both w i simplify the inner product, while the last inequality uses 1/(1 + z 2 )1/2 ≥ 1 − z 2 /2 (this can be seen as φ(x) = 1/(1 + x)1/2 , x ≥ 0, is convex, by applying φ(x) ≥ φ(0) + φ′ (0)x and plugging in x = z 2 ). By using that βi sin θi ≥ λ from our assumptions, we have that 1/ sin(θi ) ≤ βi /λ ≤ 1/λ, as βi ∈ (0, 1]. Combining with (7), we thus get that γ| cos(θi )| γ|v · wi | [ ⊥v ei − w = ≤ γ/λ ≤ 1/100. ∥w i ∥2 ≤ ⊥v | sin(θi )| wi 2 e i | ≥ 9.1/10 for Therefore, we have that CDR(Dγ,v , λ, δ) returns a list L containing ri ∈ L with |ri · w any wi satisfying βi sin(θi ) ≥ λ. Since, by the triangle inequality and Cauchy-Schwarz, [ [ [ ⊥v ⊥v ⊥v e i | ≤ |ri · w ei − w |ri · w i | + ∥ri ∥2 ∥w i ∥ ≤ |ri · wi | + 1/100, [ [ ⊥v · w ⊥v ⊥v we get that |rd i i | ≥ |ri · wi | ≥ 9.1/10 − 1/100 = 9/10. Since the rejection sampling procedure works by (approximately) localizing the marginal distribution of x to a thin band |v · x| ≤ γ, it is typically referred to as “localization”. The localization argument in Proposition 3.5 allows us to obtain components of the hidden weight vectors that our algorithm potentially has not yet discovered. The subroutine Refine(v, r, ϵ0 , δCDR ; D) formalizes the updating procedure to obtain vectors closer to the target ones. In Steps (5) through (11), we guarantee that we have created a new set that contains some vectors that are closer than before to some of the hidden vectors, by a multiplicative factor 3/5. We note that for technical reasons (Proposition 3.10 requires a large decrease between iterations), we require the update to decrease the angle by a factor larger than 3/5; for this purpose, in Refine(v, r, ϵ0 , δCDR ; D), we boost this procedure by applying three update steps instead of one.

14

Algorithm 2 Refine(v, r, ϵ0 , δCDR ; D) 1: Input: unit vector v ∈ Sd−1 , radius upper bound r ∈ (ϵ0 , π/2], accuracy ϵ0 > 0, failure probability

δCDR ∈ (0, 1), i.i.d. access to D. 2: Output: A list of vectors V satisfying the guarantees of Lemma 3.6. 3: χ ← 3/5, V ← {v}, λCDR ← ϵ0 sin(χ3 ϵ0 ), γ ← λCDR /100.

▷ λCDR : coefficient threshold; γ: rejection sampling width.

4: for j = 0, 1, 2 do

ηj ← tan(4χj r/5). for all v ∈ V do Construct Dγ,v by applying the rejection sampling procedure RSγ,v (Fact 3.2). L0 (v) ← CDR(Dγ,v , λCDR , δCDR ) ▷ Apply the CDR procedure that satisfies Assumption 1. d ⊥v 9: L(v) ← {q : q ∈ L0 (v)}. ▷ Project the vectors obtained from the initialization procedure to the tangent space of v then normalize to unit norm (Proposition 3.5). 10: for r ∈ L(v) do 11: Add projSd−1 (v + ηj r) and projSd−1 (v − ηj r) to V. ▷ Gradient descent/Refinement step 12: Return: V 5: 6: 7: 8:

The following lemma establishes that our Refine(v, r, ϵ0 , δCDR ; D) subroutine contracts the distance between the input candidate vector v and the λCDR -heavy hidden vectors wi , and bounds its sample and computational complexities. Lemma 3.6 (Properties of Refine). Fix the input parameters to Refine(v, r, ϵ0 , δCDR ; D): v ∈ Sd−1 , r ∈ [ϵ0 , π/2), ϵ0 ∈ (0, 1/2), δCDR ∈ (0, 1), and a Gaussian sum of k ReLUs distribution D, and let all other algorithm parameters be specified as in the algorithm statement. Suppose that for some i ∈ [k], βi ≥ ϵ0 and θ± (v, wi ) ≤ r. Then, with probability at least 1 − O(L3CDR δCDR ), the list V returned by Refine(v, r, ϵ0 , δCDR ; D) contains a vector vi ∈ V satisfying θ± (vi , wi ) ≤ χ3 r. e 3 NCDR log(1/δCDR )/γ) samples from D, runs Moreover, Refine(v, r, ϵ0 , δCDR ; D) draws NRefine = O(L CDR in time d poly(NRefine , TCDR ) and the size of the output list V is bounded by |V| = O(L3CDR ), where LCDR , NCDR , TCDR are the number of vectors returned, the sample complexity, and the runtime of CDR(Dγ,v ; λCDR , δCDR ), respectively. Proof. Recall that χ := 3/5. If θ± (v, wi ) ≤ χ3 r, then v already satisfies the condition. Otherwise, since we have that βi ≥ ϵ0 and r ≥ ϵ0 , we get that βi sin θ± (v, wi ) ≥ βi sin(χ3 r) ≥ ϵ0 sin(χ3 ϵ0 ) = λCDR . We argue that after each step j ∈ {0, 1, 2}, there exists a vector vi ∈ V such that χj+1 r ≥ θ± (vi , wi ). Before the algorithm starts, the input vector v satisfies θ± (v, wi ) ≤ r, by our assumption. Since βi sin θ± (v, wi ) ≥ λCDR , Proposition 3.5 implies that running CDR with NCDR samples from Dγ,v returns a list L of unit vectors that contains a vector r such that r ⊥ v and [ ⊥v |r · w i | ≥ 0.9 . The main insight is that r defines a direction that can improve the estimate of the hidden direction wi . In particular, as argued in the next claim, a projected gradient descent-style update of a candidate vector v along either r or −r leads to a new unit vector that is better aligned with (and thus closer to) wi , unless the vector v was already well-aligned with (i.e., sufficiently close to) wi . Claim 3.7 (One-step improvement in projective angle). Let 0 < r ≤ π/2. Let u, w ∈ Sd−1 be such that d ⊥u | ≥ 0.9. Let ϕ := 4r , η := tan ϕ, and define θ± (u, w) ≤ r. Let r ∈ u⊥ be a unit vector such that |r · w 5 u− := projSd−1 (u − ηr).

u+ := projSd−1 (u + ηr), Then minz∈{u,u+ ,u− } θ± (z, w) ≤ 3r/5 = χr. 15

Proof of Claim 3.7. Let θ := θ± (u, w). Assume without loss of generality that θ(u, w) ≤ π/2. If θ ≤ 3r/5, then the vector u already satisfies the conclusion. Therefore, we assume throughout the proof that 3r/5 < θ ≤ r. [ ⊥u and note that q · u = 0 and w = cos θ u + sin θ q. Recall that |r · q| ≥ 0.9. Assume Let q = w without loss of generality that r·q ≥ 0.9 and consider u′ = projSd−1 (u+ηr), as the other case is captured by the update rule u− with exactly the same analysis. Letting δ = θ(r, q), then δ ≤ arccos(0.9) < 1/2. Consider the ideal update, i.e., the update we could take if we knew exactly the unknown part q: z := cos ϕ u + sin ϕ q . Note that the angle between u and z is ϕ, and the angle between u and w is θ. Therefore θ(z, w) = |θ − ϕ| ≤ |θ − 4r/5| ≤ r/5 ,

(8)

where we used that ϕ = 4r/5 and r ≥ θ ≥ 3r/5. Recall that u′ , the output of the update, was defined by u′ = projSd−1 (u + ηr) , and observe that u + tan ϕr u′ = p = cos ϕ u + sin ϕ r . 1 + tan2 ϕ Next we bound θ(u′ , w). For that, we need to bound θ(u′ , z). Note that ∥u′ − z∥2 = 2 sin(θ(u′ , z)/2). Furthermore, by construction, ∥u′ − z∥2 = sin ϕ∥q − r∥2 = 2 sin ϕ sin(δ/2). Therefore, we have that sin(θ(u′ , z)/2) ≤ sin ϕ sin(δ/2). Noting that the sine function is concave on (0, π/2), from Jensen’s inequality we have sin((δ/2) ϕ) ≥ (δ/2) sin ϕ + (1 − δ/2) sin(0). Using that sin(δ/2) ≤ δ/2, we get that sin ϕ sin(δ/2) ≤ δ/2 sin(ϕ) ≤ sin(ϕδ/2). Therefore, from monotonicity, we get that θ(u′ , z) ≤ δϕ ≤ 2r/5 .

(9)

Since θ(u, v) is a metric on the unit sphere, by the triangle inequality, and using (8) and (9), we get θ(u′ , w) ≤ θ(u′ , z) + θ(z, w) ≤

2r r 3r + = . 5 5 5

Since u′ ∈ {u+ , u− }, this gives minz∈{u,u+ ,u− } θ± (z, w) ≤ 3r 5 , completing the proof of Claim 3.7. By Claim 3.7, the list V contains a vector vi such that χr ≥ θ± (vi , wi ). Hence, after step j = 0, we have χr ≥ θ± (vi , wi ) for some vector vi ∈ V. This is the base case of the induction argument. Now at step j = {1, 2}, by induction, we know that before updating, there exists a vector vi ∈ V such that χj r ≥ θ± (vi , wi ). Suppose θ± (vi , wi ) ≥ χj+1 r, which implies βi sin θ± (vi , wi ) ≥ λCDR . Then, using Proposition 3.5 again, the CDR procedure returns a list L(vi ) of unit vectors that contains a [ ⊥vi vector r such that r ⊥ v and |r · w | ≥ 0.9. Then Claim 3.7 implies that after updating v , inside the i

i

i

list V there is a vector vi′ so that χj+1 r ≥ θ± (vi′ , wi ). Therefore, at step j = 2, the algorithm produces a vector v ∈ V such that θ± (v, wi ) ≤ χ3 r. At each call, CDR uses NCDR samples from Dγ,v . By Fact 3.2, with probability at least 1 − δCDR , these NCDR samples from Dγ,v can be obtained by drawing O(NCDR log(NCDR /δCDR )/γ) samples from the original distribution D, and passing through the (γ, v)-rejection procedure RSγ,v . Conditioned on successfully obtaining these NCDR from Dγ,v , the CDR procedure succeeds with probability at least 1 − δCDR , by Assumption 1 and Proposition 3.5. Therefore, the total number of samples required by CDR is O(NCDR log(NCDR /δCDR )/γ), and the probability of success for one call of CDR procedure is at least (1 − δCDR )2 ≥ 1 − 2δCDR . Recall that for each query vector v ∈ V, CDR returns a list of |L(v)| ≤ LCDR vectors. Therefore, the CDR gets called T = (2|L|)3 = O(L3CDR ) times. Thus, in total, Refine(v, r, ϵ0 , δCDR ; D) uses T NCDR = O(L3CDR NCDR log(NCDR /δCDR )/γ) samples from D, and succeeds with probability 1 − 2T δCDR = 1 − O(L3CDR δCDR ). The size of V is |V| = O(L3CDR ). This completes the proof of Lemma 3.6. 3.1.2

Closeness Tester

The Refine procedure (Refine(v, r, ϵ0 , δCDR ; D)) returns a large list V of vectors with the guarantee that for each λ-heavy hidden weight vector wi , there is at least one vector v ∈ V that is sufficiently close, 16

i.e., θ± (wi , v) ≤ r for some r > 0 (Lemma 3.6). However, Refine branches over all possible updates, because it has no information on which hidden weight vector the input v approximates nor which update directions from CDR are useful for improving the approximation. As a result, the output list may contain many irrelevant vectors that provide no useful approximation of any of the hidden weight vectors. To eliminate irrelevant candidate vectors, we design a tester that distinguishes between v ∈ V that are r-close to at least one heavy hidden weight vector versus (1 + τ )r-far from all hidden weight vectors, for constant τ > 0. For an even degree p, we use the following moment test (see Fact 2.3): ap (v) :=

X 1 E[yhep (v · x)] = βi |v · wi |p . σ bp i

The main intuition behind the tester is that the small value of ap (v) identifies vectors v that are far from all hidden directions wi . To see this, recall first that |v · wi | = cos(θ± (wi , v)). If v is r-close to at least one ϵ0 -heavy wi (i.e., if θ± (wi , v) ≤ r), then ap (v) ≥ ϵ0 cosp (r). On the other hand, if θ± (wi , v) > (1 + τ )r for all i ∈ [k], then |v · wi | = cos(θ± (wi , v)) ≤ (1 − τ r2 /2) cos(r) for all i ∈ [k], and so ap (v) ≤ k(1 − τ r2 )p cosp (r). Choosing p sufficiently large (p ≫ log(k/ϵ0 )/(τ r2 )) then ensures ap (v) ≪ ϵ0 cosp (r), so the two cases can be separated by the tester. Algorithm 3 Test(v, r, ϵ0 , τ, δTest ; D) 1: Input: Unit vector v, ϵ0 , δTest > 0, τ > 0, upper bound r, sample access to D. 2: Output: TRUE if there exists i ∈ [k] such that θ± (wi , v) ≤ r and βi ≥ ϵ0 ; FALSE if for all i ∈ [k]:

θ± (wi , v) > (1 + τ )r. 3: p ← 2⌈(1/(τ r 2 )) log(k/ϵ0 )⌉.

b 4: Draw N = poly(k, 1/ϵ0 )(1/δTest ) samples from D to formulate the empirical distribution D. 5: b ap ← (1/b σp ) E(x,y)∼Db [yhep (v · x)]. 6: if b ap ≥ (7/8)ϵ0 cosp (r) then 7: Return: TRUE. 8: else 9: Return: FALSE.

Lemma 3.8 (Smallest Angle Tester). Let ϵ0 ∈ (0, 1/2], r ≥ ϵ0 /(5k), let τ > 0 be a small absolute constant, and define p = 2⌈(1/(τ r2 )) log(k/ϵ0 )⌉. Using N = poly(k, 1/ϵ0 )(1/δTest ) samples and O(dN ) time with probability at least 1 − δTest , Test(v, r, ϵ0 , τ, δTest ; D) has the following performance guarantee: 1. (Completeness) If there exists i ∈ [k] such that θ± (wi , v) ≤ r and βi ≥ ϵ0 , then Test(v, r, ϵ0 , τ, δTest ; D) returns TRUE. 2. (Soundness) If mini∈[k] θ± (wi , v) > (1 + τ )r, then Test(v, r, ϵ0 , τ, δTest ; D) returns FALSE. Proof. Denote θi := θ± (wi , v) and suppose without loss of generality that θ1 = mini∈[k] θi and β1 ≥ ϵ0 . For even degree p, by Fact 2.3, we have that ap :=

k k X 1 1 X E [yhep (v · x)] = E [βi σ(wi · x)hep (v · x)] = βi cosp θi , σ bp (x,y)∼D σ bp i=1 x∼N i=1

where σ bp is the pth Hermite coefficient of the activation σ. We first prove the following simple claim, establishing a bound between the empirical estimate b ap and the population mean ap . Note that in Claim 3.9, the number of samples required for our tester has polynomial dependence on 1/δTest . This can be replaced by logarithmic dependence using median-of-means estimation [NY83] instead of the simple empirical estimator used in Test(v, r, ϵ0 , τ, δTest ; D). Yet, for our purposes, a polynomial dependence is sufficient. Claim 3.9. With probability at least 1 − δTest , the empirical estimator b ap of ap that uses N = poly(k, 1/ϵ0 )(1/δTest ) samples satisfies |b ap − ap | ≤ (1/8)ϵ0 cosp (r).

17

Proof of Claim 3.9. Note that b ap is an unbiased estimator of ap . The proof of the claim follows by bounding the variance of the estimate and applying Chebyshev’s inequality. We start by noting that the variance of yhep (v · x) is bounded by: E[y 2 hep (v · x)2 ] ≤ k

k X i=1

E [σ(wi · x)2 hep (v · x)2 ] ≤ k

x∼N

k X i=1

E [(wi · x)2 hep (v · x)2 ] .

x∼N

(10)

√ √ Using the standard Hermite identity zhep (z) = p + 1hep+1 (z) + phep−1 (z) (see Fact 2.3) and the orthogonality of Hermite polynomials, we get that E [(v · x)2 hep (v · x)2 ] = (p + 1) E [hep+1 (v · x)2 ] + p E [hep−1 (v · x)2 ] = 2p + 1.

x∼N

x∼N

x∼N

Therefore, by (10), since p ≥ 2, we have: E[y 2 hep (v · x)2 ] ≤ k

k X i=1

≤k

[ ⊥v · x)2 he (v · x)2 ] E [(cos θ(wi , v)v · x + sin θ(wi , v)w p i

x∼N

k X [ ⊥v · x)2 ] E [he (v · x)2 ]) (2 E [(v · x)2 hep (v · x)2 ] + 2 E [(w p i x∼N

i=1

= 2k

k X

x∼N

x∼N

((2p + 1) + 1) ≤ 6k 2 p.

i=1

Since p is an even integer, σ bp ̸= 0. From Chebyshev’s inequality, we get Pr[|b ap −

E (x,y)∼D

[(1/b σp )yhep (v · x)]| ≥ t] ≤

E(x,y)∼D [y 2 hep (v · x)2 ] 6k 2 p ≤ . σ bp2 N t2 Nσ bp2 t2

Choosing t = (1/8)ϵ0 cosp (r) and N = (6k 2 p)/(b σp2 t2 δTest ) ≲ k 2 p/(b σp2 ϵ20 cos2p (r)δTest ) then leads to p |b ap − ap | ≤ (1/8)ϵ0 cos (r) with probability at least 1 − δTest . Rθ Recall that σ bp2 ≳ p−3 (see Fact 2.4) and the elementary identity cos θ = exp(− 0 tan(s) ds) ≥ exp(−θ2 ). Therefore, when τ is chosen as an absolute positive constant, plugging in our choice of p = 2⌈(1/(τ r2 )) log(k/ϵ0 )⌉, and the assumption r ≥ ϵ0 /(5k), we get that   2    2+4/τ  4  10+4/τ 1 k p 1 k log (k/ϵ0 ) k 2 4 N≲ exp(2pr ) ≲ p ≲ . 2 δTest ϵ0 σ bp δTest ϵ0 δTest ϵ0 This completes the proof of Claim 3.9. We henceforth condition on the event of Claim 3.9 and we are ready to prove the two statements of the lemma, i.e., the completeness and the soundness of the tester. Note that a direct calculation yields for θi > r:     Z θi cosp θi p(θi2 − r2 ) = exp − p tan(s) ds ≤ exp − . cosp (r) 2 r Completeness.

Assume first that θ1 ≤ r and β1 ≥ ϵ0 . Then, we have: ap =

k X

βi cosp θi ≥ β1 cosp θ1 ≥ ϵ0 cosp (r).

i=1

By Claim 3.9, this implies b ap ≥ (7/8)ϵ0 cosp (r). Thus, when θ1 ≤ r and β1 ≥ ϵ0 , Algorithm 3 would return TRUE. Soundness. ap =

k X i=1

Now assume (1 + τ )r < θ1 ≤ θi for all i ∈ [k]. Then, we have p

βi cos θi ≤

k X

2

2

2

cosp θi ≤ ke−p(θi −r )/2 cosp (r) ≤ ke−p(τ r ) cosp (r) ≤

i=1

1 ϵ0 cosp (r), 2

where the last inequality is due to our choice of p = 2⌈(1/(τ r2 )) log(k/ϵ0 )⌉. Similarly, by Claim 3.9, this implies b ap ≤ (5/8)ϵ0 cosp (r). Thus, when no angle θi is smaller than (1 + τ )r, Algorithm 3 would return FALSE. 18

(a) After Refine Vtnew

(b) After Test Vtcloser

(c) After Prune Vt+1

(radius r)

(separation 2gr)

Depth t − 1 Vt−1 (t)

Depth t Vt

Depth t + 1 next level

(t)

vi

qi Refine call. θ± (qi , wi ) ≤ r for each heavy target i.

candidate

(t)

vi

rejected by test

vi

qi

qi

Tested at radius r. Survived vectors lie within (1 + τ )r of some target (possibly light). removed by pruning

retained

zi

Rt+1 = (2g + 1)r. Only these leaves branch next.

neighbour within 2gr

Figure 2: List of candidate vectors over the different stages of CoverAndPrune. Left: after Refine is called, each v ∈ Vt produces a large list of new candidate vectors; all produced lists combined contain, for each λ-heavy wi , at least one candidate vi with θ± (vi , wi ) ≤ r. Center: Test procedure is used to prune all candidates v far from all wi , that is, all v such that θ± (v, wi ) > (1 + τ )r, ∀i ∈ [k]. Right: Prune procedure makes a single iterative pass through the candidate vectors that survive after Test, forming the output list Vt+1 ; it starts with an empty list and adds a new iterate if and only if it is at least r-far from all the candidates already in the list. 3.1.3

Covering and Pruning

The approach described thus far, based on the Refine subroutine and the Test procedure, provides a concrete way of constructing lists of candidate vectors that progressively improve approximations of ϵ0 -heavy hidden weight vectors. In particular, suppose that the algorithm starts with a list of candidates Vt such that for each ϵ0 -heavy wi , i ∈ [k], there exists at least one vi ∈ Vt such that θ± (vi , wi ) ≤ Rt . We say in this case that wi is covered by vi at resolution Rt . A call to Refine for each vi ∈ Vt creates a new list of candidate vectors, within which only one candidate vector is guaranteed to land significantly closer to wi , with θ± (vi , wi ) ≤ χ3 Rt . However, it is unclear how to identify such a candidate. What is possible to identify are the candidate vectors that are far from all hidden weight vectors; this is where our Test procedure is used. Thus, Refine and Test can be used jointly to construct a list Vt′ such that all ϵ0 -heavy wi , i ∈ [k], are covered at resolution slightly larger than χ3 Rt (as testing increases the resolution slightly), but still a constant factor smaller than Rt . Since the resolution improves by a constant factor, repeating this combination of Refine and Test geometrically reduces the resolution and leads to covering all ϵ0 -heavy hidden weight vectors at resolution ϵ/k within order-log(k/ϵ) repetitions. Unfortunately, the described approach does not run in polynomial-time, since there is no guarantee that the number of candidate vectors in our list covering the heavy hidden weight vectors is small; as a result, as described, the list Vt could grow super-polynomially. To prevent such a size blow up, we additionally introduce a simple pruning procedure (Prune(V, r)). The procedure takes as input the list Vt′ produced after Refine and Test are used. It then creates Vt+1 , starting from an empty list as follows. It traverses Vt′ element by element, and adds the traversed element to Vt+1 if and only if it is sufficiently far from all elements already in Vt+1 . The intuition behind this method is that if two vectors v, u are very close to each other, then both of them have similar correlation with the hidden weight vectors. Thus, we can remove one of them without substantially altering the coverage of the hidden weight vectors. The whole sequence of steps following Refine, Test, and Prune is illustrated in Figure 2. We argue that the last pruning step (Prune(V, r)) ensures that all heavy hidden weight vectors

19

remain covered by vectors in Vt+1 at a slightly larger resolution Rt+1 ≤ ρRt , ρ ∈ (0, 1), but only up to k candidate vectors are added to Vt+1 . More broadly, the main procedure of this section, CoverAndPrune(ϵ0 , δ; D), guarantees the following: in each iteration t ≥ 0 of CoverAndPrune, the following invariant holds: (t)

(t)

for all ϵ0 -heavy wi , there exists vi ∈ Vt such that θ± (vi , wi ) ≤ Rt , and |Vt | ≤ k.

(Invariantt )

Algorithm 4 Prune(V, r) 1: Input: A finite set V ⊆ Sd−1 and a separation radius r > 0. 2: Output: A finite set Voutput ⊆ Sd−1 such that every pair of elements is separated by r in θ± . 3: Voutput ← ∅. 4: for v ∈ V do 5: if θ± (v, v′ ) > r for every v′ ∈ Voutput then 6: Voutput ← Voutput ∪ {v}. 7: Return:

Voutput .

Algorithm 5 CoverAndPrune(ϵ0 , δ; D) 1: Input:

Coefficient threshold ϵ0 > 0, failure probability δ ∈ (0, 1), sample access to D.

b of unit vectors satisfying Proposition 3.10. 2: Output: A list V 3: Tester accuracy τ ← 1/200, ratio g ← 1 + 2τ , radius shrinking rate χ ← (3/5). 4: Initial radius r0 ← arccos(9/10), error threshold λCDR ← ϵ0 sin(χ3 ϵ0 ). 5: Relaxed initial radius R0 ← (2g + 1)r0 , relaxed radius shrinking rate ρ ← (2g + 1)χ3 . log(R /ϵ )

0 0 6: Number of iterations T ← ⌈ log(1/ρ) ⌉. 7: Failure probability of immediate steps δ ′ ← O(δ/(T kL3CDR )) 8: L ← CDR(D, λCDR , δ ′ ). ▷ Apply the Coarse Direction Recovery

9: L′ ← {v ∈ L : Test(v, r0 , ϵ0 , τ, δ ′ ; D) = TRUE}. 10: V0 ← Prune(L′ , 2gr0 ).

procedure. ▷ Initialized vector set satisfying the induction (Invariantt ).

11: for t = 0, 1, . . . , T − 1 do 12: rt+1 ← χ3 Rt and Vtnew ← Vt . 13: for v ∈ Vt do

Vtnew ← Vtnew ∪ Refine(v, Rt , ϵ0 , δ ′ ; D).

▷ “Gradient Descent/Refinement” step; Vectors in Vtnew have smaller distance to the hidden weight vectors compared to Vt 15: Vtcloser ← {v ∈ Vtnew : Test(v, rt+1 , ϵ0 , τ, δ ′ ; D) = TRUE}. ▷ Discard candidate vectors that do not cover hidden weight vectors. 16: Vt+1 ← Prune(Vtcloser , 2grt+1 ). ▷ Discard candidate vectors covering the same hidden weight vectors. 17: Rt+1 ← (2g + 1)rt+1 . ▷ Vt+1 satisfies (Invariantt ) at t + 1 with parameter Rt+1 . b ← VT . 18: Return: V 14:

We briefly discuss how CoverAndPrune(ϵ0 , δ; D) preserves the induction invariant at each step. The refinement subroutine in Line 14 generates new vectors that are closer to the hidden weights, but also produces many redundant and unwanted candidates. In order to guarantee that the vectors we keep cover the hidden weight vectors within Rt , we use our testing routine from the previous subsection (Test(v, r, ϵ0 , τ, δTest ; D)). This procedure, applied in Line 15, allows us to discard candidate vectors that 20

do not have any hidden weight vector within radius Rt (in θ± ). Then in Line 16, the Pruning subroutine removes redundant candidates, since two well-separated vectors cannot both be near the same hidden weight vector. By the triangle inequality, after the pruning step, we guarantee that each candidate vector covers hidden weight vectors within a relaxed radius, Rt+1 = (2g + 1)rt+1 = ρRt , with ρ < 1. Thus, as a result, at most k candidates remain in Vt+1 , and every hidden weight with non-negligible coefficient is covered by a candidate vector within radius Rt+1 , i.e., the invariant (Invariantt ) is satisfied after each update. Proposition 3.10 (Covering and Pruning). Assume access to an algorithm CDR that satisfies Assumption 1. Given parameters ϵ0 , δ > 0, CoverAndPrune(ϵ0 , δ; D) draws e 3 (L3 + NCDR )) poly(k, 1/ϵ0 , 1/δ) N = O(L CDR CDR b such that with probability at i.i.d. samples from D, runs in poly(d, N, TCDR ) time, and returns a list V least 1 − δ, the following holds: b ≤ k; and 1. |V| b such that θ± (b b∈V 2. For all ϵ0 -heavy wi , there exists v v, wi ) ≤ ϵ0 . Proof. Let g, χ, r0 , R0 , and ρ be defined as in CoverAndPrune(ϵ0 , δ; D). For iteration t, let Rt+1 = ρRt = (2g + 1)χ3 Rt = (2g + 1)rt+1 . By Assumption 1, we have that running the CDR routine on D returns a list of unit vectors L such that for each heavy wi , i ∈ [k], there exists vi ∈ L satisfying θ± (vi , wi ) ≤ arccos(9/10) = r0 . After Algorithm 5, the algorithm CoverAndPrune(ϵ0 , δ; D) keeps only the vectors that pass the smallest-angle test. These are the candidate vectors that cover the heavy hidden weight vectors within radius r0 , i.e., L′ := {v ∈ L : Test(v, r0 , ϵ0 , τ, δTest ; D) = TRUE} , and then in Algorithm 5 prunes them to obtain the initial list V0 : V0 := Prune(L′ , 2gr0 ) . Our goal is to show that for each t ∈ [T ], (Invariantt ) holds. We prove this claim by induction on t. Claim 3.11 (Induction Base). At t = 0, (Invariantt ) condition is satisfied. Proof. By Proposition 3.5, the call to CDR in Line 8 gives a list L such that for every ϵ0 -heavy wi , there exists vi ∈ L with the property that θ± (vi , wi ) ≤ r0 . Further, after the testing step in Algorithm 5, for any ϵ0 -heavy wi , any vector vi ∈ L such that θ± (vi , wi ) ≤ r0 will be kept in L′ , by Lemma 3.8. If vi survives pruning in Algorithm 5, it covers wi with radius r0 (i.e., θ± (vi , wi ) ≤ r0 ). If, on the other hand, vi is discarded, then there must exist some other vector vi′ ∈ V0 that satisfies θ± (vi′ , vi ) ≤ 2gr0 . Hence, by the triangle inequality, we get that θ± (vi′ , wi ) ≤ θ± (vi′ , vi ) + θ± (vi , wi ) ≤ 2gr0 + r0 = R0 . This establishes the second part of the claim. It remains to bound the size of the list. Every vector in V0 passed the smallest-angle test with cutoff r0 . By Lemma 3.8, every such vector is within (1 + τ )r0 of some hidden weight vector wj with respect to θ± . If two retained vectors vj , vj′ cover the same wj with radius (1 + τ )r0 , then, since g = 1 + 2τ , we would have θ± (vj′ , vj ) ≤ θ± (vj′ , wj ) + θ± (wj , vj ) ≤ 2(1 + τ )r0 < 2gr0 contradicting the pruning rule in Algorithm 5. Therefore, we have |V0 | ≤ k, completing the proof of Claim 3.11. 21

Now it remains to show that the invariant holds for all t > 0. This is done in the following claim. Claim 3.12 (Inductive Step). Assume that (Invariantt ) holds at step t. Then, conditioned on successful Refine and Test procedures in Algorithm 5, it also holds at step t + 1. Proof. Fix an ϵ0 -heavy wi and choose v ∈ Vt such that θ± (v, wi ) ≤ Rt . Such a v is guaranteed to exist because (Invariantt ) holds for step t. In Algorithm 5 of CoverAndPrune(ϵ0 , δ; D), the Refine subroutine (Refine(v, r, ϵ0 , δCDR ; D)) is called with radius Rt . Therefore, by Lemma 3.6, it gives a list Vtnew that contains a unit vector v′ ∈ Vtnew such that θ± (v′ , wi ) ≤ χ3 Rt = rt+1 (since wi is ϵ0 -heavy, βi ≥ ϵ0 ). Thus, by Lemma 3.8, v′ will pass the test in Algorithm 5, i.e., Test(v, rt+1 , ϵ0 , τ, δ ′ ; D), and v′ ∈ Vtcloser . It remains to show that after pruning (Algorithm 5), Vt+1 covers all ϵ0 -heavy wi ’s at resolution Rt+1 . Recall that, in Algorithm 5, Prune(V, r) is applied to V = Vtcloser with separation radius r = 2grt+1 . Fix any ϵ0 -heavy wi and let vi ∈ Vtcloser be such that θ± (vi , wi ) ≤ rt+1 = χ3 Rt (we have argued above that this vector will pass the smallest-angle test). If vi ∈ Vt+1 , it covers wi at scale χ3 Rt ≤ Rt+1 . If it is discarded, then some retained vector vi′ ∈ Vt+1 satisfies θ± (vi′ , vi ) ≤ 2grt+1 . Hence, θ± (vi′ , wi ) ≤ θ± (vi′ , zi ) + θ± (vi , wi ) ≤ (2g + 1)rt+1 = Rt+1 . It remains to bound |Vt+1 |. Similar to the base argument, since every vector in Vt+1 passed the smallest-angle test in Algorithm 5, by Lemma 3.8, for every vj ∈ Vt+1 there is a hidden weight vector wj such that θ± (vj , wj ) ≤ (1 + τ )rt+1 . For the sake of contradiction, suppose there were two distinct vectors vj , vj′ both covering wj retained in Vt+1 . Then, by the triangle inequality, θ± (vj′ , vj ) ≤ 2(1 + τ )rt+1 < 2grt+1 , indicating that one of them will not pass the pruning step in Algorithm 5. This provides the desired contradiction and establishes that |Vt+1 | ≤ k. To complete the proof of Proposition 3.10, it remains to bound the sample and computational complexities. Note that Rt = ρt R0 . Choosing   log(R0 /ϵ0 ) , T = log(1/ρ) bi ∈ VT such that θ± (b we have RT ≤ ϵ0 . Thus, for all ϵ0 -heavy wi , there exists v vi , wi ) ≤ ϵ0 . Let LCDR be the uniform upper bound on the length of the list returned by the CDR subroutine, NCDR be the corresponding sample complexity, and δCDR be the CDR subroutine’s failure probability. At each iteration t ∈ [T ], the smallest angle tester (Test(v, r, ϵ0 , τ, δTest ; D)) is called |Vtnew | = |V t | × |number of vectors returned by Refine| = O(kL3CDR ) many times, where we used the fact that number of vectors returned by Refine = O(L3CDR ) from Lemma 3.6. Therefore, by a union bound over the success events of the tester routine, the probability that all Test calls succeed at step t is 1 − O(kL3CDR δTest ). Meanwhile, at step t, the Refine subroutine is called |Vt | ≤ k times. Therefore, by a union bound over all the calls of Refine for all T rounds and by Lemma 3.6, with probability at least 1 − O(T kL3CDR δCDR ), the Refine routine successfully returns a list Vtnew containing new candidate vectors that are closer to the hidden weight vectors compared to the ones in Vt , at each step t ∈ [T ]. Thus, by setting δCDR = δ ′ = O(δ/(T kL3CDR )) and δTest = δ ′ = O(δ/(T kL3CDR )), the subroutine CoverAndPrune(ϵ0 , δ; D) succeeds with probability at least 1 − O(T kL3CDR (δCDR + δTest )) ≥ 1 − δ. By Lemma 3.8, it follows that the tester Test succeeds with (1/δTest )poly(k/ϵ0 ) = poly(k/ϵ0 )T kL3CDR /δ e 3 NCDR log(1/δCDR )/γ) samsamples for each run. By Lemma 3.6, the Refine subroutine requires O(L CDR 2 ples from D, where γ = λCDR /100 = Θ(ϵ0 ). Therefore, in total, we need N ≳ number of Test runs × NTest + number of Refine runs × NRefine ≳ T kL3CDR × poly(k/ϵ0 )T kL3CDR /δ + T k × poly(1/ϵ0 )L3CDR NCDR log(1/δCDR ) ≳ log2 (1/ϵ0 )poly(k/ϵ0 , 1/δ)L3CDR (L3CDR + NCDR ) samples. By Lemma 3.6, Refine has runtime poly(d, 1/ϵ0 , LCDR , NCDR , TCDR ) and Test has runtime O(dNTest ). Thus, in summary, CoverAndPrune(ϵ0 , δ; D) has runtime poly(d, k/ϵ0 , 1/δ, LCDR , NCDR , TCDR ) = poly(d, N, TCDR ). This completes the proof of Proposition 3.10. 22

3.2

Proof of Theorem 3.1

We are now ready to prove the main result of this section. √ Proof of Theorem 3.1. Let ϵ0 = ϵ/(4k). Run CoverAndPrune(ϵ0 , δ; D) with failure probability 1/9. By b = {v1 , . . . , vm }, where m ≤ k, such that for Proposition 3.10, with probability at least 8/9, it returns V b such that θ± (wi , v) ≤ ϵ0 . We condition on this event. Assume every i with βi ≥ ϵ0 there exists v ∈ V without loss of generality that all the weights βi ≥ ϵ0 , as otherwise, ignoring these ReLUs will only increase the L22 error by ϵ. Indeed, we have the following sequence of inequalities:  X 2  X 2 E [(y − yb(x)) ] = E βi σ(wi · x) − yb(x) + βi σ(wi · x) (x,y)∼D

x∼N (0,Id )

≤2

i:βi ≥ϵ0

 X

E x∼N (0,Id )

≤2

i:βi ≤ϵ0

2   X 2  βi σ(wi · x) − yb(x) + 2E βi σ(wi · x)

i:βi ≥ϵ0

i:βi ≤ϵ0

[(y ′ (x) − yb(x))2 ] + 2k

E

x∼N (0,Id )

≤2

k X

βi2 E[(σ(wi · x))2 ]

i:βi ≤ϵ0 ′

2

[(y (x) − yb(x)) ] + ϵ ,

E

x∼N (0,Id )

P where y ′ (x) = i:βi ≥ϵ0 βi σ(wi · x). Observe that, by standard concentration arguments, with Ng = O(dpoly(k)/ϵ2 ) samples we can b of the vector g = E[yx] up to ℓ2 error ϵ/16 with probability at least 8/9. Since get an estimate g E[σ(w · x)x] = w/2 for every unit vector w, we have that k

1X y−g·x= βi |wi · x|. 2 i=1 Let σabs (z) = σ(z) − z/2 = |z|/2. Then √ √ √ 1/2 X ϵ ϵ kϵ0 ϵ βj σabs (vj · x))2 ] ≤ + 12 βi min ∥wi − si vi ∥2 ≤ + ≤ . 4 4 4 2 (x,y)∼D si ∈{±1} j=1 P For α ∈ [0, k]m with αi ≤ k, write  2  m X b·x− R(α) = E y−g αj σabs (vj · x) 

E

b·x− [(y − g

k X

(x,y)∼D

j=1

b · x, σabs (v1 · x), . . . , σabs (vm · x)), we have that for any coordinate of ZZ⊤ , Setting Z = (y(x) − g b of D using Nquad = k 10 /(ϵ2 δ) samples. Then, E[(ZZ⊤ )2i,j ] ≲ k 4 . Construct an empirical distribution D by Chebyshev’s inequality and a union bound over all indices i, j, we have that   ϵ k 6 (k + 1)4 ⊤ ⊤ Pr max |(E[ZZ ] − E[ZZ ])i,j | > ≲ ≲ δ. i,j b (k + 1)2 Nquad ϵ2 D Let α = (1, α1 , . . . , αm ). Therefore, if the event inside the probability above is satisfied, we have that b |R(α) − R(α)| = |α⊤ (E[ZZ⊤ ] − E[ZZ⊤ ])α| b D

≤ max |(E[ZZ⊤ ] − E[ZZ⊤ ])i,j |∥α∥21 ≤ i,j

b D

ϵ (m + 1)2 ≤ ϵ. (k + 1)2

Minimizing over α is a convex program and the empirical minimizer using Nquad + Ng = d poly(k/ϵ) samples satisfies b ≲ R(α∗ ) + ϵ ≲ ϵ . R(α) 3 3 e By Proposition 3.10, the procedure CoverAndPrune(ϵ 0 , δ; D) uses Nopt = O(poly(k/ϵ)LCDR (LCDR + √ NCDR )) samples (recall that we have set ϵ0 = ϵ/(4k) and δ = 1/9). Thus, the total number of samples e poly(k/ϵ) + poly(k/ϵ)L3 (L3 + NCDR )). The required for LearnSumReLUs(d, k, ϵ; D) is N = O(d CDR CDR runtime of the algorithm is poly(d, N, TCDR ). This completes the proof.

23

4

Efficient Coarse Direction Recovery

In this section, we show that the CDR procedure postulated in Assumption 1 can be implemented in polynomial time. At a high-level, an algorithm that satisfies Assumption 1 needs to guarantee the following: Given sample access to a distribution (x, y) ∼ D such that (i) the x-marginal is the standard Pk Gaussian, i.e., x ∼ N (0, I), and (ii) y(x) = i=1 βi σ(wi · x), with σ being the ReLU activation, bi ∈ (0, 1] and ∥wi ∥2 = 1, the algorithm returns a list L of unit vectors such that for each target weight vector wi , there exists v ∈ L that correlates non-trivially with wi , i.e., |wi · v| = 1 − c, for a small constant c > 0. To this end, we prove that the top eigenvector of a carefully selected matrix (see Equation (M)) aligns well with one of the directions wi . Our main result in this section is the following polynomial-time implementation of the CDR subroutine. The proof of Proposition 4.1 is deferred to Section 4.3. Proposition 4.1 (Polynomial-Time CDR). Let k, d ∈ Z+ and λ, δ ∈ (0, 1). Let D be a Gaussian sum of k ReLUs. Then, there is an algorithm (CDRviaImplicitTensor(D, λ, δ)) that draws N = poly(d, k, 1/λ, 1/δ) i.i.d. samples from D, runs in T = poly(d, N ) time, and outputs a list L of unit vectors of size |L| = poly(k/λ) log(1/δ) such that with probability at least 1 − δ, for any wi , i ∈ [k], with βi ≥ λ, there exists a vector ri ∈ L such that |ri · wi | ≥ 0.99. We begin this section by formally arguing that the guarantee of the CDR subroutine can be reduced to finding the top eigenvector of a carefully chosen matrix (M) (Section 4.1). Then, in Section 4.2, we give a simple estimation procedure for this matrix that uses a quasi-polynomial number of samples from D. Section 4.2 highlights why more sophisticated techniques are needed for estimation and motivates the techniques based on implicit tensor computation that we employ in Section 4.3. Specifically, in Section 4.3, we present an algorithm (that uses techniques from implicit tensor computation) to estimate the matrix (M) with polynomially many samples and in polynomial time. Using this algorithm, we then prove Proposition 4.1. Throughout √ this section, we take without loss of generality that λ is sufficiently smaller than one (namely, λ ≤ ξ/10, where ξ as absolute constant). This only impacts the constants suppressed by the big-Oh notation.

4.1

Low-Degree Filtering: Revealing Hidden Directions

As discussed in Section 1.2.1, our implementation of CDR is based on the properties of Hermite tensors of ReLU functions; namely, the fact that degree-p moments of the ReLU sums are proportional to rank-tensors (recall (2)). For technical reasons, instead of working with the tensor Q2p introduced in (3) Pk in Section 1.2.1, we use a closely related matrix that is likewise proportional to i=1 βi |wi · u|2p wi wi⊤   (i) 1 1 ⊗p ⊗p M2p+2 (u) := E [yHe2p+2 (x)], u ⊗ ū Zp (U) σ b2p+2 (x,y)∼D (ii)

=

k

1 X βi |wi · u|2p wi wi⊤ , Zp (U) i=1

(M)

where Zp (U) > 0 is a fixed parameter depending on the sampling distribution U from which u is drawn and on the degree p. The probing vector u can be chosen from a complex distribution, where ū is its complex conjugate. When u is a real vector, we have ū = u, and in this case, the tensor contraction in (i) becomes ⟨E[yHe2p+2 (x)], u⊗2p ⟩. Intuitively, the tensor u⊗p ⊗ ū⊗p ∈ (Cd )⊗2p can be viewed as a “filter” that projects the signal tensor E[yHe2p+2 (x)] to the direction of u. Choosing the order p and the distribution U properly, we can guarantee that with large probability, only the target vector wi , i ∈ [k] will come through the filter—as discussed in Section 1.2.1. To see why (ii) is true, recall that by the definition of Hermite tensors and the definition of our learning problem (Definition 1.2), we Pk have E[yHe2p+2 (x)] = i=1 σ b2p+2 βi wi⊗2p+2 , which yields (ii) after contraction with the filter. The matrix M can be estimated via samples from D. Naively, one can estimate the population tensor E[yHe2p+2 (x)] using dO(p) samples, then perform contraction with the filter. However, since the target matrix only concerns the projection to the probing vector u, estimating the tensor E[yHe2p+2 (x)] accurately in all the directions is unnecessary. In Section 4.2, we show that M can be estimated approximately via a polynomial; then in Section 4.3 we give an algorithm that estimates the contraction

24

⟨E(x,y)∼D [yHe2p+2 (x)], u⊗p ⊗ ū⊗p ⟩ as a whole using polynomially many samples, avoiding a separate estimation of the tensor E[yHe2p+2 (x)]. In the following, we show that the CDR subroutine can be reduced to estimating the top eigenvector of the empirical estimator of the matrix M. Lemma 4.2. Let ξ √ ∈ (0, 1/300] be a small absolute constant. Let wi be unit vectors in Rd , βi ∈ [0, 1], i ∈ [k], and λ ∈ (0, ξ/10] be less than a sufficiently small constant. Let p = ⌈log(k/λ)/ log(1 + ξ)⌉. Consider any i∗ ∈ [k] such that βi∗ ≥ λ, and denote Ai∗ (ξ) = {i ∈ [k] : |wi · wi∗ | ≤ 1 − ξ}. Draw and fix u from the distribution U. Let M2p+2 (u) be the matrix defined in (M). Define the following “correlation” and “separation” events on U: EiCor = {u : |wi∗ · u| ≥ (Zp (U))1/(2p) }, EiSep = {u : |wi∗ · u| ≥ (1 + ξ)|wj · u|, ∀j ∈ Ai∗ (ξ)}. ∗ ∗

(11)

c c Suppose u ∈ EiCor ∩ EiSep ∗ ∗ . Let M2p+2 (u) be any real matrix such that ∥M2p+2 (u) − M2p+2 (u)∥2 ≤ √ c 2p+2 (u) satisfies |b ( ξ/3)∥M2p+2 (u)∥2 . Then, the top eigenvector b r1 of M r1 · wi∗ | ≥ 0.99. Proof. The high-level idea of the proof is that under the assumptions of Lemma 4.2, the target matrix M2p+2 (u) is approximately rank-1 and proportional to wi∗ wi⊤∗ . Let us define WA =

1 Zp (U)

X

βi |wi · u|2p , WB =

i∈Ai∗ (ξ)

1 Zp (U)

X

⊤ βi |wi · u|2p , and MB 2p+2 (u) = WB wi∗ wi∗ .

i∈A / i∗ (ξ)

(12) Observe that tr(M2p+2 (u)) = WA + WB . We establish the following: Claim 4.3. Under the assumptions of Lemma 4.2, it holds: WA ≤ (λ/k)WB and ∥M2p+2 (u) − MB 2p+2 (u)∥2 ≤ (λ/k +

p

2ξ)∥MB 2p+2 (u)∥2 .

(13)

Proof of Claim 4.3. Since i∗ ∈ / Ai∗ (ξ), βi∗ ≥ λ, and |wi∗ · u|2p ≥ Zp (U) for u ∈ EiCor , the weight WB ∗ is large, namely WB ≥ βi∗ |wi∗ · u|2p /Zp (U) ≥ λ. On the other hand, for any i ∈ Ai∗ (ξ), the probing vector u satisfies |wi · u| ≤ |wi∗ · u|/(1 + ξ) (because u ∈ EiSep ∗ ). Hence, using the assumption βi ∈ (0, 1] and our choice of p, the weight WA is small. Specifically, we have WA ≤

1 Zp (U)

X i∈Ai∗ (ξ)

|wi∗ · u|2p |wi∗ · u|2p ≤ k exp(−(2p) log(1 + ξ)) 2p (1 + ξ) Zp (U)

λ2 |wi∗ · u|2p λ βi∗ |wi∗ · u|2p λ ≤ ≤ ≤ WB . k Zp (U) k Zp (U) k Therefore, we have ∥M2p+2 (u) − MB 2p+2 (u)∥2 X 1 1 = βi |wi · u|2p wi wi⊤ + Zp (U) Zp (U) i∈Ai∗ (ξ)

≤

1 Zp (U)

X

βi |wi · u|2p ∥wi wi⊤ ∥2 +

i∈Ai∗ (ξ)

X

βi |wi · u|2p (wi wi⊤ − wi∗ wi⊤∗ )

i∈A / i∗ (ξ)

1 Zp (U)

X

2

βi |wi · u|2p ∥wi wi⊤ − wi∗ wi⊤∗ ∥2

i∈A / i∗ (ξ)

p Observe that ∥vv⊤ − ww⊤ ∥2 ≤ 1 − (v · w)2 for any unit vectors v, w. By the definition of the index √ set Ai∗ (ξ), for any i ∈ / Ai∗ (ξ), it holds |wi · wi∗ | ≥ 1 − ξ, therefore we have ∥wi wi⊤ − wi∗ wi⊤∗ ∥2 ≤ 2ξ. Thus, we conclude that: p p ∥M2p+2 (u) − MB 2ξWB ≤ (λ/k + 2ξ)WB , 2p+2 (u)∥2 ≤ WA + completing the proof of Claim 4.3.

25

c 2p+2 (u) is close to MB (u) in spectral norm, which would imply that It remains to prove that M 2p+2 c 2p+2 (u) aligns with wi∗ . As λ ∈ (0, √ξ/10], it holds that λ/k ≤ √ξ/10. By the top-eigenvector of M the triangle inequality and Claim 4.3, we have c 2p+2 (u) − MB (u)∥2 ≤ ∥M c 2p+2 (u) − M2p+2 (u)∥2 + ∥M2p+2 (u) − MB (u)∥2 ∥M 2p+2 2p+2 √ √  ξ p ξ ≤ + 2ξ WB + ∥M2p+2 (u)∥2 . 10 3 √ In addition, since ∥M2p+2 (u)∥2 ≤ WA +WB and WA ≤ (λ/k)WB ≤ ( ξ/10)WB , we have ∥M2p+2 (u)∥2 ≤ √ (1 + ( ξ/10))WB . Thus, we conclude that √  √  √ ξ p ξ ξ c 2p+2 (u) − MB (u)∥2 ≤ 1 + WB . ∥M + 2ξ + 2p+2 10 3 10 c 2p+2 by b Denote the top (unit) eigenvector of M r1 and the top eigenvector of MB 2p+2 (u) by r1 ; note that r1 = wi∗ . Then, using Wedin’s inequality (Fact 2.1), with ξ ≤ 1/300, we have c 2p+2 (u) − MB (u)∥2 ∥M 2p+2 B c ∥M2p+2 (u)∥2 − ∥M2p+2 (u) − MB 2p+2 (u)∥2 √ √ √ √ ξ ξ ξ p ( 10 + 2ξ + 3 (1 + 10 ))WB √ √ √ ≤ 1 − 0.992 . ≤ √ WB − ( 10ξ + 2ξ + 3ξ (1 + 10ξ ))WB

sin(θ(b r1 , r1 )) ≤

Thus, |b r1 · wi∗ | ≥ 0.99, completing the proof of Lemma 4.2. What remains are the following steps: (1) pick a suitable distribution of u such that the event u ∈ EiCor ∩ EiSep happens with probability at least poly(λ/k); (2) approximate the matrix M2p+2 (u) ∗ ∗ efficiently using samples from (x, y) ∼ D. In Section 4.2, we choose U to be the uniform distribution on the sphere, obtaining quasi-polynomial sample complexity; then in Section 4.3, we choose U to be a complex distribution, and by applying implicit tensor computation techniques, we argue that M2p+2 can be estimated efficiently, using poly(k/λ) samples.

4.2

Warm-up: Quasi-Polynomial Time CDR Algorithm

Pk Given a distribution (x, y) ∼ D such that y(x) = i=1 βi σ(wi · x) and x sin N (0, Id ), a simple choice of the probing vector is to sample from the uniform distribution on the unit sphere in a subspace containing Span(w1 , . . . , wk ). The latter can be estimated efficiently via the following standard result. Fact 4.4 (see, e.g., Lemma 3.4 of [DKKZ20]). Let (x, y) ∼ D be the distribution on Rd × R such that Pk x ∼ N (0, Id ) and y(x) = i=1 βi σ(wi · x), βi ∈ [0, 1], ∥wi ∥2 = 1, and σ is the ReLU activation. Let σ b2 be the second- order Hermite coefficient of σ. Let η > 0. Then, using N = poly(d, k, 1/η) samples (x(i) , y (i) ) drawn from D, it holds with high probability that: √

N k X 1 X (i) (i) (i)⊤ y (x x − Id ) − |b σ2 |βj wj wj⊤ ≤ η . 2N i=1 2 j=1

Let c W be the subspace spanned by the top-k vectors of the empirical p matrix constructed in Fact 4.4. Then, Fact 4.4 implies that for any wi , i ∈ [k], it holds ∥wi − PW w ∥ ≤ 2η/βi (with high probability, b i 2 otherwise, ascribe to the failure event of the algorithm). For any λ-heavy weights wi∗ with coefficients βi∗ ≥ λ, choosing η = poly(λ/k) yields that dist(c W, Span({wi∗ : βi∗ ≥ λ})) ≤ λ. Therefore, for a probing vector u sampled uniformly from c W ∩ Sd−1 , the probability that it satisfies Sep Cor the correlation event Ei∗ and the separation Ei∗ is almost the same as if sampled from W ∩ Sd−1 , with only differences in the absolute constants. Thus, to simplify the analysis, in the remainder of this section, we assume access to the subspace W = Span(w1 , . . . , wk ) with dimension m and we sample the probing vectors u from U = Unif(W ∩ Sd−1 ). 26

Remark 4.5 (Range of parameters.). Throughout this section, we assume that the dimension m of the subspace W = Span(w1 , . . . , wk ) satisfies m ≳ log(k). Indeed, the CDR procedure for small m admits a simpler approach. More generally, consider the regime m ≲ log(k). In this case, we can obtain the desired CDR guarantee by sampling random directions from the m-dimensional subspace W. For any fixed unit vector in W, a uniformly random unit vector has correlation at least 0.99 with probability 2−O(m) . By taking q = 2O(m) log(k/δ) ≤ poly(k, 1/ϵ) log(1/δ) random vectors u1 , . . . , uq , with probability at least 1 − δ, we have that maxj∈[q] |uj · wi | ≥ 0.99 for all i ∈ [k]. The pseudocode of our algorithm (CDRviaSpectral(D, λ, δ)) is given below. Algorithm 6 CDRviaSpectral(D, λ, δ) 1: Input: A subspace W, access to a distribution of a sum of k ReLUs with Gaussian marginals

(x, y) ∼ D, positive parameters λ, δ ∈ (0, 1). 2: Output: A list of vectors satisfying Assumption 1.

√

3: Fix an absolute constant ξ ∈ (0, 1/300], let p ← ⌈log(k/λ)/ log(1 + ξ)⌉, λ ← p min{λ, ξ/10}. ′′ 4: Let σ b2p+2 be the (2p + 2)th Hermite coefficient of ReLU and σ b2p ←σ b2p+2 (2p + 1)(2p + 2). 5: Zp (U) ← k −p , L ← ∅.

b of poly(k) log(1/δ) probing vectors from U = Unif(Sd−1 ∩ W). 6: Draw a set U b do 7: for u ∈ U b 8: Draw N = d2 poly(p)k O(log(k/λ)) samples from D and construct empirical distribution D. 9: 10:

′′ c 2p+2 (u) ← (1/(b M σ2p Zp (U))) E(x,y)∼Db [yhe2p (u · x)(x⊥u x⊥u⊤ − I⊥u )]. c 2p+2 (u)), L ← L ∪ {±r1 }. Find the top eigenvector of r1 = r1 (M

11: Return: L.

Our main result for this subsection is the following: Proposition 4.6 (Quasi-Polynomial-Time √ Implementation of CDR). Let δ ∈ (0, 1), ξ ∈ (0, 1/300] be an absolute constant and let λ ∈ (0, ξ/10]. Let (x, y) ∼ D be a distribution such that y(x) = Pk i=1 βi σ(wi · x) where σ is the ReLU activation, βi ∈ [0, 1], ∥wi ∥2 = 1, and x ∼ N (0, Id ). Let p ≥ ⌈log(k/λ)/ log(1 + ξ)⌉. Then, CDRviaSpectral(D, λ, δ) draws N = d2 k O(log(k/λ)) log(1/δ) samples from D, runs in T = poly(d, N ) time, and outputs a list L of unit vectors of size |L| = poly(k) log(1/δ) that with probability at least 1 − δ satisfies the following guarantee: for any wi , i ∈ [k] with βi ≥ λ, there exists ri ∈ L such that |ri · wi | ≥ 0.99. To prove Proposition 4.6, we need to bound from below the probability of the event EiCor ∩ EiSep for ∗ ∗ ∗ a fixed index i ∈ [k], defined in (11). Intuitively, when the probing vector u is sampled uniformly from the sphere, with high probability the values |u · wi∗ | and |u · wj | would be separated when the vectors wi∗ and wj are away from each other. This is the separation event EiSep ∗ . At the same time, Cor ∗ | cannot be too small and would be for the correlation event, E , with high probability, |u · w ∗ i i √ concentrated near 1/ k. Therefore, the probability that EiCor ∩ EiSep holds is bounded from below. We ∗ ∗ then determine the number of samples required to estimate the target matrix M. Finally, we prove Proposition 4.6 by invoking Lemma 4.2, as outlined previously. The proof of Proposition 4.6 is deferred to the end of this subsection. Determining the Probability of the Probing Event. We first determine the probability of event EiCor ∩ EiSep under the distribution U = Unif(W ∩ Sd−1 ). In fact, we prove a slightly stronger result: ∗ ∗ with probability at least poly(1/k) in addition to u ∈ EiCor ∩ EiSep ∗ ∗ , u also satisfies a “near-orthogonal” Ort Ort event E . The “near-orthogonal” event E states that u satisfies |u · wi | ≤ c for a sufficiently small absolute constant c for all hidden weight vectors wi . This event is needed in estimating the target matrix M2p+2 to high accuracy, as will become apparent in the proof of Proposition 4.8. Lemma 4.7. Let parameters p, Zp (U) and constant ξ be as prescribed in Proposition 4.6. Let w1 , . . . , wk be k unit vectors and W = Span{w1 , . . . , wk } ⊂ Rd be a linear subspace of dimension m ≤ k. Let u ∼ U where U = Unif(W ∩ Sd−1 ). Fix (any) index i∗ ∈ [k] and let Ai∗ (ξ) = {i ∈ [k] : |wi · wi∗ | ≤ 1 − ξ}. Define the “correlation” and “separation” events on U as in Lemma 4.2: EiCor = {u : |wi∗ · u| ≥ (Zp (U))1/(2p) }, EiSep = {u : |wi∗ · u| ≥ (1 + ξ)|wj · u|, ∀j ∈ Ai∗ (ξ)}. ∗ ∗ 27

In addition, let the “near-orthogonal” event on U be: p

E Ort = {u : |u · wi | ≤

ξ/20, ∀i ∈ [k]}.

∩ EiSep ∩ E Ort ] ≥ poly(1/k). Then, Pru∼U [EiCor ∗ ∗ Proof. Let p ∼ N (0, IW ) and Zj := p · wj . Then, by definition, u = p/∥p∥2 and u · wi = p · wi /∥p∥2 . Sep ∗ The roadmap of the proof is the following: we with pfirst show that Ei∗ ∩ {t ≤ |p · wi | ≤ 2t} phappens √ probability at least poly(1/k) for some t ≲ log(k) then we argue that Pr[∥p∥2 ∈ ( m/2, m)] ≥ 1 − poly(1/m). Since we have assumed without loss of generality (see Remark 4.5) that the dimension m of the subspace W satisfies m ≳ log(k), bound on these events gives a lower bound p applying a union √ Sep Cor ∩ E Ort ], completing the proof. on Pr[Ei∗ ∩ {t/ m ≤ |u · wi∗ | ≤ 2t/ m/2}] = Pr[Ei∗ ∩ EiSep ∗ ∗ ∗ ∗ ∗ For each j ∈ Ai (ξ), let ρi ,j = wi · wj = cos(θ(wi , wj )). By the definition of Ai∗ (ξ), it holds ∗ |ρi ,j | ≤ ρ := 1 − ξ. Then, we can write Zj as \ ⊥w Zj = p · cos(θ(wi∗ , wj ))wi∗ + p · sin(θ(wi∗ , wj ))wi∗ j = ρi∗ ,j Zi∗ +

q

1 − ρ2i∗ ,j Gj ,

where Gj is a standard normal random variable. By the definition of u, it holds u · wi = p · wi /∥p∥2 , and hence the event |u · wj | ≤ |u · wi∗ |/(1 + ξ) is equivalent to |Zj | ≤ |Zi∗ |/(1 + ξ). Conditional on Zi∗ = z, consider the event that EiSep is violated, i.e., there exists some Zj , j ∈ Ai∗ (ξ), such that ∗ |Zj | ≥ |Zi∗ |/(1 + ξ). Note that |Zj | ≤ |ρi∗ ,j ||z| + |Gj | and using that |ρi∗ ,j | ≤ 1 − ξ and rearranging, we get that on a violation, |Gj | ≥ |z|ξ 2 /(1 + ξ). Therefore, from a union bound, we have X Pr[∪j∈Ai∗ (ξ) {|Zj | ≥ |Zi∗ |/(1 + ξ)} | Zi∗ = z] ≤ Pr[|Gj | ≥ |z|ξ 2 /(1 + ξ) | Zi∗ = z] j∈Ai∗ (ξ)

≤ k max Pr[|Gj | ≥ |z|ξ 2 /(1 + ξ) | Zi∗ = z] . j∈Ai∗ (ξ)

Using that Pr[|Gj | ≥ q] ≤ 2 exp(−q 2 /2), we get that Pr[∪j∈Ai∗ (ξ) {|Zj | ≥ |Zi∗ |/(1 + ξ)} | Zi∗ = z] ≤ 2k exp(−z 2 ξ 4 /(2(1 + ξ)2 )) . p Let Cξ = 4(1 + ξ)2 /ξ 4 . We get that for any z ≥ t := Cξ log(2k), it holds Pr[∪j∈Ai∗ (ξ) {|Zj | ≥ |Zi∗ |/(1 + ξ)} | |Zi∗ | = z] ≤ 1/2 . Therefore, the probability of no violation conditioned on |Zi∗ | = Z ≥ t is at least Pr[∩j∈Ai∗ (ξ) {|Zj | ≤ |Zi∗ |/(1 + ξ)} | |Zi∗ | = z] ≥ 1/2. Furthermore, by the anti-concentration inequality for standard Gaussians, we have: Pr[2t ≥ |Zi∗ | ≥ t] ≥ √ Therefore, we get that (t =

p

1 −t2 /2 e ≥ e−Cξ log(2k) = (2k)−Cξ . 2πt

Cξ log(2k)):

Pr[EiSep ∩ {2t ≥ |Zi∗ | ≥ t}] = Pr[∩j∈Ai∗ (ξ) {|Zj | ≤ |Zi∗ |/(1 + ξ)} ∩ {2t ≥ |Zi∗ | ≥ t}] ∗ ≥ Pr[∩j∈Ai∗ (ξ) {|Zj | ≤ |Zi∗ |/(1 + ξ)} | t ≤ |Zi∗ | ≤ 2t] Pr[2t ≥ |Zi∗ | ≥ t] ≥ (1/2)(2k)−Cξ .

(14)

Next observe that the random variable ∥p∥22 is a sum of m squares of independent Gaussians (as W is an m-dimensional subspace). Therefore, ∥p∥22 follows the χ2m distribution. From the concentration of χ2m , we have that p √ Pr[∥p∥2 ̸∈ ( m/2, 2m)] ≤ 2 exp(−m/16) . It is without loss of generality that m is greater than a large absolute constant (Remark 4.5), and we have m ≥ 17Cξ log(2k). Thus, we obtain p √ Pr[∥p∥2 ̸∈ ( m/2, 2m)] ≤ (1/8)(2k)−Cξ . (15) 28

Finally, since Zj , j ∈ [k], are standard Gaussian random variables, by the Gaussian concentration inequality and a union bound, we obtain:   Pr max |Zj | ≥ 4t ≤ 2k exp(−8t2 ) = 2k exp(−8Cξ log(2k)) = (2k)−8Cξ +1 ≤ (1/8)(2k)−Cξ . (16) j∈[k]

Combining (14), (15) and (16), since Pr[A ∩ B ∩ C] ≥ 1 − Pr[Ac ] − Pr[B c ] − Pr[C c ], with p √ A = EiSep ∩ {2t ≥ |Zi∗ | ≥ t}, B = ∥p∥2 ∈ ( m/2, 2m), C = max |Zj | ≥ 4t , ∗ j∈[k]

we get that

p √ Pr[EiSep ∩ {2t ≥ |Zi∗ | ≥ t} ∩ {∥p∥2 ∈ ( m/2, 2m)} ∩ {max |Zj | ≤ 4t}] ≥ (1/8)(2k)−Cξ . ∗ j∈[k]

Therefore, plugging in the value t =

p

Cξ log(2k), on this event, we have r √ Cξ log(2k) ≥ 1/ k , |u · wi∗ | ≥ 2m

and since m ≥ 2 × 802 (Cξ /ξ) log(2k), we also have p √ 4 2Cξ log(2k) 4t ξ √ |u · wj | ≤ p ≤ , ∀j ∈ [k]. ≤ 20 m m/2 √ Since (Zp (U))1/(2p) = 1/ k, we conclude that Pru∼U [EiCor ∩ EiSep ∩ E Ort ] ≥ poly(1/k). ∗ ∗ Estimating the Target Matrix in Quasi-Polynomial Time. We now show that, with quasipolynomial sample and computational complexity, we can approximate the target matrix M2p+2 (u) (M) in operator norm to high accuracy. Proposition 4.8 (Sample √ complexity for the spectral method). Let 0 < ξ ≤ 1/300 be an absolute constant and let 0 < λ ≤ ξ/10. Let (x, y) ∼ D be a sum p of k ReLUs under Gaussian marginals. Let ′′ p = ⌈log(k/λ)/ log(1 + ξ)⌉, Zp (U) = k −p and σ b2p =σ b2p+2 (2p + 1)(2p + 2). Fix i∗ ∈ [k] and suppose the unit probing vector u satisfies the “correlation”, “separation” and “near orthogonality” events defined in Lemma 4.7: u ∈ EiCor ∩ EiSep ∩ E Ort . Draw N = d2 poly(k, p, 1/λ)k O(log(k/λ)) i.i.d. samples (x, y) ∼ D ∗ ∗ b Then, for and construct the empirical distribution D.   1 c 2p+2 (u)⊥u = M E yhe2p (u · x)(x⊥u x⊥u⊤ − I⊥u ) , ′′ σ b2p Zp (U) (x,y)∼Db c 2p+2 (u)⊥u − M2p+2 (u)∥2 ≤ (√ξ/3)∥M2p+2 (u)∥2 . with probability at least 2/3, it holds ∥M Proof. The high-level idea of the proof is that the target matrix M2p+2 (u) defined in (M) can be approximately written as the following expectation. Claim 4.9. Given any unit vector u, define M2p+2 (u)⊥u = P⊥u M2p+2 (u)P⊥u . Then,   1 u⊥ ⊥u M2p+2 (u) = ′′ E yhe2p (u · x)He2 (P⊥u x) . σ b2p Zp (U) (x,y)∼D Pk Proof of Claim 4.9. By our definition, the distribution (x, y) ∼ D satisfies y(x) = i=1 βi σ(wi · x) and x ∼ N (0, Id ). Therefore, denoting θiu := θ(wi , u), we have that for any unit vector u ∈ Rd , h i ⊥ E yhe2p (u · x)Heu 2 (P⊥u x) (x,y)∼D

=

k X

 βi E

u·x

i=1 k (i) X

=

βi

i=1 k (ii) X

=

i=1

E

x⊥u ∼N (0,Iu⊥ )

E

h

x∼N (0,Id )

βi

E x∼N (0,Id )

h

 i u⊥ ⊥u ⊥u [ ⊥u σ(cos(θiu )wi · x )He2 (x ) u · x he2p (u · x)

i ⊤ [ [ ⊥u · x)he (u · x) sin2 (θ )w ⊥u [ ⊥u σ ′′ (cos(θiu )u · x + sin(θiu )w 2p iu i i wi

h i Tcos(θiu ) (σ ′′ )(u · x)he2p (u · x) wi⊥u wi⊥u⊤ ,

29

(17)

where in (i) we used Fact 2.6, by which Ez∼N [f (r · z)He2 (z)] = Ez∼N [f ′′ (r · z)]rr⊤ for any unit vector r, and in (ii) Tρ (σ ′′ ), ρ ∈ [0, 1], is the smoothing operator defined in Fact 2.3. By p Fact 2.3, Part 2, we P know that σ ′′ has Hermite expansion σ ′′ (z) = i≥0 c′′i hei (z), where σ bi′′ = σ bi+2 (i + 1)(i + 2). Thus, applying Fact 2.3, Part 3 to (17) leads to: k h i 1 1 X u⊥ yhe (u · x)He E (P x) = βi |wi · u|2p wi⊥u wi⊥u⊤ 2p ⊥u 2 ′′ Z (U) σ b2p Zp (U) i=1 (x,y)∼D p

= P⊥u M2p+2 (u)P⊥u =: M2p+2 (u)⊥u , completing the proof of Claim 4.9. We next show that when u is a unit vector such that the event u ∈ EiCor ∩ EiSep ∩ E Ort holds for ∗ ∗ ⊥u λ-heavy hidden weight wi∗ , then the spectral difference ∥M2p+2 (u) − M2p+2 (u)∥2 is small. Claim 4.10. Under the assumptions of Proposition 4.8, it holds √ ξ ⊥u ∥M2p+2 (u) − M2p+2 (u)∥2 ≤ ∥M2p+2 (u)∥2 and ∥M2p+2 (u)∥2 ≥ λ . 8 Proof of Claim 4.10. Note that for any unit vector wi , by the triangle inequality, it holds ∥wi wi⊤ − wi⊥u wi⊥u⊤ ∥2 ≤ ∥(wi − wi⊥u )wi⊤ ∥2 + ∥wi⊥u (wi − wi⊥u )∥2 ≤ 2∥wi − wi⊥u ∥2 = 2|u · wi |. Since u also satisfies the near-orthogonal event E Ort , i.e, for all i ∈ [k] it holds |u · wi | ≤ √ ⊤ ⊥u ⊥u⊤ have ∥wi wi − wi wi ∥2 ≤ ξ/10. This implies the following inequality ⊥u

∥M2p+2 (u)

√

ξ/20, we

√ X k k 1 X ξ 2p ⊤ ⊥u ⊥u⊤ − M2p+2 (u)∥2 ≤ βi |wi · u| ∥wi wi − wi wi ∥2 ≤ βi |wi · u|2p . Zp (U) i=1 10 i=1

Recall that in (12) we defined: Ai∗ (ξ) = {i ∈ [k] : |wi · wi∗ | ≤ 1 − ξ}, and WA =

1 Zp (U)

X

βi |wi · u|2p , WB =

i∈Ai∗ (ξ)

1 Zp (U)

X

βi |wi · u|2p ,

i∈A / i∗ (ξ)

Pk so that WA + WB = (1/Zp (U)) i=1 βi |wi · u|2p . Then in Claim 4.3, equation (13), we proved that when u ∈ EiCor ∩ EiSep ∗ ∗ , it holds WA ≤ (λ/k)WB . Thus, we conclude that √ √ ξ ξ(λ/k + 1)WB ⊥u ∥M2p+2 (u) − M2p+2 (u)∥2 ≤ (WA + WB ) ≤ . 10 10 On the other hand, since βi ≥ 0 for all i ∈ [k], we have: ∥M2p+2 (u)∥2 ≥ wi⊤∗ M2p+2 (u)wi∗ ≥

1 Zp (U)

X

βi |wi · u|2p (wi · wi∗ )2 ≥ (1 − ξ)2 WB .

i∈A / i∗ (ξ)

√ Therefore, with the choice ξ ∈ (0, 1/300] and λ ≤ ξ/10, we conclude that: √ √ ξ(λ/k + 1) ξ ⊥u ∥M2p+2 (u) − M2p+2 (u)∥2 ≤ ∥M2p+2 (u)∥2 ≤ ∥M2p+2 (u)∥2 . 10(1 − ξ 2 ) 8 Finally, note that ∥M2p+2 (u)∥2 ≥ βi∗ |wi∗ · u|2p /Zp (U), since wi∗ is a λ-heavy hidden weight and u satisfies EiCor , it holds that ∥M2p+2 (u)∥2 ≥ λ. This completes the proof of Claim 4.10. ∗ Claim 4.10 implies that an accurate estimation of the matrix M2p+2 (u)⊥u approximates the target matrix from (M) in operator norm. The rest of the proof is devoted to determining how many samples are required to estimate M2p+2 (u)⊥u . We show that:

30

2 2p 2

Lemma 4.11. Using N ≳ (bσ′′ d)2 3λ2 Zkp (U )2 samples, with probability at least 2/3, it holds 2p

√ c 2p+2 (u)⊥u − M2p+2 (u)⊥u ≤ M 2

ξ ∥M2p+2 (u)∥2 . 5

′′ Proof of Lemma 4.11. Set Q = σ b2p Zp (U), and let Xi = Q−1 y (i) he2p (u · x(i) )(x(i)⊥u x(i)⊥u⊤ − I⊥u ) and note that E[Xi ] = M2p+2 (u)⊥u by Claim 4.9. By Chebyshev inequality, we have that h 1 X i E[∥ 1 P X − M ⊥u 2 ∥2 ] i 2p+2 (u) ⊥u ⊥u N ≥ t∥M2p+2 (u) ∥2 ≤ . Pr Xi − M2p+2 (u) 2 N t2 ∥M2p+2 (u)⊥u ∥2 2 P By the independence of samples, E[∥ N1 Xi − M2p+2 (u)⊥u ∥22 ] ≤ N2 E[∥Xi ∥2F ]. In addition, by the Pk definition of the labels, we know that y = i=1 βi σ(wi · x). Furthermore, by the independence between u · x and x⊥u and by Cauchy–Schwarz inequality, the second-order moment of Xi is bounded by   E[∥Xi ∥2F ] ≤ Q−2 E y 2 he22p (u · x)∥(x⊥u x⊥u⊤ − I⊥u )∥2F k h X 2 i 1 2 ⊥u ⊥u⊤ 2 E β σ(w · x) he (u · x)∥(x x − I )∥ i i ⊥u 2p F Q2 i=1 v u h k 4 i X 1 u βi σ(wi · x) E[he42p (u · x)∥(x⊥u x⊥u⊤ − I⊥u )∥4F ] ≤ 2 tE Q i=1 v u k 4 i 1 u h X = 2 tE βi σ(wi · x) E[he42p (u · x)] E[∥(x⊥u x⊥u⊤ − I⊥u )∥4F ] . Q i=1

=

Pk The fourth moment of y is bounded above by E[( i=1 βi σ(wi · x))4 ] ≲ k 4 . In addition, using Fact 2.2, we have that E[he42p (u · x)] ≤ 34p . Further, E[∥(x⊥u x⊥u⊤ − I⊥u )∥4F ] ≤ d4 . Therefore, we have: E[∥Xi ∥2F ] ≲

d2 32p k 2 . Q2

Plugging the upper bound on the second moment E[∥Xi ∥22 ] back into Chebyshev inequality yields √ h 1 X i d2 32p k 4 Pr Xi − M2p+2 (u)⊥u ≥ t∥M2p+2 (u)⊥u ∥2 ≲ 2 2 . N Q N t ∥M2p+2 (u)⊥u ∥22 2 When ξ ∈ (0, 1/300], by Claim 4.10 and using the triangle inequality, it holds that p ∥M2p+2 (u)⊥u ∥2 ≥ (1 − ξ/8)∥M2p+2 (u)∥2 ≥ (1/3)λ. √ Thus, choosing t ≲ ξ, if the sample size satisfies N≳

d2 32p k 2 d2 32p k 2 = ′′ 2 2 , 2 2 Q λ (b σ2p ) λ Zp (U)2

we have that with probability at least 2/3, √ ξλ ξ ≤ ≤ ∥M2p+2 (u)∥2 , 10 10 2 √

c 2p+2 (u)⊥u − M2p+2 (u)⊥u M

where in the last inequality we used the fact that conditioned on u ∈ EiCor ∩ EiSep ∩ E Ort , it holds ∗ ∗ ∥M2p+2 (u)∥2 ≥ λ by Claim 4.10. This completes the proof of Lemma 4.11. Now combining Claim 4.10, Lemma 4.11 and using the triangle inequality yields: c 2p+2 (u)⊥u − M2p+2 (u) M

2

c 2p+2 (u)⊥u − M2p+2 (u)⊥u + M2p+2 (u)⊥u − M2p+2 (u) ≤ M 2 2 √ ξ ≤ M2p+2 (u) . 3 2 31

Since p = ⌈log(k/λ)/ log(1 + ξ)⌉ and Zp (U) = k −p = k −O(log(k/λ)) , σ b2p+2 ≥ 1/(2p + 2)3 (by Fact 2.4), by Lemma 4.11 it follows that the number of samples required is N = d2 poly(p)k O(log(k/λ)) . The probability of success is at least 2/3, conditioned on the event u ∈ EiCor ∩ EiSep ∩ E Ort . ∗ ∗ Proof of the Main Result of the Subsection: Quasi-Polynomial Complexity CDR. We are now ready to prove Proposition 4.6, which we restate below: Proposition 4.6 (Quasi-Polynomial-Time √ Implementation of CDR). Let δ ∈ (0, 1), ξ ∈ (0, 1/300] be an absolute constant and let λ ∈ (0, ξ/10]. Let (x, y) ∼ D be a distribution such that y(x) = Pk i=1 βi σ(wi · x) where σ is the ReLU activation, βi ∈ [0, 1], ∥wi ∥2 = 1, and x ∼ N (0, Id ). Let p ≥ ⌈log(k/λ)/ log(1 + ξ)⌉. Then, CDRviaSpectral(D, λ, δ) draws N = d2 k O(log(k/λ)) log(1/δ) samples from D, runs in T = poly(d, N ) time, and outputs a list L of unit vectors of size |L| = poly(k) log(1/δ) that with probability at least 1 − δ satisfies the following guarantee: for any wi , i ∈ [k] with βi ≥ λ, there exists ri ∈ L such that |ri · wi | ≥ 0.99. Proof of Proposition 4.6. Fix any i∗ ∈ [k] such that βi∗ ≥ λ. By Lemma 4.7, we have that Pr[EiCor ∩ ∗ Sep Sep Cor Ort Ei∗ ] ≥ Pr[Ei∗ ∩ Ei∗ ∩ E ] ≥ poly(1/k). Combining Lemma 4.2 and Proposition 4.8, for a probing c 2p+2 (u) vector u satisfying EiCor ∩ EiSep ∩ E Ort , using N = d2 k O(log(k/λ)) samples, the empirical matrix M ∗ ∗ c 2p+2 (u)) · wi∗ | ≥ 0.99 with probability at from Algorithm 6 of CDRviaSpectral(D, λ, δ) satisfies |r1 (M least 2/3. Thus, for any fixed i∗ ∈ [k] with βi∗ ≥ λ, the probability that the top eigenvector of the c 2p+2 (u)) satisfies |ri∗ · wi∗ | ≥ 0.99 is at least poly(1/k). Hence, by a union empirical matrix ri∗ = r1 (M bound and a Chernoff bound, drawing poly(k) log(1/δ) probing vectors and constructing poly(k) log(1/δ) empirical matrices ensures that with probability at least 1 − δ, for any i∗ ∈ [k] with βi∗ ≥ λ, the list L returned by CDRviaSpectral(D, λ, δ) contains a unit vector ri∗ such that |ri∗ · wi∗ | ≥ 0.99. This completes the proof.

4.3

Polynomial Time CDR via Implicit Tensor Computation

In this section, we provide a CDR subroutine that uses polynomially many samples and runs in polynomial time. Similar to Section 4.2, the core idea is to estimate the target matrix M2p+2 (u) (M), constructed by projecting the moment tensor E(x,y)∼D [yHe2p+2 (x)] to a low-degree filter u⊗p ⊗ ū⊗p :   1 1 M2p+2 (u) = E [yHe2p+2 (x)], u⊗p ⊗ ū⊗p Zp (U) σ b2p+2 (x,y)∼D k

=

1 X bi |wi · u|2p wi wi⊤ , Zp (U) i=1

where u is a random probing vector sampled from some distribution U. Once we obtain an accurate c 2p+2 (u), we can invoke Lemma 4.2, which states that the top eigenvector of M2p+2 (u) estimation M aligns well with the target vector wi , conditions on u that are satisfied with inverse-polynomial probability. To ensure polynomial complexity, we take a significant departure from Section 4.2 in both components of the algorithm: (1) instead of using the uniform distribution on the sphere, we choose U to be a complex Gaussian distribution: u = p + iq, p, q ∼ N (0, Id ); and (2) instead of estimating the matrix M2p+2 (u) via low-degree polynomials, we use more advanced techniques from implicit tensor computation [DK25]. Choosing U to be the complex Gaussian has two main advantages: first, the operator norm of the target c 2p+2 (u) becomes larger, since |wi · u| is no longer restricted to be smaller than one, as was matrix M the case when u were sampled uniformly on the sphere; second, the covariance of u⊗p for u = p + iq is bounded above by 2p p! (cf. Lemma 4.23)—when p ≲ k, this is much smaller compared to O(k p/2 ), the covariance of u⊗p for u ∼ Unif(Sk−1 ). These benefits together ensure that we can apply the implicit tensor computation techniques to estimate the target matrix accurately using only a polynomial number of samples, in polynomial time. Finally, while it is possible to choose other types of distributions U with similar properties; the complex Gaussian distribution is the most convenient one for our analysis. The pseudo-code of the polynomial-time CDR subroutine is given in CDRviaImplicitTensor(D, λ, δ). 32

Algorithm 7 CDRviaImplicitTensor(D, λ, δ) 1: Input: Sample access to a sum of k ReLUs with Gaussian marginals (x, y) ∼ D, positive parameters

λ, δ ∈ (0, 1). A list of unit vectors satisfying Assumption 1. √ 3: Fix an absolute constant ξ ∈ (0, 1/300], λ ← min{λ, ξ/10}. 4: p ← ⌈log(k/λ)/ log(1 + ξ)⌉. ▷ Filter degree 5: Nu ← poly(k/λ) log(1/δ). ▷ Number of probing vectors/Low-degree filter 6: τ ← δ/(10d2 Nu ). ▷ Failure probability of the immediate steps 7: N0 ← poly(d, k/λ, 1/τ , 1/δ). ▷ Number of samples 8: Let the distribution I ∼ D(I) be I = (x, y, z), where (x, y) ∼ D and z ∼ N (0, Id ). b p ←TensorCompressor(p, N0 , D(I)) 9: Φ b p , D(I)) 10: Ap ← CompressEmpMoment(p, N0 , Φ 11: L ← ∅. 12: for j = 1, . . . , Nu do 13: Draw a fresh vector u = p + iq, where p, q ∼ N (0, Id ) b p , u) 14: Bp (u) ←CompressFilter(p, Φ c 15: M2p+2 (u) = ⟨Ap , Bp (u)⟩ ▷ Empirical matrix, cf. EstM c 2p+2 (u))}. 16: L ← L ∪ {±r1 (M ▷ Top eigenvectors of the empirical matrix 17: Return: L. 2: Output:

The whole section is devoted to proving Proposition 4.1, which is restated below. We defer its proof to Section 4.4, after completing all the ingredients. Proposition 4.1 (Polynomial-Time CDR). Let k, d ∈ Z+ and λ, δ ∈ (0, 1). Let D be a Gaussian sum of k ReLUs. Then, there is an algorithm (CDRviaImplicitTensor(D, λ, δ)) that draws N = poly(d, k, 1/λ, 1/δ) i.i.d. samples from D, runs in T = poly(d, N ) time, and outputs a list L of unit vectors of size |L| = poly(k/λ) log(1/δ) such that with probability at least 1 − δ, for any wi , i ∈ [k], with βi ≥ λ, there exists a vector ri ∈ L such that |ri · wi | ≥ 0.99. Remark 4.12. Throughout this section, we assume without loss of generality that k ≤ d. If k > d, before running our CDR procedure CDRviaImplicitTensor(D, λ, δ), set d′ = max{d, k}. If d′ > d, augment every example x by an independent vector z ∼ N (0, Id′ −d ), and use x′ = (x, z) with the same label. The lifted hidden weights are wi′ = (wi , 0), so the target and the Gaussian model are preserved. Run ′ the CDR procedure (CDRviaImplicitTensor(D, λ, δ)) in dimension d′ . If a unit vector r ∈ Rd is returned, retain P r/∥P r∥2 when P r ̸= 0, where P projects onto the original d coordinates; discard a zero projection. For every original unit weight wi , it holds wi ·

Pr |wi′ · r| = ≥ |wi′ · r|. ∥P r∥2 ∥P r∥2

Thus, a useful output retains its correlation guarantee after projection, the output list cannot grow, and the cost remains polynomial in d and k. In the tensor analysis below, relabel d′ as d, so d ≥ k. 4.3.1

Determining the Probability of the Probing Event

Our first step is to determine the probability that the signal of the target vectors passes through the Cor random filter—in other words, the probability that the “separation” and “correlation” events EiSep ∗ , Ei∗ (defined in (11)) happen under the complex distribution U. We prove that, with p = ⌈log(k/λ)/ log(1+ξ)⌉ and Zp (U) = p!, the event EiSep ∩ EiCor happens with probability at least poly(λ/k). ∗ ∗ iid

Lemma 4.13 (Probability of Successful Filtration). Let u = p + iq, where p, q ∼ N (0, Id ). Let ξ ∈ (0, 1/300] be a small absolute constant, p = ⌈log(k/λ)/ log(1 + ξ)⌉ and Zp (U) = p!. Let w1 , . . . , wk be k unit vectors. Fixing any i∗ ∈ [k], let Ai∗ (ξ) = {i ∈ [k] : |wi · wi∗ | ≤ 1 − ξ} and let the “correlation” and “separation” events on U be defined by (11). Then, Pru∼U [EiCor ∩ EiSep ∗ ∗ ] ≥ poly(λ/k). 33

p √ Proof. Without loss of generality, let i∗ = 1 and fix w1 . Let Cξ = 12((1+ξ)/ξ 2 ) and R = Cξ log(k/λ). By our choice of parameters, we have R2 ≥ p. Furthermore, since R2p ≥ pp ≥ p! = Zp (U), we have E1Cor = {u : |w1 · u| ≥ (Zp (U))1/(2p) } ⊃ {u : |w1 · u| ≥ R} := E1Cor,2 . We proceed by bounding the probability of the event E1Cor . For any fixed unit vector r ∈ Rd , we have |r · u|2 = (p · r)2 + (q · r)2 = X12 + X22 , where X1 , X2 are independent standard Gaussian random variables; therefore, by the anti-concentration inequality for the standard Gaussian distribution, it holds that √ Pr[|r · u| ≥ R] = Pr[|r · u|2 ≥ R2 ] ≥ Pr[X12 ≥ R2 /2, X22 ≥ R2 /2] ≥ Pr[|X1 | ≥ R/ 2]2 ≥ exp(−cR2 ) = poly(λ/k). Therefore, Pr[E1Cor ] ≥ Pr[E1Cor,2 ] ≥ poly(λ/k). It remains to bound Pr[E1Sep |E1Cor,2 ]. Denote zi := u · wi . Consider the joint distribution of z1 = w1 · u and zi = wi · u for any i ∈ A1 (ξ). Letting cos θi := cos(θ± (wi , w1 )) = |wi · w1 |, we have zi = u · (cos θi w1 + sin θi si ), where si = wi⊥w1 /∥wi⊥w1 ∥2 . Denote for simplicity gi = u · si . We condition on E1Cor,2 , i.e., |z1 | = Z ≥ R. For any i ∈ A1 (ξ), we have cos(θi ) = |wi · w1 | ≤ 1 − ξ. By the triangle inequality, whenever |wi · u| ≥ |w1 · u|/(1 + ξ) and |w1 · u| = Z, it holds √   Z cos(θi ) Z Z 1 Zξ 2 12Z |gi | ≥ − + ≥ − (1 − ξ) ≥ = , sin(θi ) (1 + ξ) sin(θi ) sin(θi ) 1 + ξ 1+ξ Cξ where we used that cos(θi ) ≤ 1 − ξ. Therefore, we have: √ √       12Z 12R |z1 | |z1 | = Z, Z ≥ R ≤ Pr |gi | ≥ Pr |zi | ≥ Z ≥ R ≤ Pr |gi | ≥ . 1+ξ Cξ Cξ

(18)

Recall that |gi |2 = |(p + iq) · si |2 = X12 + X22 ∼ χ22 , therefore, by the standard concentration inequality of χ22 random variables [LM00], we have Pr[|gi | ≥ t] = Pr[|gi |2 ≥ t2 ] ≤ exp(−t2 /6) p √ when t2 ≥ 2. Plugging in t = 12R/Cξ = 12 log(k/λ), and combining with (18), we obtain that:   |z1 | |z1 | ≥ R ≤ exp(−2 log(k/λ)) = (λ/k)2 . Pr |zi | ≥ 1+ξ Applying a union bound on all such events for i ∈ A1 (ξ), we obtain Pr[∃i ∈ A1 (ξ), |zi | ≥ |z1 |/(1 + ξ) | |z1 | ≥ R] ≤ λ2 /k. Therefore, we get that Pr[E1Sep |E1Cor,2 ] ≥ 1 − λ2 /k and hence Pr[E1Sep ∩ E1Cor ] ≥ Pr[E1Sep ∩ E1Cor,2 ] ≥ poly(λ/k), completing the proof of Lemma 4.13. 4.3.2

Estimating the Target Matrix Using Implicit Tensor Computation

In this section, we show how to use implicit tensor computation to estimate the target matrix (M). Our main result in this context is the following: Proposition 4.14 (Matrix Estimation Using Implicit Tensor Computation). Let D be a Gaussian sum of k ReLUs. Let p, N0 , Nu ∈ Z+ , ξ, λ, δ ∈ (0, 1) be the parameters prescribed in CDRviaImplicitTensor(D, λ, δ). Let M2p+2 (u) be the matrix defined in (M). Let Ub be a set of Nu random probing vectors following the iid c 2p+2 (u) computed in Algorithm 7 of distribution u = p + iq, p, q ∼ N (0, Id ). Then the matrices M CDRviaImplicitTensor(D, λ, δ) satisfy the following guarantee: with probability at least 1 − δ/2 over all c 2p+2 (u) − M2p+2 (u)∥2 ≤ √ξλ . the samples and the probing vectors, maxu∈Ub ∥M The rest of this subsection is devoted to the proof of Proposition 4.14. Our first observation is that M2p+2 (u) √ is obtained via contracting the moment tensor E(x,y)∼D [yHe2p+2 (x)] with Fp (u) ⊗ Fp (ū), Fp (u) = u⊗p / p!, where we remind the reader that ū is the complex conjugate of the vector u. 34

√ Claim 4.15. Let M2p+2 (u) be the matrix defined by (i) in (M). Let Fp (u) = u⊗p / p!. Then:     1 M2p+2 (u) = E yHe2p+2 (x) , Fp (u) ⊗ Fp (ū) . σ b2p+2 (x,y)∼D Proof. By the Hermite expansion of the ReLU activation σ, we have E (x,y)∼D

  X k yHe2p+2 (x) = βi i=1

E (x,y)∼D

  X k ⊗(2p+2) σ(wi · x)He2p+2 (x) = βi σ b2p+2 wi . i=1

Taking the tensor product and noting that ⟨wi⊗2p , u⊗p ⊗ ū⊗p ⟩ = (wi · u)p (wi · ū)p = |wi · u|2p gives the stated claim. Therefore, the problem boils down to estimating the contraction of two tensors E(x,y)∼D [yHe2p+2 (x)] and Fp (u) ⊗ Fp (ū) using samples from D. To this end, we apply the implicit tensor technique from [DK25]. Preliminaries for Implicit Tensor Computation. The central idea is the following: Recall that the target matrix M2p+2 (u) is defined by M2p+2 (u) =

1

E

σ b2p+2

(x,y)∼D

   X k yHe2p+2 (x) , Fp (u) ⊗ Fp (ū) = ⟨wi⊗p , Fp (u)⟩⟨wi⊗p , Fp (ū)⟩wi ⊗ wi . i=1

Let h1 , . . . , hk ∈ (Rd )⊗p be the basis of the subspace U ⊂ (Rd )⊗p spanned by “vectors” wi⊗p , i ∈ [k]. p d ⊗p Suppose we have the projection Φp ∈ Rk×d , Φ⊤ to p = [h1 , . . . , hk ] that maps a “vector” from (R ) ⊗p d ⊗p k the subspace U ⊂ (R ) followed by an isometry from U to R , then since wi is a “vector” in U, ⊗p ⊗p it holds Φ⊤ p Φp (wi ) = wi . In addition, since Φp is a matrix with real entries, consider its natural complex extension on a complex tensor Fp (u): Φp (Fp (u)) = Φp (Re(Fp (u))) + iΦp (Im(Fp (u))). Then its Hermitian adjoint is Φ∗p = Φ⊤ p . Therefore, we have ⊗p ⊗p ⊗p ⊤ ∗ ⟨wi⊗p , Fp (u)⟩ = ⟨Φ⊤ p Φp (wi ), Fp (u)⟩ = ⟨Φp (wi ), (Φp ) (Fp (u))⟩ = ⟨Φp (wi ), Φp (Fp (u))⟩.

By the linearity of expectation, this implies that the target matrix M2p+2 (u) can be written as: M2p+2 (u) =

k X

⟨Φp (wi⊗p ), Φp (Fp (u))⟩⟨Φp (wi⊗p ), Φp (Fp (ū))⟩wi ⊗ wi

i=1

=

k  X

(Φp ⊗ Φp ⊗ Id ⊗ Id )(wi⊗2p+2 ), Φp (Fp (u)) ⊗ Φp (Fp (ū))

i=1

=

E (x,y)∼D

  (Φp ⊗ Φp ⊗ Id ⊗ Id )

1 σ b2p+2

  yHe2p+2 (x) , Φp (Fp (u)) ⊗ Φp (Fp (ū)) . (19)

Notably, the right-hand side of the equality above is a product of a k × k × d × d tensor and a k × k matrix, rather than two higher-order tensors. This implies that the dimension of the estimation problem is greatly reduced from d2p to k 2 + (dk)2 , via the “compressor” Φp . The main idea of [DK25] is that, Pk when i=1 wi⊗t = EI [St (I)] for t = 2, 4, . . . , 2p, where St (I) are recursively-constructed random tensors depending on the random variable I, the compressor Φp can be discovered and constructed recursively and efficiently from samples I. Furthermore, Φp can also be efficiently applied to other recursively-constructed tensors. These recursively-constructed tensors are formally defined as tensors generated from Sequential Tensor Computation (STC) circuits. Definition 4.16 (Sequential Tensor Computation [DK25]). A sequential tensor computation (STC) is an arithmetic circuit St that takes as input I a set of scalars and vectors in Cd and outputs an order-t tensor of dimension d, denoted by St (I) ∈ (Cd )⊗t , for some t ∈ Z+ , by applying the following operations: scalar multiplication, addition of objects of the same order and dimension, and appending a vector by tensor product (in the last component). When I is sampled from a distribution D(I), we denote St (I) ∼ St (D(I)). The size of St is defined as the size of the underlying arithmetic circuit. 35

In short, Sequential Tensor Computation (STC) circuits capture a class of tensors that can be constructed via recursive simple operations. One simple example is order-p rank-1 tensors, i.e., tensors of the form u⊗p , since these tensors can be constructed by recursively adding a new vector in its last component via tensor product. These tensors are exactly the objectives we consider. Lemma 4.17 (Moment Tensor and Filter Tensor are STCs). Let D be a Gaussian sum of k ReLUs, and let z ∼ N (0, Id ). Let I = (x, y, z) and denote the distribution of I by D(I) = D ⊗ N (0, Id ). For even t ∈ Z+ , denote the tth Hermite coefficient of the ReLU activation σ by σ bt . Then, for any even t ∈ Z+ , there is a real STC St (D(I)) of size O(t) that given a vector I ∼ D(I), outputs: St (I) = St (x, y, z) :=

y √ Re((x + iz)⊗t ). σ bt t!

(20)

The tensor St (x, y, z) satisfies E

[St (x, y, z)] =

z∼N (0,Id )

1 yHet (x), σ bt

E

[St (I)] =

I∼D(I)

X

βi wi⊗t .

i=1

Furthermore, let u ∼ U, u = p + iq, p, q ∼ N (0, Id ). Then, for any p ∈ Z+ , there is a complex STC Fp (U) of size O(p) that given a vector u, outputs: u⊗p Fp (u) = √ . p!

(21)

√ σt t!). Define tensors Proof. We first show that St (D(I)) is a real STC of size O(t). Let a = y/(b Rj , Imj such that Rj + iImj = a(x + iz)⊗j . To construct j + 1, note that a(x + iz)⊗(j+1) = (Rj + iImj ) ⊗ (x + iz) = (Rj ⊗ x − Imj ⊗ z) + i(Rj ⊗ z + Imj ⊗ x) . This gives the recursive construction: Rj+1 = Rj ⊗ x − Imj ⊗ z and Imj+1 = Rj ⊗ z + Imj ⊗ x with R0 = a and Im0 = 0. Each step uses an STC circuit of size O(1), therefore the real STC St (D(I)) is of size O(t). Next, applying Fact 2.6, we have   y Re((x + iz)⊗t ) √ E [St (x, y, z)] = E = (1/b σt )yHet (x) . σ bt z∼N (0,Id ) z∼N (0,Id ) t! Therefore, by the law of total expectation, we have EI∼D(I) [St (I)] = (1/b σt ) E(x,y)∼D [yHet (x)] = Pk ⊗t i=1 βi wi , using the definition of Hermite expansions. For Fp (u), directly by the definition of STCs, we know that it can be generated by an STC Fp (U) of size O(p). This completes the proof of Lemma 4.17. Combining Claim 4.15 and Lemma 4.17, and using (19), the task of estimating the matrix M2p+2 (u) is reduced to estimating the tensor contraction of these two STCs mentioned above:   M2p+2 (u) = E [S2p+2 (I)], Fp (u) ⊗ Fp (ū) I∼D(I)      = E (Φp ⊗ Φp ⊗ Id ⊗ Id ) S2p+2 (I) , Φp (Fp (u)) ⊗ Φp (Fp (ū)) , (22) I∼D(I)

where Φp is the “compressor” we are looking for that maps from (Rd )⊗p to the subspace spanned by the hidden weight tensors wi⊗p . These compressors are formally called (recursive) pseudo-projections: Definition 4.18 (Recursive pseudo-projections; Definitions 3.4 and 3.6 in [DK25]). A real linear map Φ : U → V between finite-dimensional inner product spaces U and V is a pseudo-projection if ΦΦ⊤ = IV . A (d, t, k)-recursive pseudo-projection Φ has dimensions n0 = 1, n1 , . . . , nt = k and base pseudo-projections ϕj : Rnj−1 ⊗ Rd → Rnj such that Φ0 = Id ,

Φj = ϕj (Φj−1 ⊗ Id ),

Its order is t and its size is k. 36

1 ≤ j ≤ t,

Φ = Φt .

Remark 4.19 (Decomposition of Pseudo-Projections). As proved in [DK25, Lemma 3.5], every pseudoprojection Φ can be written as a projection P : U → U′ (i.e., P is a linear map satisfying PP⊤ = I) composed with an isometry R from U′ to V; and in addition, the composition of pseudo-projections is a pseudo-projection. This implies the recursive map defined above is indeed a pseudo-projection, as tensorization with an identity and composition preserves the equation Φj Φ⊤ j = Inj . Remark 4.20 (Complex Extension of the Pseudo-Projections). The mapping Φ is a linear map with real matrix entries. We also use the same notation for its complex-linear extension, defined by Φ(W + iT) = Φ(W) + iΦ(T),

W, T real tensors in real space U.

Thus, Φ can be applied to complex tensors. Since the matrix entries of Φ are real, its Hermitian adjoint coincides with its transpose Φ∗ = Φ⊤ , and ΦΦ∗ = IV , Φ(W + iT) = Φ(W + iT). In the rest of the section, whenever we apply this real linear map Φ to a complex tensor, we are applying its natural complex extension. Fact 4.21. Let Φ : U → V be a real pseudo-projection. The following all hold: 1. For any W ∈ V, ∥Φ⊤ W∥F = ∥W∥F . 2. For every T ∈ U, ∥ΦT∥F = ∥Φ⊤ ΦT∥F ≤ ∥T∥F . 3. For any V ∈ V, T ∈ U, ⟨ΦT, V⟩ = ⟨T, Φ⊤ V⟩. Proof. For Item 1, we have that ∥Φ⊤ W∥2F = ⟨Φ⊤ W, Φ⊤ W⟩ = ⟨W, ΦΦ⊤ W⟩ = ⟨W, W⟩ = ∥W∥2F . For Item 2, the first part follows by the same argument as Item 1. The second part follows from the fact that every pseudo-projection Φ can be written as a projection P : U → U′ (i.e., P is a linear map satisfying PP⊤ = I) composed with an isometry from U′ to V. To obtain the pseudo-projection Φp , we apply the results from [DK25]. However, the output of the main algorithm in [DK25] is a real-valued inner product between two projected tensors. In our application, we aim to utilize the recursive pseudo-projections explicitly. Due to this small difference, we slightly adapt the main algorithm from [DK25] to our setting, as stated in TensorCompressor(p, N0 , D(I)). Algorithm 8 TensorCompressor(p, N0 , D(I)) 1: Input: Parameters k, d, p, N0 ∈ Z+ ; sample access to D(I)

b p , represented by ϕ1 , . . . , ϕp . 2: Output: A recursive pseudo-projection Φ 3: Construct STC circuit St , t = 2, 4, . . . , 2p + 2, according to Lemma 4.17.

b 0 ← Id . 4: Φ 5: for r = 1, . . . , p do b r−1 ⊗ Id . 6: Ψr ← Φ

b r−1 : (Rd )⊗(r−1) → Rk , ▷ Φ Ψr : (Rd )⊗(r−1) ⊗ Rd → Rk ⊗ Rd b 7: Draw N0 fresh samples from D(I), construct empirical distribution D(I). 8: Let S2r (I) be the output of S2r on sample I. 9: Vr ← EI∼D(I) [(Ψr ⊗ Ψr )S2r (I)] ▷ Viewed as a d × d matrix for r = 1 b and a dk × dk matrix otherwise. 10: Ur ← Top(k)Eigenvectors[Vr ] b 11: ϕr ← U⊤ r ; Φr ← ϕr Ψr . b p. 12: return Φ TensorCompressor(p, N0 , D(I)) has the following guarantee: Proposition 4.22 (Implicit Tensor Computation). Let D be a Gaussian sum of k ReLUs and denote Pk Mp = i=1 βi wi⊗p . Let I = (x, y, z) where (x, y) ∼ D, z ∼ N (0, Id ) and let St (D(I)) be the STC defined in Lemma 4.17, Equation (20). Let p ∈ Z+ . Suppose that for all positive even integers t ≤ 2p, the output tensor St (I), I ∼ Dt (I), has covariance bounded above by VS , for some VS > 0. 37

Then TensorCompressor(p, N0 , D(I)) has the following guarantee: the algorithm draws pN0 samples b p from (Rd )⊗p from D(I), runs in poly(d, p, N0 ) time, outputs a (d, p, k)-recursive pseudo-projection Φ k b to R , such that with probability at least 1 − τ (over the samples drawn from D(I)), Φp satisfies p p ⊗p ⊗p b⊤ b b⊤ b ∥Φ p Φp (Mp ) − Mp ∥F ≤ kp ζ, and for each i ∈ [k], ∥Φp Φp (wi ) − wi ∥F ≤ p ζ/βi , √ where ζ = O(pdkVS / N0 τ ). P Proof of Proposition 4.22. Recall we proved that EI∼D(I) [St (I)] = i=1 βi wi⊗t in Lemma 4.17, which satisfies the prerequisites of Proposition 3.2 in [DK25]. Then the proof follows directly from [DK25, Lemma 3.9 and Claim 3.10]. Proposition 4.22 states that, using poly(d, k, p, 1/τ, VS ) samples, we can approximately recover b p will almost project to the the compressor Φp we are looking for: the recursive pseudo-projection Φ ⊗p subspace spanned by the hidden weight tensors wi whose coefficients are not too small. Applying Implicit Tensor Computation for Target Matrix Approximation. To apply TensorCompressor(p, N0 , D(I)) and Proposition 4.22, we need to determine the covariance of the STC output St (I) defined in Lemma 4.17. To do so, we establish the following key lemma: Lemma 4.23. Let D be a Gaussian sum of k ReLUs, and let z ∼ N (0, Id ). Let I = (x, y, z) and denote the distribution of I by D(I) = D ⊗ N (0, Id ). Let St (I) be the output of the STC circuit defined in Lemma 4.17, Equation (20). Then the covariance of St (I) is bounded by sup

[|⟨St (I) − E [St (I)], W⟩|2 ] ≤ VSt := k 2 2O(t) /b σt2 .

E

I∼D

∥W∥F =1 I∼D(I)

Proof. By the definition of tensor covariance and the distribution of the labels y, we have:  2 y √ Re((x + iz)⊗t ), W σ bt t!   4 [y 4 ] E , Re((x + iz)⊗t ), W

 sup

E

∥W∥F =1 (x,y,z)∼D(I)

≤

1

s

sup bt2 t! ∥W∥F =1 σ

E (x,y)∼D

(23)

x,z∼N (0,Id )

where the last inequality is by Cauchy-Schwarz. We prove the following claim: Claim 4.24. For any W such that ∥W∥F = 1, it holds that s   4 E Re((x + iz)⊗t ), W ≤ 2O(t) t!. x,z∼N (0,Id )

Proof. Note that |⟨(Re((x + iz)⊗t ), W⟩|2 = (⟨Re((x + iz)⊗t ), Re(W)⟩)2 + (⟨Re((x + iz)⊗t ), Im(W)⟩)2 , where a(x, z) = ⟨Re((x + iz)⊗t ), Re(W)⟩ and b(x, z) = ⟨Re((x + iz)⊗t ), Im(W)⟩ are both degree-t polynomials. Therefore, by Gaussian Hypercontractivity (Fact 2.2), we have that 1/2

∥a(x, z)2 ∥L2 (N ) =

 E

1/4  [a(x, z)4 ] ≤ 3t/2

x,z∼N (0,Id )

E x,z∼N (0,Id )

1/2 [a(x, z)2 ] = 3t/2 ∥a(x, z)∥L2 (N ) ,

and a similar inequality holds for b(x, z). Therefore, we obtain: 

 E

Re((x + iz)⊗t ), W

4 1/2

 =

x,z∼N (0,Id )

E

 2 1/2 a(x, z)2 + b(x, z)2

x,z∼N (0,Id )

= ∥a(x, z)2 + b(x, z)2 ∥L2 (N ) ≤ ∥a(x, z)2 ∥L2 (N ) + ∥b(x, z)2 ∥L2 (N ) ≤ 3t (∥a(x, z)∥2L2 (N ) + ∥b(x, z)∥2L2 (N ) ) .

38

By [DK25, Lemma 4.3], we have that ∥a(x, z)∥2L2 (N ) =

 2  ⊗t Re((x + iz) ), Re(W)

E x,z∼N (0,Id )

 ≤

sup

E

f F =1 x,z∼N (0,Id ) ∥W∥

f Re((x + iz)⊗t ), W

2

≤ t!2O(t) .

The same inequality also holds for b(x, z). Thus we conclude that: 

 E

⊗t

4 1/2



t

≤3

Re((x + iz) ), W

⊗t

2

Re((x + iz) ), W

E x,z∼N (0,Id )

x,z∼N (0,Id )

≤ t! 3t 2O(t) = t! 2O(t) .

Pk Therefore, by plugging Ex∼N (0,Id ) [( i=1 βi σ(wi · x))4 ] ≲ k 4 back into (23), we obtain:  sup

E

∥W∥F =1 (x,y,z)∼D(I)

y √ Re((x + iz)⊗t ), W ct t!

2 ≲

k 2 2O(t) k 2 2O(t) t! = . σ bt2 t! σ bt2

This completes the proof of Lemma 4.23. Combining Lemma 4.23 with Proposition 4.22 immediately implies that when p = O(log(k/λ)) (as was set in CDRviaImplicitTensor(D, λ, δ)), using N0 = poly(d, k, 1/τ )2O(p) = poly(d, k/λ, 1/τ ) samples b p maps a tensor O(λ)-close to the space of the λ-heavy hidden weight from D, the pseudo-projection Φ tensors wi⊗p (i.e., βi ≥ λ) with probability at least 1 − τ . The algorithm TensorCompressor(p, N0 , D(I)) runs efficiently, in time poly(d, p, N0 ). In short, we can efficiently approximate the compressor Φp using b p to reduce the problem dimension. Recalling from (22) that the target matrix equals Φ      E (Φp ⊗ Φp ⊗ Id ⊗ Id ) S2p+2 (I) , Φp (Fp (u)) ⊗ Φp (Fp (ū)) , M2p+2 (u) = I∼D(I)

an estimation of M2p+2 (u) can therefore be constructed via:      c b b b b M2p+2 (u) = E (Φp ⊗ Φp ⊗ Id ⊗ Id ) S2p+2 (I) , Φp (Fp (u)) ⊗ Φp (Fp (ū)) ,

(EstM)

b I∼D(I)

b where D(I) is a fresh empirical distribution of D(I) (independent of the samples used to construct b Φp ). The two tensors on the right-hand side of (EstM) are computed in the following procedures: b p , D(I)) and CompressFilter(p, Φ b p , u). In CDRviaImplicitTensor(D, λ, δ), CompressEmpMoment(p, N0 , Φ c Algorithm 7, we compute the empirical matrix M2p+2 (u) displayed in (EstM), using the tensors b p , D(I)) and CompressFilter(p, Φ b p , u). returned by CompressEmpMoment(p, N0 , Φ b p , D(I)) Algorithm 9 CompressEmpMoment(p, N0 , Φ b p , sample access to D(I). 1: Input: Moment degree p, number of samples N0 , tensor compressor Φ 2: Output: A tensor Ap ∈ Rk×k×d×d N

0 b 3: Draw N0 fresh samples D(I) = {I (i) }i=1 from D(I). 4: Construct STC circuit S2p+2 , according to Lemma 4.17.

bp ⊗Φ b p ⊗Id ⊗Id )S2p+2 (I (i) ), where S2p+2 (I (i) ) is the output of S2p+2 , using Lemma 4.25. 5: Compute (Φ PN0

(i) b b 6: Ap ← N1 i=1 (Φp ⊗ Φp ⊗ Id ⊗ Id )S2p+2 (I ). 0 7: Return: Ap .

39

b p , u) Algorithm 10 CompressFilter(p, Φ b p , probing vector u. 1: Input: Moment degree p, tensor compressor Φ k×k 2: Output: A matrix Bp (u) ∈ R 3: Construct STC circuits Fp according to Lemma 4.17.

b p (Fp (u)) and Φ b p (Fp (u)), where Fp (u) is the output of Fp , using Lemma 4.25. 4: Compute Φ b p (Fp (u)) ⊗ Φ b p (Fp (u)) 5: Bp (u) ← Φ 6: Return: Bp (u). We next show that these procedures run efficiently. Lemma 4.25 (Efficient Computation of Tensor Projections). Let St (D(I)) and Fp (U) be the STCs defined in Lemma 4.17, equations (20) and (21) respectively. Let Φp be a (d, p, k)-recursive pseudob p (Fp (u)) can be computed in time poly(d, k, p); and for any projection. Then, for any Fp (u) ∼ Fp (U), Φ b b p ⊗ Id ⊗ Id )(S2p+2 (I)) can be computed in time poly(d, k, p). S2p+2 (I) ∼ S2p+2 (D(I)), the tensor (Φp ⊗ Φ Proof. [DK25, Lemma 3.7] states that given an order-t recursive pseudo-projection Ψ along with their defining pseudo-projections θj , and an STC S of order t, dimension d, and size S, the projecb p (Fp (u)) can be tion Ψ(S) can be computed in time poly(d, t, S). This immediately implies that Φ computed in poly(d, p) time. Furthermore, according to [DK25, Lemma 3.8], given recursive pseudoprojections Φ : (Rd )⊗t → Rn and Ψ : (Rd )⊗s → Rm along with their defining pseudo-projections ϕi , θj , one can efficiently construct the recursive pseudo-projection Φ ⊗ Ψ : (Rd )⊗(t+s) → Rnm of size max(size(Φ), n size(Ψ)) in time poly(d, t, s, max(size(Φ), n size(Ψ))). Therefore, the recursive-pseudo bp ⊗ Φ b p : (Rd )⊗p ⊗ (Rd )⊗p → Rk ⊗ Rk can be efficiently constructed. Since Id : Rd → Rd projection Φ is also a pseudo-projection, we can similarly efficiently construct the recursive pseudo-projection bp ⊗ Φ b p ) ⊗ Id , and then construct ((Φ bp ⊗ Φ b p ) ⊗ Id ) ⊗ Id . The process takes time poly(d, k, p). After (Φ bp ⊗ Φ b p ⊗ Id ⊗ Id , since S2p+2 (I) is generated from a size O(p) constructing the pseudo-projection Φ b b STC, the projection (Φp ⊗ Φt ⊗ Id ⊗ Id )S2p+2 (I) can also be computed in poly(d, k, p) time. Main Result of this Subsection: Matrix Estimation Using Implicit Tensor Computation. What remains is to prove that the empirical matrix estimation in (EstM) is close to the target matrix M2p+2 (u). However, we need one last technical piece, which is to bound the covariance of the tensor Fp (u). This is the place where we crucially use the property of complex Gaussian tensors: in Lemma 4.26, we show that the covariance of the complex Gaussian filter Fp (u) is bounded above by 2O(p) , whereas in Lemma 4.27 we prove that were we using a real Gaussian filter, the covariance would be much larger. To see why a bound on the covariance of Fp (u) is required, consider the difference between the empirical matrix and the target matrix:      c 2p+2 (u) − M2p+2 (u) = bp ⊗ Φ b p ⊗ Id ⊗ Id ) S2p+2 (I) , Φ b p (Fp (u)) ⊗ Φ b p (Fp (ū)) M E (Φ b I∼D(I)

−

E

I∼D(I)

   S2p+2 (I) , Fp (u) ⊗ Fp (ū)

b p (Fp (u)) ⊗ Φ b p (Fp (ū))⟩ + ⟨∆2 , Fp (u)⟩ + ⟨∆2 , Fp (ū)⟩ , = ⟨∆1 , Φ where ∆1 , ∆2 , and ∆3 are error terms (see, e.g., (25) in the proof of Proposition 4.14). Even though b p (Fp (u)) ⊗ Φ b p (Fp (ū)) is large, the ∆1 , ∆2 , and ∆3 can be tiny, if the spectral norm of Fp (u) or Φ b p (Fp (u)) ⊗ Φ b p (Fp (ū)) difference displayed above is not necessarily small. To argue that Fp (u) and Φ have small spectral norm with large probability, we bound their covariance and then apply standard concentration inequalities. Lemma 4.26. Let u ∼ U, u = p+iq, p, q ∼ N (0, Id ), and let Fp (u) be the output of the STC Fp (U) defined in (21) in Lemma 4.17. Then, we have Eu∼U [|⟨Fp (u), W⟩|2 ] ≤ 2O(p) ∥W∥2F , Eu∼U [|⟨Fp (u), W⟩|4 ] ≤ b p be the complex extension of a real-valued pseudo-projection from (Rd )⊗p 2O(p) ∥W∥4F . In addition, let Φ k to R . Then: b p (Fp (u)) ⊗ Φ b p (Fp (ū))∥2F ] ≤ k 2 2O(p) . E [∥Φ

u∼U

40

Proof. Let Z(u) = ⟨Fp (u), W⟩. The random variables Z(u) and iZ(u) have the same distribution; to see that, note that for any θ ∈ R, we have exp(iθ)u = cos θp − sin θq + i(cos θq + sin θp) and both the real part and imaginary part are independent standard normal from orthogonality. Therefore, Z(exp(−iπ/(2p))u) = iZ(u) and because Z(exp(−iπ/(2p))u) has the same distribution as Z(u), it holds that Z(u) and i(Z(u)) have the same distribution. This further indicates that Re(Z(u)) and Im(Z(u)) have the same distribution. Similarly, Re(Fp (u)) and Im(Fp (u)) also have the same distribution. Taking imaginary parts, we conclude that Re(Z(u)) and Im(Z(u)) have the same distribution. The same rotation also shows that Re(Fp (u)) and Im(Fp (u)) have the same distribution. By [DK25, Lemma 4.3], for every real tensor T of order p, we have E

u∼U

    1 ⟨Re(Fp (u)), T⟩2 = E ⟨Re(u⊗p ), T⟩2 ≤ 2O(p) ∥T∥2F . p! u∼U

Since Re(Fp (u)) and Im(Fp (u)) have the same distribution, the same bound holds for Im(Fp (u)). Since the inner product is conjugate-linear in its first argument, we have Re(Z(u)) = ⟨Re(Fp (u)), Re(W)⟩ + ⟨Im(Fp (u)), Im(W)⟩. Using (a + b)2 ≤ 2a2 + 2b2 , it follows that       E Re(Z(u))2 ≤ 2 E ⟨Re(Fp (u)), Re(W)⟩2 + 2 E ⟨Im(Fp (u)), Im(W)⟩2 u∼U u∼U u∼U
≤ 2O(p) ∥Re(W)∥2F + ∥Im(W)∥2F = 2O(p) ∥W∥2F . The inequality also holds for Im(Z(u)). Therefore,       E |⟨Fp (u), W⟩|2 = E Re(Z(u))2 + Im(Z(u))2 = 2 E Re(Z(u))2 ≤ 2O(p) ∥W∥2F . u∼U

u∼U

u∼U

√ For the fourth moment, we apply Claim 4.24. By the definition of Fp (u) = u⊗p / p!, for every real tensor T of order p, we have E

u∼U

  |⟨Re(Fp (u)), T⟩|4 =

  1 E |⟨Re(u⊗p ), T⟩|4 ≤ 2O(p) ∥T∥4F . 2 (p!) u∼U

Since Re(Fp (u)) and Im(Fp (u)) have the same distribution, the same bound holds with Re(Fp (u)) replaced by Im(Fp (u)). Recalling that Re(Z(u)) = ⟨Re(Fp (u)), Re(W)⟩ + ⟨Im(Fp (u)), Im(W)⟩, and using (a + b)4 ≤ 8(a4 + b4 ), we obtain E [Re(Z(u))4 ] ≤ 8 E [⟨Re(Fp (u)), Re(W)⟩4 ] + 8 E [⟨Im(Fp (u)), Im(W)⟩4 ] u∼U u∼U
O(p) 4 4 ≤2 ∥Re(W)∥F + ∥Im(W)∥F ≤ 2O(p) ∥W∥4F ,

u∼U

where the last inequality uses ∥Re(W)∥4F + ∥Im(W)∥4F ≤ ∥Re(W)∥2F + ∥Im(W)∥2F

2

= ∥W∥4F .

Finally, using (a2 +b2 )2 ≤ 2a4 +2b4 and the fact that Re(Z(u)) and Im(Z(u)) have the same distribution, we conclude that E [|⟨Fp (u), W⟩|4 ] = E [(Re(Z(u))2 + Im(Z(u))2 )2 ] ≤ 2 E [Re(Z(u))4 + Im(Z(u))4 ]

u∼U

u∼U

u∼U

4

O(p)

= 4 E [Re(Z(u)) ] ≤ 2 u∼U

∥W∥4F .

(24)

b p is a map with real entries, we have For the second part of the statement, note that since Φ b p (Fp (ū)) = Φ b p (Fp (u)), therefore, it holds that Φ b p (Fp (u)) ⊗ Φ b p (Fp (ū))∥2 ] ≤ E [∥Φ b p (Fp (u))∥2 ∥Φ b p (Fp (ū))∥2 ] E [∥Φ F F F

u∼U

u∼U

b p (Fp (u))∥2F ∥Φ b p (Fp (u))∥2F ] = E [∥Φ b p (Fp (u))∥4F ]. = E [∥Φ u∼U

u∼U

41

b p maps Fp (u) to a vector in Ck , thus we have ∥Φ b p (Fp (u))∥2 = Pk |⟨Φ b p (Fp (u)), ei ⟩|2 . Recall that Φ 2 i=1 Furthermore, by Jensen’s inequality, it holds:  X 2   X  k k b p (Fp (u))∥42 ] = E b p (Fp (u)), ei ⟩|2 b p (Fp (u)), ei ⟩|4 E[∥Φ |⟨Φ ≤E k |⟨Φ u

u

u

i=1

=k

i=1 k X i=1

4 2 O(p) b⊤ E[|⟨Fp (u), Φ , p (ei )⟩| ] ≤ k 2 u

b pΦ b⊤ b⊤ where for the last inequality we use that Φ p = Ik (and therefore ∥Φp ei ∥F = 1) and (24). For completeness, in the following lemma, we demonstrate the crucial difference between complex and real Gaussian tensors in terms of their covariance bounds, which justifies the use of complex Gaussians in our approach. Lemma 4.27 (Real vs Complex Gaussian Tensor Powers). Let d, p ≥ 1 be integers, let p, q ∼ N (0, Id ), and put u = p + iq. Define FR p (p) = p

p⊗p (2p − 1)!!

,

√ u⊗p FC = Fp (u)/ 2p . p (u) = √ p 2 p!

Both normalizations give, for every real unit vector w,   ⊗p 2 ⊗p 2 E |⟨FR ⟩| = E[|⟨FC ⟩| ] = 1. p (p), w p (u), w Their uncentered second moments in general tensor directions differ as follows. 1. For every complex order-p tensor R, E[|⟨Fp (u), R⟩|2 ] = ∥Sym(R)∥2F ≤ ∥R∥2F . √ 2. For every even p ≥ 2, the real tensor Rp = (Id / d)⊗p/2 has ∥Rp ∥F = 1 and Qp−1 dp/2 2 j=0 (d + 2j) R E[ ⟨Fp (p), Rp ⟩ ] = p/2 ≥ . (2p)p d (2p − 1)!! Proof. Let p be even. Multiplicativity of the Frobenius norm under tensor products gives ∥Rp ∥F = 1, and direct contraction gives ∥p∥p ⟨p⊗p , Rp ⟩ = p/42 . d 2

The radial Gaussian density is proportional to rd−1 e−r /2 . Integration by parts gives E ∥p∥2s 2 = Qp−1 2p (d + 2s − 2) E ∥p∥2s−2 for s ≥ 1, so E ∥p∥ = (d + 2j). Squaring the contraction and dividing by 2 2 j=0 (2p − 1)!! proves the equality. Its lower bound follows from d + 2j ≥ d and (2p − 1)!! ≤ (2p)p . We are now ready to prove the main result of this subsection, Proposition 4.14. Proposition 4.14 (Matrix Estimation Using Implicit Tensor Computation). Let D be a Gaussian sum of k ReLUs. Let p, N0 , Nu ∈ Z+ , ξ, λ, δ ∈ (0, 1) be the parameters prescribed in CDRviaImplicitTensor(D, λ, δ). Let M2p+2 (u) be the matrix defined in (M). Let Ub be a set of Nu random probing vectors following the iid c 2p+2 (u) computed in Algorithm 7 of distribution u = p + iq, p, q ∼ N (0, Id ). Then the matrices M CDRviaImplicitTensor(D, λ, δ) satisfy the following guarantee: with probability at least 1 − δ/2 over all c 2p+2 (u) − M2p+2 (u)∥2 ≤ √ξλ . the samples and the probing vectors, maxu∈Ub ∥M Proof of Proposition 4.14. Let D(I) be the distribution of I = (x, y, z) where (x, y) ∼ D and z ∼ b p returned by TensorCompressor(p, N0 , D(I)) N (0, Id ). By Proposition 4.22, the pseudo-projection Φ (see Algorithm 7 in CDRviaImplicitTensor(D, λ, δ)) satisfies: with probability 1 − τ0 over the samples drawn in TensorCompressor(p, N0 , D(I)), the following event holds s   p ⊗p ⊗p ⊤ ⊤ b p satisfies ∥Φ bp Φ b p (Mp ) − Mp ∥F ≤ kp ζ, ∥Φ bp Φ b p (w ) − w ∥F ≤ p ζ , ∀i ∈ [k] . E0 := Φ i i βi

42

√ 2 where ζ = O(pdkVS / N0 τ0 ). By Lemma 4.23, we have that VS ≤ k 2 2O(p) /b σ2p+2 . From now on, we condition on this event E0 . b p , D(I)), we draw N0 new samples from D(I) and construct a new In CompressEmpMoment(p, N0 , Φ bnew (I). Then, using the pseudo-projection Φ b p returned empirical distribution, which we denote by D b from TensorCompressor(p, N0 , D(I)), CompressEmpMoment(p, N0 , Φp , D(I)) constructs b 2p+2 , bp ⊗ Φ b p ⊗ I d ⊗ I d )S Ap = (Φ

b 2p+2 = S

E

[S2p+2 (I)].

bnew (I) I∼D

b p , u) generates: Given a probing vector u from Ub (of Nu probing vectors in total), CompressFilter(p, Φ b p (Fp (u)) ⊗ Φ b p (Fp (ū)). Bp (u) = Φ Then in Algorithm 7, we compute the empirical estimation (EstM) c 2p+2 (u) = ⟨Ap , Bp (u)⟩ = ⟨(Φ b 2p+2 , Φ bp ⊗ Φ b p ⊗ I d ⊗ I d )S b p (Fp (u)) ⊗ Φ b p (Fp (ū))⟩. M Let a, b ∈ [d], and define  c 2p+2 (u))ab = (M

b 2p+2 ), Φ bp ⊗ Φ b p ⊗ Id ⊗ Id )(S b p (Fp (u)) ⊗ Φ b p (Fp (ū)) ⊗ ea ⊗ eb (Φ

,

which is the (a, b)-entry of the empirical matrix. bnew (I) Lemma 4.28. Conditioned on E0 , with probability at least 1 − 4τ0 (over the fresh N0 samples D and u), for any indices a, b ∈ [d], it holds: c 2p+2 (u))ab − (M2p+2 (u))ab | ≤ |(M

dk 3 p2 2O(p) 3/2

1/4

|b σ2p+2 |τ0 N0

.

Proof. For this proof, we use the following notation: bp ⊗ Φ b p ⊗ I d ⊗ Id Ψ2p = Φ

b p (Fp (u)) ⊗ Φ b p (Fp (ū)) Bp (u) = Φ

c 2p+2 (u))ab and (M2p+2 (u))ab can be decomposed as: First, note that the difference between (M c 2p+2 (u))ab − (M2p+2 (u))ab (M   X  ⊗(2p+2) b 2p+2 ), Bp (u) ⊗ ea ⊗ eb − = Ψ2p (S βi wi , Fp (u) ⊗ Fp (ū) ⊗ ea ⊗ eb i=1

=

 b 2p+2 − Ψ2p S

|

  E [S2p+2 (I)] , Bp (u) ⊗ ea ⊗ eb I∼D(I) {z } I1

+ Ψ2p

 E

[S2p+2 (I)] −

I∼D(I)

X

⊗(2p+2) β i wi



, Bp (u) ⊗ ea ⊗ eb

i=1

|

{z

}

I2

X   X  ⊗(2p+2) ⊗(2p+2) + Ψ2p β i wi , Bp (u) ⊗ ea ⊗ eb − βi wi , Fp (u) ⊗ Fp (ū) ⊗ ea ⊗ eb . i=1

|

i=1

{z

}

I3

|

{z I4

} (25)

We bound each term respectively. Claim 4.29. With probability at least 1 − τu over the probing vector u, and probability at least 1 − τ2p+2 bnew (I), it holds: over the samples from D |I1 | ≤

dk 2 2O(p) p . |b σ2p+2 | τ2p+2 τu N0 43

b 2p+2 − EI∼D(I) [S2p+2 (I)]. Note that Ψ2p (W) ∈ Rk×k×d×d , we have Proof. Let W = S 

X

E[∥Ψ2p (W)∥2F ] =

E

2 W, Ψ⊤ . (e ⊗ e ⊗ e ⊗ e ) i2 i3 i4 2p i1

i1 ,i2 ∈[k],i3 ,i4 ∈[d]

Note that by Item 1 from Fact 4.21, it holds ∥Ψ⊤ 2p (ei1 ⊗ ei2 ⊗ ei3 ⊗ ei4 )∥F = 1. Thus, by Lemma 4.23, for any i1 , i2 ∈ [k], i3 , i4 ∈ [d], we have:   2 ⊤ E W, Ψ2p (ei1 ⊗ ei2 ⊗ ei3 ⊗ ei4 ) ≤

1 E N0 I∼D(I)

 S2p+2 (I) −

E

2 [S2p+2 (I)], Ψ⊤ (e ⊗ e ⊗ e ⊗ e ) i i i i 2p 1 2 3 4

I∼D(I)

2 O(p)

≤

VS k 2 ≤ 2 . N0 σ b2p+2 N0

2 This implies E[∥Ψ2p (W)∥2F ] ≤ d2 k 4 2O(p) /(b σ2p+2 N0 ). Therefore, by Chebyshev’s inequality, with bnew (I): probability at least 1 − τ2p+2 over the samples D s k 4 d2 2O(p) k 2 d2O(p) p . ∥Ψ2p (W)∥F ≤ = 2 σ b2p+2 τ2p+2 N0 |b σ2p+2 | τ2p+2 N0

Let us denote this event by E∥W∥F , and condition on this event: Pr

[|I1 | ≤ t] ≥

bnew (I) u,D

Pr

  [|I1 | ≤ t, E∥W∥F ] ≥ (1 − τ2p+2 ) 1 − Pr[|I1 | ≥ t | E∥W∥F ] . u

bnew (I) u,D

(26)

Using Chebyshev’s inequality again with respect to the probing vector u, we have Pr[|I1 | ≥ t | E∥W∥F ] = Pr[|⟨Ψ2p (W), Bp (u) ⊗ ea ⊗ eb ⟩| ≥ t | E∥W∥F ] u

u

E[|⟨Ψ2p (W), Bp (u) ⊗ ea ⊗ eb ⟩|2 | E∥W∥F ] ≤ t2

(27)

By Lemma 4.26, it holds E[∥Bp (u) ⊗ ea ⊗ eb ∥2F ] ≤ E[∥Bp (u)∥2F ] ≤ k 2 2O(p) . Using the fact that ⟨W, Bp (u)⟩ ≤ ∥W∥F ∥Bp (u)∥F , we obtain  2  E Ψ2p (W), Bp (u) ⊗ ea ⊗ eb E∥W∥F 2

 ≤E

2

Bp (u) ⊗ ea ⊗ eb

Ψ2p (W) F

≤

d2 k 4 2O(p) E 2 σ b2p+2 τ2p+2 N0

 E∥W∥F

F 2

 Bp (u) ⊗ ea ⊗ eb

≤ F

d2 k 6 2O(p)

. 2 σ b2p+2 τ2p+2 N0

(28)

Therefore, combining (28) and (27) and plugging back into (26), we obtain that with probability at least (1 − τ2p+2 )(1 − τu ) ≥ 1 − τu − τ2p+2 , it holds |I1 | ≤

dk 3 2O(p) p . |b σ2p+2 | τ2p+2 τu N0

P ⊗(2p+2) For the term I2 , since EI∼D(I) [S2p+2 (I)] = i=1 βi wi by Lemma 4.23, we have I2 = 0. It remains to bound the last two terms in (25). Claim 4.30. With probability at least 1 − τu we have: s dk 5 p3 2O(p) √ |I3 − I4 | ≤ . 2 σ b2p+2 τu N0 τ0 44

b p is a composition of a Proof. First, recall that by the property of pseudo-projectors we know that Φ d ⊗p b p = Rp Pp , therefore, isometry Rp and a projection operator Pp that maps (R ) to a subspace Up : Φ b b for two tensors T and W, it holds ⟨Φp (T), Φp (W)⟩ = ⟨Rp (Pp (T)), Rp (Pp (W))⟩ = ⟨P⊤ p Pp (T), W⟩ since R⊤ R = I as R is an isometry. p k p p P ⊗(2p+2) Note that the linear operator Ψ2p acting on i=1 βi wi is simply: bp ⊗ Φ b p ⊗ Id ⊗ Id ) (Φ

X

⊗(2p+2)



βi wi

=

i=1

k X

b p (w⊗p ) ⊗ Φ b p (w⊗p ) ⊗ wi ⊗ wi , βi Φ i i

i=1

and its tensor product with Bp (u) ⊗ ea ⊗ eb is  X   ⊗(2p+2) b b I3 = (Φp ⊗ Φp ⊗ Id ⊗ Id ) βi wi , Bp (u) ⊗ ea ⊗ eb i=1

=

k X

b p (w⊗p ) ⊗ Φ b p (w⊗p ) ⊗ wi ⊗ wi , Φ b p (Fp (u)) ⊗ Φ b p (Fp (ū)) ⊗ ea ⊗ eb Φ i i

⊗p b ⊗p b b b Φp (wi ), Φp (Fp (u)) Φp (wi ), Φp (Fp (ū)) (wi )a (wi )b

βi

i=1

=

k X

βi

i=1

=

k X

b ⊤Φ b p (w⊗p ), Fp (u) b ⊤Φ b p (w⊗p ), Fp (ū) (wi )a (wi )b βi Φ Φ p p i i

i=1

=

k X

2 ⊗p b⊤ b βi ⟨Φ p Φp (wi ), Fp (u)⟩ (wi )a (wi )b ,

i=1

where we used that Fp (ū) = Fp (u) and (wi )a is the a-th entry of the vector wi . On the other hand, we have I4 =

X

⊗(2p+2) β i wi , Fp (u) ⊗ Fp (ū) ⊗ ea ⊗ eb

i=1

=

k X

2

βi ⟨wi⊗p , Fp (u)⟩ (wi )a (wi )b .

i=1

Using the complex inequality ||x|2 − |y|2 | ≤ |x − y|(|x| + |y|), the difference between these terms equals: |I3 − I4 | ≤

k X

⊗p ⊗p b⊤ b βi ⟨ Φ p Φp (wi ) − wi , Fp (u)⟩



⊗p ⊗p b⊤ b ⟨Φ p Φp (wi ), Fp (u)⟩ + ⟨wi , Fp (u)⟩



|(wi )a ||(wi )b | .

i=1

To bound the variance, we use Cauchy–Schwarz and the crude bound (X 2 + Y 2 )2 ≤ 2(X 4 + Y 4 ) to get r hD E 4i h D E 4i X ⊗p ⊗p ⊗p 2 2 ⊤Φ b⊤ b b b E[|I3 − I4 | ] ≤ 2k βi E Φ Φ (w ) − w , F (u) E Φ (w ), F (ū) |(wi )a (wi )b |2 p p p p p p u

+ 2k

i

u

i=1 k X

βi2

i

u

i

r hD E 4i h D E 4i ⊗p ⊗p b⊤ b E wi⊗p , Fp (u) E Φ |(wi )a (wi )b |2 . p Φp (wi ) − wi , Fp (ū)

i=1

u

u

⊗p ⊗p b ⊤Φ b Using Lemma 4.26, since ∥Φ p p (wi )∥F ≤ ∥wi ∥F = 1 according to Fact 4.21, we further obtain:

E[|I3 − I4 |2 ] ≤ 2k u

k X

βi2

q

⊗p ⊗p 4 O(p) b⊤ b 2O(p) ∥Φ p Φp (wi ) − wi ∥F × 2

i=1

+ ≤ 4k

k X

q

 ⊗p ⊗p 4 2 b⊤ b 2O(p) × 2O(p) ∥Φ Φ (w ) − w ∥ p p i i F |(wi )a (wi )b |

⊗p ⊗p 2 b⊤ b 2O(p) βi2 ∥Φ p Φp (wi ) − wi ∥F .

i=1

45

⊗p ⊗p b⊤ b As we have conditioned on the event E0 , we know that for every i ∈ [k], it holds ∥Φ p Φp (wi )−wi ∥F ≤ p √ 2 . Thus, as βi ∈ (0, 1), we finally get that p ζ/βi , where ζ = O(pdkVS / N0 τ0 ) and VS ≤ k 2 2O(p) /b σ2p+2 conditioned on E0 , it holds

E[|I3 − I4 |2 ] ≤ 4k u

k X

2O(p) βi2 p2

i=1

ζ dk 5 p3 2O(p) . ≤ k 2 p2 2O(p) ζ ≲ 2 √ βi σ b2p+2 N0 τ0

Then, applying Markov’s inequality, with probability at least 1 − τu we have: s dk 5 p3 2O(p) √ . |I3 − I4 | ≤ 2 τu N0 τ0 σ b2p+2

In summary, after a union bound on these events, we have that with probability at least 1 − τ2p+2 bnew and with probability at least 1 − 3τu over the random probing vector over the fresh samples from D u, we have c 2p+2 (u))ab − (M2p+2 (u))ab | ≤ |I1 | + |I2 | + |I3 − I4 | |(M dk 2 2O(p) p ≤ + |b σ2p+2 | τ2p+2 τu N0

s

dk 5 p3 2O(p) √ . 2 σ b2p+2 τu N0 τ0

After rearranging the terms, we have: c 2p+2 (u))ab − (M2p+2 (u))ab | ≲ |(M

1 dk 3 p2 2O(p) , |b σ2p+2 |(τ2p+2 τu τ0 )1/2 N01/4

with probability at least (1 − τ2p+2 )(1 − 3τu ) ≥ 1 − 3τu − τ2p+2 (by the independence between D and U), conditioned on the event E0 . Letting τ0 = τ2p+2 = τu completes the proof of Lemma 4.28. By Lemma 4.28, using a union bound over all d2 pairs of indices a, b, and recalling from Fact 2.4 that |b σ2p+2 | ≳ 1/p3 , conditioned on the event E0 , we have that with probability at least 1 − 4d2 τ0 over bnew (I) and the probing vector u, it holds: the fresh samples D c 2p+2 (u) − M2p+2 (u)∥F ≤ ∥M

d poly(k, 1/τ0 )2O(p) 1/4

.

N0

b we have that with probability at Then, using another union bound over all Nu probing vectors from U, 2 least 1 − d Nu τ0 (conditioned on E0 ), it holds c 2p+2 (u) − M2p+2 (u)∥F ≤ d2 poly(k, p, 1/τ0 )2O(p) /N 1/4 . max ∥M 0

(29)

b u∈U

Recall that E0 happens with probability at least 1 − τ0 , thus, (29) holds with probability at least b Recall that we have 1 − (d2 Nu + 1)τ0 over all N = (p + 1)N0 samples and the probing vectors set U. 2 chosen τ0 ≲ δ/(d Nu ), p ≲ log(k/λ), Nu = poly(k/λ) log(1/δ) in CDRviaImplicitTensor(D, λ, δ). Hence, with N0 = poly(d, k, Nu /δ)2O(p) = poly(d, k/λ, 1/δ), (29) yields that with probability at least 1 − δ/2, c 2p+2 (u) − M2p+2 (u)∥2 ≤ max ∥M c 2p+2 (u) − M2p+2 (u)∥F ≤ max ∥M b u∈U

b u∈U

completing the proof of Proposition 4.14.

46

d2 poly(k, p, 1/τ0 )2O(p) 1/4 N0

≤

p ξλ,

4.4

Proof of the Main Result of the Section: Polynomial Complexity CDR

We are ready to prove Proposition 4.1, which is restated below. Proposition 4.1 (Polynomial-Time CDR). Let k, d ∈ Z+ and λ, δ ∈ (0, 1). Let D be a Gaussian sum of k ReLUs. Then, there is an algorithm (CDRviaImplicitTensor(D, λ, δ)) that draws N = poly(d, k, 1/λ, 1/δ) i.i.d. samples from D, runs in T = poly(d, N ) time, and outputs a list L of unit vectors of size |L| = poly(k/λ) log(1/δ) such that with probability at least 1 − δ, for any wi , i ∈ [k], with βi ≥ λ, there exists a vector ri ∈ L such that |ri · wi | ≥ 0.99. Proof of Proposition 4.1. CDRviaImplicitTensor(D, λ, δ) uses a total of (p + 1)N0 = poly(d, k/λ, 1/δ) samples. According to Proposition 4.22, line 9 takes poly(d, N0 ) time. Further, by Lemma 4.25, lines 10 b p , D(I)) and CompressFilter(p, Φ b p , u)) take runtime and 14 (algorithms CompressEmpMoment(p, N0 , Φ only poly(d, k, p). Thus, we conclude that the runtime of CDRviaImplicitTensor(D, λ, δ) is poly(dN0 ). It remains to argue the correctness of CDRviaImplicitTensor(D, λ, δ). Recall that by Lemma 4.13, fixing any i∗ ∈ [k] such that βi∗ ≥ λ, the “separation” and “correlation” events EiSep ∩ EiCor = {u : (1 + ξ)|wj · u| ≤ |wi∗ · u|, ∀j ∈ Ai∗ (ξ)} ∩ {u : |wi∗ · u| ≥ (Zp (U))1/(2p) } ∗ ∗ happen with probability at least poly(λ/k) over the choice of the random probing vector u = p + iq, p, q ∼ N (0, Id ). Since the number of probing vectors Nu is chosen to be Nu = poly(k/λ) log(k/δ), by the Chernoff bound, we know that drawing a set Ub of Nu independent probing vectors u, with probability at least 1 − δ/2, it holds that: Cor b s.t. u ∈ E Sep for every i∗ ∈ [k] satisfying βi∗ ≥ λ, ∃u ∈ U, i∗ ∩ Ei∗ .

(30)

On the other hand, by Proposition 4.14 we know that under the choice of parameters specified in CDRviaImplicitTensor(D, λ, δ), with probability at least 1 − δ/2 over all the samples from D and the b it holds random vectors u ∈ U, p c 2p+2 (u) − M2p+2 (u)∥2 ≤ ξλ. max ∥M (31) b u∈U

Thus, the joint probability that these two events, (30) and (31) hold at the same time is at least 1 − δ. Invoking Lemma 4.2, we obtain that with probability at least 1 − δ, for every i∗ ∈ [k] satisfying βi∗ ≥ λ, there exists a vector ri∗ in the list returned by Algorithm 7 that satisfies |ri∗ · wi∗ | ≥ 0.99. Note that the size of the list is 2Nu = poly(k/λ) log(1/δ). This completes the proof of Proposition 4.1.

5

Putting Everything Together: Proof of Theorem 1.3

By putting everything together, we obtain the following theorem, which is the formal version of Theorem 1.3. Theorem 5.1. Let d, k ∈ Z+ and ϵ > 0. Let (x, y) ∼ D be a distribution supported on Rd × R such that x ∼ N (0, Id ) and y = f (x) for an unknown f ∈ Cd,k . LearnSumReLUs(d, k, ϵ; D), using subroutine CDRviaImplicitTensor(D, λ, δ) with parameters λ = poly(ϵ/k), δ = poly(ϵ/k), has the following guarantee: with probability at least 2/3, it outputs a hypothesis yb(x) satisfying E(x,y)∼D [(y − yb(x))2 ] ≤ ϵ. LearnSumReLUs(d, k, ϵ; D) uses poly(d, k, 1/ϵ) samples from D, and runs in poly(d, k, 1/ϵ) time. Proof of Theorem 5.1. By Theorem 3.1, using a Coarse Direction Recovery procedure (CDR) procedure satisfying Assumption 1 with parameters λ = poly(ϵ/k) and δ = poly(ϵ/k), the algorithm LearnSumReLUs(d, k, ϵ; D) outputs a hypothesis yb(x) with error E(x,y)∼D [(y − yb(x))2 ] ≤ ϵ with probability at least 2/3. In addition, LearnSumReLUs(d, k, ϵ; D) uses e poly(k/ϵ) + L3CDR (NCDR + L3CDR )poly(k/ϵ)) N = O(d samples from D and runs in time poly(d, N, TCDR ), where NCDR , TCDR and LCDR are the sample complexity, runtime and the output list length of the CDR procedure respectively. On the other hand, in Proposition 4.1, we proved that CDRviaImplicitTensor(D, λ, δ) is a CDR procedure satisfying Assumption 1, with sample complexity NCDR = poly(d, k, 1/λ, 1/δ), runtime TCDR = poly(d, k, 1/λ, 1/δ), and list size LCDR ≤ poly(k/λ) log(1/δ). 47

Thus, combining the results from Theorem 3.1 and Proposition 4.1, we obtain that our algorithm LearnSumReLUs(d, k, ϵ; D), using a CDR procedure CDRviaImplicitTensor(D, λ, δ) with parameters λ = poly(ϵ/k) and δ = poly(ϵ/k), with probability at least 2/3 outputs a hypothesis yb(x) with L22 error E(x,y)∼D [(y − yb(x))2 ] ≤ ϵ. In addition, considering the sample complexity, runtime, and the output list length of the CDR procedure CDRviaImplicitTensor(D, λ, δ) under the choice of parameters λ = poly(ϵ/k) and δ = poly(ϵ/k), it follows that LearnSumReLUs(d, k, ϵ; D) uses poly(d, k, 1/ϵ) samples from D, and runs in poly(d, k, 1/ϵ) time. This completes the proof of Theorem 5.1.

References [ATV21]

P. Awasthi, A. Tang, and A. Vijayaraghavan. Efficient algorithms for learning depth-2 neural networks with general ReLU activations. In Advances in Neural Information Processing Systems, volume 34, pages 13485–13496, 2021.

[BF02]

N. Bshouty and V. Feldman. On using extended statistical queries to avoid membership queries. Journal of Machine Learning Research, 2:359–395, 2002.

[BFJ+ 94]

A. Blum, M. Furst, J. Jackson, M. Kearns, Y. Mansour, and S. Rudich. Weakly learning DNF and characterizing statistical query learning using Fourier analysis. In Proceedings of the Twenty-Sixth Annual Symposium on Theory of Computing, pages 253–262, Montreal, Quebec, Canada, 1994. ACM Press.

[BJW19]

A. Bakshi, R. Jayaram, and D. P. Woodruff. Learning two layer rectified neural networks in polynomial time. In Conference on Learning Theory, COLT 2019, 2019.

[Bog98]

V. Bogachev. Gaussian measures. Mathematical surveys and monographs, vol. 62, 1998.

[CDG+ 23] S. Chen, Z. Dou, S. Goel, A. R. Klivans, and R. Meka. Learning narrow one-hidden-layer relu networks. In The Thirty Sixth Annual Conference on Learning Theory, COLT 2023, volume 195 of Proceedings of Machine Learning Research, pages 5580–5614. PMLR, 2023. [CGKM22] S. Chen, A. Gollakota, A. R. Klivans, and R. Meka. Hardness of noise-free learning for two-hidden-layer neural networks. In NeurIPS, 2022. [Che21]

S. Chen. Rethinking Algorithm Design for Modern Challenges in Data Science. PhD thesis, MIT, 2021.

[CKM21]

S. Chen, A. R. Klivans, and R. Meka. Efficiently learning one hidden layer relu networks from queries. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, pages 24087–24098, 2021.

[CKM22]

S. Chen, A. R. Klivans, and R. Meka. Learning deep relu networks is fixed-parameter tractable. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), 2022.

[CN24]

S. Chen and S. Narayanan. A faster and simpler algorithm for learning shallow networks. In The Thirty Seventh Annual Conference on Learning Theory, volume 247 of Proceedings of Machine Learning Research, pages 981–994. PMLR, 2024.

[DIKR25] I. Diakonikolas, G. Iakovidis, D. Kane, and L. Ren. Algorithms and SQ lower bounds for robustly learning real-valued multi-index models. In Advances in Neural Information Processing Systems 38, NeurIPS 2025, 2025. [DIKZ25]

I. Diakonikolas, G. Iakovidis, D. M. Kane, and N. Zarifis. Robust learning of multiindex models via iterative subspace approximation. In 66th IEEE Annual Symposium on Foundations of Computer Science, FOCS 2025, Sydney, Australia, December 14-17, 2025, pages 2231–2239. IEEE, 2025.

[DK20]

I. Diakonikolas and D. M. Kane. Small covers for near-zero sets of polynomials and learning latent variable models. In 2020 IEEE 61st Annual Symposium on Foundations of Computer Science (FOCS), pages 184–195. IEEE, 2020. 48

[DK24]

I. Diakonikolas and D. M. Kane. Efficiently learning one-hidden-layer relu networks via Schur polynomials. In The Thirty Seventh Annual Conference on Learning Theory, volume 247 of Proceedings of Machine Learning Research, pages 1364–1378. PMLR, 2024.

[DK25]

I. Diakonikolas and D. M. Kane. Implicit high-order moment tensor estimation and learning latent variable models. In 2025 IEEE 66th Annual Symposium on Foundations of Computer Science (FOCS), pages 1228–1247. IEEE, 2025.

[DKKZ20] I. Diakonikolas, D. M. Kane, V. Kontonis, and N. Zarifis. Algorithms and SQ lower bounds for pac learning one-hidden-layer ReLU networks. In Conference on Learning Theory, COLT, pages 1514–1539. PMLR, 2020. [DKS18]

I. Diakonikolas, D. M. Kane, and A. Stewart. Learning geometric concepts with nasty noise. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing, STOC 2018, pages 1061–1073, 2018.

[DLB25]

A. Damian, J. D. Lee, and J. Bruna. The generative leap: Tight sample complexity for efficiently learning Gaussian multi-index models. In Advances in Neural Information Processing Systems 38, NeurIPS 2025, 2025.

[Fel08]

V. Feldman. Evolvability from learning algorithms. In Proceedings of the 40th Annual ACM Symposium on Theory of Computing (STOC), pages 619–628. ACM, 2008.

[GGJ+ 20] S. Goel, A. Gollakota, Z. Jin, S. Karmalkar, and A. R. Klivans. Superpolynomial lower bounds for learning one-layer neural networks using gradient descent. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, volume 119 of Proceedings of Machine Learning Research, pages 3587–3596, 2020. [GGK20]

S. Goel, A. Gollakota, and A. R. Klivans. Statistical-query lower bounds via functional gradients. In Advances in Neural Information Processing Systems, NeurIPS, 2020.

[GK19]

S. Goel and A. R. Klivans. Learning neural networks with two nonlinear layers in polynomial time. In Conference on Learning Theory, COLT 2019, pages 1470–1499, 2019.

[GKLW19] R. Ge, R. Kuditipudi, Z. Li, and X. Wang. Learning two-layer neural networks with symmetric inputs. In 7th International Conference on Learning Representations, ICLR 2019, 2019. [GLM18]

R. Ge, J. D. Lee, and T. Ma. Learning one-hidden-layer neural networks with landscape design. In 6th International Conference on Learning Representations, ICLR 2018, 2018.

[JSA15]

M. Janzamin, H. Sedghi, and A. Anandkumar. Beating the perils of non-convexity: Guaranteed training of neural networks using tensor methods, 2015.

[Kea93]

M. J. Kearns. Efficient noise-tolerant learning from statistical queries. In Proceedings of the Twenty-Fifth Annual ACM Symposium on Theory of Computing, pages 392–401, 1993.

[LL22]

A. Liu and J. Li. Clustering mixtures with almost optimal separation in polynomial time. In Proceedings of the 54th Annual ACM SIGACT Symposium on Theory of Computing, pages 1248–1261. ACM, 2022.

[LM00]

B. Laurent and P. Massart. Adaptive estimation of a quadratic functional by model selection. The Annals of Statistics, 28(5):1302–1338, 2000.

[LZZ24]

S. Li, I. Zadik, and M. Zampetakis. On the hardness of learning one hidden layer neural networks. CoRR, abs/2410.03477, 2024. Available at https://arxiv.org/abs/2410.03477.

[NY83]

A. S. Nemirovsky and D. B. Yudin. Problem Complexity and Method Efficiency in Optimization. Wiley-Interscience Series in Discrete Mathematics. Wiley-Interscience, Chichester, NY, 1983. Translated from Russian by E. R. Dawson.

[O’D14]

R. O’Donnell. Analysis of Boolean Functions. Cambridge University Press, 2014.

49

[SJA16]

H. Sedghi, M. Janzamin, and A. Anandkumar. Provable tensor methods for learning mixtures of generalized linear models. In Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, AISTATS 2016, pages 1223–1231, 2016.

[VW19]

S. Vempala and J. Wilmes. Gradient descent for one-hidden-layer neural networks: Polynomial convergence and SQ lower bounds. In Conference on Learning Theory, COLT 2019, pages 3115–3117, 2019. Full version available at https://arxiv.org/abs/1805.02677.

[WZD26]

P. Wang, N. Zarifis, and J. Diakonikolas. Efficient optimization on the Euclidean sphere: Riemannian gradient alignment as a unifying principle. arXiv preprint arXiv:2609.07723, 2026.

[WZDD23] P. Wang, N. Zarifis, I. Diakonikolas, and J. Diakonikolas. Robustly learning a single neuron via sharpness. 40th International Conference on Machine Learning, 2023. [WZDD25] P. Wang, N. Zarifis, I. Diakonikolas, and J. Diakonikolas. Robustly learning monotone single-index models. In Advances in Neural Information Processing Systems 38, NeurIPS 2025, 2025. [Zhe26]

Z. Zheng. Polynomial learning of positive ReLU networks under Gaussian inputs. arXiv preprint arXiv:2610.25116, 2026.

[ZLJ16]

Y. Zhang, J. D. Lee, and M. I. Jordan. L1-regularized neural networks are improperly learnable in polynomial time. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, pages 993–1001, 2016.

[ZSJ+ 17]

K. Zhong, Z. Song, P. Jain, P. L. Bartlett, and I. S. Dhillon. Recovery guarantees for one-hidden-layer neural networks. In Proceedings of the 34th International Conference on Machine Learning, ICML 2017, pages 4140–4149, 2017.

[ZWDD25] N. Zarifis, P. Wang, I. Diakonikolas, and J. Diakonikolas. Robustly learning monotone generalized linear models via data augmentation. In Proceedings of the Thirty Eighth Conference on Learning Theory, volume 291, pages 5921–5990. PMLR, 2025.

50

Related documents

• Learning Positive Linear Combinations of ReLUs in Fully Polynomial Time#1448455
• Deriving Approximate Message Passing from the Convex Gaussian Min-Max Theorem#1342543
• The inductive bias of ReLU networks on orthogonally separable data#449033
• The Gaussian Latent Machine: Efficient Prior and Posterior Sampling for Inverse Problems#144295
• Min-Max Grassmannian Optimization for Online Subspace Tracking#1062880