Learning Single-Index Models with Shallow Neural Networks

Introduction

High-dimensional learning with both computational and statistical guarantees, which is particularly relevant given the current scaling trends, remains an outstanding challenge. One important question which has received considerable attention is on understanding the advantages of using non-linear learning models, such as neural networks, over more mature (from a theoretical standpoint) counterparts, such as kernel methods [MKAS21, MYSSS21, WGL+20, CB20]. Perhaps surprisingly, the question remains largely open even for shallow neural networks.

While approximation benefits of shallow neural networks over non-adaptive kernels have been known for decades [Bar93, Pin99], another important piece of the theoretical puzzle was provided by [Bac17a], whose analysis hinted at an inherent statistical advantage of neural networks for extracting information from high-dimensional data with a “hidden” low-dimensional structure. Providing computational guarantees, the remaining piece of this puzzle, is still mostly unresolved.

Several computational hardness results for learning functions that can be efficiently approximated by shallow neural networks have been established in the literature [DKKZ20, GGJ+20, DV20, SZB21, CGKM22], ruling out positive results in the general setting. On the other hand, progress has been made on the positive side [ABAM22, AZL20, SL21] by focusing on function classes with strong structural properties, thereby showcasing the adaptive representation learning capabilities of neural networks.

The mix of a high-dimensional parametric component (the hidden direction) with a non-parametric one in low dimension (the link function) in single-index models naturally suggests a shallow neural network architecture where the inner weights are shared and “active”, while the biases are “lazy” [COB18]. We instantiate such an architecture by freezing the biases at random initialization, and analyze gradient descent on the free parameters in the continuous-time limit.

Our main results establish that as soon as the width $N$ of the network is larger than a quantity which depends solely on smoothness properties of the (univariate) link function $f_{*}$, gradient flow recovers the unknown direction $\theta^{*}$ with near optimal sample complexity $O(d^{s})$, where $s$ is the so-called information-exponent of the link function [BAGJ21] (at least when $s\geq 3$, see Theorem 6.1 for the formal result), and approximates the univariate link function $f_{*}$ near-optimally (see Corollary 6.4). The information exponent roughly captures the signal strength, which here refers to the alignment between the network direction $\theta$ and the hidden direction $\theta^{*}$, at typical initializations.

The success of gradient flow relies on the benign optimization landscape of the empirical loss, though the presence of degenerate saddles necessitates a careful analysis leveraging uniform convergence of the empirical landscape [MBM16, FSS18]. We show that gradient flow over our proposed neural network architecture solves two distinct problems—univariate non-parametric kernel ridge regression and non-convex optimization in high dimension—simultaneously and efficiently, cementing its role as a versatile algorithm for high-dimensional learning. We illustrate our theoretical results with experiments in Section 7.

Related Work

The works [BL19, CBL+20, NBL22] show that certain neural networks trained close to initialization can learn certain sparse polynomials which take the form of multi-index models, but such networks do not directly aim to learn target directions. Recently, [ABAM22] studied the learnability of functions on the hypercube by shallow neural networks with stochastic gradient descent and introduces the merged staircase property, which provides necessary and sufficient conditions for learnability with linear sample complexity $n=O(d)$. While they learn a broader class of functions (including multi-index model) for a more efficient sample complexity regime ($O(d)$ vs $O(d^{s})$), their setup is restricted to simple discrete data distributions, while our work captures the regime of semi-parametric estimation by considering Gaussian data without the sparsity requirements on $F$ implied by their merged staircase property.

Concurrently to our work, [BES+22] and [DLS22] studied the learnability of certain single and multi-index models on Gaussian data with shallow networks, by performing a single gradient step on the first layer before fitting the second layer. While the single step is sufficient to provide a separation from kernel methods in these works, we show that optimizing both layers jointly until convergence (for a more simplistic architecture) can significantly improve the rates, by fully decoupling the non-parametric learning part from the high-dimensional inference of the hidden direction. Finally, recently [MHPG+22] studied the ability of shallow neural networks to learn certain single and multi-index models, showing in particular that SGD-trained ReLU networks can learn single-index functions with monotonic index function (corresponding in our setting to $s=1$) with linear (up to logarithmic factors) sample complexity. Our results therefore extend such positive guarantees to a broader class of index functions with arbitrary information exponent.

Teacher–student models.

In the context of neural networks, several works have considered the teacher–student setting [EVdB01], where the target function $F$ takes the form of a neural network with the same activation as the network used for learning [GAS+19, GLM17, ZSJ+17, Sol17, ZGJ21, BAGJ21, VSL+22]. In this case the problem does not involve non-parametric estimation as in our setup, but this line of work often involves studying optimization landscapes similar to ours for estimating hidden directions. In particular, the population landscape appearing in [BAGJ21] is similar to ours, based on Hermite coefficients of link functions. The follow-up work [BAGJ22] extends this to multiple student neuron directions, but still focuses on parametric rather than non-parametric statistical problems.

Kernels and random features.

In order to obtain non-parametric estimation guarantees for learning the target function $f_{*}$ of the single-index model, our work builds on the kernel methods literature for approximation and non-parametric regression [SS02, BTA11, Bac21], their links with neural networks [CS09, Bac17a], and in particular on random feature approximation [RR07, Bac17b, RR17, MMM22].

Non-convex and non-smooth optimization landscapes.

There is a vast literature studying tractable non-convex optimization landscapes, arising from high-dimensional statistics and statistical physics [MKUZ19, MBC+20, GHJY15, BAMMN19, SQW18, BVB16, GLM17, RABC19, MAB19]. A particular aspect of our setup is that the optimization landscape does not have the strict saddle property, which is often leveraged to establish global convergence [JGN+17]. [MBM16, FSS18] study concentration properties of the empirical landscape to the population one for non-convex problems including generalized linear models. Our results rely on similar concentration analyses, but depart from these previous work by also allowing optimization of the link function, and by supporting the non-smoothness arising from ReLU activations. On the algorithmic side, we consider gradient flows on non-convex and non-smooth landscapes, which require careful technical treatment, but have been studied by previous works [DIL15, DDKL20, JT20]. We refer the interested reader to Appendix E for more details on this technical issue.

Preliminaries

We focus on regression problems under a single-index model with Gaussian input data. Specifically, we assume $d$-dimensional inputs $x\sim\gamma_{d}:=\mathcal{N}(0,I_{d})$, and labels

We consider learning algorithms based on shallow neural networks of the form

The choice of ReLU activation is motivated by its popularity among practitioners. As we shall see, the fact that $\phi$ is non-smooth introduces some technical challenges, but its piece-line structure enables dedicated arguments both in terms of approximation as well as in the study of the optimization landscape. In Appendix F we discuss how our main results are affected when replacing the ReLU by a smooth activation, especially when choosing it such that $\phi^{\prime}$ is Lipschitz.

Empirical risk minimization.

The supervised learning task is to estimate $F^{*}$ (and therefore both $f_{*}$ and $\theta^{*}$) from samples $\{(x_{i},y_{i})\}_{i=1\ldots n}$. We will focus on mean-squared error with Tychonov regularisation, determined by the following losses.

Hermite decomposition.

We apply the following useful properties of Hermite polynomials [O’D14, Chapter 11.2]:

where $\langle\cdot,\cdot\rangle_{\gamma_{d}}$ is the inner product in $L^{2}(\gamma_{d})$ and $\delta$ the Kronecker delta. We will assume throughout that $\|f_{*}\|_{\gamma}^{2}=\sum_{j}\alpha_{j}^{2}$, $\|f_{*}^{\prime}\|^{2}_{\gamma}=\sum_{j}j\alpha_{j}^{2}$, and $\|f_{*}^{\prime\prime}\|^{2}_{\gamma}=\sum_{j}j(j-1)\alpha_{j}^{2}$ are all finite (see Assumption 5.2). We will also consider the weighted Sobolev space $H^{2}(\gamma)$, which contains functions $f=\sum_{j}\alpha_{j}h_{j}\in L^{2}(\gamma)$ such that $\sum_{j}j^{2}|\alpha_{j}|^{2}<\infty$.

Random features to Hermite coefficients.

Note that ${\mathcal{T}}$ has rank $N$ almost surely.

Denoting $m=\langle\theta,\theta^{*}\rangle$, the regularized population objective can be expressed as

where the term ${\sum_{j}\alpha_{j}^{2}}$ is a constant that can be ignored.

Geometry on the sphere.

Because the direction $\theta$ is constrained to lie on the sphere, our optimization algorithms rely on spherical (Riemannian) gradients, which are defined as follows:

Univariate Approximation using Random Features

Before addressing the high-dimensionality of the learning problem, we first focus on the non-parametric approximation aspects of the univariate link function. As usual, we start by deriving approximation rates of the infinitely-wide model, given by a RKHS (Section 4.1), and then establish approximation rates for our random feature model (Section 4.2).

If we fix the direction $\theta$, learning $c$ alone may be seen as a random feature model [RR07] that approximates a kernel method with the following kernel.

The following lemma characterizes the corresponding RKHS norm $\|\cdot\|_{\mathcal{H}}$, and follows from Theorem A.8, by noting that $\kappa(u,v)=\langle\psi(u),\psi(v)\rangle_{L^{2}(\gamma_{\tau})^{2}}$, with $\psi(u)=\frac{1}{\sqrt{2}}(\phi(u-\cdot),\phi(-u-\cdot))$.

The RKHS norm in $\mathcal{H}$ is given by

The choice of ReLU for the activation function gives us more explicit control over the RKHS norm, based on Sobolev representations, as already exploited by several works [OWSS19, Bac17a, SESS19].

Let $A(f,\lambda)$ be the (regularized) $L^{2}$ approximation error for functions in the space $\mathcal{H}$ with respect to the target function $f$ and measure $\gamma$. Formally,

We will now show that the approximation error of the RKHS corresponding to an infinite number of random features can be bounded in terms of the regularization $\lambda$ and the $\gamma$-norm of the second derivative of the target function. For that purpose, we consider the following ‘source’ condition to ensure a polynomial approximation error in $\lambda$.

Let $\mathcal{F}=\{f\in H^{2}(\gamma)\;|\;f^{\prime\prime}\in L^{4}(\gamma)\}$. We assume $f\in\mathcal{F}$ and define $K:=\inf\left\{B\geq 1\;|\;\|f^{\prime\prime}\|_{L^{4}(\gamma)}\leq B\|f^{\prime\prime}\|_{L^{2}(\gamma)}\right\}$.

Assumption 4.3 provides a sufficient condition for approximating $f$ with functions in the RKHS. The family of approximants $\{h_{M}\in\mathcal{H}\;|\;M>0\}$ we use in Lemma 4.4 are exactly equal to $f$ on $[-M,M]$ and are linear outside of $[-M,M]$. The $L^{4}$ assumption on $f^{\prime\prime}$ ensures control over the RKHS norm of $h_{M}$. Note that by Jensen’s inequality, $L^{4}(\gamma)\subset L^{2}(\gamma)$, so $K$ is always well-defined for $f^{\prime\prime}\in L^{4}(\gamma)$. Sigmoidal functions, compactly supported smooth functions, and, more generally, functions with polynomial growth satisfy Assumption 4.3.

Let $\lambda\in(0,1)$ and $f\in\mathcal{F}$. Then, there exists a universal constant $C>0$ such that

where $\beta=\frac{1-1/\tau^{2}}{3+1/\tau^{2}}$ and $C_{f}=\max\{\|f\|_{\gamma},\|f^{\prime}\|_{\gamma},\|f^{\prime\prime}\|_{\gamma}\}$.

The proof appears in Appendix B.2. This lemma allows us to control the RKHS approximation error of a target function in terms of their Hermite decompositions. The main technical difficulty is that the RKHS integral operator $\Sigma$ does not diagonalise in the Hermite basis; we address this with a dedicated argument exploiting the RKHS Sobolev representation of Lemma 4.2. The assumption that $f^{\prime\prime}\in L^{4}(\gamma)$ (Assumption 4.3) is sufficient for our purposes but not necessary for polynomial-in-$\lambda$ approximation rates. In Section G, we show that the ReLU function $\phi(t)=\max(0,t)$, which is Lipschitz but not in $H^{2}(\gamma)$, as the target satisfies $A(\phi,\lambda)\lesssim\tau^{2}\lambda^{2/3}$ using a direct argument. Extending the class of functions approximable by $\mathcal{H}$ with polynomial-in-$\lambda$ rate is an interesting future direction.

2 Controlling random feature approximation

We now consider (finite) random feature approximations to functions in the RKHS. We show that the best possible loss of a linear combination of sufficiently many finite features is bounded above by the best approximation with infinitely many features with high probability.

The claim follows from Lemmas B.2 and B.3 in Appendix B.3. ∎

Population Landscape under Frozen Random Biases

One of the main measures of complexity for the target link function $f_{*}$ is its information exponent (see e.g., [BAGJ21]), defined as follows.

We make the following regularity assumptions on the target link function $f_{*}$ to ensure small approximation error by random features, and benign population and empirical landscape.

We consider $f_{*}\in L^{2}(\gamma)$, with $f_{*}=\sum_{j}\alpha_{j}h_{j}$. Assume

$f_{*}^{\prime\prime}(z):=\sum_{j}\sqrt{(j+2)(j+1)}\alpha_{j+2}h_{j}(z)$ is in $L^{4}(\gamma)$ (Assumption 4.3).

We also suppose w.l.o.g. that $f_{*}$ is normalized so that $\|f_{*}\|_{\gamma}=1$. To analyze the critical points of $L(c,\theta)$, we introduce the projected population loss $\bar{L}(\theta)$,

Assume $f_{*}$ satisfies Assumption 5.2 and has information exponent $s\geq 1$. For $\tau>1$, and $\delta\in(0,1)$, there exists $\lambda^{*}\leq 1$ depending only on $\tau$ and the target link function $f_{*}$ and a universal constant $C>0$ such that if

(orientation relative to $\theta^{*}$) if $(c,\theta)\in\Omega$, then either $\theta\in\{\pm\theta^{*}\}$ or $\langle\theta,\theta^{*}\rangle=0$.

Theorem 5.3 thus establishes a benign optimization landscape in the population limit, rejoining several known non-convex objectives with similar behavior, such as tensor decomposition [GM17] or matrix completion [GLM16]. Importantly, this optimization landscape has the same topology as the one that arises from using the Hermite basis, the tailored choice for data generated by a single-index model in Gaussian space [DH18, Theorem 5], instead of random scalar features. We view this as an interesting robustness property of shallow neural networks, at least in the regime where biases are randomly frozen.

