SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Smoothing the Landscape Boosts the Signal for SGD Optimal Sample Complexity for Learning Single Index Models
ad27@princeton.edu
eshnich@princeton.edu
We focus on the task of learning a single index model  $\sigma (w^{\star}\cdot x)$  with respect to the isotropic Gaussian distribution in  $d$  dimensions. Prior work has shown that the sample complexity of learning  $w^{\star}$  is governed by the information exponent  $k^{\star}$  of the link function  $\sigma$ , which is defined as the index of the first nonzero Hermite coefficient of  $\sigma$ . Ben Arous et al. [1] showed that  $n\gtrsim d^{k^{\star} - 1}$  samples suffice for learning  $w^{\star}$  and that this is tight for online SGD. However, the CSQ lower bound for gradient based methods only shows that  $n\gtrsim d^{k^{\star} / 2}$  samples are necessary. In this work, we close the gap between the upper and lower bounds by showing that online SGD on a smoothed loss learns  $w^{\star}$  with  $n\gtrsim d^{k^{\star} / 2}$  samples. We also draw connections to statistical analyses of tensor PCA and to the implicit regularization effects of minibatch SGD on empirical losses.
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-6]. Despite the fact that the empirical loss is nonconvex and in the worst case computationally intractable 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^{\max (1,k^{\star} - 1)}$  samples; furthermore, they present a lower bound showing that for a class of online SGD algorithms,  $d^{\max (1,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^{\max(1,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  $1k$ -tensor in  $d$  dimensions, both gradient descent and the power method require  $d^{\max(1,k-1)}$  samples, whereas more sophisticated spectral algorithms can match the computational lower bound of  $d^{\max(1,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-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) = \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-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 for  $k^{\star} > 2$ , 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  $\mathrm{ReLU}(x) \coloneqq \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) = |x|$ ; this corresponds to  $k^{\star} = 2$ , and solving phase retrieval via gradient descent has been well studied [17-20]. Dudeja and Hsu [21] 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^{\max(1, k^{\star} - 1)}$  samples are necessary and sufficient for online SGD to recover  $w^{\star}$ .
Going beyond gradient-based algorithms, Chen and Meka [22] 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  $\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, 23, 24, 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. [24, 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^{\mathrm{Leap - 1}}$  samples. However, they conjecture that the optimal rate is  $n \gtrsim d^{\mathrm{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. [25] introduced the partial trace estimator which succeeds with  $n \gtrsim d^{\lceil k/2 \rceil}$  samples. Anandkumar et al. [26] 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. [27] heuristically extended this result to larger  $k$ . The success of smoothing the landscape for Tensor PCA is one of the inspirations for Algorithm 1.
3.1 Data distribution and target function
Our goal is to efficiently learn single index models of the form  $f^{\star}(x) = \sigma (w^{\star}\cdot x)$  where  $w^{\star}\in S^{d - 1}$ , the  $d$ -dimensional unit sphere. We assume that  $\sigma$  is normalized so that  $\mathbb{E}_{x\sim N(0,1)}[\sigma (x)^2 ] = 1$ . We will also assume that  $\sigma$  is differentiable and that  $\sigma^{\prime}$  has polynomial tails:
Assumption 1. There exist constants  $C_1, C_2$  such that  $|\sigma'(x)| \leq C_1 (1 + x^2)^{C_2}$ .
Our goal is to recover  $w^*$  given  $n$  samples  $(x_1, y_1), \ldots, (x_n, y_n)$  sampled i.i.d from
$$
x _ {i} \sim N (0, I _ {d}), \quad y _ {i} = f ^ {\star} \left(x _ {i}\right) + z _ {i} \quad \text {w h e r e} \quad z _ {i} \sim N \left(0, \zeta^ {2}\right).
$$
For simplicity of exposition, we assume that  $\sigma$  is known and we take our model class to be
$$
f (w, x) := \sigma (w \cdot x) \quad \text {w h e r e} \quad w \in S ^ {d - 1}.
$$
3.2 Algorithm: online SGD on a smoothed landscape
As  $w \in S^{d-1}$  we will let  $\nabla_w$  denote the spherical gradient with respect to  $w$ . That is, for a function  $g: \mathbb{R}^d \to \mathbb{R}$ , let  $\nabla_w g(w) = (I - ww^T)\nabla g(z) \big|_{z=w}$  where  $\nabla$  is the standard Euclidean gradient.
To compute the loss on a sample  $(x,y)$ , we use the correlation loss:
$$
L (w; x; y) := 1 - f (w, x) y.
$$
Furthermore, when the sample is omitted we refer to the population loss:
$$
L (w) := \mathbb {E} _ {x, y} [ L (w; x; y) ]
$$
Our primary contribution is that SGD on a smoothed loss achieves the optimal sample complexity for this problem. First, we define the smoothing operator  $\mathcal{L}_{\lambda}$ :
Definition 1. Let  $g: S^{d-1} \to \mathbb{R}$ . We define the smoothing operator  $\mathcal{L}_{\lambda}$  by
$$
\left(\mathcal {L} _ {\lambda} g\right) (w) := \mathbb {E} _ {z \sim \mu_ {w}} \left[ g \left(\frac {w + \lambda z}{\| w + \lambda z \|}\right) \right]
$$
where  $\mu_w$  is the uniform distribution over  $S^{d-1}$  conditioned on being perpendicular to  $w$ .
This choice of smoothing is natural for spherical gradient descent and can be directly related to the Riemannian exponential map on  $S^{d - 1}$ . We will often abuse notation and write  $\mathcal{L}_{\lambda}(g(w))$  rather than  $(\mathcal{L}_{\lambda}g)(w)$ . The smoothed empirical loss  $L_{\lambda}(w;x;y)$  and the population loss  $L_{\lambda}(w)$  are defined by:
$$
L _ {\lambda} (w; x; y) := \mathcal {L} _ {\lambda} (L (w; x; y)) \quad \text {a n d} \quad L _ {\lambda} (w) := \mathcal {L} _ {\lambda} (L (w)).
$$
This is equivalent to the intrinsic definition  $(\mathcal{L}_{\lambda}g)(w)\coloneqq \mathbb{E}_{z\sim UT_w(S^{d - 1})}[\exp_w(\theta z)]$  where  $\theta = \arctan (\lambda)$ ,  $UT_{w}(S^{d - 1})$  is the unit sphere in the tangent space  $T_{w}(S^{d - 1})$ , and  $\exp$  is the Riemannian exponential map.
Algorithm 1: Smoothed Online SGD
Input: learning rate schedule  $\{\eta_t\}$ , smoothing schedule  $\{\lambda_t\}$ , steps  $T$
Sample  $w_0\sim \mathrm{Unif}(S^{d - 1})$
Sample a fresh sample  $(x_{t},y_{t})$ $\hat{w}_{t + 1}\gets w_t - \eta_t\nabla_wL_{\lambda_t}(w_t;x_t;y_t)$ $w_{t + 1}\leftarrow \hat{w}_{t + 1} / \| \hat{w}_{t + 1}\|$
3.3 Hermite polynomials and information exponent
The sample complexity of Algorithm 1 depends on the Hermite coefficients of  $\sigma$ :
Definition 2 (Hermite Polynomials). The  $k$ th Hermite polynomial  $H e_{k}:\mathbb{R}\to \mathbb{R}$  is the degree  $k$ , monic polynomial defined by
$$
H e _ {k} (x) = (- 1) ^ {k} \frac {\nabla^ {k} \mu (x)}{\mu (x)},
$$
where  $\mu(x) := \frac{e^{-\frac{x^2}{2}}}{\sqrt{2\pi}}$  is the PDF of a standard Gaussian.
The first few Hermite polynomials are  $He_0(x) = 1$ ,  $He_1(x) = x$ ,  $He_2(x) = x^2 - 1$ ,  $He_3(x) = x^3 - 3x$ . For further discussion on the Hermite polynomials and their properties, refer to Appendix A.2. The Hermite polynomials form an orthogonal basis of  $L^2(\mu)$  so any function in  $L^2(\mu)$  admits a Hermite expansion. We let  $\{c_k\}_{k \geq 0}$  denote the Hermite coefficients of the link function  $\sigma$ :
Definition 3 (Hermite Expansion of  $\sigma$ ). Let  $\{c_k\}_{k\geq 0}$  be the Hermite coefficients of  $\sigma$ , i.e.
$$
\sigma (x) = \sum_ {k \geq 0} \frac {c _ {k}}{k !} H e _ {k} (x) \quad \text {w h e r e} \quad c _ {k} = \mathbb {E} _ {x \sim N (0, 1)} [ \sigma (x) H e _ {k} (x) ].
$$
The critical quantity of interest is the information exponent of  $\sigma$ :
Definition 4 (Information Exponent).  $k^{\star} = k^{\star}(\sigma)$  is the first index  $k \geq 1$  such that  $c_k \neq 0$ .
Example 1. Below are some example link functions and their information exponents:
-  $\sigma(x) = x$  and  $\sigma(x) = \mathrm{ReLU}(x) \coloneqq \max(0, x)$  have information exponents  $k^{\star} = 1$ .
-  $\sigma(x) = x^2$  and  $\sigma(x) = |x|$  have information exponents  $k^\star = 2$ .
-  $\sigma(x) = x^3 - 3x$  has information exponent  $k^\star = 3$ . More generally,  $\sigma(x) = He_k(x)$  has information exponent  $k^\star = k$ .
Throughout our main results we focus on the case  $k^{\star} \geq 3$  as when  $k^{\star} = 1,2$ , online SGD without smoothing already achieves the optimal sample complexity of  $n \asymp d$  samples (up to log factors) [1].
Our main result is a sample complexity guarantee for Algorithm 1:
Theorem 1. Assume  $w_0 \cdot w^\star \gtrsim d^{-1/2}$  and  $\lambda \in [1, d^{1/4}]$ . Then there exists a choice of  $T_1, \eta$  satisfying  $T_1 = \tilde{O}\big(d^{k^\star - 1} \lambda^{-2k^\star + 4}\big)$  and  $\eta = \tilde{O}(d^{-k^\star / 2} \lambda^{2k^\star - 2})$  such that if we run Algorithm 1 with  $\lambda_t = \lambda$  and  $\eta_t = \eta$  for  $t \leq T_1$  and  $\lambda_t = 0$  and  $\eta_t = O\big((d + t - T_1)^{-1}\big)$  for  $t > T_1$ , we have that with high probability, after  $T = T_1 + T_2$  steps, the final iterate  $w_T$  satisfies  $L(w_T) \leq O\left(\frac{d}{d + T_2}\right)$ .
Note that the condition that  $w_0 \cdot w^\star \gtrsim d^{1/2}$  can be guaranteed with probability  $1/2$ . 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$ .
Our algorithm is online SGD on the smoothed loss  $L_{\lambda}$ :
steps where  $\epsilon$  is the target accuracy. Because Algorithm 1 uses each sample exactly once, this gives the sample complexity
$$
n \gtrsim d ^ {k ^ {\star} - 1} \lambda^ {- 2 k ^ {\star} + 4} + d / \epsilon
$$
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^* - 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:
$$
n\gtrsim \underbrace{d^{\frac{k^{\star}}{2}}} _ {\substack{\text{CSQ}\\ \text{lower bound}}} + \underbrace{d / \epsilon}_{\substack{\text{information}\\ \text{lower bound}}}
$$
which matches the sum of the CSQ lower bound, which is  $d^{\frac{k^*}{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  $|\hat{q} - \mathbb{E}_{x, y}[q(x, y)]| \leq \tau$  where  $\tau$  denotes the query tolerance [28, 29]. 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. For example, if the model is  $f_{\theta}(x)$ , then the gradient with square loss can be written as
$$
\nabla L (w; x, y) = (f _ {\theta} (x) - \underbrace {y) \nabla f _ {\theta} (x)} _ {\text {q u e r y}}).
$$
Note that the other term in the gradient only depends on the distribution of  $x \sim N(0, I_d)$  which we assume is known. However, this connection with gradient descent is only heuristic as the errors in GD are random while the errors in the SQ/CSQ framework are adversarial. However, SQ/CSQ lower bounds are frequently used to argue the existence of statistical-computational gaps in statistical learning problems [30-33]. Theorem 2 measures the tolerance needed by CSQ learners to learn  $w^{\star}$ .
Theorem 2 (CSQ Lower Bound). Consider the function class  $\mathcal{F}_{\sigma} \coloneqq \{\sigma(w \cdot x) : w \in S^{d-1}\}$ . Any CSQ algorithm using  $q$  queries requires a tolerance  $\tau$  of at most
$$
\tau \lesssim \left(\frac {\log (q d)}{d}\right) ^ {k ^ {\star} / 4}
$$
to output an  $f \in \mathcal{F}_{\sigma}$  with population loss less than  $1/2$ .
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}$ .
5.1 Online SGD Analysis
To begin, we will analyze a single step of online SGD. We define  $\alpha_{t} \coloneqq w_{t} \cdot w^{\star}$  so that  $\alpha_{t} \in [-1,1]$  measures our current progress. Furthermore, let  $v_{t} \coloneqq -\nabla L_{\lambda_{t}}(w_{t};x_{t};y_{t})$ . Recall that the online SGD update is:
$$
w _ {t + 1} = \frac {w _ {t} + \eta_ {t} v _ {t}}{\| w _ {t} + \eta_ {t} v _ {t} \|} \Rightarrow \alpha_ {t + 1} = \frac {\alpha_ {t} + \eta_ {t} (v _ {t} \cdot w ^ {\star})}{\| w _ {t} + \eta_ {t} v _ {t} \|}.
$$
Using the fact that  $v_{t} \perp w_{t}$  and  $\frac{1}{\sqrt{1 + x^2}} \approx 1 - \frac{x^2}{2}$  we can Taylor expand the update for  $\alpha_{t+1}$ :
$$
\alpha_ {t + 1} = \frac {\alpha_ {t} + \eta_ {t} (v _ {t} \cdot w ^ {\star})}{\sqrt {1 + \eta_ {t} ^ {2} \| v _ {t} \| ^ {2}}} \approx \alpha_ {t} + \eta_ {t} (v _ {t} \cdot w ^ {\star}) - \frac {\eta_ {t} ^ {2} \| v _ {t} \| ^ {2} \alpha_ {t}}{2} + \mathrm {h . o . t .}
$$
As in Ben Arous et al. [1], we decompose this update into a drift term and a martingale term. Let  $\mathcal{F}_t = \sigma\{(x_0, y_0), \ldots, (x_{t-1}, y_{t-1})\}$  be the natural filtration. We focus on the drift term as the martingale term can be handled with standard concentration arguments. Taking expectations with respect to the fresh batch  $(x_t, y_t)$  gives:
$$
\mathbb {E} [ \alpha_ {t + 1} | \mathcal {F} _ {t} ] \approx \alpha_ {t} + \eta_ {t} \mathbb {E} [ v _ {t} \cdot w ^ {\star} | \mathcal {F} _ {t} ] - \eta_ {t} ^ {2} \mathbb {E} [ \| v _ {t} \| ^ {2} | \mathcal {F} _ {t} ] \alpha_ {t} / 2
$$
so to guarantee a positive drift, we need to set  $\eta_t \leq \frac{2\mathbb{E}[v_t \cdot w^\star|\mathcal{F}_t]}{\mathbb{E}[\|v_t\|^2|\mathcal{F}_t] \alpha_t}$  which gives us the value of  $\eta_t$  used in Theorem 1 for  $t \leq T_1$ . However, to simplify the proof sketch we can assume knowledge of  $\mathbb{E}[v_t \cdot w^\star |\mathcal{F}_t]$  and  $\mathbb{E}[\|v_t\|^2 |\mathcal{F}_t]$  and optimize over  $\eta_t$  to get a maximum drift of
$$
\mathbb {E} [ \alpha_ {t + 1} | w _ {t} ] \approx \alpha_ {t} + \frac {1}{2 \alpha_ {t}} \cdot \underbrace {\frac {\mathbb {E} [ v _ {t} \cdot w ^ {\star} | \mathcal {F} _ {t} ] ^ {2}}{\mathbb {E} [ \| v _ {t} \| ^ {2} | \mathcal {F} _ {t} ]}} _ {\text {S N R}}.
$$
The numerator measures the correlation of the population gradient with  $w^{\star}$  while the denominator measures the norm of the noisy gradient. Their ratio thus has a natural interpretation as the signal-to-noise ratio (SNR). Note that the SNR is a local property, i.e. the SNR can vary for different  $w_{t}$ . When the SNR can be written as a function of  $\alpha_{t} = w_{t}\cdot w^{\star}$ , the SNR directly dictates the rate of optimization through the ODE approximation:  $\alpha^{\prime}\approx \mathrm{SNR} / \alpha$ . As online SGD uses each sample exactly once, the sample complexity for online SGD can be approximated by the time it takes this ODE to reach  $\alpha \approx 1$  from  $\alpha_0\approx d^{-1 / 2}$ . The remainder of the proof sketch will therefore focus on analyzing the SNR of the minibatch gradient  $\nabla L_{\lambda}(w;x;y)$ .
5.2 Computing the Rate with Zero Smoothing
When  $\lambda = 0$ , the signal and noise terms can easily be calculated. The key property we need is:
Property 1 (Orthogonality Property). Let  $w, w^{\star} \in S^{d - 1}$  and let  $\alpha = w \cdot w^{\star}$ . Then:
$$
\mathbb {E} _ {x \sim N (0, I _ {d})} [ H e _ {j} (w \cdot x) H e _ {k} (w ^ {\star} \cdot x) ] = \delta_ {j k} k! \alpha^ {k}.
$$
Using Property 1 and the Hermite expansion of  $\sigma$  (Definition 3) we can directly compute the population loss and gradient. Letting  $P_w^\perp \coloneqq I - ww^T$  denote the projection onto the subspace orthogonal to  $w$  we have: