Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
Authors: Alex Damian
Eshaan Nichani
Rong Ge
Jason D. Lee
Venue: NeurIPS 2023 oral (2023)
Keywords: statistical learning
learning theory
single index model
gradient descent
stochastic gradient descent
TL;DR: We prove that online SGD on a smoothed loss achieves optimal sample complexity for learning single index models.

Abstract: 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. (2021) 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.

Smoothing the Landscape Boosts the Signal for SGD Optimal Sample Complexity for Learning Single Index Models
Alex Damian Princeton University ad2 11277 weprinceton....mmmm edu
Eshaan Nichani Princeton University
Rong Ge Duke University ronggelcs. duke edu
Jason D. Lee
Princeton University
jasomlee@prinnetonnccm
Abstract
We focus the task of learning single index model aw" .x) with respect the isotropic Gaussian distribution in d dimensions. Prior work has shown that the sample complexity of learning w* is governed by the information exponent ks of the link function which is defined as the index of the first nonzero Hermite coefficient of Ben Arous et al. | showed that n Z d samples suffice for learning w* and that this is tight for online SGD. However, the CSQ lower bound for gradient based methods only shows that n Z d²² /2 samples are necessary. In this work, close the gap between the upper and lower bounds by showing that online SGD smoothed loss learns w" with 2 d /2 samples. We also draw connections to statistical analyses of tensor PCA and to the implicit regularization effects of minibatch SGD empirical losses.
Introduction
Gradient descent-based algorithms are popular for deriving computational and statistical guarantees for a number of nigh-dimensional statistical learning problems 2 3 4 Despite the fact that the empirical loss is nonconvex and in the worst computationally intractible 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 overpower the noise in complexity for learning with gradient descent is determined by the landscape of the empirical loss.
uniform sense. The sample
One setting in which the empirical loss landscape showcases rich behavior is that of learning single-index model. Single index models are target functions of the form (*(x) a a 2), where "" E 5d-1 is the unknown relevant direction and o the known link function. When the covariates drawn from the standard d-dimensional Gaussian distribution, the shape of the loss landscape is governed by the information exponent ks of the link function a, 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 a_ax((1,* samples; furthermore, they present lower bound showing that for a class of online SGD algorithms, d_ samples are indeed necessary.
However, gradient descent can be suboptimal for various statistical learning problems, it only relies on local information the loss landscape and is thus prone to getting stuck local minima. For learning single index model, the Correlational Statistical Query (CSQ) lower bound only requires
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

dmaax(1, 1,k= (2) samples to recover vt 66 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 k-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_mx(11,A/2) samples.
In light of the lower bound from (. it seems hopeless for gradient- gradient-based algorithm to match the CSQ lower bound for learning single-index models. []] 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 discretization to gradient flow, but rather that has an additional implicit regularization effect. Specifically,
[9-I show that over short periods of time, SGD converges to a quasi-stationary distribution (0, AS) where 0 is an initial reference point, S is a matrix depending on the Hessian and the noise covariance and À = Yh8 measures the strength of the noise where 1) is the learning rate and B the batch size, The resulting long term dynamics therefore follow the smoothed gradient (() = Eg (0.5)) [ + >=)] which has the effect of regularizing the trace of the Hessian.
This implicit regularization effect of minibatch SGD has been shown to drastically improve general- ization and reduce the number of samples necessary for supervised learning tasks [12444. 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 [ shows that for k* > 2, online SGD on the smoothed loss learns w² in n 2 dd /2 samples, which matches the correlation statistical query (CSQ) lower bound. This improves over the n 2 d lower bound for online SGD on the unsmoothed loss from Ben Arous et al. M. Key to our analysis is the observation that smoothing the loss landscape boosts the tggnal-t-noosee ratio in a region around the initialization, which allows the iterates to avoid the poor local minima for the unsmoothed empirical loss. Our analysis inspired by the implicit regularization effect of minibatch SGD, along with the partial trace algorithm for Tensor PCA which achieves the optimal d²/2 sample complexity for computationally efficient algorithms.
The outline of our paper is as follows. In Section 3 setup, define the information exponent k", and describe our algorithm. Section contttin our main theorem, and Section Doctionn]prree 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.
we formalize the specific statistical learning
2 Related Work
There is a rich literature on learning single index models. Kakade et al. [1] 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 ReLU(x))
max(0, x) which has information exponent k-* = 1. The phase-retrieval problem is a special case of the single index
model in which the link function is a (x) x² or a (xx) |ar|: this corresponds to A= 2, and solving
phase retrieval via gradient descent has been well studied (17220]. Dudeja and Hsu [2] constructed an algorithm which explicitly uses the harmonic structure of Hermite polynomials to identify the information exponent. Ben Arous et al. M provided matching upper and lower bounds that show that n 2 _==x((1.6' -1) samples are necessary and sufficient for online SGD to recover w=.
Going beyond gradient- based algorithms, Chen and Meka 22) provide an algorithm that can learn polynomials of few relevant dimensions with n 2 d samples, including single index models with polynomial link functions. Their estimator is based
the structure of the filtered PCA matrix Ez.,, [ xxx") which relies on the heavy tails of polynomials. In particular, this upper bound does not apply bounded link functions. Furthermore, while their result achieves the information- theoretically optimal d dependence it is not a CSQ algorithm, whereas our Algorithm [ Jachieves 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] 2, 24 4 Bietti al. (5) showed that two layer neural networks are able to adapt to unknown link functions with 2 d= samples. Damian et al. 6 consider multi-index models with polynomial
2

---

link function, and under nondegeneracy assumption which corresponds to the k* = 2 case, show that SGD on a two-layer neural network requires n 2 d' + r samples. Abbe al. [24 4 provide generalization of the information exponent called the leap. They prove that in some settings, SGD can learn low dimensional target functions with n Z dLeap-1 samples. However, they conjecture that the optimal rate Z dLewp/2 and that this can be achieved by ERM rather than online SGD.
The problem of learning single index models with information exponent l is strongly related to the order k Tensor PCA problem (see Section| 7.I.. which was introduced by Richard and Montanari [8]. They conjectured the existence of cmmputational-staiiiiiii theoretic threshold for the problem is n Z d, but all known computationally efficient algorithms require n dk//2, Furthermore, simple iterative estimators including tensor power method, gradient descent, and AMP are suboptimal and require 2 d==1 samples. Hopkins et al. [25] introduced the
gap for Tensor PCA as the information-
partial trace estimator which succeeds with n Z dTk/21 samples. Anandkumar et al. [26 extended this
result to show that gradient descent on a smoothed landscape could achieve d²/²2 sample complexity when & = 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
3 Setting
3.1 Data distribution and target function
Our goal is efficiently learn single index models of the form (*) o(w* x) where wt E 5d-1, the d-dimensional unit sphere. We assume that is normalized so that (0 1. We will also assume that a is differentiable and that of has polynomial tails:
Assumption 1. There exist constants C1 C2 such that (x)( s C1(1 c
Our goal is to recover w" given samples (æ1, v1). (In, Vo.) sampled i.dd from
Ij ~ (O, y: = (x;) + + where 24 ~ N (0,
For simplicity of exposition, assume that is known and we take our model class to be
[ww x) = (w x) where w E Sd-1
3.2 Algorithm: online SGD smoothed landscape
As w E Sit-1 will let V w denote the spherical gradient with respect to w. That is, for function g IRd + R. let V (w)) = (I - ww") w g ==w where V is the standard Euclidean gradient.
To compute the loss on a sample (x, v). we use the correlation loss:
(wwr;y) 3) == f(w,z)y.
Furthermore, when the sample omitted we refer to the population loss:
L(w)= == Er,y ( y)]
Our primary contribution is that SGD on smoothed loss achieves the optimal sample complexity for this problem. First, we define the smoothing operator Lx:
Definition 1. Let g 5d RR We define the smoothing operator Lx by
w + Az
Xz )
(Lxg)(w)= = Ez g
where uw is the uniform distribution over S 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 5d-1 We will often abuse notation and write Lx (x((w)) rather than (w). The smoothed empirical loss Lx Lx(w; (w;;;;) y) and the population loss Lxww) are defined by:
(w x; y) == Cx(L x; v)) and Lx(w) = Lx(Lww))
'This is equivalent to the intrinsic definition (Cxg) B. (se-1) lexp... where 0 arctan(A), UTw (Sd-1) is the unit sphere the tangent space Ti (Sd-1) and exp is the Riemannian exponential map.
3

---

Our algorithm is online SGD on the smoothed loss Lx:
Algorithm 1: Smoothed Online SGD
Input: learning rate schedule {nt, smoothing schedule {Ar]. steps Sample Wg Unif(5*-1) for i = 0 to T - do
Sample fresh sample (24,
dit+1 f- Wi V LA, (wy; 313 W++1 114+1/10000
end
3.3 Hermite polynomials and information exponent
The sample complexity of Algorithm dependss the Hermite coefficients of a:
Definition 2 (Hermite Polynomials). The kth Hermite polynomial He R - IR is the degree k. monic polynomial defined by
(x) p(x))
m(x)
where u(x)
the PDF of standard Gaussian.
2/2
The first few Hermite polynomials are Heo(x) L, (e) He2(z)) = 1,
x² 3r. For further discussion on the Hermite polynomials and their properties, refer to Appendix A.2
The Hermite polynomials form orthogonal basis of (1² any function in [²) admits Hermite expansion. We let
{ck (okk22 denote the Hermite coefficients of the link function o:
Definition 3 (Hermite Expansion of σ). Let (ce)ao be the Hermite coefficients of i.e.
(w)
Σ 820
Ck Hex(z)
where Ck =BoooNoooooo((()Hee))))
R!
The critical quantity of interest is the information exponent of a: Definition 4 (Information Exponent). k- = k-" (a)) is the first index k 2 such that C + 0. Example 1. Below are some example link functions and their information exponents:
a (x)) and a (x) ReLU(x)) max((, x)) have information exponents k 1.
(ι) x² and (xx) = |x| have information exponents k*
(x)) x²3 3x has information exponent ke More generally, (xx) Heg (x) has information exponent k* = k.
Throughout our main results we focus on the case k- 2 3 as when 1, 2, online SGD without smoothing already achieves the optimal sample complexity of x d samples (up to log factors) MI]
4 Main Results
Our main result sample complexity guarantee for Algorithm I
Theorem 1. Assume we w 2 d-2 and À [1, d//4). Then there exists choice of T T,," satisfying
Ti = Ō(d - 1 -12-- 24*+4) and 17 02 x -2) such that if we run Algorithm I with X
and 1) = for < Ti and A = and π) = ((dd((+)))))) + -T;) for > Tr. we have that with high probability, after T Ti + T2 steps, the final iterate satisfies L(wr) < 0 7
Note that the condition that Wg = 2 1/2 can guaranteed with probability 1/2. Theorem uses large smoothing (up to = d2/4) to rapidly escape the regime which w =² x 1/2 This first 0a(1) which takes Ti = O(d="/2 /2) steps when X = =d/11 The second stage, in which À 0 and the learning rate decays linearly, lasts for an additional T3 d/c
stage continues until w w²

---

steps where e is the target accuracy. Because Algorithm ! uses each sample exactly once, this gives the sample complexity
d - 2 +4 d/c
to reach population loss L(wr)= S Setting À (()) equivalent to zero smoothing and gives sample complexity Z d -1 + d/e, which matches the results of Ben Arous et [ On the other hand, setting À to the maximal allowable value of d gives:
d/c information lower bound
CSQ lower bound
which matches the sum of the CSQ lower bound, which is , and the informattin--heorrtie lower bound, which is d/e, up poly-logarithmic factors.
To complement Theorem T we replicate the CSQ lower bound [6] for the specific function class oww x)) where
S dd11_ Statistical query learners are family of learners that query values (2, y) and receive outputs with l - [x,,, y)]l S where denotes the query tolerance 28 [29]. An important class of statistical query learners is that of correlational/ineer product statistical queries (CSQ) of the form q(x, y) yh(x). This includes wide class of algorithms including gradient descent with square loss and correlation loss. For example, the model is fo(),, then the gradient with square loss can be written as
V z, v) (fo(z) Vo(
query
Note that the other term in the gradient only depends the distribution of a (O,, La) which we assume is known. However, this connection with gradient descent is only heuristic as the errors in GD random while the errors in the SQ/CSQ framework adversarial. However, SQ/CSQ lower bounds are frequently used to argue the existence of cedistcall-ooputatiinn
gaps in statistical learning problems (30 33 Theorem measures the tolerance needed by CSQ learners to learn w*.
Theorem (CSQ Lower Bound). Consider the function class F, == {a(w x) w E Sd-1). Any CSQ algorithm using q queries requires tolerance of at most
llo(yd)
k" k4/4
TS
output f E Fo with population loss less than 1/2.
Using the standard r Rs -1/2 heuristic which comes from concentration, this implies that n 2 d samples necessary to learn ( xx) 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.
Proof Sketch
In this section we highlight the key ideas of the proof of Theorem The full proof is deferred to Appendix B. The proof sketch is broken into three parts. First, we conduct 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 () 0) to heuristically rederive the dh--1 sample complexity in Ben Arous et al. [I]. Finally, we show how smoothing boosts the SNR and leads to an improved sample complexity of dk /2 when X = 1/4
5.1
Online SGD Analysis
To begin, we will analyze a single step of online SGD. We define a == wi w² so that a [-1,1] measures our current progress. Furthermore, let v, - LL, (w, Izi B.). Recall that the online SGD update is:
W/ + Dt W++11
+ (1 w*)
a1+1
||,, + vevell
|w 1/
5

---

Using the fact that I. and VI+z7
x² 2
can Taylor expand the update for 04+11
01++1 at + (12 w*)
Ti v eellee1 h.o.t.
2 ar + ne(ue - w²)
+ n/ well
2
As in Ben Arous al. [T]. we decompose this update into drift term and martingale term. Let Fi of (To, yo)...., (#t-1, 1/-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 (Iz, Ut) gives:
E| Bla++1 Fi] Rs at + The Elu, w [FFF - Th ||
so to guarantee positive drift, we need to set The <
which gives us the value of y assume knowledge of Elve w" [F.] and B[[|=]]] 3 [FF) and optimize over f/. to get maximum drift of
2 [r
B||l,|
used in Theorem [or < Ti. However, to simplify the proof sketch
Eiee w [F
B[ox+1/w0] so ar +
2a,
E B|
SNR
The numerator measures the correlation of the population gradient with w² while the denominator measures the norm of the noisy gradient. Their ratio thus has 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 wt- When the SNR be written as function of optimization through the ODE approximation: a' Re SNR/o. As online SGD uses each sample exactly once, the sample complexity for online SGD be approximated by the time takes this ODE reach analyzing the SNR of the minibatch gradient LL x;
wy w². the SNR directly dictates the rate of
from og R d/1/2 The remainder of the proof sketch will therefore focus
5.2 Computing the Rate with Zero Smoothing
When À 0, the signal and noise terms can easily calculated. The key property we need is: Property (Orthogonality Property). w, w² E Sd and let w w². Then:
B (e H)Hee("" x)] = dyskla..
Using Property I and the Hermite expansion of a (Definition 3) we can directly compute the population loss and gradient. Letting P+ I www" denote the projection onto the subspace orthogonal to w we have:
L(w)
a
and VL(w) -(P,)) (P)) w²)
k-1
R20
As a & throughout most of the trajectory. the gradient dominated by the first nonzero Hermite coefficient so up to constants: E[u-w*] = L ze k" Similarly, standard concentration argument shows that because is random vector in d dimensions where each coordinate is O()), |][lu|]] z d. Therefore the SNR is equal to 2k(k--1) / so with optimal learning rate schedule,
Ejat+1|Fi] ~ a + 2*"-3 -3 /d.
This can be approximated by the ODE -3" /d. Solving this ODE with the initial ao x xd11//2 gives that the escape time is proportional to da, 2(h result of Ben Arous et al.
-1 -11
d which heuristically re-derives the
5.3 How Smoothing boosts the SNR
Smoothing improves the sample complexity online SGD by boosting the SNR of the stochastic gradient V L(ww, y).

---

SNR with A d
k odd
SNR
k even
No Smoothing
d
d-1/2
d-1/4
Figure 1: When À = d1/4, smoothing increases the SNR until a Ad-1/2 = 1/4/4
Computing the Signal Recall that the population loss approximately equal to 1 c a where k- is the first nonzero Hermite coefficient of Isolating the dominant a term and applying the smoothing operator Lx, we get:
[(
(1u= Az
(a(ax)
[ + Azıl
Because - w and ||||| we have that 11 + + vI+ X2 Re A. Therefore,
(x(a")) 2o A k²' Esppa (a + A(= 10 x²*))* =
=\- - -=
)aa
a -j i2 E Exnpa [( w0" y}.
Now because the terms where is odd disappear. Furthermore, for random 12 w== O(d-1/3 Therefore reindexing and ignoring all constants we have that
1-ww)=1-2xc) x ~-
25 (22/d))
1=
Differentiating gives that
X2
Biv w*] ww- VuLx(a*) X-1
2x2
As this is geometric series, it is either dominated by the first the last term depending whether a 2 1d2 or a < 1d2 Furthermore, the last term either d =1_1 if is odd or f -" k is even. Therefore the signal term is:
a 2 Ad-1/2 1/2
Ele w0] 1-
a 1/2 and k* odd
a Ad-1/2 and k* even
Computing the Noise Recall that L2 L,,,, = 1 - yo x). Differentiating through the smoothing operator gives:
Lw a, u) = V. Lx(a(w x)) R Lx (o'(o (w- x)).
We know that with high probability, y Õ(1) and || = O(vi) so suffices to bound Lx Lx (w((w)))) The variance of this term is equal to:
E. L =)) Ex Esapa o' (+++t Au
]]
(V+
To compute this expectation, we will create an i.i.d. copy z' of and rewrite this expectation as:
Ex(Cx(o'(ww x)) = Er E E. 00
(o'(vi (
w Az
)
0'
w Az'
)1}
VI+ X2
7

---

Now we can swap the expectations and compute the expectation with respect to first using the Hermite expansion of a. As the first nonzero Hermite coefficient of a' is ke this variance is approximately equal to the correlation between =+
and m+ raised to the h power:
V1+A2
v1+x²
E2 (L = E a' ( W=
(VI+X
d2
W
Az'
---17
1 [++
-1-17
Es efppra meptas
1+
V1+
As 2, z' are random unit vectors, their inner product of order d-1/2, Therefore when < d the first term in the numerator is dominant and when À >> d//4, the second term is dominant. Combining these two regimes gives that the variance is of order min(), d' d'/4jkk-1 which motivates our choice of À d Combining this with y ((1) and ||ж||| O(vd) gives that for À I d'/4, the variance of the noise is bounded by B[[|||]] S dA-2k*, 2 Note that in the high signal regime (a 2 Ad- both the signal and the noise are smaller by factors of A which cancel when computing the SNR. However, when < Ad- the smoothing shrinks the noise faster than shrinks the signal, resulting in an overall larger SNR. Explicitly,
21
2 Ad-
E/u w1 B[]lu||]] d
1
SNR
(22/d)
a Xd- and k* is odd a < Xd- 1/2 and k** is even
a² Q2(22/d)**==
For a 2 1/2 smoothing does not affect the SNR. However, when S 1/21/2, smoothing greatly increases the SNR (see Figure D
Solving the ODE: a' = SNR/a gives that it takes T' ze d 22* steps to converge to s 1 from 2 d-1/2, Once a x 1, the problem is locally strongly convex, so we can decay the learning rate and use classical analysis of strongly-convex functions to show that a 2 1 with additional d/e steps, from which Theorem | follows.
6 Experiments
... we ran minibatch variant of Algorithm with batch size B when or (x) = (x) the normalized k*th Hermite polynomial. We set: V
For 1* 3, 4, and 28, ..... 213
Bd /2 -2 B min (0.10* /2 X2) 233-2k 24 (1122992)
A d/4
1000k*
We computed the number of samples required for Algorithm to reach a² = 0.5 from d- 1/2 and we report the min, mean, and max over 10 random seeds. For each we fit a power law of the form cjd2² in order to measure how the sample complexity scales with d. For all values of k*, we find that 22 Rs k* /2 which matches Theorem The results can be found Figure 2 and additional experimental details can be found in Appendix E
7 Discussion
7.1 Tensor PCA
We next outline connections to the Tensor PCA problem. Introduced in [8), the goal of Tensor PCA to recover the hidden direction w" Id-1 from the noisy k-tensor T, E (Rd)@k given by
Tw (w (w))ak
8k
Z. VT
where Z E (Rdy0sk is a Gaussian noise tensor with each entry drawn i.i.d from 1).
The Tensor PCA problem has garnered significant interest as it exhibits staristical iiniiicallcomuaationnn gap. w" is information theoretically recoverable when n Z d. However, the best polynomial timmem algorithms require n Z dh/2; this lower bound has been shown to be tight for various notions of
"This normalization is equivalent to the original
normalization by setting -ddd.
BV8

---

Online SGD Sample Complexity, A = d'/4
231
k 3, best fit: d
(sample complexity
k best fit: d2.1
2
k hest fit: d2.5
227
2
2
2
2
2
2
2
2
2
d (input dimension)
Figure 2: For 1= 3, 4, 5, Algorithm finds w² with x dk samples. The solid lines and the shaded areas represent the mean and min/max values over 10 random seeds. For each curve, we also fit power law n C₁ d represented by the dashed lines. The value C2 is reported in the legend.
hardness such as CSQ SoS lower bounds 34 25 35) Tensor PCA also exhibits a gap between spectral methods and iterative algorithms. Algorithms that work in the x d²/2 regime rely on unfolding contracting the tensor X, or on semidefinite programming relaxations 34 225 On the other hand, iterative algorithms including gradient descent, power method, and AMP require a much larger sample complexity of n Σ d [40. The suboptimality iterative algorithms believed to be due to bad properties of the landscape of the Tensor PCA objective in the region around
the initialization. Specifically [41 42 argue that there are exponentially many local minima near the equator in the < dh regime. To overcome this, prior works have considered "smoothed" versions of gradient descent, and show that smoothing recovers the computationally optimal SNR in the k 3 case (26] and heuristically for larger & [27
7.1.1 The Partial Trace Algorithm
The smoothing algorithms above are inspired by the following partial trace algorithm for Tensor PCA
25
which can be viewed as Algorithm I in the limit as + 30 (26). Let Tn = (w*)@k @k + /n Z.
Then we will consider iteratively contracting indices of T until all that remains is a vector (if is odd) matrix (if k is even). Explicitly, we define the partial trace tensor by
M,, = Tn (1 (10
IR ddd IRd
Io is even k is odd.
When A* is odd, directly return Mn estimate for w* and when As is even return the top eigenvector of M,. A standard concentration argument shows that this succeeds when 2 dlk/2] Furthermore, this can be strengthened to d by using the partial trace vector as warm start for gradient descent tensor power method when le is odd (26] 27].
7.1.2 The Connection Between Single Index Models and Tensor PCA
For both tensor PCA and learning single index models, gradient descent succeeds when the sample complexity is n - [40 II On the other hand, the smoothing algorithms for Tensor PCA
227 26 succeed with the computationally optimal sample complexity of n = dt2 Our Theorem shows that this smoothing analysis can indeed be transferred to the single-index model setting.
In fact, one can make direct connection between learning single-index models with Gaussian covariates and Tensor PCA. Consider learning a single-index model when oxx)
Hea(#) the
normalized kth Hermite polynomial. Then minimizing the correlation loss is equivalent maximizing

---

the loss function:
Hen(w (w xi) = Tw) where T, =
He,
La (n(w) = n1 i=1 Vk!
Σ y
y6
72
Vk!
Here (x) E (R&))* denotes the kth Hermite tensor (see Appendix A.2 for background on Hermite polynomials and Hermite tensors). In addition, by the orthogonality of the Hermite tensors, Exy [Tn]] = (w*) so
decompose In (w* ysk + Zn, where by standard concentration, each entry of Zn is order n 1/2 We therefore directly apply algorithms for Tensor PCA to this problem. We remark that this connection is a heuristic, as the structure of the noise in Tensor PCA and our single index model setting different.
7.2 Empirical Risk Minimization on the Smoothed Landscape
Our main sample complexity guarantee, Theorem T is based on tight analysis of online SGD (Algorithm[]) D in which each sample used exactly once. One might expect that if the algorithm were allowed to reuse samples, is standard practice in deep learning, that the algorithm could succeed with fewer samples. particular, Abbe et al. 4 conjectured that gradient descent on the empirical loss L.(w) = n Σi=1 [ww; xzj 84) would succeed with n 2 d samples.
Our smoothing algorithm Algorithm I can be directly translated to the ERM setting to learn w" with 2 d samples. We can then Taylor expand the smoothed loss in the large À limit:
X-2
(x(L(w)) = Bs-Uni((se-)) (n( + - VLn(2)
w" LL + +(x
2
As À + oc, gradient descent on this smoothed loss will converge to EE Unit(( 1) (vn(())) which is equivalent the partial trace estimator for k- odd (see Section 7.1). .. If ** is even, this first term is zero in expectation and gradient descent will converge to the top eigenvector of Egw / L,)),, which corresponds the partial trace estimator for ks" even. Mirroring the calculation for the partial trace estimator, this succeeds with 2 dlk"/t [2] samples. When k- is odd, this can be further improved d by using this estimator as warm start from which run gradient descent with À 0 in Anandkumar et al. (26). Biroli et al. 277]
7.3 Connection Minibatch SGD
recent line of works has studied the implicit regularization effect of stochastic gradient descent I1 10]. The key idea is that over short timescales, the iterates converge quasi-stationary distribution N(0, AS) where S = I depends on the Hessian and the noise covariance at and is proportional the ratio of the learning rate and batch size. As result, over longer periods of time SGD follows the smoothed gradient of the empirical loss:
L (w)
(n + Xz)].
We therefore conjecture that minibatch SGD is also able to achieve the optimal n 2 d sample complexity without explicit smoothing if the learning rate and batch size are properly tuned.
Acknowledgments and Disclosure of Funding
AD acknowledges support from NSF Graduate Research Fellowship. EN acknowledges support from a National Defense Science & Engineering Graduate Fellowship. RG is supported by NSF Award DMS-203 849, CCF-1845177 (CAREER), CCF- (Tripods) and Sloan Research Fellowship. AD, EN, and JDL acknowledge support of the ARO under MURI Award W91 INF-II- 1-0304, the Sloan Research Fellowship. NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262. ONR Young Investigator Award, and NSF CAREER Award 2144994.
References
[1] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22()):478-44833 2021.
10

---

[2] Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. Advances neural information processing systems, 29, 2016.
[3] Tengyu Ma. Why do local methods solve nonconvex problems?, 2020.
[4] Emmanuel Abbe, Enric Boix-Adserà, and Theodor Misiakiewicz. Sgd learning neural networks: lap complexity and saddle-to-saddle dynamics. arXiv, 2023. URL https://
wexiv.org//bb.. 2tt.:/aa.b222211111 11055
[5] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[6] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
[7] Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics, 46:2747-2774, 2018.
[8] Emile Richard and Andrea Montanari. A statistical model for tensor pca. In Advances in Neural Information Processing Systems, pages 2897 2905, 2014.
[9] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an umsttiin-uleenbee like process. Conference on Learning Theory, pages 483-513, 2020.
[10] Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[11] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? - mathematical framework. In International Conference on Learning Representations, 2022.
[12] Christopher J Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism neural network training. arXiv preprint arXiv: 181 I.03600, 2018
[13] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818 2818-2826, 2016.
[14] Yeming Wen, Kevin Luk, Maxime Gazeau, Guodong Zhang, Harris Chan, and Jimmy Ba. Interplay between optimization and generalization of stochastic gradient descent with covariance noise. arXiv preprint /r02vv9902.08200 2019.
[15] Sham M Kakade, Varun Kanade, Ohad Shamir, and Adam Kalai. Efficient learning of general- ized linear and single index models with isotonic regression. Advances in Neural Information Processing Systems, 24, 2011.
[16) Mahdi Soltanolkotabi. Learning relus via gradient descent. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[17] Emmanuel Candès, Xiaodong Li, and Mahdi Soltanolkotabi. Phase retrieval via wirtinger flow: Theory and algorithms. IEEE Transactions on Information Theory, 61(4):1985-2007,
2015. doi: 10. 10.1109//T222015233992.
[18] Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. Gradient descent with random initializa- tion: fast global convergence for nonconvex phase retrieval. Mathematical Programming, 176 (1):5-37, 2019.
[19] Ju Sun, Qing Qu, and John Wright. geometric analysis of phase retrieval. Foundations of Computational Mathematics, 18(5):11310110 198, 2018.
11

---

[20] Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initial- ization solves non-smooth, non-convex phase retrieval, 2019.
[21] Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Confer- ence On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 1887-1930. PMLR. 06 09 Jul 2018. URL
dude 8a. html
[22] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 1161-1227. PMLR, 09-12 Jul 2020. URL htt.hhttt///ooooete
www.innss..........m.....e. prwss/preesvvllve 25 wwtp5//hhhee...accccmm
[23] Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang. Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How gradient step improves the representation. In Advances in Neural Information Processing Systems (NeurIPS). 2022.
[24] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782 4887. PMLR, 2022.
[25] Samuel B. Hopkins, Tselil Schramm, Jonathan Shi, and David Steurer. Fast spectral algorithms from sum-of-squares proofs: Tensor decomposition and planted sparse vectors, Proceedings of the Forty-Eighth Annual ACM Symposium Theory of Computing, STOC 16, page 191, New York. NY, USA. 2016. Association for Computing Machinery. ISBN 9781450341325. doi: 2897529
10.1145/2997518.225 S88 2897529.
http::tts////ectt http///i..org/////// 10.
[26] Anima Anandkumar, Yuan Deng, Rong Ge, and Hossein Mobahi. Homotopy analysis for tensor pea. In Satyen Kale and Ohad Shamir, editors, Proceedings of the 2017 Conference on Learning Theory, volume 65 of Proceedings of Machine Learning Research, pages 79-104. PMLR. 07-10 mlr www.sppssssee........ee / v anandk uma 7 html
Jul 2017. URL httL:tttss/spccceeeeee
[27] Giulio Biroli, Chiara Cammarota, and Federico Ricci-Tersenghi. How to iron out rough landscapes and get optimal performances: averaged gradient descent and its application to tensor pea. Journal of Physics A: Mathematical and Theoretical, 53(17):174003 apr 2020. doi: 11..088/17511812/////0 /ttp///p/////coo.....o/////.///. ab7blf
[28] Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans. Super- polynomial lower bounds for learning one-layer neural networks using gradient descent. arXiv preprint arXiv:2006. 1001, 2020.
[29] Ilias Diakonikolas, Daniel M Kane, Vasilis Kontonis, and Nikos Zarifis. Algorithms and sq lower bounds for pac learning one-hidden-layer relu networks. In Conference on Learning Theory, pages 151 1539, 2020.
[30] Vitaly Feldman, Elena Grigorescu, Lev Reyzin, Santosh Vempala, and Ying Xiao. Statistical algorithms and lower bound for detecting planted clique, 2016.
[31] Ilias Diakonikolas, Daniel M. Kane, and Alistair Stewart. Statistical query lower bounds for robust estimation of high-dimentn ligh--imensional gaussians and gaussian mixtures, 2017.
[32] Surbhi Goel, Aravind Gollakota, and Adam Klivans. Statistical-query lower bounds via functional gradients, 2020.
[33] Rishabh Dudeja and Daniel Hsu. Statistical query lower bounds for tensor pea, 2021.
[34] Samuel B. Hopkins, Jonathan Shi, and David Steurer. Tensor principal component analy- sis via sum-of-square proofs. In Peter Grünwald, Elad Hazan, and Satyen Kale, editors, Proceedings of The 28th Conference on Learning Theory, volume 40 of Proceedings of Ma- chine Learning Research, pages 956-1006, Paris, France, 03-06 Jul 2015. PMLR. URL. https: http:://prcceeeeoo https::/pproeeednggs. wwtpp/pptppppppeepppeteee html
12

---

[35] Dmitriy Kunisky, Alexander S. Wein, and Afonso S. Bandeira. Notes on computational hardness of hypothesis testing: Predictions using the low-degree likelihood ratio, 2019.
[36] Afonso S Bandeira, Ahmed El Alaoui, Samuel Hopkins, Tselil Schramm, Alexander Wein, and Ilias Zadik. The franz-parisi criterion and computational trade-off's in high dimensional statisties, In S. Koyejo, Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages Associates, Inc., 2022.
Curran
[37] Matthew S Brennan, Guy Bresler, Sam Hopkins, Jerry Li, and Tselil Schramm. Statistical query algorithms and low degree tests are almost equivalent. In Mikhail Belkin and Samory Kpotufe, editors. Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 774-774. PMLR, 15-19 Aug 2021. URL html
[38]
Rishabh Dudeja and Daniel Hsu. Statistical query lower bounds for tensor pea. Journal of Machine Learning Research, 22(83):1-51, 2021. URL htt.htt:://mmmrorrrg
20- 837. html
[39]
Rishabh Dudeja and Daniel Hsu. Statistical-computaiioon trade-offs in tensor pea and related problems via communication complexity, 2022.
[40]
Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor PCA. The Annals of Probability, 48(4):2052 - 2087, 2020. doi: 10.1214/19-A0P1115. URL
https: otg///0000r////./oo/o 1214119-A0011115
[41] Valentina Ros, Gerard Ben Arous, Giulio Biroli, and Chiara Cammarota. Complex energy
landscapes in spiked-tensor and simple glassy models: Ruggedness, arrangements of local minima, and phase transitions. Phys. Rev. X,
Jan 2019. doi: 10.1 10.1103/PhysRev ARLL:tts:: www.iiinn...........mmm..mmm.... orq /ww./o///100 1103/11000 wwt.//ywyyeev..eeeee 9. 011003
01 1003.
[42] Gérard Ben Arous, Song Mei, Andrea Montanari, and Mihai Nica. The landscape of the spiked tensor model. Communications on Pure and Applied Mathematics, 72(11):2282-2333 2019. doi: https://doi.org/10.0000000000/ hts/abbs/10.1002/pp...
URL
com/
[43] Balázs Szörényi. Characterizing statistical query learning: Simplified notions and proofs. In ALT, 2009.
[44]
losif Pinelis. Optimum Bounds for the Distributions of Martingales in Banach Spaces, The Annals of Probability, 22(()) 1706, 1994. doi: 10. 10.1214/app/1176696 URL https 7/doi.org 1tt.://211/////11699111
[45] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye MandeermnnMiinge and
Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:
[46] Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:/ www wandb.com com/ Software available from wandb.com.
A Background and Notation
A.1 Tensor Notation
Throughout this section let T' e (IRd) be k-tensor.
Definition 5 (Tensor Action). For j tensor A E (Re))s) with j S k, we define the action T[A] of on A by
(T(AA) T Ti Ath- At E (0)0))))-)))
We will also use (T, A) to denote T[A] = A/T] when A, are both & tensors. Note that this corresponds to the standard dot product after flattening A.T.
13

