The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks

Introduction

Major research activity has recently been devoted to understanding what function classes can be learned by SGD on neural networks. Two extremal cases are well understood. On one extreme, neural networks can be parametrized to collapse under SGD to linear models, for which a clear picture has been drawn [JGH18, LL18, Du+18, Du+19, ALS19, ALL19, Aro+19, Zou+20, OS20]. On the other extreme, neural networks with zero parametrization constraint (besides polynomial size) have been shown to be able to emulate essentially any efficient learning algorithm [AS20, Abb+21] albeit with non-regularHere we refer to ‘regular’ for architectures used in tangent kernel results or more generally architectures used in neural network applications. architectures. So both of these extremes admit a fairly complete characterization. However, none of these seem to capture the right behavior behind deep learning, or more specifically, behind non-linear but regular networks. Such networks are known to go beyond linear learning [Bac17, Gho+21a, DM20, Ref+21, AL19, Gho+19, YS19, AL20, LMZ20] (even though the NTK can be competitive on several instances [Gei+20]), and seem to exploit structural properties of the target functions in order to build efficiently their features.

Can we thus characterize learning in the non-linear regime for regular networks? Various important results have been developed in this direction, we focus here on the most relevant to us. [CB18, MMN18, RV18, SS20] show that for a certain scaling at initialization, the SGD dynamics on large-width neural networks concentrates on a fully non-linear dynamics, the mean-field dynamics, described by a Wasserstein gradient flow, contrasting with the linear dynamics of the NTK regime [JGH18]. In [AL19, AL20a], the power of deep networks is demonstrated by showing how SGD and quadratic activations can efficiently learn a non-trivial teacher class hierarchically, with the notion of backward feature correction [AL20a].

However, no tight necessary and sufficient characterization of what functions are learnable emerges from these works. The difficulty being that tight necessity results are difficult to obtain in such a setting since SQ-like arguments [Blu+94, FGV17, Kea98, BKW03, Fel16, Yan05, Fel+17, SVW15, AS20, Abb+21, Goe+20] are not expected to be tight (besides for the extreme case of unconstrained networks [AS20, Abb+21]), and sufficiency results are significantly more difficult to obtain due to the more complex (non-linear) dynamics of SGD training.

Is there hope to characterize tight necessary and sufficient conditions for function classes to be learnable by standard SGD on standard neural networks?

As a first attempt in that direction, we focus in this paper on a natural setting: learning sparse functions on the $d$-dimensional hypercube, i.e., functions that depend on a small latent (unknown) subset of coordinates of the input. We further restrict the optimization regime considered to two layers neural networks trained by one-pass batch-SGD in the mean-field regime. This allows us to study a regime of optimization that goes beyond the linear regime while averaging out some of the complexity of studying non-linear SGD.

The motivation for the setting of learning sparse function is three-fold: (1) Linear (fixed features) methods do not adapt to latent sparsity, and suffer from the curse of dimensionality [Bac17]. (2) On the contrary, [Bac17, Sch20] shows that neural networks can overcome this curse and learn sparse functions sample-efficiently. However, these works do not provide tractable algorithms and the question of when SGD-trained neural networks can adapt to sparsity remains largely open. (3) Some sparse functions, such as monomials, are known to be much harder to learn than others from SQ-like lower bounds [Kea98, Blu+94, Abb+21], and we expect SGD to inherit some of this complex behavior. Therefore, the problem of learning sparse functions presents a clear-cut separation between fixed-feature and feature learning methods, and can help understand the limits of SGD-training on neural networks.

Both of these functions are $3$-sparse, and both present the same tight sample-complexity lower bound of $\Omega(d^{3})$ to be learned by any linear method. However, are these functions equivalent for SGD-trained neural networks? If not, can we obtain a fine-grained analysis that separates them?

In this paper, we introduce the following notion: we say that a latent function $h_{*}$ is strongly SGD-learnable in $O(d)$-scaling, if $O(d)$ samples are enough to learn $f_{*}({\bm{x}})=h_{*}({\bm{z}})$ for arbitrary latent subspace ${\bm{z}}$ and dimension $d$, using batch-SGD on a two-layer neural network in the mean-field regime. The main contribution of this paper is then to characterize with a necessary and nearly sufficient condition the class of functions that are strongly SGD-learnable in $O(d)$-scaling. This is achieved with the merged-staircase property (MSP), stating that the non-zero Fourier coefficients of $g$ can be ordered as subsets $\{S_{1},\ldots,S_{r}\}$ such that for any $i\in[r]$,

For instance, $h_{2}({\bm{z}})=z_{1}+z_{1}z_{2}+z_{1}z_{2}z_{3}$ has Fourier coefficients (monomials) that can be ordered as $\{\{1\},\{1,2\},\{1,2,3\}\}$, and each new set is incremented only by one element at each time. So $h_{2}$ satisfies the MSP (or is an MSP function with a slight abuse of terminology) and so is the function $z_{1}+z_{1}z_{2}+z_{2}z_{3}+z_{1}z_{2}z_{3}$. However, the function $h_{1}({\bm{z}})=z_{1}z_{2}z_{3}$ directly makes a leap to a degree-3 Fourier coefficient and does not therefore satisfy the MSP. Our main results thus imply that $h_{2}$ can be learned with $O(d)$ samples in this regime, but not $h_{1}$. The near sufficiency part in our result stands for the fact that the sufficiency result is proved for “generic” merged-staircase functions, i.e., excluding a measure zero subclass. This ‘genericity’ is in fact needed, as we provide degenerate examples in Section 4 for which the strong SGD-learnability in $O(d)$-scaling is indeed not achievable.

The terminology MSP comes from the fact that this condition generalizes the basic staircase property introduced in [Abb+21a], which only encompasses nested chains of coefficients with $|S_{i}\setminus S_{i-1}|=1$, such as the vanilla staircase function (e.g., $h_{2}$) and slight generalizations with multiple chains. In [Abb+21a] it is shown that staircase functions are learnable by neural nets that are deep bur sparse, and with an unconventional gradient-based training algorithm (see Section 1.2 for further discussion). Further [Abb+21a] does not provide necessary conditions for learning, nor fine-grained complexity guarantees (beyond ‘polynomial’).

Finally, while strong SGD-learnability is defined for a fixed latent function and fixed $P$, the number of samples required to fit MSP functions remains polynomial in $d$ for $P$ growing sufficiently slowly in $d$. This is of interest because in this regime, we can show that the considered functions are not learnable by any linear methods with any sample complexity (or feature space dimension) that is polynomial (using contribution (4) below). Thus the merged-staircase functions of such degree are efficiently learnable by SGD on networks of depth two but not by linear methods.

We now formally define the Merged-Staircase Property. Let us call any ${\mathcal{S}}\subseteq 2^{[P]}$ a set structure.

We say that ${\mathcal{S}}=\{S_{1},\ldots,S_{m}\}\subseteq 2^{[P]}$ is a Merged-Staircase Property (MSP) set structure if the sets can be ordered so that for each $i\in[m]$, $|S_{i}\setminus(S_{i}\cap(\cup_{i^{\prime}<i}S_{i^{\prime}}))|\leq 1$.

Let ${\mathcal{S}}\subset 2^{[P]}$ be the non-zero Fourier coefficients of $h_{*}$, i.e., $\hat{h}_{*}(S)\neq 0$ iff $S\in{\mathcal{S}}$. We say that $h_{*}$ satisfies the merged-staircase property (MSP) if ${\mathcal{S}}$ is a MSP set structure.

