Limitations of Lazy Training of Two-layers Neural Networks

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, Andrea Montanari

Introduction

The function class of two-layers neural networks (with NN neurons) is defined by:

These facts lead to the following central question in neural network theory:

Here ‘efficiently’ can be formalized in multiple ways: in this paper we will focus on learning via stochastic gradient descent.

Significant amount of work has been devoted to two subclasses of FNN,N{\mathcal{F}}_{{\sf NN},N} which we will refer to as the random feature model (RF{\sf RF}) [RR08], and the neural tangent model (NT{\sf NT}) [JGH18]:

We can think of RF{\sf RF} and NT{\sf NT} as tractable inner bounds of the class of neural networks NN{\sf NN}:

Tractable. Both FRF,N(W){\mathcal{F}}_{{\sf RF},N}({\boldsymbol{W}}), FNT,N(W){\mathcal{F}}_{{\sf NT},N}({\boldsymbol{W}}) are finite-dimensional linear spaces, and minimizing the empirical risk over these classes can be performed efficiently.

Inner bounds. Indeed FRF,N(W)⊆FNN,N{\mathcal{F}}_{{\sf RF},N}({\boldsymbol{W}})\subseteq{\mathcal{F}}_{{\sf NN},N}: the random feature model is simply obtained by fixing all the first layer weights. Further FNT(W)⊆cl(FNN,2N){\mathcal{F}}_{{\sf NT}}({\boldsymbol{W}})\subseteq{\rm cl}({\mathcal{F}}_{{\sf NN},2N}) (the closure of the class of neural networks with 2N2N neurons). This follows from ε−1[σ(⟨wi+εai,x⟩)−σ(⟨wi,x⟩)]=⟨ai,x⟩σ′(⟨wi,x⟩)+o(1){\varepsilon}^{-1}[\sigma(\langle{\boldsymbol{w}}_{i}+{\varepsilon}{\boldsymbol{a}}_{i},{\boldsymbol{x}}\rangle)-\sigma(\langle{\boldsymbol{w}}_{i},{\boldsymbol{x}}\rangle)]=\langle{\boldsymbol{a}}_{i},{\boldsymbol{x}}\rangle\sigma^{\prime}(\langle{\boldsymbol{w}}_{i},{\boldsymbol{x}}\rangle)+o(1) as ε→0{\varepsilon}\to 0.

It is possible to show that the class of neural networks NN{\sf NN} is significantly more expressive than the two linearization RF{\sf RF}, NT{\sf NT}, see e.g. [YS19, GMMM19]. In particular, [GMMM19] shows that, if the feature vectors xi{\boldsymbol{x}}_{i} are uniformly random over the dd-dimensional sphere, and N,dN,d are large with N=O(d)N=O(d), then FRF,N(W){\mathcal{F}}_{{\sf RF},N}({\boldsymbol{W}}) can only capture linear functions, while FNT,N(W){\mathcal{F}}_{{\sf NT},N}({\boldsymbol{W}}) can only capture quadratic functions.

In this paper we explore systematically the gap between RF{\sf RF}, NT{\sf NT} and NN{\sf NN}, by considering two specific data distributions:

Quadratic functions: feature vectors are distributed according to xi∼N(0,Id){\boldsymbol{x}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}) and responses are quadratic functions yi=f∗(xi)≡b0+⟨xi,Bxi⟩y_{i}=f_{*}({\boldsymbol{x}}_{i})\equiv b_{0}+\langle{\boldsymbol{x}}_{i},{\boldsymbol{B}}{\boldsymbol{x}}_{i}\rangle with B⪰0{\boldsymbol{B}}\succeq 0.

Mixture of Gaussians: yi=±1y_{i}=\pm 1 with equal probability 1/21/2, and xi∣yi=+1∼N(0,Σ(1)){\boldsymbol{x}}_{i}|y_{i}=+1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(1)}), xi∣yi=−1∼N(0,Σ(2)){\boldsymbol{x}}_{i}|y_{i}=-1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(2)}).

Let us emphasize that the choice of quadratic functions in model qf is not arbitrary: in a sense, it is the most favorable case for NT{\sf NT} training. Indeed [GMMM19] proves thatNote that [GMMM19] considers feature vectors xi{\boldsymbol{x}}_{i} uniformly random over the sphere rather than Gaussian. However, the results of [GMMM19] can be generalized, with certain modifications, to the Gaussian case. Roughly speaking, for Gaussian features, NT{\sf NT} with N=O(d)N=O(d) neurons can represent quadratic functions, and a low-dimensional subspace of higher order polynomials. (when N=O(d)N=O(d)): (i)(i) Third- and higher-order polynomials cannot be approximated nontrivially by FNT,N(W){\mathcal{F}}_{{\sf NT},N}({\boldsymbol{W}}); (ii)(ii) Linear functions are already well approximated within FRF,N(W){\mathcal{F}}_{{\sf RF},N}({\boldsymbol{W}}).

For clarity, we will first summarize our result for the model qf, and then discuss generalizations to mg. The prediction risk achieved within any of the regimes RF{\sf RF}, NT{\sf NT}, NN{\sf NN} is defined by

Our results are summarized by Figure 1, which compares the risk achieved by the three approaches above in the population limit n→∞n\to\infty, using quadratic activations σ(u)=u2+c0\sigma(u)=u^{2}+c_{0}. We consider the large-network, high-dimensional regime N,d→∞N,d\to\infty, with N/d→ρ∈(0,∞)N/d\to\rho\in(0,\infty). Figure 1 reports the risk achieved by various approaches in numerical simulations, and compares them with our theoretical predictions for each of three regimes RF{\sf RF}, NT{\sf NT}, and NN{\sf NN}, which are detailed in the next sections.

The agreement between analytical predictions and simulations is excellent but, more importantly, a clear picture emerges. We can highlight a few phenomena that are illustrated in this figure:

Random features do not capture quadratic functions. The random features risk RRF,N(f∗)R_{{\sf RF},N}(f_{*}) remains generally bounded away from zero for all values of ρ=N/d\rho=N/d. It is further highly dependent on the distribution of the weight vectors wi∼N(0,Γ){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}). Section 2.1 characterizes explicitly this dependence, for general activation functions σ\sigma. For large ρ=N/d\rho=N/d, the optimal distribution of the weight vectors uses covariance Γ∗∝B{\boldsymbol{\Gamma}}^{*}\propto{\boldsymbol{B}}, but even in this case the risk is bounded away from zero unless ρ→∞\rho\to\infty.

Fully trained neural networks achieve vanishing risk on quadratic functions for N>dN>d: this is to be expected on the basis of the previous point. For N/d→ρ∈(0,1)N/d\to\rho\in(0,1) the risk is generally bounded away from , but its value is smaller than for the neural tangent model. Namely, in Section 2.3 we give an explicit expression for the asymptotic risk (holding for B⪰0{\boldsymbol{B}}\succeq{\boldsymbol{0}}) implying that, for some GAP(ρ)>0{\rm GAP}(\rho)>0 (independent of N,dN,d),

The picture emerging from these findings is remarkably simple. The fully trained network learns the most important eigendirections of the quadratic function f∗(x)f_{*}({\boldsymbol{x}}) and fits them, hence surpassing the NT{\sf NT} model which is confined to a random set of directions.

Let us emphasize that the above separation between NT{\sf NT} and NN{\sf NN} is established only for N≤dN\leq d. It is natural to wonder whether this separation generalizes to N>dN>d for more complicated classes of functions, or if instead it always vanishes for wide networks. We expect the separation to generalize to N>dN>d by considering higher order polynomial, instead of quadratic functions. Partial evidence in this direction is provided by [GMMM19]: for third- or higher-order polynomials NT{\sf NT} does not achieve vanishing risk at any ρ∈(0,∞)\rho\in(0,\infty). The mechanism unveiled by our analysis of quadratic functions is potentially more general: neural networks are superior to linearized models such as RF{\sf RF} or NT{\sf NT}, because they can learn a good representation of the data.

