Title: Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: statistical learning

learning theory

single index model

gradient descent

stochastic gradient descent
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

Abstract

We focus on the task of learning a single index model σ ( w ⋆ · x ) with respect to the isotropic Gaussian distribution in d dimensions. Prior work has shown that the sample complexity of learning w ⋆ is governed by the information exponent k ⋆ of the link function σ , which is defined as the index of the first nonzero Hermite coefficient of σ . Ben Arous et al. [1] showed that n ≳ d k ⋆ -1 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 ≳ d k ⋆ / 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 ⋆ with n ≳ d k ⋆ / 2 samples. We also draw connections to statistical analyses of tensor PCA and to the implicit regularization effects of minibatch SGD on empirical losses.

1 Introduction

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

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

However, gradient descent can be suboptimal for various statistical learning problems, as it only relies on local information in the loss landscape and is thus prone to getting stuck in local minima. For learning a single index model, the Correlational Statistical Query (CSQ) lower bound only requires

d max(1 ,k ⋆ / 2) samples to recover w ⋆ [6, 4], which is far fewer than the number of samples required by online SGD. This gap between gradient-based methods and the CSQ lower bound is also present in the Tensor PCA problem [8]; for recovering a rank 1 k -tensor in d dimensions, both gradient descent and the power method require d max(1 ,k -1) samples, whereas more sophisticated spectral algorithms can match the computational lower bound of d max(1 ,k/ 2) samples.

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

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

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

2 Related Work