---

Definition 6 Cremmtaaiio/Transpsiiiii Given k-tensor T and a permutation π E Sk, we use (T) denote the result of permuting the axes T by the permutation т. i.e.
Definition 7 (Symmetrization). We define Sym, E (Rd) by
(Symg))a....dee k! TES
Σ
Sraayd ..
where Sy is the symmetric group on l,...,k k. Note that Symyg acts on k tensors T by
(Syma (T)
Σ (TT).
hd
WESk
Lee Sym; [T] is the symmetrized version T.
We will also overload notation and use Sym to denote the symmetrization operator, i.e. if T is k-tensor, (Tm( == [y)
Definition 8 (Symmetric Tensor Product), For k tensor T and a j tensor we define the symmetric tensor product of I and A by
T A = Sym(T & A).
Lemma 1. For any tensor T
[Sym(T)ll S F
Proof.
1 J S
[Sym(TT)l
Σ
T(T)
I 1|
k!
TES
RESk
because permuting the indices T does not change the Frobenius norm.
We will use the following two lemmas for tensor moments of the Gaussian distribution and the uniform distribution over the sphere:
Definition 9. For integers ki d > 0, define the quantity Vis as
u == (2k 1)!!
d 2j O(d
(d) Note that Vi (n 1>(=**]
Lemma (Tensorized Moments).
E 1a) 026 (2k: 1)!!7 and 1=-0 v((0) T
Proof. For the Gaussian moment, (6). The spherical moment follows from the decomposition 2r where r x()..
Lemma 3.
@ 13 = (1)
Proof. By Lemma 2 we have
= E B= [(E]] (a 1811 1 1|
14

---

A.2 Hermite Polynomials and Hermite Tensors
We provide brief review of the properties of Hermite polynomials and Hermite tensors. Definition 10. We define the kth Hermite polynomial H by
Hex(x) (2) == (-1)* V V*((x)
((x)
where p(x) := (2 is the PDF of a standard Gaussian in d dimensions. Note that when d = this definition reduces the standard univariate Hermite polynomials.
We begin with the classical properties of the scalar Hermite polynomials:
Lemma 4 (Properties of Hermite Polynomials). When d = 1,
Orthogonality:
E B...No..jHHee((eeeeee klōsa
Derivatives:
d (e)(x) kHe
dx
Correlations: If If N(0, 1) correlated Gaussians with correlation E [ryy]
F (e))(()))) kl6
Hermite Expansion: If f (2())) where μ is the PDF of standard Gaussian,
(c)
Ck Hex(x) where C ====NNo.)/((eeeeett
f(x)
k
k20
These properties also have tensor analogues:
Lemma (Hermite Polynomials in Higher Dimensions).
Relationship Univariante Hermite Polynomials: If If||w||
Her (w x) = (Hex(x), 0K)
Orthogonality:
E (Hez(x) @ Hex(x)] = 5nn Sym or equivalently, for any j tensor A and k tensor B:
Ex-N(o.L) (x), (Hex(z), BB] = 8jnk! (Sym(A), Sym(B))
Hermite Expansions: If R IR satisfies E xc,
(x), Ck) where C = Co= kd
K20
B Proof of Theorem
The proof of Theorem iis divided into four parts. First, Appendix [.1 introoduuucs some notation that will used throughout the proof. Next, Appendix [B22 matching upper and lower bounds for the gradient of the smoothed population loss. Similarly, Appendix [B3 concenttess the empirical gradient of the smoothed loss. Finally, Appendix [B.4 combines the bounds in Appendix B.2 and Appendix B.3 with standard online SGD analysis to arrive the final rate.
15

---

B.1 Additional Notation
Throughout the proof we will assume that w E 8 so that V denotes the spherical gradient of ws. In particular, (w)(w) w for any g S + R. We will also use a to denote w w that we can write expressions such as:
V a P w" w" a...
We will use the following assumption on without reference throughout the proof: Assumption 2. X2 - d/C for sufficiently large constant C.
note that this is satisfied for the optimal choice d²/4.
We will use Ö(-) to hide polylog(d) dependencies. Explicitly, ==001 (() (1 if there exists C1, C2 such that [X| & Ci log(d) We will also use the following shorthand for denoting high probability events:
Definition 11. We say event E happens with high probability for every k 2 0 there exists d(k) such that for all d d(k), P[E) 2 1 d
Note that high probability events are closed under polynomially sized union bounds. As an example, X (O, 1) then X S log(d) with high probability because
P[r log(d) S exp (- log(dd)//2) < d
for sufficiently large d. In general, Lemma 24 shows that if X is mean zero and has polynomial tails, i.c. there exists C' such that /|XP]// s p then X = Ō(1) with high probability.
B.2
Computing the Smoothed Population Gradient
Recall that
(xx)
(a)
Ck
Heg(x)
where := Be No.
R!
#20
In addition, because we assumed that Ez-No.1)(0(x)) 1 we have Parseval's identity:
= 1=
M
k20
This Hermite decomposition immmediately implies closed form for the population loss: Lemma 6 (Population Loss). Let w w² Then,
I(w)=
a+).
k20
Lemma 6 implies that understand the smoothed population [x(w) (L))) it suffices understand Lx (a+)) for & 0. First, will show that the set of single index models is closed under smoothing operator Σx:
Lemma 7. Let g [-1,1] - IR and let u E Sd-1
Then
Ex(g(w u)) = gx(w - u)
where
a aa VI+ X2
(x(a) = E2-Unit(S
and denotes the first coordinate of z.
16

---

Proof. Expanding the definition of Lx gives:
w + A
Cx(g(w u)) Ez~Ea
1| + Az 1
w u + X2 1
E Egnps
X2
Now I claim that when Max. 2- complete the proof. To this, note that we can decompose R into span(w) span(w} 1 Under this decomposition we have the polyspherical decomposition 2 - (0, where s' ~ Unif(S Then
(w- uu)² where Unif(~ Uni which would
d u = x' Pat i P = zIV1 (w u)².
Of central interest the quantities Lx (a as these terms show up when smoothing the population loss (see Lemma 6). We begin by defining the quantity s (a; A) which will provide matching upper and lower bounds on Lx(a²) when a 2 0:
Definition 12. We define sx A) by
a
a²
2 x
sk(0; A)
A2
a² < xa
and k is even
(1 X2)
2
a²
X² and ls is odd
Lemma 8. For all k 2 0 and 0, there exist constants c(k), C(k) such that clk)sx(a; S ( (C(K)ss(a;))
Proof. Using Lemma 7 we have that
a+>avvl- a²
Lx(a*) Ez-Uni((Su-))
+1
kk'
-11 4/2
k x² a2552 Un- 2)[=1]-
Now note that when j is odd, E Bew
[2
0 re-index this sum to get
= (1 (1 k/2
k 2 22
a²y Ez- 2
a
2j
1=
k
(1 =1 X k/2
x 21 (1
a
2j
1=
Note that every term in this sum is non-negative. Now we can ignore constants depending on & and use that y (d-1) x d to get
k L
fx(a
a
X2
a²)
VI+N
2d
==
Now when a² > 22 T this is decreasing geometric series which is dominated by the first term so
(x(at) x (Ø)* Next, when a² < 2 we have by Assumption that S 1/2 so - a² is bounded
17

---

away from O. Therefore the geometric series is dominated by the last term which is
1 (1 + R//
A2
7
k is even
k/2
A2
1 is odd
which completes the proof.
Next, in order to understand the population gradient, need to understand how the smoothing operator LA commutes with differentiation. We note that these do not directly commute because the smoothing distribution lAw depends so this term must he differentiated well. However, smoothing and differentiation almost commute, which is described the following lemma:
Lemma 9, Define the dimension- univariate smoothing operator by:
Y2y - 2
c/ah g(a)= = E EniUni(SS 2) g
VI+ X2
Then,
d
< 2i0 (o(a))
(d)
X2) - 1) 2x (g"(a)).
X²a (d++2)
da
V1+ (1+
Proof. Directly differentiating the definition for (/ gives
c/(0) C (g)(a)
+ a²
da
VI+
(1+ 12) aF)
vI+
By Lemma
25 this is equal to
d (d) (h
X²a
a + λ2ι Aeidivi-1 1 \
1E Uniif((()))
9"
da
V1+
(1 + X2) -
v1+ 12
c (g'())
Xaa
[ (g((t))
V
(1+
Now we are ready to analyze the population gradient: Lemma 10.
V Vwls(w)) = -(w (ow)
where for 2 1-1/4d-1/2
8k - A)
ca(a)
VI+X 1
Proof. Recall that
Lww)= -1
h!
Because 7et the index of the first nonzero Hermite coefficient, we can start this sum at l: k Smoothing and differentiating gives:
V Lyww) -ww" caw)cx(a) where cx(a) = Σ a d
kl da Ex(a²).
12k
18

---

We will break this into the h = k² term and the ls > k- tail. First when k = k- we can use Lemma 9
and Lemma 8 to get:
C. (kr)) do (((aa)))
C (k*)
ks Laaaaa--1)
k*(k** 1)A a)
VI+ X2
(1 (1+
fex- (aid)
The first term is equal up to constants to Xaa
while the second term is equal up to constants to
V1+A²
(1+x²) Sht -2a; A). However, we have that
Aaa [I+ Sk 2((;
A2
a²
1a
X2 1
a²
and k- is even 1
$
*k* (1=s))
d C'
VI+X*
a²
a²
and k odd
Therefore the k = k* term in (x(a) is equal up to constants to ( V1+1²
Next, we handle the k > k* tail. By Lemma thhi is equal to
(ax(a*11) k:k 1) (())) 2)
vI+ X2
(1 + x2)
Now recall that from Lemma tail in absolute value by
Cx(a*) is always non-negative so we can use of S k! to bound this
kLxla=-1)
k(k - 1) (()+)) (at
Σ k>k
V1+ X2
(1+ (2) 1)
X²a
Cx
Σ ka x
+
(1 A2)d c(d+2)
Σ k(k 1)a
V1 + X2
A
1
Xaa
cld+2) (
a+-
VI+
(1-
(1 (2
(1 - a)3
Now by Corollary this bounded for by
8k (ar; A)
2aa
+
-1 (a;
VI+ X2
(1+ + 22)d
For the first term, we have
a² >
2²
84- (a; X)
d
a² < 2
and k* is even C-11/4
Sk (a;
da
a² < d
and k* is odd
The second term is trivially bounded by
X²a
X2
(2d C(1+ S CVI+X
x2dd
which completes the proof.
B.3
Concentrating the Empirical Gradient
We cannot directly apply Lemma 7 to a a(w .z) as || # 1. Instead, we will use the properties of the Hermite tensors to directly smooth a (w- x:).
Lemma 11.
Lx Lx(Heu x))==)))== = ((lex())TT((w))
where
Tu(w) - (1+
k
((k-22)) (d-1)
2j
0=
19

---

Proof. Using Lemma 5
we can write
Lx x)) = (Hex(e), (w)k)))
Now
(www) = 8x
8k
w + du
= =Egupa
x2 k -
kh) w)((k-)) & 1 1200
(1+
f=U
Now by Lemma 2 this is equal to
Tx(w)= (1+ 12)
Σ9
k 2j
w w 25 )(Pt))) 1
j=0
which completes the proof.
Lemma 12. For any u E Sid-1 with u 1 w,
(n..
Bg-N(0.1a))
1 1+ (x
1-2
chat2) (at
$ k!l
+ 12+
V1+X2
Proof. Recall that by Lemma 11 we have
Lx (HHHex(w x)) = [HHe(())TT((())
where
h w 201 1 (dd-))
T= =
Differentiating this with respect to w gives
u. C ())))))))))))))))))))))) (Hez(v), VuTa(w)(u))
Now note that by Lemma 5 EzwN(o.La)
[tu- VwLx ()))) = EowN(0.La) (H)) Vwrc)))))) m))}²}
Therefore it suffices to compute the Frobenius norm of (..T( We first explicitly differentiate (w):
V. V. (1+x))
Σ
k
(k- 2j)u
2j
-==
[
21
$
(2))u 3 2 1) X 23 (d-1)
2j
1=
الجاا
k (1 (1+ X2)
Σ
k 1-
u
2
w
@K 23
3P
2j
0=ر
الجي
x²k(k- - 1)
k 2
Σ
@k 1-2j ((P)) و (225
u
2 w
1)
2j
1=
20

---

Taking Frobenius norms gives
|I [VTTu(W)
k-1
u u w 10 1-1)
\3+
2j
Σ
hs k-
u 2 w
1 21 (P. /80 5 22 (d+1)
(d - 1) 223A
2j
Now we can use Lemma 1 to pull out and get:
22
Σ
k
1)
w 1 (PP)) (( (d-1)
X3y
2j
0
k-2
X4111 d3 X2)
Σ
k (k-
w 1
(P
22 ا (d+
2j
Now note that the terms in each sum are orthogonal as at least one w will need to be contracted with Pal. Therefore this is equivalent to:
V [uw)
k 1)
X 11 ) 2
X1+
2j
j=0
x+
k 2 x ( 1 23 ) 15
2
dI
2j
F
Next, note that for any k tensor A, [|Sym(A)]]
1 Eres, (A, (A)).
When A =
2k )( the only permutations that don't give 0 are the ones which pair up all of the ws of which there are (2j)1(k: 2j)!. Therefore, by Lemma
Ok- ((P
35 F
1
(P
65
1
k
F
Vj
(a)
Plugging this in gives:
I [v.Te(/)
22
k- 1
\ \ (d))
22T
2j
ابيل
k - 2
Σ
2j
j=0
Now note that
1
kg
Lx(a*)
(1 + A23
2j
V14A3
which completes the proof.
21

---

Corollary 1. For any u E S with u 1 u,
min 11 X2, vd) -1
EewN(o,la) [ VwLx(o(w--a) (()(())))))))
X2
Proof. Note that
(x(o(w-x)) =
Ck (Hex(t),T(uw))
k!
80
VwLx(o(w-2)) x)) Σ
Ck (H))(1), V.TT((
k!
&
By Lemma these terms are orthogonal so by Lemma
12
(u (wLx(o(wx2)) =)))²|
E B
Lx (a=
+
2x
(d++)
(a (a²-2)
K!
X2
(1 + 22)2
V1+22
1+ X2 Lx
1
(1-a) a4'-1
+
X
C
c(d+2)
at'
(1+ 22)
(1 - aa)s
A
8k+
v1 V1+ x2
+
8k+
VI+12i0
1 X2
(1 + x2)2d
A
2A
1 X2 Sk -1 1x> +
+ x2)2d Sk 8k- (A)
Now plugging in the formula for sn (a; A) gives:
(1+ (1+ X2)
1 <
I+ Vd
(wu vwLx(o(w =)))² S (1+ (1+
1+x2
1+ Vd
-Ik* -1)
min 1 1 X2, vd
1+ X2
The following lemma shows that V ( inherits polynomial tails from a':
Lemma 13. There exists an absolute constant such that for any u E Sd-1 with w and any p E [0,d/C].
(o ((u. Vu x)))
v1+ + 12
Proof. Following the proof of Lemma we have
u V..Lxlow--ar
= =uu (u-x) Engi 2) w A2ι P. x
A xx) x1 X2
V1 + X2
V1+ + X²
First, consider the first term. Its p norm is bounded by
2p
1 E. [uu 2)30|LB Ex
o' w x + X2y |21||P P x
E Enj 2)
/1 + X2
v1+ X²
22

---

By Jensen we can pull out the expectation over 21 and use Assumption I to get
1
Ezl(u x) " 1o. o' w w..r + )2 P P
vI X2
v1 12
poly(p)
V1+ X2
Similarly, the of the second term is bounded by
2y #
A
21(22 w)
A
poly(p)
poly(p) Err
poly(p)
P
dv + x
VI+
Finally, we can use Corollary a Lemma 13 bound the p norms of the gradient:
L (w; x; Then there exists constant Lemma 14. Let (z, y) be a fresh sample and let u = L C' such that for any u E 5d-1 with u 1 w, any X - //44 and all 2 S S d/C,
Exy ku w/ S poly(p) Ō 13) 13)
Proof. First,
Ex w)///// Ex,p(uu (La (1)
= EEzoly'u CwCx(o(w 2//p
Applying Lemma 23 with x= (u v x)))² and Y yp VwLx x))) gives:
min 1 + 12 Vd
En (aul(u-e)0 w/p/// [e)](u-))j/// z poly(p)ō
V1 - X²
$ poly(p) ((1
which completes the proof.
Corollary 2. Let e be as Lemma 14 Then for all 25 SS sdd/
Ex Ellullll| 20 $ poly(p)
x1
Proof. By Jensen's inequality,
( d || EE
Σ (0-))))))
S 111 E
Σ (v. e,)00
S dl max (e))()))))))))))))))))))))) e.)
Taking pth roots and using Lemma
14 finishes the proof.
B.4 Analyzing the Dynamics
Throughout this section will assume 1 < X < d/44, The proof of the dynamics is split into three stages.
In the first stage, we analyze the regime a E [oo, Ad-1/22 In this regime, the signal is dominated by the smoothing.
23

---

In the second stage, we analyze the regime a E [Y 1/2 1 [1]] This analysis is similar to the analysis in Ben Arous et al. [ and could be equivalently carried out with À =
Finally in the third stage, we decay the learning rate linearly to achieve the optimal rate
n 2 d d
All three stages will use the following progress lemma:
Lemma 15. Let E Sd-1 and let a :- w w'. Let (a,, v) be fresh batch and define
+ Au
= -u 2:= Ex,lll, w' 1 + Xlll
and a' == w' w
Then "insavi+>'
a' = + +2 2 (1+
nida
where V and for all 2 pspdd/c
ER /011/0 s (poly(()) (1(1 (2) -
(+--1)
v1 a²
nda
v1+ + X2
Furthermore, if O(1) the Ö(-) can be replaced with O(.).
Proof. Because 1 and 1 >
2 1
VI+2"
a + n7v w²)
a'
a + ( (() +
v 1 + 1/1
where |r| s " [1/1 + n/u w*ll Note that by Lemma 14 n7u w²) has moments bounded by
(ply(p) S poly(p). Therefore by Lemma 23 with = ||ll and Y = a + m/l w').
Exy s Ō n B|lbb|(aa.
Plugging in the bound on |[[||| 2] from Corollary 2 gives
E [ay v/(1 a² ())((a)) + + (7 x2)-k* x*)
In addition, by Lemma 14
Eny [[v(v w²) - Ba.yli w "11 z poly(p) n [|Pl'*l| 0(1 + x2) 12)---- -
= poly(p) V1-aa a² 0(11 (1)
Similarly, by Lemma 23 with x = ||a||2 and Y = ||11| (a + 7/ . wIp Lemma 14 and
Corollary 2
Ez, [Ire Ez,,
Ex,,pr:s]P [r:]1p]// 1/0 $ En.,y [1ret]1//
)
1/p
S n² poly(p)
(1
p
poly(p) d2 δ (1 x3)- -
We can now analyze the first stage in which a E [d-1/2 X d-1/2]. This stage is dominated by the signal from the smoothing.
24

---

Lemma 16 (Stage 1). Assume that À 2 and Oo 2 d/11/2 Set
7 d (1+ and Ti C(1+ 12)+ log(d)
log(d)
for sufficiently large constant C. Then with high probability, there exists t Ti such that 0g 2 Ad 1/2
Proof. Let T be the hitting time for a,
1/2 1/2 For t T1, let E, be the event that
1-1
00 2
00
+ Σ 1=0
cx(as)
2
We will prove by induction that for any Ti, the event: {E, t 7 happens with high probability. The base case of f = 0 is trivial so let 2 0 and assume the result for all s < t. Note that
n/> d-1/2 S aj so by Lemma 15 and the fact that 1,
at = 00 +
Σ
n11 (()))(())))))))))))))))))))))() + +2 + + (1)d))2222)))
2M- )
j=U
Now note that PlE, or / 2 2 - IP El,, and / < 7] so let us condition on the event < т. Then by the induction hypothesis, with high probability we have as E 1 Ad-1/2 for all 8 < t. Plugging in the value of y) gives:
(11 - (11-3))a(aa)) + Ō
nd 2 x 2-2
(1 - az)oc(a))) n7((a))
Similarly, because probability,
Σ/ Zj a martingale we have by Lemma 22 and Lemma 24 that with high
1-1 Σ ==
z, 2 0
[T np- +nd
max nda; jet X
# +nd-1) 2 VT, . n
1/2 1/2
where we used that mda,/A S n << 1. Therefore conditioned on t S we have with high probability that for all S
2
+ Σ calaz)
1=0
Now we split into two cases depending on the parity 1. First, if ler is odd we have that with high probability, for all t < Ti:
04 Σ ap + πバλ - - or t t2T. T.
Now let t Ti. Then have that with high probability,
S Ti which implies that r < Ti with high probability, Next, if kt even we have that with high probability
04 2 1/2
a₁ ス ap
17
Σ as or t2t.
As above, by Lemma the first event implies that or,, 2 Ad2 so we must have 7 S Ti with high probability.
25