2 Further related work

The connection (and differences) between two-layers neural networks and random features models has been the object of several papers since the original work of Rahimi and Recht [RR08]. An incomplete list of references includes [Bac13, AM15, Bac17a, Bac17b, RR17]. Our analysis contributes to this line of work by establishing a sharp asymptotic characterization, although in more specific data distributions. Sharp results have recently been proven in [GMMM19], for the special case of random weights wi{\boldsymbol{w}}_{i} uniformly distributed over a dd-dimensional sphere. Here we consider the more general case of anisotropic random features with covariance Γ∝̸I{\boldsymbol{\Gamma}}\not\propto{\mathbf{I}}. This clarifies a key reason for suboptimality of random features: the data representation is not adapted to the target function f∗f_{*}. We focus on the population limit n→∞n\to\infty. Complementary results characterizing the variance as a function of nn are given in [HMRT19].

The NT{\sf NT} model (3) is much more recent [JGH18]. Several papers show that SGD optimization within the original neural network is well approximated by optimization within the model NT{\sf NT} as long as the number of neurons is large compared to a polynomial in the sample size N≫nc0N\gg n^{c_{0}} [DZPS18, DLL+18, AZLS18, ZCZG18]. Empirical evidence in the same direction was presented in [LXS+19, ADH+19].

Chizat and Bach [CB18] clarified that any nonlinear statistical model can be approximated by a linear one in an early (lazy) training regime. The basic argument is quite simple. Given a model x↦f(x;θ){\boldsymbol{x}}\mapsto f({\boldsymbol{x}};{\boldsymbol{\theta}}) with parameters θ{\boldsymbol{\theta}}, we can Taylor-expand around a random initialization θ0{\boldsymbol{\theta}}_{0}. Setting θ=θ0+β{\boldsymbol{\theta}}={\boldsymbol{\theta}}_{0}+{\boldsymbol{\beta}}, we get

Here the second approximation holds since, for many random initializations, f(x;θ0)≈0f({\boldsymbol{x}};{\boldsymbol{\theta}}_{0})\approx 0 because of random cancellations. The resulting model βT∇θf(x;θ0){\boldsymbol{\beta}}^{{\mathsf{T}}}\nabla_{{\boldsymbol{\theta}}}f({\boldsymbol{x}};{\boldsymbol{\theta}}_{0}) is linear, with random features.

Our objective is complementary to this literature: we prove that RF{\sf RF} and NT{\sf NT} have limited approximation power, and significant gain can be achieved by full training.

Finally, our analysis of fully trained networks connects to the ample literature on non-convex statistical estimation. For two layers neural networks with quadratic activations, Soltanolkotabi, Javanmard and Lee [SJL19] showed that, as long as the number of neurons satisfies N≥2dN\geq 2d there are no spurious local minimizers. Du and Lee [DL18] showed that the same holds as long as N≥d∧2nN\geq d\wedge\sqrt{2n} where nn is the sample size. Zhong et. al. [ZSJ+17] established local convexity properties around global optima. Further related landscape results include [GLM17, HYV14, GJZ17].

Main results: quadratic functions

As mentioned in the previous section, our results for quadratic functions (qf) assume xi∼N(0,Id){\boldsymbol{x}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}) and yi=f∗(xi)y_{i}=f_{*}({\boldsymbol{x}}_{i}) where

We consider random feature model with first-layer weights (wi)i≤N∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\leq N}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}). We make the following assumptions:

Then, the following holds as N,d→∞N,d\to\infty with N/d→ρN/d\to\rho:

Moreover, assuming ⟨Γ,B⟩2/∥Γ∥F2∥B∥F2\langle{\boldsymbol{\Gamma}},{\boldsymbol{B}}\rangle^{2}/\|{\boldsymbol{\Gamma}}\|_{F}^{2}\|{\boldsymbol{B}}\|_{F}^{2} to have a limit as d→∞d\to\infty, (10) simplifies as follows for ρ→∞\rho\rightarrow\infty:

Notice that RRF,N(f∗)/∥f∗∥L22R_{{\sf RF},N}(f_{*})/\|f_{*}\|_{L_{2}}^{2} is the RF{\sf RF} risk normalized by the risk of the trivial predictor f(x)=0f({\boldsymbol{x}})=0. The asymptotic result in (11) is remarkably simple. By Cauchy-Schwartz, the normalized risk is bounded away from zero even as the number of neurons per dimension diverges ρ=N/d→∞\rho=N/d\to\infty, unless Γ∝B{\boldsymbol{\Gamma}}\propto{\boldsymbol{B}}, i.e. the random features are perfectly aligned with the function to be learned. For isotropic random features, the right-hand side of Eq. (11) reduces to 1−Tr(B)2/(d∥B∥F2)1-\text{\rm Tr}({\boldsymbol{B}})^{2}/(d\|{\boldsymbol{B}}\|_{F}^{2}). In particular, RF{\sf RF} performs very poorly when Tr(B)≪d∥B∥F\text{\rm Tr}({\boldsymbol{B}})\ll\sqrt{d}\|{\boldsymbol{B}}\|_{F}, and no better than the trivial predictor f(x)=0f({\boldsymbol{x}})=0 if Tr(B)=0\text{\rm Tr}({\boldsymbol{B}})=0.

Notice that the above result applies to quite general activation functions. The formulas simplify significantly for quadratic activations.

Under the assumptions of Theorem 1, further assume σ(x)=x2−1\sigma(x)=x^{2}-1. Then we have, as N,d→∞N,d\to\infty with N/d→ρN/d\to\rho:

The right-hand side of Eq. (12) is plotted in Fig. 1 for isotropic features Γ=I/d{\boldsymbol{\Gamma}}={\mathbf{I}}/d, and for optimal features Γ=Γ∗∝B{\boldsymbol{\Gamma}}={\boldsymbol{\Gamma}}^{*}\propto{\boldsymbol{B}}.

2 Neural tangent

For the NT{\sf NT} regime, we focus on quadratic activations and isotropic weights wi∼N(0,Id/d){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}/d).

where the expectation is taken over wi∼i.i.dN(0,Id/d){\boldsymbol{w}}_{i}\sim_{i.i.d}{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}/d).

3 Neural network

For the analysis of SGD-trained neural networks, we assume f∗f_{*} to be a quadratic function as per Eq. (8), but we will now restrict to the positive semidefinite case B⪰0{\boldsymbol{B}}\succeq 0. We consider quadratic activations σ(x)=x2\sigma(x)=x^{2}, and we fix the second layers weights to be 11:

Notice that we use an explicit offset to account for the mismatch in means between f∗f_{*} and f^\hat{f}. It is useful to introduce the population risk, as a function of the network parameters W,c{\boldsymbol{W}},c:

Here expectation is with respect to x∼N(0,Id){\boldsymbol{x}}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}). We will study a one-pass version of SGD, whereby at each iteration kk we perform a stochastic gradient step with respect to a fresh sample (xk,f∗(xk))({\boldsymbol{x}}_{k},f_{*}({\boldsymbol{x}}_{k}))

Then we have (probability is over the initialization (W0,c0)({\boldsymbol{W}}_{0},c_{0}) and the samples)

where λ1(B)≥λ2(B)≥⋯≥λd(B)\lambda_{1}({\boldsymbol{B}})\geq\lambda_{2}({\boldsymbol{B}})\geq\dots\geq\lambda_{d}({\boldsymbol{B}}) are the ordered eigenvalues of B{\boldsymbol{B}}.