There is a rich literature on learning single index models. Kakade et al. [15] showed that gradient descent can learn single index models when the link function is Lipschitz and monotonic and designed an alternative algorithm to handle the case when the link function is unknown. Soltanolkotabi [16] focused on learning single index models where the link function is 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 σ ( x ) = x 2 or σ ( x ) = | x | ; this corresponds to k ⋆ = 2 , and solving phase retrieval via gradient descent has been well studied [17-20]. Dudeja and Hsu [21] constructed an algorithm which explicitly uses the harmonic structure of Hermite polynomials to identify the information exponent. Ben Arous et al. [1] provided matching upper and lower bounds that show that n ≳ d max(1 ,k ⋆ -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 ≳ d samples, including single index models with polynomial link functions. Their estimator is based on the structure of the filtered PCA matrix E x,y [ 1 | y |≥ τ xx T ] , which relies on the heavy tails of polynomials. In particular, this upper bound does not apply to bounded link functions. Furthermore, while their result achieves the informationtheoretically optimal d dependence it is not a CSQ algorithm, whereas our Algorithm 1 achieves the optimal sample complexity over the class of CSQ algorithms (which contains gradient descent).

Recent work has also studied the ability of neural networks to learn single or multi-index models [5, 6, 23, 24, 4]. Bietti et al. [5] showed that two layer neural networks are able to adapt to unknown link functions with n ≳ d k ⋆ samples. Damian et al. [6] consider multi-index models with polynomial

link function, and under a nondegeneracy assumption which corresponds to the k ⋆ = 2 case, show that SGD on a two-layer neural network requires n ≳ d 2 + r p samples. Abbe et al. [24, 4] provide a generalization of the information exponent called the leap . They prove that in some settings, SGD can learn low dimensional target functions with n ≳ d Leap -1 samples. However, they conjecture that the optimal rate is n ≳ d Leap / 2 and that this can be achieved by ERM rather than online SGD.

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

3 Setting

3.1 Data distribution and target function

Our goal is to efficiently learn single index models of the form f ⋆ ( x ) = σ ( w ⋆ · x ) where w ⋆ ∈ S d -1 , the d -dimensional unit sphere. We assume that σ is normalized so that E x ∼ N (0 , 1) [ σ ( x ) 2 ] = 1 . We will also assume that σ is differentiable and that σ ′ has polynomial tails:

Assumption 1. There exist constants C 1 , C 2 such that | σ ′ ( x ) | ≤ C 1 (1 + x 2 ) C 2 .

Our goal is to recover w ⋆ given n samples ( x 1 , y 1 ) , . . . , ( x n , y n ) sampled i.i.d from

For simplicity of exposition, we assume that σ is known and we take our model class to be

3.2 Algorithm: online SGD on a smoothed landscape

As w ∈ S d -1 we will let ∇ w denote the spherical gradient with respect to w . That is, for a function g : R d → R , let ∇ w g ( w ) = ( I -ww T ) ∇ g ( z ) ∣ ∣ z = w where ∇ is the standard Euclidean gradient.

To compute the loss on a sample ( x, y ) , we use the correlation loss:

Furthermore, when the sample is omitted we refer to the population loss:

Our primary contribution is that SGD on a smoothed loss achieves the optimal sample complexity for this problem. First, we define the smoothing operator L λ :

Definition 1. Let g : S d -1 → R . We define the smoothing operator L λ by

where µ w is the uniform distribution over S d -1 conditioned on being perpendicular to w .

This choice of smoothing is natural for spherical gradient descent and can be directly related 1 to the Riemannian exponential map on S d -1 . We will often abuse notation and write L λ ( g ( w )) rather than ( L λ g )( w ) . The smoothed empirical loss L λ ( w ; x ; y ) and the population loss L λ ( w ) are defined by:

Our algorithm is online SGD on the smoothed loss L λ :

Algorithm 1: Smoothed Online SGD

3.3 Hermite polynomials and information exponent

The sample complexity of Algorithm 1 depends on the Hermite coefficients of σ :

Definition 2 (Hermite Polynomials) . The k th Hermite polynomial He k : R → R is the degree k , monic polynomial defined by

where µ ( x ) := e -x 2 2 √ 2 π is the PDF of a standard Gaussian.

The first few Hermite polynomials are He 0 ( x ) = 1 , He 1 ( x ) = x, He 2 ( x ) = x 2 -1 , He 3 ( x ) = x 3 -3 x . For further discussion on the Hermite polynomials and their properties, refer to Appendix A.2. The Hermite polynomials form an orthogonal basis of L 2 ( µ ) so any function in L 2 ( µ ) admits a Hermite expansion. We let { c k } k ≥ 0 denote the Hermite coefficients of the link function σ :

Definition 3 (Hermite Expansion of σ ) . Let { c k } k ≥ 0 be the Hermite coefficients of σ , i.e.

The critical quantity of interest is the information exponent of σ :

̸

Definition 4 (Information Exponent) . k ⋆ = k ⋆ ( σ ) is the first index k ≥ 1 such that c k = 0 .

Example 1. Below are some example link functions and their information exponents:

σ ( x ) = x and σ ( x ) = ReLU( x ) := max(0 , x ) have information exponents k ⋆ = 1 .

σ ( x ) = x 2 and σ ( x ) = | x | have information exponents k ⋆ = 2 .

σ ( x ) = x 3 -3 x has information exponent k ⋆ = 3 . More generally, σ ( x ) = He k ( x ) has information exponent k ⋆ = k .

Throughout our main results we focus on the case k ⋆ ≥ 3 as when k ⋆ = 1 , 2 , online SGD without smoothing already achieves the optimal sample complexity of n ≍ d samples (up to log factors) [1].

4 Main Results

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

Theorem 1. Assume w 0 · w ⋆ ≳ d -1 / 2 and λ ∈ [1 , d 1 / 4 ] . Then there exists a choice of T 1 , η satisfying T 1 = ˜ O ( d k ⋆ -1 λ -2 k ⋆ +4 ) and η = ˜ O ( d -k ⋆ / 2 λ 2 k ⋆ -2 ) such that if we run Algorithm 1 with λ t = λ and η t = η for t ≤ T 1 and λ t = 0 and η t = O ( ( d + t -T 1 ) -1 ) for t > T 1 , we have that with high probability, after T = T 1 + T 2 steps, the final iterate w T satisfies L ( w T ) ≤ O ( d d + T 2 ) .

Note that the condition that w 0 · w ⋆ ≳ d 1 / 2 can be guaranteed with probability 1 / 2 . Theorem 1 uses large smoothing (up to λ = d 1 / 4 ) to rapidly escape the regime in which w · w ⋆ ≍ d -1 / 2 . This first stage continues until w · w ⋆ = 1 -o d (1) which takes T 1 = ˜ O ( d k ⋆ / 2 ) steps when λ = d 1 / 4 . The second stage, in which λ = 0 and the learning rate decays linearly, lasts for an additional T 2 = d/ϵ

steps where ϵ is the target accuracy. Because Algorithm 1 uses each sample exactly once, this gives the sample complexity

to reach population loss L ( w T ) ≤ ϵ . Setting λ = O (1) is equivalent to zero smoothing and gives a sample complexity of n ≳ d k ⋆ -1 + d/ϵ , which matches the results of Ben Arous et al. [1]. On the other hand, setting λ to the maximal allowable value of d 1 / 4 gives:

which matches the sum of the CSQ lower bound, which is d k ⋆ 2 , and the information-theoretic lower bound, which is d/ϵ , up to poly-logarithmic factors.

To complement Theorem 1, we replicate the CSQ lower bound in [6] for the specific function class σ ( w · x ) where w ∈ S d -1 . Statistical query learners are a family of learners that can query values q ( x, y ) and receive outputs ˆ q with | ˆ q -E x,y [ q ( x, y )] | ≤ τ where τ denotes the query tolerance [28, 29]. An important class of statistical query learners is that of correlational/inner product statistical queries (CSQ) of the form q ( x, y ) = yh ( x ) . This includes a wide class of algorithms including gradient descent with square loss and correlation loss. For example, if the model is f θ ( x ) , then the gradient with square loss can be written as

Note that the other term in the gradient only depends on the distribution of x ∼ N (0 , I d ) which we assume is known. However, this connection with gradient descent is only heuristic as the errors in GD are random while the errors in the SQ/CSQ framework are adversarial. However, SQ/CSQ lower bounds are frequently used to argue the existence of statistical-computational gaps in statistical learning problems [30-33]. Theorem 2 measures the tolerance needed by CSQ learners to learn w ⋆ .

Theorem 2 (CSQ Lower Bound) . Consider the function class F σ := { σ ( w · x ) : w ∈ S d -1 } . Any CSQ algorithm using q queries requires a tolerance τ of at most

to output an f ∈ F σ with population loss less than 1 / 2 .

Using the standard τ ≈ n -1 / 2 heuristic which comes from concentration, this implies that n ≳ d k ⋆ 2 samples are necessary to learn σ ( w · x ) unless the algorithm makes exponentially many queries. In the context of gradient descent, this is equivalent to either requiring exponentially many parameters or exponentially many steps of gradient descent.

5 Proof Sketch

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

5.1 Online SGD Analysis

To begin, we will analyze a single step of online SGD. We define α t := w t · w ⋆ so that α t ∈ [ -1 , 1] measures our current progress. Furthermore, let v t := -∇ L λ t ( w t ; x t ; y t ) . Recall that the online SGD update is:

Using the fact that v t ⊥ w t and 1 √ 1+ x 2 ≈ 1 -x 2 2 we can Taylor expand the update for α t +1 :

As in Ben Arous et al. [1], we decompose this update into a drift term and a martingale term. Let F t = σ { ( x 0 , y 0 ) , . . . , ( x t -1 , y t -1 ) } be the natural filtration. We focus on the drift term as the martingale term can be handled with standard concentration arguments. Taking expectations with respect to the fresh batch ( x t , y t ) gives:

so to guarantee a positive drift, we need to set η t ≤ 2 E [ v t · w ⋆ |F t ] E [ ∥ v t ∥ 2 |F t ] α t which gives us the value of η t used in Theorem 1 for t ≤ T 1 . However, to simplify the proof sketch we can assume knowledge of E [ v t · w ⋆ |F t ] and E [ ∥ v t ∥ 2 |F t ] and optimize over η t to get a maximum drift of

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 a natural interpretation as the signalto-noise ratio (SNR). Note that the SNR is a local property, i.e. the SNR can vary for different w t . When the SNR can be written as a function of α t = w t · w ⋆ , the SNR directly dictates the rate of optimization through the ODE approximation: α ′ ≈ SNR /α . As online SGD uses each sample exactly once, the sample complexity for online SGD can be approximated by the time it takes this ODE to reach α ≈ 1 from α 0 ≈ d -1 / 2 . The remainder of the proof sketch will therefore focus on analyzing the SNR of the minibatch gradient ∇ L λ ( w ; x ; y ) .

5.2 Computing the Rate with Zero Smoothing

When λ = 0 , the signal and noise terms can easily be calculated. The key property we need is:

Property 1 (Orthogonality Property) . Let w,w ⋆ ∈ S d -1 and let α = w · w ⋆ . Then:

Using Property 1 and the Hermite expansion of σ (Definition 3) we can directly compute the population loss and gradient. Letting P ⊥ w := I -ww T denote the projection onto the subspace orthogonal to w we have:

As α ≪ 1 throughout most of the trajectory, the gradient is dominated by the first nonzero Hermite coefficient so up to constants: E [ v · w ⋆ ] = -∇ L ( w ) · w ⋆ ≈ α k ⋆ -1 . Similarly, a standard concentration argument shows that because v is a random vector in d dimensions where each coordinate is O (1) , E [ ∥ v ∥ ] 2 ≈ d . Therefore the SNR is equal to α 2( k ⋆ -1) /d so with an optimal learning rate schedule,

This can be approximated by the ODE α ′ = α 2 k ⋆ -3 /d . Solving this ODE with the initial α 0 ≍ d -1 / 2 gives that the escape time is proportional to dα -2( k ⋆ -1) 0 = d k ⋆ -1 which heuristically re-derives the result of Ben Arous et al. [1].

5.3 How Smoothing boosts the SNR

Smoothing improves the sample complexity of online SGD by boosting the SNR of the stochastic gradient ∇ L ( w ; x ; y ) .

d

-

1

/

2

d

-

1

/

4

1

d

-

k

⋆

d

-

k

2

+1

d

-

k

2

+

1

2

1

SNR

SNR with

λ

=

d

1

/

4

k

⋆

odd

k

⋆

even

No Smoothing

Computing the Signal Recall that the population loss was approximately equal to 1 -c 2 k ⋆ k ⋆ ! α k ⋆ where k ⋆ is the first nonzero Hermite coefficient of σ . Isolating the dominant α k ⋆ term and applying the smoothing operator L λ , we get:

Because z ⊥ w and ∥ z ∥ = 1 we have that ∥ w + λz ∥ = √ 1 + λ 2 ≈ λ . Therefore,

Now because z d = -z , the terms where j is odd disappear. Furthermore, for a random z , | z · w ⋆ | = Θ( d -1 / 2 ) . Therefore reindexing and ignoring all constants we have that

Differentiating gives that

As this is a geometric series, it is either dominated by the first or the last term depending on whether α ≥ λd -1 / 2 or α ≤ λd -1 / 2 . Furthermore, the last term is either d -k ⋆ -1 2 if k ⋆ is odd or α λ d -k ⋆ -2 2 if k ⋆ is even. Therefore the signal term is:

Computing the Noise Recall that L λ ( w ; x, y ) = 1 -yσ ( w · x ) . Differentiating through the smoothing operator gives:

Weknowthat with high probability, y = ˜ O (1) and ∥ x ∥ = O ( √ d ) so it suffices to bound L λ ( σ ′ ( w · x )) . The variance of this term is equal to:

To compute this expectation, we will create an i.i.d. copy z ′ of z and rewrite this expectation as:

Now we can swap the expectations and compute the expectation with respect to x first using the Hermite expansion of σ . As the first nonzero Hermite coefficient of σ ′ is k ⋆ -1 , this variance is approximately equal to the correlation between w + λz √ 1+ λ 2 and w + λz ′ √ 1+ λ 2 raised to the k ⋆ -1 power:

As z, z ′ are random unit vectors, their inner product is of order d -1 / 2 . Therefore when λ ≪ d 1 / 4 , the first term in the numerator is dominant and when λ ≫ d 1 / 4 , the second term is dominant. Combining these two regimes gives that the variance is of order min( λ, d 1 / 4 ) k ⋆ -1 which motivates our choice of λ = d 1 / 4 . Combining this with y = ˜ O (1) and ∥ x ∥ = O ( √ d ) gives that for λ ≤ d 1 / 4 , the variance of the noise is bounded by E [ ∥ v ∥ 2 ] ≲ dλ -2 k ⋆ . Note that in the high signal regime ( α ≥ λd -1 / 2 ), both the signal and the noise are smaller by factors of λ k ⋆ which cancel when computing the SNR. However, when α ≤ λd -1 / 2 the smoothing shrinks the noise faster than it shrinks the signal, resulting in an overall larger SNR. Explicitly,

For α ≥ λd -1 / 2 , smoothing does not affect the SNR. However, when α ≤ λd -1 / 2 , smoothing greatly increases the SNR (see Figure 1).

Solving the ODE: α ′ = SNR /α gives that it takes T ≈ d k ⋆ -1 λ -2 k ⋆ +4 steps to converge to α ≈ 1 from α ≈ d -1 / 2 . Once α ≈ 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 α ≥ 1 -ϵ with an additional d/ϵ steps, from which Theorem 1 follows.

6 Experiments

For k ⋆ = 3 , 4 , 5 and d = 2 8 , . . . , 2 13 we ran a minibatch variant of Algorithm 1 with batch size B when σ ( x ) = He k ⋆ ( x ) √ k ! , the normalized k ⋆ th Hermite polynomial. We set:

We computed the number of samples required for Algorithm 1 to reach α 2 = 0 . 5 from α = d -1 / 2 and we report the min, mean, and max over 10 random seeds. For each k we fit a power law of the form n = c 1 d c 2 in order to measure how the sample complexity scales with d . For all values of k ⋆ , we find that c 2 ≈ k ⋆ / 2 which matches Theorem 1. The results can be found in 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 is to recover the hidden direction w ∗ ∈ S d -1 from the noisy k -tensor T n ∈ ( R d ) ⊗ k given by 2

where Z ∈ ( R d ) ⊗ k is a Gaussian noise tensor with each entry drawn i.i.d from N (0 , 1) .

The Tensor PCA problem has garnered significant interest as it exhibits a statistical-computational gap . w ∗ is information theoretically recoverable when n ≳ d . However, the best polynomial-time algorithms require n ≳ d k/ 2 ; this lower bound has been shown to be tight for various notions of

Online SGD Sample Complexity, l = d 1 / 4

2 8

2 9

2 10

2 11

2 12

2 13

d (input dimension)

2 19

2 21

2 23

2 25

2 27

2 29

2 31

n (sample complexity)

k

/star

=

3, best fit:

n

∝

d

1

.

5

k

/star

=

4, best fit:

n

∝

d

2

.

1

k

/star

=

5, best fit:

n

∝

d

2

.

5

hardness such as CSQ or SoS lower bounds [34, 25, 35-39]. Tensor PCA also exhibits a gap between spectral methods and iterative algorithms. Algorithms that work in the n ≍ d k/ 2 regime rely on unfolding or contracting the tensor X , or on semidefinite programming relaxations [34, 25]. On the other hand, iterative algorithms including gradient descent, power method, and AMP require a much larger sample complexity of n ≳ d k -1 [40]. The suboptimality of iterative algorithms is 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 n ≪ d k -1 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 k [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 1 in the limit as λ →∞ [26]. Let T n = ( w ⋆ ) ⊗ k + 1 √ n Z . Then we will consider iteratively contracting indices of T until all that remains is a vector (if k is odd) or a matrix (if k is even). Explicitly, we define the partial trace tensor by

When k ⋆ is odd, we can directly return M n as our estimate for w ⋆ and when k ⋆ is even we return the top eigenvector of M n . Astandard concentration argument shows that this succeeds when n ≳ d ⌈ k/ 2 ⌉ . Furthermore, this can be strengthened to d k/ 2 by using the partial trace vector as a warm start for gradient descent or tensor power method when k 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 = d k -1 [40, 1]. On the other hand, the smoothing algorithms for Tensor PCA [27, 26] succeed with the computationally optimal sample complexity of n = d k/ 2 . Our Theorem 1 shows that this smoothing analysis can indeed be transferred to the single-index model setting.

In fact, one can make a direct connection between learning single-index models with Gaussian covariates and Tensor PCA. Consider learning a single-index model when σ ( x ) = He k ( x ) √ k ! , the normalized k th Hermite polynomial. Then minimizing the correlation loss is equivalent to maximizing

the loss function:

Here He k ( x i ) ∈ ( R d ) ⊗ k denotes the k th Hermite tensor (see Appendix A.2 for background on Hermite polynomials and Hermite tensors). In addition, by the orthogonality of the Hermite tensors, E x,y [ T n ] = ( w ⋆ ) ⊗ k so we can decompose T n = ( w ⋆ ) ⊗ k + Z n where by standard concentration, each entry of Z n is order n -1 / 2 . We can 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 are different.

7.2 Empirical Risk Minimization on the Smoothed Landscape

Our main sample complexity guarantee, Theorem 1, is based on a tight analysis of online SGD (Algorithm 1) in which each sample is used exactly once. One might expect that if the algorithm were allowed to reuse samples, as is standard practice in deep learning, that the algorithm could succeed with fewer samples. In particular, Abbe et al. [4] conjectured that gradient descent on the empirical loss L n ( w ) := 1 n ∑ n i =1 L ( w ; x i ; y i ) would succeed with n ≳ d k ⋆ / 2 samples.

Our smoothing algorithm Algorithm 1 can be directly translated to the ERM setting to learn w ⋆ with n ≳ d k ⋆ / 2 samples. We can then Taylor expand the smoothed loss in the large λ limit:

As λ →∞ , gradient descent on this smoothed loss will converge to E z ∼ Unif( S d -1 ) [ ∇ L n ( z )] which is equivalent to the partial trace estimator for k ⋆ odd (see Section 7.1). If k ⋆ is even, this first term is zero in expectation and gradient descent will converge to the top eigenvector of E z ∼ Unif( S d -1 ) [ ∇ 2 L n ( z )] , which corresponds to the partial trace estimator for k ⋆ even. Mirroring the calculation for the partial trace estimator, this succeeds with n ≳ d ⌈ k ⋆ / 2 ⌉ samples. When k ⋆ is odd, this can be further improved to d k ⋆ / 2 by using this estimator as a warm start from which to run gradient descent with λ = 0 as in Anandkumar et al. [26], Biroli et al. [27].

7.3 Connection to Minibatch SGD

A recent line of works has studied the implicit regularization effect of stochastic gradient descent [9, 11, 10]. The key idea is that over short timescales, the iterates converge to a quasi-stationary distribution N ( θ, λS ) where S ≈ I depends on the Hessian and the noise covariance at θ and λ is proportional to the ratio of the learning rate and batch size. As a result, over longer periods of time SGD follows the smoothed gradient of the empirical loss:

We therefore conjecture that minibatch SGD is also able to achieve the optimal n ≳ d k ⋆ / 2 sample complexity without explicit smoothing if the learning rate and batch size are properly tuned.

Acknowledgments and Disclosure of Funding

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

References

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(1):4788-4838, 2021.

Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. Advances in neural information processing systems , 29, 2016.

Tengyu Ma. Why do local methods solve nonconvex problems?, 2020.

Emmanuel Abbe, Enric Boix-Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv , 2023. URL https:// arxiv.org/abs/2302.11055 .

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.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistics , 46:2747-2774, 2018.

Emile Richard and Andrea Montanari. A statistical model for tensor pca. In Advances in Neural Information Processing Systems , pages 2897 - 2905, 2014.

Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on Learning Theory , pages 483-513, 2020.

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.

Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations , 2022.

Christopher J Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv:1811.03600 , 2018.

Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2818-2826, 2016.

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 arXiv:1902.08234 , 2019.

Sham M Kakade, Varun Kanade, Ohad Shamir, and Adam Kalai. Efficient learning of generalized linear and single index models with isotonic regression. Advances in Neural Information Processing Systems , 24, 2011.

Mahdi Soltanolkotabi. Learning relus via gradient descent. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Emmanuel J. 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.1109/TIT.2015.2399924.

Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. Gradient descent with random initialization: fast global convergence for nonconvex phase retrieval. Mathematical Programming , 176 (1):5-37, 2019.

Ju Sun, Qing Qu, and John Wright. A geometric analysis of phase retrieval. Foundations of Computational Mathematics , 18(5):1131001198, 2018.

Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval, 2019.

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 Conference On Learning Theory , volume 75 of Proceedings of Machine Learning Research , pages 1887-1930. PMLR, 06-09 Jul 2018. URL https://proceedings.mlr.press/v75/ dudeja18a.html .

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 https://proceedings.mlr.press/v125/chen20a.html .

Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

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.

Samuel B. Hopkins, Tselil Schramm, Jonathan Shi, and David Steurer. Fast spectral algorithms from sum-of-squares proofs: Tensor decomposition and planted sparse vectors. In Proceedings of the Forty-Eighth Annual ACM Symposium on Theory of Computing , STOC '16, page 178-191, New York, NY, USA, 2016. Association for Computing Machinery. ISBN 9781450341325. doi: 10.1145/2897518.2897529. URL https://doi.org/10.1145/2897518.2897529 .

Anima Anandkumar, Yuan Deng, Rong Ge, and Hossein Mobahi. Homotopy analysis for tensor pca. 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 Jul 2017. URL https://proceedings.mlr.press/v65/anandkumar17a.html .

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 pca. Journal of Physics A: Mathematical and Theoretical , 53(17):174003, apr 2020. doi: 10.1088/1751-8121/ab7b1f. URL https://dx.doi.org/10.1088/1751-8121/ ab7b1f .

Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans. Superpolynomial lower bounds for learning one-layer neural networks using gradient descent. arXiv preprint arXiv:2006.12011 , 2020.

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 1514-1539, 2020.

Vitaly Feldman, Elena Grigorescu, Lev Reyzin, Santosh Vempala, and Ying Xiao. Statistical algorithms and a lower bound for detecting planted clique, 2016.

Ilias Diakonikolas, Daniel M. Kane, and Alistair Stewart. Statistical query lower bounds for robust estimation of high-dimensional gaussians and gaussian mixtures, 2017.

Surbhi Goel, Aravind Gollakota, and Adam Klivans. Statistical-query lower bounds via functional gradients, 2020.

Rishabh Dudeja and Daniel Hsu. Statistical query lower bounds for tensor pca, 2021.

Samuel B. Hopkins, Jonathan Shi, and David Steurer. Tensor principal component analysis 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 Machine Learning Research , pages 956-1006, Paris, France, 03-06 Jul 2015. PMLR. URL https://proceedings.mlr.press/v40/Hopkins15.html .

Dmitriy Kunisky, Alexander S. Wein, and Afonso S. Bandeira. Notes on computational hardness of hypothesis testing: Predictions using the low-degree likelihood ratio, 2019.

Afonso S Bandeira, Ahmed El Alaoui, Samuel Hopkins, Tselil Schramm, Alexander S Wein, and Ilias Zadik. The franz-parisi criterion and computational trade-offs in high dimensional statistics. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 33831-33844. Curran Associates, Inc., 2022.

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 https://proceedings.mlr.press/v134/brennan21a.html .

Rishabh Dudeja and Daniel Hsu. Statistical query lower bounds for tensor pca. Journal of Machine Learning Research , 22(83):1-51, 2021. URL http://jmlr.org/papers/v22/ 20-837.html .

Rishabh Dudeja and Daniel Hsu. Statistical-computational trade-offs in tensor pca and related problems via communication complexity, 2022.

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-AOP1415. URL https://doi.org/10.1214/19-AOP1415 .

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 , 9:011003, Jan 2019. doi: 10.1103/PhysRevX.9. 011003. URL https://link.aps.org/doi/10.1103/PhysRevX.9.011003 .

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-2330, 2019. doi: https://doi.org/10.1002/cpa.21861. URL https://onlinelibrary.wiley.com/ doi/abs/10.1002/cpa.21861 .

Balázs Szörényi. Characterizing statistical query learning: Simplified notions and proofs. In ALT , 2009.

Iosif Pinelis. Optimum Bounds for the Distributions of Martingales in Banach Spaces. The Annals of Probability , 22(4):1679 - 1706, 1994. doi: 10.1214/aop/1176988477. URL https: //doi.org/10.1214/aop/1176988477 .

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https://www. wandb.com/ . Software available from wandb.com.

A Background and Notation

A.1 Tensor Notation

Throughout this section let T ∈ ( R d ) ⊗ k be a k -tensor.

Definition 5 (Tensor Action) . For a j tensor A ∈ ( R d ) ⊗ j with j ≤ k , we define the action T [ A ] of T on A by

We will also use ⟨ T, A ⟩ to denote T [ A ] = A [ T ] when A,T are both k tensors. Note that this corresponds to the standard dot product after flattening A,T .

Definition 6 (Permutation/Transposition) . Given a k -tensor T and a permutation π ∈ S k , we use π ( T ) to denote the result of permuting the axes of T by the permutation π , i.e.

Definition 7 (Symmetrization) . We define Sym k ∈ ( R d ) ⊗ 2 k by

where S k is the symmetric group on 1 , . . . , k . Note that Sym k acts on k tensors T by

i.e. Sym k [ T ] is the symmetrized version of T .

We will also overload notation and use Sym to denote the symmetrization operator, i.e. if T is a k -tensor, Sym( T ) := Sym k [ T ] .

Definition 8 (Symmetric Tensor Product) . For a k tensor T and a j tensor A we define the symmetric tensor product of T and A by

Lemma 1. For any tensor T ,

Proof.

because permuting the indices of 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 k, d > 0 , define the quantity ν ( d ) k as

Note that ν ( d ) k = E z ∼ Unif( S d -1 ) [ z 2 k 1 ] .

Lemma 2 (Tensorized Moments) .

Proof. For the Gaussian moment, see [6]. The spherical moment follows from the decomposition x d = zr where r ∼ χ ( d ) .

Lemma 3.

Proof. By Lemma 2 we have

A.2 Hermite Polynomials and Hermite Tensors

We provide a brief review of the properties of Hermite polynomials and Hermite tensors.

Definition 10. We define the k th Hermite polynomial He k ( x ) by

where µ ( x ) := e - ∥ x ∥ 2 2 (2 π ) d/ 2 is the PDF of a standard Gaussian in d dimensions. Note that when d = 1 , this definition reduces to 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:

Derivatives:

Orthogonality:

or equivalently, for any j tensor A and k tensor B :

Hermite Expansions: If f : R d → R satisfies E x ∼ N (0 ,I d ) [ f ( x ) 2 ] < ∞ ,

B Proof of Theorem 1

The proof of Theorem 1 is divided into four parts. First, Appendix B.1 introduces some notation that will be used throughout the proof. Next, Appendix B.2 computes matching upper and lower bounds for the gradient of the smoothed population loss. Similarly, Appendix B.3 concentrates the empirical gradient of the smoothed loss. Finally, Appendix B.4 combines the bounds in Appendix B.2 and Appendix B.3 with a standard online SGD analysis to arrive at the final rate.

Correlations: If x, y ∼ N (0 , 1) are correlated Gaussians with correlation α := E [ xy ] ,

Hermite Expansion: If f ∈ L 2 ( µ ) where µ is the PDF of a standard Gaussian,

These properties also have tensor analogues:

Lemma 5 (Hermite Polynomials in Higher Dimensions) .

Relationship to Univariante Hermite Polynomials: If ∥ w ∥ = 1 ,

B.1 Additional Notation

Throughout the proof we will assume that w ∈ S d -1 so that ∇ w denotes the spherical gradient of w . In particular, ∇ w g ( w ) ⊥ w for any g : S d -1 → R . We will also use α to denote w · w ⋆ so that we can write expressions such as:

We will use the following assumption on λ without reference throughout the proof:

Assumption 2. λ 2 ≤ d/C for a sufficiently large constant C .

We note that this is satisfied for the optimal choice of λ = d 1 / 4 .

We will use ˜ O ( · ) to hide polylog( d ) dependencies. Explicitly, X = ˜ O (1) if there exists C 1 , C 2 > 0 such that | X | ≤ C 1 log( d ) C 2 . We will also use the following shorthand for denoting high probability events:

Definition 11. We say an event E happens with high probability if for every k ≥ 0 there exists d ( k ) such that for all d ≥ d ( k ) , P [ E ] ≥ 1 -d -k .

Note that high probability events are closed under polynomially sized union bounds. As an example, if X ∼ N (0 , 1) then X ≤ log( d ) with high probability because

for sufficiently large d . In general, Lemma 24 shows that if X is mean zero and has polynomial tails, i.e. there exists C such that E [ X p ] 1 /p ≤ p C , then X = ˜ O (1) with high probability.

B.2 Computing the Smoothed Population Gradient

Recall that

In addition, because we assumed that E x ∼ N (0 , 1) [ σ ( x ) 2 ] = 1 we have Parseval's identity:

This Hermite decomposition immmediately implies a closed form for the population loss:

Lemma 6 (Population Loss) . Let α = w · w ⋆ . Then,

Lemma 6 implies that to understand the smoothed population L λ ( w ) = ( L λ L )( w ) , it suffices to understand L λ ( α k ) for k ≥ 0 . First, we will show that the set of single index models is closed under smoothing operator L λ :

Lemma 7. Let g : [ -1 , 1] → R and let u ∈ S d -1 . Then

and z 1 denotes the first coordinate of z .

where

Proof. Expanding the definition of L λ gives:

Now I claim that when z ∼ µ w , z · u d = z 1 √ 1 -( w · u ) 2 where z Unif( ∼ Unif( S d -2 )) which would complete the proof. To see this, note that we can decompose R d into span { w } ⊕ span { w } ⊥ . Under this decomposition we have the polyspherical decomposition z d = (0 , z ′ ) where z ′ ∼ Unif( S d -2 ) . Then

Of central interest are the quantities L λ ( α k ) as these terms show up when smoothing the population loss (see Lemma 6). We begin by defining the quantity s k ( α ; λ ) which will provide matching upper and lower bounds on L λ ( α k ) when α ≥ 0 :

Definition 12. We define s k ( α ; λ ) by

Lemma 8. For all k ≥ 0 and α ≥ 0 , there exist constants c ( k ) , C ( k ) such that

Proof. Using Lemma 7 we have that

Now note that when j is odd, E z ∼ Unif( S d -2 ) [ z j 1 ] = 0 so we can re-index this sum to get

Note that every term in this sum is non-negative. Now we can ignore constants depending on k and use that ν ( d -1) j ≍ d -j to get

Now when α 2 ≥ λ 2 d , this is a decreasing geometric series which is dominated by the first term so L λ ( α k ) ≍ ( α λ ) k . Next, when α 2 ≤ λ 2 d we have by Assumption 2 that α ≤ 1 C so 1 -α 2 is bounded

away from 0 . Therefore the geometric series is dominated by the last term which is

which completes the proof.

Next, in order to understand the population gradient, we need to understand how the smoothing operator L λ commutes with differentiation. We note that these do not directly commute because the smoothing distribution µ w depends on w so this term must be differentiated as well. However, smoothing and differentiation almost commute, which is described in the following lemma:

Lemma 9. Define the dimension-dependent univariate smoothing operator by:

Then,

Proof. Directly differentiating the definition for L ( d ) λ gives

By Lemma 25, this is equal to

Now we are ready to analyze the population gradient:

Lemma 10.

where for α ≥ C -1 / 4 d -1 / 2 ,

Proof. Recall that

Because k ⋆ is the index of the first nonzero Hermite coefficient, we can start this sum at k = k ⋆ . Smoothing and differentiating gives:

We will break this into the k = k ⋆ term and the k > k ⋆ tail. First when k = k ⋆ we can use Lemma 9 and Lemma 8 to get:

The first term is equal up to constants to s k ⋆ -1 ( α ; λ ) √ 1+ λ 2 while the second term is equal up to constants to λ 2 α (1+ λ 2 ) d s k ⋆ -2 ( α ; λ ) . However, we have that

Therefore the k = k ⋆ term in c λ ( α ) is equal up to constants to s k ⋆ -1 ( α ; λ ) √ 1+ λ 2 .

Next, we handle the k > k ⋆ tail. By Lemma 9 this is equal to

Now recall that from Lemma 8, L λ ( α k ) is always non-negative so we can use c 2 k ≤ k ! to bound this tail in absolute value by

Now by Corollary 3, this is bounded for d ≥ 5 by

For the first term, we have

The second term is trivially bounded by

which completes the proof.

B.3 Concentrating the Empirical Gradient

We cannot directly apply Lemma 7 to σ ( w · x ) as ∥ x ∥ ̸ = 1 . Instead, we will use the properties of the Hermite tensors to directly smooth σ ( w · x ) .

Lemma 11.

where

Proof. Using Lemma 5, we can write

Now

Now by Lemma 2, this is equal to

which completes the proof.

Lemma 12. For any u ∈ S d -1 with u ⊥ w ,

Proof. Recall that by Lemma 11 we have

where

Differentiating this with respect to w gives

Now note that by Lemma 5:

Therefore it suffices to compute the Frobenius norm of ∇ w T k ( w )[ u ] . We first explicitly differentiate T k ( w ) :

Taking Frobenius norms gives

Now we can use Lemma 1 to pull out u and get:

Now note that the terms in each sum are orthogonal as at least one w will need to be contracted with a P ⊥ w . Therefore this is equivalent to:

Next, note that for any k tensor A , ∥ Sym( A ) ∥ 2 F = 1 k ! ∑ π ∈ S k ⟨ A,π ( A ) ⟩ . When A = w ⊗ k -2 j ˜ ⊗ ( P ⊥ w ) ˜ ⊗ j , the only permutations that don't give 0 are the ones which pair up all of the w s of which there are (2 j )!( k -2 j )! . Therefore, by Lemma 3,

Plugging this in gives:

Now note that

which completes the proof.

Corollary 1. For any u ∈ S d -1 with u ⊥ w ,

Proof. Note that so

By Lemma 4, these terms are orthogonal so by Lemma 12,

Now plugging in the formula for s k ( α ; λ ) gives:

The following lemma shows that ∇ w L λ ( σ ( w · x )) inherits polynomial tails from σ ′ :

Lemma 13. There exists an absolute constant C such that for any u ∈ S d -1 with u ⊥ w and any p ∈ [0 , d/C ] ,

Proof. Following the proof of Lemma 9 we have

First, we consider the first term. Its p norm is bounded by

By Jensen we can pull out the expectation over z 1 and use Assumption 1 to get

Similarly, the p norm of the second term is bounded by

Finally, we can use Corollary 1 and Lemma 13 to bound the p norms of the gradient:

Lemma 14. Let ( x, y ) be a fresh sample and let v = -∇ L λ ( w ; x ; y ) . Then there exists a constant C such that for any u ∈ S d -1 with u ⊥ w , any λ ≤ d 1 / 4 and all 2 ≤ p ≤ d/C ,

Proof. First,

Applying Lemma 23 with X = ( u · ∇ w L λ ( σ ( w · x ))) 2 and Y = y p ( u · ∇ w L λ ( σ ( w · x ))) p -2 gives:

which completes the proof.

Corollary 2. Let v, ϵ be as in Lemma 14. Then for all 2 ≤ p ≤ d/C ,

Proof. By Jensen's inequality,

Taking p th roots and using Lemma 14 finishes the proof.

B.4 Analyzing the Dynamics

Throughout this section we will assume 1 ≤ λ ≤ d 1 / 4 . The proof of the dynamics is split into three stages.

In the first stage, we analyze the regime α ∈ [ α 0 , λd -1 / 2 ] . In this regime, the signal is dominated by the smoothing.

In the second stage, we analyze the regime α ∈ [ λd -1 / 2 , 1 -o d (1)] . This analysis is similar to the analysis in Ben Arous et al. [1] and could be equivalently carried out with λ = 0 .

Finally in the third stage, we decay the learning rate linearly to achieve the optimal rate

All three stages will use the following progress lemma:

Lemma 15. Let w ∈ S d -1 and let α := w · w ⋆ . Let ( x, y ) be a fresh batch and define

Then if η ≲ α √ 1 + λ 2 ,

where E x,y [ Z ] = 0 and for all 2 ≤ p ≤ d/C ,

Furthermore, if λ = O (1) the ˜ O ( · ) can be replaced with O ( · ) .

Proof. Because v ⊥ w and 1 ≥ 1 √ 1+ x 2 ≥ 1 -x 2 2 ,

where | r | ≤ η 2 2 ∥ v ∥ 2 [ α + η | v · w ⋆ | ] . Note that by Lemma 14, η ( v · w ⋆ ) has moments bounded by η λ poly( p ) ≲ α poly( p ) . Therefore by Lemma 23 with X = ∥ v ∥ 2 and Y = α + η | v · w ⋆ | ,

Plugging in the bound on E [ ∥ v ∥ 2 ] from Corollary 2 gives

Similarly, by Lemma 23 with X = ∥ v ∥ 2 and Y = ∥ v ∥ 2( p -1) [ α + η | v · w | ] p , Lemma 14, and Corollary 2,

We can now analyze the first stage in which α ∈ [ d -1 / 2 , λ · d -1 / 2 ] . This stage is dominated by the signal from the smoothing.

Lemma 16 (Stage 1) . Assume that λ ≥ 1 and α 0 ≥ 1 C d -1 / 2 . Set

for a sufficiently large constant C . Then with high probability, there exists t ≤ T 1 such that α t ≥ λd -1 / 2 .

Proof. Let τ be the hitting time for α τ ≥ λd -1 / 2 . For t ≤ T 1 , let E t be the event that

We will prove by induction that for any t ≤ T 1 , the event: { E t or t ≥ τ } happens with high probability. The base case of t = 0 is trivial so let t ≥ 0 and assume the result for all s < t . Note that η/λ ≪ d -1 / 2 C ≤ α j so by Lemma 15 and the fact that λ ≥ 1 ,

Now note that P [ E t or t ≥ τ ] = 1 -P [! E t and t < τ ] so let us condition on the event t < τ . Then by the induction hypothesis, with high probability we have α s ∈ [ α 0 2 , λd -1 / 2 ] for all s < t . Plugging in the value of η gives:

Similarly, because ∑ t -1 j =0 Z j is a martingale we have by Lemma 22 and Lemma 24 that with high probability,

where we used that ηdα j /λ ≤ η d ≪ 1 . Therefore conditioned on t ≤ τ we have with high probability that for all s ≤ t :

Now we split into two cases depending on the parity of k ⋆ . First, if k ⋆ is odd we have that with high probability, for all t ≤ T 1 :

Now let t = T 1 . Then we have that with high probability,

which implies that τ ≤ T 1 with high probability. Next, if k ⋆ is even we have that with high probability

As above, by Lemma 27 the first event implies that α T 1 ≥ λd -1 / 2 so we must have τ ≤ T 1 with high probability.

Next, we consider what happens when α ≥ λd -1 / 2 . The analysis in this stage is similar to the online SGD analysis in [1].

Lemma 17 (Stage 2) . Assume that α 0 ≥ λd -1 / 2 . Set η, T 1 as in Lemma 16. Then with high probability, α T 1 ≥ 1 -d -1 / 4 .

Proof. The proof is almost identical to Lemma 16. We again have from Lemma 15

First, from martingale concentration we have that

where we used that η ≪ λ 2 d . Therefore with high probability,

Therefore while α t ≤ 1 -1 k ⋆ , for sufficiently large C we have

Therefore by Lemma 27, we have that there exists t ≤ T 1 / 2 such that α t ≥ 1 -1 k ⋆ . Next, let p t = 1 -α t . Then applying Lemma 15 to p t and using (1 -1 k ⋆ ) k ⋆ ≳ 1 /e gives that if

With high probability, the martingale term is bounded by λd -1 / 2 /C as before as long as s ≤ T 1 , so for s ∈ [ C log( d ) ηc , T 1 ] we have that p t + s ≲ C -1 / 2 ( ( λ 2 d ) k ⋆ -2 2 + λd -1 / 2 ) ≲ C -1 / 2 d -1 / 4 . Setting s = T 1 -t and choosing C appropriately yields p T 1 ≤ d -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 ϵ :

then

Therefore,

Lemma 18 (Stage 3) . Assume that α 0 ≥ 1 -d -1 / 4 . Set λ = 0 and

for a sufficiently large constant C . Then for any t ≤ exp ( d 1 /C ) , we have that with high probability,

Proof. Let p t = 1 -α t . By Lemma 15, while p t ≤ 1 /k ⋆ :

where the moments of W t , Z t are each bounded by poly( p ) . We will prove by induction that with probability at least 1 -t exp ( -Cd 1 /C /e ) , we have for all s ≤ t :

The base case is clear so assume the result for all s ≤ t . Then from the recurrence above,

First, because p 0 ≤ d -1 / 4 ,

Next,

The next error term is:

Fix p = d 1 /C e . Then we will bound the p th moment of p t :

Now note that because t ≤ exp ( d 1 /C ) ,

Therefore E [ p p t ] 1 /p ≤ 4 C 3 d C 4 d + t . Therefore the p norm of the predictable quadratic variation of the next error term is bounded by:

In addition, the p norm of the largest term in this sum is bounded by

Therefore by Lemma 22 and Lemma 24, we have with probability at least 1 -exp ( Cd -1 /C /e ) , this term is bounded by

Finally, the last term is similarly bounded with probability at least 1 -exp ( -Cd -1 /C /e ) by

which completes the induction.

We can now combine the above lemmas to prove Theorem 1:

Proof of Theorem 1. By Lemmas 16 to 18, if T = T 1 + T 2 we have with high probability for all T 2 ≤ exp ( d 1 /C ) :

Next, note that by Bernoulli's inequality ( (1 + x ) n ≥ 1 + nx ), we have that 1 -α k ≤ k (1 -α ) . Therefore,

which completes the proof of Theorem 1.

B.5 Proof of Theorem 2

We directly follow the proof of Theorem 2 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 a class of functions and D be a data distribution such that

̸

Then any correlational statistical query learner requires at least |F| ( τ 2 -ϵ ) 2 queries of tolerance τ to output a function in F with L 2 ( D ) loss at most 2 -2 ϵ .

C Concentration Inequalities

Lemma 20 (Rosenthal-Burkholder-Pinelis Inequality [44]) . Let { Y i } n i =0 be a martingale with martingale difference sequence { X i } n i =1 where X i = Y i -Y i -1 . Let

denote the predictable quadratic variation. Then there exists an absolute constant C such that for all p ,

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 X 1 , . . . , X n ,

This has the immediate corollary:

Lemma 22. Let { Y i } n i =0 be a martingale with martingale difference sequence { X i } n i =1 where X i = Y i -Y i -1 . Let ⟨ Y ⟩ = ∑ n i =1 E [ ∥ X i ∥ 2 |F i -1 ] denote the predictable quadratic variation. Then there exists an absolute constant C such that for all p ,

We will often use the following corollary of Holder's inequality to bound the operator norm of a product of two random variables when one has polynomial tails:

Lemma 23. Let X,Y be random variables with ∥ Y ∥ p ≤ σ Y p C . Then,

̸

First, we will construct a function class from a subset of F := { σ ( w · x ) : w ∈ S d -1 } . By [6, Lemma 3], for any ϵ there exist 1 2 e cϵ 2 d unit vectors w 1 , . . . , w s such that their pairwise inner products are all bounded by ϵ . Let ̂ F := { σ ( w i · x ) : i ∈ [ s ] } . Then for i = j ,

Now set which gives

Proof. Fix ϵ ∈ [0 , 1] . Then using Holder's inequality with 1 = 1 -ϵ + ϵ 2 + ϵ 2 gives:

Using the fact that X,Y have polynomial tails we can bound this by

First, if ∥ X ∥ 2 ≥ e C ∥ X ∥ 1 , we can set ϵ = C log ( ∥ X ∥ 2 ∥ X ∥ 1 ) which gives

Next, if ∥ X ∥ 2 ≤ e C ∥ X ∥ 1 we can set ϵ = 1 which gives

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 δ ≥ 0 and let X be a mean zero random variable satisfying

for some C . Then with probability at least 1 -δ , | X | ≤ σ X ( ep ) C .

Proof. Let ϵ = σ X ( ep ) C . Then,

Therefore,

D Additional Technical Lemmas

The following lemma extends Steins's lemma ( E x ∼ N (0 , 1) [ xg ( x )] = E x ∼ N (0 , 1) [ g ′ ( x )] ) to the ultraspherical distribution µ ( d ) where µ ( d ) is the distribution of z 1 when z ∼ Unif( S d -1 ) :

Lemma 25 (Spherical Stein's Lemma) . For any g ∈ L 2 ( µ ( d ) ) ,

where z 1 denotes the first coordinate of z .

Proof. Recall that the density of z 1 is equal to

Now we can integrate by parts to get

Lemma 26. For j ≤ d/ 4 ,

Proof. Recall that the PDF of µ ( d ) is

Using this we have that:

We have the following generalization of Lemma 8:

Corollary 3. For any k, j ≥ 0 with d ≥ 2 j +1 and α ≥ C -1 / 4 d 1 / 2 , there exist c ( j, k ) , C ( j, k ) such that

Proof. Expanding the definition of L λ gives:

Now let X = λ √ 1 -α 2 √ 1+ λ 2 · ( 1 -α √ 1+ λ 2 ) -1 and note that by Cauchy-Schwarz, X ≤ 1 . Then,

Now we can use the binomial theorem to expand this. Ignoring constants only depending on j, k :

By Lemma 26, the z 1 term is bounded by d -i 2 when i is even and Xd -i +1 2 when i is odd. Therefore this expression is bounded by

Now note that

Therefore, s k ( α ; λ ) is the dominant term which completes the proof.

Lemma 27 (Adapted from Abbe et al. [4]) . Let η, a 0 ≥ 0 be positive constants, and let u t be a sequence satisfying

Then, if max 0 ≤ s ≤ t -1 ηu k -1 s ≤ log 2 k , we have the lower bound

Proof. Consider the auxiliary sequence w t = a 0 + η ∑ t -1 s =0 w k s . By induction, u t ≥ w t . To lower bound w t , we have that

Therefore

Altogether, we get or

as desired.

E Additional Experimental Details

To compute the smoothed loss L λ ( w ; x ; y ) we used the closed form for L λ ( He k ( w · x )) (see Appendix B.3). Experiments were run on 8 NVIDIA A6000 GPUs. Our code is written in JAX [45] and we used Weights and Biases [46] for experiment tracking.

Peer reviews:
comments: 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.

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.

[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.

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?

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThere are a few spots where I would like some clarifications, see questions below.\n\nQUESTIONS:\n*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\"\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\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- 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\nLIMITATIONS:\nThe authors should discuss the limitations of using the correlation loss for learning single index functions.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOverall, 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\n1. [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.\n\nLIMITATIONS:\nThe limitations of the work are well addressed by the authors. I do not believe this work has any particular negative social impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}