---

Next, we consider what happens when a 2 Ad-1/2 The analysis in this stage is similar to the online SGD analysis in [T]
Lemma 17 (Stage 2). Assume that oo 2 Ad-1/2 1/2 Set 1) T1 as in Lemma 16 Then with high probability, on, 2 1 d-1/4
Proof. The proof almost identical to Lemma
We again have from Lemma 15
L-1
at
2 ag +
Σ n1
-3)x(as)) + 2, 001 day 1 --26))
(
j=U
First, from martingale concentration we have that
|t ad
Σ 25
Ti nd- + +nd
x
0 TT
n +nd-t].1 :2)
nn
1d-1/2
C
where we used that rj < x
1²
Therefore with high probability,
Cu
n11 a3)ox(as) Ō 001² dajd
)
2
2
1=0
(1 - (1-0a))a(as)) ayd- 2'2
00 2
C22
1=
Therefore while a < 1
1.
for sufficiently large C we have
ge,
a0 2
j=0
Therefore by Lemma 27 we have that there exists t < T2/2 such that os 2 - ke. Next, let
1.
Pt Ou- Then applying Lemma
15 Pt and using E )** 2 1/e gives that if
d and c: x1/2X
r: -:
CX2
then
Pi+1 E Pt ncpp + npr + Z₁
= pc 1p + Zz-
Therefore,
Pets S (1 17)) pe + rr/c +
(1 n7c)* 7tt--1-i-
==0
With high probability, the martingale term bounded by Ad /2/C/ as before as long as S Tı, so for
Ad-1/22CC kー22
C log(d)
Ti] we have that Pc+. S C-1/2
A2
+X2 SC-1/2 1/4 Setting
ne
s Ti t and choosing aaprropriately yields pr, I 1/4 which completes the proof.
Finally, the third stage guarantees not only a hitting time but a last iterate guarantee. It also achieves the optimal sample complexity in terms of the target accuracy
26

