High-dimensional limit theorems for SGD: Effective dynamics and critical scaling
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 trajectories of summary statistics (i.e., ﬁnite-dimensional functions) of SGD as the dimension goes to inﬁnity.
Our approach allows one to choose the summary statistics that are tracked, the initialization, and the step-size.
It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices.
We ﬁnd a critical scaling regime for the step-size below which this “effective dynamics" matches gradient ﬂow for the population loss, but at which, a new correction term appears which changes the phase diagram.
About the ﬁxed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate.
We demonstrate our approach on popular examples including estimation for spiked matrix and tensor models and classiﬁcation via two-layer networks for binary and XOR-type Gaussian mixture models.
These examples exhibit surprising phenomena including multimodal timescales to convergence as well as convergence to sub-optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations.
Introduction
Background.
Stochastic gradient descent (SGD) is the go-to method for large-scale optimization problems 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, is by now classical.
There have been works on path-wise limit theorems, such as functional central limit theorems and even large deviations principles [39, 31, 28, 22, 17, 7, 16, 6].
At the core of this line of work 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 interest in quantifying the rate of this trajectory-wise convergence to higher order, in terms of a diffusion approximation.
Namely, there are many works developing asymptotic expansions of the trajectory in 36th Conference on Neural Information Processing Systems (NeurIPS the learning rate [29, 23, 25, 1, 26].
Motivated by this, there is a rich line of work bounding the time to 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 work obtaining PDE limits in the “shallow network” regime where the dimension of the parameter space diverges 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 highdimensional nature of the data and the complexity of the model being trained.
In these regimes, one cannot simply take the learning rate to be arbitrarily small as this would force an unlimited sample size and run-time.
This is a common issue in high-dimensional statistics and the standard analytic approach 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 dimension dependence 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 been important work on scaling limits as the dimension tends to inﬁnity for the speciﬁc problems of linear regression [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 warm starts.
We also note that the study of high-dimensional regimes of gradient descent and Langevin dynamics 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-dimensions with constant learning rate that allows us to understand a broad range of estimation tasks.
One of course cannot develop a high-dimensional scaling limit for the full trajectory of SGD as the dimension of the underlying parameter space is growing.
On the other hand, in practice, one is rarely interested in the full trajectory; instead one typically tracks the trajectory of various summary statistics of the algorithm’s evolution, such as the loss, the amplitude of various weights, or correlations between the classiﬁer and the ground truth (in a supervised setting).
We show in Theorem 2.2 that under mild regularity assumptions, the evolution of these summary statistics converges as the dimension grows to the solution of a system of (possibly stochastic) differential equations.
These effective dynamics depend dramatically on the initializations (warm vs. random or cold), the parameter regions in which one 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 summary statistics 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 between these 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) and the ﬁnite-dimensional intuition that the summary statistics evolve under the gradient ﬂow for the population 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 the high-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, in microscopic neighborhoods of the ﬁxed points of this ODE, the effective dynamics become diffusive and are given by SDEs which can exhibit a wide range of (possibly degenerate) behaviors.
We note that the appearance of the correction term in the ballistic phase was ﬁrst observed in the setting of teacher-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 regime from 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 critically with the dimension, in the ballistic phase the dynamics have additional ﬁxed points as compared to the population gradient ﬂow.
The stability of these ﬁxed points exhibit sharp transitions at special signal-to-noise ratios.
When initialized randomly, the SGD starts in a microscopic neighborhood of an uninformative such ﬁxed point, within which its effective dynamics become diffusive and exhibit a sharp transition between stable and unstable Ornstein–Uhlenbeck processes.
To demonstrate our approach on more complex classiﬁcation tasks typically studied using neural networks, 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 least two-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 39 ﬁxed point regions 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ﬁer with 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 Gaussian mixture model in Section 4.
This task is of course best solved using a one-layer network i.e., logistic regression, but with a two-layer network it exhibits some similar phenomenologies to the XOR problem while being more amenable to ﬁner analysis.
Here, we again ﬁnd that if with random initial weights, with probability 1 / 2 the SGD will ﬁrst converge to a classiﬁer with macroscopic generalization error, and then follow a degenerate diffusion in a microscopic neighborhood of that set of unstable ﬁxed points.
We demonstrate this both empirically for positive signal-to-noise ratio and theoretically 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 parameter space.
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 this evolution, ( X ` ) , in the regime where both p n and d n ! 1 as n ! 1 .
To this end, suppose that there is a ﬁnite collection of summary statistics of ( X ` ) whose evolution we are interested in.
More precisely, 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 the step-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 suppress the 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 by compact 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 3.
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 , corresponding to linear sample complexity, the scaling relation is the same as what one would get e.g., if r H were a 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ﬁnite k ⇥ 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 second-order differential operators,
[Equation]
We then have the following convergence result.
Theorem 2.2.
Let ( X  n ` ) ` be the SGD initialized from X 0 ⇠ µ n for µ n 2 M 1 ( R p n ) with learning rate  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 exists locally 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 ) t weakly 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 of this discussion, suppose that not only does (2.1) hold, but each of the two terms A n u and  n L n u individually admit n !
1 limits: namely that there exists f , g : R k !
R k
such that sup
[Equation]
in which case, evidently (2.1) holds with h
=  f + g .
When (2.4) both hold, we call f , g and ⌃ the population 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 the population 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 dynamics agree with the population dynamics (2.5)—call this the sub-critical scaling regime, where the classical perspective applies—and a critical scaling regime in which g and ⌃ may be non-zero, and the 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 , g and ⌃ , 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 smaller scales 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 of summary statistics is zero even in the critical scaling regime where h 6 = f .
We call this the ballistic limit .
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 evolution of summary statistics on length-scales that are order 1 and time-scales that are order 1 /  n .
If one is then 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 version of 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 -localizable and to pick out the next order drifts for ˜ u —the ﬁrst order term being zero microscopically close to u ? —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 ﬁnds for the evolution of SGD near critical points in ﬁxed dimensions.
However, we note two important differences as compared to this perspective.
Firstly, since this is a high-dimensional limit of general summary 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 typically used 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, a correlation observable and a radial term.
Even with this relative simplicity, we encounter a wide range of ODE and SDE limits.
Interestingly, by means of these SDE limits, we can sharply identify the 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 .
We take as loss the (negative) log-likelihood namely, L ( x, Y ) =
||
Y  x ⌦ k || 2 .
The pair of summary statistics 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 our normalization 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 to the 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 the dynamics transits ballistically between these ﬁxed points in   1 n many steps.)
We focus on the critical step-size regime with c  = 1 where one sees from (3.1) that u 2 ! 1 , which is where a random vector in 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.
[Figure]
If  >  c ( k ) : when k = 2 , two stable ﬁxed points at ( ± m ?
When k  3 , two unstable points at ( m
Remark 1 .
The presence of two pairs of ﬁxed points when k  3 with non-zero correlation with v may seem surprising—indeed it indicates that even some warm starts will fail to attain good correlation with the signal when  is ﬁnite.
This is an interesting consequence of the corrector in (3.1) and if one tracked the c  dependence in the above, the ﬁxed point m † goes to zero as c  !
and this barrier to recovery from warm starts vanishes as one approaches sub-critical step-sizes.
Let us now consider a rescaling of u n in a microscopic neighborhood of the saddle set m = 0 .
This captures 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 solution of the following SDE initialized from ˜ ⌫ :
[Equation]
We see that ˜ u
2 solves an autonomous ODE which converges exponentially to 1 .
When k = 2 , the equation for ˜ u
1 then converges to 4(   1)˜ u 1 dt + 2 p 2 dB t for large t .
This is an OU process which is stable when  < 1 and unstable when  > 1 .
By stitching together the prelimits of these OU processes at a sequence of scales interpolating between that of ˜ u n and u n , one could in principle establish 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 suggests it needs a much larger number of samples in order to correlate with the vector v .
See Figure 1 for numerical 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  is kept ﬁnite in n , the expression for ˜ u 1 in (3.2) is always a stable OU-type process.
Interestingly, one can also capture the (diverging) signal-to-noise threshold for SGD to recover v in tensor PCA by our methods.
Indeed, for k  3 if one considers  n = ⇤ n ( k  2) / 2 (matching the predicted gradient-based algorithm 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 Gaussian mixture model (binary GMM) using a two-layer network.
Our goal here is to demonstrate how our approach can be used to analyze the performance of SGD for multi-layer networks, and indeed we will ﬁ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 our techniques 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 and sigmoid activation,  ( x ) = 1 / (1+ e  x ) .
Our parameter space is then X n = R 2 N +
and we therefore take n = 2 N + 2 when applying Theorem 2.2.
As we are interested in supervised classiﬁcation, we take the usual binary cross-entropy loss with ` 2 regularization,
[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 , depends only 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 orthogonal to µ .
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 to the solution of the ODE system, ˙ u t =
f ( u t ) + g ( u t ) , initialized from lim n !1 ( u n ) ⇤ µ n , with:
[Equation]
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  !
corresponding to a small noise regime for the Gaussian mixture.
We emphasize that this limit is taken after n !
and therefore is still approximately on the critical scale of  = ⇥ (1) at which there is a transition in the existence of any ﬁxed point which is a good classiﬁer.
In particular, if  =  n is any diverging sequence, then the limiting effective dynamics would exactly match that attained by now sending  !
In Figure 2, we demonstrate numerically that the following predicted ﬁxed points from 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 by 1.
A ﬁxed point at ( v 1 , v 2 ) =
(0 , 0) that is stable if ↵ > 1 / 4 .
[Figure]
If ↵ < 1 / 4 , two unstable sets of ﬁxed points at the quarter-circles given by ( v 1 , v 2 ) having v 1 v 2 > 0 such that v 21 + v 22 = C ↵ for C ↵ := log(1  2 ↵ )  log(2 ↵ ) .
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 ) ⇤ µ n is 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 the following two scenarios occur, each with probability 1 / 2 (w.r.t.
the initialization): the algorithm will converge to the optimal classiﬁer in linear time or it will appear to have converged to a macroscopically sub-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 the sub-optimal classiﬁers which lie on the aforementioned quarter-circles.
Proposition 4.2 rigorously justiﬁed the exchange of n ! 1 and  !
1 limits in the ballistic phase.
In the diffusive phase, one could in principle ﬁnd the quarter circle of ﬁxed points of the ODE in Proposition 4.1 and consider rescaled observables ˜ v i , ˜ m i corresponding to blowing up v i , m i in diffusive O ( n  1 / 2 ) neighborhoods about 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 precise locations of these ﬁxed points from Proposition 4.2.
This also captures the limit obtained by taking any  n diverging faster than O ( n 1 / 2 ) ; the numerics of Figure 2(c)–(d) demonstrate its qualitative consistency 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 j = ˜
⌃ ˜ 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 the rescaled coordinates (˜ v i , ˜ m i ) .
Moreover, the entries of ˜ ⌃ vanish on the axes a 1 = 0 or a 2 = 0 .
In particular, crossing from the unstable quarter ring into the quadrants v 1 v 2 < 0 where the stable ﬁxed points 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].
[Figure]
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 µ, ⌫ are orthogonal 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 a 4 -vector, and consider the binary cross-entropy loss as before.
It is shown in Lemma F.1 that the law of the loss at a point ( v, W ) depends only on the following 22 variables: for 1  i  j  4 , v i ,
[Equation]
where W ?
i
= W i  m µi µ  m ⌫
i ⌫ is the part perpendicular to µ, ⌫ .
With the choice of u n given by these variables, for any ﬁxed  > 0 , the localizability criterion of Deﬁnition 2.1 can be veriﬁed to hold as long as  n = O (1 /n ) .
In particular, we can apply Theorem 2.2 to obtain limits in both the ballistic 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 as follows (see Proposition F.3).
If ↵ > 1 / 8 , then the only ﬁxed point is at u n = 0 .
If 0 < ↵ < 1 / 8 , then let ( I 0 , I + µ , I  µ , I + ⌫ , I  ⌫ ) be any disjoint (possibly empty) subsets whose union is { 1 , . .
Each such 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 39 connected components of ﬁxed points, 4! of which are stable, one for each permutation 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 converges to a stable ﬁxed point (as opposed to an unstable one).
From a Gaussian initialization µ n where v 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 ﬁxed points, let us consider the unstable ﬁxed points in which v has the correct signature (two positive, two negative) but for each of those we are at a corresponding quarter-ring.
Here, the dynamics effectively becomes a pair of 2 two-layer GMM’s on quarter-rings (as in Section 4), that are anti-correlated.
More precisely, 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,µ >
and v i
= m ⌫ i = a i, ⌫ < 0 for i = 3 , 4 .
Namely, we let
[Equation]
and 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 be applied 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, ˜ h R ?
ij = 2 ↵ R ?
ij , and ˜ ⌃ is the constant rank-2 matrix whose non-zero entries are
[Equation]
m µi ˜ m ⌫ j = ˜
⌃ m µi ,v j = ˜
⌃ ˜ v i ˜ m ⌫ j =  ↵ 2 a i,µ a j, ⌫ if i 2 { 1 , 2 } , j 2 { 3 , 4 } .
Acknowledgements
The authors thank the anonymous referees for their careful readings of the manuscript and helpful comments.
R.G. thanks the Miller Institute for Basic Research in Science for its support.
A
J. acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC). Cette recherche a été ﬁnancée par le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG), [RGPIN-2020-04597, DGECR-2020-00199]