The proof of this theorem depends on the following proposition concerning the landscape of the population risk, which is of independent interest.

Let f∗f_{*} be a quadratic function as per Eq. (8), with B⪰0{\boldsymbol{B}}\succeq 0. For any sub-level set of the risk function Ω(B0)={x=(W,c):L(W,c)≤B0}\Omega(B_{0})=\{{\boldsymbol{x}}=({\boldsymbol{W}},c):L({\boldsymbol{W}},c)\leq B_{0}\}, there exists constants ε,δ>0\varepsilon,\delta>0 such that LL is (ε,δ)(\varepsilon,\delta)-strict saddle in the region Ω(B0)\Omega(B_{0}). Namely, for any x∈Ω(B0){\boldsymbol{x}}\in\Omega(B_{0}) with ∥∇L(x)∥2≤ε\|\nabla L({\boldsymbol{x}})\|_{2}\leq\varepsilon, we have λmin⁡(∇2L(x))<−δ\lambda_{\min}(\nabla^{2}L({\boldsymbol{x}}))<-\delta.

We can now compare the risk achieved within the regimes RF{\sf RF}, NT{\sf NT} and NN{\sf NN}. Gathering the results of Corollary 1, and Theorems 2, 3 (using wi∼N(0,I/d){\boldsymbol{w}}_{i}\sim{\sf N}(0,{\mathbf{I}}/d) for RF{\sf RF} and NT{\sf NT}), we obtain

As anticipated, NN{\sf NN} learns the most important directions in f∗f_{*}, while RF{\sf RF}, NT{\sf NT} do not.

Main results: mixture of Gaussians

In this section, we consider the mixture of Gaussian setting (mg): yi=±1y_{i}=\pm 1 with equal probability 1/21/2, and xi∣yi=+1∼N(0,Σ(1)){\boldsymbol{x}}_{i}|y_{i}=+1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(1)}), xi∣yi=−1∼N(0,Σ(2)){\boldsymbol{x}}_{i}|y_{i}=-1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(2)}). We parametrize the covariances as Σ(1)=Σ−Δ{\boldsymbol{\Sigma}}^{(1)}={\boldsymbol{\Sigma}}-{\boldsymbol{\Delta}} and Σ(2)=Σ+Δ{\boldsymbol{\Sigma}}^{(2)}={\boldsymbol{\Sigma}}+{\boldsymbol{\Delta}}, and will make the following assumptions:

There exists constants 0<c1<c20<c_{1}<c_{2} such that c1Id⪯Σ⪯c2Idc_{1}{\mathbf{I}}_{d}\preceq{\boldsymbol{\Sigma}}\preceq c_{2}{\mathbf{I}}_{d};

∥Δ∥op=Θd(1/d)\|{\boldsymbol{\Delta}}\|_{\rm op}=\Theta_{d}(1/\sqrt{d}).

The scaling in assumption M2 ensures the signal-to-noise ratio to be of order one. If the eigenvalues ofΔ{\boldsymbol{\Delta}} are much larger than 1/d1/\sqrt{d}, then it is easy to distinguish the two classes with high probability (they are asymptotically mutually singular). If ∥Δ∥op=od(1/sqrtd)\|{\boldsymbol{\Delta}}\|_{\rm op}=o_{d}(1/sqrt{d}) then no non-trivial classifier exists.

As mentioned in the introduction, the picture emerging from our analysis of the mg model is aligned with the results obtained in the previous section. We will limit ourselves to stating the results without repeating comments that were made above. Our results are compared with simulations in Figure 2. Notice that, in this case, the Bayes error (MMSE) is not achieved even for very wide networks N/d≫1N/d\gg 1 either by NT{\sf NT} or NN{\sf NN}.

As in the previous section, we generate random first-layer weights (wi)i≤N∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\leq N}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}). We consider a general activation function satisfying condition A1{\bf A1}. We make the following assumption on Γ,Σ{\boldsymbol{\Gamma}},{\boldsymbol{\Sigma}}:

Define ζ1(d)≡d Tr(ΣΓΣΓ)/2\zeta_{1}(d)\equiv d\,\text{\rm Tr}({\boldsymbol{\Sigma}}{\boldsymbol{\Gamma}}{\boldsymbol{\Sigma}}{\boldsymbol{\Gamma}})/2, ζ2(d)≡d Tr(ΔΓ)2/4\zeta_{2}(d)\equiv d\,\text{\rm Tr}({\boldsymbol{\Delta}}{\boldsymbol{\Gamma}})^{2}/4. Then, the following holds as N,d→∞N,d\to\infty with N/d→ρN/d\to\rho:

Moreover, assume ζ1(d)\zeta_{1}(d) ζ2(d)\zeta_{2}(d) to have limits as d→∞d\to\infty, i.e. we have lim⁡d→∞ζj(d)=ζj,∗\lim_{d\to\infty}\zeta_{j}(d)=\zeta_{j,*} for j=1,2j=1,2. Then the following holds as ρ→∞\rho\rightarrow\infty:

2 Neural tangent

For the NT{\sf NT} model, we first state our theorem for general Σ{\boldsymbol{\Sigma}} and wi∼N(0,Γ){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) and then give an explicit concentration result in the case Σ=I{\boldsymbol{\Sigma}}={\mathbf{I}} and isotropic weights wi∼N(0,I/d){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}/d).

Assuming further that Σ=I{\boldsymbol{\Sigma}}={\mathbf{I}} and wi∼i.i.d.N(0,Id/d){\boldsymbol{w}}_{i}\sim_{i.i.d.}{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}/d), we have as N,d→∞N,d\to\infty with N/d→ρN/d\to\rho:

In particular, for ρ≥1\rho\geq 1, we have (for almost every W{\boldsymbol{W}})

3 Neural network

We consider quadratic activations with general offset and coefficients f^(x;W,a,c)=∑i=1Nai⟨wi,x⟩2+c\hat{f}({\boldsymbol{x}};{\boldsymbol{W}},{\boldsymbol{a}},c)=\sum_{i=1}^{N}a_{i}\langle{\boldsymbol{w}}_{i},{\boldsymbol{x}}\rangle^{2}+c. This is optimized over (ai,wi)i≤N(a_{i},{\boldsymbol{w}}_{i})_{i\leq N} and cc.

Let us emphasize that, for this setting, we do not have a convergence result for SGD as for the model qf, cf. Theorem 3. However, because of certain analogies between the two models, we expect a similar result to hold for mixtures of Gaussians.

We can now compare the risks achieved within the regimes RF{\sf RF}, NT{\sf NT} and NN{\sf NN}. Gathering the results of Theorems 4, 5 and 6 for Σ=I{\boldsymbol{\Sigma}}={\mathbf{I}} and σ(x)=x2−1\sigma(x)=x^{2}-1 (using wi∼N(0,I/d){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}/d) for RF and NT), we obtain

We recover a similar behavior as in the case of the (qf) model: NN{\sf NN} learns the most important directions of Δ{\boldsymbol{\Delta}}, while RF{\sf RF}, NT{\sf NT} do not. Note that the Bayes error is not achieved in this model.

Numerical Experiments

For the experiments illustrated in Figures 1 and 2, we use feature size of d=450d=450, and number of hidden units N∈{45,⋯ ,4500}N\in\{45,\cdots,4500\}. NT{\sf NT} and NN{\sf NN} models are trained with SGD in TensorFlow [ABC+16]. We run a total of 2.0×1052.0\times 10^{5} SGD steps for each (qf) model and 1.4×1051.4\times 10^{5} steps for each (mg) model. The SGD batch size is fixed at 100100 and the step size is chosen from the grid {0.001,⋯ ,0.03}\{0.001,\cdots,0.03\} where the hyper-parameter that achieves the best fit is used for the figures. RF{\sf RF} models are fitted directly by solving KKT conditions with 5.0×1055.0\times 10^{5} observations. After fitting the model, the test error is evaluated on 1.0×1041.0\times 10^{4} fresh samples. In our figures, each RF{\sf RF} data point corresponds to the test error averaged over 1010 models with independent realizations of W{\boldsymbol{W}}.

