Source-linked AI summary
The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
Emmanuel Abbe, Enric Boix-Adsera, Theodor Misiakiewicz
TL;DR
The paper addresses the lack of a tight characterization for learning sparse latent functions with nonlinear but regular neural networks. It analyzes two-layer mean-field SGD through dimension-free dynamics and proves that the merged-staircase property is necessary and nearly sufficient, while linear methods require superpolynomial sample complexity in a corresponding regime.
Problem
No tight necessary-and-sufficient characterization was known for functions learnable by standard SGD on nonlinear regular networks, especially for sparse functions in large ambient dimensions.
Method
The paper studies sparse binary-input functions with latent dimension P using two-layer batch-SGD networks in the mean-field regime and introduces equivalent dimension-free dynamics.
Results
The merged-staircase property is necessary and nearly sufficient for strong O(d)-SGD-learnability, with generic MSP functions learnable and non-MSP functions excluded.
Takeaways & Limitations
The results provide a regime where nonlinear SGD on depth-two networks learns functions efficiently that linear methods cannot learn with polynomial sample complexity.
Takeaways & Limitations
The MSP characterization is tied to fixed or slowly growing P, O(d) sample complexity, and two-layer networks, and excludes a measure-zero degenerate subset.
Abstract
from arXiv · showhide
It is currently known how to characterize functions that neural networks can learn with SGD for two extremal parameterizations: neural networks in the linear regime, and neural networks with no structural constraints. However, for the main parametrization of interest (non-linear but regular networks) no tight characterization has yet been achieved, despite significant developments. We take a step in this direction by considering depth-2 neural networks trained by SGD in the mean-field regime. We consider functions on binary inputs that depend on a latent low-dimensional subspace (i.e., small number of coordinates). This regime is of interest since it is poorly understood how neural networks routinely tackle high-dimensional datasets and adapt to latent low-dimensional structure without suffering from the curse of dimensionality. Accordingly, we study SGD-learnability with $O(d)$ sample complexity in a large ambient dimension $d$. Our main results characterize a hierarchical property, the "merged-staircase property", that is both necessary and nearly sufficient for learning in this setting. We further show that non-linear training is necessary: for this class of functions, linear methods on any feature map (e.g., the NTK) are not capable of learning efficiently. The key tools are a new "dimension-free" dynamics approximation result that applies to functions defined on a latent space of low-dimension, a proof of global convergence based on polynomial identity testing, and an improvement of lower bounds against linear methods for non-almost orthogonal functions.
1 Introduction
The paper studies when standard SGD on two-layer mean-field networks can learn sparse functions in arbitrarily large ambient dimensions, where existing characterizations do not cover nonlinear regular networks. It introduces the merged-staircase property as a necessary and nearly sufficient condition, with efficient learning for generic functions and a separation from linear methods.
- Setting: The paper focuses on sparse functions on the d-dimensional hypercube that depend on an unknown P-dimensional latent coordinate subset, trained with one-pass batch-SGD on two-layer mean-field networks.This setting goes beyond linear dynamics while averaging some complexity of nonlinear SGD analysis.
- Merged-staircase property: The merged-staircase property orders nonzero Fourier supports so that each successive support adds at most one new coordinate.It generalizes staircase structures based on nested chains of Fourier coefficients.
- Examples: For example, z1 + z1z2 + z1z2z3 satisfies MSP and is learnable with O(d) samples, whereas z1z2z3 is not MSP and is not learnable in this regime.The distinction is whether the Fourier supports grow incrementally or jump directly to degree three.
- Main results: MSP is necessary for strong O(d)-SGD-learnability, as non-MSP functions have dimension-free dynamics that remain bounded away from zero risk.The dimension-free dynamics provides the equivalence used to derive this necessity result.
- Main results: Generic MSP functions are strongly O(d)-SGD-learnable, while symmetric degenerate MSP functions can fail, making sufficiency only nearly universal.Vanilla staircases require no genericity assumption and work under weaker activation conditions.
- Separation from linear methods: MSP functions with slowly growing latent dimension remain polynomial-sample learnable, while any linear method requires d^ωd(1) samples.Thus two-layer nonlinear SGD separates from linear methods in this regime.
2 Strong SGD-learnability in O(d)-scaling
The paper formalizes strong SGD-learnability for sparse latent functions in O(d) samples and specifies the two-layer mean-field SGD setting and regularity assumptions used for analysis.
- Training setup: The model uses a two-layer fully connected network trained with square-loss batch-SGD, time-varying step sizes, and regularization.The parameters are represented through hidden-unit weights and updated at each batch-SGD step.
- Definition: The learnability definition permits hyperparameter tuning but requires performance uniformly over every latent coordinate subset of size P.This uniformity excludes irregular scenarios that work only for selected latent subspaces.
- Assumptions: The strong regime fixes P and the target function, uses O(d) scaling, and restricts hyperparameters to continuous or discrete regular regimes.The continuous regime uses small batches and step sizes over Θ(1/η) steps, while the discrete regime uses larger batches and fixed-step dynamics.
- Assumptions: Smooth bounded activations, fresh bounded-noise samples, controlled initialization, and bounded Lipschitz hyperparameters ensure convergence to continuous mean-field dynamics.Under these assumptions, batch-SGD is approximated by a limiting dynamics for sufficiently large d, n, and N.
- Definition: Strong O(d)-SGD-learnability requires batch-SGD to achieve any target error with constant probability uniformly over unknown P-coordinate latent subspaces.The definition allows sufficiently large d, sample size n proportional to d, and network width at least a constant depending on the target function and error.
- Linear-method contrast: For any degree-k sparse function, linear methods require Ω(d^k) samples uniformly over latent subspaces, creating a separation whenever nonlinear SGD learns it with O(d) samples.The lower bound is not adaptive to the sparsity parameter P.
3 Continuous dimension-free dynamics and necessary condition
The analysis replaces mean-field SGD for sparse targets with dimension-free dynamics on the latent coordinates, proves approximation under finite-width sampling, and derives MSP necessity from the limiting flow.
- Dimension-free dynamics: The dimension-free dynamics reduces the ambient-dimensional mean-field network to effective parameters aligned with the P latent coordinates plus an adaptive Gaussian smoothing variable.The reduction uses symmetry and Gaussian approximation of the nuisance-coordinate contribution when d≫P.
- Dimension-free dynamics: The reduced predictor is a P-dimensional two-layer network with adaptive Gaussian smoothing, while its dynamics remains a Wasserstein gradient flow.The effective flow is defined over the reduced parameter distribution and latent-space test error.
- Approximation theorem: Theorem 5 gives a high-probability non-asymptotic approximation between batch-SGD, mean-field PDE, and dimension-free PDE trajectories under explicit width, batch-size, and step-size conditions.The approximation holds through times k≤T/η with probability at least 1−1/N.
- Illustration: Figure 1 reports that dimension-free PDE dynamics tracks batch-SGD until convergence and sequentially develops increasing-degree monomials for a four-term staircase target.The plot compares test error and Fourier coefficients, with batch-SGD averages and 95% confidence intervals over 10 experiments.
- Learnability equivalence: Theorem 6 characterizes strong O(d)-SGD-learnability exactly by whether the dimension-free gradient flow can reach arbitrarily small risk for suitable regularization and Lipschitz learning-rate schedules.For generic activations with zero attainable population risk, this is dynamic reachability of a global minimizer from the specified initialization.
- Necessary condition: MSP is necessary because without it some initialized coordinates remain zero throughout the flow, preventing every Fourier coefficient containing those coordinates from being learned.Theorem 7 therefore lower-bounds the limiting risk away from zero for all allowed learning-rate schedules and regularizations.
4 Sufficient conditions for strong SGD-learnability
The paper proves that MSP is nearly sufficient: generic MSP functions are learnable, while symmetries create exceptional degenerate cases; vanilla staircases require weaker conditions.
- Degenerate cases: Some MSP functions remain unlearnable because permutation symmetries force equal first-layer coordinates and keep the risk bounded away from zero.The example h*(z)=z1+z2+z1z3+z2z4 has a symmetry that constrains the learned Fourier coefficients.
- Generic MSP functions: For every MSP set structure, almost every choice of nonzero Fourier coefficients is strongly O(d)-SGD-learnable with a polynomial activation of degree L=28P.The exceptional degenerate coefficient choices form a measure-zero set under the defined coefficient distribution.
- Activation scope: The sufficiency theorem extends beyond the displayed polynomial activation to almost all degree-L polynomial activations and, with a technical caveat, generic smooth non-polynomial activations.The extension requires the relevant polynomial not to be identically zero for the given set structure.
- Vanilla staircases: Vanilla staircase functions are strongly O(d)-SGD-learnable without a genericity assumption using sufficiently smooth activations with nonzero derivatives through order P.The activation condition is σ∈C^(2P−1+1)(R) with σ^(r)(0)≠0 for r=0,…,P.
- Proof strategy: The sufficiency proof constructs a nonlinear first phase that produces informative features, followed by linear kernel training whose risk decays exponentially once the kernel matrix is nonsingular.Polynomial identity testing shows the simplified kernel determinant is nonzero almost surely for random second-layer weights and Fourier coefficients.
5 Separation with linear methods
The paper establishes that linear methods face dimension-dependent lower bounds on sparse nonlinear targets, whereas SGD-trained two-layer networks can exploit staircase structure with near-linear sample complexity.
- Linear-method model: The lower-bound framework models a linear method through a Hilbert-space feature map and shows its fitted predictor lies in a subspace of dimension at most min(n,q).Here q is the feature-space dimension, including the effective number of random features or the potentially infinite kernel case.
- Lower bounds: The improved lower bounds apply to both degree-k polynomial families and degree-P staircase families, including targets that are not almost orthogonal.They combine subspace-projection arguments with averaged prediction error over permuted latent embeddings.
- Lower bounds: For degree-k polynomial families, Proposition 12 requires min(n,q) to meet the stated degree-dependent lower bound for linear methods to achieve the target error.The corresponding staircase-family bound is stated separately for F_Str-P.
- Separation: SGD-trained two-layer networks outperform linear methods almost surely on nonlinear MSP functions under the paper’s respective learnability and lower-bound results.This separation follows by comparing O(d)-sample SGD learnability with the linear-method lower bounds.
- Growing-degree staircases: For vanilla staircases with growing degree P, linear methods require d^{ω_d(1)} samples when P=ω_d(1), while the neural-network result remains d^{O(1)} for sufficiently slow growth of P.This comparison addresses staircase functions that evade standard SQ lower-bound arguments because their hierarchy permits sequential queries.
6 Conclusion and future directions
The paper establishes MSP as necessary and nearly sufficient for O(d)-sample SGD learning of sparse functions, while identifying degenerate MSP cases and broader settings as open directions.
- Conclusion: MSP is necessary and nearly sufficient for learning sparse functions with two-layer mean-field networks in O(d) samples.The near-sufficiency result excludes a measure-zero subset of MSP functions.
- Conclusion: Degenerate symmetric MSP functions prevent a fully exact characterization because some are not strongly SGD-learnable in O(d)-scaling.
- Future directions: The current characterization is specific to fixed or slowly growing latent dimension, O(d) sample complexity, and two-layer networks.
- Future directions: Future analyses may incorporate leap size, compositionality order, network depth, and other function spaces such as Gaussian L2 spaces.
A Additional numerical simulations
The simulations illustrate that MSP structure alone does not guarantee convergence: generic MSP examples can succeed, while symmetry-induced degeneracies trap the dimension-free dynamics.
- Dimension-free dynamics: The simulations use a low-dimensional dimension-free dynamics that can be solved numerically, while its global convergence is analyzed through a Wasserstein gradient flow.
- Global convergence: Four MSP examples split between global convergence and nonzero limiting risk, with the latter corresponding to G-invariant functions.The convergent examples are h1 and h2; h3 and h4 do not converge.
- G-invariant MSP functions: Symmetry constrains the weight distribution to a lower-dimensional linear subspace, which can exclude all global minimizers and keep risk bounded away from zero.
- Perturbed MSP functions: A small perturbation breaking h4's permutation symmetry makes the perturbed MSP strongly O(d)-SGD-learnable.
- Escaping the saddle-space: When stairs are missing, the dimension-free dynamics remains trapped while one-pass batch-SGD escapes only after n approximately d^2 or d^2.7 iterations.
B.2.1 Proof of Theorem 5
The proof of Theorem 5 establishes that batch-SGD is accurately approximated by dimension-free dynamics over the relevant time scale, enabling the paper’s learnability equivalence and MSP necessity results.
- Approximation strategy: For constant time and sufficiently large width, batch size, and inverse step size, the mean-field PDE accurately approximates batch-SGD while requiring n=O(d) samples.
- Consequences: The resulting approximation supports the equivalence between strong O(d)-SGD learnability and convergence of the dimension-free dynamics, from which non-MSP lower bounds follow.
- Coupling argument: The proof couples mean-field and dimension-free trajectories through their initialization and bounds parameter, distribution, and prediction-function discrepancies.
C.1 Discrete time mean-field and dimension-free dynamics
The discrete-time analysis constructs discrete mean-field and dimension-free dynamics, proves their approximation to batch-SGD, and transfers the continuous-time learnability characterization to bounded-step discrete training.
- Discrete dynamics: Discrete batch-SGD is approximated by a discrete mean-field dynamics and then by a discrete dimension-free dynamics as ambient dimension and batch size grow.
- Approximation theorem: The discrete approximation has a worse dependence on the number of iterations than the continuous-time result.
- MSP necessity: The discrete MSP necessity result follows by transferring the continuous argument through the discrete approximation bounds.
- Learnability equivalence: Strong discrete-time O(d)-SGD learnability is equivalent to convergence of the discrete dimension-free dynamics for suitable regularization and bounded step sizes.
- Proof strategy: The proof combines batch-SGD-to-mean-field concentration with a discrete Gronwall bound controlling the mean-field-to-dimension-free discrepancy.
D Vanilla staircase functions are strongly O(d)-SGD-learnable: Proof of Theorem 10
The proof establishes strong O(d)-SGD-learnability for vanilla staircases by analyzing two training phases: nonlinear first-layer learning followed by linear second-layer convergence. Under suitable activation assumptions and a two-phase schedule, the dimension-free dynamics reaches arbitrarily small risk.
- Hyperparameters: The schedule initializes first-layer weights at zero, uses random second-layer signs, then alternates first-layer and second-layer training.Phase 1 trains u while fixing a; Phase 2 trains a while fixing u.
- Theorem 10: Vanilla staircase functions are strongly O(d)-SGD-learnable when the activation has nonzero derivatives through order P and sufficient local smoothness.The theorem requires L > 2^P−1 and provides hyperparameters yielding arbitrarily small risk.
- Proof strategy: The proof separates training into nonlinear first-layer dynamics and linear second-layer dynamics.The first phase creates structured first-layer weights; the second phase keeps them fixed and trains the output layer through a kernel evolution.
- Phase 1: First-layer training remains analyzable because the staircase structure lets the proof track the lowest-order terms in the evolving weights.The activation is approximated locally by a polynomial while the weights remain small.
- Phase 2: Second-layer training converges once the phase-1 kernel has a dimension-independent positive minimum eigenvalue.If λmin(K_T1) > c, the required phase-2 duration is T2 = T1 + log(K/ε)/c.
E Generic MSP functions are strongly O(d)-SGD-learnable: Proof of Theorem 9 (discrete-time regime)
The discrete-time analysis shows that generic functions with any merged-staircase Fourier support are strongly O(d)-SGD-learnable almost surely. The result uses polynomial activations, two-phase training, and a kernel argument after a finite nonlinear phase, while excluding a measure-zero set of degenerate coefficients.
- MSP structure: An MSP set structure orders Fourier supports so each new support introduces at most one previously unseen variable.Functions with MSP structure have exactly those sets as their nonzero Fourier coefficients.
- Scope: The near-sufficiency result is not universal over MSP functions because explicit degenerate MSP functions fail strong SGD-learnability, although such functions form a measure-zero set.Thus the theorem is almost-sure rather than pointwise over all nonzero Fourier coefficients.
- Main theorem: For any MSP set structure, almost every function with that Fourier support is strongly O(d)-SGD-learnable.The theorem uses σ(x) = (1 + x)^L with L = 2^8P.
- Training setup: The proof uses zero first-layer initialization, random second-layer weights, and two phases with Θ(d) batch size and constant-scale steps.The first phase lasts P steps, while the second lasts a constant number of steps depending on the target accuracy and latent dimension.
- Phase 2: After first-layer training, the second-layer phase is kernel gradient descent and reaches ε-risk when the phase-1 kernel has a dimension-independent positive minimum eigenvalue.A sufficiently long constant second phase then contracts the risk to ε.
E.1.2 Phase 1 (nonlinear training)
Phase 1 converts nonlinear weight evolution into a finite polynomial system and reduces learnability to showing that a determinant is not identically zero. Polynomial identity testing and linear independence of large powers then establish this condition for generic MSP functions under sufficiently high-degree activations.
- Polynomial representation: After a constant number of first-layer steps, the weights can be represented as constant-degree polynomials in the second-layer weights.This follows because the activation is polynomial on the bounded input range.
- Simplification: The nonlinear dynamics can be simplified by dropping the evolving network-output interaction term, whose effect is O(η^k) for small learning rate.The simplified dynamics remains close enough to the original dynamics for the kernel analysis.
- Polynomial identity test: Learnability reduces to proving that the determinant of a structured matrix is a nonzero polynomial in the second-layer weights and Fourier coefficients.A nonzero determinant yields a positive kernel eigenvalue and therefore strong O(d)-SGD-learnability.
- Generic MSP functions: For MSP supports, the proof enforces shared Fourier coefficients across time steps and shows the resulting determinant is almost surely nonzero over generic coefficients.The constrained matrix N incorporates the MSP coefficient structure before applying the determinant argument.
- Smooth activations: For smooth activations beyond polynomial ones, polynomial approximation and a random activation perturbation extend the argument, with the perturbation serving as a technical device for the Vandermonde proof.The extension requires σ^(r)(0) ≠ 0 for r = 0,...,P.
F.1 Statement of the result
The section establishes a two-phase proof strategy for generic MSP functions: first-layer nonlinear dynamics create a nondegenerate kernel, then perturbed-activation second-layer training converges. Polynomial approximation controls the dynamics sufficiently to obtain dimension-free learnability.
- Training phases: Phase 1 trains first-layer weights with fixed second-layer weights, while Phase 2 trains the second layer using the perturbed activation.The two phases separate nonlinear feature formation from linear kernel-based fitting.
- Theorem 56: Generic MSP functions are shown to be strongly O(d)-SGD-learnable almost surely after a one-time activation perturbation.The result assumes a smooth activation with nonzero derivatives through order P and allows the activation to change between training phases.
- Dimension-free dynamics: The first-layer dynamics admit a polynomial approximation with error O(t^L) for sufficiently small training times.This reduces the analysis to a finite-order dynamical system while preserving the required accuracy.
- Technical challenge: The simplified dynamics are easier to analyze but may approximate the true dynamics only to order Ω(t^2) in L2 norm.The proof therefore tracks determinant and higher-order terms rather than relying on a uniformly sharp L2 approximation.
- Kernel nondegeneracy: A nonzero low-order term in the kernel determinant yields a minimum-eigenvalue lower bound and a condition number independent of ambient dimension d.The kernel remains sufficiently nondegenerate for second-layer convergence.
F.8 Proving learnability of generic MSP functions, Theorem 56
Theorem 56 proves learnability for generic functions with an MSP Fourier support by showing that activation perturbations almost surely make the relevant kernel nondegenerate. A constant-duration second phase then reduces the residual to any target accuracy.
- Polynomial identity testing: For any MSP set structure, generic Fourier coefficients and random activation perturbations make the determinant polynomial nonzero almost surely.The proof specializes the determinant to the target’s MSP support and invokes polynomial nonvanishing under random parameter choices.
- Kernel lower bound: The resulting kernel satisfies λ_min(K_T1) ≥ c(T1)^(2l_S) for sufficiently small T1.Because the relevant degrees depend only on the latent dimension P, the lower bound is independent of ambient dimension d.
- Convergence: Choosing T1 as a constant and T2 = T1 + log(∥g_T1∥2/ε)/c^2 drives the residual to error ε.The residual is bounded before Phase 2, so both training times depend on latent-level quantities rather than d.
F.9 Proof of Lemma 73
The proof of Lemma 73 reduces general MSP structures to minimal ones and constructs a low-order nonzero determinant term through polynomial and Vandermonde arguments. This establishes the kernel nondegeneracy needed for Theorem 56.
- Reduction to minimal structures: It suffices to prove a low-degree nonzero determinant term for minimal MSP structures, because larger MSP structures preserve such terms under specialization.Setting additional Fourier coefficients to zero cannot remove the existing nonzero polynomial term.
- Recurrence analysis: The recurrence analysis shows that each coefficient vanishes below its structure-dependent order and has a nonzero leading contribution at that order.The leading terms are identified inductively from the MSP ordering and its set differences.
- Row separation: Distinct hypercube inputs produce derivative differences with nonzero terms of degree at most 2^P−1 in T1.This separates the rows needed for the auxiliary determinant argument.
- Determinant argument: A Vandermonde determinant argument transfers the low-order nonzero term from the auxiliary matrix to the kernel determinant.The determinant remains a nonzero polynomial in the activation and Fourier parameters.
G Explicit sample-complexity bounds in all parameters
The appendix makes the sample-complexity dependence explicit beyond the central O(d) scaling, while showing that sharper bounds require additional structure or smoothing. It also identifies substantial dependence on the latent dimension P and activation choices.
- Scope of bounds: The main theorem gives n = C(ε, δ, h*)d samples, while the appendix leaves improving explicit dependencies on ε, δ, and h* for future work.The paper’s primary focus is dependence on ambient dimension d.
- Vanilla staircase bounds: For vanilla staircases, explicit bounds depend exponentially on training time and can have at least doubly exponential dependence on P.The authors conjecture an optimal exp(O(P)) dependence, but the displayed propositions are looser.
- Smoothed complexity: A tighter analysis of second-layer training improves complexity bounds for MSP functions under a smoothed model.The refinement targets dependencies on parameters other than ambient dimension d.
- Proposition 86: Under stronger nondegeneracy assumptions on degree-1 Fourier coefficients and an ad hoc activation, SGD succeeds with d · poly(2^P log(1/δ)/(με)) samples.The corresponding network width is polynomial in the same quantity.
- Corollary 87: A small perturbation of degree-1 Fourier coefficients converts any normalized function into one satisfying the learnability conditions with μ = μ′/2^O(P).This provides the smoothed-complexity bridge to Proposition 86.
H Lower bounds on learning with linear methods
This section develops dimension-based lower bounds showing that linear methods need large feature-space or sample dimensions to learn difficult function classes. The bounds apply broadly, including non-almost-orthogonal targets and staircase functions.
- Lower-bound method: The lower-bound framework improves prior results for target functions that are not almost orthogonal by using the Gram matrix operator norm.
- Dimension lower bounds: Linear methods represent predictions in a subspace whose dimension is at most min(n, q), limiting approximation of many target functions.The bound does not assume a particular statistical model and applies when the estimator lies in a finite-dimensional subspace.
- Lower-bound method: The construction targets subspaces containing difficult high-degree monomials, separating the complexity of individual components from that of the full staircase function.
- Monomial targets: For fixed-degree monomial targets, learning permutation-invariant function classes requires n = Ω_d(d^k/m) samples.
- Staircase functions: For vanilla staircase functions of degree P, any linear method requires a superpolynomial sample complexity n ≥ d^ω_d(1).The lower bound concerns average generalization error over the class of staircase functions.
I.2 Discrete-time analysis of SGD on second layer
The discrete-time analysis separates training into first-layer and second-layer phases and reduces second-layer convergence to certificate existence and kernel conditioning. A direct bias-variance analysis yields tighter sample-complexity requirements than propagation-of-chaos bounds.
- Tighter second-layer analysis: n = d/(ελ_min)^C suffices in the direct second-layer SGD analysis, improving the earlier requirement n = d log(1/ε)^7/λ^7.Here C is a positive constant, and λ_min is the relevant kernel-matrix minimum eigenvalue.
- Second-layer dynamics: With first-layer weights fixed, second-layer training is linear in the output weights under a feature map determined by the initialization.
- Certificate reduction: A low-norm certificate achieving small loss is sufficient for batch-SGD to reach approximately the certificate's loss after few iterations.
- Bias-variance analysis: The analysis decomposes the second-layer error into a noiseless optimization term and a batch-SGD noise term.
- Two-phase training: The two-phase proof trains first-layer weights before fixing them and training second-layer weights, then bounds final loss through the resulting kernel matrix.
- Kernel conditioning: Sufficient network width and hyperparameter conditions reduce the convergence guarantee to proving that the matrix M is well-conditioned.