We first derive critical point equations for the regularized population loss.

Then, denoting $m=\langle\theta^{*},\theta\rangle$, we have

Furthermore, the critical points of $L(c,\theta)$ satisfy the following equations:

We prove Claim 5.4 in Appendix C.1 by analyzing the population gradient and relating the critical points of $\bar{L}(\theta)$ to those of $L(c,\theta)$. Note that the function $g_{m}$ corresponds to the minimizer of $\left\|f_{*}(\langle\theta^{*},\cdot\rangle)-g(\langle\theta,\cdot\rangle)\right\|_{\gamma_{d}}$, which is essentially the optimal function we may learn from fitting the second layer $c$ with no regularization when $\theta$ is fixed.

2 Proof of Theorem 5.3

where $m=\langle\theta^{*},\theta\rangle$. $\bar{L}(\theta)$ in the ideal case is strictly decreasing in $|m|\in(0,1]$. Using the expression for $\bar{L}(\theta)$ in Eq. (12), we observe that by setting $\lambda>0$ sufficiently small (and $N$ proportional to $1/\lambda$), the projection $\hat{P}_{\lambda}$ onto the subspace $\hat{\mathcal{H}}$ spanned by random features approximates the identity map in the operator norm, thereby preserving the strict monotonicity of $\bar{L}(\theta)$ with respect to $|m|$. We formalize this intuition in the following proof.

We first establish the following lemma (proved in Section C.2) which controls the approximation error for the functions $g_{m}$ defined in (11).

where $\beta=\frac{1-1/\tau^{2}}{3+1/\tau^{2}}$ and $C_{g_{m}}=\max\{\|g_{m}\|_{\gamma},\|g_{m}^{\prime}\|_{\gamma},\|g_{m}^{\prime\prime}\|_{\gamma}\}$.

If $\lambda<\lambda^{*}$ and $N\geq\frac{C}{\lambda}\log\frac{1}{\lambda\delta}$, then with probability greater than $1-\delta$ it holds

which contradicts (14). Therefore, the existence of critical points satisfying $m\notin\{-1,0,1\}$ is ruled out with probability at least $1-\delta$ over the random features. ∎

Empirical Landscape and Generalization Guarantees

Section 5 shows that the population landscape has a relatively simple structure given $N=\Omega_{d}(1)$ random features. We now study the optimization properties of its finite-sample counterpart.

Prior works have obtained sample complexity of $n={O}(d^{s})$, where we recall that $s$ is the information exponent of the target function $f_{*}$, for recovering $\theta^{*}$ either by employing a learning algorithm that explicitly learns individual Hermite polynomials [DH18] or by assuming that $f_{*}$ is known a priori [BAGJ21].Actually, in [BAGJ21] the authors obtain a slightly improved sample complexity of $\widetilde{O}(d^{s-1})$ for $s\geq 3$ by directly analyzing SGD with fixed step-size, as well as a matching lower bound (for SGD in the small step-size regime) up to polylogarithmic factors. The intuition behind this sample complexity is roughly as follows.

The empirical optimization landscape (when regarded as a function only of the direction $\theta$) near the equator ($|m|\ll 1$) is of the form $L(\theta)\asymp m^{s}$.

In order to certify that the optimization algorithm does not converge to a suboptimal critical point (i.e., $\|\nabla L(\theta)\|\leq\epsilon$) on the equator, one requires that $m\geq\epsilon^{1/(s-1)}$.

A uniform gradient convergence bound of the form $\|\nabla L(\theta)-\nabla L_{n}(\theta)\|=O(\sqrt{d/n})$ and the fact that $m=\Theta(1/\sqrt{d})$ at initialization together imply that $n=O(d^{s})$ samples are sufficient to escape from the “influence” of the equator.

In order to repurpose these arguments to our setting, the relative scaling of the top-layer weights $c$ relative to the direction vector $\theta$ is crucial, as has also been observed in the literature on lazy-vs-rich regimes [COB18, WGL+20] in the context of overparametrized neural networks.

We consider an idealized version of Gradient Descent over the empirical loss $L_{n}$ in the infinitesimally small learning rate regime. This results in a gradient flow ODE of the form:

Our main result, proved in Appendix D, establishes that this gradient flow efficiently finds an approximate minimizer of the population loss, with an error (explicitly quantified as a function of $n$) that reveals the fundamental role of the information exponent $s$ of $f_{*}$. On top of the regularity conditions on the target link function of Assumption 5.2, the upper bound on $\lambda$, and the lower bound on $N$ from Eq. (10) of Theorem 5.3, the main result imposes a (compatible) upper bound on $N$ and an appropriate choice of initial norm ($\rho$) and sparsity ($N_{0}/N$) for $c(0)$. In this section, we are interested in behavior as $n$, $d$, and $N$ grow asymptotically and hence treat the target function $f_{*}$ and terms derived from it (including Hermite coefficients $\alpha_{j}$ and information exponent $s$), along with the bias parameter $\tau$ and regularity parameter $\beta$, as constants and omit them from asymptotic notation.

We note that the dependence on $d+N$ in the recovery guarantee can likely be improved to $d$ using a more refined norm-based landscape concentration analysis. We also remark that if we chose the number of random features $N=\Theta(\frac{1}{\lambda}\log\frac{1}{\lambda\delta})$, then the requirement $\lambda=\Omega(\sqrt{\Delta_{\text{crit}}})$ for large enough sample size $n$ imposes $\lambda\gg n^{-1/5}$. This lower bound on $\lambda$ guarantees that critical points near initialization can be escaped, but may slow down learning. Nevertheless, this is sufficient to obtain an excess risk that vanishes with $n$, with a rate independent of $d$, as we now show.

Under the assumptions of Theorem 6.1, and further assuming $n\gtrsim d^{3}$, the choice of $\lambda\sim n^{-1/(\beta+5)}$ yields an excess risk guarantee of the form

where $\beta$ is defined as in Lemma 4.4.

Since $N=\widetilde{\Omega}(\lambda^{-1})$, this result indicates that one needs to increase the width of the network to $N=\widetilde{\Omega}(n^{1/(\beta+5)})\gg d^{s/(1+\beta)}$ in order to achieve vanishing excess error. With the joint training we are thus paying a dependency in the ambient dimension in terms of width, yet this width increase is driven by the non-parametric approximation error of $f_{*}$, which is itself independent of $d$. A simple mechanism to break this inefficiency is by considering a fine-tuning step of the second-layer terms.

After running Algorithm 1, we may include a final fine-tuning phase of training for second layer weights $c$ alone, using a separate training sample $(x_{i}^{\prime},y_{i}^{\prime})$, $i=1,\ldots,n^{\prime}$ and a possibly different regularization parameter $\lambda_{n^{\prime}}$. More precisely, we set

where $\hat{\theta}$ denotes the output of the previous gradient descent phase. Note that this is a strongly convex optimization problem, and can thus be optimized efficiently using gradient methods or by solving a linear system. While this may not be needed in practice, we use a different training sample for technical reasons, namely to break the dependence between the data and the kernel, which depends on the initial training sample through $\hat{\theta}$. We note that such sample splitting strategies are commonly used in other contexts in the statistics literature (e.g., [BBSMW21, CCD+18]). We obtain the following guarantee.