For (qf) experiments, we choose B{\boldsymbol{B}} to be diagonal with diagonal elements chosen i.i.d from standard exponential distribution with parameter 11. For (mg) experiments, Δ{\boldsymbol{\Delta}} is also diagonal with the diagonal element chosen uniformly from the set {2d,1.5d,1d}\{\frac{2}{\sqrt{d}},\frac{1.5}{\sqrt{d}},\frac{1}{\sqrt{d}}\}.

Acknowledgements

This work was partially supported by grants NSF DMS-1613091, CCF-1714305, IIS-1741162, and ONR N00014-18-1-2729, NSF DMS-1418362, NSF DMS-1407813.

References

Appendix A Technical background

A.2 Notations

Appendix B Proofs for quadratic functions

Our results for quadratic functions (qf) assume xi∼N(0,Id){\boldsymbol{x}}_{i}\sim{\sf N}(0,{\mathbf{I}}_{d}) and yi=f∗(xi)y_{i}=f_{*}({\boldsymbol{x}}_{i}) where

Note that it is easy to see from the proof that the result stays the same if we add an offset cc.

where V=[V1,…,VN]T{\boldsymbol{V}}=[V_{1},\ldots,V_{N}]^{\mathsf{T}}, and U=(Uij)i,j∈[N]{\boldsymbol{U}}=(U_{ij})_{i,j\in[N]}, with

Simply write the KKT conditions. The optimum is achieved at a=U−1V{\boldsymbol{a}}={\boldsymbol{U}}^{-1}{\boldsymbol{V}}. ∎

B.1.2 Approximation of kernel matrix 𝑼𝑼{\boldsymbol{U}}

where (wi)i∈[N]∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\in[N]}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) independently. Assume conditions A1 and A2 hold.

Then we have as N/d=ρN/d=\rho and d→∞d\to\infty,

Step 1. Hermite expansion of σ\sigma for ∥wi∥2≠1\|{\boldsymbol{w}}_{i}\|_{2}\neq 1. Denote σi(x)=σ(∥wi∥2⋅x)\sigma_{i}(x)=\sigma(\|{\boldsymbol{w}}_{i}\|_{2}\cdot x). First notice that by a change of variables, we get

By Assumption A1, there exists c1<1c_{1}<1 such that

Denote the Hermite expansion of σ\sigma to be

By dominated convergence theorem, we have

In addition, by sub-Gaussianity of the norm of a multivariate Gaussian random variable (see [Ver10]), it is easy to show that

Step 2. Expansion of U{\boldsymbol{U}}. Denote ui=wi/∥wi∥2{\boldsymbol{u}}_{i}={\boldsymbol{w}}_{i}/\|{\boldsymbol{w}}_{i}\|_{2}, then we have

Step 3. Term T0{\boldsymbol{T}}_{0}. By definition of μi\mu_{i}, we have

Recall the change of variable (22) and do a first order Taylor expansion of the exponential: there exists a function ξ(G)∈[0,G]\xi(G)\in[0,G] such that

We see that the integrand goes to zero as t→1t\to 1. For ∣t−1∣|t-1| sufficiently small, we have

Furthermore, for μ=(μi)i∈[N]{\boldsymbol{\mu}}=(\mu_{i})_{i\in[N]} with μi=λ2(∥wi∥22−1)/2\mu_{i}=\lambda_{2}(\|{\boldsymbol{w}}_{i}\|_{2}^{2}-1)/2, we have

where the last equality comes from assumption A2. We get

Step 4. Term T1{\boldsymbol{T}}_{1}. For T1{\boldsymbol{T}}_{1}, we have

By the uniform convergence of ζ1(σi)\zeta_{1}(\sigma_{i}) to λ1(σ)\lambda_{1}(\sigma), cf Eq. (24), we have

where we denoted by G{\boldsymbol{G}} the matrix with columns gi∼N(0,Id/d){\boldsymbol{g}}_{i}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}/d). Hence, we have

Step 5. Term T2{\boldsymbol{T}}_{2}. We have

By the uniform convergence of ζ2(σi)\zeta_{2}(\sigma_{i}) to λ2(σ)\lambda_{2}(\sigma), we have

Moreover, by the estimates in proof of Theorem 2.1 in [EK+10], we have

Step 6. Term ∑k≥3ddiag(Tk)\sum_{k\geq 3}{\rm ddiag}({\boldsymbol{T}}_{k}). Denote ddiag(Tk){\rm ddiag}({\boldsymbol{T}}_{k}) the diagonal matrix composed of diagonal entries of Tk{\boldsymbol{T}}_{k}. We have

Step 7. Term ∑k≥3[Tk−ddiag(Tk)]\sum_{k\geq 3}[{\boldsymbol{T}}_{k}-{\rm ddiag}({\boldsymbol{T}}_{k})]. We have

Combining the bounds (27), (28), (29), (30) and (31) into the decomposition (25) proves the lemma. ∎

B.1.3 Approximation of the 𝑽𝑽{\boldsymbol{V}} vector

Under the assumptions of Theorem 1, define V=(V1,…,VN)T{\boldsymbol{V}}=(V_{1},\ldots,V_{N})^{\mathsf{T}} with

where (wi)i∈[N]∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\in[N]}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) independently. Then as N/d=ρN/d=\rho with d→∞d\to\infty, we have

where P⊥wi{\boldsymbol{P}}_{\perp{\boldsymbol{w}}_{i}} is the projection on the hyperplane orthogonal to wi{\boldsymbol{w}}_{i}, and we recall the definition of ζ2(σi)\zeta_{2}(\sigma_{i}) of Lemma 2:

with GG a standard normal random variable.

We define the following interpolating variables:

and the associated vectors V(1){\boldsymbol{V}}^{(1)}, V(2){\boldsymbol{V}}^{(2)} and V(3){\boldsymbol{V}}^{(3)}. We bound successively the distance between these vectors. We will denote by Pwi{\boldsymbol{P}}_{{\boldsymbol{w}}_{i}} the projection onto vector wi{\boldsymbol{w}}_{i}. First, we consider:

One can check, using a similar argument as for Eq. (26) and dominated convergence, that

Let us first show that the sum is bounded with high probability: denoting g∼N(0,Id){\boldsymbol{g}}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}_{d}), classical sub-Gaussian concentration inequalities (see for example Theorem 6.3.2 in [Ver10]) shows that

where ∥⋅∥ψ2\|\cdot\|_{\psi_{2}} denotes the sub-Gaussian Orlicz norm. By assumption, we have ∥Γ1/2∥op=∥Γ∥op1/2=Od(d−1/2)\|{\boldsymbol{\Gamma}}^{1/2}\|_{\text{op}}=\|{\boldsymbol{\Gamma}}\|_{\text{op}}^{1/2}=O_{d}(d^{-1/2}), and ∥Γ1/2∥F=TrΓ=1\|{\boldsymbol{\Gamma}}^{1/2}\|_{F}=\sqrt{\text{\rm Tr}{\boldsymbol{\Gamma}}}=1. Hence, for wi∼N(0,Γ){\boldsymbol{w}}_{i}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}), we have

Furthermore, we readily have (for example from (23))

