Mean-Field Analysis for Learning Subspace-Sparse Polynomials with Gaussian Input

Authors: ziang  chen
rong  ge
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 20:08:35

Abstract: In this work, we study the mean-field flow for learning subspace-sparse polynomials using stochastic gradient descent and two-layer neural networks, where the input distribution is standard Gaussian and the output only depends on the projection of the input onto a low-dimensional subspace. We establish a necessary condition for SGD-learnability, involving both the characteristics of the target function and the expressiveness of the activation function. In addition, we prove that the condition is almost sufficient, in the sense that a condition slightly stronger than the necessary condition can guarantee the exponential decay of the loss functional to zero.

Main paper: 
 introduction: 
 Neural Networks (NNs) are powerful in practice to approximate mappings on certain data structures, such as Convectional Neural Networks (CNNs) for image data, Graph Neural Networks (GNNs) for graph data, and Recurrent Neural Networks (RNNs) for sequential data, stimulating numerous breakthroughs in application of machine learning in many branches of science, engineering, etc. The surprising performance of neural networks is often explained by arguing that neural networks automatically learns useful representations of the data. However, how simple training procedures such as stochastic gradient descent (SGD) extract features remains a major open problem.Optimization of neural networks has received lots of attention. For simpler networks such as linear neural networks, local minima are also globally optimal [14, 15, 17] . However, this is not true for nonlinear networks even of depth 2 [23] . Neural Tangent Kernel (NTK, [4, 12, 13] ) is a line of work that establishes strong convergence results for wide neural networks. However, in the NTK regime, neural network is equivalent to a kernel, which cannot learn useful features based on the target function. Such limitation prevents neural networks in NTK regime from efficiently learning even simple single index models [28] .As an alternative, the behavior of SGD can also be understood via mean-field analysis, for both two-layer neural networks [9, 19, 20, 22, 24, 25] and multi-layer neural networks [5, 21, 22] . Neural networks in the mean-field regime have the potential to do feature learning. Recently, [1] showed an interesting setup where a two-layer neural network can learn representations if the target function satisfies a merged-staircase property. More precisely, [1] considers a sparse polynomial as a polynomial f * : R d → R defined on the hypercube {-1, 1} d , i.e., f * (x) = h * (z) = h * (x I ) where z = x I = (x i ) i∈I , I is an unknown subset of {1, 2, . . . , d} with |I| = p, and h * : {-1, 1} p → R is a function on the subset of coordinates in I. They prove that a condition called the merged-staircase property is necessary and in some sense sufficient for learning such f * using SGD and two-layer neural networks. The merged-staircase property proposed in [1] states that all monomials of h * can be ordered such that each monomial contains at most one z i that does not appear in any previous monomial. For example, h * (z) = z 1 + z 1 z 2 + z 1 z 2 z 3 satisfies the merged-staircase property while h * (z) = z 1 + z 1 z 2 z 3 does not. Results on similar structures can also be found in [3] . The work [2] proposes the concept of leap complexity and generalizes the results in [1] to a larger family of sparse polynomials.In this work, we consider "subspace-sparse" polynomial that is more general. Concretely, let f * (x) = h * (z) = h * (x V ), where V is a subspace of R d with dim(V ) = p ≪ d, x V is the orthogonal projection of x onto the subspace V , and h * : V → R is an underlying polynomial map. In other words, the sparsity is in the sense that f * (x) only depends on the projection of the input x ∈ R d in a low-dimensional subspace. Throughout this paper, the input data distribution is the standard d-dimensional normal distribution, i.e., x ∼ N (0, I d ), which is rotation-invariant in the sense that Ox ∼ N (0, I d ) for any orthogonal matrix O ∈ R d×d . Similar rotation-invariant/basis-free settings are also considered in some recent studies, including [2, 8, 10, 11] .Our contribution and related works Our first contribution is a basis-free necessary condition for SGD-learnability. More specially, we propose the reflective property of the underlying polynomial h * : V → R with respect to some subspace S ⊂ V , which also involves the expressiveness of the activation function. We prove that as long as the reflective property is satisfied with respect to nontrivial S, the training dynamics cannot learn any information about the behavior of h * on S (see Theorem 3.4) . Therefore the loss functional will be bounded away from 0 during the whole training procedure.One key point is that our reflective property precisely characterizes the necessary expressiveness of the activation function. If the activation function is expressive enough, the reflective property equivalently recovers a necessary condition characterized by isoLeap [2] that is the maximal leap complexity over all orthonormal basis and can be viewed as a basis-free generalization of the mergedstaircase property. This also indicates that our necessary condition is a bit weaker. Other related rotation-invariant conditions in the previous literature include leap exponent/index [8, 10] , subspace conditioning [10] and even-symmetric directions [11] . The analysis in [10, 11] is for training the first layer for finitely many iterations with fixed second-layer, and [8] studies the joint learning dynamics where they assume that for any fixed first layer, the optimal parameters in second layer can be found efficiently and reformulate the loss as a function of the first layer. Differently and more generally, our analysis for the necessary condition does not require specific learning strategies and works for any learning rates satisfying some mild conditions.Our second contribution is a sufficient condition for SGD-learnability that is also basis-free and is slightly stronger than the necessary condition. In particular, we show that if the training dynamics cannot be trapped in any proper subspace of V , then one can choose the initial parameter distribution and the learning rate such that the loss functional decays to zero exponentially fast with dimensionfree rates (see Theorem 4.3) . Our training strategy is inspired by [1] with the difference that we take the average of p independent training trajectories, which can lift some linear independence property required for polynomials on hypercube to algebraic independence in the general polynomial setting.Technical challenges It may seem simple to leave the standard basis and generalize the results of [1, 2] to learn subspaces, because SGD itself is independent of the basis, and we can consider a symmetric Gaussian input distribution. However, there are some significant barriers that motivated our training process. The condition and the analysis in [1, 2] rely on an orthonormal basis of the input space R d . This is natural for polynomials on the hypercube {-1, 1} d , but not for general polynomials on R d . Particularly, their theory does not work for Gaussian input data x ∼ N (0, I d ), which is probably the most common distribution in data science, unless an orthonormal basis of R d is specified and V is known to be spanned by p elements in the basis. In this work, we consider a more general setting in which specifying a basis is not required and the space V can be any p-dimensional subspace of R d . This setting is consistent with the rotation-invariant property of N (0, I d ) and introduces more difficulties since less knowledge of V is available prior to training.Organization The rest of this paper will be organized as follows. We introduce some preliminaries on mean-field dynamics in Section 2. The basis-free necessary and sufficient conditions for SGDlearnability are discussed in Section 3 and Section 4, respectively. We conclude in Section 5. 
preliminaries on mean-field dynamics: 
 The mean-field dynamics describes the limiting behavior of the training procedure when the stepsize/learning rate converges to zero, i.e., the evolution of a neuron converges to the solution of a differential equation with continuous time, and when the number of neurons converges to infinity, i.e., the empirical distribution of all neurons converges to some limiting probability distribution. For two-layer neural networks, some quantitative results are established in [20] that characterize the distance between the SGD trajectory and the mean-field evolution flow, and these results are further improved as dimension-free in [19] . Such results suggest that analyzing the mean-field flow is sufficient for understanding the SGD trajectory in some settings. In this section, we briefly review the setup of two-layer neural networks, SGD, and their mean-field versions, following [19, 20] .Two-layer neural network and SGD The two-layer neural network is of the following form:EQUATIONwhere N is the number of neurons, Θ = (θ 1 , θ 2 , . . . , θ N ) with θ i = (a i , w i ) ∈ R d+1 is the set of parameters, and σ : R → R is the activation functions with τ (x; θ) := aσ(w ⊤ x) for θ = (a, w).Then the task is to find some parameter Θ such that the ℓ 2 -distance between f * and f NN is minimized:min Θ E N (Θ) := 1 2 E x∼N (0,I d ) |f * (x) -f NN (x; Θ)| 2 . (2.2)In practice, a widely used algorithm for solving (2.2) is the stochastic gradient descent (SGD) that iterates as θEQUATIONwhere x k , k = 1, 2, . . . are the i.i.d. samples drawn from N (0, I d ) and γ (k) = diag(γEQUATIONw I d ) ⪰ 0 is the stepsize or the learning rate. In this paper, we only consider the one-pass model with each data point being used exactly once, following [19] .Mean-field dynamics One can generalize (2.1) to an infinite-width two-layer neural network:f NN (x; ρ) := τ (x; θ)ρ(dθ) = aσ(w ⊤ x)ρ(da, dw),where ρ ∈ P(R d+1 ) is a probability measure on the parameter space R d+1 , and generalize the loss/energy functional (2.2) toE(ρ) := 1 2 E x∼N (0,I d ) |f * (x) -f NN (x; ρ)| 2 .We will use P(X) to denote the collection of probability measures on a space X throughout this paper. The limiting behavior of the SGD trajectory (2.3) when γ (k) → 0 and N → ∞ can be described by the following mean-field dynamics:EQUATIONwhere d+1) with ξ a (t) ≥ 0 and ξ w (t) ≥ 0 being the learning rates andξ(t) = diag(ξ a (t), ξ w (t)I d ) ∈ R (d+1)×(Φ(θ; ρ) = aE x∼N (0,I d ) (f NN (x; ρ) -f * (x)) σ(w ⊤ x) .One can also write Φ(θ; ρ) asΦ(θ; ρ) = V (θ) + U (θ, θ ′ )ρ(dθ ′ ),whereV (θ) = -aE x f * (x)σ(w ⊤ x) and U (θ, θ ′ ) = aa ′ E x σ(w ⊤ x)σ((w ′ ) ⊤ x) . (2.5)The PDE (2.4) is understood in the weak sense, i.e., ρ t is a solution to (2.4) if and only if ρ t t=0 = ρ 0 and(-∂ t η + ∇ θ η • (ξ(t)∇ θ Φ(θ; ρ t ))) ρ t (dθ)dt = 0, ∀ η ∈ C ∞ c (R d+1 × (0, +∞)),where C ∞ c (R d+1 × (0, +∞)) is the collection of all smooth and compactly supported functions on R d+1 × (0, +∞). It can also be computed that the energy functional is non-increasing along ρ t :d dt E(ρ t ) = -∇ θ Φ(θ; ρ t ) ⊤ ξ(t)∇ θ Φ(θ; ρ t )ρ t (dθ) ≤ 0. (2.6)There have been standard results in the existing literature that provide dimension-free bounds for the distance between the empirical distribution of the parameters generalized by (2.3) and the solution to (2.4) . For the simplicity of reading, we will not present those results and the proof; interested readers are referred to [19] . In the rest of this paper, we will focus on the analysis of (2.4) and briefly discuss the sample complexity results implied by our mean-field analysis. 
necessary condition for sgd-learnability: 
 This section introduces a condition that can prevent SGD from recovering all information about f * , or in other words, prevent the loss functional E(ρ t ) decaying to a value sufficiently close to 0.Before rigorously presenting our main theorem, we state the assumptions used in this section.Assumption 3.1. Assume that the followings hold:(i) The activation function σ : R → R is twice continuously differentiable with ∥σ∥ L ∞ (R) ≤ K σ , ∥σ ′ ∥ L ∞ (R) ≤ K σ , and ∥σ ′′ ∥ L ∞ (R) ≤ K σ for some constant K σ > 0.(ii) The learning rates ξ a , ξ w : (iii) The initialization is ρ 0 = ρ a × ρ w such that ρ a is symmetric and is supported inR ≥0 → R satisfy that ∥ξ a ∥ L ∞ (R ≥0 ) ≤ K ξ and ∥ξ w ∥ L ∞ (R ≥0 ) ≤ K ξ for some constant K ξ > 0.[-K ρ , K ρ ]for some constant K ρ > 0.In Assumption 3.1, the Condition (i) is satisfied by some commonly used activation functions, such as σ(x) = 1 1+e -x and σ(x) = cos(x), and is required for establishing the existence and uniqueness of the solution to (2.4) . The Condition (ii) and (iii) are also standard and easy to satisfy in practice.Remark 3.2. The symmetry of ρ a implies that f NN (x; ρ 0 ) = 0. Therefore, the initial loss E(ρ 0 ) =1 2 E x [|f * (x)| 2 ] = 1 2 E x V [|h * (x V )| 2 ] = 1 2 E z [|h * (z)| 2 ], where x V = z ∼ N (0, I V ), can be viewed as a constant depending only on h * and p, independent of d. Noticing also the decay property (2.6), the loss at any time t can be bounded asE(ρ t ) ≤ E(ρ 0 ) = 1 2 E z [|h * (z)| 2 ].The main goal of this section is to generalize the merged-staircase property in a basis-free setting for general polynomials. Without a standard basis, it is hard to talk about having a "staircase" of monomials. Even with a fixed basis, it is still nontrial to define the merged-staircase property for general polynomials since the analysis in [1] highly depends on z 2 i = 1 that is only true for polynomials on the hypercube. Instead, we use the observation that when a function does not satisfy the merged staircase property, it implies that two of the variables will behave the same in the training dynamics. Such a symmetry can be generalized to the basis-free setting for general polynomials and we summarize this as the following reflective property: Definition 3.3 (Reflective property). Let S ⊂ V ⊂ R d be a subspace of V . We say that the underlying polynomial h * : V → R satisfies the reflective property with respect to the subspace S and the activation function σ ifEQUATIONwhere z S = P V S (z) and z ⊥ S = z -P V S (z), with P V S : V → S being the orthogonal projection from V onto S.The reflective property defined above is closely related to the merged-staircase property in [1] . Let us illustrate the intuition using a simple example. Consider V = R 3 and h * (z) = z 1 + z 1 z 2 z 3 . Then h * does not satisfy the merged-staircase property since z 1 z 2 z 3 involves two new coordinates that do not appear in the first monomial z 1 . In our setting, this h * satisfies the reflective with respect to S = span{e 2 , e 3 }, where e i is the vector in R 3 with the i-th entry being 1 and other entries being 0. More specifically, for z = (z 1 , z 2 , z 3 ), one has that z S = (0, z 2 , z 3 ) and z ⊥ S = (z 1 , 0, 0). Thus, one has for any u ∈ R and v ∈ V that σ ′ u + v ⊤ z ⊥ S is independent of z 2 , z 3 and thatE z2,z3 h * (z)σ ′ u + v ⊤ z ⊥ S z S = σ ′ u + v ⊤ z ⊥ S E z2,z3 0, z 1 z 2 + z 1 z 2 2 z 3 , z 1 z 3 + z 1 z 2 z 2 3= (0, 0, 0), which leads to (3.1). One can see from this example that satisfying the reflective property with respect to a nontrivial subspace S ⊂ V is in the same spirit as not satisfying the merged-staircase property. Furthermore, the reflective property is rotation-invariant, meaning that using a different orthonormal basis does not change the property. In this sense, our proposed condition is more general than that in [1] . We also remark that there have been other rotation-invariant conditions generalizing [1] , see e.g., [2, 8, 10, 11] .Another comment is that the reflective property (3.1) depends on the activation function σ, while conditions in previous works [1, 2, 8, 10, 11] are all defined for the target function f * or h * itself.There does exist a variant of our reflective property that is independent of σ ′ , namely, 2) with some nontrivial S is equivalent to isoLeap(h * ) ≥ 2 that means h * : V → R does not satisfy the merged-staircase property for some orthonormal basis of V [2] , and the idea of leaps is used in [8, 10] . We include the proof of equivalence in Appendix B.1.EQUATIONOur main result in this section is that the reflective property with nontrivial S would lead to a positive lower bound of E(ρ t ) along the training dynamics, which provides a necessary condition for the SGD-learnability and is formally stated as follows. Theorem 3.4. Suppose that Assumption 3.1 holds with ρ w ∼ N (0, 1 d I d ), and that h * : V → R satisfies the reflective property with respect to some subspace S ⊂ V and activation function σ. Then for any T > 0, there exists a constant C > 0 depending only on p, h * , K σ , K ξ , K ρ , and T , such thatEQUATION)where h * S ⊥ (z ⊥ S ) = E z S [h * (z)].In particular, if h * (z) is not independent of z S , then for any T > 0, there exists d(T ) > 0 depending only on p, h * , K σ , K ξ , K ρ , and T , such that for any d > d(T ), we haveinf 0≤t≤T E(ρ t ) ≥ 1 4 E z∼N (0,I V ) |h * (z) -h * S ⊥ (z ⊥ S )| 2 > 0. (3.4)It is worth remarking that in Theorem 3.4, the training time T is a constant independent of the dimension d. If a longer d-dependent training beyond a constant time is allowed, then E(ρ t ) might be reasonably small even if the necessary condition is not satisfied, see e.g. [2, 18, 26] .In Appendix B.2, we include a brief discussion of the sample complexity result of SGD implied by Theorem 3.4. In particular, SGD with O(d) samples cannot recover f * reliably if the refelctive property holds, which is consistent with observations in previous works such as [1, 2] . We also remark that our result in Theorem 3.4 is established for the mean-field dynamics corresponding to the one-pass SGD (2.3), any may not apply for other variants of SGD. In particular, some recent works [6, 11, 16] prove that multi-pass SGD with batch-reuse mechanism can learn some target functions with fewer samples than one-pass SGD.To prove Theorem 3.4, the main intuition is that under some mild assumptions, if (3.1) is satisfied and the initial distribution ρ 0 is supported in {(a, w) ∈ R d+1 : w S = 0}, where w S is the orthogonal projection of w ∈ R d onto S, then ρ t is supported in {(a, w) ∈ R d+1 : w S = 0} for all t ≥ 0. This means that the trained neural network f NN (x; ρ t ) learns no information about x S , the orthogonal projection of x ∈ R d onto S, and hence cannot approximate f * (x) = h * (x V ) with arbitrarily smaller error if h * (z) is dependent on z S . We formulate this observation in the following theorem. Theorem 3.5. Suppose that Assumption 3.1 hold and let ρ t be the solution to (2.4) . Let S ⊂ R d be a subspace with the projection map P S : R d+1 → S that maps (a, w) to w S . If (P S ) # ρ 0 = δ S , where δ S is the delta measure on S andEQUATIONwhere x ⊥ S = x -x S , then it holds for any t ≥ 0 thatEQUATIONHere, the delta measure δ S on S is a probability measure on S such that for any continuous and compactly supported function φ : S → R, it holds that S φ(x)δ S (dx) = φ(0). In Theorem 3.5, the condition (3.5) is stated in terms of f * . We will show later that it is closely related to and is actually implied by (3.1), via a decompositionw ⊤ x ⊥ S = w ⊤ (x -x V ) + w ⊤ (x V -x S ), with w ⊤ (x -x V ) and w ⊤ (x V -x S ) corresponding to u and v ⊤ z ⊤ S in (3.1), respectively. The main idea in the proof of Theorem 3.5 is to construct a flow ρt in the space P(R × S ⊥ ), where S ⊥ is the orthogonal complement of S in R d , and then show that ρ t = ρt × δ S is the solution to (2.4). More specifically, the flow ρt is constructed as the solution to the following evolution equation in P(R × S ⊥ ):∂ t ρt = ∇ θ • ρt ξ(t)∇ θ Φ( θ; ρt ) , ρt t=0 = ρ0 , (3.7)where ρ0 ∈ P(R × S ⊥ ) satisfies ρ 0 = ρ0 × δ S , θ = (a, w ⊥ S ), ξ(t) = diag(ξ a (t), ξ w (t)I S ⊥ ), andEQUATIONV ( θ) = -aE x f * (x)σ (w ⊥ S ) ⊤ x ⊥ S , Û (θ, θ ′ ) = aa ′ E x σ (w ⊥ S ) ⊤ x ⊥ S σ ((w ⊥ S ) ′ ) ⊤ x ⊥ S .The detailed proof will be presented in Appendix A.1.In practice, both V and S are unknown and it is nontrivial to choose an initialization ρ 0 supported in {(a, w) ∈ R d+1 : w S = 0}. However, one can set ρ 0 = ρ a × ρ w with ρ w ∼ N (0, 1 d I d ) and this can make the marginal distribution of ρ 0 on S very close to the the delta measure δ S if d >> p = dim(V ) ≥ dim(S), which fits the setting of subspace-sparse polynomials. Rigorously, we have the following theorem stating dimension-free stability with respect to initial distribution, with the proof deferred to Appendix A.2. Theorem 3.6. Suppose that Assumption 3.1 holds for both ρ 0 and ρ0 . Let ρ t solve ∂ t ρ t = ∇ θ • (ρ t ξ(t)∇ θ Φ(θ; ρ t )) and let ρt solve ∂ t ρt = ∇ θ • (ρ t ξ(t)∇ θ Φ(θ; ρt )). Then for any T ∈ (0, +∞), there exists a constant C s > 0 depending only on p, h * , K σ , K ξ , K ρ , and T , such thatEQUATIONwhere W 2 (•, •) is the 2-Wasserstein metric.Based on Theorem 3.5 and Theorem 3.6, Theorem 3.4 can be proved by some straightforward computation, for which the details can be found in Appendix A.3.1: Set the initial distribution asρ 0 = ρ a × δ R d , where ρ a = U([-1, 1]) and δ R d is the delta measure on R d . 2: Set ξ a (t) = 0 and ξ w (t) = 1 for 0 ≤ t ≤ T , and train the neural network with activation function σ. Denote by (a, w(a, t)), 0 ≤ t ≤ T the trajectory of a single particle that starts at (a, 0). 3: Repeat Step 2 for p times independently and obtain p copies of parameters at T , say (a i , w(a i , T ))with a i ∼ U([-1, 1]), i = 1, 2, . . . , p.  
