Title: Mean-Field Analysis for Learning Subspace-Sparse Polynomials with Gaussian Input
Venue: NeurIPS 2024
Accepted: yes
Keywords: Subspace-sparse polynomial

merged-staircase property

algebraic independence

mean-field analysis

stochastic gradient descent
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.

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

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.

1 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.

2 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:

where 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:

In practice, a widely used algorithm for solving (2.2) is the stochastic gradient descent (SGD) that iterates as

where x k , k = 1 , 2 , . . . are the i.i.d. samples drawn from N (0 , I d ) and γ ( k ) = diag ( γ ( k ) a , γ ( k ) w 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:

where ρ ∈ P ( R d +1 ) is a probability measure on the parameter space R d +1 , and generalize the loss/energy functional (2.2) to

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:

where ξ ( t ) = diag ( ξ a ( t ) , ξ w ( t ) I d ) ∈ R ( d +1) × ( d +1) with ξ a ( t ) ≥ 0 and ξ w ( t ) ≥ 0 being the learning rates and

One can also write Φ( θ ; ρ ) as

where

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

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 :

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.

3 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 .

3.1 Reflective Property

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 : R ≥ 0 → R satisfy that ∥ ξ a ∥ L ∞ ( R ≥ 0 ) ≤ K ξ and ∥ ξ w ∥ L ∞ ( R ≥ 0 ) ≤ K ξ for some constant K ξ > 0 . Furthermore, ξ a and ξ w are Lipschitz continuous with ∫ + ∞ 0 ξ a ( t ) dt = + ∞ and ∫ + ∞ 0 ξ w ( t ) dt = + ∞ .

(iii) The initialization is ρ 0 = ρ a × ρ w such that ρ a is symmetric and is supported in [ -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 as E ( ρ 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 σ if

where 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 that

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,

which actually implies (3.1). But these two conditions are different: h ∗ ( z ) = z 1 + z 1 z 2 + z 1 z 2 z 3 does not satisfy (3.2) but still satisfies (3.1) if σ ( ζ ) = ζ . Weuse (3.1) with σ ′ because we want to emphasize that the SGD learnability depends on the activation function σ . If σ is less expressive, then SGD may not learn the target function even if h ∗ itself satisfies the merged-staircase property. Typically people use activation functions that are expressive enough, for which (3.1) and (3.2) are similar. In addition, it can be verified that (3.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.

Our 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 that

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 have

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.

3.2 Proof Sketch for Theorem 3.4

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 small 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 and

where x ⊥ S = x -x S , then it holds for any t ≥ 0 that

Here, 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 decomposition w ⊤ 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 ⊥ ) :

where ˆ ρ 0 ∈ P ( R × S ⊥ ) satisfies ρ 0 = ˆ ρ 0 × δ S , ˆ θ = ( a, w ⊥ S ) , ˆ ξ ( t ) = diag ( ξ a ( t ) , ξ w ( t ) I S ⊥ ) , and

with and

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 that

where 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.

Algorithm 1 Training strategy

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 .

4: Reset ρ T as the distribution of (0 , u ( a 1 , . . . , a p , T )) = ( 0 , 1 p ∑ p i =1 w ( a i , T ) ) . Train the neural network with ξ a ( t ) = 1 , ξ w ( t ) = 0 , and a new activation function ˆ σ ( ζ ) = (1 + ζ ) n , where n = deg ( f ∗ ) , for t ≥ T starting at ρ T .

4 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 .

4.1 Training Procedure and Convergence Guarantee

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 ˆ w V ( t ) in V :

We assume that for some s ∈ N + , the Taylor's expansion up to s -th order of ˆ w 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, if V = R p and h ∗ ( z ) = z 1 + z 1 z 2 + z 1 z 2 z 3 + · · · + z 1 z 2 · · · z p which satisfies the merged-staircase property, then it can be computed that the leading order terms of the coordinates of ˆ w V ( t ) are given by ( c 1 t, c 2 t 2 , c 3 t 2 2 , . . . , c p t 2 p -1 ) with nonzero constants c 1 , c 2 , . . . , c p if σ ∈ C s ( R ) with s = 2 p -1 and σ (1) (0) , σ (2) (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 ˆ w 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 = 2 sn ( 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 that

Let us also remark that it is possible to use the original dynamics ˆ w 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 ˆ w 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].

4.2 Proof Sketch for Theorem 4.3

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 matrix

where ˆ σ ( ξ ) = (1 + ξ ) n , ( a 1 , . . . , a p ) ∼ U ([ -1 , 1] p ) , 1 ≤ i 1 , i 2 ≤ ( n + p p ) , 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 that

where λ 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 ˆ w ( a i , t ) and ˆ u ( a 1 , . . . , a p , t ) = 1 p ∑ p i =1 ˆ w ( 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 ) with

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

which is true if ˆ u i ( a 1 , . . . , a p , t ) , 1 ≤ i ≤ p are R -algebraically independent polynomials in a 1 , a 2 , . . . , a p , where ˆ u i is the i -th coefficient of ˆ u 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 ˆ u i . With Assumption 4.1, ˆ w 1 ( a, t ) , ˆ w 2 ( a, t ) , . . . , ˆ w 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 ˆ w i , whose constant term is zero since initialization in training is set as ρ 0 = ρ a × δ R d , to the algebraic independence of ˆ u 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 , a 2 , . . . , a p ] are R -algebraically independent if and only if

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 matrix X ( q ) ∈ R ( n + p p ) × ( n + p p ) via

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 form

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 = ˆ u ( a i 2 , t ) leads to

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

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.

5 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.

Acknowledgments and Disclosure of Funding

The work of R. Ge is supported by NSF Award DMS-2031849 and CCF-1845171 (CAREER). We thank Joan Bruna for helpful comments and discussion.

References

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

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

Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems , 34:26989-27002, 2021.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International conference on machine learning , pages 242-252. PMLR, 2019.

Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.00193 , 2019.

Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows SGD to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405.15459 , 2024.

Malte Beecken, Johannes Mittmann, and Nitin Saxena. Algebraic independence and blackbox identity testing. Information and Computation , 222:2-19, 2013.

Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning Gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborova, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. In Forty-first International Conference on Machine Learning , 2024.

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2018.

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

Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems , 29, 2016.

Kenji Kawaguchi and Yoshua Bengio. Depth with nonlinearity creates no bad local minima in resnets. Neural Networks , 118:167-174, 2019.

Jason D Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns lowdimensional polynomials with SGD near the information-theoretic limit. arXiv preprint arXiv:2406.01581 , 2024.

Haihao Lu and Kenji Kawaguchi. Depth creates no bad local minima. arXiv preprint arXiv:1702.08580 , 2017.

Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond NTK with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. Advances in Neural Information Processing Systems , 36, 2024.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665E7671, 2018.

Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv:1902.02880 , 2019.

Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics , 75(9):1889-1935, 2022.

Itay Safran and Ohad Shamir. Spurious local minima are common in two-layer ReLU neural networks. In International conference on machine learning , pages 4433-4441. PMLR, 2018.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics , 80(2):725-752, 2020.

Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond. Advances in Neural Information Processing Systems , 36, 2024.

Alain-Sol Sznitman. Topics in propagation of chaos. Lecture notes in mathematics , pages 165-251, 1991.

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

A Proofs for Section 3

This section collects the proofs of Theorem 3.5, Theorem 3.6, and Theorem 3.4.

A.1 Proof of Theorem 3.5

Existence 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

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) that and hence that