Let $\delta\in(0,1/4)$. Let $m=\langle\theta^{*},\hat{\theta}\rangle$, where $\hat{\theta}$ is obtained from the previous gradient descent phase, and let $\hat{c}$ be the ridge regression estimator obtained from a fresh dataset $\mathcal{D}^{\prime}$ of $n^{\prime}$ samples, $N$ random features,These can be the same as in the previous phase. and regularization parameter $\lambda_{n^{\prime}}:=(\sigma^{2}\tau^{2}/\|f_{*}^{\prime\prime}\|_{\gamma}^{2}n^{\prime})^{1/(\beta+1)}$, and let $\hat{F}(x)=\hat{c}^{\top}\Phi(\langle\hat{\theta},x\rangle)$. Assume

Then with probability at least $1-\delta$ over the random features, we have

where the expectation is over the $n^{\prime}$ fresh samples, and is conditioned on the previously obtained $\hat{\theta}$.

Decoupling the regularization parameters of the two phases (along with number of random features $N$) allows us to keep a large $\lambda$ in the first phase, leading to fast recovery as per Theorem 6.1, while obtaining vanishing excess risk through a decreasing $\lambda_{n^{\prime}}$. This is illustrated in the result on the excess risk for Algorithm 2.

Let $\delta\in(0,1/4)$. As in Theorem 6.1, let $\mu_{s}=\langle h_{s},\Sigma h_{s}\rangle>0$, and let $f_{*}$ satisfy Assumption 5.2 on $f_{*}$ with a constant $C_{f_{*}}>0$. Let $\lambda=(s^{2}\alpha_{s}^{2}/2C_{f_{*}})^{2/\beta}$, and assume the following on the sample sizes and number of random features for the first phase ($n,N,N_{0}$) and fine-tuning phase ($n^{\prime},N^{\prime}$):

and let $\rho$ be as in Theorem 6.1. With probability at least $1/2-2\delta$ over the initial $n$ samples, initialization, random features, we have

where the constants in $\widetilde{O}$ do not depend on $d$ other than through logarithmic factors.

By comparing Corollaries 6.2 and 6.4, we observe that the fine-tuning stage recovers the optimal sample complexity, where the non-parametric rate no longer depends on the ambient dimension $d$.

We can make the following additional remarks:

The time-scale separation schedule for $\zeta$ in Theorem 6.1 is sufficient but possibly not necessary. The analysis of vanilla dynamics ($\zeta(t)\equiv 1$) is challenging, since during the initial phase of training there may be adverse interaction effects between $c$ and $\theta$, which under naive analysis lead to sub-optimal sample complexity of $n\geq O(d^{2s})$. Observe that this separate analysis of ‘weak’ and ‘strong’ recovery phases of learning appears in most contemporary related work [DLS22, ABAM22, BES+22, BAGJ21, BAGJ22].

The time discretization to turn Procedure 1 into a proper algorithm should follow from standard time discretization arguments, although the case where $\phi=\text{ReLU}$ requires special care due to the non-smoothness of the loss (see Appendix E for further discussion). In such setting, such discretization arguments do not hold for vanilla gradient descent in the worst-case [KS21], although these may be recovered by appropriately smoothing the objective prior to computing the gradient, or by using instead a smooth activation function (see Appendix F).

Numerical Experiments

Our experiment results are shown in Figure 2. For $s\geq 3$, only some of the 10 runs were successful in recovering the target direction, and we thus show the best performing run for such curves (indeed, our theory suggests that there may be a non-negligible probability of failure). We observe that full recovery ($|m|\to 1$) requires more samples when the dimension $d$ increases, while the excess risk curves have approximately the same rate for large enough $n$, regardless of the dimension or information exponent, as predicted by our theory. The bottom plots for $d=50$ suggest that $s=3$ requires more samples than smaller $s$ for perfect recovery, while the remaining curves are somewhat comparable. This similarity between $s=1$ and $s=2$ is reminiscent of the situation in [BAGJ21], where the rates for these two cases only differ by a logarithmic factor, and suggests that it may be possible to improve the $O(d^{s})$ rates in our results for $s\geq 2$.

Conclusion

This work studies the ability of shallow neural networks to learn single-index models with gradient descent. Our main results are positive, and demonstrate their ability to solve a semi-parametric problem with nearly optimal guarantees. Interestingly, this success story combines elements from the feature-learning regime, i.e., the ability to efficiently identify the hidden direction in high-dimensions under a non-convex objective, with ingredients from the lazy-regime, which offer better computational tradeoffs to solve low-dimensional non-parametric problems. Our technical analysis leverages tools from high-dimensional probability (such as uniform gradient concentration) and RKHS approximation theory, and complements the growing body of theoretical work certifying the efficiency of gradient methods on non-convex objectives. We have followed the standard approach of first establishing benign topological properties of the population loss, and then extending them to the empirical loss.

There are nonetheless several unanswered questions that our work has not addressed. Below, we provide a list of potentially interesting future directions.

Our approximation rate for ReLU as the target (see Appendix G) suggests that the polynomial-in-$\lambda$ approximation rate may be extended to function classes beyond $\mathcal{F}\subset H^{2}(\gamma)$, such as Lipschitz functions with smooth tail behavior. Thus, it would be interesting to extend our empirical landscape concentration results to such functions satisfying weaker regularity assumptions, which currently rely on certain polynomial decay of the Hermite coefficients (see Assumption 5.2). Additionally, by using a smooth activation function (see Appendix F), our Gradient Flow dynamics can be discretized and turned into Gradient Descent (GD) with analogous sample and time complexity. In that context, a natural goal is to compare quantitatively the differences between GD with multiple passes over the training data and SGD by adapting tools from [BAGJ21, BAGJ22].

Trainable biases and untied directions.

Our proposed neural network architecture is non-standard, in the sense that its biases are frozen at initialization and all neurons share the same inner weight. For the purposes of learning single-index models, removing these restrictions would not bring any statistical benefits. However, it would be interesting to extend our analysis to the general setting where the first layer weights are not tied and biases are not frozen.

Extension to Multi-index Models.

Multi-index models are natural extensions of single-index models where the hidden direction $\theta^{*}$ is replaced by a hidden low-dimensional subspace. Typically, multi-index models enjoy similar statistical guarantees as single-index models [DH18, Bac17a], and thus a natural question is whether the same algorithmic tools developed here extend to the multi-index setting.

Gradient dynamics without warm-start.

An unsatisfactory aspect of our results is the requirement that the algorithm starts by only optimizing $\theta$ for $t<T_{0}$. It would be interesting to understand whether the vanilla dynamics can also succeed provably.

Acknowledgements.

We are thankful to Enric Boix-Adserà, Alex Damian, Cédric Gerbelot, Daniel Hsu, Jason Lee, Theodor Misiakiewicz, Matus Telgarsky, Eric Vanden-Eijnden, and Denny Wu for useful discussions. We also thank the anonymous NeurIPS reviewers and area chair for helpful feedback. JB, AB and MJ are partially supported by NSF RI-1816753, NSF CAREER CIF 1845360, NSF CHS-1901091, NSF Scale MoDL DMS 2134216, Capital One and Samsung Electronics. CS is supported by an NSF GRFP and by NSF grants CCF-1814873 and IIS-1838154.