In words, $h_{*}$ satisfies the MSP if the monomials in its Fourier decomposition can be ordered sequentially such that the supports of the monomials grow by at most one at a time. Examples of MSP functions include vanilla staircases (i.e., $z_{1}+z_{1}z_{2}+\dots+\prod_{i=1}^{P}z_{i}$), $z_{1}+z_{1}z_{2}+z_{2}z_{3}+z_{3}z_{4}$, or $z_{1}+z_{2}+z_{3}+z_{4}+z_{1}z_{2}z_{3}z_{4}$, but not $z_{1}+z_{1}z_{2}z_{3}+z_{1}z_{2}z_{3}z_{4}$, $z_{1}+z_{1}z_{2}+z_{3}z_{4}$, or $z_{1}z_{2}z_{3}$. We briefly summarize our results here:

We show that for non-MSP $h_{*}$, the associated dimension-free dynamics stays bounded away from . From the previous equivalence, we deduce that MSP is necessary for a function to be strongly $O(d)$-SGD-learnable.

We first show that vanilla staircases are strongly $O(d)$-SGD-learnable for smooth activation functions as long as $\sigma^{(r)}(0)\neq 0$ for $r=0,\ldots,P$.

For general MSP functions, however, some symmetric MSP functions have degenerate dynamics and are not strongly $O(d)$-SGD-learnable (see Section 4). We show instead that MSP $h_{*}$ are almost surely strongly $O(d)$-SGD-learnable. I.e., the degenerate examples are a measure-zero set. This is proved for generic degree-$L$ polynomial activations, and we explain how one can extend this result to generic smooth activations in the appendix.

One can take MSP functions (e.g., vanilla staircases) with $P$ slowly growing with $d$ so that the overall sample complexity of the above neural network results stay as $d^{O(1)}$, while we show that any linear method requires a sample complexity of $d^{\omega_{d}(1)}$.

These main results are further achieved with several side results of independent interest: (i) The approximation of the standard mean-field dynamics by the dimension-free dynamics, valid for $P$-sparse target functions and $d\gg P$. We provide a new version of the non-asymptotic bounds from [MMN18, MMM19], which now compares SGD with this dimension-independent dynamics; (ii) A new proof technique to study layer-wise SGD dynamics which reduces the proof of global convergence to a polynomial identity testing problem, i.e., whether a certain polynomial is non-identically zero; (iii) An improvement of prior dimension lower-bounds for linear (kernel) methods [Hsu+21, Hsu, KMS20] that is tighter for function classes that are non-almost orthogonal (such as staircase functions, allowing for contribution (4) above).

The rest of the paper is organized as follows. The next section overviews related work. Section 2 provides a formal definition of strong SGD-learnability in $O(d)$-scaling. In Section 3, we introduce the dimension-free dynamics and the equivalence with strong $O(d)$-SGD-learnability. The MSP necessary condition is then derived as a direct consequence of this equivalence. In Section 4, we provide our sufficient conditions for strong $O(d)$-SGD-learnability. In Section 5, we discuss how this implies a separation with linear methods.

2 Further related literature

[Abb+21a] introduces a class of staircase functions, which our merged-staircase function class generalizes. They show that staircase functions are learnable by some neural nets with a gradient-based training algorithm. However, the approach remains non-standard: (i) the network’s layers are sparse in order to guide the construction of the features; (ii) a coordinate descent variant of SGD is used that differs from the classical SGD algorithm. Further, the analysis is carried in the ‘polynomial scaling lens’ rather than a finer sample complexity, and no necessity results are derived. In contrast, we provide here both a necessary and nearly sufficient characterization for SGD-learning on a two-layer neural networks in the fine-grained $O(d)$-scaling.

Multiple works have used mean-field (also called distributional) dynamics to approximate the SGD trajectory. Relevant to us is [CB20] which showed that neural networks trained in the mean-field regime converge to a max-margin classifier that is independent of the dimension for latent low-dimensional target functions. However, these works do not provide quantitative results in terms of sample-complexity. A notable exception is [MMN18] which studies classifying anisotropic gaussians: they show that the mean-field dynamics concentrates on a simplified low-dimensional dynamics as $d\to\infty$. However, this simplification is due to rotational invariance of the problem and not the sparsity of the target function.

In approximation theory, it has been understood for a long time that sparse functions are naturally well approximated by neural networks [Bar93]. Recent work [Bac17, Sch20, Gho+21a, CMM21] have shown that neural networks can learn sparse functions more sample-efficiently than linear methods. However, these works do not provide tractable algorithms.

Finally, a string of works [YS19, AL19, AL20, LMZ20, DM20, Ref+21, Gho+21, Gho+21a, Mal+21, Kar+21, SA20] have shown separation results between gradient-trained neural networks and fixed-features models. We refer to Appendix B of [Mal+21] for a detailed survey. In particular, [DM20] considers the learning of parity functions, with a modified input distribution that gives correlation to the response and allows for domain extraction; it also uses the population dynamics (infinite samples). In [MS20], the learning of Boolean circuits of logarithmic depth is considered via neural networks with layer-wise gradient descent, but with an architecture that is required to match the Boolean circuit being learned, i.e., not with a ‘regular’ or ‘blackbox’ architecture. Lastly, [Bas+19, Cao+21] show that during training, SGD on 2-layer $\mathsf{ReLU}$ networks learns faster the lower frequency components of a target function, in similar spirit to low degree monomials, but the approach relies on the linear regime rather than the non-linear regime of interest here, and suffers from an exponential dependency on the degree.

Strong SGD-learnability in O​(d)𝑂𝑑O(d)-scaling

We first consider a general definition for a class of sparse functions to be learnable. We take $\{P(d)\}_{d\geq 1}$ a sequence of integers (here, we allow the sparsity parameter $P$ to grow with $d$) and consider a general class of functions defined as ${\mathcal{H}}=\{{\mathcal{H}}_{P(d)}\}_{d\geq 1}$ with ${\mathcal{H}}_{P(d)}\subseteq L^{2}(\{+1,-1\}^{P(d)})$.

This definition covers many scenarios that occur in practice where the practitioner is allowed to tune the hyperparameters of the dynamics. While this choice leaves the question of tractability open, we note that the requirement that learnability must hold uniformly over all possible latent subspaces excludes many irregular scenarios. Furthermore, the next definition will require strong regularity on the hyperparameters, and our sufficiency results will hold for simple choices of hyperparameters.

In order to introduce strong SGD-learnability, we will restrict the previous definition in three major ways: (1) we consider a fixed dimension $P$ and function ${\mathcal{H}}=\{h_{*}\}$; (2) we consider the scalingExtending our results to $\alpha>1$, and establishing how this relates to the ‘leap’ in the staircase definition (i.e., how can one jump monomial degrees) is a natural future direction to this work. of $\alpha=1$; (3) we restrain the hyperparameters to be in either two regimes (i) small batch size $b=o(d)$ and step size $\eta=o(1)$ trained for $\Theta(1/\eta)$ steps (“continuous”); and (ii) large batch size $b=\Theta(d)$ and step size $\eta=\Theta(1)$ trained for a total number of $\Theta(1)$ steps (“discrete”). For the sake of presentation, we will only present the continuous regime in the main text and defer the presentation of the discrete regime to Appendix C. We will assume that the hyperparameters obey the following for some constant $K$ (independent of $d$):

(One-pass) We have fresh samples at each steps, meaning $\{({\bm{x}}_{k_{i}},y_{k_{i}})\}_{k\geq 0,i\in[b]}$ are iid. Furthermore, the response variable is bounded $|y|\leq K$.

