Title: On Single-Index Models beyond Gaussian Data
Venue: NeurIPS 2023
Accepted: yes
Keywords: gradient descent

shallow neural networks
Abstract: Sparse high-dimensional functions have arisen as a rich framework to study the behavior of gradient-descent methods using shallow neural networks, and showcasing its ability to perform feature learning beyond linear models. 
Amongst those functions, the simplest are single-index models $f(x) = \phi( x \cdot \theta^*)$, where the labels are generated by an arbitrary non-linear link function $\phi$ of an unknown one-dimensional projection $\theta^*$ of the input data. By focusing on Gaussian data, several recent works have built a remarkable picture, where the so-called information exponent (related to the regularity of the link function) controls the required sample complexity. In essence, these tools exploit the stability and spherical symmetry of Gaussian distributions.

In this work, we explore extensions of this picture beyond the Gaussian setting, where both stability or symmetry might be violated. Focusing on the planted setting where $\phi$ is known, our main results establish that Stochastic Gradient Descent recovers the unknown direction $\theta^*$ with constant probability in the high-dimensional regime, under mild assumptions that significantly extend ~[Yehudai and Shamir,20].

On Single Index Models beyond Gaussian Data

Abstract

Sparse high-dimensional functions have arisen as a rich framework to study the behavior of gradient-descent methods using shallow neural networks, showcasing their ability to perform feature learning beyond linear models. Amongst those functions, the simplest are single-index models f ( x ) = ϕ ( x · θ ∗ ) , where the labels are generated by an arbitrary non-linear scalar link function ϕ applied to an unknown one-dimensional projection θ ∗ of the input data. By focusing on Gaussian data, several recent works have built a remarkable picture, where the socalled information exponent (related to the regularity of the link function) controls the required sample complexity. In essence, these tools exploit the stability and spherical symmetry of Gaussian distributions. In this work, building from the framework of Ben Arous et al. [2021], we explore extensions of this picture beyond the Gaussian setting, where both stability or symmetry might be violated. Focusing on the planted setting where ϕ is known, our main results establish that Stochastic Gradient Descent can efficiently recover the unknown direction θ ∗ in the highdimensional regime, under assumptions that extend previous works Yehudai and Shamir [2020], Wu [2022].

1 Introduction

Over the past years, there has been sustained effort to enlarge our mathematical understanding of high-dimensional learning, particularly when using neural networks trained with gradient-descent methods - highlighting the interplay between algorithmic, statistical and approximation questions. An essential, distinctive aspect of such models is their ability to perform feature learning , or to extract useful low-dimensional features out of high-dimensional observations.

An appealing framework to rigorously analyze this behavior are sparse functions of the form f ( x ) = ϕ (Θ ⊤ ∗ x ) , where the labels are generated by a generic non-linear, low-dimensional function ϕ : R k → R of linear features Θ ⊤ ∗ x , with Θ ∗ ∈ R d × k with k ≪ d . While the statistical and approximation aspects of such function classes are by now well-understood Barron [1993], Bach [2017], the outstanding challenge remains computational, in particular in understanding the ability of gradientdescent methods to succeed. Even in the simplest setting of single-index models ( k = 1 ), and assuming that ϕ is known, the success of gradient-based learning depends on an intricate interaction between the data distribution x ∼ ν and the 'link' function ϕ ; and in fact computational lower bounds are known for certain such choices Yehudai and Shamir [2020], Song et al. [2021], Goel et al. [2020], Diakonikolas et al. [2017], Shamir [2018].

Positive results thus require to make specific assumptions, either about the data, or about the link function, or both. On one end, there is a long literature, starting at least with Kalai and Sastry [2009], Shalev-Shwartz et al. [2010], Kakade et al. [2011], that exploits certain properties of ϕ , such as invertibility or monotonicity, under generic data distributions satisfying mild anti-concentration properties Soltanolkotabi [2017], Frei et al. [2020], Yehudai and Shamir [2020], Wu [2022]. On the other end, by focusing on canonical high-dimensional measures such as the Gaussian distribution, the seminal works Ben Arous et al. [2021], Dudeja and Hsu [2018] built a harmonic analysis

framework of SGD, resulting in a fairly complete picture of the sample complexity required to learn generic link functions ϕ , and revealing a rich asymptotic landscape beyond the proportional regime n ≍ d , characterized by the number of vanishing moments, or information exponent s of ϕ , whereby n ≍ d s -1 samples are needed for recovery. Since then, several authors have built and enriched this setting to multi-index models Abbe et al. [2022, 2023], Damian et al. [2022], Arnaboldi et al. [2023], addressing the semi-parametric learning of the link function Bietti et al. [2022], as well as exploring SGD-variants Ben Arous et al. [2022], Barak et al. [2022], Berthier et al. [2023], Chen et al. [2023]. This harmonic analysis framework relies on two key properties of the Gaussian measure and their interplay with SGD: its spherical symmetry and its stability by linear projection. Together, they provide an optimization landscape that is well-behaved in the limit of infinite data, and enable SGD to escape the 'mediocrity' of initialisation, where the initial direction θ 0 , in the high-dimensional setting, has vanishingly small correlation | θ 0 · θ ∗ | ≃ 1 / √ d with the planted direction θ ∗ .

In this work, we study to what extent the 'Gaussian picture' is robust to perturbations, focusing on the planted setting where ϕ is known. Our motivation comes from the fact that real data is rarely Gaussian, yet amenable to being approximately Gaussian via CLT-type arguments. We establish novel positive results along two main directions: (i) when spherical symmetry is preserved but stability is lost, and (ii) when spherical symmetry is lost altogether. In the former, we show that spherical harmonics can be leveraged to provide a benign optimization landscape for SGD under mild regularity assumptions, for initialisations that can be reached with constant probability with the same sample complexity as in the Gaussian case. In the latter, we quantify the lack of symmetry with robust projected Wasserstein distances, and show that for 'quasi-symmetric' measures with small distance to the Gaussian reference, SGD efficiently succeeds for link functions with information exponent s ≤ 2 . Finally, using Stein's method, we address substantially 'non-symmetric' distributions, demonstrating the strength and versatility of the harmonic analysis framework.

2 Preliminaries and Problem Setup

The focus of this work is to understand regression problems with input/output data ( x, y ) ∈ R d × R generated by single-index models. This is a class of problems where the data labels are produced by a non-linear map of a one-dimensional projection of the input, that is

where ϕ : R → R is also known as the link function , and θ ∗ ∈ S d -1 , the sphere of R d , is the hidden direction that the models wants to learn. Quite naturally, the learning is made through the family of generalized linear predictors H = { ϕ θ : x → ϕ ( x · θ ) , for θ ∈ S d -1 } , built upon the link function (which is assumed known) and parametrized by the sphere.

Loss function. We assume that the input data is distributed according to a probability ν ∈ P ( R d ) . Equation (1) implies that the target function that produces the labels, ϕ θ ∗ , lies in this parametric class. The overall loss classically corresponds to the average over all the data of the square penalisation l ( θ, x ) := ( ϕ θ ( x ) -ϕ θ ∗ ( x )) 2 so that the population loss writes

where we used the notation ∥ f ∥ p L p ν = E ν [ | f | p ] , valid for all p ∈ N ∗ . Let us put emphasis on the fact that the loss L is a non-convex function of the parameter θ , hence it is not a priori guaranteed that gradient-based method are able to retrieve the ground-truth θ ∗ . This often requires a precise analysis of the loss landscape , and where the high-dimensionality can play a role of paramount importance: we place ourselves in this high-dimensional setting for which the dimension is fixed but considered very large d ≫ 1 . Finally, we assume throughout the article that ϕ θ belongs to the weighted Sobolev space W 1 , 4 ν := { ϕ, such that sup θ ∈S d -1 [ ∥ ϕ θ ∥ L 4 ν + ∥ ϕ ′ θ ∥ L 4 ν ] < ∞} .