We also have that

if w S = 0 . In addition, it holds also for w S = 0 that

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

where the last equality holds by applying the test function η ( · , 0 , · ) ∈ C ∞ c ( R × S ⊥ × (0 , + ∞ )) to (3.7). The proof is hence completed.

A.2 Proof of Theorem 3.6

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 by

which implies that

Therefore, one has | a t | ≤ | a 0 | + K ξ K σ E z [ | h ∗ ( z ) | 2 ] 1 / 2 t , 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

Then it holds for any 0 ≤ t ≤ T that

and where 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

Define γ t ∈ Γ( ρ t , ˜ ρ t ) and ∆( t ) as in Lemma A.2. According (A.6) and the Grönwall's inequality, it holds that

which combined with (A.5) yields that

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 has

Proof. It can be computed that

which implies that

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

and hence that

We then bound the two terms above as follows:

and

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 ) 2 by combining all estimations above.

We then head into the proof of (A.6), for which we need the particle dynamics and

The distance between the particle dynamics can be decomposed as

and

Therefore, it can be computed that

where we used Remark 3.2, and that

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 that

where

This proves (A.6).

A.3 Proof of Theorem 3.4

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 that

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:

where z S = P V S z , z ⊥ S = z -z S , and h ∗ 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 that

Then applying Corollary A.3, we can conclude for any T > 0 that

where C s > 0 is the constant in Theorem 3.6 depending only on p , h ∗ , K σ , K ξ , K ρ , and T . Therefore, we can obtain (3.3) and (3.4) immediately.

