Technical Summary: Full-Batch Gradient Descent Outperforms One-Pass SGD in Single-Index Learning

1. Problem Setup and Motivation

The paper investigates the statistical and optimization efficiency of gradient-based learning in non-convex, nonlinear settings, specifically focusing on the Gaussian single-index model. The goal is to recover an unknown unit vector $\theta^\star \in \mathbb{S}^{d-1}$ from $n$ i.i.d. samples $(x_i, y_i)$, where:
$$x_i \sim \mathcal{N}(0, I_d), \quad y_i = \sigma(\langle x_i, \theta^\star \rangle)$$
The activation function $\sigma$ is either quadratic ($\sigma(z) = z^2$) or a truncated quadratic. This setting is a standard sandbox for understanding feature learning in shallow neural networks and is closely related to phase retrieval.

The Core Question: While it is folklore that reusing data (multi-pass) improves statistical efficiency in linear regression, this phenomenon is not well-understood in nonlinear settings. Specifically, does full-batch Gradient Descent (GD) (which reuses all data at every step) outperform one-pass Stochastic Gradient Descent (SGD) (which uses each data point exactly once) in terms of sample complexity?

Prior work established that for activations with an "information exponent" of 2 (including quadratic functions), one-pass SGD requires $n \gtrsim d \log d$ samples to achieve weak recovery. This leaves a logarithmic gap from the information-theoretic limit of $n \gtrsim d$. The paper seeks to determine if full-batch GD can close this gap without algorithmic modifications (like spectral initialization or loss modification).

2. Methodology and Analysis Framework

The authors analyze two distinct algorithmic setups to characterize the sample and iteration complexity:

A. Spherical Gradient Flow on Correlation Loss

• Setup: The algorithm minimizes the correlation loss $\hat{L}(\theta) = -\frac{1}{n}\sum y_i \sigma(\langle x_i, \theta \rangle)$ constrained to the unit sphere $\mathbb{S}^{d-1}$.
• Technique: The authors analyze the dynamics of the gradient flow ODE. They relate the convergence to the spectral properties of the matrix $A(\theta)$, which acts as a time-varying operator in the flow.
• Key Analytical Tool: They establish a uniform BBP (Baik-Ben Arous-Péché) phase transition for the Hessian along the optimization trajectory. This involves proving that the top eigenvector of the empirical matrix remains aligned with the signal $\theta^\star$ and that a spectral gap exists uniformly over the sphere.

B. Euclidean Gradient Descent on Squared Loss

• Setup: The algorithm minimizes the squared loss $\hat{L}(\theta) = \frac{1}{2n}\sum (\sigma(\langle x_i, \theta \rangle) - \sigma(\langle x_i, \theta^\star \rangle))^2$ in Euclidean space $\mathbb{R}^d$.
• Initialization: The analysis assumes small initialization ($\|\theta_0\| \ll 1$), a regime inspired by recent work on feature learning where early dynamics approximate a power iteration.
• Technique: The authors decompose the training trajectory into two phases:
  1. Search Phase: The angle between $\theta_t$ and $\theta^\star$ shrinks, and the norm $\|\theta_t\|$ grows. During this phase, the truncation of the activation is inactive, and dynamics resemble a power method on a fixed matrix.
  2. Refinement Phase: Once the iterate is close to the optimum, the authors establish one-point strong convexity of the empirical gradient to prove geometric convergence to the exact solution.

3. Key Contributions and Results

Contribution 1: The Necessity of Truncation for Full-Batch GD

The paper first addresses the standard quadratic activation $\sigma(z) = z^2$ under the correlation loss.

• Negative Result: For unbounded quadratic activation, full-batch spherical gradient flow fails to achieve weak recovery when $n \ll d \log d$. The estimator converges to a trivial direction uncorrelated with $\theta^\star$.
• Implication: In this specific setup, full-batch GD offers no statistical advantage over one-pass SGD; both require $n \gtrsim d \log d$ samples.

Contribution 2: Sample Complexity Separation via Truncation

The authors introduce a simple modification: truncating the quadratic activation (either smoothly or via a hard cutoff).

• Positive Result: With a truncated activation, full-batch spherical gradient flow achieves weak recovery with $n \gtrsim d$ samples.
• Separation: This matches the information-theoretic limit and strictly outperforms one-pass SGD, which still requires $n \gtrsim d \log d$ for the same truncated link function.
• Mechanism: The truncation ensures the empirical loss landscape is benign (possessing a uniform spectral gap and alignment with $\theta^\star$) even at finite $n/d$, allowing the gradient flow to escape trivial solutions.

Contribution 3: Strong Recovery and Iteration Complexity

Moving beyond landscape analysis, the paper analyzes full-batch GD on the squared loss with small initialization.

• Strong Recovery: The authors prove that with $n \gtrsim d$ samples, full-batch GD achieves strong (exact) recovery (convergence to $\theta^\star$ up to sign).
• Iteration Complexity: The algorithm converges in $T \gtrsim \log d$ gradient steps.
• Significance: This is the first convergence rate guarantee for standard full-batch GD (without algorithmic or loss modifications) in the proportional regime ($n/d$ constant) for single-index models with information exponent 2.

4. Significance and Claims

The paper claims to provide the first rigorous theoretical evidence that data reuse (full-batch GD) can strictly improve statistical efficiency in a non-convex feature learning problem, provided the activation function is appropriately truncated.

• Bridging the Gap: The work demonstrates that the $\log d$ gap between the sample complexity of one-pass SGD and the information-theoretic limit can be closed by full-batch GD, but only if the optimization landscape is sufficiently regularized (via truncation).
• Beyond Landscape Analysis: Unlike prior works that rely solely on static landscape properties, this paper provides a trajectory analysis that characterizes both the sample complexity and the iteration complexity (time to convergence).
• Limitations: The authors note that their refined trajectory analysis is currently limited to squared loss dynamics from small initialization. They do not claim that full-batch GD outperforms one-pass SGD for all nonlinear settings or without the specific truncation mechanism.

In summary, the paper establishes that while full-batch GD does not inherently solve the "hard" learning problem of quadratic single-index models, a minimal modification (truncation) allows it to achieve optimal sample complexity ($O(d)$) and logarithmic iteration complexity ($O(\log d)$), thereby outperforming one-pass SGD in both dimensions.