PaperlayerView exact revision on arXiv

Source-linked AI summary

Online stochastic gradient descent on non-convex losses from high-dimensional inference

Gerard Ben Arous, Reza Gheissari, Aukosh Jagannath

arXiv:2003.10409v4stat.MLcs.LGmath.PRmath.ST

TL;DR

The paper asks how online SGD recovers high-dimensional parameters from random starts when inference losses are random and non-convex. It classifies these problems by the population loss’s information exponent and derives nearly sharp sample thresholds, finding that search dominates the cost while descent is rapid once correlation is attained.

  • Problem

    High-dimensional random starts concentrate near uninformative points, while common dimension-free smoothness and convexity assumptions are unrealistic for many inference losses.

  • Method

    The paper analyzes online SGD for rank-one estimation using the information exponent and sample-wise error scaling, without requiring convexity or uniform derivative control.

  • Results

    Weak-recovery time is linear for k<2, quasi-linear for k=2, and polynomial for k>2; for k≥2, most data is used in search and descent takes O(N) samples.

  • Takeaways & Limitations

    The information exponent determines the dimension dependence of online-SGD sample complexity, while post-search descent follows population-gradient behavior rapidly.

  • Takeaways & Limitations

    The results restrict initializations to the upper half-sphere with CLT-scale positive correlation; an exact equatorial start can remain trapped without anti-concentration.

Abstract

from arXiv · show

Stochastic gradient descent (SGD) is a popular algorithm for optimization problems arising in high-dimensional inference tasks. Here one produces an estimator of an unknown parameter from independent samples of data by iteratively optimizing a loss function. This loss function is random and often non-convex. We study the performance of the simplest version of SGD, namely online SGD, from a random start in the setting where the parameter space is high-dimensional. We develop nearly sharp thresholds for the number of samples needed for consistent estimation as one varies the dimension. Our thresholds depend only on an intrinsic property of the population loss which we call the information exponent. In particular, our results do not assume uniform control on the loss itself, such as convexity or uniform derivative bounds. The thresholds we obtain are polynomial in the dimension and the precise exponent depends explicitly on the information exponent. As a consequence of our results, we find that except for the simplest tasks, almost all of the data is used simply in the initial search phase to obtain non-trivial correlation with the ground truth. Upon attaining non-trivial correlation, the descent is rapid and exhibits law of large numbers type behavior. We illustrate our approach by applying it to a wide set of inference tasks such as phase retrieval, and parameter estimation for generalized linear models, online PCA, and spiked tensor models, as well as to supervised learning for single-layer networks with general activation functions.

1. Introduction

The paper classifies high-dimensional rank-one estimation problems by an information exponent and derives nearly sharp online-SGD sample-complexity thresholds without convexity or uniform derivative assumptions. It shows that search dominates the data cost for information exponent at least 2, while descent is rapid afterward.

  • Applications: Online SGD succeeds with polynomially many samples across phase retrieval, generalized linear models, online PCA, spiked matrix and tensor models, mixtures, and single-layer networks.These problems can have non-convex losses and exponentially many critical points.
  • Classification: The information exponent is the first non-zero Taylor-term degree of the population loss near the equator, classifying sample complexity into linear, quasi-linear, and polynomial regimes.The classification depends on this population-loss geometry rather than uniform loss-landscape control.
  • Search and descent: For k≥2, essentially all samples are spent attaining non-trivial correlation, whereas descent to near-full correlation takes O(N) samples.The descent-to-search sample ratio is O(αc(N,k)^-1), which vanishes with dimension.
  • Assumptions: The analysis uses scaling assumptions on sample-wise gradient errors rather than convexity, Hessian control, or uniform derivative bounds.The approach allows losses whose initialized Hessians may simultaneously have positive and negative eigenvalues.
  • Main results: Theorem 1.3 gives consistent-estimation thresholds scaling as αc(N,1), αc(N,2) log N, and αc(N,k)(log N)^2 for k≥3.The corresponding conditions are αN≫αc(N,1), αN≫αc(N,2)·log N, and αN≫αc(N,k)·(log N)^2.
  • Main results: Theorem 1.4 shows these thresholds are sharp up to logarithmic factors: below them, online SGD fails to exit the search phase.The lower bound applies under the stated step-size regimes and implies at least αc(N,k)N samples are necessary.

