Smoothing the Landscape Boosts the Signal for SGD
Optimal Sample Complexity for Learning Single Index Models

Gradient descent-based algorithms are popular for deriving computational and statistical guarantees for a number of high-dimensional statistical learning problems [2, 3, 1, 4, 5, 6]. Despite the fact that the empirical loss is nonconvex and in the worst case computationally intractible to optimize, for a number of statistical learning tasks gradient-based methods still converge to good solutions with polynomial runtime and sample complexity. Analyses in these settings typically study properties of the empirical loss landscape [7], and in particular the number of samples needed for the signal of the gradient arising from the population loss to overpower the noise in some uniform sense. The sample complexity for learning with gradient descent is determined by the landscape of the empirical loss.

One setting in which the empirical loss landscape showcases rich behavior is that of learning a single-index model. Single index models are target functions of the form $f^{*}(x)=\sigma(w^{\star}\cdot x)$, where $w^{\star}\in S^{d-1}$ is the unknown relevant direction and $\sigma$ is the known link function. When the covariates are drawn from the standard $d$-dimensional Gaussian distribution, the shape of the loss landscape is governed by the information exponent ${k^{\star}}$ of the link function $\sigma$, which characterizes the curvature of the loss landscape around the origin. Ben Arous et al. [1] show that online stochastic gradient descent on the empirical loss can recover $w^{*}$ with $n\gtrsim d^{{k^{\star}}-1}$ samples; furthermore, they present a lower bound showing that for a class of online SGD algorithms, $d^{{k^{\star}}-1}$ samples are indeed necessary.

However, gradient descent can be suboptimal for various statistical learning problems, as it only relies on local information in the loss landscape and is thus prone to getting stuck in local minima. For learning a single index model, the Correlational Statistical Query (CSQ) lower bound only requires $d^{{k^{\star}}/2}$ samples to recover $w^{\star}$ [6, 4], which is far fewer than the number of samples required by online SGD. This gap between gradient-based methods and the CSQ lower bound is also present in the Tensor PCA problem [8]; for recovering a rank 1 $k$-tensor in $d$ dimensions, both gradient descent and the power method require $d^{k-1}$ samples, whereas more sophisticated spectral algorithms can match the computational lower bound of $d^{k/2}$ samples.

In light of the lower bound from [1], it seems hopeless for a gradient-based algorithm to match the CSQ lower bound for learning single-index models. [1] considers the regime in which SGD is simply a discretization of gradient flow, in which case the poor properties of the loss landscape with insufficient samples imply a lower bound. However, recent work has shown that SGD is not just a discretization to gradient flow, but rather that it has an additional implicit regularization effect. Specifically, [9, 10, 11] show that over short periods of time, SGD converges to a quasi-stationary distribution $N(\theta,\lambda S)$ where $\theta$ is an initial reference point, $S$ is a matrix depending on the Hessian and the noise covariance and $\lambda=\frac{\eta}{B}$ measures the strength of the noise where $\eta$ is the learning rate and $B$ is the batch size. The resulting long term dynamics therefore follow the smoothed gradient $\widetilde{\nabla}L(\theta)=\operatorname{\mathbb{E}}_{z\sim N(0,S)}[\nabla L(\theta+\lambda z)]$ which has the effect of regularizing the trace of the Hessian.

This implicit regularization effect of minibatch SGD has been shown to drastically improve generalization and reduce the number of samples necessary for supervised learning tasks [12, 13, 14]. However, the connection between the smoothed landscape and the resulting sample complexity is poorly understood. Towards closing this gap, we consider directly smoothing the loss landscape in order to efficiently learn single index models. Our main result, Theorem 1, shows that online SGD on the smoothed loss learns $w^{\star}$ in $n\gtrsim d^{{k^{\star}}/2}$ samples, which matches the correlation statistical query (CSQ) lower bound. This improves over the $n\gtrsim d^{{k^{\star}}-1}$ lower bound for online SGD on the unsmoothed loss from Ben Arous et al. [1]. Key to our analysis is the observation that smoothing the loss landscape boosts the signal-to-noise ratio in a region around the initialization, which allows the iterates to avoid the poor local minima for the unsmoothed empirical loss. Our analysis is inspired by the implicit regularization effect of minibatch SGD, along with the partial trace algorithm for Tensor PCA which achieves the optimal $d^{k/2}$ sample complexity for computationally efficient algorithms.