---

Lemma 18 (Stage 3), Assume that ag 21 d d-1/4 Set X = 0 and
C
11 = C'd
for a sufficiently large constant C. Then for any S exp(d'/C), we have that with high probability,
2 1 00 (d+1) )
d
Proof. Let = =1 04. By Lemma 15 while Pt S 11/k:
Pc+1 S 1tPt
+ C1 11 Pp Wt + n/ d. Zt
20
= -(
C ((
Pt + C d pt Wr + 1/ d Zt.
C4d+t +t
where the moments of W,, Z, are each bounded by poly(p). We will prove by induction that with probability at least 1 - cxp (-Cd /e), we have for all s I t:
ad
S Ctd +
- C k*
The base case is clear so assume the result for all s S t. Then from the recurrence above,
d
Pu+1 S PU (C4d+t)2
E ( 1)
nvpi W. + nidd Zr).
(Cad+t)²
1=
First, because pa < 1/4,
Pod < C*pod
C²d+t & d
(CAd+1J' 1)2
Cad+t'
Next,
1
Σ I(Cd++))CChd (C+d+ + t)2
Σ Cdd
(C*d + t)2
5=0
3=0
da
(Cd++t²
Ctd+t
The next error term is:
Σ
(A
m Wt.
(C4d + t)²
1=0
Fix a//c c.
Then we will bound the pth moment of pr:
d0td C'd+t +
Elpll <
E
p.Z
203d Ctdee
203d 8
2Cad
c4d+1)
+2 P
pt 2 Ctd+t
203d
+ (PPtexp(Ccd100
C4d+t)
Now note that because / < (x)(
log (Lexp(- ((exp(-cot/0)) = log(() - Ca/lc= < - - - < - -plog()).
27

---

Therefore B/r < 40dd C4dtz. Therefore the p norm of the predictable quadratic variation of the next
eπror term is bounded by:
(ply(p) Σ (C*d+ Blyl] poly(p) Σ cd(cA++
S Co +).
In addition, the norm of the largest term in this sum is bounded by
poly(p) (Ccl()))""((C+++))
Therefore by Lemma 22 and Lemma 24 term is bounded by
we have with probability at least - exp(Cd=1/e/e) this
dt (C + t)3/2
Cc CCd+t'
d=/2 S
Finally, the last term is similarly bounded with probability at least 1 - eepppppppddcpp by
C-dvi (CAd+1y)
Cd
d < Cid+1
which completes the induction.
We can now combine the above lemmas to prove Theorem
Proof of Theorem 7 By Lemmas 16 to 18 if T
Ti + T2 have with high probability for all
Tr S cxp(d'/)):
aT 2 1 2000 0
d
d+ T2
Next, note that by Bernoulli's incquality ((1 x)" 2 1 + nx), have that 1 a* < k(1 a). Therefore,
2
L(WT) = Σ
④T)
h!
k20
$
(T)
k 1)!
(1 ar))
d+T2
which completes the proof of Thcorem
B.5
Proof of Theorem 2
We directly follow the proof of Theorem in Damian et al. 6 which is reproduced here for completeness. We begin with the following general CSQ lemma which can be found in Szörényi 43 Damian et al. [6]:
Lemma 19. Let F be class of functions and D be a data distribution such that
Exw Ee-Dif())] and [E==ol(()gg<))] Se Vf
Then any correlational statistical query learner requires at least [=((=-c) output function in F with (D) loss at most 2 2c.
queries of tolerance T to
2
28

---

First, we will construct a function class from subset of F := (a(w- W By 6 Lemma 3], for any there exist je unit vectors WI,.... W, such that their pairwise inner products all bounded by e. Let F == (o x) [s]]. Then for i $
Bo-aao 2w0 x)]1
Σ hd w w₂)x S lw, . wzl" S Se:
Therefore by Lemma 19
4q 2 c
a (02
Now set
log (4g(cd)k* 12)
cd
which gives
1 log k/2 log
(Aq(ed) 12)
log" 12g"'1((g))
72 x5
(ed) /2
d
C
Concentration Inequalities
Lemma 20
Pinells Inequality 44 Let (Y,)= be martingale with
martingale difference sequence {X:]== where X; Yi Y,-1. Let
(X)
Σ E
denote the predictable quadratic variation. Then there exists an absolute constant C such that for all
1 s c ((plltr) + max |
The above inequality is found in Pinelis 44 Theorem 4.1]. It is often combined with the following simple lemma:
Lemma 21. For any random variables XI...., Xm.
1/p
max ||X:|||
S
Σ
|xxilp
This has the immediate corollary:
Lemma 22. Let {Yi be martingale with martingale difference sequence {Xa) where
X; Y. Y,-1. Let (Y) = Dim xIl denote the predictable quadratic variation. Then there exists absolute constant C. such that for all p.
|Y s c (VellY 2/2 + pr//r max [|x:llp]
We will often use the following corollary of Holder's inequality to bound the operator norm of product of two random variables when one has polynomial tails:
Lemma 23. Let X,Y be random variables with [|Y|I, I sypp. Then,
|| C
EXXY] $ X By (26) max 1, (1.2) log
||
29

---

Proof. e [0, 1]. Then using Holder's 's inequality with 1 = c+ + gives:
X| E[x-~Xx-1] - ||||| |X 12ye
Using the fact that X, Y have polynomial tails we can bound this by
= - || | 1 (2/e)
First, if | l 2 e | we can set
C
which gives
log [TXT
B[XY] S |X||, oy (2c log
|| 2 |||||
Next, if || S x1l set which gives
E|XY] S |X | s oy
which completes the proof.
Finally, the following basic lemma will allow is to easily convert between p-norm bounds and concentration inequalities:
Lemma 24. Let 8 2 0 and let X be mean zero random variable satisfying
E B|x S Sxxp for log(1/5)
C
for some C. Then with probability at least - 6, | S (ep)
Proof. Let Cxpe) Then,
P| 2 = P 2 ee
B||X P1
PD
eP
e Cp
8.
D Additional Technical Lemmas
The following lemma extends Steins's 's lemma (f (o pherical distribution (p(d) where (u()) is the distribution of 21 when = ~ Unif(S
(z)) to the ultras-
Lemma 25 (Spherical Stein's 's Lemma). For any g E (1(()),
(=n-tmi(((t-))111((0))) E Uswunii(Sur)) (z)))
d
where 21 denotes the first coordinate of z.
Proof. Recall that the density of 2 is equal to
(1 - (1-x2) _²
C(d) = VT I(
where
C(d)
I(g)
Therefore,
Ex-Unif(S" 1][=1([2]]]= = C(d) [ 219(21)(1 - -22) = dzy
30

---

Now we can integrate by parts to get
C(d) [" ((21)(1 - 222) 1
dz1
d-1
C(
B2nUnit(sets))l((0 (21))
C(d)(d-1
d B
Lemma 26. For S d/4,
2 (1 27}
Es-Unif((S4-1)
S E Unit(S 25-1) (zt).
Proof. Recall that the PDF of (()) is
(1 - x d=3
where C(d) = Vi ((
C(d)
r(g)
Using this we have that:
zk
C(d) J_, (1 - x2) (1 x²2) d=3
x
Ex Uniffi
1)
(1 2
dr
C(d) [" x0 x²2 s- dx
1
C(d - 2ÿ)
E [21]
C(d)
r( ( 2)-1
E Unifii 25 1)12"]
$ (52-0
We have the following generalization of Lemma g
Corollary 3. For any k.j with d > 2j + 11 and a 2 C/1/4 there exist clj, k), C(i. k) such that
a+
Lx
(1 (1-a)² a|j
S ((1,k)ss(a;)))
Proof. Expanding the definition of Lx gives:
a+tazvi-a vI+x"
( a
L2
[1- ar5
E Unittis 2)
(1
(1- at-a
VI+X*
Now let X
Av VI+
(1 (1- -1-
and note that by Cauchy- X
Then,
vI+xz
a+A=vI-a² V1+x²
a
1
Cx
E Unif 2)
(1 aa)s
(1- VI+x²
X
(1 + Xz)
a-Aavi-a
V1+A"
x Bew-nit(SS
(1- X222)
31

---

Now we can use the binomial theorem to expand this. Ignoring constants only depending on j, k:
a+
Lx (1 a}3) (1a=))))==
a²
V1+A2
E B2-Uni((SM-2)
(1 vI+x?
(1 xx2) )
(1 + (21)2
k- 11i(1 -02)0/21E Unif 2)
(1- x²2))
b S
(1+xx1)'
(1-1-111 - a2) E Unif 2)
(1 - 2
==
By Lemmaa 26 the 21 term is bounded by d when is even and Xd + when i odd. Therefore this expression is bounded by
X 12
Σ
21 x2 2 (1- (1+1)
a²d
1=0
x sk A) + jsh (0;
Now note that
2
2
X²
1 5k (0;)1
da
11
and k- is even & C 1/4
8 (a: A)
and k odd
do
Therefore, $a(a; A) is the dominant term which completes the proof.
Lemma 27 (Adapted from Abbe et al. sequence satisfying
Let 1): 0g 2 be positive constants, and let un be
44 2 00 +
Then, if maxoosstt-1
mu 8-1
log2 <
we have the lower bound
x
Ut
2
(k-1)
2 2 - 1)t
1
Proof. Consider the auxiliary sequence w; = ag + 17 1 w By induction, u, 2 wy To lower
bound We. we have that
Wt Wt--1
W/ - W
w
w
w
<
dr
w/
x
w/
10-1 -w, 11-
-(k-1)
w/ (k - 1)
(1 + mw
k-1
-(k-1) - (k-
12,
(kc - 1)
(1 + 1k k
(k-11 (k-
w/-1
(k 1)
2
Wi -k--1)
wi, - (k-1))
32

---

Therefore
1 (k-1))
S 1
(k-11
gı(k 1).
Altogether, we get
w
(k-1)
S au
(k-11)
şπイ( 1)6,
or
-
(k-1)
w/ >
00
77(k - 1)6
2
as desired.
E Additional Experimental Details
To compute the smoothed loss (x æ; y) we used the closed form for Lx (H (w z)) (see Ap-
pendix [.3). Experiments were run on 8 NVIDIA A6000 GPUs. Our code is written in JAX 45) and we used Weights and Biases [66 for experiment tracking.
33


Reviews:
[{"review_id": "vkJxiRflXC", "reviewer": "Reviewer_68nG", "summary": "This paper aims to fill the gap between the sample size complexity of online SGD and CSQ lower bound for learning a single index model. Inspired by implicit regularization of minibatch SGD, the authors show that online SGD on a smoothed correlation loss only needs sample size $n=\\Omega(d^{k^*/2})$ to efficiently learn the feature direction. This smoothed loss helps us to avoid poor local minima of previous unsmoothed loss which reduces the number of samples for learning a single index model and also matches the CSQ lower bound. The authors also present a connection with a partial trace algorithm for tensor PCA.", "questions": "1. [Tan and Vershynin, 2019] also studied phase retrieval via online SGD and got a sharp bound similar to [1].\n2. How about the misspecification case when the link target function is different from the activation function? Most of the analysis in Section 5 still works well in this case.\n3. Footnote 1 on page 3: you should mention $T_w$ is the tangent space at $w$.\n4. Line 125: typo $He_0(x)=1$.\n5. Line 178: should it be $v_t\\perp w_t$? In the analysis below Line 178, how do we ensure that the gradient norm $||v_t||=O(1)$?\n6. Line 212, should it be $z\\perp w$?\n7. Equations below Lines 212 and 214, index $k$ should be changed into $k^*$.\n8. In Section 6, for the experiments, how do you choose the batch size and learning rate? And in Figure 2, if we use the square loss for minibatch SGD training, do we need a larger sample size? It would be better to compare these two cases in the simulation to visualize the difference.\n9. Line 433, typo\n10. Maybe you should provide references for Lemmas 4 and 5 in Section A.2. Similar to the equation below Line 272, you should distinguish Hermite polynomial and Hermite tensor in the multivariate case.\n11. In Lemma 7, you should explain the notion $z\\sim\\mathbb{S}^{d-2}$ and its relationship with $z_1$. Or use $\\text{Unif}(\\mathbb{S}^{d-2})$. Is $z_1$ the first entry of vector $z$? Similar issue for Lemma 25.\n12. Equation below Line 537: should be $\\nabla_w L_\\lambda (w)$\n13. Lemma 14, what is $\\mathbb{E}_{\\mathcal{B}}$?\n14. Below Line 692, what is $m$? It should be $q$ in Theorem 2.\n15. Equation below Line 726, $z$ in the integral should be $z_1$.\n\n=====================================================================================================\n- Tan, Y.S. and Vershynin, R., 2019. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval. arXiv preprint arXiv:1910.12837.", "limitations": "The limitations of the work are well addressed by the authors. I do not believe this work has any particular negative social impact.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, I found the paper well-written and easy to follow. The proof sketch in Section 5 provides a clear relation among SNR of the feature alignment, the smoothed loss, and sample complexity. The paper presents many interesting insights into bridging between smoothed loss landscape and sample complexity. I believe this paper will further help us to understand the training dynamics of minibatch SGD for learning a single-index model in future work.", "weaknesses": "The main concern is the CSQ lower bound. The gradient-based algorithm for correlation loss is a correlation statistical query learner. But can gradient descent of square loss be fully described by correlation statistical query? There should be some additional term in gradient descent of square loss which is not a correlation query. It would be better to have a clarification on this problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed and thoughtful review. We’ve corrected the typos you identified and we’ll try to clarify the higher level questions below. Please let us know if you have any further questions or if anything is still unclear.\n\n> The main concern is the CSQ lower bound. The gradient-based algorithm for correlation loss is a correlation statistical query learner. But can gradient descent of square loss be fully described by correlation statistical query? There should be some additional term in gradient descent of square loss which is not a correlation query. It would be better to have a clarification on this problem.\n\nYou are correct that there is an extra term, but this extra term does not interact with the labels $y$. As we are in the setting where the distribution over $x \\sim N(0,I_d)$ is known, this extra term can be directly estimated by a CSQ learner. We’ve included additional details in the CSQ section of the rebuttal.\n\n> [Tan and Vershynin, 2019] also studied phase retrieval via online SGD and got a sharp bound similar to [1].\n\nThank you, we will add this reference to the revision.\n\n> How about the misspecification case when the link target function is different from the activation function? Most of the analysis in Section 5 still works well in this case.\n\nOur result naturally extends to the misspecified setting but in a somewhat nontrivial way. Let the target link function be $\\phi$ with information exponent $k^\\star$ and the learner activation be $\\sigma$ with information exponent $s^\\star$. In the SNR calculation, the signal remains unchanged but the noise is equal to $d \\lambda^{-2 s^\\star}$. Going through the rest of the proof gives that the final sample complexity is $d^{k^\\star - \\max(1,s^\\star/2)}$. Therefore when the information exponents of $\\sigma$ and $\\phi$ are equal, the analysis and final result are unchanged. However in the case when $\\sigma$ has a lower information exponent, the sample complexity is strictly worse. We will include this discussion in the next revision of our paper.\n\n> In the analysis below Line 178, how do we ensure that the gradient norm $\\|v_t\\| = O(1)$?\n\nThe derivation in this part of the proof sketch is heuristic so we do not directly track all of the error terms. In fact, $\\|v_t\\| = O(d)$ so this should be $\\eta^3 d$. We believe explicitly tracking these error terms does not add to the proof sketch but to avoid being misleading we’ve replaced the $O(\\eta_t^3)$ with $\\ldots$ to represent the higher order terms which need to be carefully bounded.\n> In Section 6, for the experiments, how do you choose the batch size and learning rate?\n\nWe give explicit formulas that we used to pick $\\eta$ and the batch size (line 231-232) for all of the experiments.\n\n> And in Figure 2, if we use the square loss for minibatch SGD training, do we need a larger sample size? It would be better to compare these two cases in the simulation to visualize the difference.\n\nWe believe that the sample complexity would remain unchanged so long as the smoothing is applied to the function $f$ rather than to the loss.\n\n > Lemma 14, what is $\\mathbb{E}\\_\\mathcal{B}$\n\nOur apologies, $\\mathcal{B} = (x,y)$ denotes the current sample; however, you are correct that this was never defined. We’ve replaced this notation throughout the paper with $\\mathbb{E}\\_{x,y}$."}, {"review_id": "4erMDYr5Uh", "reviewer": "Reviewer_2zse", "summary": "This paper studies the sample complexity of learning a single-index function \\sigma(w*^Tx) via SGD on a smoothed correlation loss. The authors show that when k* is the first non-zero Hermite coefficient of \\sigma, with optimally tuned smoothing, defined as averaging the loss over a sphere of radius \\lambda centered at the iterate, the SGD will converge to error epsilon in d^{k*/2} + d/epsilon iterations. This improves over the iteration complexity d^{k* - 1}. This is a tight analysis, since it meets the CSQ lower bound.\n\nThe analysis is based on analyzing a certain signal-to-noise ratio which arises from comparing the alignment of the gradient with the ground truth direction, and the norm of the gradient. The authors show that when the smoothing increases, both of these terms decrease, but the norm decreases more.\n", "questions": "- The authors say in line 81 that the class of CSQ algorithms contains gradient descent. I am not certain this is true? Could the authors include a citation. Or perhaps it requires some qualifications? \n- In line 86, I do not see a definition of r and p.\n- In line 89 could the author explain (or at least define) what they mean by \"ERM\" or point to section 7.2\n- Section 7.2 is somewhat confusing, because it skips from discussion general objectives to the Tensor PCA. The paragraph starting on line 287 is hard to understand (why will GD converge to the expectation over z of the gradient?\n- Line  212 there is a v that should be a z?\n", "limitations": "The authors should discuss the limitations of using the correlation loss for learning single index functions. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper achieves a significant result which enhances our understanding of SGD and achieves a known lower-bound. \n- The paper is written very clearly in a way that highlights the main analysis techniques in the main body. It is useful that the authors first explain the vanilla SGD and then move on to the smoothed case.\n- The paper explains the connection to related work and particularly Tensor PCA very well.", "weaknesses": "- Discussion of the analysis of E[|v|^2] which seems a very key part of the proof is lacking in the main body. If the authors do not have space for many more details, could they include at least some intuition or a simple example (perhaps in low dimension or for a simple sigma?) for why E[|v|^2] has the stated dependence on lambda? And then if would be helpful if there were some pointers to the appendix for where the real proofs of the main steps can be found.\n- The paper only studies the correlation loss which is not frequently used in practice. Could the authors at least state why they use this, and if they believe their techniques would extend to the MSE loss?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed and thoughtful review. We’ve corrected the typos you identified and we’ll try to clarify the higher level questions below. Please let us know if you have any further questions or if anything is still unclear.\n\n> Discussion of the analysis of E[|v|^2] which seems a very key part of the proof is lacking in the main body.\n\nWe originally omitted the variance calculation due to space constraints but as this is a crucial part of the proof sketch, we’ve added a brief sketch for the $\\lambda$ dependence. See the “Computing the Variance” section of our rebuttal.\n\n> The paper only studies the correlation loss which is not frequently used in practice. Could the authors at least state why they use this, and if they believe their techniques would extend to the MSE loss?\n\nBecause our parameters are constrained to the sphere, there is no difference between the correlation and MSE losses in population. In particular,\n$$\n\\mathbb{E}\\_{x,y}\\left[\\frac{1}{2}(y - f_\\theta(x))^2\\right] = \\mathbb{E}\\_{x,y}\\left[1 - yf_\\theta(x\\right].\n$$\nThe difference is that the noise structure is changed due to the additional $f_\\theta(x)^2$ term. We believe that if you directly smooth the model, rather than smoothing the loss, our analysis should still go through. However, the correlation loss significantly simplifies the computations.\n\n> The authors say in line 81 that the class of CSQ algorithms contains gradient descent. I am not certain this is true? Could the authors include a citation. Or perhaps it requires some qualifications? \n\nThe key connection is that GD with square loss only interacts with the labels $y$ through correlational queries. For example for GD with model $f_\\theta$ and with square loss you have $$\\nabla L(w;x,y) = (f_\\theta(x) - \\underbrace{y)\\nabla f_\\theta(x)}_{\\text{query}}.$$\nThe other term in the gradient only depends on the distribution of $x \\sim N(0,I_d)$ which is known so it does not enter the sample complexity. See the CSQ section of the general rebuttal for additional discussion.\n\n> In line 89 could the author explain (or at least define) what they mean by \"ERM\" or point to section 7.2\n\nWe used ERM to refer to empirical risk minimization where the goal is to directly minimize the empirical loss $L_n := \\frac{1}{n} \\sum_i L(w;x_i,y_i)$ using an algorithm like gradient descent or minibatch SGD. In this setting, each sample is seen multiple times, in contrast to the online setting studied in the rest of the paper."}, {"review_id": "eMjWYg6qko", "reviewer": "Reviewer_mDyK", "summary": "This paper considers the problem of learning single index models in high dimensions, i.e., functions of a high-dimensional parameter that only depend on a one-dimensional projection of this parameter. This paper is interested in the case where the link function (the function of the one-dimensional projection) is known to the statistician, but the direction of the projection is unknown. \n\nThis paper studies the sample complexity of stochastic gradient descent (SGD). Its contribution is to show that smoothing the gradients allows to take larger stepsizes, and thus improves the sample complexity. The papers draws the connection with similar results for tensor decomposition. Finally, as previous results suggests that SGD has the effect of locally averaging the gradients, the authors present their work as a first step towards improving the sample complexity bounds for SGD. ", "questions": "*Remarks*\n\n(1) Running Algorithm 1 requires to be able to compute smoothed gradients. Is it easy to compute these? Are there closed-form formulas or do we need to use sampling over $z$? It would be nice to make this explicit as it would help to understand the significance of the results: is it a suggestion for better practical performances when learning single-index models, or is it an intermediary results towards proving improved convergence bounds for SGD?\n\n(2) Could you detail how stochastic gradient descent is a CSQ algorithm? We can show that if we use $n$ samples, then $|\\hat{q} - E_{x,y}[q(x,y)]| \\leq \\tau / n^{1/2}$ w.h.p., but not almost surely, right? Relatedly, you write that $\\tau \\approx n^{-1/2}$ is a \"heuristic\" (l.163), thus how rigorous is the CSQ lower bound? How rigorous are the claims of \"optimality\"?\n\n(3) From the proof sketch, I did not understand what is special about $\\lambda = d^{1/4}$. My intuition is that smoothing with a diverging $\\lambda$ should be very close to smoothing with $\\lambda = \\infty$. How would an algorithm with $\\lambda = \\infty$ (in the first steps) perform?\n\n\n*Minor remarks.*\n\n- When referring to [1], it is written many times that $n \\geq d^{k^*-1}$ samples are necessary (up to log factors), however my understanding is that this assumes $k^* \\geq 2$.\n- l.86: what is $r$?\n- Why do you consider the correlation loss rather than the square loss? I think the two would be equivalent; however it is quite unusual to optimize the correlation loss. \n- Thm 1: maybe it would be good to remind that $w_0 \\cdot w^* \\geq d^{-1/2}$ can be guaranteed with probability 1/2. \n- The statement of Thm 1 is a bit confusing. Certainly it is false that all stepsizes satisfying the O(.) bounds do not satisfy the statement (I can take the stepsizes to be 0), but maybe what you mean is that there exists a choice of stepsizes with these bounds that satisfy the statement? \n- I did not understand how you count the sample complexity in the parallel with the CSQ lower bound. My understanding is that to get one query with precision $n^{-1/2}$, you need $n$ samples. So the total sample complexity should depend on the number $q$ of queries? In light of this, I did not understand l.163-164. \n- Eq. below l.259: what does this tensor notation mean? \n\n\n*Typo.*\n- \"nad\" -> \"and\"", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper is well-written; in particular, the proof sketch does a good job at explaining the technical contributions of the paper. This paper proves rigorously the remarkable phenomenon that smoothing improves the sample complexity. This is a significant steps towards understanding the benefits of stochasticity in SGD.", "weaknesses": "There are a few spots where I would like some clarifications, see questions below. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed and thoughtful review. We’ve corrected the typos you identified and we’ll try to clarify the higher level questions below. Please let us know if you have any further questions or if anything is still unclear.\n\n> Running Algorithm 1 requires to be able to compute smoothed gradients. Is it easy to compute these?\n\nThere are two ways to compute the smoothed gradients. The most naive way is to approximate the expectation in $L_\\lambda$ using Monte Carlo. This requires roughly $d^{k^\\star/2}$ draws of z at every step; however, it does not affect the sample complexity.\n\nThe more efficient method is the one that we used for the experiments in section 6. It uses the observation from Lemma 7 that the smoothed gradient is a function of only two parameters: $w \\cdot x$ and $\\|x\\|$. Explicitly, there exists a $g$ such that $\\nabla L_t = g(w \\cdot x, \\|x\\|)$ where the closed form for $g$ follows from Lemma 7. This $g$ can be efficiently numerically computed for any activation function $\\sigma$ in $O(1)$ time. Furthermore, this computation only needs to be done once after which this $g$ can be reused. When $\\sigma$ is a polynomial, this $g$ has a closed form (see Appendix B.3) which we used for the experiments in section 6. We’ve extended Appendix E to include a discussion of the various methods for computing the smoothed gradients.\n\n> Could you detail how stochastic gradient descent is a CSQ algorithm? How rigorous are the claims of \"optimality\"?\n\nYou are correct and the CSQ lower bound does not constitute a rigorous lower bound for gradient descent. We’ve included additional discussion in the CSQ section of our rebuttal.\n\n> From the proof sketch, I did not understand what is special about $\\lambda = d^{1/4}$.\n\nThis threshold shows up in the computation for the noise term in the SNR after smoothing. We’ve added an additional section in the proof sketch that computes this variance and derives the $\\lambda = d^{1/4}$ threshold. For now, we’ve added this sketch to the “Computing the Variance” section of the general rebuttal. However, we will include it in the proof sketch section of the next revision of our paper.\n\n> I did not understand how you count the sample complexity in the parallel with the CSQ lower bound.\n\nThe key is that each query can be answered with the same $n$ samples. In particular, any SQ algorithm can be implemented with probability $1-\\delta$ with $n \\ge \\tau^{-2} \\log(\\text{queries}/\\delta)$ samples: for every query $q(x,y)$ return $\\frac{1}{n} \\sum_{i=1}^n q(x_i,y_i)$. For fixed $q$ this will be of order $\\sqrt{\\frac{\\log(1/\\delta)}{n}}$ so the result follows from a union bound. As the query complexity only appears through a log and is generally assumed to be polynomially large in $d$, it is often omitted. See the CSQ section in our general rebuttal for additional discussion.\n\n> Eq. below l.259: what does this tensor notation mean?\n\n$M_n$ is defined by taking the $k$ tensor $T$ and iteratively contracting indices until you are left with a vector or a matrix. This is sometimes written using $\\mathrm{Tr}$ notation, for example if $T$ is a 6-tensor then\n$$\nM_n = \\mathrm{Tr}_{(3,4),(5,6)}(T)\n$$\nmeaning that you contract the third/fourth indices and the fifth/sixth indices of $T$ to get a matrix. For our tensor notation, we use $T[A]$ to denote the contraction of $T$ with $A$ along the last $dim(A)$ dimensions of $T$ (Definition 5). In this case $A = I_d^{\\otimes 2}$ is the tensor product of the identity matrix with itself and $M_n = T[A]$. This is the higher dimensional analogue of $M[I] = \\langle M, I \\rangle = \\mathrm{tr}(M)$ when $M$ is a matrix."}, {"review_id": "AQXjRQn5AC", "reviewer": "Reviewer_xzE5", "summary": "Summary: The paper closes the gap between the previously established lower bound on learning a single index model with Gaussian inputs, giving a modified learning algorithm that matches the best possible sample complexity.\n", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Strengths: The alteration to SGD is simple but clever, and the given arguments are quite clear.  I particularly appreciate the proof sketch showing how the choice of smoothing impacts the drift term.  Tightening the gap left by [1] is a considerable contribution.\n\n[1] Arous, Gerard Ben, Reza Gheissari, and Aukosh Jagannath. \"Online stochastic gradient descent on non-convex losses from high-dimensional inference.\" The Journal of Machine Learning Research 22.1 (2021): 4788-4838.\n", "weaknesses": "Weaknesses: I may be misunderstanding, but it appears the experiments are using an analytic formulation of the smoothed loss, based on the hidden knowledge that the link functions (in the experiments) are hermite polynomials.  If I understand correctly, this corresponds to an “infinite synthetic sample” setting where the expectation in the smoothed loss is calculated exactly.  This wouldn’t impact the sample complexity in terms of querying the true single index model, but it raises a computational question.  Is there any reason to suspect that the analysis would change if, say, at every timestep the smoothed loss was approximated empirically with fresh samples drawn from the appropriate subset of the sphere?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the detailed and thoughtful review. Please let us know if you have any further questions or if anything is still unclear.\n\n> experiments are using an analytic formulation of the smoothed loss, based on the hidden knowledge that the link functions (in the experiments) are hermite polynomials. If I understand correctly, this corresponds to an “infinite synthetic sample” setting where the expectation in the smoothed loss is calculated exactly. This wouldn’t impact the sample complexity in terms of querying the true single index model, but it raises a computational question. Is there any reason to suspect that the analysis would change if, say, at every timestep the smoothed loss was approximated empirically with fresh samples drawn from the appropriate subset of the sphere?\n\nThere are two ways to compute the smoothed gradients. The most naive way is to approximate the expectation in $L_\\lambda$ using Monte Carlo. This requires roughly $d^{k^\\star/2}$ draws of z at every step; however, it does not affect the sample complexity.\n\nThe more efficient method is the one that we used for the experiments in section 6. It uses the observation from Lemma 7 that the smoothed gradient is a function of only two parameters: $w \\cdot x$ and $\\|x\\|$. Explicitly, there exists a $g$ such that $\\nabla L_t = g(w \\cdot x, \\|x\\|)$ where the closed form for $g$ follows from Lemma 7. This $g$ can be efficiently numerically computed for any activation function $\\sigma$ in $O(1)$ time. Furthermore, this computation only needs to be done once after which this $g$ can be reused. When $\\sigma$ is a polynomial, this $g$ has a closed form (see Appendix B.3) which we used for the experiments in section 6. We’ve extended Appendix E to include a discussion of the various methods for computing the smoothed gradients."}]

Author rebuttal:
We would like to thank the reviewers for their detailed and thoughtful reviews. We’ve addressed the most common questions in this rebuttal section.

## On the CSQ lower bound

The connection between the CSQ framework and gradient descent with square loss is that GD only interacts with the labels $y$ through correlational queries. For example if the model is $f_\theta$, the gradient is equal to $$\nabla L(w;x,y) = (f_\theta(x) - \underbrace{y)\nabla f_\theta(x)}_{\text{query}}.$$
The other term in the gradient only depends on the distribution of $x \sim N(0,I_d)$ which is known so it does not enter the sample complexity. However, we emphasize that this connection is only heuristic as the errors in GD are random while the errors in the SQ/CSQ framework are adversarial. However, such SQ/CSQ bounds have been commonly used to argue the existence of statistical-computational gaps in various learning problems [1,2,3,4].

[1] Feldman et al., 2012, Statistical Algorithms and a Lower Bound for Detecting Planted Cliques

[2] Diakonikolas & Kane, 2017, Statistical query lower bounds for robust estimation of high-dimensional gaussians and gaussian mixtures

[3] Goel et al., 2020, Statistical-Query Lower Bounds via Functional Gradients

[4] Dudeja & Hsu, 2020, Statistical Query Lower Bounds for Tensor PCA

## Computing the Variance and why we need $\lambda \le d^{1/4}$.

We originally omitted the variance calculation due to space constraints but as this is a crucial part of the proof sketch, we’ve added a brief sketch for the $\lambda$ dependence. This is also the part of the proof that introduces the fundamental constraint $\lambda \le d^{1/4}$ which is crucial for the final sample complexity. Here is the sketch:

Recall that $L_\lambda(w;x,y) = 1-y\sigma(w \cdot x)$. Differentiating through the smoothing operator gives:
$$
	\nabla_w L_\lambda(w;x,y)
	= -y \nabla_w~ \mathcal{L}\_\lambda(\sigma(w \cdot x)) \approx \lambda^{-1} x \mathcal{L}\_\lambda(\sigma'(w \cdot x)).
$$
We have that $y = O(1)$ and $\|x\| = O(\sqrt{d})$ so it suffices to bound $\mathcal{L}\_\lambda(\sigma'(w \cdot x))$. The variance of this term is equal to:
$$
	\mathbb{E}\_{x}[\mathcal{L}\_\lambda(\sigma'(w \cdot x))^2] = \mathbb{E}\_x \left[ \mathbb{E}\_{z \sim \mu\_w} \left[\sigma'\left(\frac{w + \lambda z}{\sqrt{1 + \lambda^2}} \cdot x\right)\right]^2\right].
$$
To compute this expectation, we will create an i.i.d. copy $z'$ of $z$ and rewrite this expectation as:
$$
	\mathbb{E}\_{x}[\mathcal{L}\_\lambda(\sigma'(w \cdot x))^2] = \mathbb{E}\_x\left[\mathbb{E}\_{z,z' \sim \mu\_w}\left[\sigma'\left(\frac{w + \lambda z}{\sqrt{1 + \lambda^2}} \cdot x\right)\sigma'\left(\frac{w + \lambda z'}{\sqrt{1 + \lambda^2}} \cdot x\right)\right]\right].
$$
Now we can swap the expectations and compute the expectation with respect to $x$ first using the Hermite expansion of $\sigma$. As the first nonzero Hermite coefficient of $\sigma'$ is $k^\star - 1$, this variance is approximately equal to the correlation between $\frac{w + \lambda z}{\sqrt{1 + \lambda^2}}$ and $\frac{w + \lambda z'}{\sqrt{1 + \lambda^2}}$ raised to the $k^\star-1$ power:
$$
	\mathbb{E}\_{x}[\mathcal{L}\_\lambda(\sigma'(w \cdot x))^2] \approx
	\mathbb{E}\_{z,z' \sim \mu_w}\left[\left(
		\frac{w + \lambda z}{\sqrt{1 + \lambda^2}} \cdot \frac{w + \lambda z'}{\sqrt{1 + \lambda^2}}
	\right)^{k^\star-1}\right]
	= \mathbb{E}\_{z,z' \sim \mu\_w}\left[\left(\frac{1 + \lambda^2 z \cdot z'}{1 + \lambda^2}\right)^{k^\star-1}\right]
$$
As $z,z'$ are random unit vectors, their inner product is order $d^{-1/2}$. Therefore when $\lambda \le d^{1/4}$, the first term in the numerator is dominant while when $\lambda \ge d^{1/4}$, the second term is dominant. Combining these regimes gives that the variance is of order $\min(\lambda,d^{1/4})^{k^\star-1}$ which motivates our optimal choice of $\lambda = d^{1/4}$. Combining this with the fact that $y = O(1)$ and $\|x\| = O(\sqrt{d})$ gives that the gradient of the smoothed loss has variance $d\lambda^{-2k^\star}$ for $\lambda \le d^{1/4}$.

Meta-review: This paper studies the task of learning a single index model, where the covariates follow an isotropic Gaussian distribution. The main goal of the paper is to fill the gap between the sample size complexity of online SGD and CSQ lower bound for learning this model. On the one hand, the existing literature shows that online SGD applied directly to the loss function has a sample complexity of $d^{k^*-1}$. On the other hand, the existing CSQ lower bound shows that $d^{k^*/2}$ samples are necessary to recover the true parameters of the problem. This paper closes this gap in the sample complexity, proving that online SGD applied to the smoothed loss function can recover the true solution with a sample complexity that matches the CSQ lower bound.

Five reviewers have reviewed the paper, and their overall assessment of the paper was extremely positive. I agree with this assessment and believe the paper answers a fundamental and impactful question.
Decision: Accept (oral)