B Further Discussion and Characterization of the Reflective Property and Theorem 3.4

B.1 Equivalence between (3.2) and isoLeap ( h ∗ ) ≥ 2

We prove the following equivalence where isoLeap ( h ∗ ) is the isotropic leap complexity defined in [2, Appendix B.2].

Proposition B.1. For any polynomial h ∗ : V → R , then it satisfies (3.2) with some nontrivial subspace S ⊂ V if and only if isoLeap ( h ∗ ) ≥ 2 .

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 ∗ as

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 } that

where

Thus, by the orthogonality of Hermite polynomials, we have for i ∈ { m 1 +1 , . . . , m } that

The above also holds for i ∈ { 1 , 2 , . . . , m 1 } by the definition of J m 1 . Therefore, we can conclude that

i.e., (3.2) holds, for S = span { e j : j ∈ J m 1 } . Moreover, S is nontrivial since (B.2) implies J m 1 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 p 1 , 0 , . . . , 0) : z 1 , . . . , z p 1 ∈ R } with 1 ≤ p 1 ≤ p . We still consider the Hermite decomposition as in (B.1) and rewrite it as

where c ′ 1 , c ′ 2 , . . . , c ′ m 2 are nonzero coefficients, α ′ 1 , α ′ 2 , . . . , α ′ m 2 are pairwise distinct elements in N p 1 , and h 1 , h 2 , . . . , h m 2 are nonzero polynomials defined on R p -p ′ . Then it follows from (3.2) that

̸

for any j ∈ { 1 , 2 , . . . , p 1 } and z p 1 +1 , . . . , z p ∈ R , which implies that

̸

Therefore, the leap complexity of h ∗ is at least 2 with respect to this basis, which leads to isoLeap ( h ∗ ) ≥ 2 .

B.2 Discretization Results Implied by Theorem 3.4

We 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

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 that E x ∼N (0 ,I d ) [ | f ∗ ( x ) -˜ f ∗ ( x ) | 2 ] < δ . The associated mean-field dynamics is

where the learning rate ξ ( t ) = diag ( ξ a ( t ) , ξ w ( t ) I d ) and the initialization ρ 0 are shared with (2.4), and

The corresponding SGD is given by

where N is the number of neurons and γ ( k ) = diag ( γ ( k ) a , γ ( k ) w 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 constant C b f such that

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 :

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.

C Proofs for Section 4

C.1 Approximation of w ( a, t ) by Polynomials

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:

is the residual. The first observation is that

By Theorem 3.5, we have that ρ t is supported in { ( a, w ) : w ⊥ V = 0 } , and hence that

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 } .

with m l = σ ( l ) (0) and then an approximated solution to (C.1) (by polynomial expansion with high-order terms omitted) can be written as

where Q ( t ) is given by Q (0) = 0 and the following dynamics:

Let us remark that even if every single ˜ w i ( a, t ) depends on the basis { e 1 , e 2 , . . . , e p } , the linear combination

is basis-independent. To see this, let Q j ( t ) ∈ R d be the j -th column of Q ( t ) and it holds that

The distance between w i ( a, t ) and ˜ w i ( a, t ) for i ∈ I can be bounded as follows.

Proposition C.1. Suppose that Assumption 4.2 holds. Then

We need the following two lemmas to prove Proposition C.1.

where

Lemma C.2. For any i ∈ { 1 , 2 . . . , p } and j ∈ { 1 , 2 , . . . , L } , it holds that

Proof. The non-increasing property of the energy functional implies that

Then (C.6) can be proved by induction. For j = 1 , it follows from the boundedness of

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 that

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 } that

Proof. There exists a constant C > 0 and a open subset A ⊂ { w ∈ R d : w ⊥ V = 0 } containing 0 , such that

Thus, we have

as long as w ( a, t ) does not leave A . This implies (C.7).

Now we can proceed to prove Proposition C.1.

Proof of Proposition C.1. Set ˜ w ⊥ V ( a, t ) = 0 . It can be estimated for any i ∈ { 1 , 2 , . . . , p } that

which combined with (C.2) yields that

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 replacing g ( x, t ) = f ∗ ( x ) -f NN ( x ; ρ t ) by f ∗ ( x ) in (C.3):

where ˆ Q ( t ) is given by ˆ Q (0) = 0 and

Similar to ˜ w ( a, t ) , the linear combination ˆ w ( a, t ) defined as

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 ˆ q i,j depending only on h ∗ and σ , such that

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

We need the following lemma for the proof of Proposition C.5.