2. Applications to some important inference problems

The paper applies its information-exponent framework across supervised learning, generalized linear models, linear regression, PCA, spiked matrix and tensor models, and Gaussian mixtures. These applications show that activation structure and model alignment determine whether online SGD achieves weak or strong recovery and how many samples random initialization requires.

  • The information exponent is determined by the first nonzero Hermite coefficient: exponents 1, 2, or at least 3 correspond respectively to nonzero first, zero first but nonzero second, or zero first and second coefficients.
  • 2.1. Supervised learning for single-layer networks.: For supervised learning, f(x)=x and f(x)=x3 have exponent 1, whereas f(x)=x3−3x has exponent 3 and requires Ω(N3/2 log N) samples just to exit search.
  • 2.1. Supervised learning for single-layer networks.: From a warm start, descent rapidly reaches the ground truth with linearly many samples, largely independently of the activation function’s information exponent.
  • 2.2–2.3. Generalized linear models.: Invertible increasing activations in Gaussian generalized linear models always have information exponent 1 under the stated growth and moment conditions.
  • 2.3. Linear regression with random covariates.: Linear regression satisfies the assumptions with information exponent 1 without requiring independent covariate entries or Gaussian covariates and errors.
  • 2.4–2.5. Online PCA and spiked matrix and tensor models.: Spiked matrix models have exponent 2, spiked tensor models have exponent p, and online SGD matches conjecturally optimal first-order thresholds in the tensor setting while incurring logarithmic overhead for online matrices.The tensor threshold is reported as α∼N^{k−2}; for p=2, the online matrix threshold lies between α∼log N and α∼(log N)^2.

3. Analysis of two stages of performance

The analysis separates online SGD into a noisy search phase and a concentrated descent phase. Random initialization requires dimension- and exponent-dependent time to obtain weak recovery, after which SGD follows population dynamics and reaches strong recovery in linear time.

  • SGD reaches order-one correlation in time ˜O(αc(N,k)N), then well-approximates population dynamics until attaining 1−o(1) correlation.
  • 3.1. Search phase.: During search, the weak-recovery timescale applies from initial correlations m0=Ω(N−1/2), including the natural scale of random high-dimensional initialization.
  • 3.1. Search phase.: Search trajectories can differ substantially from population dynamics because directional-error fluctuations are comparable to the correlation, although their recovery timescales agree.
  • 3.2. Descent phase.: Once correlation is non-trivial, SGD rapidly reaches the ground truth in linear time and its trajectory is well approximated by population dynamics.
  • 3.3. The k=1 case.: For information exponent k=1, the entire trajectory obeys a law-of-large-numbers behavior, allowing sharper order-one sample results for weak recovery.

4. A Difference inequality for Online SGD

The proof controls the correlation process through a difference inequality that decomposes online SGD into population drift, martingale noise, and projection corrections. This yields uniform high-probability bounds under moment and step-size conditions across information-exponent regimes.

  • The key technical step shows that m_t is, with high probability, a super-solution to an integral equation for the population dynamics.
  • The difference inequalities distinguish k=1 from k≥2, with separate bounds controlling the corresponding drift regimes under αδ^2≤ε and polynomially growing α.
  • The correlation increment is decomposed into population-loss drift, directional martingale noise, and a second-order correction from the nonlinear projection step.
  • The analysis bounds projection corrections and higher-order terms using truncation, moment inequalities, and martingale concentration uniformly over initial points.
  • The directional-error process is a martingale because sample-wise gradient errors have conditional mean zero, enabling Doob- and Freedman-type control.

5. Attaining weak recovery