(Boundedness and lipschitzness of hyperparameters) There exists a constant $\eta>0$ such that $\eta_{k}^{a},\eta_{k}^{w}\leq\eta K$, $|\eta_{k+1}^{a}-\eta_{k}^{a}|\leq\eta^{2}K$ and $|\eta_{k+1}^{w}-\eta_{k}^{w}|\leq\eta^{2}K$. Furthermore, $\lambda^{a},\lambda^{w}\leq K$.

Conditions ${\rm A}0$-${\rm A}3$ guarantee that as long as $d,n,N$ are taken sufficiently large, there exists a continuous mean-field dynamics that well-approximates batch-SGD up to (continuous) time $T$ depending on $\varepsilon,h_{*}$. An analogous statement is true for strong-SGD-learnability in the “discrete regime”, except convergence is to a family of limiting discrete-time dynamics (deferred to Appendix C). This allows us to get a necessary condition for strong-learnability by studying the limiting dynamics (see next section).

Finally, we note that for any degree-$k$ sparse function $h_{*}$, any linear method (e.g., arbitrary kernel or random feature methods) will require $\Omega(d^{k})$ samples to fit functions $f_{*}({\bm{x}})=h_{*}({\bm{z}})$ uniformly well over all latent subspaces ${\bm{z}}={\bm{x}}_{{\mathcal{I}}}$ (see Section 5 for a formal statement). As emphasized in the introduction, this bound is not adaptive to the sparsity parameter $P$. In particular, any non-linear $h_{*}$ that is strongly $O(d)$-SGD-learnable provides a separation result between SGD-trained neural networks and linear methods.

Continuous dimension-free dynamics and necessary condition

For simplicity, the results in this section are stated in the ‘continuous regime’ of strong SGD-learnability. Discrete versions can be found, with little modification, in Appendix C.

$\xi^{a},\xi^{w}$ are bounded Lipschitz $\|\xi^{a}\|_{\infty},\|\xi^{w}\|_{\infty},\|\xi^{a}\|_{\text{Lip}},\|\xi^{w}\|_{\text{Lip}}\leq K$ and $\lambda^{a},\lambda^{w}\leq K$.

Note that for any $(\eta_{k}^{a},\eta_{k}^{w})_{k\geq 0}$ obeying ${\rm A}3$, there exists functions $\xi^{a},\xi^{w}$ such that ${\rm A}3^{\prime}$ holds with same constant $K$. Conversely, any $\eta$ discretization of $\xi^{a},\xi^{w}$ obeys ${\rm A}3$ with constants $\eta,K$.

and we denote with a slight abuse of notation, $\hat{f}_{{\sf NN}}({\bm{z}};\rho_{t}):=\hat{f}_{{\sf NN}}({\bm{x}};\rho_{t})$.

We see that $\hat{f}_{{\sf NN}}(\cdot;\overline{\rho}_{t})$ can be seen as a two layer neural network in dimension $P$, with adaptive Gaussian smoothing. Taking $d\to\infty$ with $P$ fixed, $(a^{0},{\bm{u}}^{0},\|{\bm{v}}^{0}\|_{2})$ (with distribution $\rho_{0}$ satisfying ${\rm A}2$) converges in distribution to $(\overline{a}^{0},\overline{{\bm{u}}}^{0},\overline{s}^{0})\sim\overline{\rho}_{0}$ with $\overline{a}^{0}\sim\mu_{a}$, $\overline{{\bm{u}}}^{0}={\bm{0}}$ and $\overline{s}^{0}=m_{2}^{w}$, and the dynamics (MF-PDE) simplifies into the following dimension-free dynamics

The following theorem provides a non-asymptotic bound between the (bSGD) solution $\hat{f}_{{\sf NN}}(\cdot;\hat{\rho}^{(N)}_{t})$ and the (DF-PDE) solution $\hat{f}_{{\sf NN}}(\cdot;\overline{\rho}_{t})$:

Assume conditions ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$ hold, and let $T\geq 1$. There exist constants $K_{0}$ and $K_{1}$ depending only on the constants in ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$ (in particular, independent of $d,P,T$), such that for any $b\leq d$, $N\leq e^{d}$, $\eta\leq e^{-K_{0}T^{3}}b/(d+\log(N))$, we have

The proof of Thm. 3.1 can be found in App. B.2.1. An extension of the results in [MMM19] bounds the difference between (bSGD) and (MF-PDE) dynamics, and then we use a propagation-of-chaos argument to bound the distance between the (MF-PDE) and (DF-PDE) solutions.

From Theorem 3.1, (DF-PDE) is a good approximation of (bSGD) as long as $d,N,1/\eta$ are taken sufficiently large while keeping $T=\eta n/b$ bounded. This leads to the equivalence described in the introduction (the proof can be found in Appendix B.2.1):

For generic activation, we have $\inf_{\overline{\rho}}R(\overline{\rho})=0$. Hence, Theorem 3.2 states that $h_{*}$ is strongly $O(d)$-SGD-learnable if and only if the global minimizer is dynamically reachable by a gradient flow initialized at $\overline{\rho}_{0}=\mu_{a}\otimes\delta_{\overline{{\bm{u}}}^{0}={\bm{0}}}\otimes\delta_{\overline{s}^{0}=c}$. See Appendix A for additional discussions and numerical illustrations. In Figure 1, we plotted a comparison between (bSGD) and (DF-PDE) for $h_{*}({\bm{z}})=z_{1}+z_{1}z_{2}+z_{1}z_{2}z_{3}+z_{1}z_{2}z_{3}z_{4}$ and shifted sigmoid activation $\sigma(x)=(1+e^{-x+0.5})^{-1}$. We fix $d=N=100$, $b=150$, $\lambda^{a}=\lambda^{w}=0$, $\eta^{a}_{k}=\eta^{w}_{k}=1/2$, $\mu_{a}={\rm Unif}([+1,-1])$ and $\mu_{w}={\sf N}(0,1)$. Let us emphasize a few prominent features of this plot: 1) The (DF-PDE) approximation tracks well (bSGD) until convergence even for moderate $d,N,b/\eta$, despite a convergence with nontrivial structure. 2) The monomials are picked up sequentially with increasing degree, which agrees with the intuition that lower-degree monomials guide SGD to learn higher degree monomials. 3) (DF-PDE) reaches a global minimum, which by Theorem 3.2 implies that $h_{*}$ is strongly SGD-learnable in $O(d)$-scaling.

We can show that the (DF-PDE) dynamics with $h_{*}$ without MSP cannot reach arbitrarily small test error when initialized with $\overline{\rho}_{0}$. By Theorem 3.2, this implies that MSP is necessary for strong SGD-learnability in $O(d)$-scaling.

for some $r(u_{23}^{t})\in[-2u_{23}^{t},2u_{23}^{t}]$ using the mean value theorem. Recalling that $u_{23}^{0}=0$, we deduce that $u^{t}_{23}=\overline{u}_{2}^{t}=\overline{u}_{3}^{t}=0$.

Sufficient conditions for strong SGD-learnability

In the previous section, we saw that having MSP is necessary for strong $O(d)$-SGD-learnability. Is the converse true? Is any MSP function strongly SGD-learnable in the $O(d)$-scaling?

To bypass this difficulty, we prove a learnability result that holds for “generic” MSP functions – i.e., that holds almost surely over a random choice of non-zero Fourier coefficients. Formally, for any set structure ${\mathcal{S}}=\{S_{1},\ldots,S_{m}\}\subseteq 2^{[P]}$, let us define a measure over functions that have those Fourier coefficients.