sufficient condition for sgd-learnability: 
 In this section, we propose a sufficient condition and a training strategy that can guarantee the exponential decay of E(ρ t ) with constants independent of the dimension d.We prove in Section 3 that if the trained parameters always stay in a proper subspace {(a, w) ∈ R d+1 : w S = 0}, then f NN (x; ρ t ) cannot learn all information about f * or h * . Ideally, one would expect the negation to be a sufficient condition for the SGD-learnability, i.e., the existence of a choice of learning rates and initial distribution that guarantees lim t→∞ E(ρ t ) with dimension-free rate. This is almost true but we need a slightly stronger condition due to technical issues. More specifically, we need that the Taylor's expansion of some dynamics (not the dynamics itself) is not trapped in any proper subspace. Assumption 4.1. Consider the following flow ŵV (t) in V :d dt ŵV (t) = E z zh * (z)σ ′ ( ŵV (t) ⊤ z) , ŵV (0) = 0. (4.1)We assume that for some s ∈ N + , the Taylor's expansion up to s-th order of ŵV (t) at t = 0 is not contained in any proper subspace of V .Assumption 4.1 aims to state the same observation as the merged-staircase property in [1] . As a simple example, ifV = R p and h * (z) = z 1 + z 1 z 2 + z 1 z 2 z 3 + • • • + z 1 z 2 • • • z pwhich satisfies the merged-staircase property, then it can be computed that the leading order terms of the coordinates of ŵV (t) are given by (cEQUATION(0), . . . , σ (p) (0) are all nonzero (see Proposition 33 in [1] ). This is to say that Assumption 4.1 with s = 2 p-1 is satisfied for this example. We provide further characterization of Assumption 4.1 by verifying it in a more general setting in Appendix D.1.We also remark that the Taylor's expansion of the flow ŵV (t) that solves (4.1) depends only on the h * and σ (1) (0), σ (2) (0), . . . , σ (s) (0). We require some additional regularity assumption on higher-order derivatives of σ.Assumption 4.2. Assume that σ satisfies σ ∈ C L+1 (R) and σ, σ ′ , σ ′′ , σ (L+1) ∈ L ∞ (R), where L = 2sn n+p p with n = deg(f * ) = deg(h *) and s being the positive integer in Assumption 4.1.Our proposed training strategy is stated in Algorithm 1. The training strategy is inspired by the two-stage strategy proposed in [1] that trains the parameters w with fixed a for t ∈ [0, T ] and then trains the parameter a with fixed w and a perturbed activation function for t ≥ T . Several important modifications are made since we consider general polynomials, rather than polynomials on hypercubes as in [1] . In particular,• We need to repeat Step 2 (training w) for p times and use their average as the initialization of training a, while this step only needs to be done once in [1] . The reason is that the space of polynomials on the hypercube {±1} p is essentially a linear space with dimension 2 p . However, the space of general polynomials on V is an R-algebra that is also a linear space but is of infinite dimension. Therefore, to make the kernel matrix in training a non-degenerate, we require some algebraic independence which can be guaranteed by u(a 1 , . . . , a p , t)) = 1 p p i=1 w(a i , t), 0 < t ≤ T , though linear independence suffices for [1] . Let us also emphasize that each run of Step 2 involves training an interacting particle system instead of training a single particle.• In Step 4, we use a new activation function σ(ζ) = (1 + ζ) n that is a polynomial of the same degree as f * and h * . The reason is still that we work with the space general polynomials whose dimension as a linear space is infinite. Thus, we need the specific formσ(ζ) = (1 + ζ) n to guarantee the trained neural network f NN (x; ρ t ) is a polynomial with degree at most n = deg(f * ) = deg(h * ).As a comparison, in the setting of [1] , all functions on {±1} p can be understood as a polynomial, and no specific format of the new activation function is needed.Our main theorem in this section is as follows, stating that the loss functional E(ρ t ) can decay to 0 exponentially fast, with rates independent of the dimension d. Theorem 4.3. Suppose that Assumption 4.1 and 4.2 hold and let ρ t be the flow generated by Algorithm 1. There exist constants C 1 , C 2 > 0 depending on h * , σ, n, p, s, such thatE(ρ t ) ≤ C 1 exp(-C 2 t), ∀ t ≥ 0.Let us also remark that it is possible to use the original dynamics ŵV defined in (4.1) when we state Assumption 4.1, which can actually imply its Taylor's expansion up to some order is not trapped in any proper subspace of V if we further assume ŵV (t) is analytic. We choose to directly use Taylor's expansion in Assumption 4.1 since we want to avoid the additional analytic assumption and to emphasize that the constants C 1 , C 2 in Theorem 4.3 depend on the order s of the Tayler's expansion satisfying Assumption 4.1.Discussion about the sample complexity implied by Theorem 4.3 is included in Appendix D.2, suggesting that O(d) samples suffices for SGD to learn f * reliably if conditions in Theorem 4.3 are true. This is also consistent with previous works such as [1, 2] .To prove Theorem 4.3 we follow the same general strategy as [1] , though some technical analysis is significantly different due to the roatation-invariant setting. The main goal here is to show before Step 4, the algorithm already learned a diverse set of features. After that, note that Step 4 in Algorithm 1 is essentially a convex/quadratic optimization problem (since we only train a and set ξ w (t) = 0). In addition, thanks to the new activation function σ(ζ) = (1 + ζ) n , one only needs to consider P V,n that is the space of of all polynomials on V with degree at most n = deg(h * ) = deg(f * ). The dimension of P V,n as a linear space is n+p p . Let p 1 , p 2 , . . . , p ( n+p p ) be the orthonormal basis of P V,n with input z ∼ N (0, I V ) and define the kernel matrixK i1,i2 (t) = E a1,...,ap E z,z ′ p i1 (z)σ(u(a 1 , . . . , a p , t) ⊤ z)σ(u(a 1 , . . . , a p , t) ⊤ z ′ )p i2 (z ′ ) , (4.2) where σ(ξ) = (1 + ξ) n , (a 1 , . . . , a p ) ∼ U([-1, 1] p ), 1 ≤ i 1 , i 2 ≤ n+pp , and 0 ≤ t ≤ T . As long as this kernel matrix is non-degenerate, we know that the loss functional is strongly convex with respect to the parameters in the second layer when fixing the first layer, and thus, it can be computed straightforwardly that the loss decays to 0 exponentially fast for t ≥ T , leading to the desired convergence rate in Theorem 4.3; see Appendix C.3 for details. Thus, the main part in the proof of Theorem 4.3 is to establish the non-degeneracy of the kernel matrix. Proposition 4.4. Suppose that Assumption 4.1 and 4.2 hold. There exist constants C, T > 0 depending on h * , σ, n, p, s, such thatEQUATIONwhere λ min (K(t)) is the smallest eigenvalue of K(t).In the rest of this subsection, we sketch the main ideas in the proof of Proposition 4.4, with the details of the proof being deferred to Appendix C. We first show that w(a i , t) and u(a 1 , . . . , a p , t) can be approximated well by ŵ(a i , t) and û(a 1 , . . . , a p , t) = 1 p p i=1 ŵ(a i , t) that are polynomials in a i and a 1 , . . . , a p respectively. This approximation step follows [1] closely and is analyzed detailedly in Appendix C.1. Therefore, to give a positive lower bound of λ min (K(t)), one only needs to show the non-degeneracy of the matrix M (a, t) ∈ R ( n+p p )×( n+p p ) withMi1,i2 (a, t) = E z p i1 (z)σ(û(a i2 , t) ⊤ z) ,where a = a 1 , a 2 , . . . , a ( n+p p ) and a i ∈ R p for i = 1, 2, . . . , n+p p . Intuitively, this nondegeneracy can be implied by span σ(û(a 1 , . . . , a p , t) ⊤ z) : a 1 , a 2 , . . . , a p ∈ [-1, 1] = P V,n , which is true if ûi (a 1 , . . . , a p , t), 1 ≤ i ≤ p are R-algebraically independent polynomials in a 1 , a 2 , . . . , a p , where ûi is the i-th coefficient of û under some basis of V , and algebraic independence can be obtained from linear independence by taking the average of independent copies. We illustrate this intuition with a bit more detail.Algebraic independence of ûi . With Assumption 4.1, ŵ1 (a, t), ŵ2 (a, t), . . . , ŵ1 (a, t) can be proved as R-linear independent polynomials in a ∈ R. Then one can apply the following theorem to boost the linear independence of ŵi , whose constant term is zero since initialization in training is set asρ 0 = ρ a × δ R d , to the algebraic independence of ûi . Theorem 4.5. Let v 1 , v 2 , . . . , v p ∈ R[a] be R-linearly independent polynomials with the constant terms being zero. Then 1 p (v 1 (a 1 ) + • • • + v 1 (a p )), . . . , 1 p (v p (a 1 ) + • • • + v p (a p )) ∈ R[a 1 , a 2 , . . . , a p ] are R-algebraically independent.The proof of Theorem 4.5 is deferred to Appendix C.2 and is based on the celebrated Jacobian criterion stated as follows. Theorem 4.6 (Jacobian criterion [7] ). . . .v 1 , v 2 , . . . , v p ∈ R[a 1 ,∂vp ∂ap       is a nonzero polynomial in R[a 1 , a 2 , . . . , a p ].Non-degeneracy of M (a, t). With the observation that span σ(q ⊤ z) : q ∈ V = P V,n (see Lemma C.13), we define another matrixX(q) ∈ R ( n+p p )×( n+p p ) via X i1,i2 (q) = E z p i1 (z)σ(q ⊤ i2 z) ,where q = q 1 , q 2 , . . . , q ( n+p p ) with q i ∈ V , and prove that det(X(q)) is a non-zero polynomial in q of the formdet(X(q)) = ( n+p p ) i=1 0≤∥ji∥1≤n X j q j = ( n+p p ) i=1 0≤∥ji∥1≤n X j ( n+p p ) l=1 q j l l ,where q i is understood as a (coefficient) vector in R p associated with a fixed orthonormal basis of V and j = j 1 , j 2 , . . . , j ( n+p p ) with j i ∈ N p . Then setting q = û(a i2 , t) leads todet( M (a, t)) = ( n+p p ) i=1 0≤∥ji∥1≤n X j ( n+p p ) l=1 û(a l , t) j l .To prove that det( M (a, t)) is a non-zero polynomial in a, i.e., M (a, t) is non-degenerate, we use the following lemma linking algebraic independence back to linear independence.Lemma 4.7. Suppose that v 1 , v 2 , . . . , v p ∈ R[a 1 , a 2 , . . . , a p ] are R-algebraically independent. For any m ≥ 1, the following polynomials in a = (a 1 , a 2 , . . . ,a m ) ∈ (R p ) m are R-linearly independent m l=1 v(a l ) j l , 1 ≤ ∥j i ∥ 1 ≤ n, 1 ≤ i ≤ m, where v = (v 1 , v 2 , . . . , v p ).The proof of Lemma 4.7 and some other related analysis are deferred to Appendix C.3. 
conclusion and discussions: 
 In this work, we generalize the merged-staircase property in [1] to a basis-free version and establish a necessary condition for learning a subspace-sparse polynomial on Gaussian input with arbitrarily small error. Moreover, we prove the exponential decay property of the loss functional under a sufficient condition that is slightly stronger than the necessary one. The bounds and rates are all dimension-free. Our work provides some understanding of the mean-field dynamics, though its general behavior is extremely difficult to characterize due to the non-convexity of the loss functional.Let us also make some comments on limitations and future directions. Firstly, there is still a gap between the necessary condition and the sufficient condition, which is basically from the fact that the sufficient condition is built on the Taylor's expansion of the flow (4.1). One future research question is whether we can fill the gap by considering the original flow (4.1) rather than its Taylor's expansion. Secondly, Algorithm 1 repeats training w for p times and takes the average of parameters, which is different from the usual strategy for training neural networks. This step is used to guarantee the algebraic independence. We conjecture that this step can be removed since the general algebraic independence is too strong when we have some preknowledge on the degree of f * or h * , which deserves future research. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [NA] Justification: This paper is purely theoretical and does not include experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [NA] Justification: This paper is purely theoretical and does not include experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [NA] Justification: This paper is purely theoretical and does not include experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]Justification: The authors have reviewed the NeurIPS Code of Ethics and confirm that this paper conforms with it in every respect.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper conducts fundamental research and is purely theoretical. Thus, it has no societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper is purely theoretical and poses no such risks.Answer: [NA] Justification: This paper is purely theoretical/mathematical and does not involve crowdsourcing or research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This section collects the proofs of Theorem 3.5, Theorem 3.6, and Theorem 3.4.A.1 Proof of Theorem 3.5Existence and uniqueness of solutions to (2.4) Before proving Theorem 3.5, let us remark on the existence and uniqueness of solution to the mean-field dynamics (2.4). According to Remark 7.1 in [20] and Theorem 1.1 in [27] , the PDE (2.4) admits a unique solution if Assumption 3.1 (ii) holds and both ∇V (θ) and ∇ θ U (θ, θ ′ ) are bounded and Lipschitz continous. Here we recall that V and U are defined in (2.5) . With Assumption 3.1 (i) and (iii), it is not hard to verify the boundedness and Lipschitz continuity of ∇V (θ) and ∇ θ U (θ, θ ′ ) by noticing that any finite-order moment of N (0, I d ) is finite.Proof of Theorem 3.5. Since (P S ) # ρ 0 = δ S , the initial distribution ρ 0 can be decomposed as ρ 0 = ρ0 × δ S , where ρ0 ∈ P(R × S ⊥ ). Consider the following evolution equation (3.7) in P(R × S ⊥ ).By the discussion at the beginning of Section A.1, we know that ρ t is the unique solution to (2.4) . Similar arguments also leads to the existence and uniquess of the solution to (3.7).We will show that the solution to (2.4) must be of the form ρ t = ρt × δ S , (A.1) where ρt solves (3.7), and this decomposition can immediatel imply (3.6) . By the uniquess of the solution, it suffices to verify that ρ t = ρt × δ S is a solution to (2.4). It follows directly from (A.1) thatf NN (x; ρ t ) = f NN (x; ρt × δ S ) = fNN (x ⊥ S ; ρt ), and hence that ∂ a Φ(θ; ρ t ) = ∂ a Φ( θ; ρt ), if w S = 0. (A.2) We also have that ∇ w ⊥ S Φ(θ; ρ t ) = aE x (f NN (x; ρ t ) -f * (x)) σ ′ (w ⊤ x)x ⊥ S = aE x fNN (x ⊥ S ; ρt ) -f * (x) σ ′ (w ⊥ S ) ⊤ x ⊥ S x ⊥ S = ∇ w ⊥ S Φ( θ; ρt ), (A.3)if w S = 0. In addition, it holds also for w S = 0 that ∇ w S Φ(θ;ρ t ) = aE x (f NN (x; ρ t ) -f * (x)) σ ′ (w ⊤ x)x S = aE x fNN (x ⊥ S ; ρt )σ ′ (w ⊥ S ) ⊤ x ⊥ S x S -aE x f * (x)σ ′ (w ⊥ S ) ⊤ x ⊥ S x S = 0, (A.4)where we used E x S [x S ] = 0 and (3.5). Combining (A.2), (A.3), and (A.4), we have for anyη ∈ C ∞ c (R d+1 × (0, +∞)) = C ∞ c (R × S ⊥ × S × (0, +∞)) that (-∂ t η + ∇ θ η • (ξ(t)∇ θ Φ(θ; ρ t ))) ρ t (dθ)dt = -∂ t η(θ, t) + ξ a (t)∂ a η(θ, t) • ∂ a Φ(θ; ρ t ) + ξ w (t)∇ w ⊥ S η(θ, t) • ∇ w ⊥ S Φ(θ; ρ t ) + ξ w (t)∇ w S η(θ, t) • ∇ w S Φ(θ; ρ t ) ρt (d θ)δ S (dw S )dt = -∂ t η( θ, 0, t) + ξ a (t)∂ a η( θ, 0, t) • ∂ a Φ( θ; ρt ) + ξ w (t)∇ w ⊥ S η( θ, 0, t) • ∇ w ⊥ S Φ( θ; ρt ) ρt (d θ)dt = -∂ t η( θ, 0, t) + ∇ θ η( θ, 0, t) • ( ξ(t)∇ θ Φ( θ; ρt )) ρt (d θ)dt =0,where the last equality holds by applying the test function η(•, 0, •) ∈ C ∞ c (R × S ⊥ × (0, +∞)) to (3.7). The proof is hence completed.The proof of Theorem 3.6 uses some ideas from the proof of Theorem 16 in [1] . Similar ideas also exist in earlier works (see e.g., [19, 20] ). Lemma A.1. Suppose that Assumption 3.1 holds and let ρ t solve (2.4). Then for any t > 0,ρ t is supported in θ = (a, w) ∈ R d+1 : |a| ≤ K ρ + K ξ K σ E z |h * (z)| 2 1/2 t .Proof. The particle dynamics for a t associated with (2.4) is given byd dt a t = ξ a (t)E x (f NN (x; ρ t ) -f * (x))σ(w ⊤ t x) ,which implies thatd dt a t ≤ K ξ K σ |E x [f NN (x; ρ t ) -f * (x)]| ≤ K ξ K σ (2E(ρ 0 )) 1/2 = K ξ K σ E z |h * (z)| 2 1/2 .Therefore, one has|a t | ≤ |a 0 | + K ξ K σ E z |h * (z)| 2 1/2t, which completes the proof.Lemma A.2. Suppose that Assumption 3.1 holds for both ρ 0 and ρ0 . Letρ t solve ∂ t ρ t = ∇ θ • (ρ t ξ(t)∇ θ Φ(θ; ρ t )) and let ρt solve ∂ t ρt = ∇ θ • (ρ t ξ(t)∇ θ Φ(θ; ρt )).For any coupling γ 0 ∈ Γ(ρ 0 , ρ0 ), let γ t ∈ Γ(ρ t , ρt ) be the associated coupling during the evolution and define∆(t) = |a -ã| 2 + ∥w -w∥ 2 γ t (dθ, d θ).Then it holds for any 0 ≤ t ≤ T thatEQUATIONwhere C f and C ∆ are constants depending only on p, h * , K σ , K ξ , K ρ , and T .Proof of Theorem 3.6. Let C f and C ∆ be the constants in Lemma A.2. For any ϵ > 0, there exists a coupling γ 0 ∈ Γ(ρ 0 , ρ0 ) such that(|a -ã| 2 + ∥w -w∥ 2 )γ 0 (dθ, d θ) ≤ W 2 2 (ρ 0 , ρ0 ) + ϵ.Define γ t ∈ Γ(ρ t , ρt ) and ∆(t) as in Lemma A.2. According (A.6) and the Grönwall's inequality, it holds that ∆(t) ≤ ∆(0)e C∆t ≤ W 2 2 (ρ 0 , ρ0 ) + ϵ e C∆t , ∀ 0 ≤ t ≤ T, which combined with (A.5) yields thatsup 0≤t≤T E x |f NN (x; ρ t ) -f NN (x; ρt )| 2 ≤ W 2 2 (ρ 0 , ρ0 ) + ϵ C f e C∆T .Then we can conclude (3.8) be setting ϵ → 0 and C s = C f e C∆T .Corollary A.3. Under the same setting as in Theorem 3.6, one hassup 0≤t≤T |E(ρ t ) -E(ρ t )| ≤ C s E z |h * (z)| 2 1/2 W 2 (ρ 0 , ρ0 ) + 1 2 C s W 2 2 (ρ 0 , ρ0 ). (A.7)Proof. It can be computed thatE(ρ t ) = 1 2 E x |(f NN (x; ρt ) -f * (x)) + (f NN (x; ρ t ) -f NN (x; ρt ))| 2 =E(ρ t ) + E x [(f NN (x; ρt ) -f * (x))(f NN (x; ρ t ) -f NN (x; ρt ))] + 1 2 E x |f NN (x; ρ t ) -f NN (x; ρt )| 2 ,which implies thatsup 0≤t≤T |E(ρ t ) -E(ρ t )| ≤E x |f NN (x; ρt ) -f * (x)| 2 1/2 E x |f NN (x; ρ t ) -f NN (x; ρt )| 2 1/2 + 1 2 E x |f NN (x; ρ t ) -f NN (x; ρt )| 2 ≤ C s E x V |h * (x V )| 2 1/2 W 2 (ρ 0 , ρ0 ) + 1 2 C s W 2 2 (ρ 0 , ρ0 ),where we used Theorem 3.6 and Remark 3.2.Proof of Lemma A.2. We first prove (A.5). It can be computed that|f NN (x; ρ t ) -f NN (x; ρt )| = aσ(w ⊤ x)ρ t (dθ) -ãσ( w⊤ x)ρ t (d θ) ≤ (a -ã)σ(w ⊤ x)γ t (dθ, d θ) + ã σ(w ⊤ x) -σ( w⊤ x) γ t (dθ, d θ) ,and hence thatE x |f NN (x; ρ t ) -f NN (x; ρt )| 2 ≤2E x (a -ã)σ(w ⊤ x)γ t (dθ, d θ) 2 + 2E x ã σ(w ⊤ x) -σ( w⊤ x) γ t (dθ, d θ) 2 .We then bound the two terms above as follows:EQUATIONE x ã σ(w ⊤ x) -σ( w⊤ x) γ t (dθ, d θ) 2 ≤K 2 σ K ρ + √ 2K ξ K σ E(ρ 0 ) 1/2 T 2 E x |(w -w) ⊤ x|γ t (dθ, d θ) 2 ≤K 2 σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T 2 E x |(w -w) ⊤ x| 2 γ t (dθ, d θ) =K 2 σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T 2 ∥w -w∥ 2 γ t (dθ, d θ) ≤K 2 σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T 2 ∆(t),where we used Lemma A.1 and(w -w) ⊤ x ∼ N (0, ∥w -w∥ 2 ) if x ∼ N (0, I d ). Then we can conclude (A.5) with C f = K 2 σ + K 2 σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T 2by combining all estimations above.We then head into the proof of (A.6), for which we need the particle dynamicsEQUATIONd dt w t = ξ w a t E x (f NN (x; ρ t ) -f * (x))σ ′ (w ⊤ t x)x , d dt wt = ξ w ãt E x (f NN (x; ρt ) -f * (x))σ ′ ( w⊤ t x)x . The distance between the particle dynamics can be decomposed asd dt a t - d dt ãt ≤K ξ E x (f NN (x; ρ t ) -f NN (x; ρt ))σ(w ⊤ t x) + K ξ E x (f NN (x; ρt ) -f * (x))(σ(w ⊤ t x) -σ( w⊤ t x)) , (A.8)andw t -wt , d dt w t - d dt wt ≤K ξ (a t -ãt )E x (f NN (x; ρ t ) -f * (x))σ ′ (w ⊤ t x)(w t -wt ) ⊤ x + K ξ ãt E x (f NN (x; ρ t ) -f NN (x; ρt ))σ ′ (w ⊤ t x)(w t -wt ) ⊤ x + K ξ ãt E x (f NN (x; ρt ) -f * (x)) σ ′ (w ⊤ t x) -σ ′ ( w⊤ t x) (w t -wt ) ⊤ x .(A.9)Therefore, it can be computed thatd dt |a -ã| 2 γ t (dθ, d θ) = d dt |a t -ãt | 2 γ 0 (dθ 0 , d θ0 ) = 2 (a t -ãt ) d dt a t - d dt ãt γ 0 (dθ 0 , d θ0 ) ≤ |a t -ãt | 2 γ 0 (dθ 0 , d θ0 ) + d dt a t - d dt ãt 2 γ 0 (dθ 0 , d θ0 ) (A.8) ≤ ∆(t) + 2K 2 ξ K 2 σ E x [|f NN (x; ρ t ) -f NN (x; ρt )|] 2 γ 0 (dθ 0 , d θ0 ) + 2K 2 ξ K 2 σ E x |f NN (x; ρt ) -f * (x)| • (w t -wt ) ⊤ x 2 γ 0 (dθ 0 , d θ0 ) (A.5) ≤ ∆(t) + 2K 2 ξ K 2 σ C f ∆(t) + 2K 2 ξ K 2 σ E x |f NN (x; ρt ) -f * (x)| • (w t -wt ) ⊤ x 2 γ 0 (dθ 0 , d θ0 ) ≤ ∆(t) + 2K 2 ξ K 2 σ C f ∆(t) + 2K 2 ξ K 2 σ E x |f NN (x; ρt ) -f * (x)| 2 E x (w t -wt ) ⊤ x 2 γ 0 (dθ 0 , d θ0 ) ≤ ∆(t) + 2K 2 ξ K 2 σ C f ∆(t) + 2K 2 ξ K 2 σ E z [|h * (z)| 2 ] ∥w t -wt ∥ 2 γ 0 (dθ 0 , d θ0 ) ≤ 1 + 2K 2 ξ K 2 σ C f + 2K 2 ξ K 2 σ E z [|h * (z)| 2 ] ∆(t), where we used Remark 3.2, and thatd dt ∥w -w∥ 2 γ t (dθ, d θ) = d dt ∥w t -wt ∥ 2 γ 0 (dθ 0 , d θ0 ) = 2 w t -wt , d dt w t - d dt wt γ 0 (dθ 0 , d θ0 ) (A.9) ≤ 2K ξ K σ |a t -ãt | • E x |f NN (x; ρ t ) -f * (x)| • (w t -wt ) ⊤ x γ 0 (dθ 0 , d θ0 ) + 2K ξ K σ |ã t | • E x |f NN (x; ρ t ) -f NN (x; ρt ))| • (w t -wt ) ⊤ x γ 0 (dθ 0 , d θ0 ) + 2K ξ K σ |ã t | • E x |f NN (x; ρt ) -f * (x)| • (w t -wt ) ⊤ x 2 γ 0 (dθ 0 , d θ0 ) ≤ 2K ξ K σ |a t -ãt | • E z |h * (z)| 2 1/2 E x (w t -wt ) ⊤ x 2 1/2 γ 0 (dθ 0 , d θ0 ) + 2K ξ K σ |ã t | • (C f ∆(t)) 1/2 E x (w t -wt ) ⊤ x 2 1/2 γ 0 (dθ 0 , d θ0 ) + 2K ξ K σ |ã t | • E z |h * (z)| 2 1/2 E x (w t -wt ) ⊤ x ≤ 2K ξ K σ E z |h * (z)| 2 1/2 |a t -ãt | • ∥w t -wt ∥γ 0 (dθ 0 , d θ0 ) + 2K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T (C f ∆(t))1/2 ∥w t -wt ∥γ 0 (dθ 0 , d θ0 )+ 2K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T E z |h * (z)| 2 1/2 • √ 3∥w t -wt ∥ 2 γ 0 (dθ 0 , d θ0 ) ≤ K ξ K σ E z |h * (z)| 2 1/2 ∆(t) + 2K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T C 1/2 f ∆(t) + 2 √ 3K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T E z |h * (z)| 2 1/2 ∆(t),where we used Remark 3.2, (A.5), Lemma A.1, and(w -w) ⊤ x ∼ N (0, ∥w -w∥ 2 ) if x ∼ N (0, I d ).Therefore, we can conclude thatd dt ∆(t) = d dt |a -ã| 2 γ t (dθ, d θ) + d dt ∥w -w∥ 2 γ t (dθ, d θ) ≤ C ∆ ∆(t),whereC ∆ =1 + 2K 2 ξ K 2 σ C f + 2K 2 ξ K 2 σ E x V [|h * (x V )| 2 ] + K ξ K σ E z |h * (z)| 2 1/2 + 2K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T C 1/2 f + 2 √ 3K ξ K σ K ρ + K ξ K σ E z |h * (z)| 2 1/2 T E z |h * (z)| 2 1/2 .This proves (A.6).The proof of Theorem 3.4 is based on Theorem 3.5 and Theorem 3.6.Proof of Theorem 3.4. Let P S : R d+1 → S be the projection in Theorem 3.5 and let P ⊥ S = I d+1 -P S . Let ρt solve (2.4) with ρ0 = P ⊥ S ρ w × δ S . It is clear that (P S ) # ρ0 = δ S . In addition, with the decomposition x = x 1 + x 2 + x 3 where x 1 = x ⊥ V = x -x V , x 2 = x V -x S , and x 3 = x S are independent Gaussian random variables, we have for any w ∈ R d thatE x f * (x)σ ′ w ⊤ x ⊥ S x S = E x1 E x2,x3 h * (x 2 + x 3 )σ ′ w ⊤ x 1 + w ⊤ x 2 x 3 = 0, where we used (3.1). Then according to Theorem 3.5, for any t ≥ 0 that (P S ) # ρt = δ S , which implies that f NN (x; ρ t ) is a constant function in x S for any fixed x ⊥ S , giving a lowerbound on its loss:E(ρ t ) = 1 2 E x ∥f * (x) -f NN (x; ρt )∥ 2 = 1 2 E x ⊥ S E x S ∥f * (x) -f NN (x; ρt )∥ 2 ≥ 1 2 E x ⊥ S E x S ∥f * (x) -E x S [f * (x)]∥ 2 = 1 2 E z ⊥ S E z S ∥h * (z) -h * S ⊥ (z ⊥ S )∥ 2 = 1 2 E z ∥h * (z) -h * S ⊥ (z ⊥ S )∥ 2 ,where z S = P V S z, z ⊥ S = z -z S , andh * S ⊥ (z ⊥ S ) = E z S [h * (z)]. Now we show the actual flow is not very different. For ρ 0 = ρ a × ρ w with ρ w ∼ N (0, I d ) and ρ0 = P ⊥ S ρ w × δ S , it can be estimated thatW 2 2 (ρ 0 , ρ0 ) ≤ dim S d ≤ p d .Then applying Corollary A.3, we can conclude for any T > 0 that Proof. Without loss of generality, we assume that V = R p . Suppose that isoLeap(h * ) ≥ 2, which means that the leap complexity of h * (as defined in [2, Definition 1] ) is greater than one for some orthonormal basis of V . We can assume that the basis is {e 1 , e 2 , . . . , e p }, where e j is the vector in R p with the j-th entry being 1 and other entries being 0. Denote the Hermite decomposition of h * asinf 0≤t≤T E(ρ t ) ≥ inf 0≤t≤T E(ρ t ) -C s E z |h * (z)| 2 1/2 W 2 (ρ 0 , ρ0 ) - 1 2 C s W 2 2 (ρ 0 , ρ0 ) ≥ 1 2 E z ∥h * (z) -h * S ⊥ (z ⊥ S )∥ 2 - pC s E z |h * (z)| 2 1/h * (z) = m i=1 c i p j=1He αi(j) (z j ), (where c 1 , c 2 , . . . , c m are nonzero coefficients, α 1 , α 2 , . . . , α m are pairwise distinct elements in N p with α i (j) being the j-th entry of α i , and He k is the k-th order Hermite polynomial. Since the leap complexity of h * is at least two, the following is true after applying some permutation on {1, 2, . . . , m}: There exists some m 1 ∈ {1, 2, . . . , m -1}, such that it holds for any i ∈ {m 1 + 1, . . . , m} thatj∈Jm 1 α i (j) ≥ 2, (B.2)where J m1 = {j ∈ {1, 2, . . . , p} : α i (j) = 0, ∀ i ∈ {1, 2, . . . , m 1 }}.Thus, by the orthogonality of Hermite polynomials, we have for i ∈ {m 1 + 1, . . . , m} thatE zj ∼N (0,1)   z j p j=1He αi(j) (z j )  = 0, ∀ j ∈ J m1 .The above also holds for i ∈ {1, 2, . . . , m 1 } by the definition of J m1 . Therefore, we can conclude that E z S ∼N (0,I S ) [h * (z)z S ] = 0, ∀ z ⊥ S , i.e., (3.2) holds, for S = span{e j : j ∈ J m1 }. Moreover, S is nontrivial since (B.2) implies J m1 is not empty.On the other hand, suppose that (3.2) is satisfied with some nontrivial subspace S ⊂ V that can be assumed as S = {(z 1 , . . . , z p1 , 0, . . . , 0) : z 1 , . . . , z p1 ∈ R} with 1 ≤ p 1 ≤ p. We still consider the Hermite decomposition as in (B.1) and rewrite it as h * (z) = m2 i=1 c ′ i p1 j=1 He α ′ i (j) (z j )h i (c ′ i h i (z p1+1 , . . . , z p )E zj ∼N (0,1) z j He α ′ i (j) (z j ) p1 j ′ =1,j ′ ̸ =j E z j ′ ∼N (0,1) He α ′ i (j ′ ) (z j ′ ) = 0,for any j ∈ {1, 2, . . . , p 1 } and z p1+1 , . . . , z p ∈ R, which implies thatp1 j=1 α i (j) ′ ̸ = 1, ∀ i ∈ {1, 2, . . . , m 2 }.Therefore, the leap complexity of h * is at least 2 with respect to this basis, which leads to isoLeap(h * ) ≥ 2.Implied by Theorem 3.4We discuss the sample complexity result of SGD implied by Theorem 3.4 in this subsection. Recall that there have been standard dimension-free results for bounding the distance between SGD and the mean-field dynamics; see e.g., [19] . So the result in this subsection is somehow a direct corollary. However, one needs to make minor modifications to guarantee that all boundedness assumptions in [19] are satisfied.Given a constant C b f > 0, define f * (x) = sign(f * (x)) min{|f * (x), C b f |}, (B.3) which is bounded with | f * (x)| ≤ C b f .One observation is that the subspace-sparse structure of f * implies that for any δ > 0, there exists a dimension-free constant C b f depending on h * and δ such thatE x∼N (0,I d ) [|f * (x) -f * (x)| 2 ] < δ. The associated mean-field dynamics is ∂ t ρt = ∇ θ • ρt ξ(t)∇ θ Φ(θ; ρt ) , ρt t=0 = ρ 0 , (B.4)where the learning rate ξ(t) = diag(ξ a (t), ξ w (t)I d ) and the initialization ρ 0 are shared with (2.4), andΦ(θ; ρ) = aE x∼N (0,I d ) f NN (x; ρ) -f * (x) σ(w ⊤ x) .The corresponding SGD is given byθ (k+1) i = θ (k) i + γ (k) f * (x k ) -f NN (x k ; Θ (k) ) ∇ θ τ (x k ; θ (k) i ), i = 1, 2, . . . , N, (B.5)where N is the number of neurons andEQUATIONw I d ) ⪰ 0 is the learning rate with γ (k) a = ϵξ a (kϵ) and γ (k) w = ϵξ w (kϵ).Suppose that assumptions made in Theorem 3.4 hold and fix T > 0. Using similar analysis as in Appendix A.2, one can conclude that for any δ > 0, there exists a dimension-free constantC b f such that sup 0≤t≤T |E(ρ t ) -E(ρ t )| < δ.Applying Theorem 3.4 and [19, Theorem 1], we can conclude that for any µ ∈ (0, 1), there exists dimension-free constants N 0 , d 0 , C ϵ , such that for any N ≥ N 0 and d ≥ d 0 , the following holds with probability at least µ for any ϵ ≤ Cϵ d+log N :inf k∈[0,T /ϵ]∩N E N (Θ (k) ) ≥ 1 8 E z∼N (0,I V ) |h * (z) -h * S ⊥ (z ⊥ S )| 2 > 0.If we further assume that N = O(e d ), this indicates that SGD as in (B.5) cannot learn the subspacesparse polynomial f * within finite time horizon and with O(d) samples/data points.This subsection follows [1] closely to approximate and analyze the behavior of w(a, t) for 0 ≤ t ≤ T with ξ a (t) = 0 and ξ w (t) = 1. The dynamics of a single particle starting at θ = (a, 0) ∈ R d+1 can be described by the following ODE:∂ ∂t w(a, t) = aE x g(x, t)σ ′ (w(a, t) ⊤ x)x , w(a, 0) = 0, (C.1) where g(x, t) = f * (x) -f NN (x; ρ t )is the residual. The first observation is thatE x f * (x)σ(w ⊤ x V )x ⊥ V = E x V h * (x V )σ(w ⊤ x V )E x ⊥ V x ⊥ V = 0, ∀ w ∈ R d .By Theorem 3.5, we have that ρ t is supported in (a, w) : w ⊥ V = 0 , and hence thatw ⊥ V (a, t) = 0, ∀ 0 ≤ t ≤ T.We then analyze the behaviour of w V (a, t). Let {e 1 , e 2 , . . . , e p } be an orthonormal basis of V and we denote w i = w ⊤ e i and x i = x ⊤ e i for any w, x ∈ R d and i ∈ {1, 2, . . . , p}.σ ′ w(a, t) ⊤ x = m 1 + L-1 l=1 m l+1 l! (w(a, t) ⊤ x) l + O (w(a, t) ⊤ x) L (C.2)with m l = σ (l) (0) and then an approximated solution to (C.1) (by polynomial expansion with high-order terms omitted) can be written aswi (a, t) = 1≤j≤L Q i,j (t)a j , 1 ≤ i ≤ p,where Q(t) is given by Q(0) = 0 and the following dynamics:       d dt Q i,1 (t) = E x [x i g(x, t)m 1 ], d dt Q i,j (t) = E x   x i g(x, t) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is   , 2 ≤ j ≤ L.(C.3) Let us remark that even if every single wi (a, t) depends on the basis {e 1 , e 2 , . . . , e p }, the linear combinationw(a, t) = p i=1 e i wi (a, t) (C.4)is basis-independent. To see this, let Q j (t) ∈ R d be the j-th column of Q(t) and it holds that       d dt Q 1 (t) = E x [xg(x, t)m 1 ], d dt Q j (t) = E x   xg(x, t) L-1 l=1 m l+1 l! j1+•••+j l =j-1 l s=1 x ⊤ Q js (t)   ,The distance between w i (a, t) and wi (a, t) for i ∈ I can be bounded as follows.Proposition C.1. Suppose that Assumption 4.2 holds. Then|w i (a, t) -wi (a, t)| = O((|a|t) L+1 ), ∀ i ∈ {1, 2, . . . , p}. (C.5)We need the following two lemmas to prove Proposition C.1.Lemma C.2. For any i ∈ {1, 2 . . . , p} and j ∈ {1, 2, . . . , L}, it holds thatQ i,j (t) = O(t j ). (C.6)Proof. The non-increasing property of the energy functional implies that2E(ρ t ) = E x [|g(x, t)| 2 ] ≤ E x [|f * (x) -f NN (x, ρ 0 )| 2 ] = E z [|h * (z)| 2 ] < +∞.Then (C.6) can be proved by induction. For j = 1, it follows from the boundedness ofd dt Q i,1 (t) = |E x [x i g(x, t)m 1 ]| ≤ |m 1 | E x [x 2 i ] • E x [|g(x, t)| 2 ] 1/2 that Q i,e l (t) = O(t).Consider any 2 ≤ j ≤ L and assume that Q i,j ′ = O(t j ′ ) holds for any 1 ≤ i ≤ p and 1 ≤ j ′ < j. Then one has thatd dt Q i,j (t) ≤   Ex[|g(x, t)| 2 ] • E x    x i L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is 2       1/2 =O(t j-1 ), which implies that Q i,j (t) = O(t j ).Lemma C.3. Suppose that Assumption 4.2 holds. We have for all i ∈ {1, 2, . . . , p} thatw i (a, t) = O(|a|t). (C.Proof. There exists a constant C > 0 and a open subset A ⊂ {w ∈ R d : w ⊥ V = 0} containing 0, such that E x x i g(x, t)σ ′ (w ⊤ x) ≤ C, ∀ w ∈ A, t ≥ 0, i ∈ {1,∂ ∂t wi (a, t) -aE x x i g(x, t) m 1 + L-1 l=1 m l+1 l! ( w(a, t) ⊤ x) l ≤ 1≤j≤L d dt Q i,j (t)a j -aE x   xig(x, t)   m1 + L-1 l=1 m l+1 l!   1≤i ′ ≤p 1≤j≤L Q i ′ ,j (t)a j x i ′   l       ≤ L+1≤j≤L L-1 +1 a j • E x   x i g(x, t) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is   ≤ L+1≤j≤L L-1 +1 |a| j •   E[|g(x, t)| 2 ] • E x    x i L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is 2       1/2 =O(|a| L+1 t L ),which combined with (C.2) yields that∂ ∂t p i=1 |w i (a, t) -wi (a, t)| ≤ p i=1 ∂ ∂t w i (a, t) - ∂ ∂t wi (a, t) ≤ p i=1 aE x x i g(x, t) m 1 + L-1 l=1 m l+1 l! (w(a, t) ⊤ x) l -aE x x i g(x, t) m 1 + L-1 l=1 m l+1 l! ( w(a, t) ⊤ x) l + |aE x [x i g(x, t)]| • O (w(a, t) ⊤ x) L + O |a| L+1 t L ≤ p i=1 E x x i a ⊤ g(x, t) L-1 l=1 m l+1 l! (w(a, t) ⊤ x) l -( w(a, t) ⊤ x) l + O |a| L+1 t L =O p i=1 |w i (a, t) -wi (a, t)| + O |a| L+1 t L .Then one can conclude (C.5) from Gronwall's inequality.Even if w(a, t) approximates w(a, t) using polynomial expansion, the coefficients Q(t) are still very difficult to analyze. Thus, we follow [1] to consider the following dynamics that is obtained by replacingg(x, t) = f * (x) -f NN (x; ρ t ) by f * (x) in (C.3): ŵi (a, t) = 1≤j≤LQi,j (t)a j , i ∈ {1, 2, . . . , p}, (C.8)where Q(t) is given by Q(0) = 0 and       d dt Qi,1 (t) = E x [x i f * (x)m 1 ], d dt Qi,j (t) = E x   x i f * (x) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Qis,js (t)x is   , 2 ≤ j ≤ L.(C.9) Similar to w(a, t), the linear combination ŵ(a, t) defined asŵ(a, t) = p i=1 e i ŵi (a, t)is also independent of the orthogonal basis {e 1 , e 2 , . . . , e p }. Q(t) can be understood clearly as follows. Proposition C.4. For any i ∈ {1, 2, . . . , p} and j ∈ {1, 2, . . . , L}, there exists a constant qi,j depending only on h * and σ, such that Qi,j (t) = qi,j t j .(C.10)Proof. The proof is straightforward by induction on j.The next proposition quantifies the distance between Q(t) and Q(t). Proposition C.5. It holds for any i ∈ {1, 2, . . . , p} and j ∈ {1, 2, . . . , L} that|Q i,j (t) -Qi,j (t)| = O(t j+1 ). (C.11)We need the following lemma for the proof of Proposition C.5.Lemma C.6. It holds thatE x |f NN (x; ρ t )| 2 1/2 = O(t). (C.12)Proof. Noticing that f NN (x; ρ 0 ) = 0 by the symmetry of ρ a , one has thatf NN (x; ρ t ) = f NN (x; ρ t ) -f NN (x; ρ 0 ) = a σ(w(a, t) ⊤ x) -σ(0) ρ a (da) = a L l=1 m l l! (w(a, t) ⊤ x) l + O (w(a, t) ⊤ x) L+1 ρ a (da),and hence by Lemma C.3 and w ⊥ V (a, t) = 0 thatE x |f NN (x, ρ t )| 2 = O(t 2 ),which implies (C.12).Proof of Proposition C.5. We prove (C.11) by introduction on j. For j = 1, one hasd dt Q i,1 (t) - d dt Qi,1 (t) = |E x [x i f NN (x, ρ t )m 1 ]| = |m 1 | E x [x 2 i ] • E x [|f NN (x, ρ t )| 2 ] 1/2 = O(t),which leads to |Q i,e l (t) -Qi,e l (t)| = O(t 2 ). Then we consider 2 ≤ j ≤ L and assume that|Q i,j ′ (t) -Qi,j ′ (t)| = O(t j ′ +1 ) holds for 1 ≤ j ′ < j. It can be estimated that d dt Q i,j (t) - d dt Qi,j (t) = E x   x i g(x, t) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is   -E x   x i f * (x) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Qis,js (t)x is   = E x   x i f NN (x; ρ t ) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t)x is   -E x   x i f * (x) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 Q is,js (t) - l s=1 Qis,js (t) x is   = E x   x i f NN (x, ρ t ) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 O(t js )x is   -E x   x i f * (x) L-1 l=1 m l+1 l! 1≤i1,...,i l ≤p j1+•••+j l =j-1 l s=1 qis,js t js + O(t js+1 ) - l s=1 qis,js t js x is =O(t j ),where we used Lemma C.6. Then one can conclude (C.11).We prove Theorem 4.5 in this subsection. Definition C.7 (Algebraic independence). Let v 1 , v 2 , . . . , v m ∈ R[a 1 , a 2 , . . . , a p ] be polynomials in a 1 , a 2 , . . . , a p . We say that v 1 , v 2 , . . . , v m are algebraically independent if for any nonzero polynomial F : R m → R,F (v 1 (a 1 , a 2 , . . . , a p ), . . . , v m (a 1 , a 2 , . . . , a p )) ̸ = 0 ∈ R[a 1 , a 2 , . . . , a p ]. Lemma C.8. If v 1 , v 2 , . . . , v p ∈ R[a] are R-linearly independent, then det     v 1 (a 1 ) v 1 (a 2 ) • • • v 1 (a p ) v 2 (a 1 ) v 2 (a 2 ) • • • v 2 (a p ) . . . . . . . . . . . . v p (a 1 ) v p (a 2 ) • • • v p (a p )     (C.13) is a non-zero polynomial in R[a 1 , a 2 , . . . , a p ].Proof. Let n i be the smallest degree of nonzero monomials of v i and let c i be the accociated coefficient for i = 1, 2, . . . , p. Without loss of generality, we assume that n 0 < n 1 < • • • < n p (otherwise one can perform some row reductions or row permutations). The polynomial defined in (C.13) consists of monomials of degree at leastn 1 + n 2 + • • • + n p .So it suffices to prove that the sum of monomials with degree beingn 1 + n 2 + • • • + n p is nonzero, which is true since det      c 1 a n1 1 c 1 a n1 2 • • • c 1 a n1 p c 2 a n2 1 c 2 a n2 2 • • • c 2 a n2 p . . . . . . . . . . . . c p a np 1 c p a np 2 • • • c p a np p      = c 1 c 2 . . . c p • det      a n1 1 a n1 2 • • • a n1 p a n2 1 a n2 2 • • • a n2 p . . . . . . . . . . . . a np 1 a np 2 • • • a np p     is nonzero as a generalized Vandermonde matrix.Proof of Theorem 4.5. Since v 1 , v 2 , . . . , v p ∈ R[a] be R-linearly independent with the constant terms being zero, we can see thatv ′ 1 , v ′ 2 , . . . , v ′ p ∈ R[a] are also R-linearly independent. Noticing that ∂ ∂a j 1 p (v i (a 1 ) + v i (a 2 ) + • • • + v i (a p )) = 1 p v ′ i (a j ), one can conclude that 1 p (v 1 (a 1 ) + • • • + v 1 (a p )), . . . , 1 p (v p (a 1 ) + • • • + v p (a p )) ∈ R[a 1 , .. . , a p ] are R-algebraically independent by using Theorem 4.6 and Lemma C.8.Some ideas in this subsection are from [1] , but the proofs are significantly different since we need to show the algebraic independence to obtain a non-degenerate kernel, as discussed in Section 4. Lemma C.9. Suppose that Assumption 4.1 holds with s ∈ N + . There exists some orthonormal basis {e 1 , e 2 , . . . , e p } of V such that the coefficients qi,j ,1 ≤ i ≤ p, 1 ≤ j ≤ L in (C.10) satisfies s 1 < s 2 < • • • < s p ≤ s, (C.14)where s i = min{j : qi,j ̸ = 0}, i = 1, 2, . . . , p.Proof. Let ŵV (t) be the dynamics defined in (4.1). It can be seen that the Taylor's expansion of ŵV (t) at t = 0 up to s-th order is given byp i=1 s j=1 e i Qi,j (t) = p i=1 s j=1 e i qi,j t j ,where Qi,j and qi,j are as in (C.9) and (C.5). According to Assumption 4.1, the matrix (q i,j ) 1≤i≤p,1≤j≤s is of full-row-rank. One can thus perform the QR decomposition, or equivalently choose some orthogonal basis, to obtain (C.14).In the rest of this subsection, we will always denote s = (s 1 , s 2 , . . . , s P ) andu(a 1 , . . . , a p , t) = 1 p (w(a 1 , t) + w(a 2 , t) + • • • + w(a p , t)), ũ(a 1 , . . . , a p , t) = 1 p ( w(a 1 , t) + w(a 2 , t) + • • • + w(a p , t)), û(a 1 , . . . , a p , t) = 1 p ( ŵ(a 1 , t) + ŵ(a 2 , t) + • • • + ŵ(a p , t)), for t ∈ [0, T ],where w(a, t), w(a, t), and ŵ(a, t) are defined in (C.1), (C.4), and (C.8), respectively. Recall that p 1 , p 2 , . . . , p ( n+p p ) are the orthonormal basis of P V,n with input z ∼ N (0, I V ), where P V,n is the collection of all polynomials on V with degree at most n = deg(h * ) = deg(f * ). Proposition 4.4 aims to bound from below the smallest eigenvalue of the kernel matrix (4.2) whose definition is restated as followsK i1,i2 (t) = E a1,...,ap E z,z ′ p i1 (z)σ(u(a 1 , . . . , a p , t) ⊤ z)σ(u(a 1 , . . . , a p , t) ⊤ z ′ )p i2 (z ′ ) , where σ(ξ) = (1 + ξ) n , (a 1 , . . . , a p ) ∼ U([-1, 1] p ), and 1 ≤ i 1 , i 2 ≤ n+p p .To do this, we define three n+p pEQUATIONMi1,i2 (a, t) = E z p i1 (z)σ(ũ(a i2 , t) ⊤ z) , Mi1,i2 (a, t) = E z p i1 (z)σ(û(a i2 , t) ⊤ z) ,where a = a 1 , a 2 , . . . , a ( n+p p ) anda i ∈ R p for i = 1, 2, . . . , n+p p . Lemma C.10. It holds that det( M (a, t)) = ( n+p p ) i=1 0≤∥ji∥1≤Lnĥj t ∥j∥1 a j , (C.15)where j = j 1 , j 2 , . . . , j ( n+p p ) with j i ∈ N p for i = 1, 2, . . . , n+p p , ĥj is a constant depending on h * and σ, and a j represents the product of entrywise powers.Proof. The result follows directly from Proposition C.4.Lemma C.11. It holds that det( M (a, t)) = ( n+p p ) i=1 0≤∥ji∥1≤Lnhj (t)a j , with hj (t) = ĥj t ∥j∥1 + O(t ∥j∥1+1 ).Proof. The result follows directly from Proposition C.4 and Proposition C.5.Lemma C.12. For any m ∈ N, one has thatEQUATIONwhere P h V,m is the collection of all homogeneous polynomials in z ∈ V with degree m.Proof. Without loss of generality, we assume that V = R p and prove the result by induction on p. (C. 16 ) is clearly true for p = 1. Then we consider p ≥ 2 and assume that (C.16) holds for p -1 and any m.For any q, z ∈ R p , we denote that q = (q 2 , . . . , q p ) and z = (z 2 , . . . , z p ). Let t 0 , t 1 , . . . , t m ∈ R be distinct. Then it follows from the invertibility of (t j i ) 0≤i,j≤m and     (t 0 z 1 + q⊤ z) m (t 1 z 1 + q⊤ z) m . . . (t n z 1 + q⊤ z) m      =      m i=0 m i t i 0 z i 1 (q ⊤ z) m-i m i=0 m i t i 1 z i 1 (q ⊤ z) m-i . . . m i=0 m i t i m z i 1 (q ⊤ z) m-i      =     1 t 0 • • • t m 0 1 t 1 • • • t m 1 . . . . . . . . . . . . 1 t m • • • t m m          m 0 (q ⊤ z) m m 1 z 1 (q ⊤ z) m-1 . . . m m z m 1      that z i 1 (q ⊤ z) m-i ∈ span (r ⊤ z) m : r ∈ R p , ∀ q ∈ R p-1 , i ∈ {0, 1, .. . , m}. Then using the induction hypothesis that (C. 16 ) is true for p -1 and m -i, i = 0, 1, . . . , m, one can conclude that (C.16) is also true for p and m.Lemma C.13. For σ(ξ) = (1 + ξ) n , one has that span σ(q ⊤ z) : q ∈ V = P V,n .Proof. One can still assume that V = R p without loss of generality. Consider any q ∈ R p and any distinct t 0 , t 1 , . . . , t n ∈ R. Then     σ((t 0 q) ⊤ z) σ((t 1 q) ⊤ z) . . . σ((t n q) ⊤ z)      =      n i=0 n i t i 0 (q ⊤ z) i n i=0 n i t i 1 (q ⊤ z) i . . . n i=0 n i t i 0 (q ⊤ z) i      =     1 t 0 • • • t n 0 1 t 1 • • • t n 1 . . . . . . . . . . . . 1 t n • • • t n n          n 0 n 1 q ⊤ z . . . n n (q ⊤ z) n     .Note that the matrix (t j i ) 0≤i,j≤n is invertible when t 0 , t 1 , . . . , t N are distinct. Therefore, one can conclude that(q ⊤ z) i ∈ span σ(r ⊤ z) : r ∈ R p , ∀ q ∈ R p , 0 ≤ i ≤ n,which combined with Lemma C.12 implies thatP V,n ⊃ span σ(q ⊤ z) : q ∈ R p ⊃ P h V,0 ⊕ P h V,1 ⊕ • • • ⊕ P h V,n = P V,n , which completes the proof.Proof of Lemma 4.7. We prove the result by induction. When m = 1, the result follows directly from the R-algebraic independence of v 1 , v 2 , . . . , v p . Now we assume that the result is true for m -1 and consider the case of m. Suppose that 0 = j1,...,jmX j m l=1 v(a l ) j l = jm   j1...,jm-1 X j m-1 l=1 v(a l ) j l   v(a m ) jm .By the R-algebraic independence of v 1 , v 2 , . . . , v p , one must have j1...,jm-1X j m-1 l=1 v(a l ) j l = 0, ∀ j m ,which then leads to X j = 0, ∀ j by the induction hypothesis.Lemma C.14. Suppose that Assumption 4.1 and 4.2 hold and let ĥj be the coefficient of det( M (a, t)) in (C. 15) . Then there exists some ĵ = ĵ1 , ĵ2 , . . . , ĵ( n+p p ) with ∥ ĵ∥ 1 ≤ sn n+p p , such that ĥĵ ̸ = 0.Proof. Let X(q) ∈ R ( n+p p )×( n+p p ) be defined asX i1,i2 (q) = E z p i1 (z)σ(q ⊤ i2 z) ,where q = q 1 , q 2 , . . . , q ( n+p p ) with q i ∈ V . Then det(X(q)) is a polynomial in q of the form det(X(q)) = ( n+p p ) i=1 0≤∥ji∥1≤nX j q j , (C.17)where we understand q i as a (coefficient) vector in R p associated with a fixed orthonormal basis {e 1 , e 2 , . . . , e p } of V that satisfies Lemma C.9. By Lemma C.13, there exists some q such that σ(q ⊤ 1 z), σ(q ⊤ 2 z), . . . , σ(q ⊤ ( n+p p )z) form a basis of P V,n , which implies that det(X(q)) ̸ = 0 for this q. Thus, (C.17) is a non-zero polynomial in q. Let s = (s 1 , s 2 , . . . , s p ) collect all indices s i from Lemma C.9. DenoteEQUATIONJ S =      j : ( n+p p ) l=1 s ⊤ j l = S, X j ̸ = 0, 0 ≤ ∥j i ∥ ≤ n, 1 ≤ i ≤ n + p p     .Then we have thatdet( M (a, t)) = det X û(a 1 , t), û(a 2 , t), . . . , û(a ( n+p p ) , t) = ( n+p p ) i=1 0≤∥ji∥1≤n X j ( n+p p ) l=1 û(a l , t) j l = j∈J S X j ( n+p p ) l=1 ûs (a l , t) j l + O(t S+1 ),where ûs (a l , t) = (û 2,s1 (a l , t), . . . , ûp,sp (a l , t)) and ûi,si (a 1 , . . . , a p , t) = 1 p p k=1 qi,si t si a si k collects the leading order terms of ûi (a 1 , . . . , a p , t). According to Assumption C.9, Theorem 4.5, and Lemma 4.7, we have that j∈J S X j ( n+p p ) l=1 ûs (a l , t) j l ̸ = 0, which provides at least one non-zero term in det( M (a, t)) whose degree in t isS ≤ sn n + p p .This completes the proof.Proof of Proposition 4.4. According to Lemma C.14, there exist some ĵ with ∥ ĵ∥ 1 ≤ sn n+p p such that ĥĵ ̸ = 0. According to Lemma C.11 and Lemma 103 in [1] , it holds thatE a det( M (a, t)) 2 ≥ C 1 hĵ (t) 2 = C 1 ĥĵ t ∥ ĵ∥1 + O(t ∥ ĵ∥1+1 ) 2 ≥ C 1 2 ĥĵ t 2∥ ĵ∥1 , (C.18)where a ∼ U ([-1, 1] p ) ( n+p p ) for some constant C 1 depending only on n, p, s, and for sufficiently small t. It can be seen that Then we can obtain (4.3) as we set L = 2sn n+p p .K(t) = 1 n+p p E a M (Proof of Theorem 4.3. Let C, T be the constants in Proposition 4.4 and consider t > T . It follows from Theorem 3.5 that w ⊥ V (a, T ) = 0, which leads to that u ⊥ V (a 1 , . . . , a p , T ) = 0 and hence that f NN (x; ρ t ) and g(x, t) only depend on x V for all t ≥ T . Define g(t) ∈ R ( n+p p ) via g i (t) = E x [g(x, t)p i (x V )]. Then according to (2.6) and ( 4.3), one has for t > T that We provide some verification or characterization of Assumption 4.1. Without loss of generality, we fix an orthonormal basis and V and view that V = R p . The results in this subsection are independent of the choice of the orthonormal basis.d dt E(ρ t ) = - P i=1 E a E x,x ′ g i (x, t)σ(w(a, T ) ⊤ x)σ(w(a, T ) ⊤ x ′ )g i (x ′ , t) = -g(t) ⊤ K(T )g(t) ≤ -λ min (K(T ))∥g(t)∥ 2 = -2λ min (K(T ))E(ρ t ) ≤ -2CT 2sn( n+p p ) E(ρ t ), which implies that E(ρ t ) ≤ E(ρ T ) exp -2CT 2snIt has been discussed in Section 4.1 that if σ ∈ C s (R) with s = 2 p-1 and σ (1) (0), σ (2) (0), . . . , σ (p) (0) are all nonzero, then Assumption 4.1 can be verified with s for h * (z) = z 1 + z 1 z 2 + • • • + z 1 z 2 • • • z p ,using the calculation in [1, Proposition 33] . Similarly, the same result is also true for h Proof. We assume that α 1 = (1, 0, . . . , 0), α 2 = (1, 1, 0, . . . , 0), . . . , α m = (1, 1, . . . , 1). As in the proof of Lemma C.9, Assumption 4.1 is true with s if and only if the matrix q := (q i,j ) 1≤i≤p,1≤j≤s is of full-row-rank, i.e., det(q q⊤ ) ̸ = 0. By (C.9), each entry in q is a polynomial in (c 1 , c 2 , . . . , c m ), which implies that det(q q⊤ ) is also a polynomial in (c 1 , c 2 , . . . , c m ). This polynomial is nonzero since it takes nonzero value at (1, 1, . . . , 1, 0, . . . , 0) with the p entries being 1 and all other entries being 0, which is because that Assumption 4.1 is true for hEQUATION* (z) = z 1 + z 1 z 2 + • • • + z 1 z 2 • • • z p .Finally, the conclusion of Proposition D.1 is true since the set of roots of a nonzero polynomial is of measure zero with respect to the Lebesgue measure.Implied by Theorem 4.3The discussion in this subsection is similar to those in Appendix B.2. We slightly modified the flow ρ t generated by Algorithm 1 to guarantee some boundedness conditions, and then use the standard dimension-free estimate [19] to derive a sample complexity result implied by Theorem 4.3.We use the same bounded modification f * as in (B.3) for a given constant C b f > 0. Similarly, for any δ > 0 and C w > 0, there exists dimension-free constant C b σ > 0 depending on h * , δ, C w , such that one can modify the activation function σ(ζ) = (1 + ζ) n to σ satisfying that ∥σ∥ L ∞ (R) ≤ C b σ , ∥σ ′ ∥ L ∞ (R) ≤ C b σ , ∥σ ′′ ∥ L ∞ (R) ≤ C b σ , and E x∼N (0,I d ) |σ(w ⊤ x) -σ(w ⊤ x)| 2 < δ, E x∼N (0,I d ) |σ(w ⊤ x) -σ(w ⊤ x)| 2 < δ, for all w ∈ R d with ∥w∥ ≤ C w .The associated mean-field dynamics ρt , that can be viewed as a slight modification of ρ t generated by Algorithm 1, is given by (B.4) for 0 ≤ t ≤ T and follows∂ t ρt = ∇ θ • ρt ξ(t)∇ θ Φ′ (θ; ρt ) , ρt t=T = ρT , (D.1) with Φ′ (θ; ρ) = aE x∼N (0,I d ) fNN (x; ρ) -f * (x) σ(w ⊤ x) , fNN (x; ρ) = aσ(w ⊤ x)ρ(da, dw),for t ≥ T . Here, the learning rate ξ(t) = diag(ξ a (t), ξ w (t)I d ) is shared with Algorithm 1, namely ξ a (t) = 0, ξ w (t) = 1 for 0 ≤ t ≤ T and ξ a (t) = 1, ξ w (t) = 0 for t ≥ T .The SGD associated to the modified mean-field dynamics is given byEQUATIONEQUATIONfor i = 1, 2, . . . , N and k = 0, 1, . . . , T /ϵ-1, where N is the number of neurons and T /ϵ is assumed to be an integer, and EQUATIONb σ such that sup 0≤t≤T ′ |E(ρ t ) -E(ρ t )| < δ.Applying Theorem 4.3 and [19, Theorem 1], we can conclude that for any µ ∈ (0, 1) and any δ > 0, there exists dimension-free constants N 0 , d 0 , C ϵ , such that for any N ≥ N 0 and d ≥ d 0 , the following holds with probability at least µ for any ϵ ≤ Cϵ d+log N with T /ϵ ∈ N:inf k∈[0,T ′ /ϵ]∩N E N (Θ (k) ) < δ.If we further assume that N = O(e d ), this indicates that SGD with (D.2) and (D.3) can learn the subspace-sparse polynomial f * within finite time horizon and with O(d) samples/data points.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.1/2 γ 0 (dθ 0 , d θ0 ) 


Decision: The reviewers initially raised several important concerns and the extensive responses of the authors have partially addressed them, which lead the reviewers increase their scores. Given all the material (paper, reviews, rebuttal and discussion), I am recommending an acceptance by trusting the authors that they will implement all the promised changes and clarifications one by one (there are many points to be clarified & elaborated as revealed during the discussion period). Please make sure you go over all the requested changes carefully and update the camera-ready version accordingly.

Reviews: review: paper_summary: paper_summary: This paper studies how well two-layer neural networks can learn subspace-sparse polynomials using stochastic gradient descent when the input is Gaussian. It proposes a necessary condition for learning, showing that if the target function has certain symmetries, the neural network cannot learn it perfectly. The paper also provides a sufficient condition and training strategy that guarantee the loss decays exponentially fast to zero, with rates independent of the input dimension.
strength_weakness: strength_weakness: No obvious weakness.
questions: questions: No questions
limitations: limitations: N/A
paper_summary: paper_summary: This paper considers learning a subspace sparse polynomial with Gaussian data, using a two-layer neural network trained in the mean-field regime. Their main results are two fold:
1)They introduce a reflective property that generalizes the non- merged-staircase property from [1] to Gaussian input. They show that this condition is sufficient for the dynamics to be stuck in a suboptimal saddle-subspace for a constant time horizon.
2) They show that a slightly stronger condition is sufficient for some modified dynamics to decay to $0$ exponentially fast.
strength_weakness: strength_weakness: - I would be careful when claiming novelty for the basis-free generalization of the merged-staircase property: as correctly pointed out in the related works paragraph, several papers [2,7,9,10] have proposed such properties. It should be easy to check that Eq. (3.2) is equivalent to [2,7,9] for leap=1. The difference is that Eq. (3.1) takes into account that the activation might not be expressive enough (see questions).