The weak-recovery proof integrates the difference inequalities to show that online SGD exits the search phase within the predicted exponent-dependent timescale. The argument treats k=1, k=2, and k≥3 separately through discrete comparison inequalities.

  • The proof targets weak recovery before the guaranteed recovery time by applying the search-phase estimates to the correlation process.
  • The argument analyzes k=1, k=2, and k≥3 separately because their difference inequalities have different growth forms.
  • k=1: For k=1, discrete Grönwall comparison yields weak recovery under conditions involving αδ log N→∞ and sufficiently small δ.
  • k≥2: For higher information exponents, discrete Bihari–LaSalle comparisons produce exponent-dependent recovery conditions, including αδ N^{(k−2)/2}→∞.

6. Strong recovery and the descent phase

The descent phase is controlled by proving a law of large numbers for the SGD trajectory and comparing it with linearized population dynamics. These estimates establish strong recovery under the theorem’s parameter conditions, including linear sample complexity when k=1.

  • 6. Strong recovery and the descent phase: Theorem 3.2 proves a law of large numbers for the descent trajectory, which combines with Theorem 3.1 to establish the main recovery result.The proof controls the stochastic trajectory relative to its linearized population dynamics.
  • 6. Strong recovery and the descent phase: With probability 1−o(1), the nonlinear and linearized trajectories remain uniformly within ε over the analyzed time horizon when αδ2=o(1).The uniform approximation is obtained through maximal inequalities and a discrete Gronwall argument.
  • 6. Strong recovery and the descent phase: The proof shows that the final correlation m(XM) converges to 1 in probability under the parameters specified by Theorem 1.3.The argument bounds the probability that |m(XM)−1| exceeds any fixed ε.
  • 6. Strong recovery and the descent phase: For k=1, the theorem yields linear sample complexity by taking α as a sufficiently large constant while keeping ε and ζ fixed at order one.The passage identifies this as item (a) of Theorem 3.3.

7. Online SGD does not recover with smaller sample complexity

This section proves that online SGD cannot weakly recover when α is below the threshold of Theorem 3.1, up to logarithmic factors. The refutation uses excursion-time control to show that the process lacks enough time to build positive correlation in one pass.

  • 7. Online SGD does not recover with smaller sample complexity: If α is smaller than the Theorem 3.1 threshold up to log N factors, online SGD does not have enough time to weakly recover during one pass through M samples.The result is stated as an accompanying refutation theorem.
  • 7. Online SGD does not recover with smaller sample complexity: For k=1, choosing α=o(1) and δ=O(1) yields, with probability 1−o(1), an excursion-time lower bound ˜t≥M.The conclusion holds for every fixed η>0.
  • 7. Online SGD does not recover with smaller sample complexity: For k>2, a sufficiently slowly diverging dN again yields ˜t≥M under αδ2=O(1).The required conditions combine √α=o(N^(k−2)/2) with √α=O(δ−1).
  • 7. Online SGD does not recover with smaller sample complexity: The refutation applies because the process remains non-positive during negative excursions, preventing the correlation from reaching any fixed positive η before the sample pass ends.The excursion intervals satisfy m(Xt)≤0<η.

Appendix B. Deferred proofs from Section 2

Appendix B verifies that the examples from Section 2 satisfy the paper’s assumptions. The proof uses polynomial-growth conditions and Gaussian rotational invariance to control the relevant population-loss quantities.

  • Appendix B. Deferred proofs from Section 2: The appendix verifies that the examples introduced in Section 2 satisfy Assumptions A–B.This section serves as the deferred assumption-checking argument.
  • Appendix B. Deferred proofs from Section 2: Polynomial growth of f and f′ ensures the required Gaussian and noise expectations are finite under the stated finite-moment condition on ϵ.The proof uses Gaussianity of a·x and the assumption that ϵ has a finite 4+δ moment.

B.1. Proof of Proposition 2.1.