Lemma C.6. It holds that

which implies (C.12).

Proof of Proposition C.5. We prove (C.11) by introduction on j . For j = 1 , one has

which leads to | Q i,e l ( t ) -ˆ Q i,e l ( t ) | = O ( t 2 ) . Then we consider 2 ≤ j ≤ L and assume that | Q i,j ′ ( t ) -ˆ Q i,j ′ ( t ) | = O ( t j ′ +1 ) holds for 1 ≤ j ′ < j . It can be estimated that

where we used Lemma C.6. Then one can conclude (C.11).

Proof. Noticing that f NN ( x ; ρ 0 ) = 0 by the symmetry of ρ a , one has that

and hence by Lemma C.3 and w ⊥ V ( a, t ) = 0 that

C.2 From Linear Independence to Algebraic Independence

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 ,

̸

Lemma C.8. If v 1 , v 2 , . . . , v p ∈ R [ a ] are R -linearly independent, then

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 least n 1 + n 2 + · · · + n p . So it suffices to prove that the sum of monomials with degree being n 1 + n 2 + · · · + n p is nonzero, which is true since

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 that v ′ 1 , v ′ 2 , . . . , v ′ p ∈ R [ a ] are also R -linearly independent. Noticing that

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.

C.3 Proofs of Proposition 4.4 and Theorem 4.3

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 ˆ q i,j , 1 ≤ i ≤ p, 1 ≤ j ≤ L in (C.10) satisfies

̸

Proof. Let ˆ w V ( t ) be the dynamics defined in (4.1). It can be seen that the Taylor's expansion of ˆ w V ( t ) at t = 0 up to s -th order is given by

where ˆ Q i,j and ˆ q i,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).

where

In the rest of this subsection, we will always denote s = ( s 1 , s 2 , . . . , s P ) and

for t ∈ [0 , T ] , where w ( a, t ) , ˜ w ( a, t ) , and ˆ w ( 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 follows

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 p ) × ( n + p p ) matrices

and

where a = ( a 1 , a 2 , . . . , a ( n + p p ) ) and a i ∈ R p for i = 1 , 2 , . . . , ( n + p p ) .

where j = ( j 1 , j 2 , . . . , j ( n + p p ) ) with j i ∈ N p for i = 1 , 2 , . . . , ( n + p p ) , ˆ h 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.

Proof. The result follows directly from Proposition C.4 and Proposition C.5.

Lemma C.12. For any m ∈ N , one has that

where 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

that

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

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

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

which combined with Lemma C.12 implies that

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

By the R -algebraic independence of v 1 , v 2 , . . . , v p , one must have

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 ˆ h j be the coefficient of det ( ˆ M ( a , t )) in (C.15) . Then there exists some ˆ j = ( ˆ j 1 , ˆ j 2 , . . . , ˆ j ( n + p p ) ) with ∥ ˆ j ∥ 1 ≤ sn ( n + p p ) , such that ˆ h ˆ j = 0 .

̸

Proof. Let X ( q ) ∈ R ( n + p p ) × ( n + p p ) be defined as

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

̸

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. Denote

̸

and

Then we have that

where ˆ u s ( a l , t ) = (ˆ u 2 ,s 1 ( a l , t ) , . . . , ˆ u p,s p ( a l , t )) and ˆ u i,s i ( a 1 , . . . , a p , t ) = 1 p ∑ p k =1 ˆ q i,s i t s i a s i k collects the leading order terms of ˆ u i ( a 1 , . . . , a p , t ) . According to Assumption C.9, Theorem 4.5, and Lemma 4.7, we have that

̸

which provides at least one non-zero term in det( ˆ M ( a , t )) whose degree in t is

This completes the proof.

̸

Proof of Proposition 4.4. According to Lemma C.14, there exist some ˆ j with ∥ ˆ j ∥ 1 ≤ sn ( n + p p ) such that ˆ h ˆ j = 0 . According to Lemma C.11 and Lemma 103 in [1], it holds that

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

̸

By Jensen's inequality, one has that

Since entries of ˜ M ( a, t ) are all O (1) by Lemma C.2, which implies the boundedness of the eigenvalues λ i ( ˜ M ( a , t ) ˜ M ( a , t ) ⊤ ) , i = 1 , 2 , . . . , ( n + p p ) , one has that

Combining (C.18), (C.19), and (C.20), one can conclude that

where we used ∥ ˆ j ∥ ≤ N ( n + p p ) max i ∈ I s i and only considered t ≤ 1 . By Proposition C.1, we have that