- Theorem 3.4 seems to be very similar to [1] ([1] goes further and show concentration on a dimensionless dynamics, so the proof does not directly compare to a dynamics that is $0$ on the subspace)

- There are several unnatural simplifications made in the algorithm. 
First the dynamics is actually $p$ particles that are trained independently, to avoid dealing with a PDE (then it is just an ODE on a single particle). This does not seem necessary (e.g., [1]). Is there a reason to not directly consider the PDE?
Second, the activation is changed between the two phases of the algorithm.
Third, the result is proved under Assumption 4.1, which looks very much like assuming that the algorithm works (phase 1 of training is roughly this ODE, and therefore it looks like we are assuming that the dynamics converge to prove that the dynamics converge)... This seems to be a major flaw, but could be clarified during discussions. Could a smooth analysis, similar to [1] work here? (see questions)

- The exponential convergence in Theorem 4.3 is confusing. When seeing such a result in the PDE literature, I assume another type of analysis, e.g., log-sobolev inequality. Here, if I am not mistaken, the exponential convergence is simply due to the fact that we switch in the second phase of the algorithm to a linear dynamics, where the kernel has eigenvalues bounded away from 0.
questions: questions: - Is there no hope to prove Assumption 4.1, even at the price of randomizing some coefficients of $h_*$ or $\sigma’$ and only proving an almost sure result?