Noticing that Tr(wiwiTB)=∥B1/2wi∥22\text{\rm Tr}({\boldsymbol{w}}_{i}{\boldsymbol{w}}_{i}^{\mathsf{T}}{\boldsymbol{B}})=\|{\boldsymbol{B}}^{1/2}{\boldsymbol{w}}_{i}\|^{2}_{2} and by the same argument as for (34), we have:

By assumption A2, we have ∥B1/2Γ1/2∥op≤∥B1/2∥op∥Γ1/2∥op=Od(d−1/2)\|{\boldsymbol{B}}^{1/2}{\boldsymbol{\Gamma}}^{1/2}\|_{\text{op}}\leq\|{\boldsymbol{B}}^{1/2}\|_{\text{op}}\|{\boldsymbol{\Gamma}}^{1/2}\|_{\text{op}}=O_{d}(d^{-1/2}) and

Combining the bounds (36), (37) and (39) into (33), we get

which, combined with (39) and (41), yields

where V(3)=λ2Tr(ΓB)1{\boldsymbol{V}}^{(3)}=\lambda_{2}\text{\rm Tr}({\boldsymbol{\Gamma}}{\boldsymbol{B}}){\mathbf{1}}. Combining the above three bounds (33), (42) and (43) yields the desired result. ∎

The following proposition is stated in slightly more general terms, in order to be used in both the proofs of Theorem 1 and Theorem 4.

where D\mathcal{D} is the empirical distribution of eigenvalues of d⋅Γd\cdot{\boldsymbol{\Gamma}}.

The proof of Proposition 2 is a direct combination of Lemma 4, 5, and 6 below.

Let (wi)i∈[N]∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\in[N]}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) independently. Assume condition A2 holds (resp. B2). Let μ=(∥wi∥22−1)i∈[N]{\boldsymbol{\mu}}=(\|{\boldsymbol{w}}_{i}\|_{2}^{2}-1)_{i\in[N]}, and A0=c1IN+c2WTW{\boldsymbol{A}}_{0}=c_{1}{\mathbf{I}}_{N}+c_{2}{\boldsymbol{W}}^{\mathsf{T}}{\boldsymbol{W}}, where c1≡c1(d)c_{1}\equiv c_{1}(d) and c2≡c2(d)c_{2}\equiv c_{2}(d) are constants that are asymptotically upper and lower bounded by strictly positive constants. Then as d→∞d\rightarrow\infty and N/d→ρN/d\to\rho, we have

We first prove the lemma under the following extra assumption on the covariance matrix: there exists a (fixed) integer KK such that

for some orthogonal matrix Q{\boldsymbol{Q}} and d⋅γi≤Cd\cdot\gamma_{i}\leq C. Furthermore, there exists an ε>0{\varepsilon}>0 such that dk/d≥εd_{k}/d\geq{\varepsilon} for dd sufficiently large.

Without loss of generality, we assume Γ=diag(γ1Id1,…,γKIdK){\boldsymbol{\Gamma}}={\rm diag}(\gamma_{1}{\mathbf{I}}_{d_{1}},\ldots,\gamma_{K}{\mathbf{I}}_{d_{K}}), and we divide wi{\boldsymbol{w}}_{i} into vectors corresponding to each block

Using the fact that ∥g∥2\|{\boldsymbol{g}}\|_{2} is independent of g/∥g∥2{\boldsymbol{g}}/\|{\boldsymbol{g}}\|_{2} for g∼N(0,I){\boldsymbol{g}}\sim{\sf N}({\boldsymbol{0}},{\mathbf{I}}), the following two sets of random variables have the same distribution:

Since dk→∞d_{k}\to\infty as d→∞d\to\infty, we have

Combining (47), (48) and (49) proves the lemma in the case of a covariance of the form (46):

Step 4. From discrete to continuous spectrum.

We consider Γ{\boldsymbol{\Gamma}} a covariance matrix verifying assumption A2. For a given ε>0{\varepsilon}>0 and KK sufficiently large, we consider Γε{\boldsymbol{\Gamma}}_{{\varepsilon}} a matrix obtained from Γ{\boldsymbol{\Gamma}} by binning its eigenvalues to at most KK points of [0,C/d][0,C/d], such that we have Tr(Γε)=1\text{\rm Tr}({\boldsymbol{\Gamma}}_{{\varepsilon}})=1 and lim⁡d→∞d⋅∥Γ−Γε∥op≤ε\lim_{d\to\infty}d\cdot\|{\boldsymbol{\Gamma}}-{\boldsymbol{\Gamma}}_{{\varepsilon}}\|_{{\rm op}}\leq{\varepsilon} (recall that ∥Γ∥op≤C/d\|{\boldsymbol{\Gamma}}\|_{{\rm op}}\leq C/d by assumption). Such a matrix always exists from the condition Tr(Γ)=1\text{\rm Tr}({\boldsymbol{\Gamma}})=1 and the weak convergence of the spectrum of d⋅Γd\cdot{\boldsymbol{\Gamma}}.

Furthermore, using Tr(Γ−Γε)=0\text{\rm Tr}({\boldsymbol{\Gamma}}-{\boldsymbol{\Gamma}}_{{\varepsilon}})=0, we have

Noticing that ∥A0−1∥op,∥A0,ε−1∥op≤c1−1\|{\boldsymbol{A}}_{0}^{-1}\|_{{\rm op}},\|{\boldsymbol{A}}_{0,{\varepsilon}}^{-1}\|_{{\rm op}}\leq c_{1}^{-1}, and using (50) applied to Γε{\boldsymbol{\Gamma}}_{{\varepsilon}}, we get for dd sufficiently large:

Taking a sequence δ→0\delta\to 0 and ε{\varepsilon} such that ε∝Cδ−1{\varepsilon}\propto C^{-1}_{\delta} shows that this is equivalent to

Taking ε∝δ{\varepsilon}\propto\sqrt{\delta}, we deduce that this is equivalent to

Substituting (52) and (53) in (51) concludes the proof. ∎

Under the same setting as Proposition 2, we have

Define z=κ1/d{\boldsymbol{z}}=\sqrt{\kappa}{\mathbf{1}}/\sqrt{d}. Then we have

By Sherman Morrison Woodbury formula, we have

In the following, we give an asymptotic expression for ⟨1,A0−11⟩/d\langle{\mathbf{1}},{\boldsymbol{A}}_{0}^{-1}{\mathbf{1}}\rangle/d.

Let ρ∈(0,∞)\rho\in(0,\infty). We have almost surely

In addition, assume D\mathcal{D} is the limiting spectral distribution of d⋅Γd\cdot{\boldsymbol{\Gamma}}. Then, we have almost surely

We know due to fast concentration of ∥z∥2/N\|z\|^{2}/N around one (see e.g. [BLM13]), P1P_{1} vanish exponentially fast in NN (equivalently in dd since N/dN/d is fixed to be ρ\rho).

Now, let’s consider P2P_{2}. ∣zTA0−1z/N−Tr(A0−1)/N∣|{\boldsymbol{z}}^{\mathsf{T}}{\boldsymbol{A}}_{0}^{-1}{\boldsymbol{z}}/N-\text{\rm Tr}({\boldsymbol{A}}_{0}^{-1})/N|. By Hanson-Wright inequality (see e.g. [BLM13]), we have

B.1.5 Proof of Theorem 1

By Lemma 1, the risk has a representation

B.2 Neural Tangent model: proof of Theorem 2

We can rewrite the neural tangent model with a squared non-linearity σ(x)=x2\sigma(x)=x^{2} as

where Bij=PiTBPj{\boldsymbol{B}}_{ij}={\boldsymbol{P}}_{i}^{\mathsf{T}}{\boldsymbol{B}}{\boldsymbol{P}}_{j} for i,j=1,2i,j=1,2. We readily deduce that