References

Appendix A Additional Preliminaries and Concentration Bounds

We introduce several well-known concentration bounds that we apply throughout the appendix. Borrowing notation from [Ver18], we first introduce notation of sub-gaussian and sub-exponential random variables, vectors, and matrices.

We note several key properties of sub-gaussian and sub-exponential random variables that we repeatedly rely on.

Let $z_{1},\dots,z_{N}$ and $y_{1},\dots,y_{N}$ be sub-gaussian and sub-exponential random variables respectively. Then the following hold for some universal constant $C$.

Products of sub-gaussian random variables are sub-exponential: $z_{1}z_{2}$ is sub-exponential and $\left\|z_{1}z_{2}\right\|_{\psi_{1}}\leq\left\|z_{1}\right\|_{\psi_{2}}\left\|z_{2}\right\|_{\psi_{2}}$ [Ver18, Lemma 2.7.7].

Sums of independent sub-gaussian random variables are sub-gaussian. If $z_{1},\dots,z_{N}$ are independent, then,

Sums of pairs of sub-exponential random variables are sub-exponential. $\left\|y_{1}+y_{2}\right\|_{\psi_{1}}\leq C(\left\|y_{1}\right\|_{\psi_{1}}+\left\|y_{2}\right\|_{\psi_{1}})$ [MBM16, Lemma 2].

For independent, mean-zero, sub-exponential random variables $x_{1},\dots,x_{n}$ and any $t\geq 0$,

for universal $C$ and $K=\max_{i}\left\|x_{i}\right\|_{\psi_{1}}$.

We also include several basic facts about $\epsilon$-covers, which are useful in several proofs.

We show this using an elementary argument. Let $\epsilon\in$. Define

By Gautschi’s inequality [DLMF, Eq. 5.6.4] for the Gamma function, we have

Finally, we recall the following result on reproducing kernel Hilbert spaces, which describes the RKHS for kernels defined from explicit features maps.

Let $\psi:\mathcal{X}\to\mathcal{F}$ be a mapping into a Hilbert space $\mathcal{F}$, and for $x,x^{\prime}\in\mathcal{X}$, define the kernel $\kappa(x,x^{\prime})=\langle\psi(x),\psi(x^{\prime})\rangle_{\mathcal{F}}$. The RKHS $\mathcal{H}$ of $\kappa$ consists of functions of the form $f(x)=\langle g,\psi(x)\rangle_{\mathcal{F}}$, and for any $f\in\mathcal{H}$, the RKHS norm of $f$ is defined by

Appendix B Proofs of Section 4

Moreover, the RKHS norm of $f$ is upper bounded as follows.

By the Fundamental Theorem of Calculus and Fubini’s Theorem,

The upper bound on the RKHS norm follows from the above representation and Lemma 4.1. ∎

For general $f$, the boundary conditions of Claim B.1, i.e., $\lim_{t\rightarrow-\infty}f(t)=\lim_{t\rightarrow-\infty}f^{\prime}(t)=0$, do not hold. However, we can reduce to the case considered in Claim B.1 by decomposing $f$ into 2 parts, i.e., $f=f_{1}+f_{2}$, where $f_{1}(t)$ and $f_{2}(-t)$ individually satisfy the assumptions of Claim B.1.

Let $\varphi(t)=\int_{-\infty}^{t}\gamma(u)du$. We decompose

To upper bound the RHS of Eq. (23) in terms of $f$ and its derivatives, we derive explicit expressions for $f_{1}^{\prime\prime}$ and $f_{2}^{\prime\prime}$. Since $\varphi^{\prime}(t)=\gamma(t)$ and $\varphi^{\prime\prime}(t)=-t\gamma(t)$, we have

The same upper bound holds for $f_{2}$. Thus, from Eq. (23) and the fact that $\tau>1$, we have

Hence, by the assumption that $f,f^{\prime}$ have polynomial growth and applying the identity Eq. (24) twice, we have

B.2 Proof of Lemma 4.4

and recall that $A(f,\lambda)=\min_{h\in\mathcal{H}}\mathcal{E}(h)$.

Thus, $h_{M}$ matches $f$ exactly on $[-M,M]$ and is linear with slope $f^{\prime}(M)$ (resp. $f^{\prime}(-M)$) for $t\geq M$ (resp. $t\leq-M$). We first show that $h_{M}\in\mathcal{H}$, which implies $A(f,\lambda)\leq\inf_{M>0}\mathcal{E}(h_{M})$, and then show that for an explicit choice of $M$, $\mathcal{E}(h_{M})$ has the desired upper bound.

For any finite $M>0$, $h_{M}$ satisfies the assumptions of Lemma 4.2; both $f$ and $f^{\prime}$ have polynomial growth (linear and zero growth, respectively) and

where $r_{M}=f-h_{M}$ and we used the triangle inequality and $2ab\leq a^{2}+b^{2}$ in Eq. (26). Note that since $\tau>1$, the first term of Eq. (26) is upper bounded by $\lesssim\tau^{2}\|f^{\prime\prime}\|_{\gamma}^{2}e^{(\tau^{2}+1)M^{2}/(2\tau^{2})}$.

We now upper bound $\|r_{M}\|_{\gamma}$ and $\|r_{M}^{\prime}\|_{\gamma}$. Note that both $r_{M}$ and its derivative $r_{M}^{\prime}$ are identically zero on $[-M,M]$ and that $r^{\prime\prime}_{M}(t)=f^{\prime\prime}(t)$ for $|t|>M$. Thus, for $t>M$ (same holds for $t>-M)$,

Next, we decompose $\|r_{M}^{\prime}\|_{\gamma}^{2}$ into two terms, the positive part $\|r_{M}^{\prime}\|_{\gamma,+}^{2}$ and the negative part $\|r_{M}^{\prime}\|_{\gamma,-}^{2}$. That is,

An upper bound on $\|r_{M}\|_{\gamma,+}^{2}$ follows from similar calculations.

where we used the fact that $\tau>\max\{1,\lambda\tau\}$ in Eq. (28).

It remains to balance the terms in Eq. (28) by choosing an appropriate value for $M>0$. We choose $M$ by balancing $\tau\|f^{\prime\prime}\|_{4}^{2}\cdot e^{-\frac{M^{2}}{4}}$ and $\lambda\tau^{2}\|f^{\prime\prime}\|_{\gamma}^{2}\cdot e^{\frac{1+\tau^{2}}{2\tau^{2}}\cdot M^{2}}$.

Let $\beta=1-\frac{1+\tau^{2}}{2\tau^{2}}\cdot\frac{4\tau^{2}}{1+3\tau^{2}}=\frac{\tau^{2}-1}{1+3\tau^{2}}=\frac{1-1/\tau^{2}}{3+1/\tau^{2}}$. Plugging the above value of $M$ into Eq. (28),

where $C_{f}=\max\{\|f\|_{\gamma},\|f^{\prime}\|_{\gamma},\|f^{\prime\prime}\|_{\gamma}\}$.

B.3 Proof of Lemma 4.5

Before stating and proving the two supporting lemmas, we introduce several terms are used to study the similarity of a finite random feature model to its infinite counterpart. Let $\hat{\kappa}(u,v)$ be the random empirical kernel associated with $N$ random features:

where $\phi^{\varepsilon}_{b}(u):=\phi(\varepsilon u-b)$, $b_{i},\varepsilon_{i}$ are i.i.d. random variables drawn from $\gamma_{\tau}\otimes\text{Rad}$. Its associated integral operator in $L_{2}(\gamma)$ is given by $\hat{\Sigma}$.

In the following, we consider $\Sigma$ the integral operator corresponding to the kernel $\kappa$:

By a technical lemma adapted from [Bac17b], the approximation error of the random feature model is controlled via the regularization parameter $\lambda>0$.

There exists a constant $C>0$ such that if $N\geq Cd_{\max}(\lambda)\log(d_{\max}(\lambda)/\delta)$, we have, with probability at least $1-\delta$, for any $f\in L^{2}(\gamma)$,

as long as $N\geq Cd_{\max}(\lambda)\log(d_{\max}(\lambda)/\delta)$. Now note that we have

where the last equality follows from [Bac21, Lemma 7.2]. Thus, we have proved the result for $f\in\mathcal{H}$. Given that (9) does not require $f$ to be in $\mathcal{H}$, we may conclude by limiting arguments that the result holds for any $f$ in the closure of $\mathcal{H}$, which includes $L^{2}(\gamma)$, since the kernel is universal, given that its associated RKHS is a weighted Sobolev Space, which is dense in $L^{2}(\gamma)$. ∎

We have $d_{\max}(\lambda)\leq C/\lambda$, for an absolute constant $C>0$.

If $b\geq 0$, then we have $|\phi_{b}^{\varepsilon}(u)|\leq|u|$ for all $u$, thus

with $C=2\int_{0}^{\infty}u^{2}\gamma(u)du$.

with $g(u)=\epsilon u-b$ a linear function, using the relation $\max(0,u)=u+\max(0,-u)$. Then we have

Appendix C Proofs for Section 5

To characterize the critical points of $L(c,\theta)$ for fixed random features and prove Claim 5.4, we derive exact expressions for $L$ and its gradients. We observe that the population loss depends on the student direction $\theta$ only via its angle to the teacher direction $\theta^{*}$.

Recall the decomposition $c^{\top}\Phi(z)=\sum_{j=0}^{\infty}\langle\mathcal{T}_{j},c\rangle h_{j}(z)$. Straightforward calculation gives

Recall that the criticality of $\theta$ depends on the spherical gradient being zero, not the standard one. Since the gradient $\nabla_{\theta}L$ is colinear with $\theta^{*}$, we stipulate necessary and sufficent conditions for $\theta$ to be critical.

Recall that $g_{m}(z)=\sum_{j=s}^{\infty}\alpha_{j}m^{j}h_{j}(z)$. Then, the projected population loss $\bar{L}$ is given by

Furthermore, critical points of $\bar{L}$ satisfy the following equation.

Now we plug in $c=Q_{\lambda}^{-1}(\sum_{j}\alpha_{j}m^{j}\mathcal{T}_{j})$ into Eq. (32). Then, we have

Differentiating $\bar{L}(\theta)$ with respect to $m=\langle\theta,\theta^{*}\rangle$, we obtain the following critical point equation.

As discussed above $(c,\theta)$ is a critical point of $L$ if and only if $\theta$ is a critical point of $\bar{L}$ and $c=Q_{\lambda}^{-1}(\sum_{j}\alpha_{j}m^{j}\mathcal{T}_{j})$. By applying Lemma C.4, we separate the diagonal and off-diagonal terms to rewrite Eq. (35) as

Dividing both sides by $2$ gives the claim. ∎

C.2 Proof of Lemma 5.5

For simplicity, we denote the $L^{p}(\gamma)$ norms by $\|\cdot\|_{p}$. For any $\rho\in$, we define the noise operator $U_{\rho}$ by

This is a reparametrisation of the Ornstein-Uhlenbeck semigroup, The Ornstein–Uhlenbeck semigroup $\mathcal{P}_{t}$ is given by $\mathcal{P}_{t}f(x)=\int f(e^{-t}x+\sqrt{1-e^{-2t}}z)d\gamma(z)~{}.$ We thus have $\mathcal{P}_{t}=U_{e^{-t}}$ and we have from [O’D14, Prop 11.33] that $U_{\rho}h_{j}=\rho^{j}h_{j}$. In other words, the Hermite polynomials are eigenfunctions of the semigroup. As a consequence, we have from (11) that $g_{m}=U_{m}f_{*}$.

By Nelson’s Gaussian hypercontractivity [Nel73] (reproduced in [O’D14, Theorem 11.23]),

Let us now consider $m\geq 1/\sqrt{3}$. Since $U_{\rho}$ is an averaging operator for all $\rho\leq 1$, from Jensen’s inequality (reproduced in [O’D14, Proposition 11.15]) it holds that $\|U_{\rho}f\|_{p}\leq\|f\|_{p}$ for any $p\geq 1$. We thus have

where $\beta=\frac{1-1/\tau^{2}}{3+1/\tau^{2}}$, $C>0$ is a universal constant, and $C_{g_{m}}=\max\{\|g_{m}\|_{\gamma},\|g_{m}^{\prime}\|_{\gamma},\|g_{m}^{\prime\prime}\|_{\gamma}\}$.

C.3 Other lemmas for the proof of Theorem 5.3

Let $f\in H^{2}(\gamma)$ be such that $f^{\prime},f^{\prime\prime}\in L^{2}(\gamma)$ and let $s\geq 1$ be its information exponent. Furthermore, let $f_{*}=\sum_{j}\alpha_{j}h_{j}$ be the Hermite expansion of $f_{*}$, and let $g_{m}$ and $\bar{g}_{m}$ be defined as in Theorem 5.3. Then,

By definition of $g_{m}$ and Holder’s inequality,

Let $f\in H^{2}(\gamma)$ be a function satisfying assumptions of Lemma C.5, let $g_{m}$ and $\bar{g}_{m}$ be defined as in Theorem 5.3, and let $C_{g_{m}}=\max\{\|g_{m}\|,\|g_{m}^{\prime}\|_{\gamma},\|g_{m}^{\prime\prime}\|_{\gamma}\}$. Then,

The first inequality follows from the following.

The proof of the second inequality is via straightforward algebraic manipulation.

where we used the inequalities $ab\leq a^{2}+b^{2}$, and $\sqrt{a^{2}+b^{2}}\leq a+b$ which apply to any $a,b\geq 0$. ∎

Appendix D Proofs for Section 6

The proof of this theorem has two separate parts: we first prove that our gradient flow procedure escapes the neighborhood of the equator, and then show that it converges to a neighborhood of the north pole. Define the set of approximate-first-order critical points of the empirical landscape in the sublevel set $L_{n}(c,\theta)\leq\nu$.

with $c^{*}=\hat{P}_{\lambda}f_{*}$. One would expect that for $n$ sufficiently large, these topological properties should be transferred to the empirical landscape. This intuition is indeed correct, and relies on the following uniform convergence result, proved in Appendix D.2.

Equipped with this uniform gradient concentration, we can first establish the analogous classification of first-order critical points for the empirical landscape (proof in Section D.3):

We consider the gradient flow procedure of Algorithm 1, that we restate here for convenience:

We establish the following fact, proved in Appendix D.4:

D.2 Proof of Lemma D.1

Let $C_{1}=\max\{C,1\}$, where $C>0$ is the constant from Fact D.4. Then, with probability at least $1-\delta$ over the random features,

The tail of this random vector is subexponential, as stated in the following lemma.

Define $W=c^{\top}\Phi(\langle\theta,x\rangle)-f_{*}(\langle\theta^{*},x\rangle)-\xi$. Using Fact A.3,

By Corollary D.5, the following holds with probability at least $1-\delta$ over the random features.

By the union bound (over $j\in[N]$) and basic properties of Gaussian random variables,

Finally, we bound the remaining $n-q$ terms.

As a result, with probability at least $1-\delta/3$,

where we set $q=4d\log(1/(\epsilon_{\theta}\delta))$ and recall that $\epsilon_{c}=r/n^{2}$, $\epsilon_{\theta}=1/(4n^{2}N)$.

where we used the assumption $d\leq n$ in the last inequality.

We now bound third term. Note that this term involves only populational quantities, so discontinuity of the sample gradients is not an issue here.

We upper bound the two terms individually. By Eq. (42), the second term is bounded as follows.

The following observation gives an upper bound on the Lipschitz constant of $g^{\prime}(m;c)$, which depends only on $\|c\|$ and is thus constant if $r$ is fixed.

Lemma D.7 gives a simple expression for $\sum_{j}j^{2}(j-1)^{2}\alpha_{j}^{2}$. Its proof can be found in Section D.8.

Let $C_{f_{*}}=\max\{\|f^{(1)}\|_{\gamma},\ldots,\|f^{(4)}\|_{\gamma},1\}$, which is well-defined thanks to our regularity assumption on the target link function (Assumption 5.2). Then,

Putting everything together, we have that with probability at least $1-\delta$,

The proof is similar to the one for $\nabla_{\theta}L_{n}$.

There exists a universal constant $C_{0}>0$ such that under the same assumptions as Lemma D.6, with probability at least $1-\delta$ over the random features,

We bound the first and third terms by bounding the discretization error for each samplewise gradient. First observe that

We bound the first term on the RHS as follows.

Taking $\epsilon_{c}=r/n^{2}$ and $\epsilon_{\theta}=1/n^{2}$, we observe that with probability $1-(2/3)\delta$,

D.3 Proof of Lemma D.2

By definition, $L_{n}$ can be expressed as

where $C_{\text{data}}>0$ is some constant independent of $c$ and $\theta$, and

Using the fact that $\nabla_{c}L_{n}(c,\theta)=\hat{Q}_{\lambda}(\theta)c-\hat{Y}(\theta),\|\nabla_{c}L_{n}(c,\theta)\|\leq\epsilon$, and $\hat{Q}_{\lambda}^{-1}(\theta)\preceq\lambda^{-1}I$ uniformly in $\theta$, we obtain

where $x\sim\mathcal{N}(0,I_{d})$ and $\xi\sim\mathcal{N}(0,\sigma^{2})$.

where we used the fact that $\tau\sqrt{\log(1/\delta)}>1$ in the last line.

By Bernstein’s inequality (Theorem A.4) and the union bound over $\mathcal{N}_{1/2}$, the following holds with probability at least $1-\delta$.

where we used the assumption $d\leq n$ for the last inequality.

where $C_{f_{*}}=\max\{\|f^{(1)}\|_{\gamma},\ldots,\|f^{(4)}\|_{\gamma}\}$.

Since we assumed $N\geq\frac{C}{\lambda}\log\frac{1}{\lambda\delta}$, the conditions of Lemma D.11 are satisfied. Hence,

As a result, by Lemma D.9, which applies since $\lambda<\lambda^{*}$ and $N\geq\frac{C}{\lambda}\log\frac{1}{\lambda\delta}$, either one of the following must be true.

We use the representation of the restricted population loss $\bar{L}(\theta)$ from Lemma C.4 Eq. (34), the notation $\rho_{m}=2\langle\hat{P}_{\lambda}g_{m},\bar{g}_{m}\rangle$, and the definition of the Riemannian gradient to obtain

If $|m|\leq 1/2$, then $|\rho_{m}|\leq\sqrt{4/3}\epsilon$. Hence, our lower bound on $\left|\rho_{m}\right|$ implies that

where $C_{f_{*}}=\max\{\|f_{*}^{\prime}\|_{\gamma},\|f_{*}^{\prime\prime}\|_{\gamma}\}$.

Define $c_{\theta}=\arg\min_{c}L(c,\theta)$, so that $\bar{L}(\theta)=L(c_{\theta},\theta)$. We first show that if $\left\|\nabla_{c}L(c,\theta)\right\|\leq\epsilon_{c}$, then $c$ and $c_{\theta}$ are nearby. Because $\nabla_{c}L(c_{\theta},\theta)=0$ and $Q_{\lambda}\succeq\lambda I_{N}$,

Thus, $\left\|c-c_{\theta}\right\|\leq\frac{\epsilon_{c}}{2\lambda}$. We now recall the Riemannian gradient for $\theta$,

and use it to bound the norm of the projected gradient.

We conclude by employing Lemmas C.5 and D.12 to obtain a bound on the final term that holds with probability at least $1-\delta$.

We conclude by selecting a sufficiently large $N$. ∎

By Fact A.3, it suffices to bound $\left\|c^{\top}\Phi(\langle x,\theta\rangle)-f_{*}(\langle x,\theta_{*}\rangle)-\xi\right\|_{\psi_{2}}$. Note that this quantity identically equals $\left\|W\right\|_{\psi_{2}}$ for the random variable $W$ defined in the proof of Lemma D.6. Thus, with probability at least $1-\delta$,

We bound the first and last terms by considering the discretization error of samplewise loss.

We bound the first factor, relying on the event of Corollary D.5 with probability at least $1-(\delta/6)$.

We use the same event to bound the second factor.

Hence, by taking $\epsilon_{c}=r\epsilon_{\theta}=r/n$, we have

By applying Fact D.4 on all $\xi_{i}$ and the fact that $\left\|x_{i}\right\|^{2}\lesssim d\log(n/\delta)$ for all $i$ with overwhelming probability, we conclude that with probability at least $1-\delta/3$

Likewise, bounds on the expectations of $\left|\xi\right|$ and $\left\|x\right\|$ similarly give

We conclude by bounding the second term using Bernstein’s inequality with the sub-exponential norm bound of Lemma D.14. Recall that $D=\max\{d,N\}$. Then, for sufficiently large $C_{0}$ (and thus sufficiently large $C_{1}$),

D.4 Proof of Lemma D.3

Recall our gradient flow dynamics in the first phase:

Thanks to the concentration results from Lemma D.13 and Lemma D.1, we can first compute the correlation trajectory $m(t)$ for the population loss, and then extend them to the empirical gradients.

Assume that $m(0)$ and $c(0)$ are such that that $\text{sign}(\alpha_{s}c(0)^{\top}\mathcal{T}_{s})=\text{sign}(m(0))$, which occurs with probability $1/2$ over the randomness of $c(0)$ and $m(0)$. By symmetry, we will assume $m(0)>0$ and $\alpha_{s}c(0)^{\top}\mathcal{T}_{s}>0$ for the rest of the proof. Let us express the population objective without offset as

