High-dimensional limit theorems for SGD: Effective dynamics and critical scaling
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We study the scaling limits of stochastic gradient descent (SGD) with constant stepsize in the high-dimensional regime.
We prove limit theorems for the trajectoriesof summary statistics (i.e., ﬁnite-dimensional functions) of SGD as the dimensiongoes to inﬁnity.
Our approach allows one to choose the summary statistics thatare tracked, the initialization, and the step-size.
It yields both ballistic (ODE) anddiffusive (SDE) limits, with the limit depending dramatically on the former choices.
We ﬁnd a critical scaling regime for the step-size below which this “effectivedynamics" matches gradient ﬂow for the population loss, but at which, a newcorrection term appears which changes the phase diagram.
About the ﬁxed pointsof this effective dynamics, the corresponding diffusive limits can be quite complexand even degenerate.
We demonstrate our approach on popular examples includingestimation for spiked matrix and tensor models and classiﬁcation via two-layernetworks for binary and XOR-type Gaussian mixture models.
These examplesexhibit surprising phenomena including multimodal timescales to convergence aswell as convergence to sub-optimal solutions with probability bounded away fromzero from random (e.g., Gaussian) initializations.
Introduction
Background. Stochastic gradient descent (SGD) is the go-to method for large-scale optimizationproblems in data science.
It is used to train complex parametric models on high-dimensional data.
Since its introduction in [39], there has been a tremendous amount of work in analyzing its evolution.
In ﬁxed dimensions, the asymptotic theory of SGD and stochastic approximations more broadly, isby now classical.
There have been works on path-wise limit theorems, such as functional central limittheorems and even large deviations principles [39, 31, 28, 22, 17, 7, 16, 6].
At the core of this line ofwork is the idea that in the limit where the step-size, or learning rate, tends to zero, the trajectory of
SGD with a ﬁxed loss function (appropriately rescaled in time) converges to the solution of gradientﬂow for the population loss with the same initialization.
Recently there has been considerable interestin quantifying the rate of this trajectory-wise convergence to higher order, in terms of a diffusionapproximation.
Namely, there are many works developing asymptotic expansions of the trajectory inthe learning rate [29, 23, 25, 1, 26].
Motivated by this, there is a rich line of work bounding the timeto equilibrium for the associated diffusion approximation (as well as Langevin–type modiﬁcations)under uniform ellipticity assumptions [29, 36, 11, 53].
There is also an interesting line of workobtaining PDE limits in the “shallow network” regime where the dimension of the parameter spacediverges but the dimension of the data remains constant: see e.g., [32, 40, 12, 45, 2].
In recent years, there has been considerable interest in understanding the high-dimensional setting ,where one is constrained in the amount of data or the run-time of the algorithm due to the high- dimensional nature of the data and the complexity of the model being trained.
In these regimes, onecannot simply take the learning rate to be arbitrarily small as this would force an unlimited samplesize and run-time.
This is a common issue in high-dimensional statistics and the standard analyticapproach is to study regimes where the sample size scales with the dimension of the problem [50, 51].
For SGD with constant learning rate, there has been recent progress on quantifying the dimensiondependence of the sample complexity for various tasks on general (pseudo or quasi-) convex objectives[8, 9, 44, 34, 21, 15] and special classes of non-convex objectives [19, 47, 3].
There has also beenimportant work on scaling limits as the dimension tends to inﬁnity for the speciﬁc problems of linearregression [52, 35], Online PCA [52, 24], and phase retrieval [47] from random starts, and teacherstudent networks [41, 42, 20, 49] and two-layer networks for XOR Gaussian mixtures [37] from warmstarts.
We also note that the study of high-dimensional regimes of gradient descent and Langevindynamics have a history from the statistical physics perspective, e.g., in [13, 14, 43, 30, 10, 27]. Our contributions.
We develop a uniﬁed approach to the scaling limits of SGD in high-dimensionswith constant learning rate that allows us to understand a broad range of estimation tasks.
One ofcourse cannot develop a high-dimensional scaling limit for the full trajectory of SGD as the dimensionof the underlying parameter space is growing.
On the other hand, in practice, one is rarely interestedin the full trajectory; instead one typically tracks the trajectory of various summary statistics of thealgorithm’s evolution, such as the loss, the amplitude of various weights, or correlations between theclassiﬁer and the ground truth (in a supervised setting).
We show in Theorem 2.2 that under mildregularity assumptions, the evolution of these summary statistics converges as the dimension growsto the solution of a system of (possibly stochastic) differential equations.
These effective dynamicsdepend dramatically on the initializations (warm vs. random or cold), the parameter regions in whichone is developing the scaling limit, and the scaling of the step-size with the dimension.
In practice, SGD often exhibits two types of phases in training: ballistic phases where the summarystatistics macroscopically change in value, and diffusive phases , where they ﬂuctuate microscopically.(During training, the evolution can start with either, and can even alternate multiple times betweenthese phases.)
Our approach allows us to develop scaling limits for both types of phases.
In ballistic phases, the effective dynamics are given by an ordinary differential equation (ODE) andthe ﬁnite-dimensional intuition that the summary statistics evolve under the gradient ﬂow for thepopulation loss is correct provided the (constant) learning rate is sufﬁciently small in the dimension.
When the learning rate follows a certain critical scaling—matching scalings commonly used in thehigh-dimensional statistics literature—an additional correction term appears.
At this critical scaling,the phase portrait deviates signiﬁcantly from that of the population gradient ﬂow.
Furthermore, inmicroscopic neighborhoods of the ﬁxed points of this ODE, the effective dynamics become diffusiveand are given by SDEs which can exhibit a wide range of (possibly degenerate) behaviors.
We notethat the appearance of the correction term in the ballistic phase was ﬁrst observed in the setting ofteacher-student networks in [41, 42] and very recently investigated in detail in [49].
As a simple, ﬁrst example of the departure of the effective dynamics in the critical step-size regimefrom the classical perspective, we study estimation for spiked matrix and tensor models in Section 3.
In these models, the effective dynamics are exactly solvable and when the step-size scales criticallywith the dimension, in the ballistic phase the dynamics have additional ﬁxed points as compared tothe population gradient ﬂow.
The stability of these ﬁxed points exhibit sharp transitions at specialsignal-to-noise ratios.
When initialized randomly, the SGD starts in a microscopic neighborhood ofan uninformative such ﬁxed point, within which its effective dynamics become diffusive and exhibita sharp transition between stable and unstable Ornstein–Uhlenbeck processes.
To demonstrate our approach on more complex classiﬁcation tasks typically studied using neuralnetworks, we study a Gaussian mixture model analogue of the classical XOR problem in Section 5.(The XOR problem is arguably the canonical example of a decision boundary requiring at leasttwo-layers to represent [33].)
Here we ﬁnd that the natural summary statistics are 22 dimensional,and their (ballistic) effective dynamics exhibit a rich phenomenology between some 625 ﬁxed pointregions of varying topological dimension.
Surprisingly, we ﬁnd that if we initialize the weights of the network randomly (following a Gaussian distribution), then the algorithm will converge to a classiﬁerwith macroscopic generalization error with probability 29 / 32 and then follow a degenerate diffusion.
Before delving into the XOR problem, we ﬁrst analyze the classiﬁcation of a two component Gaussianmixture model in Section 4.
This task is of course best solved using a one-layer network i.e., logisticregression, but with a two-layer network it exhibits some similar phenomenologies to the XORproblem while being more amenable to ﬁner analysis.
Here, we again ﬁnd that if with randominitial weights, with probability 1 / 2 the SGD will ﬁrst converge to a classiﬁer with macroscopicgeneralization error, and then follow a degenerate diffusion in a microscopic neighborhood of that setof unstable ﬁxed points.
We demonstrate this both empirically for positive signal-to-noise ratio andtheoretically in the limit where the SNR tends to zero after the dimension tends to inﬁnity.
Main result
Suppose that we are given a sequence of i.i.d.
data Y 1 , Y 2 , . .
taking values in Y n ✓ R d n with law
P n 2 M 1 ( R d n ) , and a loss function L n : X n ⇥ Y n !
R , where here X n ✓ R p n is the parameterspace.
Consider online stochastic gradient descent with constant learning rate,  n , which is given by
[Equation]
with possibly random initialization X
⇠ µ n 2 M 1 ( X n ) .
Our interest is in understanding thisevolution, ( X ` ) , in the regime where both p n and d n ! 1 as n ! 1 .
To this end, suppose thatthere is a ﬁnite collection of summary statistics of ( X ` ) whose evolution we are interested in.
Moreprecisely, suppose that we are given a sequence of functions u n 2 C 1 ( R p n ; R k ) for some ﬁxed k,where u n ( x ) =
( u n 1 ( x ) , ..., u nk ( x )) , and our goal is to understand the evolution of u n ( X ` ) .
To develop a scaling limit, we need some assumptions on the relationship between how thestep-size scales in relation to the loss, its gradients, and the data distribution.
To this end let
H ( x, Y ) =
L n ( x, Y )   ( x ) , where  ( x ) =
E [ L n ( x, Y )] .
Throughout the following, we suppressthe dependence of H on Y and simply write H ( x ) , and instead view H as a random function of x .
Deﬁnition 2.1.
We say that a triple ( u n , L n , P n ) is  n -localizable if there is an exhaustion bycompact sets of R k , call it ( E K ) K , and constants 0 < C ( K ) <
1 (independent of n ) such that
max 1  i  k sup x 2 u  1 n ( E K ) ||
r j u
i || op  C ( K ) ·   (3  j ) / 2 n for j = 2 , 3 ;2.
sup x 2 u  1 n ( E K )
kr  k  C ( K ) , and sup x 2 u  1 n ( E K ) E
[ kr H k 3 ]  C ( K )   3 / 2 n ;
max 1  i  k sup x 2 u  1 n ( E K ) E [ hr H, r u
i
i 4 ]  C ( K )   2 n , and max 1  i  k sup x 2 u  1 n ( E K ) E [ hr 2 u i , r H ⌦ r H  V i 2 ]
= o (   1 n ) . When these hold we call the sequence ( E K ) the localizing sequence of ( u n , L n , P n ) .
Localizability is wider than uniform Lipchitz or smoothness assumptions common to the literature.
In particular, it does not imply that the population loss is Lipschitz everywhere, as we may have that
S
K u  1 n ( E K ) does not cover R p n , nor does it imply uniform smoothness of L as we will be taking n ! 0 with n .
To motivate the scaling relations between 2–3, note that if  n ⇣ p  1 n , correspondingto linear sample complexity, the scaling relation is the same as what one would get e.g., if r H werea random vector with independent entries of bounded variance in R p n and e.g., r u i =   1 / 2 n e 1 .
We now turn to the statement of our main result.
Let P k denote the space of positive semi-deﬁnitek ⇥ k matrices and for a function f and measure µ we let f ⇤ µ denote the push-forward of µ .
Let
J n = ( r u ` ) denote the Jacobian of the summary statistics.
Also, let V ( x ) =
E [ r H ( x ) ⌦ r H ( x )]denote the covariance matrix for r H at a point x and deﬁne the corresponding ﬁrst and secondorder differential operators, A n =
[Equation] and [Equation]
j .
We then have the followingconvergence result.
Theorem 2.2.
Let ( X  n ` ) ` be the SGD initialized from X 0 ⇠ µ n for µ n 2 M 1 ( R p n ) with learningrate  n for the loss L n ( · , · ) and data distribution P n .
Suppose that u n is such that the triple( u n , L n , P n ) is  n -localizable with localizing sequence ( E K ) .
Suppose furthermore that there existslocally Lipschitz h : R k !
R k and ⌃ : R k !
P k , such that for every K ,
[Equation]
where the suprema are over x 2 u  1 n ( E K ) . Then if we let ( u n ( t )) t be the linear interpolation of( u n ( X  n b t   1 n c )) t , and the initial data µ n , is such that ( u n ) ⇤ µ n !
⌫ weakly, then, ( u n ( t )) t !
( u t ) tweakly as n ! 1 , where u t is the solution to
[Equation]
initialized from ⌫ , where B t is a standard Brownian motion in R k .
Theorem 2.2 is proved in Appendix C. The proof can be seen as a version of the martingale problem(see [46]) for high-dimensional SGD.
2.1 Comparison to ﬁxed dimensional perspective: critical v.s. subcritical step-sizes
Let us compare this with the classical limit theory of SGD in ﬁxed dimension.
For the sake ofthis discussion, suppose that not only does (2.1) hold, but each of the two terms A n u and  n L n uindividually admit n !
1 limits: namely that there exists f , g : R k !
R k such thatsup [Equation] in which case, evidently (2.1) holds with h
=  f + g .
When (2.4) both hold, we call f , g and ⌃ thepopulation drift , the population corrector , and the diffusion matrix of u respectively.
From the ﬁxed dimensional perspective, when (2.4) holds, one predicts u to asymptotically solve
[Equation]
with initial data u 0 ⇠ u ⇤ µ .
as this is its evolution under gradient descent on the population loss .
Evidently this perspective only applies in the high-dimensional limit of Theorem 2.2 if both thepopulation corrector g and the diffusion matrix ⌃ are zero.
We ﬁnd that for any triple ( u n , L n , P n ) ,there is a scaling of the learning rate  n with n below which g = ⌃ = 0 , and the effective dynamicsagree with the population dynamics (2.5)—call this the sub-critical scaling regime, where theclassical perspective applies—and a critical scaling regime in which g and ⌃ may be non-zero, andthe high-dimensionality induces non-trivial corrections to f . (In the case of teacher–student networks,the terms f and g can be compared to the “learning" and “variance" terms in Eq.
(14a) of [49].)
To see this, notice that if the triple ( u n , L n , P n ) is  n -localizable for some  n ! 0 , then it is also 0 n -localizable for every sequence  0 n
= O (  n ) .
If furthermore (2.4)-(2.2) hold for  n with some f , gand ⌃ , then these limits also exists for  0 n = o (  n ) with the same f but with g = ⌃ = 0 .
As such,there can be exactly one scaling of  n with n at which g or ⌃ may be non-zero, and for all smallerscales of learning rate, the ﬁxed-dimensional perspective of (2.5) applies.
2.2 Ballistic vs. diffusive behavior of effective dynamics
In all of our examples, the diffusion matrix for the effective dynamics of the most natural choice ofsummary statistics is zero even in the critical scaling regime where h 6 = f .
We call this the ballisticlimit .
In this case, the effective dynamics of the summary statistics is given by the ODE system
[Equation]
In these settings, the phase portrait of the summary statistics is asymptotically that of this ﬂow.
By construction of the scaling limit, the phase portrait of the ballistic limit only describes the evolutionof summary statistics on length-scales that are order 1 and time-scales that are order 1 /  n .
If one isthen interested in the evolution of u n in microscopic o (1) neighborhoods of the ﬁxed points of (2.6), Theorem 2.2 also allows one to develop separate diffusive limits there.
To understand diffusive regimes, one must apply Theorem 2.2 to a re-centered and re-scaled versionof the summary parameters, ˜ u n ( t ) =
↵ n ( u n ( t )  u ? ) where u ? is a ﬁxed point of (2.6).
2 To apply
Theorem 2.2, ↵ must be chosen appropriately so that the triple (˜ u n ( t ) , L n , P n ) is  n -localizableand to pick out the next order drifts for ˜ u —the ﬁrst order term being zero microscopically close tou ? —and such that the initial data still converges (˜ u n ) ⇤ µ n !
˜ ⌫ .
This then leads to the rescaled effective dynamics of the summary statistics u n near u ? : [Equation] with
[Equation]
The rescaled effective dynamics are similar in spirit to diffusion approximations typically one ﬁndsfor the evolution of SGD near critical points in ﬁxed dimensions.
However, we note two importantdifferences as compared to this perspective.
Firstly, since this is a high-dimensional limit of generalsummary statistics, (2.7) applies in a neighborhood of a ﬁxed point of the effective ODE system (2.6),rather than the population dynamics (2.5).
Secondly, in many examples (indeed all the ones we study)the SDE’s we get are degenerate to some degree, so that uniform ellipticity assumptions typicallyused to understand hitting and mixing times in these regimes do not apply.
Matrix and Tensor PCA
As our ﬁrst example, we consider the problems of spiked matrix models and spiked tensor models[38] using SGD.
These examples are exactly solvable and only require two summary statistics, acorrelation observable and a radial term.
Even with this relative simplicity, we encounter a widerange of ODE and SDE limits.
Interestingly, by means of these SDE limits, we can sharply identifythe signal-to-noise thresholds for solving the recovery problem by means of the SGD.
Suppose that we are given data of the form Y =  v ⌦ k + W where W is an i.i.d.
Gaussian k -tensor,v 2 R n is a unit vector, and  =  n > 0 is the signal-to-noise ratio.
Our goal is infer v .
Wetake as loss the (negative) log-likelihood namely, L ( x, Y ) =
||
Y  x ⌦ k || 2 .
The pair of summarystatistics m = m ( x ) :
= h x, v i and r 2 ?
=
r 2 ?
( x ) :=
k x  mv k
=
k x k 2  m 2 are such that ( x ) =
2  m k + ( r 2 ?
+ m 2 )
k
+ c , and the law of L only depends on them: see Section D.1.
For the pair u n = ( u 1 , u 2 ) = ( m, r 2 ? ) ,
Theorem 2.2 yields the following effective dynamics.
In ournormalization with  > 0 ﬁxed, the regime  n = o (1 /n ) is sub-critical and the regime  n = ⇥ (1 /n )is critical; we focus on this normalization for presentation, but note that with different scalings of  n ,the critical learning rates change.
For notational simplicity, let R 2 :
= m 2 + r 2 ? .
Proposition 3.1.
Fix k  2 ,  > 0 , c  > 0 and let  n = c  / n .
Then u n ( t ) converges as n !
1 tothe solution of the following ODE initialized from lim n !1 ( u n ) ⇤ µ n :
[Equation]
We are able to identify and classify the set of ﬁxed points of this effective dynamics.
(Recall that thedynamics transits ballistically between these ﬁxed points in   1 n many steps.)
We focus on the criticalstep-size regime with c  = 1 where one sees from (3.1) that u 2 ! 1 , which is where a random vectorin R n lies, and where the problem in the matrix case is most directly related to an eigenvalue problem(see Appendix D for the generic c  dependencies).
Proposition 3.2.
Eq.
(3.1) has isolated ﬁxed points classiﬁed as follows.
Let  c ( k ) be as in (D.4)and m † ( k,  )  m ?
( k,  ) be as in (D.5)
(if k = 2 ,  c = 1 and m † = m ?
=
p   1 ):
An unstable ﬁxed point at (0 , 0) and a ﬁxed point at (0 , 1) ; if k = 2 , (0 , 1) is stable if <
c (2) and unstable if  >  c (2) ; if k > 2 (0 , 1) is always stable.
If  >  c ( k ) : when k = 2 , two stable ﬁxed points at ( ± m ?
When k  3 , twounstable ﬁxed points at ( ± m † ( k ) , 1) and two stable ﬁxed points at ( ± m ?
( k ) , 1) .
Remark 1 .
The presence of two pairs of ﬁxed points when k  3 with non-zero correlation with v mayseem surprising—indeed it indicates that even some warm starts will fail to attain good correlationwith the signal when  is ﬁnite.
This is an interesting consequence of the corrector in (3.1) and if onetracked the c  dependence in the above, the ﬁxed point m † goes to zero as c  !
and this barrier torecovery from warm starts vanishes as one approaches sub-critical step-sizes.
[Figure]
Let us now consider a rescaling of u n in a microscopic neighborhood of the saddle set m = 0 .
Thiscaptures the initial phase from a random start: if µ n ⇠ N (0 , I n /n ) , then ( u n ) ⇤ µ n !
(0 , 1) weakly.
Now rescale and let ˜ u n = ( p nm, r 2 ? ) .
Evidently, ˜ ⌫ = lim n (˜ u n ) ⇤ µ n = N (0 , 1) ⌦  1 .
Proposition 3.3.
Fix k  2 ,  > 0 and  n = 1 /n .
Then ˜ u n ( t ) converges as n !
1 to the solutionof the following SDE initialized from ˜ ⌫ :
[Equation]
We see that ˜ u
2 solves an autonomous ODE which converges exponentially to 1 .
When k = 2 , theequation for ˜ u
1 then converges to 4(   1)˜ u 1 dt + 2 p 2 dB t for large t .
This is an OU process whichis stable when  < 1 and unstable when  > 1 .
By stitching together the prelimits of these OUprocesses at a sequence of scales interpolating between that of ˜ u n and u n , one could in principleestablish that for any  > 1 , SGD reaches the stable ﬁxed points at ( ± m ?
(2) , 1) in O ( n log n ) steps(with precise asymptotics, etc.)
, while when  < 1 , the mean-reverting nature of the OU suggestsit needs a much larger number of samples in order to correlate with the vector v .
See Figure 1 fornumerical veriﬁcation of this intuition.
When k  3 , the tensor PCA problem is known to be hard for
SGD to solve without a polynomially diverging sample complexity or  [3].
Accordingly, when  iskept ﬁnite in n , the expression for ˜ u 1 in (3.2) is always a stable OU-type process.
Interestingly, onecan also capture the (diverging) signal-to-noise threshold for SGD to recover v in tensor PCA by ourmethods.
Indeed, for k  3 if one considers  n = ⇤ n ( k  2) / 2 (matching the predicted gradient-basedalgorithm threshold from [4]), ˜ u n would instead converge to the solution of
[Equation]
which transitions between stable and unstable OU processes at ⇤ c ( k ) = 1 , as in the matrix case.
Two-layer networks for classifying a binary Gaussian mixture
As our second example, we consider the problem of supervised classiﬁcation of a binary Gaussianmixture model (binary GMM) using a two-layer network.
Our goal here is to demonstrate how ourapproach can be used to analyze the performance of SGD for multi-layer networks, and indeed wewill ﬁnd the calculations here to be relevant in Section 5 where we consider XOR-type GMM’s.
Let us now formalize the problem.
Suppose that we are given i.i.d.
samples of the form Y =
( y, X ) ,where y is a Ber (1 / 2) random variables and, conditionally on y , we have X ⇠ N ((2 y  1) µ, I/  ) ,where µ 2 R N is a ﬁxed unit vector, I is the identity on R N , and  > 0 is the signal-to-noise ratio.
Here, y is the class label and X is the data.
For the sake of concreteness, we consider classiﬁcation via the following architecture (though ourtechniques generalize to other settings mutatis mutandis ): The ﬁrst layer has weights ( W 1 , W 2 )
R N ⇥ R N and ReLu activation, g ( x ) = x _ 0 ; and the second layer has weights v 1 , v 2 2 R andsigmoid activation,  ( x ) = 1 / (1+ e  x ) .
Our parameter space is then X n = R 2 N +
and we thereforetake n = 2 N + 2 when applying Theorem 2.2.
As we are interested in supervised classiﬁcation, wetake the usual binary cross-entropy loss with ` 2 regularization,
[Equation]
[Equation]
where g is applied component wise and p ( v, W ) :
= ( ↵ / 2)( || v ||
+ || W || 2 ) .
It can be shown (see Lemma E.1) that the law of the loss at a given point, ( v, W ) 2 X n , dependsonly on the 7 summary statistics,
[Equation]
where m i
= W i · µ and R ?
ij =
W ?
i · W ? j with W ?
i
= W i
m i µ denoting the part of W i orthogonalto µ .
For a point, ( v, W ) 2 X n , let
[Equation]
By Lemma E.1, these are functions only of u n , and we denote them as such, e.g., A µi =
A µi ( u n ) .
The critical scaling for  is then of order ⇥ (1 /n ) and we obtain the following effective dynamics.
Proposition 4.1.
Let u n be as in (4.2) and ﬁx any  > 0 and  n =
c  / N .
Then u n ( t ) converges tothe solution of the ODE system, ˙ u t
=  f ( u t ) + g ( u t ) , initialized from lim n !1 ( u n ) ⇤ µ n , with:
[Equation]
) + 2 ↵ R ?
ij , and correctors g v i =
g m i = 0 , g R ?
ij =
c  v i v j  B ij for i, j = 1 , 2 .
Due to the Gaussian integrals deﬁning f , g , it is difﬁcult to analyze the ODE system deﬁned by
Proposition 4.1, let alone any rescaled effective dynamics.
For ease of analysis, we next send  !
1corresponding to a small noise regime for the Gaussian mixture.
We emphasize that this limit is takenafter n ! 1 and therefore is still approximately on the critical scale of  = ⇥ (1) at which there is atransition in the existence of any ﬁxed point which is a good classiﬁer.
In particular, if  =  n is anydiverging sequence, then the limiting effective dynamics would exactly match that attained by nowsending  !
In Figure 2, we demonstrate numerically that the following predicted ﬁxed pointsfrom the  !
1 limit match those arising at ﬁnite large n and  > 0 .
Proposition 4.2.
The  !
1 limit of the ODE system of Proposition 4.1 is given by
[Equation]
else
[Equation]
and ˙ R ?
ij =
2 ↵ R ?
ij .
The ﬁxed points of this system are classiﬁed as follows.
All ﬁxed points have
R ? ij
=
and m i
= v i for i,
j = { 1 , 2 } .
In ( v 1 , v 2 ) , the coordinates are classiﬁed by1.
A ﬁxed point at ( v 1 , v 2 ) =
(0 , 0) that is stable if ↵ > 1 / 42.
If ↵ < 1 / 4 , two unstable sets of ﬁxed points at the quarter-circles given by ( v 1 , v 2 ) havingv 1 v 2 > 0 such that v 21 + v 22 = C ↵ for C ↵ := log(1  2 ↵ )  log(
↵
If ↵ < 1 / 4 , two stable ﬁxed points at ( v 1 , v 2 ) equals ( p C ↵ , p C ↵ ) and ( p C ↵ , p C ↵ ) .
If µ n is e.g., given by ( v 1 , v 2 )
⇠ N (0 , I 2 ) and W 1 , W 2 ⇠ N (0 , I N / (  N ))
then ⌫ :=
lim( u n ) ⇤ µ nis N (0 , I 2 ) in the v 1 , v 2 coordinates, and is in the basin of attraction of the quarter-circles of item (2)with probability 1 / 2 and the basin of attraction of the stable ﬁxed points of (3) with probability 1 / 2 .
Let us pause to interpret this result.
The stable ﬁxed points when ↵ < 1 / 4 are the optimal classiﬁers,whereas the unstable set of ﬁxed points given by item (2) misclassify half of the data.
Therefore,the above indicates that when solving the above task with randomly initialized weights, one of thefollowing two scenarios occur, each with probability 1 / 2 (w.r.t.
the initialization): the algorithm willconverge to the optimal classiﬁer in linear time or it will appear to have converged to a macroscopicallysub-optimal classiﬁer on the same timescale, see Figure 2(a)–(b).
It is then natural to ask about the behaviour of the SGD in the latter regime, after it converges to thesub-optimal classiﬁers which lie on the aforementioned quarter-circles.
Proposition 4.2 rigorously
Lipschitz
[Figure]
justiﬁed the exchange of n ! 1 and  !
1 limits in the ballistic phase.
In the diffusive phase, onecould in principle ﬁnd the quarter circle of ﬁxed points of the ODE in Proposition 4.1 and considerrescaled observables ˜ v i , ˜ m i corresponding to blowing up v i , m i in diffusive O ( n  1 / 2 ) neighborhoodsabout them to get SDE limits from Theorem 2.2.
In order to have explicit formulae, in what follows,we consider the diffusive limits obtained when taking  = 1 , for which we know the preciselocations of these ﬁxed points from Proposition 4.2.
This also captures the limit obtained by takingany  n diverging faster than O ( n 1 / 2 ) ; the numerics of Figure 2(c)–(d) demonstrate its qualitativeconsistency with the behavior in microscopic neighborhoods of ﬁxed points even at  ﬁnite.
Proposition 4.3.
Let  n = 1 / N , ( a 1 , a 2 ) 2 R 2+ be such that a 21 + a 22 = C ↵ and let ˜ v i
= p N ( v i  a i )and ˜ m i = p N ( m i  a i ) .
When  = 1 , the SDE system obtained by applying Theorem 2.2 to ˜ u n is
[Equation]
where ˜ ⌃ is a constant matrix whose only non-zero entries are ˜ ⌃ ˜ v i ˜ v j = ˜ ⌃
˜ m i ˜ m [Equation] = ˜
⌃ ˜ v i ˜ m j = ↵ 2 a i a j .
Notice that this diffusion matrix is rank 1, so this diffusion is non-trivial but degenerate even in therescaled coordinates (˜ v i , ˜ m i ) .
Moreover, the entries of ˜ ⌃ vanish on the axes a 1 = 0 or a 2 = 0 .
Inparticular, crossing from the unstable quarter ring into the quadrants v 1 v 2 < 0 where the stable ﬁxedpoints lie is impossible in the noiseless setting, and happens on a much larger timescale at ﬁnite  .
Two-layer networks for the XOR Gaussian mixture
We end by applying our methods to supervised learning for an XOR-type Gaussian mixture model in R N with a two-layer network. The ballistic phase of this problem was recently studied in [37].
Suppose that we are given i.i.d.
samples of the form Y =
( y, X ) , where y is Ber (1 / 2) as before and
X has the following distribution: if y = 1 then X is a 1 / 2 - 1 / 2 mixture of N ( µ, I/  ) and N (  µ, I/  )and if y =
it is a 1 / 2 - 1 / 2 mixture of N ( ⌫ , I/  ) and N (  ⌫ , I/  ) , where  > 0 , and µ, ⌫ areorthogonal unit vectors.
Here, y is the class membership label and X is the data.
Consider the corresponding classiﬁcation problem using a two-layer neural network, ˆ y ( X ) =
( v · g ( WX )) , where  and g are as in Section 4.
We take W to be a 4 ⇥ N matrix and v to be a4 -vector, and consider the binary cross-entropy loss as before.
It is shown in Lemma F.1 that the lawof the loss at a point ( v, W ) depends only on the following 22 variables: for 1  i  j  4 ,
[Equation]
i
[Equation]
where W ?
i
= W i  m µi µ  m ⌫
i ⌫ is the part perpendicular to µ, ⌫ .
With the choice of u n given bythese variables, for any ﬁxed  > 0 , the localizability criterion of Deﬁnition 2.1 can be veriﬁed tohold as long as  n = O (1 /n ) .
In particular, we can apply Theorem 2.2 to obtain limits in both theballistic and diffusive phases.
For the precise equations in the ballistic phase, see Proposition F.2.
The ﬁxed points of the ballistic dynamics, again in the limit  ! 1 after n ! 1 , are classiﬁed asfollows (see Proposition F.3).
If ↵ > 1 / 8 , then the only ﬁxed point is at u n = 0 .
If 0 < ↵ < 1 / 8 , then [Figure] let ( I 0 , I + µ , I  µ , I + ⌫ , I  ⌫ ) be any disjoint (possibly empty) subsets whose union is { 1 , . . .
Eachsuch partition corresponds to a connected component of ﬁxed points.
Corresponding to a such tuple,the connected component of ﬁxed points has R ?
ij =
for all i, j , and
1. m µi
= m ⌫ i =
v i =
for i 2 I 0 ,2.
m µ i = v i >
such that P
i
I + µ v
i =  logit (4 ↵ ) and m ⌫ i = 0 for all i 2 I + µ ,3.
m µi
= v i > 0 such that P
i
I  µ v
i =  logit (4 ↵ ) and m ⌫ i = 0 for all i
I  µ ,4.
m ⌫ i
=
v i < 0 such that P
i
I + ⌫
v
i
=
logit (4 ↵ ) and m µi =
for all i 2 I + ⌫ ,5.
m ⌫ i
=
v i < 0 such that P
i
I  ⌫ v 2 i
=  logit (4 ↵ ) and m µi =
for all i 2 I  ⌫ .
There are 5 4 = 625 connected components of ﬁxed points, 4! of which are stable, one for eachpermutation where I + µ , I  µ , I + ⌫ , I  ⌫ are all singletons.
For the proof of this limit, see Proposition F.3.
We can also compute the probability that the effective dynamics in the ballistic phase convergesto a stable ﬁxed point (as opposed to an unstable one).
From a Gaussian initialization µ n wherev i ⇠ N (0 , 1) and W i ⇠ N (0 , I N /N )
independently, this probability will converge to 3 / 32 .
As an example of the diffusions that can arise in the rescaled effective dynamics at the unstable ﬁxedpoints, let us consider the unstable ﬁxed points in which v has the correct signature (two positive, twonegative) but for each of those we are at a corresponding quarter-ring.
Here, the dynamics effectivelybecomes a pair of 2 two-layer GMM’s on quarter-rings (as in Section 4), that are anti-correlated.
Moreprecisely, let ( a 1 ,µ , a 2 ,µ ) be such that a 21 ,µ + a 22 ,µ = C ↵ and ( a 3 , ⌫ , a 4 , ⌫ ) such that a 23 , ⌫ + a 24 , ⌫ = C ↵ ,for C ↵ =
logit (4 ↵ ) .
Take as ﬁxed points about which we expand to be v i = m µi = a i,µ > 0 andv i
= m ⌫ i = a i, ⌫ < 0 for i = 3 , 4 .
Namely, we let
[Equation]
(We set ˜ m ⌫ i =
for i = 1 , 2 and ˜ m µi =
for i = 3 , 4 in ˜ u n effectively removing those variables.)
Proposition 5.1.
Let  n = 1 / N and let ˜ u n = (˜ v i , ˜ m µi , ˜ m ⌫ i , R ? ij ) .
When  = 1 , Theorem 2.2 can beapplied and ˜ u n ( t ) converges to the solution of the SDE d ˜u ( t ) =
˜ h (˜ u ) dt + p ⌃
(˜ u )
d B t where
[Equation]
˜ h ˜ m µi (resp., ˜ h ˜ m ⌫ i ) is like h ˜ v i for i = 1 , 2 (resp.
, i = 3 , 4 ) with ˜ v i and ˜ m µi (resp.,
˜ m ⌫ i ) swapped,˜ [Equation] R ?
[Equation]
= 2 ↵ R ?
ij , and ˜ ⌃ is the constant rank-2 matrix whose non-zero entries are
[Equation]