- Is definition 3.1 that much different from 3.2? Except in the edge case of $\sigma’$ identically $0$ (which is the only counterexample given), it seems that we can always condition on $z_s^\perp$.
limitations: limitations: Yes
paper_summary: paper_summary: This manuscript extends previous work on learning sparse polynomials on the hypercube to Gaussian space. To achieve this, the authors introduce a generalization of the merged staircase property studied in previous work. Specifically, they demonstrate that their "reflective property" is necessary to learn polynomials in O(1) time, while a slightly stronger condition, defined in Assumption 4.1, is sufficient for learning.

As is common in the literature, the authors propose a two-stage training process. First, they train the input layer while keeping the output layer fixed. After learning the representation, they fix the input layer and train the output layer to learn the link function.
strength_weakness: strength_weakness: The main weakness of the proposed algorithm is its requirement to run the training algorithm several times to average the weights and change the activation function in the second stage, which differs significantly from practical training algorithms. At the same time, existing work addresses this issue with random feature approximation in the second stage, though this is also not an ideal solution. Therefore, the question of whether we can train input-output layers together remains an open problem in the field.
questions: questions: The authors did not discuss their "Reflective property" in the context of leap complexity suggested in [1]. It is known that the merged staircase property corresponds to leap 1 functions. In this sense, is there a way to relate the reflective properties with the leap complexity?