Stochastic gradient descent. To recover the signal given by θ ∗ ∈ S d -1 , we run online stochastic gradient descent (SGD) on the sphere S d -1 . This corresponds to having at each iteration t ∈ N ∗ a fresh sample x t drawn from ν and independent of the filtration F t = σ ( x 1 , . . . , x t -1 ) and performing a spherical gradient step, with step-size δ > 0 , with respect to θ → l ( θ, x t ) :

An important scalar function that enables to track the progress of the SGD iterates is the correlation with the signal m θ := θ · θ ∗ ∈ [ -1 , 1] . We will drop the subscript in case there is no ambiguity. Note that, due to the high-dimensionality of the setup, we have the following lemma: √

Lemma 2.1. For all a > 0 , we have P θ 0 ( m θ 0 ≥ a/ d ) ≤ a -1 e -a 2 / 4 . Additionally, for any δ > 0 such that max { a, δ } ≤ √ d/ 4 , we have the lower bound: P θ 0 ( m θ 0 ≥ a/ √ d ) ≥ δ 4 e -( a + δ ) 2 .

This fact implies that, when running the algorithm in practice, it is initialized with high probability near the equator of S d -1 , or at least in a band of typical size 1 / √ d (see Figure 1 for a schematic illustration of this fact). Finally, we use the notation ∇ S θ to denote the spherical gradient, that is ∇ S θ l ( θ, x ) = ∇ θ l ( θ, x ) -( ∇ θ l ( θ, x ) · θ ) θ . As ∇ S θ l ( · , x t ) is an unbiased estimate of ∇ S θ L , it is expected that the latter gradient field rules how the SGD iterates travel across the loss landscape.

Loss landscape in the Gaussian case. As stressed in the introduction, this set-up has been studied by Dudeja and Hsu [2018], Ben Arous et al. [2021] in the case where ν is the standard Gaussian, noted as γ here to avoid any confusion for later. Let us comment a bit this case to understand what can be the typical landscape of this single-index problem. Thanks to the spherical symmetry, the loss admits a scalar summary statistic, given precisely by the correlation m θ . Moreover, the loss admits an explicit representation in terms of the Hermite decomposition of the link function ϕ : if { h j } j denotes the orthonormal basis of Hermite polynomials of L 2 γ , then L ( θ ) = 2 ∑ j |⟨ ϕ, h j ⟩| 2 (1 -m j ) := ¯ ℓ ( m ) . As a result, the gradient field projected along the signal is a (locally simple) positive function of the correlation that behaves similarly to

where s ∈ N ∗ is the index of the first non-zero of the Hermite coefficients {⟨ ϕ, h j ⟩} j . This has at least three important consequences for the gradient flow: (i) if initialized positively, the correlation is an increasing function along the dynamics and there is no bad local minima in the loss landscape, (ii) the parameter s ∈ N ∗ controls the flatness of the loss landscape near the origin and therefore controls the optimization speed of SGD in this region (iii) as soon as m is large enough, the contractive term 1 -m makes the dynamics converge exponentially fast.

Loss landscape in general cases. Obviously for general distributions ν , the calculation presented in Eq.(4) is no-longer valid. However, the crux of the present paper is that properties (i)-(ii)-(iii) are robust to the change of distribution and can be shown to be preserved under small adaptations. More precisely, we have the following definition.

Definition 2.2 (Local Polynomial Growth) . We say that L has the local polynomial growth of order k ∈ N ∗ and scale b ≥ 0 , if there exists C > 0 such that for all m θ ≥ b ,

In such a case we say that L satisfies LPG ( k, b ) .

In this definition, and as showed later in specific examples given in Section 4, the scale parameter b should be thought as a small parameter proportional to 1 / √ d . If ν is Gaussian, we can rewrite Eq.(4) and show that L verifies LPG ( s, 0) for s ∈ N ∗ , referred to as the information exponent of the problem in Ben Arous et al. [2021]. An important consequence of satisfying LPG ( k, b ) is that the the population landscape is free of bad local minima outside the equatorial band Σ b := { θ ∈ S d -1 , m θ ≤ b } . Therefore, when b is of scale 1 / √ d , Lemma 2.1 indicates that one can efficiently produce initializations that avoid it. Hence, this property is the fundamental ingredient that enables the description of the path taken by SGD that we derive it in the next Section. Section 4 is devoted to showcasing generic examples when this property is satisfied.

3 Stochastic Gradient Descent under LPG

In this section, we derive the main results on the trajectory of the stochastic gradient descent. They state that the property LPG ( s, b/ √ d ) is in fact sufficient to recover the same quantitative guarantees as the one depicted in Ben Arous et al. [2021], despite the lack of Gaussianity of the distribution ν . Recall that the recursion satisfied by the SGD iterates is given by Eq.(3). To describe their movement, let us introduce the following notations: for all t ∈ N ∗ , we denote the normalization by r t := ∣ ∣ θ t -δ ∇ S θ ℓ ( θ t , x t ) ∣ ∣ and the martingale induced by the stochastic gradient descent as M t := l ( θ t , x t ) -E ν [ l ( θ t , x )] .

Moment growth assumptions. To be able to analyse the SGD dynamics, we make the following assumptions on the moments of the martingale increments induced by the random sampling. Let us denote for all θ ∈ R d , x θ = x · θ ∈ R and C ( u, v ) = ϕ ′ ( u ) ϕ ( v ) , for u, v ∈ R .

Assumption 3.1 (Moment Growth Assumption) . There exists a constant K > 0 , independent of the dimension d , such that

A precise care is given to the dependency in the dimension in the upper bound to match the practical cases that we later discuss in Section 4. Note that these assumptions are typically true for sub-gaussian random variables if ϕ belongs to a Sobolev regularity class. These assumptions are similar to the one given in Eqs. (1.3)-(1.4) in Arous et al. [2021] for the Gaussian case. In all the remainder of the section we assume that Assumption 3.1 is satisfied.

Tracking the correlation. Recall that the relevant signature of the dynamics is the one-dimensional correlation m t = θ t · θ ∗ . For infinitesimal step-sizes δ → 0 , it is expected that θ t follow the spherical gradient flow ˙ θ t = -∇ S θ L ( θ t ) , that translates to the evolution on the summary statistics m t whereby

The main idea behind the result of this section is to show that, even if the energy landscape near m = m 0 is rough at scale 1 / √ d , the noise induced by SGD does not prevent m to grow as the idealized dynamics described by the ODE (8). Let us write the iterative recursion followed by ( m t ) t ≥ 0 : with the notation recalled above, for t ∈ N ∗ , we have

With this dynamics at hand, the proof consists in controlling both the discretization error through r t and the directional martingale induced by the term δ ∇ S θ M t · θ ∗ .

Weak recovery. As it is the case for the gradient flow, most of the time spent by the SGD dynamics is near the equator, or more precisely in a band of the type Σ b,c = { θ ∈ S d -1 , b/ √ d ≤ m θ ≤ c/ √ d } , where b < c are constants independent of the dimension. Hence, the real first step of the dynamics is to go out any of these bands. This is the reason why it is natural to define S a := { θ ∈ S d -1 , m θ ≥ a } , the spherical cap of level a ∈ (0 , 1) as well as the hitting time

which corresponds to the first time ( θ t ) t ≥ 0 enters in S a . We fix a dimension-independent constant a = 1 / 2 , and refer to the related hitting time τ + 1 / 2 as the weak recovery time of the algorithm.

Theorem 3.2 (Weak Recovery) . Let ( θ t ) t ≥ 0 follow the SGD dynamics of Eq. (3) and let L satisfy LPG ( s, b/ √ d ) , with b > 0 and s ∈ N ∗ , then, conditionally on the fact that m 0 ≥ 5 b/ √ d , for any 0 < ε ≤ ε ∗ , we have

with probability at least 1 -K ε , for generic constants K , ε ∗ > 0 that depend solely on the link function ϕ and the distribution ν .