From Lemma A.7, we know that the correlation $m(0)$ at initialization cannot be too small. More precisely,

Moreover, the change in correlation according to the population gradient is given by

The following lemma, proved below, shows there exists $\gamma=\gamma(c(0))>0$ and $\bar{\gamma}$ such that $R(m)>R(0)/2$ for $m\in[0,\gamma)$ and $\bar{R}(m)>\bar{R}(0)/2$ for $m\in[0,\bar{\gamma})$.

Let $C_{f*,\tau}=2\tau\left(\sum_{j>0}(j+s)^{2}j^{2}\alpha_{j+s}^{2}\right)^{1/2}$ and $\bar{C}_{f^{*},\tau}=2\tau\left(\sum_{j>0}(j)^{2}\alpha_{j+s}^{2}\right)^{1/2}$. Then

$R(m)>\frac{1}{2}s\alpha_{s}c^{\top}\mathcal{T}_{s}$ for $m\in[0,\gamma)$, where

$\bar{R}(m)>\frac{1}{2}\alpha_{s}c^{\top}\mathcal{T}_{s}$ for $m\in[0,\bar{\gamma})$, where

In other words, the gradient flow under the population loss sees a monotonically increasing correlation $m$ (since its time derivative under the population gradient flow is positive), until $m(t)$ reaches a value $\gamma=C\frac{\alpha_{s}c^{\top}\mathcal{T}_{s}}{\rho}$.

Let $\gamma_{*}=\min(\gamma,\bar{\gamma})$ and $\rho_{0}=\rho\sqrt{\frac{N_{0}}{N}}$. As the correlation reaches the value $m=\gamma_{*}$, using Lemma D.15 to lower bound $\bar{R}$, one can verify that the population loss obeys the following upper bound:

where, denoting by $\mathcal{S}$ the support of $c$, we defined

Let us now verify that the empirical correlation trajectory and loss have the same behavior. Observe that

and therefore from Lemma D.15 we deduce that $\dot{m}(0)>0$, and $m(t)$ keeps increasing at least until it reaches $\gamma_{*}$. From Lemma D.13, the empirical loss at this correlation level is with probability greater than $1-\delta$

In order to ensure that this initial training phase escapes the ‘bad’ empirical points near the equator $|m|\approx 0$, by Eq. (38), it is sufficient to show that

with $\Delta_{\text{crit}}:=\max\left\{\sqrt{\frac{D}{n}},\left(\frac{d^{2}}{n}\right)^{\frac{2s}{2s-1}}\right\}$.

where the probability is over both the initial draw of $c$ and the draw of the random features.

Taking $\rho_{0}=\frac{C^{\prime}N_{0}^{-\frac{1+s}{2}}}{2(2\tau^{2}+\lambda N/N_{0})}$ yields the new condition

In particular, since we assume $\lambda^{-2}\Delta_{\text{crit}}\ll 1$, we may take $N_{0}=\Theta(1)$ and

Finally, let us upper bound the escape time $T$ needed to reach $m(T)=\gamma_{*}$. Denote $\Delta_{n}:=\lambda^{-2}\max\left\{\sqrt{\frac{D}{n}},\frac{d^{2}}{n}\right\}$. Observe that for $t\leq T$,

which leads to a Gronwall-type inequality of the form

Thus, if we assume that $m\leq\gamma$ as specified in the theorem statement,

The derivation for $\bar{R}(m)$ is analogous. ∎

Observe that $Z_{s}=\psi(Z)$ with $\psi(x)=\langle h_{s},\sigma(\cdot-x)\rangle_{\gamma}$ satisfying

which shows that $Z_{s}$ is $\frac{1}{2\tau^{2}}$-subgaussian, and thus that the random vector

Finally, using again the anticoncentration of the correlation of a uniform direction with a fixed direction (Lemma A.7), we obtain with a union bound that

D.5 Proof of Corollary 6.2

We restate Corollary 6.2 here for convenience.

Let $\hat{F}(x)=\hat{f}(\langle x,\hat{\theta}\rangle)$, and $G_{m,\hat{\theta}}(x)=g_{m}(\langle\hat{\theta},x\rangle)=\sum_{j}\alpha_{j}m^{j}h_{j}(\langle\hat{\theta},x\rangle)$, where $m=\langle\hat{\theta},\theta^{*}\rangle$. We have

Denoting $c_{\theta}=Q_{\lambda}^{-1}\mathcal{T}g_{m}$, and considering $N=\Theta(\frac{1}{\lambda}\log\frac{1}{\lambda})$, recall that we have

using that $\lambda=o(1)$ when $n$ grows (as confirmed later), and hiding the dimension $d$ in the $O(\cdot)$ notation. Thus, we obtain

As a consequence, using Lemmas 4.5 and 4.4 we obtain

where the $\widetilde{O}$ follows from Lemma D.2. We thus obtain

Optimizing the second and third term over $\lambda$ yields $\lambda\sim n^{-1/(\beta+5)}$. Note that with this choice, the upper bound on $N$ required by Theorem 6.1 is of order $\lambda\Delta_{\text{crit}}^{-1}=\widetilde{\Theta}\left(\sqrt{\lambda^{3}n}\right)=\widetilde{\Theta}\left(n^{\frac{1}{\beta+5}+\frac{\beta}{2(\beta+5)}}\right)\gg N=\widetilde{\Theta}\left(n^{\frac{1}{\beta+5}}\right)$, so that the condition is satisfied. With this choice of $\lambda$, the first term is negligible compared to the other terms, which are of order $n^{-\frac{\beta}{\beta+5}}$. Overall this leads to a final rate

D.6 Proof of Proposition 6.3

Then, we have for $n^{\prime}\geq\frac{R^{2}}{2\lambda}\log\frac{R^{2}}{\lambda}$ and $\lambda\leq R^{2}$, following [Bac21, Proposition 7.1]

where the expectation is over the $n^{\prime}$ fresh samples.

We have the following upper bound on the first (variance) term

The approximation error may be controlled as follows:

where we assume $N\geq Cd_{\max}(\tau,\lambda)\log(d_{\max}(\tau,\lambda)/\delta)$ in order to apply Lemma B.2.

By limiting arguments, we may show that this holds for any $F^{*}$ in the closure of $\mathcal{H}_{\theta}$. Now consider the true $F^{*}(x)=f_{*}(\langle\theta^{*},x\rangle)$. When $\theta^{*}\neq\theta$, $F^{*}$ does not belong to the closure of $\mathcal{H}_{\theta}$, but we may consider the projection $F^{*}_{\mathcal{H}_{\theta}}$ on this closure. Then, following the arguments of [Bac21, Section 7.6.4], we obtain

We may take $F^{*}_{\mathcal{H}_{\theta}}(x)=g(\langle\theta,x\rangle)$ for some $g$, since all functions in $\mathcal{H}_{\theta}$ and its closure take this form. Then, we may consider $g$ of the form $g=\sum_{j}b_{j}h_{j}$, since such functions are dense in the closure of $\mathcal{H}_{\theta}$. Optimizing the approximation error $\|F^{*}_{\mathcal{H}_{\theta}}-F^{*}\|_{\gamma_{d}}$ over such $g$ yields $b_{j}=\alpha_{j}m^{j}$, so that the approximation error becomes