Then we can obtain (4.3) as we set L = 2 sn ( n + p p ) .

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

which implies that

by Gronwall's inequality. Then we obtain the desired exponential decay property by noticing that E ( ρ T ) = 1 2 E z [ ∥ h ∗ ( z ) ∥ 2 ] .

D Further Discussion and Characterization of Assumption 4.1 and Theorem 4.3

D.1 Verification of Assumption 4.1

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.

It 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 ∗ ( z ) = c 1 z 1 + c 2 z 1 z 2 + · · · + c p z 1 z 2 · · · z p if c 1 , c 2 , . . . , c p are nonzero. More generally, we have the following.

Proposition D.1. Let { α 1 , α 2 , . . . , α m } be a set of pairwise distinct elements in N p that contains (1 , 0 , . . . , 0) , (1 , 1 , 0 , . . . , 0) , . . . , (1 , 1 , . . . , 1) . If σ ∈ C s ( R ) with s = 2 p -1 and σ (1) (0) , σ (2) (0) , . . . , σ ( p ) (0) are all nonzero, then

satisfies Assumption 4.1 with s unless ( c 1 , c 2 , . . . , c m ) is in some measure-zero subset of R m with respect to the Lebesgue measure.

with

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 by

for 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

for i = 1 , 2 , . . . , N and k = T/ϵ, T/ϵ +1 , . . .

Suppose that assumptions made in Theorem 4.3 hold and fix T ′ > T > 0 . Using similar analysis as in Appendix A.2, one can conclude that for any δ > 0 , there exist dimension-free constants C b f and C b σ such that

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 :

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.

̸

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 h ∗ ( 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.

D.2 Discretization Results Implied by Theorem 4.3

The 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

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The theoretical results claimed in the abstract and the introduction are established in Section 3 and Section 4.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations are discussed in the last paragraph of Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The assumptions and the proof sketches are included in Section 3 and Section 4. The complete proofs are in the appendices.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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.

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. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., 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.

Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: This paper is purely theoretical and does not include experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. 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.

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

8. 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).

9. 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).

10. 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).

11. 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.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: This paper is purely theoretical and does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper is purely theoretical and does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

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.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\nThe main contribution of this submission is the introduction of the reflective property of the target function which entails whether the parameter distribution evolved under the mean-field gradient flow can achieve nontrivial overlap with a given subspace in finite time. This \"basis-free\" extension of MSP to Gaussian data is an interesting result.\n\nWEAKNESSES:\nI have reviewed this paper in another venue. The authors addressed some of my concerns in the revision.  \nI shall reiterate the remaining weaknesses from my previous review. \n\n- 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. \n\n- 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. \n\nAfter 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. \n\n- 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. \n\n- 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.  \n(Lee et al. 2024) *Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit*.  \n(Arnaboldi et al. 2024) *Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions*.\n\nQUESTIONS:\nSee Weaknesses above.\n\nLIMITATIONS:\nSee Weaknesses above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nAs 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.\n\nSTRENGTHS:\nOverall, the work effectively extends the previously studied merged staircase properties to a basis-independent setting. The manuscript is well-written and contributes valuable insights to the multi-index setting, which is an active area of research in the neural network theory community.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nThe 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?\n\nAdditionally, 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?\n\nFurthermore, 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?\n\n[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.\n\nLIMITATIONS:\nSee the Weaknesses part.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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:\n1)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.\n2) They show that a slightly stronger condition is sufficient for some modified dynamics to decay to $0$ exponentially fast.\n\nSTRENGTHS:\n- The paper is reasonably well written, and provides ample details on their technical results.\n\n- The paper extends the results from coordinate sparse on hypercube, to subspace-sparse on Gaussian data (also known as multi-index models in the statistics literature).\n\n- The proof techniques are involved and require careful arguments, in particular for the algebraic independence.\n\nWEAKNESSES:\n- 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).\n\n- 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)\n\n- There are several unnatural simplifications made in the algorithm. \nFirst 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?\nSecond, the activation is changed between the two phases of the algorithm.\nThird, 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)\n\n- 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.\n\nQUESTIONS:\n- 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?\n\n- 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$.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

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

STRENGTHS:
1. This paper is well written, with detailed explanations, proofs, and algorithms.

2. It gives both necessary and sufficient conditions for when neural networks can learn these functions well, which helps understand the limits of what neural networks can do.

3. The results don't depend on the dimension of the input, which means they work even for very high-dimensional data.

WEAKNESSES:
No obvious weakness.

QUESTIONS:
No questions

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13