The outline of our paper is as follows. In Section 3 we formalize the specific statistical learning setup, define the information exponent ${k^{\star}}$, and describe our algorithm. Section 4 contains our main theorem, and Section 5 presents a heuristic derivation for how smoothing the loss landscape increases the signal-to-noise ratio. We present empirical verification in Section 6, and in Section 7 we detail connections to tensor PCA nad minibatch SGD.

There is a rich literature on learning single index models. Kakade et al. [15] showed that gradient descent can learn single index models when the link function is Lipschitz and monotonic and designed an alternative algorithm to handle the case when the link function is unknown. Soltanolkotabi [16] focused on learning single index models where the link function is $\operatorname{ReLU}(x):=\max(0,x)$ which has information exponent ${k^{\star}}=1$. The phase-retrieval problem is a special case of the single index model in which the link function is $\sigma(x)=x^{2}$ or $\sigma(x)=\absolutevalue{x}$; this corresponds to ${k^{\star}}=2$, and solving phase retrieval via gradient descent has been well studied [17, 18, 19]. Dudeja and Hsu [20] constructed an algorithm which explicitly uses the harmonic structure of Hermite polynomials to identify the information exponent. Ben Arous et al. [1] provided matching upper and lower bounds that show that $n\gtrsim d^{{k^{\star}}-1}$ samples are necessary and sufficient for online SGD to recover $w^{\star}$.

Going beyond gradient-based algorithms, Chen and Meka [21] provide an algorithm that can learn polynomials of few relevant dimensions with $n\gtrsim d$ samples, including single index models with polynomial link functions. Their estimator is based on the structure of the filtered PCA matrix $\operatorname{\mathbb{E}}_{x,y}[\mathbf{1}_{|y|\geq\tau}xx^{T}]$, which relies on the heavy tails of polynomials. In particular, this upper bound does not apply to bounded link functions. Furthermore, while their result achieves the information-theoretically optimal $d$ dependence it is not a CSQ algorithm, whereas our Algorithm 1 achieves the optimal sample complexity over the class of CSQ algorithms (which contains gradient descent).

Recent work has also studied the ability of neural networks to learn single or multi-index models [5, 6, 22, 23, 4]. Bietti et al. [5] showed that two layer neural networks are able to adapt to unknown link functions with $n\gtrsim d^{{k^{\star}}}$ samples. Damian et al. [6] consider multi-index models with polynomial link function, and under a nondegeneracy assumption which corresponds to the ${k^{\star}}=2$ case, show that SGD on a two-layer neural network requires $n\gtrsim d^{2}+r^{p}$ samples. Abbe et al. [23, 4] provide a generalization of the information exponent called the leap. They prove that in some settings, SGD can learn low dimensional target functions with $n\gtrsim d^{\text{Leap}-1}$ samples. However, they conjecture that the optimal rate is $n\gtrsim d^{\text{Leap}/2}$ and that this can be achieved by ERM rather than online SGD.

The problem of learning single index models with information exponent $k$ is strongly related to the order $k$ Tensor PCA problem (see Section 7.1), which was introduced by Richard and Montanari [8]. They conjectured the existence of a computational-statistical gap for Tensor PCA as the information-theoretic threshold for the problem is $n\gtrsim d$, but all known computationally efficient algorithms require $n\gtrsim d^{k/2}$. Furthermore, simple iterative estimators including tensor power method, gradient descent, and AMP are suboptimal and require $n\gtrsim d^{k-1}$ samples. Hopkins et al. [24] introduced the partial trace estimator which succeeds with $n\gtrsim d^{\lceil k/2\rceil}$ samples. Anandkumar et al. [25] extended this result to show that gradient descent on a smoothed landscape could achieve $d^{k/2}$ sample complexity when $k=3$ and Biroli et al. [26] heuristically extended this result to larger $k$. The success of smoothing the landscape for Tensor PCA is one of the inspirations for Algorithm 1.

Our main result is a sample complexity guarantee for Algorithm 1:

Theorem 1 uses large smoothing (up to $\lambda=d^{1/4}$) to rapidly escape the regime in which $w\cdot w^{\star}\asymp d^{-1/2}$. This first stage continues until $w\cdot w^{\star}=1-o_{d}(1)$ which takes $T_{1}=\tilde{O}(d^{{k^{\star}}/2})$ steps when $\lambda=d^{1/4}$. The second stage, in which $\lambda=0$ and the learning rate decays linearly, lasts for an additional $T_{2}=d/\epsilon$ steps where $\epsilon$ is the target accuracy. Because Algorithm 1 uses each sample exactly once, this gives the sample complexity

to reach population loss $L(w_{T})\leq\epsilon$. Setting $\lambda=O(1)$ is equivalent to zero smoothing and gives a sample complexity of $n\gtrsim d^{{k^{\star}}-1}+d/\epsilon$, which matches the results of Ben Arous et al. [1]. On the other hand, setting $\lambda$ to the maximal allowable value of $d^{1/4}$ gives:

which matches the sum of the CSQ lower bound, which is $d^{\frac{{k^{\star}}}{2}}$, and the information-theoretic lower bound, which is $d/\epsilon$, up to poly-logarithmic factors.

To complement Theorem 1, we replicate the CSQ lower bound in [6] for the specific function class $\sigma(w\cdot x)$ where $w\in S^{d-1}$. Statistical query learners are a family of learners that can query values $q(x,y)$ and receive outputs $\hat{q}$ with $\absolutevalue{\hat{q}-\operatorname{\mathbb{E}}_{x,y}[q(x,y)]}\leq\tau$ where $\tau$ denotes the query tolerance [27, 28]. An important class of statistical query learners is that of correlational/inner product statistical queries (CSQ) of the form $q(x,y)=yh(x)$. This includes a wide class of algorithms including gradient descent with square loss and correlation loss.

Using the standard $\tau\approx n^{-1/2}$ heuristic which comes from concentration, this implies that $n\gtrsim d^{\frac{{k^{\star}}}{2}}$ samples are necessary to learn $\sigma(w\cdot x)$ unless the algorithm makes exponentially many queries. In the context of gradient descent, this is equivalent to either requiring exponentially many parameters or exponentially many steps of gradient descent.

In this section we highlight the key ideas of the proof of Theorem 1. The full proof is deferred to Appendix B. The proof sketch is broken into three parts. First, we conduct a general analysis on online SGD to show how the signal-to-noise ratio (SNR) affects the sample complexity. Next, we compute the SNR for the unsmoothed objective ($\lambda=0$) to heuristically rederive the $d^{{k^{\star}}-1}$ sample complexity in Ben Arous et al. [1]. Finally, we show how smoothing boosts the SNR and leads to an improved sample complexity of $d^{{k^{\star}}/2}$ when $\lambda=d^{1/4}$.

For ${k^{\star}}=3,4,5$ and $d=2^{8},\ldots,2^{13}$ we ran a minibatch variant of Algorithm 1 with batch size $B$ when $\sigma(x)=\frac{He_{{k^{\star}}}(x)}{\sqrt{k!}}$, the normalized ${k^{\star}}$th Hermite polynomial. We set:

We computed the number of samples required for Algorithm 1 to reach $\alpha^{2}=0.5$ from $\alpha=d^{-1/2}$ and we report the min, mean, and max over $10$ random seeds. For each $k$ we fit a power law of the form $n=c_{1}d^{c_{2}}$ in order to measure how the sample complexity scales with $d$. For all values of ${k^{\star}}$, we find that $c_{2}\approx{k^{\star}}/2$ which matches Theorem 1. The results can be found in Figure 2 and additional experimental details can be found in Appendix E.

AD acknowledges support from a NSF Graduate Research Fellowship. EN acknowledges support from a National Defense Science & Engineering Graduate Fellowship. RG is supported by NSF Award DMS-2031849, CCF-1845171 (CAREER), CCF-1934964 (Tripods) and a Sloan Research Fellowship. AD, EN, and JDL acknowledge support of the ARO under MURI Award W911NF-11-1-0304, the Sloan Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994.