Additionally, in the definition of the "Reflective property," they require the condition to hold for any shifted version of the activation (i.e., for any $u \in \mathbb{R}$  in Definition 3.3). However, in Theorem 3.5, they did not use a bias in the activation. What is the reason for this discrepancy?

Furthermore, I am not familiar with the terminology of Dirac delta measure defined on a subspace, and I could not find its definition in the paper either. Could the authors include this definition in the revised version of the paper?

[1] Abbe, E., Boix-Adserà, E., & Misiakiewicz, T. (2022). The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks. ArXiv, abs/2202.08658.
limitations: limitations: See the Weaknesses part.
paper_summary: paper_summary: This submission studies the training dynamics of a two-layer neural network in the mean-field limit, when the data is isotropic Gaussian and the target function is a low-dimensional multi-index model. The authors consider the population gradient flow on the squared error, and identify a "reflective property" that is necessary for GD to learn the target function in finite time. A multi-phase gradient-based training algorithm is also proposed to demonstrate that such property is (almost) sufficient for GD learnability. This provides a coordinate-free generalization of the merged staircase property (MSP) in (Abbe et al. 2023).
strength_weakness: strength_weakness: I have reviewed this paper in another venue. The authors addressed some of my concerns in the revision.  
I shall reiterate the remaining weaknesses from my previous review. 