The proof of Proposition 2.1 analyzes the population loss using Gaussian noise operators and Hermite expansions, then verifies the required gradient moment bounds. Polynomial-growth assumptions on the activation and its derivative support these estimates.

  • B.1. Proof of Proposition 2.1.: The proof uses the Gaussian noise operator’s action on Hermite polynomials, Tshk=skhk, to compute the expectations required for Proposition 2.1.The argument extends the identity to negative s by symmetry.
  • B.1. Proof of Proposition 2.1.: Assumption B follows from f′ having at most polynomial growth, which places f in the Gaussian Sobolev space H1.The proof then uses the resulting Sobolev control in the population-loss analysis.
  • B.1. Proof of Proposition 2.1.: The sample gradient is expressed as 2(f(a·x)+ϵ−f(a1))f′(a·x)πxa, separating the residual, activation derivative, and tangent-space projection.This representation is used to establish the moment conditions for the loss gradient.
  • B.1. Proof of Proposition 2.1.: Hölder’s inequality and the noise moment assumption yield the required bounds for the loss and gradient moments.The proof also uses polynomial growth of f and f′ and standard inequalities for sums.

B.2. Proof of Proposition 2.4.

The proof verifies the assumptions for a population loss with information exponent 1 and establishes the required bounds for the true loss using projection and moment estimates.

  • Gaussian integration by parts gives u1(f) = E[f′(Z)] > 0 when f is increasing, invertible, and differentiable.
  • The true-loss gradient is expressed as ∇xL(y; a, x) = (y − b′(a · x))πxa, enabling Cauchy–Schwarz bounds.
  • Finite moments and at-most-exponential growth of b′ = f provide the finiteness needed for the gradient estimates.
  • The population loss has information exponent 1 and satisfies Assumption A.

B.3. Proof of Proposition 2.5.

The proof shows that the linear model has information exponent 2 and verifies the required directional and norm bounds using projection identities, Cauchy–Schwarz, and moment assumptions.

  • The model’s directional derivative satisfies ∇L(x; a, ϵ) · v = (a · (v − x) + ϵ)(a · πxv).
  • The proof reduces the required bounds to moment estimates for projected vectors and the noise term.
  • Choosing 4 + ι = 5 yields the norm bound under finite 10-th moments for a and finite 5-th moment for ϵ.
  • The population loss satisfies Assumption A and has information exponent 2.

B.4. Proof of Proposition 2.7.

The proof identifies information exponent p for the spiked tensor model and verifies Assumption B through moment bounds on the loss gradient.

  • Taking expectations gives Φ(x) = −λm(x)^p, so Assumption A holds and the problem has information exponent p.

B.5. Proof of Proposition 2.8.

The supplied passages introduce the tensor-gradient setup and indicate that the proof controls its derivatives through centered independent components and finite moments.

  • Writing H(x) = (J, x⊗p) identifies the tensor polynomial whose gradient is analyzed.
  • The derivative expansion uses multi-indices and tensor contractions such as J(x⊗p−1, ·).
  • Finite 6-th moments of J’s entries yield E||DH||^6 ≤ CN^3 and the second half of Assumption B.
  • The directional gradient is a weighted sum of centered i.i.d. variables with uniformly bounded 6th moments.

B.6. Proof of Proposition 2.10.

The proof combines two sign and symmetry-based results to identify the information exponent and establish positivity of φ′(m) for m > 0. The appendix also proves a discrete Bihari–LaSalle inequality by induction and lower-bounding the resulting sequence.

  • Rotation invariance of Z and the identification of Z1 and Z2 as its first two entries support the preceding calculation.
  • The information exponent is 1 when p ≠ 1/2 and 2 when p = 1/2, with φ′(m) > 0 for m > 0.
  • The discrete Bihari–LaSalle argument inductively shows m_t ≥ b_t ≥ a_t by comparing recursive sequences.
  • The proof then uses monotonicity of a_t and a uniform bound holding for every t to complete the lower bound on a_t.
Loading 2003.10409v4…