Let us comment on this result. It says that that the integer s coming from the growth condition controls the hardness of exiting the equator of the sphere. Indeed, as can be seen in LPG ( s, b/ √ d ) , the larger the s , the smaller the gradient projection is and hence the less information the SGD dynamics has to move from the initialization. This result can be seen as an extension of [Ben Arous et al., 2021, Theorem 1.3] valid only in the Gaussian case ( b = 0 ). Furthermore, the Gaussian case shows that Theorem 3.2 is tight up to log( d ) factors. Finally, note that the result is conditional to the fact that the

•

θ

w

θ

∗

θ

0

√

d

initialization is larger that some constant factor of 1 / √ d , which has at least constant probability to happen in virtue of Lemma 2.1. This probability can be lowered by any constant factor by sampling offline a constant factor (independent of d ) of i.i.d. initializations and keeping the one that maximizes its correlation with θ ∗ . Finally, note that in all the cases covered by the analysis, it is possible to keep track of the constant C , and show that overall it depends only (i) on the property of ν w.r.t. the Gaussian on the one hand and (ii) on the Sobolev norm of the link function ∥ ϕ ∥ W 1 , 4 ν on the other.

Strong recovery. We place ourselves after the weak recovery time described in the previous question and want to understand if the ( θ t ) t ≥ 0 dynamics goes to θ ∗ and if yes, how fast it does so. This is what we refer to as the strong recovery question, captured by the fact that the one-dimensional summary statistics m go towards 1 . Thanks to the Markovian property of the SGD dynamics, we have the equality between all time s > 0 marginal laws of

and hence the strong recovery question is equivalent to study the dynamics with initialization that has already weakly recovered the signal, i.e. such that m θ = 1 / 2 . We show that this part of the loss landscape is very different that the equator band in which the dynamics spends most of its times: in all the cases, we can choose stepsizes independent of the dimension and show that the time to reach the vicinity of θ ∗ will be independent of d .

Theorem 3.3 (Strong Recovery) . Let ( θ t ) t ≥ 0 follow the SGD dynamics of Eq. (3) and let L satisfy LPG ( s, b/ √ d ) , with b > 0 and s ∈ N ∗ , then, for any ε > 0 , taking δ = ε/d , we have that there exists a time T > 0 , such that

with probability at least 1 -K ε , for some generic K > 0 that depends solely of the link function ϕ .

As introduced above, the important messages conveyed by this theorem are that (i) there is no difference between the different parameters setups captured by the information exponent s , and (ii) the time it takes to reach an ε -vicinity of θ ∗ is always strictly smaller than the one needed to exit the weak recovery phase (e.g. d compared to d s -1 when s ≥ 3 ). This means that the dynamics spends most of its time escaping the mediocrity . Remark that we decided to present Theorem 3.3 resetting the step-size δ to put emphasis on the intrinsic difference between the two phases. Yet, we could have kept the same stepsize at the expense of slowing the second phase.

4 Typical cases of loss landscape with LPG property

In this section, we showcase two prototypical cases where the LPG holds true: the section 4.1 deals with the spherically symmetric setting, whereas the section 4.2 describes a perturbative regime where the distribution is approximately Gaussian is a quantitative sense.

4.1 The symmetric case

We start our analysis with the spherically symmetric setting. We show that a spherical harmonic decomposition provides a valid extension of the Hermite decomposition in the Gaussian case, leading to essentially the same quantitative performance up to constant (in dimension) factors.

Spherical Harmonic Representation of the Population Loss. We express the data distribution ν as a mixture of uniform measures ν = ∫ ∞ 0 τ r,d ρ ( dr ) , where τ r,d = Unif( r S d -1 ) . Let τ d = τ 1 ,d and u d ∈ P ([ -1 , 1]) be the projection of τ d onto one direction, with density given in close form by u d ( dt ) = Z -1 (1 -t 2 ) ( d -3) / 2 1 ( | t | ≤ 1) dt , where Z is a normarlizing factor. Let { P j,d } j ∈ N be the orthogonal basis of Gegenbauer polynomials of L 2 u d ([ -1 , 1]) , normalized such that P j,d (1) = 1 for all j, d . For each r > 0 , consider

where we define ϕ ( r ) : [ -1 , 1] → R such that ϕ ( r ) ( t ) := ϕ ( rt ) . We write its decomposition in L 2 u d ([ -1 , 1]) as ϕ ( r ) = ∑ j α j,r,d P j,d , with α j,r,d = ⟨ ϕ ( r ) ,P j,d ⟩ ∥ P j,d ∥ 2 . Let Ω d be the Lebesgue measure of S d -1 , and N ( j, d ) = 2 d + j -2 d ( d + j -3 d -1 ) the so-called dimension on the spherical harmonics of degree j in dimension d . From the Hecke-Funk representation formula [Frye and Efthimiou, 2012, Lemma 4.23] and the chosen normalization P j,d (1) = 1 , we have ∥ P j,d ∥ 2 = ( Ω d -1 Ω d -2 N ( j,d ) ) 1 / 2 [Frye and Efthimiou, 2012, Proposition 4.15] and obtain finally

where we defined for convenience ¯ α j,r,d = α j,r,d / √ N ( j, d ) . As a result, it follows that the overall loss writes as solely the correlaton m θ = θ · θ ∗ as

where β j,d = ∫ ∞ 0 ¯ α 2 j,r,d ρ ( dr ) ≥ 0 . Unsurprisingly, we observe that, thanks to the spherical symmetry and analogous to the Gaussian case, the loss still admits m = θ · θ ∗ as a summary statistics. Yet, it is represented in terms of Gegenbauer polynomials, rather than monomials as in the Gaussian case. The monomial representation is a consequence of the stability of the Gaussian measure, as seen by the fact that ⟨ ϕ θ , ϕ ˜ θ ⟩ γ d = ⟨ ϕ, A θ · ˜ θ ϕ ⟩ γ , where ( A m f )( t ) = E z ∼ γ [ f ( mt + √ 1 -m 2 z )] has a semi-group structure (it is even known in fact as the Ornstein-Ulhenbeck semi-group).

Let η ∈ P ( R ) be the marginal of ν along any direction. The following proposition charaterizes the coefficients β j,d as integrals over the radial distribution ρ and projections of the link function ϕ :

Proposition 4.1 (Loss representation) . The β j,d defined in (14) have the integral representation

where K j is a positive semi-definite integral operator of L 2 η that depends solely on ρ and ϕ .

Note that a closed form expression of K j can be found in Appendix D.2. The above proposition is in fact the stepping stone to calculate properly the information exponent that plays a crucial role in the property LPG . This is given through the link between the spectrum β j,d and the decomposition of ϕ in the L 2 η orthogonal basis of polynomials, that we denote by { q j } j .

̸

Thus, the number of vanishing moments of ϕ with respect to the data marginal η provides an upper bound on the 'effective' information exponent of the problem s , as we will see next.

Local Polynomial Growth. From (14), and as ∇ S θ L ( θ ) = ℓ ′ ( m )( θ ∗ -mθ ) , we directly obtain

which is the quantity we want to understand to exhibit the property LPG in this case. Hence, we now turn into the question of obtaining sufficient guarantees on the coefficients ( β j,d ) j that ensure

√

local polynomial growth. Since the typical scale of initialization for m is Θ(1 / d ) , our goal is to characterize sufficient conditions of local polynomial growth with b = O (1 / √ d ) .

For that purpose, let us define two key quantities of Gegenbauer polynomials:

We have the following sufficient condition based on the spectrum ( β j,d ) j :

Proposition 4.3 (Spectral characterization of LPG ) . Suppose there exist constants K,C > 0 and s ∈ N such that we both have β s,d ≥ C and ∑ j>s β j,d j ( j + d -2) υ j -1 ,d +2 ≤ Kd (3 -s ) / 2 . Then, taking s ∗ as the infimum of such s , L has the property LPG ( s ∗ -1 , z s ∗ ,d ) . In particular, whenever s ∗ ≪ d , we have z s ∗ ,d ≤ 2 √ s ∗ /d .