Our main sufficiency result shows that the degenerate cases are a measure-zero set. In this sense, there are very few bad examples, and so MSP structure is “nearly” sufficient for strong $O(d)$-SGD-learnability.

For any MSP set structure ${\mathcal{S}}\subseteq 2^{[P]}$, $h_{*}$ is strongly $O(d)$-SGD-learnable almost surely with respect to $\mu_{{\mathcal{S}}}$, using activation function $\sigma(x)=(1+x)^{L}$ where $L=2^{8P}$.Technically speaking, for the strong SGD-learnability definition we cannot take $\sigma(x)=(1+x)^{L}$ as it is not bounded. However, we take an activation function that equals $(1+x)^{L}$ on the interval $(-1,1)$ and is bounded elsewhere.

The converse to this result is implied by the necessity result of the previous section, which states that for any $h_{*}$ with non-zero Fourier coefficients (set structure) ${\mathcal{S}}$ that is not MSP, $h_{*}$ is not strongly $O(d)$-SGD-learnable. While we prove Theorem 4.2 for a particular activation, we note that the proof implies that the same is true for any degree-$L$ polynomial activation almost surely over its $(L+1)$-coefficients (see Theorem E.5 in Appendix E). In Appendix F we show how this result extends to generic smooth (non-polynomial) activations as long as a certain polynomial is not identically for a given set structure (which we show with a small technical caveat).

In the special case of functions with “vanilla staircase” structure we do not need a genericity assumption, and we require weaker assumptions on the activation function.

The proofs for Theorems 4.2 and 4.3 follow a similar approach. From the equivalence stated in Theorem 3.2, it is sufficient to display, for each $\varepsilon>0$, hyperparameters such that the (DF-PDE) dynamics reaches $\varepsilon$-risk. We choose $\lambda^{a}=\lambda^{w}=0$ (no regularization) and initialization $\mu_{a}={\rm Unif}([-1,+1])$ and $\mu_{w}=\delta_{0}$ (this choice simplifies the analysis as $\overline{s}^{t}=\overline{s}^{0}=0$). We split the learning in two phases: in Phase 1, we train the first layer weights $\overline{{\bm{u}}}^{t}$ for time $t\in[0,T_{1}]$ while keeping $\overline{a}^{t}=\overline{a}^{0}$ fixed, and in Phase 2, we train the second layer weights $\overline{a}^{t}$ for time $t\in[T_{1},T_{2}]$ while keeping $\overline{{\bm{u}}}^{t}=\overline{{\bm{u}}}^{T_{1}}$ fixed.

The goal of the analysis in Phase 1 is therefore to prove this lower bound on the eigenvalues of the kernel matrix. Phase 1 corresponds to a nonlinear dynamics, and is a priori unclear how to analyze. In the case of vanilla staircases, we show that it is enough to track the leading order in $t$ for each coordinates $(\overline{u}_{i}^{t})_{i\in[P]}$ and take $T_{1}$ small enough: the lower bound on the eigenvalues of ${\bm{K}}^{T_{1}}$ follows from a simple algebraic fact (see Appendix D for the detailed proof). For general MSP set structure, it is not enough to only track the leading order term. We show instead that it is enough to lower bound a kernel matrix $\hat{{\bm{K}}}^{T_{1}}$ obtained from a simplified dynamics $\hat{{\bm{u}}}^{t}(\overline{a}^{0})$. The weights $\hat{u}_{i}^{t}(\overline{a}^{0})$ can be written in terms of polynomials in $\overline{a}^{0}$, $(\alpha_{S})_{S\in{\mathcal{S}}}$ and $(\sigma^{(r)}(0))_{r=0,\ldots,L}$, with coefficients defined explicitly by a recurrence relation and only depending on the set structure ${\mathcal{S}}$. Using algebraic facts about the linear independence of large powers of polynomials and plugging $\overline{a}^{0}\sim{\rm Unif}()$, we show that $\det(\hat{{\bm{K}}}^{T_{1}})>0$ for $T_{1}$ small enough, by anti-concentration of polynomials, which implies the lower bound on $\lambda_{\min}({\bm{K}}^{T_{1}})$. The proof can be found in Appendix E. We present a more general argument in Appendix F for non-polynomial activations, where we show instead that the polynomials are not identically by adding a perturbation to the activation.

Separation with linear methods

It is well known that any SQ algorithm (adaptive or not) with polynomially-many queries and polynomial query precision cannot learn the class of degree $P$ monomials if $P$ grows with the input dimension $d$ [Kea98, Blu+94], and likewise no linear method with polynomially many features or samples can learn this function class [Hsu+21, Hsu, KMS20]. But the implication is not obvious for staircase functions of growing degree. These functions contain monomials of growing degree, but the hierarchical structure could potentially allow a sequential learning of these monomials; in fact, staircases of degree $P$ are efficiently SQ learnable with adaptive queries as one can sequentially query the monomials of increasing degree (making at most $d$ queries per degree, e.g., at most $d\log d$ queries for $P=\log d$ vanilla staircases). We thus need a lower-bound on linear methods that goes beyond general SQ lower-bounds, which we obtain by separating the analysis using subspace projections.

We will further denote $q=\dim({\mathcal{H}})$. Popular examples include random feature models ($q$ is equal to the number of random features) and kernel methods ($q=\infty$ typically). While the optimization problem (72) is over a (potentially) infinite dimensional space ${\mathcal{H}}$, it is an easy exercise to verify that ${\hat{\bm{a}}}\in{\rm span}\{\psi({\bm{x}}_{i}):i\in[n]\}$ which has dimension bounded by $\min(n,q)$.

Let $\Omega\subseteq L^{2}({\mathcal{X}})$ a linear subspace. Let ${\mathcal{F}}_{M}=\{f_{1},\ldots,f_{M}\}\subset L^{2}({\mathcal{X}})$ such that $\|{\mathsf{P}}_{\Omega}f_{j}\|_{L^{2}}^{2}=1-\kappa$ and $\|{\mathsf{P}}_{\Omega}^{\perp}f_{j}\|_{L^{2}}^{2}=\kappa$ for all $f_{j}\in{\mathcal{F}}_{M}$. For any linear method, if $\overline{R}_{n}({\mathcal{F}}_{M})\leq 1-\eta$, then we must have

For any linear method, if $\overline{R}_{n}({\mathcal{F}}_{{\text{Poly-}k}})\leq 1/2\cdot(1-\eta)$ then we must have $\min(n,q)\geq\frac{\eta}{m}{{d}\choose{k}}$. Similarly, if $\overline{R}_{n}({\mathcal{F}}_{{\text{Str-}P}})\leq 1-\eta$ then we must have $\min(n,q)\geq\frac{\eta}{2}{{d}\choose{\lfloor\frac{\eta P}{2}\rfloor}}$.

Note that kernel and random features methods achieve the lower bound for ${\mathcal{F}}_{{\text{Poly-}k}}$ [Gho+21, MMM21]. Comparing Proposition 5.1 with the result of Section 4, we get the following separation results between SGD-trained neural networks and linear methods:

SGD on two-layer neural networks outperforms linear methods almost surely on non-linear MSP functions ($n=O(d)$ versus $n=\Omega_{d}(d^{k})$ for degree-$k$ MSP).

We obtain from Proposition 5.2 that for any $P=\omega_{d}(1)$, linear methods must have $\min(n,q)=d^{\omega_{d}(1)}$ to learn the vanilla staircase of degree $P$, while Theorem 4.3 can still guarantee a sample complexity of $d^{O(1)}$ for $P$ growing slowly enough with $d$.

Conclusion and future directions