The last term of the sum (65) is also derived by first conditioning on v1{\boldsymbol{v}}_{1}. Let us denote z1=P⊥v1e1{\boldsymbol{z}}_{1}={\boldsymbol{P}}_{\perp{\boldsymbol{v}}_{1}}{\boldsymbol{e}}_{1} and z2=P⊥v1e2{\boldsymbol{z}}_{2}={\boldsymbol{P}}_{\perp{\boldsymbol{v}}_{1}}{\boldsymbol{e}}_{2} the projections of (e1,e2)({\boldsymbol{e}}_{1},{\boldsymbol{e}}_{2}) on the hyperplane perpendicular to v1{\boldsymbol{v}}_{1}, on which v2{\boldsymbol{v}}_{2} is uniformly distributed over the unit sphere. We decompose z2{\boldsymbol{z}}_{2} into two components: one along z1{\boldsymbol{z}}_{1} that we denote z2(1)=P∥z1z2{\boldsymbol{z}}_{2}^{(1)}={\boldsymbol{P}}_{\parallel{\boldsymbol{z}}_{1}}{\boldsymbol{z}}_{2} and one perpendicular to z1{\boldsymbol{z}}_{1}, denoted z2(2)=P⊥z1z2{\boldsymbol{z}}_{2}^{(2)}={\boldsymbol{P}}_{\perp{\boldsymbol{z}}_{1}}{\boldsymbol{z}}_{2}. Then we have:

where we used the same argument as for (66). Plugging the above limits (66), (67) and (68) in the expansion (65), we get

The above formula for the RF{\sf RF} risk Eq. (69) has two terms that corresponds to the two limits Tr(B)/∥B∥F=od(d)\text{\rm Tr}({\boldsymbol{B}})/\|{\boldsymbol{B}}\|_{F}=o_{d}(\sqrt{d}) (e.g. spiked matrix)

and Tr(B)2=d∥B∥F2\text{\rm Tr}({\boldsymbol{B}})^{2}=d\|{\boldsymbol{B}}\|_{F}^{2} (i.e. B∝I{\boldsymbol{B}}\propto{\mathbf{I}})

B.3 Neural Network model: proof of Theorem 3

We consider two-layers neural networks with quadratic activation function σ(x)=x2\sigma(x)=x^{2} and we fix the second layer weights to 11,

We consider the ground truth function f∗f_{*} to be a quadratic function as per Eq. (20), and the risk function defined by

We consider running SGD dynamics upon the risk function for a fresh sample (xk,f∗(xk))({\boldsymbol{x}}_{k},f_{*}({\boldsymbol{x}}_{k})) for each iteration

The infimum of LL over W{\boldsymbol{W}} is equivalent to the low-rank approximation problem of matrix B{\boldsymbol{B}} in Frobenius norm, with rank less or equal to max⁡(d,N)\max(d,N), and is given by the Eckart-Young-Mirsky theorem (see [EY36]). ∎

B.3.2 Landscape: proof of Proposition 1

Without loss of generality, throughout the proof, we assume that B{\boldsymbol{B}} is diagonal and b0=0b_{0}=0. Our first proposition characterizes the critical points of L(W,c)L({\boldsymbol{W}},c).

Then for any critical point (W0,c0)({\boldsymbol{W}}_{0},c_{0}) of L(W,c)L({\boldsymbol{W}},c), there exists a projection matrix P=∑i=1keτ(i)eτ(i)T{\boldsymbol{P}}=\sum_{i=1}^{k}{\boldsymbol{e}}_{\tau(i)}{\boldsymbol{e}}_{\tau(i)}^{\mathsf{T}} for some injection τ:[k]→[d]\tau:[k]\to[d], such that Γ0=W0W0T{\boldsymbol{\Gamma}}_{0}={\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}} is diagonal and satisfy

We consider the gradient of this function. We get:

By the stationary condition, at a critical point (W0,c0)({\boldsymbol{W}}_{0},c_{0}), we must have:

This is of the form of the eigenvalue equation of matrix B{\boldsymbol{B}}. Hence we must have the columns of U1{\boldsymbol{U}}_{1} to be a set of eigenvectors and S12{\boldsymbol{S}}_{1}^{2} to be positive eigenvalues of B{\boldsymbol{B}}. This proves the proposition. ∎

Note the global minimizers are attained for Γ0=W0W0T{\boldsymbol{\Gamma}}_{0}={\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}} corresponding to the min⁡(N,d)\min(N,d) directions of B{\boldsymbol{B}} with the largest eigenvalues. We prove in the following proposition that stationary points that are not global minimizers are strict saddle points.

Define the spectral separation of B{\boldsymbol{B}} as

and δeig\delta^{\textup{eig}} the minimum strictly positive eigenvalue of B{\boldsymbol{B}}.

Consider (W0,c0)({\boldsymbol{W}}_{0},c_{0}) a stationary point of L(W,c)L({\boldsymbol{W}},c) but not a global minimizer. Then, we have

Let us first compute the Hessian of the risk with respect to the W{\boldsymbol{W}} variable. We have

Plugging the value of c0c_{0} at a critical point (cf Eq. (70)), we get

Case 1: Consider the case rank(W0)<min⁡{rank(B),N}\text{rank}({\boldsymbol{W}}_{0})<\min\{\text{rank}({\boldsymbol{B}}),N\}. Then there exists an i∈[d]i\in[d] such that Bii>0{\boldsymbol{B}}_{ii}>0 (recall that we assumed B{\boldsymbol{B}} diagonal , with diagonal elements given by the positive eigenvalues of B{\boldsymbol{B}}) and (W0W0T)ii=0({\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}})_{ii}=0. For simplicity, let us permute the coordinates so that i=1i=1. The singular value decomposition of W0{\boldsymbol{W}}_{0} verifies

We have ∥Z∥F=1\|{\boldsymbol{Z}}\|_{F}=1 and W0ZT=0{\boldsymbol{W}}_{0}{\boldsymbol{Z}}^{\mathsf{T}}=0. Plugging these matrices in the above expression of the Hessian, see Eq. (73), we get

Case 2: Consider the case when rank(W0W0T)=N<rank(B)\text{rank}({\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}})=N<\text{rank}({\boldsymbol{B}}) and W0W0T{\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}} does not correspond to the NN largest eigenvalues of B{\boldsymbol{B}}. Then there exists i≠j∈[n]i\neq j\in[n], such that Bii>Bjj{\boldsymbol{B}}_{ii}>{\boldsymbol{B}}_{jj}, (W0W0T)ii=0({\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}})_{ii}=0 and (W0W0T)jj=Bjj({\boldsymbol{W}}_{0}{\boldsymbol{W}}_{0}^{\mathsf{T}})_{jj}={\boldsymbol{B}}_{jj}. For simplicity, let us permute the coordinates such that i=1i=1 and j=2j=2. The SVD decomposition of W0{\boldsymbol{W}}_{0} now verifies:

We have ∥Z∥F=1\|{\boldsymbol{Z}}\|_{F}=1. Plugging these matrices in the above expression of the Hessian (73), note

First, remark that L(W,c)L({\boldsymbol{W}},c) has compact sub-level sets. The proposition then follows from Proposition 4 and the continuity of the gradient ∇L(x)\nabla L({\boldsymbol{x}}) and of the minimum eigenvalue of the Hessian λmin⁡(∇2L(x))\lambda_{\min}(\nabla^{2}L({\boldsymbol{x}})). ∎

B.3.3 Dynamics

The following lemma is a standard combination of Lojasiewicz inequality and center and stable manifold theorem. We prove it for completeness.

Then for (Lebesgue) almost all initialization x0{\boldsymbol{x}}_{0}, there exists a second order local minimizer x∗{\boldsymbol{x}}_{*}, such that