This proposition thus establishes that, modulo a mild regularity assumption expressed though the decay of the coefficients β j,d , the spherically symmetric non-Gaussian setting has the same geometry as the Gaussian setting, for correlations slightly above the equator. Crucially, the required amount of correlation to 'feel' the local polynomial growth is a O ( √ s ) factor from the typical initialization, and can be thus obtained with probability ≃ e -s over the initialization, according to Lemma 2.1, a lower bound which is independent of d .

The sufficient condition for β j,d appearing in Proposition 4.3 involves the minimum values υ j,d of Gegenbauer polynomials P j,d , as well as sums of the form ∑ j j 2 β j,d . In order to obtain a more user-friendly condition, we now provide an explicit control of υ j,d , and leverage mild regularity of ϕ to control ∑ j j 2 β j,d . This motivates the following assumption on ϕ and ν :

Assumption 4.4. The link function ϕ satisfies ϕ ∈ L 2 η and ϕ ′ ∈ L 4 η , and the radial distribution ρ has finite fourth moment E ρ [ r 4 ] < ∞ .

̸

Theorem 4.5 ( LPG for symmetric distributions) . Assume that ϕ and ν satisfy Assumption 4.4, and let s ∗ = inf { j ; ⟨ ϕ, q j ⟩ η = 0 } . Then L has the property LPG ( s ∗ -1 , 2 √ s ∗ /d ) .

The proof is provided in Appendix D.5. At a technical level, the main challenge in proving Theorem 4.5 is to achieve a uniform control of υ j,d in j , a result which may be of independent interest. We address it by combining state-of-the-art bounds on the roots of the Gegenbauer polynomials, allowing us to cover the regime where j is small or comparable to d , together with integral representations via the Cauchy integral formula, providing control in the regime of large j . On the other hand, we relate the sum ∑ j j 2 β j,d to a norm of ϕ ′ using a Cauchy-Schwartz argument, where we leverage the fourth moments from Assumption 4.4.

Remark 4.6. Since we are in a setting where ϕ is known, an alternative to the original recovery problem from Eq (2) is to consider a pure Gegenbauer 'student' link function of the form ˜ ϕ = P s,d , where s is the information exponent from Proposition 4.2. Indeed, the resulting population loss ˜ L ( θ ) = E [( ˜ ϕ ( x · θ ) -˜ ϕ ( x · θ ∗ )) 2 ] satisfies the LPG property, as easily shown in Fact D.5.

For the sake of completeness, we describe more precisely two concrete case studies below.

̸

Example 4.7 (Uniform Measure on the Sphere) . When ν = Unif ( d S d -1 ) , we have ρ = δ √ d , and therefore β j,d = ¯ α 2 j, √ d,d . In that case, the orthogonal polynomial basis { q j ( t ) } j of L 2 η coincides with the rescaled Gegenbauer polynomials, q j ( t ) = P j,d ( t/ √ d ) . Consider now a link function ϕ with s -1 vanishing moments with respect to L 2 η , i.e. such that ¯ α j,d = ⟨ ϕ, q j ⟩ η = 0 for j < s and ¯ α s,d = ⟨ ϕ, q s ⟩ η = 0 ; and with sufficient decay in the higher harmonics as to satisfy the bound on the sum presented in Proposition 4.3 (for example, ϕ ( t ) = q s ( t ) trivially satisfies this condition). Then Proposition 4.3 applies and we conclude that the resulting population landscape satisfies LPG ( s -1 , O ( √ s/d )) .

In Yehudai and Shamir [2020], Wu [2022] it is shown that monotonically increasing link functions 1 lead to a benign population landscape, provided the data distribution ν satisfies mild anti-concentration properties. We verify that in our framework.

√

̸

Example 4.8 (Non-decreasing ϕ ) . Indeed, Proposition 4.3 is verified with s = 1 , provided ϕ ′ ∈ L 2 η . Indeed, if ϕ = 0 is monotonic, then we have β 1 = ⟨ ϕ, K 1 ϕ ⟩ η = C d ( E η [ tϕ ( t )]) 2 = 0 , since we can assume without loss of generality that ϕ ( t ) ≥ 0 for t ≥ 0 and ϕ ( t ) ≤ 0 for t ≤ 0 .

̸

Weemphasize that the results of Yehudai and Shamir [2020], Wu [2022] extend beyond the spherically symmetric setting, which is precisely the focus of next section.

4.2 Non-Spherically Symmetric Case

We now turn to the setting where ν is no longer assumed to have spherical symmetry. By making further regularity assumptions on ϕ , our main insight is that distributions that are approximately symmetric (defined in an appropriate sense) still benefit from a well-behaved optimization landscape.

Two-dimensional Wasserstein Distance. When ν is not spherically symmetric, the machinery of spherical harmonics does not apply, and we thus need to rely on another structural property. Consider a centered and isometric data distribution ν ∈ P 2 ( R d ) , i.e. such that E ν x = 0 and Σ ν = E ν [ xx ⊤ ] = I d . We consider the two-dimensional 1 -Wassertein distance [Niles-Weed and Rigollet, 2022, Definition 1] -see also Paty and Cuturi [2019]- between a pair of distributions ν a , ν b ∈ P ( R d ) , defined as

where the supremum runs for any two-dimensional subspace P ∈ Gr(2 , d ) , and P # ν ∈ P ( R 2 ) is the projection (or marginal) of ν onto the span of P . ˜ W 1 , 2 is a distance ([Paty and Cuturi, 2019, Proposition 1]) and measures the largest 1 -Wasserstein distance over all two-dimensional marginals.

Weare in particular interested in the setting where ν a = ν is our data distribution, and ν b is a reference symmetric measure - for instance the standard Gaussian measure γ d . Consider the fluctuations

where ¯ ℓ ( m ) is the Gaussian loss defined in Section 2 and L ( θ ) = E ν [ | ϕ ( x · θ ) -ϕ ( x · θ ∗ ) | 2 ] . ∆ L and ∆ ∇ L thus measure respectively the fluctuations of the population loss and the relevant (spherical) gradient direction. By making additional mild regularity assumptions on the link function ϕ , we can obtain a uniform control of the population loss geometry using the dual representation of the 1 -Wasserstein distance.

Assumption 4.9 (Regularity of link function) . We assume that ϕ, ϕ ′ are both B -Lipschitz, and that ϕ ′′ ( t ) = O (1 /t ) .