In this paper, we considered learning sparse functions in arbitrarily large ambient dimension, using two-layer neural networks trained by batch-SGD in the mean-field regime. We proved that the merged-staircase property is a necessary and nearly-sufficient condition for such functions to be learnable on such models in $O(d)$ sample-complexity. The near sufficiency part, which excludes a measure-zero subset, is unavoidable as there exist symmetric MSP functions with degenerate dynamics that are not strongly SGD-learnable in $O(d)$-scaling. This provides a regime where one can achieve a tight characterization of functions that are learnable by regular SGD on regular neural networks, while going beyond the linear regime.

Finally, it is natural to seek counterparts of the results in this work and counterparts of the staircase notions for other Hilbert spaces than the one considered here.

Acknowledgements

We thank Guy Bresler, Dheeraj Nagaraj, and Nati Srebro for stimulating discussions. We thank the Simons Foundations and the NSF for supporting us through the Collaboration on the Theoretical Foundations of Deep Learning (deepfoundations.ai). This work was done (in part) while E.B. and T.M. were visiting the Simons Institute for the Theory of Computing and while E.B. was visiting the Bernoulli Center at EPFL.

References

Appendix A Additional numerical simulations

In this Appendix, we provide further background and numerical illustrations on the strong $O(d)$-SGD learning setting, merged-staircase functions and the dimension-free dynamics.

While global convergence proofs are challenging, the (DF-PDE) dynamics is a low dimensional problem and can be efficiently solved numerically. In the rest of this section, we provide a few numerical simulations to illustrate phenomena alluded to in the main text. We will fix the activation to be a shifted sigmoid $\sigma(x)=(1+e^{-x+1})^{-1}$, and choose learning schedules $\xi^{a}(t)=\xi^{w}(t)=1$, zero regularization parameters $\lambda^{a}=\lambda^{w}=0$, and initialization $\mu_{a}={\rm Unif}([+1,-1])$ and $\overline{s}^{0}=1$. In Figure 2, we consider four MSP functions and plot the evolution of their Fourier coefficients during the (DF-PDE) dynamics. In particular, the two top row examples, $h_{1}({\bm{z}})=z_{1}+z_{1}z_{2}+z_{3}+z_{1}z_{2}z_{3}z_{4}$ and $h_{2}({\bm{z}})=z_{1}+z_{1}z_{2}+z_{2}z_{3}+z_{3}z_{4}+z_{1}z_{2}z_{3}z_{4}$, converge to the global minimum and are therefore strongly $O(d)$-SGD-learnable. The bottom row examples, $h_{3}({\bm{z}})=z_{1}+z_{1}z_{2}+z_{3}+z_{3}z_{4}$ and $h_{4}({\bm{z}})=z_{1}+z_{2}+z_{3}+z_{1}z_{2}z_{3}$, do not converge and have risks bounded away from . Functions $h_{3}$ and $h_{4}$ are two examples of $G$-invariant MSP functions.

In this paper, we only prove that the set of MSP functions that are not strongly $O(d)$-SGD-learnable is of Lebesgue measure . We do not characterize this set beyond this and do not prove that $G$-invariant MSP functions coincides with this set (in particular, we do not show that $G$-invariant MSP functions are the only functions that might not be strongly $O(d)$-SGD-learnable).

We conclude this section with a final comment about the necessity condition of MSP, which holds only when considering arbitrarily large $d$.

The proof that non-MSP functions are not strongly $O(d)$-SGD-learnable relies on the fact that, when $d$ goes to infinity, the initialization $u_{i}^{0}\to 0$ for $i\in[P]$. However, for $d$ fixed, $u_{i}^{0}\approx 1/\sqrt{d}$ and waiting sufficiently long, one-pass (bSGD) escapes the neighborhood of the subspace $u_{i}^{0}=0$. In this case, the time to escape the subspace has to grow with $d$, and we are not in the $O(d)$-scaling anymore (indeed $n=Tb/\eta\approx Td$ for one pass (bSGD)). In Figure 4, we consider the same experimental setting as Figure 1 but with $h_{*}$ missing one (left) or two (right) stairs. We see that (DF-PDE) remains trapped in the saddle-space, while one-pass (bSGD) escapes around $n\approx d^{2}$ and $n\approx d^{2.7}$ respectively. This agrees with the intuition that staircases with larger leaps are harder to learn with SGD.

Appendix B Proofs for continuous mean-field and dimension-free dynamics

In this appendix, we provide proofs and discussions for the results presented in Section 3, which corresponds to the ‘continuous-time regime’ of strong $O(d)$-SGD-learnability. A discrete version of these results and proofs are presented in Appendix C and require little modifications.

Throughout this section, we will denote by $K$ a constant that depends only on the constants in Assumptions ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$ (in particular, $K$ is independent of $d,P,T$). The value of this constant is allowed to change from line to line.

Here, we provide more details and intuition on how to derive the equations of the dimension-free dynamics (DF-PDE). We report to Section B.2.1 a rigorous proof of Theorem 3.1, which shows a non-asymptotic bound between (bSGD) and (DF-PDE) dynamics.

First, by Assumption ${\rm A}2$, the coordinates of ${\bm{w}}^{0}$ are iid and symmetric and therefore

By symmetry of (MF-PDE), the following lemma shows that the neural network stays independent of the uninformative part ${\bm{r}}$ of the input during the whole trajectory.

The solution $(\rho_{t})_{t\geq 0}$ of (MF-PDE) with initialization $\rho_{0}$ satisfying ${\rm A}2$, obeys:

where we used in the third line that $\psi(\varphi_{{\bm{r}}}({\bm{\theta}});\rho_{t})=\psi({\bm{\theta}};\rho_{t}^{\#{\bm{r}}})$. Hence $\rho_{t}^{\#{\bm{r}}}$ is the solution of the (MF-PDE) dynamics with initialization $\rho^{\#{\bm{r}}}_{0}=\rho_{0}$. Hence by uniqueness of the solution, we deduce that $\rho^{\#{\bm{r}}}_{t}=\rho_{t}$ for any $t\geq 0$. ∎

The evolution equations of $(a^{t},{\bm{u}}^{t},\|{\bm{v}}^{t}\|_{2})$ associated to the (MF-PDE) dynamics are given by

where we used that $\frac{{\rm d}}{{\rm d}t}\|{\bm{v}}^{t}\|_{2}=\frac{1}{\|{\bm{v}}^{t}\|_{2}}\langle{\bm{v}}^{t},\frac{{\rm d}}{{\rm d}t}{\bm{v}}^{t}\rangle$ to write the last equation.

Equivalently, this PDE corresponds to the gradient flow in the Wasserstein space (with $\overline{{\bm{H}}}(t)$ rescaling) over the regularized risk functional:

B.2 Proof of the results in Section 3

In this section, we gather the proofs for the results on the dimension free dynamics and the necessity condition. The longer and more technical arguments are deferred to Sections B.3 and B.4.

We use the mean-field dynamics (MF-PDE) as an intermediary dynamics for the bound. Theorem 3.1 is a direct consequence of the following two bounds:

Assume conditions ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$, and let $T\geq 1$. There exists constants $K_{0}$ and $K_{1}$ depending only on the constants in ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$ (in particular, independent of $d,P,T$), such that for any \eta\leq e^{-K_{0}T^{3}}\big{[}\{b/(d+\log(N))\}\wedge 1\big{]}, we have

This proposition follows from a straightforward extension of [MMM19] to batch-SGD and anisotropic step sizes, and can be found in Section B.3. In particular, Proposition B.2 implies that, if we consider $T,K=O_{d}(1)$, then $N=\Omega_{d}(1)$ and $1/\eta=\Omega_{d}(d/b)$ are sufficient for the mean-field PDE to be an accurate approximation of batch-SGD up to time $T$ (recall that $T=\eta n/b$ by one-pass assumption and therefore $n=O_{d}(d)$).

Assume conditions ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$, and let $T\geq 1$. There exists a constant $K_{1}$ depending only on the constants in ${\rm A}0$-${\rm A}2$,${\rm A}3^{\prime}$ (in particular, independent of $d,P,T$), such that

The proof of Theorem B.3 can be found in Section B.4.

B.2.2 Proof of Theorem 3.2

Conversely, assume that $h_{*}$ is strongly SGD-learnable in $O(d)$-scaling. Let $(b,\lambda^{a},\lambda^{w},\{\eta_{k}^{a},\eta_{k}^{w}\}_{k\in[0,k_{0}]})$ be the hyperparameters that satisfy strong learnability for $\varepsilon/4$: in particular, $k_{0}=n/b$ and $R({\bm{\Theta}}^{k_{0}})\leq\varepsilon/4$ with probability at least $9/10$. Take $T=n\eta/b\leq T(h_{*},\varepsilon/4)$ and let $\eta^{a},\eta^{w}$ be piecewise linear functions such that $\eta^{a}(\eta k)=\eta_{k}^{a}/\eta$ and $\eta^{w}(\eta k)=\eta_{k}^{w}/\eta$. Consider $(\overline{\rho}_{t})_{t\geq 0}$ the solution of (DF-PDE) with $\eta^{a},\eta^{w},\lambda^{a},\lambda^{w}$. From Theorem 3.1, there exists constants $K_{0}$ and $K_{1}$ that only depend on $\eta^{a},\eta^{w},\lambda^{a},\lambda^{w}$ through the constants in assumption ${\rm A}3$, such that

with probability at least $1-1/N$. We can therefore take $d,N,n$ sufficiently big such that the right-hand side is less than $\sqrt{\varepsilon}/2$. On the intersection of this event and the event $R({\bm{\Theta}}^{k_{0}})\leq\varepsilon/4$ (which happens with positive probability), we have

B.2.3 Proof of Theorem 3.3

This lower bound does not depend on the details of the dynamics (parameters $\xi^{a},\xi^{w},\lambda^{a},\lambda^{w}$, activation and initialization $\mu_{a}$,$m_{2}^{w}$). Let ${\bm{z}}_{i,+}$ and ${\bm{z}}_{i,-}$ denote the vector ${\bm{z}}\in\{+1,-1\}^{P}$, with $z_{i}=+1$ and $z_{i}=-1$ respectively, and note that by Lemma B.6, $|\overline{a}^{t}|\leq K(1+t)$. Using Assumption ${\rm A}0$, we have by integrating out $z_{i}$:

Similarly, for $i\in\Omega$ and $S\in\overline{{\mathcal{S}}}_{*}$ (in particular, since $i\not\in S$)

while if $S\in{\mathcal{S}}_{*}\setminus\overline{{\mathcal{S}}}_{*}$, then there exists $j\in\Omega\cap S$ with $j\neq i$, hence

Denoting $m_{\Omega}^{t}=\max_{i\in\Omega}|\overline{u}_{i}^{t}|$ (recall $m_{\Omega}^{0}=0$), we conclude that for any $i\in\Omega$:

and therefore $m_{\Omega}^{t}=0$ during the whole dynamics.

B.3 Proof of Proposition B.2

The proof is an application of an extension of Theorem 1.(B) in [MMM19] to batch-SGD and anisotropic step sizes. This extension is straightforward and we simply list below the two main differences with the proof in Appendix C of [MMM19]:

Recall that we defined the regularized risk $E(\rho):=\frac{1}{2}R(\rho)+\frac{1}{2}\int{\bm{\theta}}^{\mathsf{T}}{\bm{\Lambda}}{\bm{\theta}}\rho({\rm d}{\bm{\theta}})$. We have

where $\psi({\bm{\theta}}^{t};\rho_{t})$ is defined in Eq. (MF-PDE) and we denoted $\|{\bm{v}}\|_{{\bm{A}}}=\|{\bm{A}}^{1/2}{\bm{v}}\|_{2}$. We conclude that $E(\rho_{t})$ is nonincreasing. The rest of the proof only uses that ${\bm{H}}(t)$ verifies $\|{\bm{H}}\|_{\infty},\|{\bm{H}}\|_{\text{Lip}}\leq K$.

The concentration between the batch-SGD and gradient descent (Appendix C.5 in [MMM19]) uses that there is an extra $1/b$ factor in the sub-Gaussian constant.

The proof of Proposition B.2 simply amounts to checking that our setting (with Assumptions ${\rm A}0$-${\rm A}2,{\rm A}3^{\prime}$) falls under the general framework of Theorem 1.(B) in [MMM19].

We conclude that $\|\nabla v({\bm{w}})\|_{2},\|\nabla u({\bm{w}}_{1},{\bm{w}}_{2})\|_{2},\|\nabla^{2}v({\bm{w}})\|_{\rm op},\|\nabla^{2}u({\bm{w}}_{1},{\bm{w}}_{2})\|_{{\rm op}}\leq K$, and assumption ${\rm A3}$ in [MMM19] is verified. ∎

B.4 Proof of Theorem B.3: bound between (MF-PDE) and (DF-PDE) dynamics

We will assume throughout this section that the assumptions and the setting of Theorem B.3 hold. In particular, we will use Assumptions ${\rm A}0$-${\rm A}2,{\rm A}3^{\prime}$ without mention when clear from context. For clarity, we will write the proof in the case $\xi^{a}(t)=\xi^{w}(t)=1$ and $\lambda^{a}=\lambda^{w}=0$. The general case follows easily, using $\|\xi^{a}\|_{\infty},\|\xi^{w}\|_{\infty},\lambda^{a},\lambda^{w}\leq K$ by Assumption ${\rm A}3^{\prime}$.

The value of the integrand in Eq. (11) only depends on $\{(a^{t},{\bm{u}}^{t},\langle{\bm{v}}^{t},{\bm{r}}\rangle)\}_{t\geq 0}$ and $\{(\overline{a}^{t},\overline{{\bm{u}}}^{t},\overline{s}^{t}G)\}_{t\geq 0}$ with ${\bm{r}}\sim{\rm Unif}(\{-1,+1\}^{d-P})$ and $G\sim{\sf N}(0,1)$ independent of $\{({\bm{\theta}}^{t},\overline{{\bm{\theta}}}^{t})\}_{t\geq 0}$. Conditioning on $({\bm{\theta}}^{t},\overline{{\bm{\theta}}}^{t})$, we consider the $1$-Wasserstein distance

where we recall that $W_{1}$ is defined by

The following lemma bounds the right hand-side through the value of ${\bm{v}}^{t}$ at initialization:

Consider the same setting and assumptions as Theorem B.3. There exists a constant $K$ independent of $d,P$ and depending only on the Assumptions ${\rm A}0$-${\rm A}2,{\rm A}3^{\prime}$ such that for any $T\geq 0$,

where we expanded the expectation on $r_{i}$ in the second line and used the mean value theorem, and used Eq. (40) in Lemma B.6 in the last line. We deduce that

Using Lemma B.4 in the bound (13) yields (conditional on ${\bm{\theta}}^{0}$):

Using Eq. (15) and the coupling described above, we will bound (11). Introduce the random quantity

We will show the following technical bounds:

Consider the same setting and assumptions as Theorem B.3. There exists a constant $K$ independent of $d,P$ and depending only on the Assumptions ${\rm A}0$-${\rm A}2,{\rm A}3^{\prime}$ such that for any $T\geq 0$,

From this lemma, we can now complete the proof of Theorem B.3:

From Gronwall’s lemma applied to Eq. (52) in Lemma B.5, we have

Injecting this bound in Eq. (18) concludes the proof. ∎

The proof consists in carefully bounding the evolution of the distance between the parameters in the two dynamics.

Step 1. Bound on $\|\hat{f}_{{\sf NN}}(\cdot;\rho_{t})-\hat{f}_{{\sf NN}}(\cdot;\overline{\rho}_{t})\|_{L^{2}}$.

We can bound the difference between the two functions with

while we use Eq. (20) for the second term

where we used Eq. (40) in Lemma B.6 and Eq. (15) with $q=1$.

Combining bounds (21) and (22) and by Jensen’s inequality,

Step 2. Bound on $(a^{t}-\overline{a}^{t})^{2}$.

Noting that $\hat{g}({\bm{z}};\rho_{t})-\hat{g}({\bm{z}};\overline{\rho}_{t})=\hat{f}_{{\sf NN}}({\bm{z}};\overline{\rho}_{t})-\hat{f}_{{\sf NN}}({\bm{z}};\rho_{t})$, the first term can be bounded as in step 1 by

For the second term, we use Eq. (41) in Lemma B.6 and the decomposition (20):

Combining Eqs. (23) and (24) and applying Cauchy-Schwarz inequality yield

where we used the bound (13) on $W_{1}$ and Eq. (15) in Lemma B.8 in the last line. We deduce that for $t\in[0,T]$,

where we used that $a^{0}=\overline{a}^{0}$ at initialization.

Step 3. Bound on $\|{\bm{u}}^{t}-\overline{{\bm{u}}}^{t}\|_{2}^{2}$.

Combining inequalities (27) and (28) yields

where we again used the bound (15). We deduce that for $t\in[0,T]$,

Step 4. Bound on \big{|}\overline{s}^{t}-\|{\bm{v}}^{t}\|_{2}\big{|}^{2}.

First, notice that we have the following simple upper bounds on the evolution of $\|{\bm{v}}^{t}\|_{2}$ and $\overline{s}^{t}$:

Furthermore, we have by Gaussian integration by part

Similarly, we have by expanding the expectation over the $r_{i}$’s and using the mean-value theorem:

We can now bound the evolution in time of $(\|{\bm{v}}^{t}\|_{2}-\overline{s}^{t})$. Using the expressions in Eqs. (32) and (33), we decompose

These four quantities can be bounded as previously:

For the last term, we use Eq. (34) and that $|\xi_{i}|\leq|v_{i}^{t}|$:

Combining Eqs. (35) and (36) and applying Cauchy-Schwarz inequality yield

We can now combine inequalities (26), (30) and (30) to get

B.5 Auxiliary lemmas

Denote the residuals of the dynamics $\hat{g}({\bm{z}};\rho_{t})=h_{*}({\bm{z}})-\hat{f}_{{\sf NN}}({\bm{z}};\rho_{t})$ and $\hat{g}({\bm{z}};\overline{\rho}_{t})=h_{*}({\bm{z}})-\hat{f}_{{\sf NN}}({\bm{z}};\overline{\rho}_{t})$. By the properties of gradient flows, the risks

By definition $(\rho_{t})_{t\geq 0}$ and $(\overline{\rho}_{t})_{t\geq 0}$ are the solutions of a gradient flow:

and similarly for $\|\hat{g}(\cdot;\overline{\rho}_{t})\|_{L^{2}}$.

A similar result holds for $\overline{a}^{t}$. Finally,

This is a simple application of Stein’s method. Consider $f$ twice differentiable such that $\|f\|_{\infty}\leq 1$, $\|f^{\prime}\|_{\infty}\leq\sqrt{2/\pi}$ and $\|f^{\prime\prime}\|_{\infty}\leq 2$. Introduce $S_{i}=S-v_{n}^{-1/2}X_{i}=v_{n}^{-1/2}\sum_{j\neq i}X_{j}$. By expanding, we get

where we used Jensen’s inequality in the last line. Combining bounds Eqs. (44) and (45) in the identity (43) yields

By sub-Gaussianity, there exists a universal constant $c>0$ such that

Consider t_{c}=\Big{(}\tau^{2}\frac{\log(2d)}{\kappa_{q}cd}\Big{)}^{q/2} with $\kappa_{q}=2^{2(q-2)/q^{2}\vee 0}$, such that $2de^{-cd\kappa_{q}t_{c}^{2/q}/\tau^{2}}=1$. Then, we have the following upper bound:

Appendix C Strong SGD-learnability in the discrete-time regime

In this appendix, we define strong SGD-learnability in the discrete-time regime, i.e., for large batch size $b$ and large $\eta$. We keep the same assumptions ${\rm A}0$-${\rm A}2$, and replace Assumption ${\rm A}3$ by

(Boundedness of hyperparameters) We have $\eta_{k}^{a},\eta_{k}^{w}\leq K$ and $\lambda^{a},\lambda^{w}\leq K$.

While the continuous-time regime requires step size $\eta$ to be small enough compared to $n/b$, the discrete-time regime requires the batch size $b$ to be big enough compared to $cn$ for $c\ll 1$ (recall $b\leq n$ by one-pass assumption) in the discrete regime.

Again, conditions ${\rm A}0$-${\rm A}2,{\rm D}3$ guarantee that as long as $d,n,N$ are taken sufficiently large, there exist a discrete mean-field dynamics that well-approximates batch-SGD up to a constant number of steps that depends on $\varepsilon,h_{*}$.

where $(a^{k},{\bm{w}}^{k})\sim\rho_{k}$.

We have the new non-asymptotic bound between the (bSGD) and (d-DF-PDE) dynamics, analogous to Theorem 3.1, but with a worse dependency on the number of iterations.

Assume conditions ${\rm A}0$-${\rm A}2$,${\rm D}3$ hold, and let $k_{0}\geq 0$. There exists a constant $K$ depending only on the constants in ${\rm A}0$-${\rm A}2$,${\rm D}3$ (in particular, independent of $d,P,k_{0}$), such that

From there, it is straightforward, following the same arguments as for Theorems 3.2 and 3.3, to get the equivalence of strong $O(d)$-SGD-learnability in the discrete-time regime and global convergence of the discrete (d-DF-PDE) dynamics, and the MSP necessary condition:

C.2 Proof of Theorem C.2

The proof relies on first comparing the (bSGD) dynamics to the discrete mean-field dynamics (d-MF-PDE), using an extension of the results in [MMM19] to the discrete (d-DF-PDE) dynamics (see Appendix H.1).

The proof of this proposition follows from applying Proposition H.1, with the assumptions already verified in Appendix B.3.

The proof of Theorem C.2 then follows by combining the above result with the following bound between the discrete mean-field dynamics (d-MF-PDE) and the discrete dimension-free dynamics (d-DF-PDE):

Assume conditions ${\rm A}0$-${\rm A}2$,${\rm D}3$, and let $k_{0}\geq 0$. There exists a constant $K$ depending only on the constants in ${\rm A}0$-${\rm A}2$,${\rm D}3$ (in particular, independent of $d,P,T$), such that

The proof follows similarly to the proof in the continuous case (see Section B.4) and we will simply highlight the differences. First, by the same argument as in the proof of Proposition H.1, we replace the bounds from Lemma B.6 by

The proof follows by using discrete Grönwall lemma in Lemma C.7 stated in the next section, which is the analogous of Lemma B.5 in discrete time. ∎

C.3 Auxiliary lemma

The proof proceeds similarly to the proof of Lemma B.5 in Section B.4.1, where we use discrete Grönwall instead. Step 1 to Step 3 are very similar, using that

Note that $\langle{\bm{v}}_{k}/\|{\bm{v}}^{k}\|_{2},{\bm{T}}^{k}\rangle=M_{k}$ as defined in Eq. (34) and we can use the bound in Eq. (36):

We can further bound $({\rm II})$ using the same decomposition as in Eq. (35):

Appendix D Vanilla staircase functions are strongly O​(d)𝑂𝑑O(d)-SGD-learnable: Proof of Theorem 4.3

We start by providing the proof that vanilla staircases are strongly $O(d)$-SGD-learnabile, as described in Theorem 4.3. This proof will outline the main ideas behind our global convergence results, without the technical complexity of dealing with general MSP set structure.

We will assume the following hold for the activation $\sigma$:

In particular, this assumption implies that we have the following polynomial approximations of $\sigma$ and $\sigma^{\prime}$ around : for any $x\in(-\eta,\eta)$,

Recall from the equivalence with (DF-PDE) (Theorem 3.2) that it is sufficient to show for any $\varepsilon>0$, there exist hyperparameters satisfying ${\rm A}0$-${\rm A}2$, ${\rm A}3^{\prime}$ such that (DF-PDE) dynamics reaches $\varepsilon$-risk. We consider the following hyperparameters:

We do not regularize, i.e., $\lambda^{a}=\lambda^{w}=0$.

We initialize the first layer to deterministically weights, and the second layer to uniform random weights. I.e., we take $\mu_{a}={\rm Unif}([+1,-1])$ and $\mu_{W}=\delta_{0}$. Although initializing the first layer to 0 may at first glance seem restrictive, there turns out to be enough randomness in the initialization of the second layer to ensure that the neural network learns. For the dimension-free dynamics, this corresponds to taking $(\overline{a}^{0},\overline{{\bm{u}}}^{0},\overline{s}^{0})\sim\overline{\rho}_{0}$ with $\overline{a}^{0}\sim{\rm Unif}([+1,-1])$, $\overline{{\bm{u}}}^{0}={\bm{0}}$ and $\overline{s}^{0}=0$. In particular, $\overline{s}^{t}=0$ during the whole dynamics, which allows for a simpler analysis.

Our learning rate schedule has two phases:

We train the first layer weights $\overline{{\bm{u}}}^{t}$ while keeping the second layer weights fixed $\overline{a}^{t}=\overline{a}^{0}$. We set $\xi^{a}(t)=0$ and $\xi^{w}(t)=1$ for $t\in[0,T_{1}]$.

We train the second layer weights $\overline{a}^{t}$ while keeping the first layer weights fixed at $\overline{{\bm{u}}}^{t}=\overline{{\bm{u}}}^{T_{1}}$. We set $\xi^{a}(t)=1$ and $\xi^{w}(t)=0$ for $t\in[T_{1},T_{2}]$.

We restate the sufficient condition in the case of the vanilla staircase.

D.1 Outline of the proof

The proof analyzes Phase 1 and Phase 2 of training separately.

In this phase, we train the first layer, which has nonlinear dynamics, and so it is a priori unclear how to analyze. Nevertheless, since $h_{*}$ is specially structured, the structure in the weights during training is particularly simple and it is enough to track the smallest order terms in the weights.

Specifically, in Proposition D.5 (see next section), we prove that there exist constants $c,C>0$ such that for all $t\leq c$ and $k\in[P]$, we have $|u_{k}^{t}(a)-\hat{u}_{k}^{t}(a)|\leq Ct^{2^{k-1}+1}$, where

Denote $\nu_{k}(t)=2^{1-2^{k-1}}t^{2^{k-1}}\prod_{i\in[k]}(m_{i}\alpha_{\{1,\ldots,i\}})^{2^{\max((k-1-i),0)}}$ such that $u_{k}^{t}(a)=\nu_{k}(t)(a^{2^{k-1}}+O(t))$.

In this phase, we train the second layer, and the training has linear dynamics. Denote $g_{t}({\bm{z}}):=h_{*}({\bm{z}})-\hat{f}_{{\sf NN}}({\bm{z}};\overline{\rho}_{t})$ the residual function at time $t$. During this phase, we have the following evolution on the risk:

(This is indeed the kernel, since at the end of Phase 1, the distribution $\overline{\rho}_{T_{1}}$ of the parameters is given by $(a,{\bm{u}}^{T_{1}}(a))$ with $a\sim\mu_{a}$, and the first-layer weights are kept constant during Phase 2.)

It only remains to lower-bound $\lambda_{\min}({\bm{K}}^{T_{1}})$. For this we use the structure on ${\bm{u}}^{T_{1}}$ that we prove holds in Phase 1. For all $S\subseteq[P]$, denote

From Lemma D.4 (see next section), there exists a constant $C>0$ depending only on $P$ (and independent of $T_{1}$) such that for any $S\subseteq[P]$,

Denote $D_{S}=m_{|S|}\nu_{S}(T_{1})$ and ${\bm{D}}={\rm diag}((D_{S})_{S\subseteq[P]})$. We have

where $\|{\bm{\Delta}}\|_{{\rm op}}\leq CT_{1}P$.

Note that $\beta(S)$ takes value $\{0,\ldots,2^{P}-1\}$, and ${\bm{M}}$ is the Gram matrix of the monomials $[1,X,\ldots,X^{2^{P}-1}]$ in $L^{2}([+1,-1],{\rm Unif})$, which are linearly independent. We deduce that $\lambda_{\min}({\bm{M}})$ is bounded away from (independent of $T_{1}$). We can therefore take $T_{1}\leq\lambda_{\min}({\bm{M}})/(2P)$, so that $\lambda_{\min}({\bm{M}}+{\bm{\Delta}})\geq\lambda_{\min}({\bm{M}})/2$, and $\lambda_{\min}({\bm{K}}^{T_{1}})\geq\{\min_{S}D_{S}^{2}\}\lambda_{\min}({\bm{M}})/2>0$.

First, we have the following simple bound on $\|{\bm{u}}^{t}\|_{1}$:

There exists a constant $C$ depending on $K,P$ such that $\|{\bm{u}}^{t}\|_{1}\leq Ct$.

By Assumptions ${\rm A}0$ and ${\rm A}_{1}$, we have $\|\sigma^{\prime}\|_{\infty},\|h_{*}\|_{\infty}\leq K$ and $\|\hat{f}_{{\sf NN}}(\cdot;\overline{\rho}_{t})\|_{\infty}\leq\int|a|\|\sigma\|_{\infty}{\rm d}\mu_{a}\leq K$. Combining these bounds, we get for $t\leq T_{1}$:

and therefore $\frac{{\rm d}}{{\rm d}t}\|{\bm{u}}^{t}\|_{1}\leq PK$. Recalling, ${\bm{u}}^{0}={\bm{0}}$, we conclude $\|{\bm{u}}^{t}\|_{1}\leq KPt$. ∎

The following lemma give the leading order in $t$ approximation of the Fourier coefficients of $\sigma(\langle{\bm{u}}^{t},{\bm{z}}\rangle)$:

There exists a constant $c>0$ that depend on $\eta,K,P$ such that for any $t\leq c$, $S\subseteq[P]$ and $i\in\{0,1\}$,