- This submission only analyzes the mean-field dynamics up to finite time. If $T$ scales with the dimensionality, then it is known that SGD can learn functions with higher leap / information exponent; these functions may not satisfy the reflective property, yet the mean-field dynamics can achieve low population loss. I do not hold this against the authors since this point is already acknowledged in the main text. 

- The learning algorithm in Section 4 is non-standard as it requires an averaging over trajectories and a change of activation function during training. Moreover, the authors do not present any learning guarantee for the discretized dynamics, and therefore a sample complexity cannot be inferred from the analysis. 

After reading the revised manuscript, my main concern is that the relation between the proposed reflective property and prior definitions of SGD hardness (in the Related Works paragraph starting line 73) is not clearly explained. Specifically, the authors only listed the alternative rotation-invariant conditions on SGD learnability, but did not discuss the advantages of the proposed definition. 

- Is the reflective property equivalent to isoLeap > 1, if we do not take the activation $\sigma$ into account (that is, assuming $\sigma$ is sufficiently expressive)? If not, how does it differ from the sufficient conditions in (Abbe et al. 2023) (Bietti et al. 2023)? More explanation is needed. 

- The authors brought up the even symmetry condition in (Dandi et al. 2024), which does not exclude high isoLeap functions. In fact, (Lee et al. 2024) (Arnaboldi et al. 2024) have shown that any non-even single-index functions can be learned by variants of online SGD in $T=O(1)$ time horizon (and the discretized dynamics require $n=O(d)$ samples). The authors should clarify if this is consistent with the proposed necessary condition for SGD.  
(Lee et al. 2024) *Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit*.  
(Arnaboldi et al. 2024) *Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions*.
questions: questions: See Weaknesses above.
limitations: limitations: See Weaknesses above.