Assumption 4.10 (Subgaussianity) . The data distribution ν is M -subgaussian: for any v ∈ S d -1 , we have ∥ x · v ∥ ψ 2 ≤ M , where ∥ z ∥ ψ 2 := inf { t > 0; E [exp( z 2 /t 2 ) ≤ 2 } is the Orlitz-2 norm.

Proposition 4.11 (Uniform gradient approximation) . Under Assumptions 4.9 and 4.10, for all θ ∈ S d -1 ,

where the O ( · ) notation only hides constants appearing in Assumptions 4.9 and 4.10.

In words, the population gradient under ν is viewed as a perturbation of the population gradient under γ , which has the well-behaved geometry already described in Section 2. These perturbations can be uniformly controlled by the projected 1-Wasserstein distance, thanks to the subgaussian tails of ν .

Our focus will be in situations where ˜ W 1 , 2 ( ν, γ ) = O (1 / √ d ) . This happens to be the 'natural' optimistic scale for this metric in the class of isotropic distributions Σ ν = I d , as can be seen for instance when ν = Unif( √ d S d -1 ) . Under such conditions, it turns out that link functions with information exponent s ≤ 2 can be recovered with simple gradient-based methods, by paying an additional polynomial (in d ) cost in time complexity.

Assumption 4.12. The Gaussian information exponent of ϕ , s := arg min { j ; ⟨ ϕ, H j ⟩ ̸ = 0 } satisfies s ≤ 2 .

Assumption 4.13. The projected Wasserstein distance satisfies ˜ W 1 , 2 ( ν, γ ) ≤ M ′ / √ d .

Proposition 4.14 ( LPG , non-symmetric setting) . Under Assumptions 4.9, 4.10, 4.12 and 4.13, L verifies LPG ( 1 , O (√ (log d κ ) d )) , where κ depends only on B,M,M ′ .

This proposition illustrates the cost of breaking spherical symmetry in two aspects: (i) it requires additional regularity on ϕ , and notably restricts its (Gaussian) information exponent to s = 2 , and (ii) the scale to reach LPG is now no longer dimension-free, but has a polynomial dependency on dimension, since from Lemma 2.1, picking δ any positive constant we have

We are not able to rule this out as a limitation of our proof; whether this polynomial dependency on dimension is inherent to the symmetry breaking is an interesting question for future work.

While assumptions 4.9, 4.10 and 4.12 are transparent and impose only mild conditions on the link function and tails of ν , the 'real' assumption of Proposition 4.14 is the concentration of ˜ W 1 , 2 ( ν, γ ) (Assumption 4.13). The ball { ν ; ˜ W 1 , 2 ( ν, γ ) = O (1 / √ d ) } contains many non-symmetric measures, for instance empirical measures sampled from γ with n = ω ( d 2 ) [Niles-Weed and Rigollet, 2022, Proposition 8], and we suspect it contains many other examples, such as convolutions of the form ν ∗ γ σ arising for instance in diffusion models. That said, one should not expect the distance ˜ W 1 , 2 ( ν, γ ) to be of order 1 / √ d for generic 'nice' distributions ν ; for instance, log-concave distributions are expected to satisfy W 1 ( P # ν, γ ) ≃ 1 / √ d for most subspaces P , as captured in Klartag's CLT for convex bodies Klartag [2007].

Localized Growth via Stein's method. To illustrate the mileage of the previous techniques beyond this 'quasi-symmetric' case, we consider now an idealised setting where the data is drawn from a product measure ν = η ⊗ d , with η ∈ P ( R ) and W 1 ( η, γ 1 ) = Θ(1) . In other words, x = ( x 1 , . . . , x d ) ∼ ν if x i ∼ η are i.i.d. In this setting, the distances W 1 ( P # ν, γ ) reflect a CLT phenomena, which requires the subspace P to 'mix' across independent variables. Consequently, one may expect the expression of the hidden direction θ ∗ in the canonical basis to play a certain role so we make the following additional regularity assumption on the tails of ϕ :

Assumption 4.15 (Additional Regularity in third derivatives) . ϕ admits four derivatives bounded by L , with | ϕ (3) ( t ) | = O (1 /t ) and | ϕ (4) ( t ) | = O (1 /t 2 ) . Moreover, the third moment of the data distribution is finite: τ 3 = E t ∼ η [ t 3 ] < ∞ .

Stein's method provides a powerful control on ∆ L ( θ ) and ∆ ∇ L ( θ ) , as shown by the following result:

Proposition 4.16 (Stein's method for product measure) . Let χ ( θ, θ ∗ ) := ∥ θ ∥ 2 4 + ∥ θ ∗ ∥ 2 4 . Under Assumptions 4.9, 4.10 and 4.15, there exists a universal constant C = C ( M,B,τ 3 ) such that

The proof is based on the Stein's method for multivariate variables [Röllin, 2013, Theorem 3.1] with independent entries, which provides a quantitative CLT bound. Contrary to the quasi-symmetric case, here the concentration is not uniform over the sphere, but crucially depends on the sparsity of both θ and θ ∗ , measured via the ℓ 4 norms ∥ θ ∥ 4 , ∥ θ ∗ ∥ 4 : for incoherent, non-sparse directions, we have ∥ θ ∥ 2 4 ≃ 1 / √ d , recovering the concentration rate that led to Proposition 4.14, while for sparse directions we have ∥ θ ∥ 2 4 = Θ(1) , indicating an absence of concentration to the Gaussian landscape.

Therefore, the natural conclusion is to assume a planted model where θ ∗ is incoherent with the data distribution, i.e. ∥ θ ∗ ∥ 4 = O ( d 1 / 4 ) . While the LPG property does not directly apply in this setting, we outline an argument that suggests that the single-index model can still be efficiently solved using gradient-based methods. For that purpose, we assume that θ ∗ is drawn uniformly in S d -1 , which implies that its squaredL 4 norm ∥ θ ∗ ∥ 2 4 is of order d -1 / 2 with high probability:

Because θ 0 is also drawn uniformly on the sphere, the typical value of χ ( θ, θ ∗ ) is of order d -1 / 2 . For Gaussian information exponent s = 2 , the population gradient under Gaussian data satisfies

-∇ S d -1 θ L γ ( θ ) · θ ∗ ≥ Cm θ . As a consequence, by Proposition 4.16, whenever | θ 0 · θ ∗ | > c √ d (which happens with constant probability lower bounded by e -c ), we enter a 'local' LPG region where -∇ S d -1 θ L ( θ ) · θ ∗ ≥ C ( m -c/ √ d ) > 0 . While this condition is sufficient in the quasi-symmetric setting to start accumulating correlation (Theorem 3.2), now this event is conditional on θ being dense, ie so that χ ( θ, θ ∗ ) = O (1 / √ d ) .

Since the typical value of χ ( θ, θ ∗ ) is of scale 1 / d , one would expect that SGD will rarely visit sparse points where χ ( θ, θ ∗ ) ≫ O (1 / √ d ) , and thus that the local LPG property will be valid for most times during the entropic phase of weak recovery, summarized by the following conjecture:

√

Conjecture 4.18 (SGD avoids sparse points) . Assume θ ∗ , θ 0 are drawn from the uniform measure, and let θ t be the t -th iterate of SGD with δ ≃ 1 / ( d log d ) . There exists a universal constant C such that for any ξ > 0 , we have

Since the time to escape mediocrity in the case s = 2 is T ≃ d log( d ) 2 , this conjecture would imply that SGD does not effectively 'see' any sparse points, and thus escapes mediocrity. If one assumed that in this phase the dynamics is purely noisy, now pretending that θ i were drawn independently from the uniform measure, and that ∥ θ ∥ 4 4 is approximately Gaussian with mean d -1 and variance d -3 , the result follows by simple concentration. The challenging aspect of Conjecture 4.18 is precisely to handle the dependencies across iterates, as well as the spherical projection steps.

5 Experiments

In order to validate our theory, and inspect the degree to which our bounds may be pessimistic, we consider empirical evaluation of the training process in our two primary settings. Specifically, we consider random initialization on the half-sphere (with the sign chosen to induce positive correlation as in Arous et al. [2021]), and investigate how often strong recovery occurs relative to the information exponent of the link function. Ultimately the empirical evidence corroborates the theory, showing the symmetric case succeeds when the initial correlation is above a polynomial threshold and the nonsymmetric case succeeds as long as s = 2 . Due to space constraints, the details of the experimental setup and further discussion appear in Section A.

6 Conclusion and Perspectives

In this work, we have asked whether the remarkable properties of high-dimensional Gaussian SGD regression of single-index models are preserved as one loses some key aspects that make Gaussian distributions so special (and so appealing for theorists). Our results are mostly positive, indicating a robustness of the Gaussian theory, especially within the class of spherically symmetric distributions, where a rich spherical harmonic structure is still available. As one loses spherical symmetry, the situation becomes more dire, motivating a perturbative analysis that we have shown is effective via projected Wasserstein and Stein couplings.

That said, there are several open and relevant avenues that our work has barely touched upon, such as understanding whether the robustness can be transferred to other algorithms beyond SGD, or addressing the semi-parametric problem when the link function is unknown, along the lines of Bietti et al. [2022], Abbe et al. [2023], Damian et al. [2022], Berthier et al. [2023]. A particularly interesting direction of future work is to extend the analysis of product measures to 'weakly dependent' distributions, motivated by natural images where locality in pixels captures most (but not all) of the statistical dependencies. Stein's method appears to be a powerful framework that can accommodate such weak dependencies, and deserves future investigation.

Acknowledgements This work was partially supported by NSF DMS 2134216, NSF CAREER CIF 1845360, NSF IIS 1901091 and the Alfred P Sloan Foundation.

References

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. arXiv preprint arXiv:2202.08658 , 2022.

Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055 , 2023.

Iván Area, Dimitar K. Dimitrov, Eduardo Godoy, and André Ronveaux. Zeros of gegenbauer and hermite polynomials and connection coefficients. Math. Comput. , 73:1937-1951, 2004.

Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. arXiv preprint arXiv:2302.05882 , 2023.

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.

Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research , 18(1):629-681, 2017.

Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35:21750-21764, 2022.

Andrew R Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information theory , 39(3):930-945, 1993.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research (JMLR) , 22:106-1, 2021.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. arXiv preprint arXiv:2206.04030 , 2022.

Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055 , 2023.

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 , 2022.

Sitan Chen, Zehao Dou, Surbhi Goel, Adam R Klivans, and Raghu Meka. Learning narrow onehidden-layer relu networks. arXiv preprint arXiv:2304.10524 , 2023.

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

Laura De Carli. Local lp inequalities for gegenbauer polynomials. In Topics in classical analysis and applications in honor of Daniel Waterman , pages 73-87. World Scientific, 2008.

Ilias Diakonikolas, Daniel M Kane, and Alistair Stewart. Statistical query lower bounds for robust estimation of high-dimensional gaussians and gaussian mixtures. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) , pages 73-84. IEEE, 2017.

Dimitar K. Dimitrov and Geno P. Nikolov. Sharp bounds for the extreme zeros of classical orthogonal polynomials. Journal of Approximation Theory , 162(10):1793-1804, 2010. ISSN 0021-9045. doi: https://doi.org/10.1016/j.jat.2009.11.006. URL https://www.sciencedirect.com/science/ article/pii/S0021904509001993 . Special Issue dedicated to the memory of Borislav Bojanov.

DLMF. NIST Digital Library of Mathematical Functions . https://dlmf.nist.gov/ , Release 1.1.10 of 2023-06-15, 2022. URL https://dlmf.nist.gov/ . F. W. J. Olver, A. B. Olde Daalhuis, D. W. Lozier, B. I. Schneider, R. F. Boisvert, C. W. Clark, B. R. Miller, B. V. Saunders, H. S. Cohl, and M. A. McClain, eds.

Driver and K. Jordaan. Bounds for extreme zeros of some classical orthogonal polynomials. Journal of Approximation Theory , 164(9):1200-1204, 2012. ISSN 0021-9045. doi: https://doi. org/10.1016/j.jat.2012.05.014. URL https://www.sciencedirect.com/science/article/ pii/S0021904512001050 .

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 .

David A Freedman. On tail probabilities for martingales. the Annals of Probability , pages 100-118, 1975.

Spencer Frei, Yuan Cao, and Quanquan Gu. Agnostic learning of a single neuron with gradient descent. Advances in Neural Information Processing Systems , 33:5417-5428, 2020.

Christopher Frye and Costas J Efthimiou. Spherical harmonics in p dimensions. arXiv preprint arXiv:1205.3548 , 2012.

Walter Gautschi. Some elementary inequalities relating to the gamma and incomplete gamma function. J. Math. Phys , 38(1):77-81, 1959.

Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans. Superpolynomial lower bounds for learning one-layer neural networks using gradient descent. In International Conference on Machine Learning , pages 3587-3596. PMLR, 2020.

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.

Adam Tauman Kalai and Ravi Sastry. The isotron algorithm: High-dimensional isotonic regression. In COLT , 2009.

Boaz Klartag. A central limit theorem for convex sets. Inventiones mathematicae , 168(1):91-131, 2007.

Andrea Laforgia and Pierpaolo Natalini. On some inequalities for the gamma function. Advances in Dynamical Systems and Applications , 8(2):261-267, 2013.

Jonathan Niles-Weed and Philippe Rigollet. Estimation of wasserstein distances in the spiked transport model. Bernoulli , 28(4):2663-2688, 2022.

François-Pierre Paty and Marco Cuturi. Subspace robust wasserstein distances. In International conference on machine learning , pages 5072-5081. PMLR, 2019.

Adrian Röllin. Stein's method in high dimensions with applications. In Annales de l'IHP Probabilités et statistiques , volume 49, pages 529-549, 2013.

Shai Shalev-Shwartz, Ohad Shamir, and Karthik Sridharan. Learning kernel-based halfspaces with the zero-one loss. arXiv preprint arXiv:1005.3681 , 2010.

Ohad Shamir. Distribution-specific hardness of learning neural networks. The Journal of Machine Learning Research , 19(1):1135-1163, 2018.

Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information processing systems , 30, 2017.

Min Jae Song, Ilias Zadik, and Joan Bruna. On the cryptographic hardness of learning single periodic neurons. Advances in Neural Processing Systems (NeurIPS) , 2021.

Adriaan J Stam. Limit theorems for uniform distributions on spheres in high-dimensional euclidean spaces. Journal of Applied probability , 19(1):221-228, 1982.

Gabor Szego. Orthogonal polynomials , volume 23. American Mathematical Soc., 1939.

F Ursell. Integrals with nearly coincident branch points: Gegenbauer polynomials of large degree. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences , 463(2079): 697-710, 2007.

George Neville Watson. A treatise on the theory of Bessel functions , volume 3. The University Press, 1922.

Lei Wu. Learning a single neuron for non-monotonic activation functions. In International Conference on Artificial Intelligence and Statistics , pages 4178-4197. PMLR, 2022.

Gilad Yehudai and Ohad Shamir. Learning a single neuron with gradient methods. In Conference on Learning Theory , pages 3756-3786. PMLR, 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the problem of learning a single-index model beyond Gaussian data. Prior work (Ben Arous et al. 2021) studies the optimization landscape of recovering the hidden direction $\\theta^*$ given a known link function $\\phi$ when the covariates are drawn from the standard Gaussian distribution; in this case, the sample complexity is governed by the information exponent of $\\phi$. This paper’s key contribution is defining the “Local Polynomial Growth” condition, which characterizes the geometry of the population loss landscape. The main results, Theorems 3 and 4, show that online SGD recovers the hidden direction with sample complexity depending on the order parameter in the LPG condition. The proof of this main theorem relies on the same martingale-plus-drift decomposition from (Ben Arous et al. 2021). The paper concludes with specific examples of data distributions satisfying the LPG condition; for example, spherically symmetric distributions.\n\nSTRENGTHS:\n- The paper is well-written and easy to understand.\n- Learning a single-index model is a well studied problem in the statistics literature and is a primitive for neural network learning. Prior work in this area either assumes Gaussian data or that the link function satisfies basic properties such as monotonicity. This paper initiates the study of the landscape properties of learning single-index models with arbitrary data distribution and link function, which is to the best of my knowledge novel and an important direction of future work.\n- The local polynomial growth definition is an interesting and useful way to abstract away the key ideas of the proof in (Ben Arous et al. 2021), and has not appeared previously in the literature. This definition allows for their key idea of escaping a higher-order saddle at initialization to apply to a larger family of distributions more generally.\n- Section 4.1 shows how the LPG framework can be used to learn single index models over arbitrary spherically symmetric distributions, which is a novel and interesting result.\n\nWEAKNESSES:\n- There is one major weakness I am concerned about. Theorem 3 depends on this constant $C$, which in turn depends on a constant $K$ defined in Appendix B. This constant is used to upper bound quantities such as $\\sup_\\theta E_x\\|\\nabla M(x, \\theta)\\|^2$ which show up when concentrating the martingale term. However, it appears to me that this constant is not actually dimension independent. In the proof of Lemma 22, before line 558, this constant depends on $|x|^2$ which is typically $\\Theta(d)$ for many examples (such as Gaussian or uniform on sphere of radius $\\sqrt{d}$). Thus it seems that there is an additional dimension dependence hidden in the constant $C$. Can you please clarify whether there is indeed a hidden dimension dependence here? This is a key issue because it changes the exponent of $d$ in Theorem 3. In my opinion this is a major issue that makes me recommend reject, but if you can clarify this point or I have misunderstood I am wiling to increase my score.\n- A more minor weakness is that I find the examples in Section 4.2 to be a bit restrictive. For instance, requiring *all* 2d marginals of $\\nu$ to be close to Gaussian in Wasserstein distance seems to be a very restrictive assumption, and I am not sure what distributions satisfy this property (beyond uniform on the sphere of radius $\\sqrt{d}$). Next, proposition 15 shows that LPG is satisfied only for $b = \\log d/\\sqrt{d}$, which seems to imply that polynomially many initializations would need to be sampled before a good one is reached. Finally, as the authors indeed point out, the Stein’s method argument only shows the LPG along dense directions and does not yield an end-to-end guarantee for online SGD.\n\nQUESTIONS:\n- As mentioned above, can you please comment on the dimension dependence of the constant $C$ in Theorem 3?\n- A key requirement in the paper is for the initialization to have sufficient correlation with $\\theta^*$. How does one check whether an initialization has sufficient correlation? This is particularly of interest for the 2D-Wasserstein distance example, where the probability of achieving a good initialization is polynomially small.\n- Do you have examples of specific distributions which satisfy Assumption 14?\n- lines 283-284: “by paying an additional $O(d^2)$ cost in complexity” — what do you mean by this? is this just the complexity of resampling until a good initialization is achieved?\n\nLIMITATIONS:\nLimitations/broader impact are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the results of Ben Arous et al., 2021 for learning a single-index model with a known link function via online SGD beyond the Gaussian input setting. The authors introduce a condition called Local Polynomial Growth (LPG) and establish the convergence rate of online SGD under LPG for spherically symmetric input distributions and certain perturbations from the isotropic Gaussian distribution.\n\nSTRENGTHS:\n* The paper has an excellent presentation with sufficient intuition given for different regimes of SGD convergence.\n* Learning single-index models for distributions without spherical symmetry is an interesting and significant open problem.\n* The calculations performed for spherically symmetric distributions can serve as a useful example of how to perform a harmonic analysis for such distributions in the context of learning with online SGD.\n\nWEAKNESSES:\n* My main concern is that LPG might not be the right definition to analyze this problem with non-Gaussian distributions, due to the following reasons:\n    * For spherically symmetric distributions, I'm not sure if LPG is required. Specifically, I believe the *information exponent* framework of Ben Arous et al., 2021, which only depends on spherical symmetry and Taylor expansion of the population loss, must be able to handle such distributions, and the contributions in this section might be interpreted as calculating the information exponent for non-Gaussian examples.\n\n    * For distributions without spherical symmetry, it appears that LPG is not the suitable definition. Particularly, as the authors have pointed out, the validity of Assumption 14 for many distributions of interest is not clear, and the probability that SGD has a successful initialization is of order $O(1/d^2)$, making it inapplicable in high-dimensional settings. Moreover, there appears to be no direct connection between the Stein's-method-based argument and LPG.\n\nQUESTIONS:\n* I believe it would be helpful to have a discussion on the connection between the calculations of Section 3 and the framework of Ben Arous et al., 2021. In particular, is showing LPG strictly required for establishing SGD convergence for spherically symmetric distributions, or is it possible to use the Taylor expansion of Ben Arous et al., 2021 to arrive at the same results?\n\n* In line 161, it is mentioned that \"This probability can be lowered by any constant factor ...\". Is this not the success probability? If that is the case, is this sentence meant to say the probability can be increased?\n\n* In the same sentence, how can one measure the alignment of different initializations with $\\theta^*$ in order to maximize such alignment, without having knowledge of $\\theta^*$?\n\nLIMITATIONS:\nThe authors have mostly adequately discussed the limitations of the work, and additional details are provided above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper deals with the study of online SGD in single-index models of the type $f(\\mathbf x)=\\phi(\\boldsymbol\\theta^\\intercal\\mathbf x)$ for a generic non-linear scalar link function $\\phi$. The work builds on previous results of Ben Arous et al and aims at going beyond the Gaussian-data assumptions therein. A theorem extending the Gaussian result is given where a crucial condition on the population loss, namely the local polynomial growth, is identified. Finally, the theoretical results are applied to the spherically symmetric setting and to specific non-spherical symmetric settings.\n\nSTRENGTHS:\nThe paper is an interesting attempt to go beyond the Gaussian data assumption in the study of online SGD in single-index models and provides a criterion for preserving the scenario observed in the Gaussian setting, characterizing therefore its robustness. If on one hand, the theoretical analysis seems to become extremely complex going beyond the Gaussian realm, it is nevertheless essential to better understand the SGD dynamics on real data. The paper is quite clear although the analysis in Section 4 is more cumbersome (but I have the feeling that this is unavoidable).\n\nWEAKNESSES:\nI have the impression that a weakness of the paper lies in the fact that the main information on the loss landscape in the general setting is concentrated on a property, the local polynomial growth condition, which is in general not straightforwardly analyzable. If, on one hand, the typical case discussed in Section 4.1 exemplifies the application of the theorem, on the other hand, it also seems to show that verifying the LPG property is not trivial and requires a setting in which a proper basis (in the case, the Gegenbauer polynomials) is available. The study in Section 4.2, mostly instead relies on a Gaussian reference measure, which, as the authors themselves stress, is also a limitation. I wonder therefore if the theorem is *practically* useful beyond a small set of adequate measures that actually satisfy the LPG property (in this sense the title sounds partially misleading a\n\nQUESTIONS:\nAs discussed in the *Weaknesses* section, it appears that verifying the LPG property requires a case-by-case, very nontrivial analysis of the distribution involved. Could the author comment on this point? Do they expect that a general (even heuristic) strategy is feasible to expect to be within the hypothesis of Theorem 3 in the case of a general measure? Finally, with respect to future developments, do the authors expect the LPG condition to be necessary for Theorem 3 to hold?\n\nAs a final comment, the authors might consider citing in their work the very classical contribution\n> David Saad and Sara A. Solla, Phys. Rev. E 52, 4225 (1995)\n\nwhere the dynamic of $m$ was studied within the framework of statistical mechanics.\n\nTypos and other comments.\n* The authors refer to the goal function as a *Sparse high-dimensional function* but there is no sparsity assumption on, e.g., $\\theta$ and, if referring to the fact that the goal is to learn a single direction, the terminology appears slightly misleading.\n* In line 70, $\\theta^*$ in the definition of $\\mathcal H$ should be $\\theta$.\n* In line 76 $\\|f\\|^p_{L_\\mu^p}$ should be $\\|f\\|^p_{L_\\nu^p}$.\n* In lines 248-249, two subsequent \"Indeed\" appear.\n* In line 260 \"Wassertein\" in place of \"Wasserstein\".\n\nLIMITATIONS:\nSee the *Weaknesses* section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main contributions of the authors are as follows:\n\n\t(a) Rather than assuming the features are i.i.d. Gaussian, the authors identify a condition called Local Polynomial Growth (LPG, see Def. 2) on the population loss, under which the analysis of Ben Arous et. al. still goes through without any change in their conclusions (see Theorem 3 and Theorem 4 in the paper). Roughly speaking this condition posits a lower bound on the directional derivative of the population loss along the hidden direction at a given point $\\theta$ in terms of the overlap between $\\theta$ and $\\theta_\\star$. \n\n\t(b) In Proposition 7, the authors then verify that the LPG condition holds when the features are sampled from a possibly non-Gaussian, rotationally invariant distribution (for e.g., the uniform distribution on the sphere). \n\n\t(c) In Proposition 15, the authors show that if the 2D projected Wasserstein distance (see eq. 18) between the feature distribution and the standard Gaussian scales like $1/sqrt(d)$ ($d$ is the feature dimension), then for population losses whose Gaussian information exponent is at most $2$ satisfy the LPG condition. \n\n\t(d) In Proposition 17, the authors make some partial progress in verifying a relaxed form of the LPG condition for i.i.d. non-Gaussian feature vectors. The basic idea here is that since the LPG condition only depends on a two-dimensional projection of the feature vectors, because of the central limit theorem, the LPG condition should hold. The authors conjecture that SGD path avoids localized vectors, and based on this conjecture, sketch an informal argument that would show that the conclusions of Ben Arous et. al. should also extend to feature vectors with i.i.d. coordinates.\n\nSTRENGTHS:\nI found the paper to be well-written and the proofs (as far as I checked them) seem correct to me. The SGD analysis of Ben Arous et. al. does not assume that the features are i.i.d. Gaussian. specifically, but assume a condition on the population loss (see Eq. 1.1 in the ``Online stochastic gradient descent on non-convex losses from high-dimensional inference'', JMLR 2021),  and analyze SGD dynamics under this condition. However, in most of their examples, the condition can only be verified under the i.i.d. Gaussian feature assumption. The condition proposed by the authors is weaker than the condition in the work of Ben Arous et. al. and can be verified in examples beyond those considered in the work of Ben Arous et. al. This is a key strength of this paper.\n\nWEAKNESSES:\nThe key weakness of this paper seems to be that most of the examples discussed by the authors seem to be small perturbations of Gaussian sensing vectors. \n\nAt the same time, the abstract and the introduction seem to overclaim the generality of the results obtained in the paper. For instance, in the abstract, the authors write ``our main results establish that Stochastic Gradient Descent can efficiently recover the unknown direction in the high dimensional regime, under mild assumptions that significantly extend Yehudai and Ohad (2020) and Wu (2022)''. However, the papers cited by the authors appear to make weaker assumptions on the feature vectors (compared to rotationally invariant features or i.i.d. features), but stronger assumptions on the non-linear link functions, so it appears that the results of the authors are not stronger or weaker - rather complementary. \n\nA second weakness of the paper is that the analysis of the authors for i.i.d. feature vectors seems incomplete. It does seem quite plausible that the conclusions of Ben Arous et al. should extend to this case. However, the claims of the authors are conditional on a conjecture (see Conjecture 35 in the supplement). Furthermore, the fact that their results are conditional on this unproven conjecture is not transparently indicated in the main paper, where the authors simply write ``While the LPG property does not directly apply in this setting, in Appendix B.4 we outline an energy-barrier argument that demonstrates that the single-index model 324 can still be efficiently solved using gradient-based methods'', without mentioning that the result relies on a conjecture they haven't proven.\n\nQUESTIONS:\nIt would be great if the authors could clarify the following:\n\n\t1) In Theorem 4, is the scaling of the iteration bound with the desired recovery accuracy $\\epsilon$ expected to be optimal? This seems important since it would clarify if SGD achieves the optimal statistical rate of convergence. In addition, the result stated in Theorem 4 seems inconsistent with the result obtained in the proof (Page 16 in the supplement). \n\n\t2) In Prop 7, the authors consider the setting where the features are rotationally invariant. In this situation, I wonder if it might be possible to recover this result from the results of Ben Arous et. al. since it seems to me that the main assumption of Ben Arous et. al. holds in this case (that is, the population loss at a point $\\theta$ is only a function of the dot product between $\\theta$ and the hidden direction). \n\n\t3) A relevant reference to cite is the work ``One-bit compressed sensing with non-Gaussian measurements'' by Albert Ai, Alex Lapanowski, Yaniv Plan, and, Roman Vershynin. This work also uses a central limit theorem-type argument to go beyond i.i.d. Gaussian measurements to i.i.d. non-Gaussian measurements, as in Section 4.2 of this paper.  \n\nI have the following minor comments:\n\n\t(1) I think the discussion of the work of Ben Arous et. al. can be improved to clarify that their analysis of SGD doesn't use Gaussianity per se, but is under a general condition, which can be verified for i.i.d. Gaussian features and rotationally invariant features. \n\n\t(2) In the introduction, lines 55-57, I didn't quite understand what the authors mean by \"stability\". \n\n\t(3) The discussion in lines 162-164 seems a bit inappropriate: the authors propose drawing many random initializations and keeping the one that has the maximum correlation with the hidden direction - but the hidden direction is unknown. \n\n\t(4) Below Eq: 18: Has the notation $Gr(2,d)$ been defined?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studied the online SGD of learning a single-index teacher model with the link function as the student model. This work generalized the results of [Ben Arous et al. 2021], where the information exponent of the link function governs the sample complexity, to the non-Gaussian dataset, where the author raised a local polynomial growth (LPG) to describe the loss landscape.  This LPG prevents the loss landscape to have bad local minima outside a band of the equator. The author provided two cases of loss landscape with LPG: (i) in the spherically symmetric case, similar behavior as the Gaussian case is exhibited; (ii) in the non-symmetric case, more assumptions on link function and feature direction $\\theta^*$ are needed and the deviation between the non-Gaussian and Gaussian distribution in certain metric characterizes the LPG property.\n\nSTRENGTHS:\nOverall, the presentation of this paper is clear, and the motivation of the paper is strong. It is important to generalize online SGD with non-Gaussian data distributions and extend the information exponent of the link function in this general setting. The results obtained by the authors are as sharp as can be expected given the problem under consideration. Although the model studied is simple, the insights derived are non-trivial and the technical approach is original. Especially, in general cases, this paper introduces a two-dim Wasserstein distance to measure the deviation from the Gaussian reference and Stein method for product measures, which are interesting directions for future work.\n\nWEAKNESSES:\n1. Simulations like [Ben Arous et al. 2021] did should be included in the paper. Especially certain non-Gaussian distributions either spherically symmetric or not, should be presented in the simulations to see how information exponent and LPG control the weak recovery and training dynamic.\n2. More examples of non-Gaussian distribution could be mentioned in Section 4. Like Example 8 and 9, there should be some examples for non-spherically symmetric cases in Section 4.2.\n3. Unlike [Ben Arous et al. 2021], there is no lower bound showing the smallest sample complexity for efficient online SGD using LPG in non-Gaussian cases. In the Gaussian case, $n\\ge d^{s-1}$ is nearly tight for online SGD where $s$ is the information exponent.\n\nQUESTIONS:\n1. [Ba et al. 2022] is another paper addressing a semi-parametric problem when the link function is unknown with certain GD training. [Arnaboldi et al. 2023] also studied the SGD training of two-layer neural networks in high-dimensional settings.\n2. Can we consider noisy labels in (1)? \n3. Line 106: information exponent $s$ should be explicitly defined here.\n4. [Tan and Vershynin 2019] considered uniform sphere data distribution, although the link function is only quadratic in this phase retrieval case. The SGD dynamics for the non-Gaussian phase retrieval model may be analyzed in the literature.\n5. Line 154 double \"that\"\n6. Between lines 189 and 190, what is distribution $\\rho$ when we define measure $\\nu$?\n7. In Proposition 7, is there some relation between $s^*$ information exponent $s$ as $d\\to\\infty$? There is a formula describing the relation between Hermite coefficients and spherical harmonic coefficients. Hence, asymptotically, will the results of Example 8 return to the Gaussian case or not?\n8. How about the problem of model misspecification, where the link function of the target is not known?\n\n===================================================================================================\n- Ba, J., Erdogdu, M.A., Suzuki, T., Wang, Z., Wu, D. and Yang, G., 2022. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, pp.37932-37946.\n- Arnaboldi, L., Stephan, L., Krzakala, F. and Loureiro, B., 2023. From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks. arXiv preprint arXiv:2302.05882.\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:\nI think the authors have adequately addressed the potential negative social impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}