Step 1. Show convergence to a critical point. Since ff is an analytic function, by Lojasiewicz inequality [Loj82], and the fact that the level set of ff is compact, we have

for x∗{\boldsymbol{x}}_{*} some critical point of ff.

Step 2. Show convergence to a local minimizer. In this step, we proceed similarly to the proof of Theorem 3 in [PP16]. First, consider a sublevel set

Then we have Ω(K)\Omega(K) compact. Since ff is an analytic function, ∇f\nabla f is Lipschitz in the compact set Ω(K)\Omega(K). We define the map ϕt:Ω(K)→ϕt(Ω(K))\phi_{t}:\Omega(K)\to\phi_{t}(\Omega(K)), x↦xt{\boldsymbol{x}}\mapsto{\boldsymbol{x}}_{t} where xt{\boldsymbol{x}}_{t} is defined as the solution of

By Picard’s existence and uniqueness theorem, we have ϕt\phi_{t} is a diffeomorphism from Ω(K)\Omega(K) to ϕ(Ω(K))\phi(\Omega(K)) for any t>0t>0. Fix an ε0>0\varepsilon_{0}>0, and we define g=ϕε0:Ω(K)→Ω(K)g=\phi_{\varepsilon_{0}}:\Omega(K)\to\Omega(K).

Let r{\boldsymbol{r}} be a strict saddle point of ff, then r{\boldsymbol{r}} must be an unstable fixed point of the diffeomorphism g=ϕε0g=\phi_{{\varepsilon}_{0}}. By center and stable manifold theorem (such as Theorem 9 in [PP16]), there exists a manifold Wlocsc(r)W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}}) of dimension at most d−1d-1, and a ball B(r,ε(r)){\mathsf{B}}({\boldsymbol{r}},\varepsilon({\boldsymbol{r}})) centered at r{\boldsymbol{r}} with radius ε(r)\varepsilon({\boldsymbol{r}}), such that we have the following facts:

g(Wlocsc(r)∩B(r,ε(r)))⊆Wlocsc(r)g\left(W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}})\cap{\mathsf{B}}({\boldsymbol{r}},\varepsilon({\boldsymbol{r}}))\right)\subseteq W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}});

If gn(x)∈B(r,ε(r))g^{n}({\boldsymbol{x}})\in{\mathsf{B}}({\boldsymbol{r}},\varepsilon({\boldsymbol{r}})) for all n≥0n\geq 0, we have x∈Wlocsc(r){\boldsymbol{x}}\in W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}}) (here gng^{n} means composition of gg for nn times).

We consider the union of the balls associated to all the strict saddle points of ff in Ω(K)\Omega(K)

Due to Lindelof’s lemma, we can find a countable subcover for AA, i.e., there exists fixed-points r1,r2,…{\boldsymbol{r}}_{1},{\boldsymbol{r}}_{2},\ldots such that A=∪m=1∞B(rm,ε(rm))A=\cup_{m=1}^{\infty}{\mathsf{B}}({\boldsymbol{r}}_{m},\varepsilon({\boldsymbol{r}}_{m})). If gradient descent converges to a strict saddle point, starting from a point v∈Ω(K){\boldsymbol{v}}\in\Omega(K), there must exist a t0t_{0} and mm such that ϕt(v)∈B(rm,ε(rm))\phi_{t}({\boldsymbol{v}})\in{\mathsf{B}}({\boldsymbol{r}}_{m},\varepsilon({\boldsymbol{r}}_{m})) for all t≥t0t\geq t_{0}. By center and stable manifold theorem, we get that ϕt(v)∈Wlocsc(rm)∩Ω(K)\phi_{t}({\boldsymbol{v}})\in W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}}_{m})\cap\Omega(K). By setting D1(rm)=g−1(Wlocsc(rm)∩Ω(K))D_{1}({\boldsymbol{r}}_{m})=g^{-1}(W_{{\rm loc}}^{{\rm sc}}({\boldsymbol{r}}_{m})\cap\Omega(K)) and Di+1(rm)=g−1(Di(rm)∩Ω(K))D_{i+1}({\boldsymbol{r}}_{m})=g^{-1}(D_{i}({\boldsymbol{r}}_{m})\cap\Omega(K)) we get that v∈Dk(rm){\boldsymbol{v}}\in D_{k}({\boldsymbol{r}}_{m}) for all kε0≥t0k\varepsilon_{0}\geq t_{0}. Hence the set of initial points in Ω(K)\Omega(K) such that gradient descent converges to a strict saddle point is a subset of

The following lemma is standard, and a corollary of Theorem 2.11 in [Kur70].

Let xt{\boldsymbol{x}}_{t} be the trajectory of

with initialization x0∈Ω{\boldsymbol{x}}_{0}\in\Omega. Further assume that there exists η>0\eta>0, such that ∪t≥0B(xt,η)⊆Ω\cup_{t\geq 0}{\mathsf{B}}({\boldsymbol{x}}_{t},\eta)\subseteq\Omega.

Consider the following Markov jump process xt,ε{\boldsymbol{x}}_{t,\varepsilon} starting from x0{\boldsymbol{x}}_{0}, with jump time to be an exponential random variable with fixed mean ε\varepsilon, and jump direction −ε∇f(x;z)-\varepsilon\nabla f({\boldsymbol{x}};{\boldsymbol{z}}) where x{\boldsymbol{x}} is the current state, and z{\boldsymbol{z}} an independent sample. Then we have for any fixed T>0T>0 and δ>0\delta>0,

B.3.4 Proof of Theorem 3

By Proposition 4, we know that for L(W,c)L({\boldsymbol{W}},c), any critical point that is not a global minimizer is a strict saddle point. Consider the gradient flow

with random initialization (W0,c0)∼ν0({\boldsymbol{W}}_{0},c_{0})\sim\nu_{0} where ν0\nu_{0} is a distribution that is absolutely continuous with respect to Lebesgue measure. Since L(W,c)L({\boldsymbol{W}},c) is an analytic function, by Lemma 8, we have (Wt,ct)({\boldsymbol{W}}_{t},c_{t}) converges to a global minimizer of L(W,c)L({\boldsymbol{W}},c). That is, we have almost surely (over ν0\nu_{0})

where inf⁡W,cL(W,c)\inf_{{\boldsymbol{W}},c}L({\boldsymbol{W}},c) is calculated in Lemma 7.

Consider the following Markov jump process (Wt,ε,ct,ε)({\boldsymbol{W}}_{t,\varepsilon},c_{t,\varepsilon}) starting from (W0,ct)∼ν0({\boldsymbol{W}}_{0},c_{t})\sim\nu_{0}, with jump time to be an exponential random variable with fixed mean ε\varepsilon, and jump direction to be −ε∇L(W,c;z)-\varepsilon\nabla L({\boldsymbol{W}},c;{\boldsymbol{z}}) where

with (W,c)({\boldsymbol{W}},c) the current state, and z{\boldsymbol{z}} an independent sample. By Lemma 9, we have for any fixed T>0T>0 and δ>0\delta>0,

Note the sequence of Markov jump process at jump time is exactly the SGD iterates. Hence the SGD iterates with properly scaled number of iterations is uniformly close to (Wt,ct)({\boldsymbol{W}}_{t},c_{t}) over finite horizon as ε→0{\varepsilon}\to 0. This proves the Theorem.

Appendix C Proofs for Mixture of Gaussians