Comments: comment: Dear Reviewer 7m62,

Thank you very much for increasing the score! We appreciate all your valuable comments and will revise our manuscript accordingly.
Thank you for taking into account my comments.
I increased my score to 5. The setting is interesting and can be of interest to the NeurIPS community. However, I believe that the paper still falls short of solving the main difficulty of going from hypercube to Gaussian, which is the algebraic independence, and requires an unatural construction (repeating step 2).
Dear Reviewer kar1,

Thank you very much for your valuable comment and for your support. As discussed in the comment of Reviewer 7m62, our equation 3.2 without $\sigma$ is equivalent to IsoLeap$\geq 2$, which we agree with after a careful check. Hence, the contribution on the necessary condition is a condition considering the expressiveness of $\sigma$ and a different analysis. We will clearly discuss the equivalence and modify the description of our contribution accordingly.
I thank the authors for their detailed response.

**Connection between isoleap and reflective property:** As Reviewer 7m62 suggests, the isoleap property can be extended to $\ell_2$ integrable functions with respect to the standard Gaussian measure. Moreover, after spending some time on the suggested definition, I am fairly confident that the reflective property is equivalent to isoleap > 1,  if the activation function can be chosen freely in the definition of reflective property. In this case, the major distinction lies in incorporating the Hermite expansion of the activation function, which somewhat weakens the contribution of “proposing a basis-free generalization of the merged-staircase property”.

**Definition 3.2 and Theorem 3.5:** Thank you for the clarification.

Despite my concern regarding the first point, I would be happy to see this paper accepted at NeurIPS. Therefore, I will keep my score unchanged. However, I strongly urge the authors to address and clarify the issue regarding isoleap in their final submission.
Dear Reviewer JjNH,

Thank you very much for your insightful and valuable comments! We are sorry for the unclarity and we will revise our manuscript according to your suggestions:
* We will discuss the discretization of the mean-field gradient flow with online SGD, using standard discretization error bound in Mei et al. (2019). We did not include it in our original submission due to a technical issue from the unboundedness of $f$ and $\hat{\sigma}$. This issue can be resolved with slight modifications, as described in our rebuttal.
* We will clearly mention in our revision that our condition and analysis are only for the mean-field gradient and the online SGD (that is the "original SGD" mentioned in our reply) and that the recovery ability of SGD can be improved in some cases when people apply other techniques such as batch reuse, as in Dandi et al. (2024), Lee et al. (2024), and Arnaboldi et al. (2024). We will include the discussion of all these different but related settings and conditions.

Please let us know if you have other comments or questions. Thank you again!
> Lee et al. (2024) run the training algorithm for $\Theta(d\text{polylog}d)$ iterations

For functions that do not satisfy even symmetry, including $\text{He}_5$, finite time is sufficient as seen in Dandi et al. (2024), Lee et al. (2024) and Arnaboldi et al. (2024). This is *not* the reason that the proposed reflective property fails to predict weak recovery.  

My point in the original review is that when the authors cite related conditions on SGD learnability, they need to first understand the settings under which these conditions have predictive power, and compare these against the setting studied in this submission.  
As it currently stands, this submission only considers the mean-field *gradient flow* at finite time, and there's no notion of "original SGD" yet. It is expected that if you discretize this flow with online SGD, then the proposed condition is informative. But if you use SGD to minimize the empirical loss, there's a chance that the data-reuse mechanism in Dandi et al. (2024) could appear.
Dear Reviewer 7m62,

Thank you very much for your reply and for the insightful comments. I agree with you that IsoLeap can be defined beyond polynomials and that Equation 3.2 should be equivalent to IsoLeap$\geq 2$. We will add the discussion and clarify this point in our revision. I also agree with you that we probably need to modify the claim of our contribution given the equivalence. But I would still consider the reflective property as interesting -- This property (Equation 3.1) is derived when we compare the original mean-field dynamics (Equation 2.4) and the flow in a subspace (Equation 3.7); after removing the dependence of $\sigma$, a stronger condition (Equation 3.2) equivalent to IsoLeap$\geq 2$ is derived, which re-discovers an interesting condition from a different perspective/analysis. We also give additional insight as Equation 3.1 considers the expressiveness of $\sigma$.
Dear Reviewer JjNH,

Thank you very much for your reply and for the insightful comment. The function $f^*(x) = \text{He}_5(x^\top \theta)$ indeed satisfies the reflective property and our results are consistent with Lee et al. (2024) and Dandi et al. (2024):
* If I understand correctly, both Lee et al. (2024) and Dandi et al. (2024) are based on reusing batches. Differently, our work analyzes the mean-field dynamics induced by the original SGD without reusing batches.
* Lee et al. (2024) run the training algorithm for $\Theta(d\cdot\text{polylog}d)$ iterations and the stepsize is chosen as $\eta = \Theta(d^{-1})$. So the training time is $\Theta(d\cdot\text{polylog}d)\cdot \Theta(d^{-1}) = \Theta(\text{polylog}d)$. I would consider this as "almost" finite time, not exactly finite time.

Please let us know if you have further comments or questions.
I thank the authors for their detailed response and for clearing out some of my confusion. Here are come comments:

- **Concerning isoleap:** The isoleap is defined in terms of the Hermite coefficients in the orthonormal decomposition of the target function, which always exist for squared integrable functions (which is always the case when using squared loss), so it is not limited to polynomials. (The isoleap is always defined and finite for any squared integrable function.)  

- **Concerning Assumption (3.2):** Let's consider $S$ to be one direction $z_1$. Then $ h_* (z) = \sum_{k \geq 0} He_k (z_1) h_k (z_{-1}).$ Therefore the condition (3.2) shows that $h_1(z_{-1}) = 0$ for all $z_{-1}$. Hence, in this coordinate basis, $z_1$ only appears in $He_k$ with $k\geq2$, so its leap $\geq 2$ in that basis, and therefore the isoleap is $\geq 2$ (this generalizes straighforwardly to $S$ of any dimension). Hence Assumption (3.2) is equivalent to isoleap $\geq 2$ (if I am not missing anything!). Indeed, Assumption 3.1 takes into account the expressivity of the activation function. However, note that in your example, one cannot fit $h_*$ anyway (degree $n$ activation and degree $n+1$ activation). More generally, for any $\sigma$ degree $k$ polynomial, I would expect (didn't check) $\sigma (u + v^T z_{S^\perp})$ to span  all polynomials of degree-$k$ and to be equivalent to Assumption (3.2) whenever $h_*$ is degree $\leq k$, which is anyway the only part of the target function that can be fitted by the neural network. Hence, for $\sigma$ degree-$k$ polynomial, we could imagine the following: decompose $h_* = h_{\leq k} + h_{>k}$, assume Assumption (3.2) on $h_{\leq k}$ and disregard $h_{>k}$, which does not contribute to the dynamics by orthogonality. (Note that $h_{\leq k}$ can only have higher leap than $h_*$.) All in all, I am still not convinced that Assumption (3.1) is significantly different from isoleap to claim ``we propose a basis-free generalization of the merged-staircase property in [1]''. I apologize for being nit-picking (and I do not consider it to be the most important part of your paper), but several other groups, including [2,7,9,10] cited in your paper, have worked on this basis-free leap complexity. Please let me know if I misunderstood anything.

- **Concerning Algorithm 1 and Assumption 4.1:** Thank you for the clarification, I mistakenly thought the initialization was a point mass. I acknowledge that the Gaussian case is much harder than the hypercube case, due to the need of an algebraic independence. Thank you for making the effort of verifying Assumption 4.1 in both these cases, I believe it already makes a strong case for Assumption 4.1!
> Consistency with Lee et al. (2024) and Arnaboldi et al. (2024). Lee et al. (2024) require a sample complexity $\mathcal{O}(d\cdot \text{polylog} d)$.  

The authors need to elaborate. Consider single index model $f^*(x) = \text{He}_5(x^\top \theta)$. Does this function satisfy the proposed reflective property?  
This function is learnable under finite $T$ using the algorithm in Lee et al. (2024). In fact, Dandi et al. (2024) (which the authors cited) showed that two GD steps with batch size $\mathcal{O}(d)$ suffice.