In this section, we consider the mixture of Gaussian setting (mg): yi=±1y_{i}=\pm 1 with equal probability 1/21/2, and xi∣yi=+1∼N(0,Σ(1)){\boldsymbol{x}}_{i}|y_{i}=+1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(1)}), xi∣yi=−1∼N(0,Σ(2)){\boldsymbol{x}}_{i}|y_{i}=-1\sim{\sf N}(0,{\boldsymbol{\Sigma}}^{(2)}) where Σ(1)=Σ−Δ{\boldsymbol{\Sigma}}^{(1)}={\boldsymbol{\Sigma}}-{\boldsymbol{\Delta}} and Σ(2)=Σ+Δ{\boldsymbol{\Sigma}}^{(2)}={\boldsymbol{\Sigma}}+{\boldsymbol{\Delta}}. With these notations,

Throughout this section, we will make the following assumptions:

There exists constants 0<c1<c20<c_{1}<c_{2} such that c1Id⪯Σ⪯c2Idc_{1}{\mathbf{I}}_{d}\preceq{\boldsymbol{\Sigma}}\preceq c_{2}{\mathbf{I}}_{d};

∥Δ∥op=Θd(1/d)\|{\boldsymbol{\Delta}}\|_{\rm op}=\Theta_{d}(1/\sqrt{d}).

Note that it is easy to see from the proof that the result stays the same if we add an offset cc.

Consider the RF model introduced above. We have

where V=[V1,…,VN]T{\boldsymbol{V}}=[V_{1},\ldots,V_{N}]^{\mathsf{T}}, and U=(Uij)i,j∈[N]{\boldsymbol{U}}=(U_{ij})_{i,j\in[N]}, with

Simply write the KKT conditions. The optimum is achieved at a=U−1V{\boldsymbol{a}}={\boldsymbol{U}}^{-1}{\boldsymbol{V}}. ∎

C.1.2 Approximation of kernel matrix 𝑼𝑼{\boldsymbol{U}}

where (wi)i∈[N]∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\in[N]}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) independently. Assume conditions A1 and B2 hold.

Then we have as N/d=ρN/d=\rho and d→∞d\to\infty, we have

Recalling that in the (mg{\sf mg}) model, we have x∼(1/2)⋅N(0,I−Δ)+(1/2)⋅N(0,I+Δ){\boldsymbol{x}}\sim(1/2)\cdot{\sf N}({\boldsymbol{0}},{\mathbf{I}}-{\boldsymbol{\Delta}})+(1/2)\cdot{\sf N}({\boldsymbol{0}},{\mathbf{I}}+{\boldsymbol{\Delta}}), we have

We also have Tr(ΔΓΔΓ)2=od(d−1)\text{\rm Tr}({\boldsymbol{\Delta}}{\boldsymbol{\Gamma}}{\boldsymbol{\Delta}}{\boldsymbol{\Gamma}})^{2}=o_{d}(d^{-1}) by assumptions M2 and B2, hence

Therefore, combining (76) and (77), we get:

Combining (75) and (78) concludes the proof. ∎

C.1.3 Approximation of the 𝑽𝑽{\boldsymbol{V}} vector

Under the assumption of Theorem 4, define V=(V1,…,VN)T{\boldsymbol{V}}=(V_{1},\ldots,V_{N})^{\mathsf{T}} with

where (wi)i∈[N]∼N(0,Γ)({\boldsymbol{w}}_{i})_{i\in[N]}\sim{\sf N}({\boldsymbol{0}},{\boldsymbol{\Gamma}}) independently. Then as N/d=ρN/d=\rho with d→∞d\to\infty, we have

Using dominated convergence theorem and arguments similar to those used to prove (26), one can check that

The same arguments as in the proofs of Lemma 2 and Lemma 3 show

Combining (80) with (81) in (79), we get:

Bounding similarly the term depending on (I+Δ)1/2wi({\mathbf{I}}+{\boldsymbol{\Delta}})^{1/2}{\boldsymbol{w}}_{i} in Vi(1)V_{i}^{(1)}, we get

Now, consider the difference between V(1){\boldsymbol{V}}^{(1)} and V(2){\boldsymbol{V}}^{(2)}. We use the fact for xx on a neighborhood of , there exists cc such that

where the last equality is due to assumptions M2 and B2. We conclude that

For the last comparison between V(2){\boldsymbol{V}}^{(2)} and V(3){\boldsymbol{V}}^{(3)}, we take the expectation:

Combining the above three bounds (82), (83) and (84) yields the desired result. ∎

C.1.4 Proof of Theorem 4

By Lemma 10, the risk has a representation

C.2 Neural Tangent model: proof of Theorem 5

Assume conditions M1 and M2 hold. Consider the function

Define the risk function optimized over a,ca,c while Γ{\boldsymbol{\Gamma}} is fixed

Minimizing successively over cc and aa, we get the following formula:

By Assumptions M1 and M2, we have Σ⪰cId{\boldsymbol{\Sigma}}\succeq c{\mathbf{I}}_{d} and ∥Δ∥op≤C/d\|{\boldsymbol{\Delta}}\|_{{\rm op}}\leq C/\sqrt{d} for some constants cc and CC. We get

C.2.2 Proof of Theorem 5

where the minimizer G=Δ{\boldsymbol{G}}={\boldsymbol{\Delta}} is obtained by Cauchy-Schwarz inequality.

Case N/d→ρ<1N/d\to\rho<1. Consider now the case when N<dN<d. From (88), the optimal G{\boldsymbol{G}} is the one maximizing

which we rewrite as the following convex problem

which yields, using P1TP1=IN{\boldsymbol{P}}_{1}^{\mathsf{T}}{\boldsymbol{P}}_{1}={\mathbf{I}}_{N},

where Δij=PiTΔPj{\boldsymbol{\Delta}}_{ij}={\boldsymbol{P}}_{i}^{\mathsf{T}}{\boldsymbol{\Delta}}{\boldsymbol{P}}_{j} for i,j=1,2i,j=1,2. The constraint reads in the P{\boldsymbol{P}} basis

Considering the (unique) symmetric optimizer G1{\boldsymbol{G}}_{1} and substituting (92) in (90), we get the minimizer

Substituting (94) in (88), we then obtain

where Δ22=PW⊥ΔPW⊥{\boldsymbol{\Delta}}_{22}={\boldsymbol{P}}_{{\boldsymbol{W}}^{\perp}}{\boldsymbol{\Delta}}{\boldsymbol{P}}_{{\boldsymbol{W}}^{\perp}} with PW⊥=Id−W(WTW)−1WT{\boldsymbol{P}}_{{\boldsymbol{W}}^{\perp}}={\mathbf{I}}_{d}-{\boldsymbol{W}}({\boldsymbol{W}}^{\mathsf{T}}{\boldsymbol{W}})^{-1}{\boldsymbol{W}}^{\mathsf{T}} is the random projection along the orthogonal subspace to the columns of W{\boldsymbol{W}}. From Theorem 2, we know that

by assumption M2 on Δ{\boldsymbol{\Delta}}. We deduce that there exists a constant cc (that depends on ρ\rho and CC) such that:

Using (97) and (95), we deduce the final high probability formula for the risk of the NT model:

C.3 Neural Network model: proof of Theorem 6

where we consider the function class of two-layers neural networks (with NN neurons) with quadratic activation function and general offset and coefficients

We define the risk function for a given set of parameters as

The risk is optimized over (ai,wi)i≤N(a_{i},{\boldsymbol{w}}_{i})_{i\leq N} and cc.

where A=diag(a){\boldsymbol{A}}={\rm diag}({\boldsymbol{a}}). Define Γ=WAWT{\boldsymbol{\Gamma}}={\boldsymbol{W}}{\boldsymbol{A}}{\boldsymbol{W}}^{\mathsf{T}} and using Eq. (87) in Lemma 13, the minimizer Γ∗{\boldsymbol{\Gamma}}^{*} is the solution of

where the λi\lambda_{i}’s are the singular values of Δ{\boldsymbol{\Delta}} in descending order. Plugging this expression in Eq. (87) concludes the proof. ∎