extfieldsofstudy:
Computer Science
Mathematics

text:
High-dimensional limit theorems for SGD: Effective dynamics and critical scaling

We study the scaling limits of stochastic gradient descent (SGD) with constant step-size in the high-dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite-dimensional functions) of SGD as the dimension goes to infinity. 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 show a critical scaling regime for the step-size, below which the effective ballistic dynamics matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed 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 classification 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. At the same time, we demonstrate the benefit of overparametrization by showing that the latter probability goes to zero as the second layer width grows.

Introduction
Stochastic gradient descent (SGD) is the go-to method for large-scale optimization problems in modern data science. It is often used to train complex parametric models on high-dimensional data. Since its introduction in [40], there has been a tremendous amount of work in analyzing its evolution.
In fixed 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 [40,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 fixed loss function (appropriately rescaled in time) converges to the solution of gradient flow 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 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 modifications) under uniform ellipticity assumptions [29,37,11,55]. 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,41,12,47,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 [52,53].
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,46,34,21,15] and special classes of non-convex objectives [19,49,3]. There has also been important work on scaling limits as the dimension tends to infinity for the specific problems of linear regression [54,36], Online PCA [54,24], and phase retrieval [49] from random starts, and teacherstudent networks [42,43,20,51] and two-layer networks for XOR Gaussian mixtures [38] 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,45,30,10,27].
We develop a unified 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 classifier 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 fluctuate 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 finite-dimensional intuition that the summary statistics evolve under the gradient flow for the population loss is correct provided the (constant) learning rate is sufficiently 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 significantly from that of the population gradient flow. Furthermore, in microscopic neighborhoods of the fixed 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 first observed in the setting of teacher-student networks in [42,43] and very recently investigated in detail in [51].
As a simple, first 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 fixed points as compared to the population gradient flow. The stability of these fixed points exhibit sharp transitions at special signal-to-noise ratios. When initialized randomly, the SGD starts in a microscopic neighborhood of an uninformative such fixed 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 classification 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 find that the natural summary statistics are 22 dimensional, and their (ballistic) effective dynamics exhibit a rich phenomenology between some 39 connected fixed point regions of varying topological dimension. Surprisingly, we find that if we initialize the weights of the network randomly (following a Gaussian distribution), then the algorithm will converge to a classifier with macroscopic generalization error with probability 29 /32 and then follow a degenerate diffusion. On the other hand, we demonstrate the benefit of overparametrization, showing that as the width of the second layer grows, the probability of ballistically converging to a Bayes optimal classifier goes to 1.
Before delving into the XOR problem, we first analyze the classification of a two component Gaussian mixture model in Section 8. 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 finer analysis. Here, we again find that if with random initial weights, with probability 1/2 the SGD will first converge to a classifier with macroscopic generalization error, and then follow a degenerate diffusion in a microscopic neighborhood of that set of unstable fixed 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 infinity.
While the above are a few examples that we are able to solve in detail for both their ballistic and diffusive limits, we expect our main theorem to be applicable and lend new insights into a host of other problems including SGD for finite-rank matrix and tensor PCA, and one and two-layer neural networks applied to mixtures of k-Gaussians for fixed k ≥ 2.

Main result
Suppose that we are given a sequence of i.i.d. data Y 1 , Y 2 , . . . taking values in Y n ⊆ R dn with law P n ∈ M 1 (R dn ), and a loss function L n : X n × Y n → R, where here X n ⊆ R pn is the parameter space. Consider online stochastic gradient descent with constant learning rate, δ n , which is given by with possibly random initialization X 0 ∼ µ n ∈ M 1 (X n ). Our interest is in understanding this evolution, (X ), in the regime where both p n and d n → ∞ as n → ∞. To this end, suppose that there is a finite 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 ∈ C 1 (R pn ; R k ) for some fixed k, where u n (x) = (u n 1 (x), ..., u n k (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 . Throughout the following, we suppress the dependence of H on Y and instead view H as a random function of x, denoted H(x). We let V (x) = E [∇H(x) ⊗ ∇H(x)] denote the covariance matrix for ∇H at x. Definition 2.1. A triple (u n , L n , P n ) is δ n -localizable if there is an exhaustion by compact sets (E K ) K of R k , and constants 0 < C(K) < ∞ (independent of n) such that When these hold we call the sequence (E K ) the localizing sequence of (u n , L n , P n ) We will motivate the localizability condition shortly, but let us first present our main result. Let P k denote the space of positive semi-definite k × k matrices and for a function f and measure µ we let f * µ denote the push-forward of µ. Let J n = (∇u ) denote the Jacobian of the summary statistics. Define the corresponding first and second-order differential operators, Alternatively written, A n = ∇Φ, ∇ and L n = 1 2 V, ∇ 2 . Our main result is the following.
Theorem 2.2. Let (X δn ) be the SGD initialized from X 0 ∼ µ n for µ n ∈ M 1 (R pn ) 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, Then if we let (u n (t)) t be the linear interpolation of (u n (X δn tδ −1 n )) t , and the initial data µ n , is such that (u n ) * µ n → ν weakly, then, (u n (t)) t → (u t ) t weakly as n → ∞, where u t solves The proof of Theorem 2.2 is provided in Appendix 6 and can be seen as a version of the classical martingale problem (see [48]) for high-dimensional stochastic gradient descent.
We end this subsection with a discussion of the various scalings appearing in Definition 2.1.
Remark 1. The kinds of summary statistics that we most frequently have in mind for application are (1) linear functions of the parameter space X n , for instance the correlation with a unit vector, or some ground truth; (2) radial statistics, like the 2 -norm of the parameters, or some subset of the parameters; and (3) rescaled versions (usually blown up by δ −1/2 n of these near their fixed points, as described in Section 2.2. Regarding the item (1) in Definition 2.1, for linear functions, it trivially holds; for radial statistics, the Hessian is a block identity matrix, so item (1) holds as long as δ n is O(1); therefore item (1) is most restrictive for rescalings of non-linear statistics, e.g., where it prevents consideration of this statistic with α > 1/2. Turning to item (2) of Definition 2.1, we comment that the regularity assumptions made on L here are less restrictive than uniform Lipchitz or smoothness assumptions common to the literature. In particular, we do not assume the population loss is Lipschitz everywhere, as we may have that K u −1 n (E K ) does not cover X n , nor does it imply uniform smoothness of H (and in turn L) as we will be taking δ n → 0 with n.
Let us lastly motivate the scalings appearing in item (3), which ensure there is some independence between H and the values of ∇u and ∇ 2 u at x. As a testbed, suppose that ∇H(x) is a random vector with i.i.d. entries all of order 1. Then, if u is a rescaled linear statistic, e.g., δ −1/2 n x, e 1 then the first bound of item (3) is saturated, and the second of course is trivial due to the linearity of u. The second bound is saturated by taking a rescaling of a radial statistic, e.g., δ −1/2 n x 2 , again assuming for maximal simplicity that ∇H is an i.i.d. random vector with order one entries. In fact, the second part of item (3) could be dropped at the expense of more complicated diffusion coefficients in limiting SDE's: see Remark 2.
Remark 2. While we discussed above the reasons for which the various scalings of Definition 2.1 were selected, it is interesting to ask what changes in Theorem 2.2 should certain of the assumptions of Definition 2.1 be violated. Most of the assumed bounds in the definition of localizability are used to establish tightness and ensure higher order terms in Taylor expansions vanish in the n → ∞ limit. One assumption which could in principle be relaxed is the second assumption in item (3) of Definition 2.1; indeed if that assumption is dropped, Theorem 2.2 would still essentially apply, except the limiting diffusion matrix would be the n → ∞ limit (assuming it exists) of as opposed to simply the limit of δJV J T .
The choice of summary statistics to which to apply Theorem 2.2 depends both on the quantities one is interested in, and the specifics of the task. In our examples, the choices are natural: correlations with ground truth vectors, finite numbers of final layer weights, and 2 norms of the parameters. In less structured settings, the choice of summary statistics may be more open-ended. One could start with a summary statistic of interest, like the projection in a principal subspace of an empirical matrix (a covariance or, as suggested experimentally in e.g., [44,35], a Hessian), or the population loss itself. Then from that statistic, one would determine the other statistics needed to close (or bound) (2.2)-(2.3) in the n → ∞ limit.
2.1. Comparison to fixed dimensional perspective: critical v.s. subcritical step-sizes. Let us compare this with the classical limit theory of SGD in fixed dimension. For the sake of this discussion, suppose that not only does (2.2) hold, but each of the two terms A n u and δ n L n u (recall (2.1)) individually admit n → ∞ limits: namely that there exists f , g : in which case, evidently (2.2) holds with h = −f + g. When (2.5) and (2.6) both hold, we call f , g and Σ the population drift, the population corrector, and the diffusion matrix of u respectively. From the fixed dimensional perspective, when (2.5) holds, one predicts u to solve 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 find 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.7) (we 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 [51].) 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 δ n -localizable for every sequence δ n = O(δ n ). If furthermore (2.3) and (2.5)-(2.6) hold for δ n with some f , g and Σ, then these limits also exists for δ 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 δ n , the fixed-dimensional perspective of (2.7) applies. 1 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 = f . We call this the ballistic limit. In this case, the effective dynamics of the summary statistics is given by the ODE system (2.8) In these settings, the phase portrait of the summary statistics is asymptotically that of this flow. 1 Note that if δn = o(δ n ), then limiting g, Σ may not exist for δ n , so there is no super-critical regime.
Note that 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 fixed points of the ballistic effective dynamics of (2.8), Theorem 2.2 also allows one to develop separate diffusive limits there.
To study diffusive regimes, one must apply Theorem 2.2 to re-centered and re-scaled summary statistics,ũ n (t) = δ −α n (u n (t) − u ) where u is a fixed point of (2.8). 2 To apply Theorem 2.2, α must be chosen appropriately so that the triple (ũ n (t), L n , P n ) is δ nlocalizable and to pick out the next order drifts forũ-the first order term being zero microscopically close to u -and such that the initial data still converges (ũ n ) * µ n →ν.
This then leads to the rescaled effective dynamics of the summary statistics u n near u : The rescaled effective dynamics are similar in spirit to diffusion approximations typically one finds for the evolution of SGD near critical points in fixed 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.9) applies in a neighborhood of a fixed point of the effective ODE system (2.8), rather than the population dynamics (2.7). Secondly, in many examples (indeed all the ones we study) the SDE's we get are degenerate, so that uniform ellipticity assumptions typically used to understand hitting and mixing times in these regimes do not apply. The degeneracies can take various forms, withΣ sometimes being rank deficient in the entireũ-space, and sometimes vanishing completely asΣ approaches certain distinguished points, for instance u . The implications of such degeneracies can be severe, as degenerate diffusions can be absorbed for arbitrarily long times by their unstable fixed points (c.f. the simple case of a 1D geometric Brownian motion).

Matrix and Tensor PCA
As our first example, we consider the problems of spiked matrix models and spiked tensor models [39] 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 with SGD.
Suppose we are given data Y = λv ⊗k + W where W is an i.i.d. Gaussian k-tensor, v ∈ R n is a unit vector, and λ = λ n > 0 is the signal-to-noise ratio. Our goal is to infer v.
We take as loss the (negative) log-likelihood namely, ⊥ + m 2 ) k + c, and the law of L only depends on them. 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 fixed, 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. We also present our results in the most interesting critical regime of δ n = c δ /n for some constant c δ , the limits in the subcritical regime are in all cases recovered by taking the c δ ↓ 0 limits of the ODE's/SDE's of the critical regime. For notational simplicity, let R 2 := m 2 + r 2 ⊥ .
2 One might also wish to rescale time like δ −β n , where β may depend on t; we leave this to future work. 3 Note that one might also add additional penalty terms. The case of a ridge penalty is treated in Section 7. Proposition 3.1. Fix k ≥ 2, λ > 0, c δ > 0 and let δ n = cδ /n. Then u n (t) converges as n → ∞ to the solution of the following ODE initialized from lim n→∞ (u n ) * µ n : We are able to identify and classify the set of fixed points of this effective dynamics. (Recall that the dynamics transits ballistically between these fixed points in δ −1 n many steps.) We focus on the critical step-size regime with c δ = 1 where one sees from (3.1) that r 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 7 for the generic c δ dependencies).
Remark 3. The presence of two pairs of fixed 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 finite. This is an interesting consequence of the corrector in (3.1) and if one tracked the c δ dependence in the above, the fixed point m † goes to zero as c δ → 0 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ũ n = (m, r 2 ⊥ ) = ( √ nm, r 2 ⊥ ). Evidently,ν = lim n (ũ n ) * µ n = N (0, 1) ⊗ δ 1 . Proposition 3.3. Fix k ≥ 2, λ > 0 and δ n = 1/n. Thenũ n (t) converges as n → ∞ to the solution of the following SDE initialized fromν: We see that r 2 ⊥ now solves an autonomous ODE which converges exponentially to 1. When k = 2, the equation form then converges as t → ∞ to 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ũ n and u n , one could in principle establish that for any λ > 1, SGD reaches the stable fixed 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 an overview, and Figures 2-3 for more refined numerical verification 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 finite in n, the expression for m 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 which transitions between stable and unstable OU processes at Λ c (k) = 1, as in k = 2.
Of course we only considered a few specific choices of summary statistics in the above, and the strength of Theorem 2.2 derives from its general applicability. As demonstrations, let us mention a few other examples that we would expect to be of interest in the study of SGD for matrix and tensor PCA. The first example is a limiting ballistic limit theorem for the evolution of the population loss Φ(x). The population loss can be taken as a summary statistic in our δ n -localizable triple; in the case of k-tensor PCA, this yields, The second example is a diffusive limit theorem near the fixed point (m, r 2 ⊥ ) = (0, 1) (as opposed to (3.3) where we blew up only the m variable about the saddle set m = 0 and therefore onlym was moving diffusively). In order to do so, we consider the scaling limit of the pair (m,r ⊥ ) = ( √ nm, √ n(r 2 ⊥ − 1)) and find the following limit: Interestingly, with this double rescaling, the n → ∞ limit yields a pair of OU processes that are decoupled, namely, each of their drifts are autonomous and their stochastic parts independent. This pair of independent OU processes is depicted in . Two-layer networks for classifying a binary Gaussian mixture As our second example, we consider the problem of supervised classification of a binary Gaussian mixture model (binary GMM) using a two-layer network. Our goal is to demonstrate how our approach can be used to analyze the performance of SGD for multi-layer networks, and indeed we will find the calculations relevant in Section 5 where we consider XOR-type GMM's.  Depicted is the evolution of summary statistic r 2 ⊥ (t) for 10n steps of SGD initialized randomly. This follows a stable OU process independent of λ.
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 ((2y − 1)µ, I/λ), where µ ∈ R N is a fixed 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 classification via the following architecture (though our techniques generalize to other settings mutatis mutandis): The first 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 ∈ R and sigmoid activation, σ(x) = 1/(1 + e −x ). Our parameter space is then X n = R 2N +2 and we therefore take n = 2N + 2 when applying Theorem 2.2. As we are interested in supervised classification, we take the usual binary cross-entropy loss with 2 regularization, where g is applied component wise and p(v, W ) := (α/2)(||v|| 2 + ||W || 2 ). It can be shown (see Lemma 8.1) that the law of the loss at a given point, (v, W ) ∈ X n , depends only on the 7 summary statistics, By Lemma 8.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.
Proposition 4.1. Let u n be as in (4.2) and fix 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→∞ (u n ) * µ n , with: Due to the Gaussian integrals defining f , g, it is difficult to analyze the ODE system defined 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 fixed point which is a good classifier. In particular, if λ = λ n is any diverging sequence, then the limiting effective dynamics would exactly match that attained by now sending λ → ∞. In Figure 6, we demonstrate numerically that the following predicted fixed points from the λ → ∞ limit match those arising at finite large N and λ > 0. 4 Proposition 4.2. The λ → ∞ limit of the ODE system of Proposition 4.1 is given bẏ The fixed points of this system are classified as follows. All fixed points have (2) If α < 1 /4, two unstable sets of fixed points at the quarter-circles given by 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 fixed points of (3) with probability 1 /2. values that 500 runs of SGD converge to after 100N steps from a random Gaussian initialization. The − and × are the unstable and stable fixed points of the λ = ∞ ballistic effective dynamics. The fixed points of the limiting effective dynamics have the same structure at finite λ as λ = ∞, and that as λ gets large quantitatively approach the λ = ∞ ones. Let us pause to interpret this result. The stable fixed points when α < 1/4 are the optimal classifiers, whereas the unstable set of fixed 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 classifier in linear time or it will appear to have converged to a macroscopically sub-optimal classifier on the same timescale, see Figures 6-7 for numerical verification of this at finite N and λ.
It is then natural to ask about the behaviour of the SGD in the latter regime, after it converges to the sub-optimal classifiers which lie on the aforementioned quarter-circles. Proposition 4.2 rigorously justified the exchange of n → ∞ and λ → ∞ limits in the ballistic phase. In the diffusive phase, one could in principle find the quarter circle of fixed points of the ODE in Proposition 4.1 and consider rescaled observablesṽ 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 λ = ∞, for which we know the precise locations of these fixed 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 8 demonstrate its qualitative consistency with the behavior in microscopic neighborhoods of fixed points at λ finite. where the diffusions can be seen to not be of full rank.
When λ = ∞, the SDE system obtained by applying Theorem 2.2 toũ n is whereΣ is a matrix whose only non-zero entries areΣṽ iṽj =Σm imj =Σṽ imj = α 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 (ṽ 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 fixed points lie is impossible in the noiseless setting, and happens on a much larger timescale at finite λ.

Two-layer networks for the XOR Gaussian mixture
We end with a discussion of the application of our methods to the problem of supervised learning for an XOR-type Gaussian mixture model in R N with a two-layer network.
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 = 0 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 label and X is the data.
Consider the corresponding classification problem using a two-layer neural network,ŷ(X) = σ(v · g(W X)), where σ and g are as in Section 8. We take W to be a K × N matrix and v to be a K-vector, and consider the binary cross-entropy loss (K ≥ 4 so that a good classifier is expressible by the network). It is shown in Lemma 9.1 that the law of the loss at a point (v, W ) depends only on the following 4K + K 2 variables: With the choice of u n given by these variables, for any fixed λ > 0, the localizability criterion of Definition 2.1 can be verified 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, see Proposition 9.2.
The fixed points of the ballistic dynamics, again in the limit λ → ∞ after n → ∞, are classified as follows (see Proposition 9.3). Similar to the binary GMM, in Figures 9-10, we demonstrate numerically that the following predicted fixed points from the λ → ∞ limit match those arising at finite large n and λ > 0. If α > 1 /8, then the only fixed 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, . . . , K}. Each such values that 500 runs of SGD converge to after 100N steps from a random Gaussian initialization. The − and × are the unstable and stable fixed points of the λ = ∞ ballistic effective dynamics. This demonstrates that the fixed points of the limiting effective dynamics have the same qualitative structure at finite λ as λ = ∞, and approach the λ = ∞ ones as λ gets large. partition corresponds to a collection of fixed points. For each tuple, the collection of fixed points has R ⊥ ij = 0 for all i, j, and In the minimal case of K = 4, analyzing the connectivity structure of these fixed points, there are 39 connected sets of fixed points, 4! = 24 of which are stable, one for each permutation where I + µ , I − µ , I + ν , I − ν are all singletons. For the proof, see Proposition 9.3. In the K = 4 case, we can also exactly calculate the probability that the effective dynamics in the ballistic phase converges to a stable fixed 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 converges to 3 /32. We refer the reader to Section 9.4 for the proof.

5.1.
Overparametrization in the XOR GMM. Since the the derivations of the ballistic limiting equations apply for general K, we can also study the probability of ballistic convergence to a stable (a) λ = 100 vs. unstable fixed point as one varies K. This addresses the regime of overparametrization for the XOR GMM since K = 4 suffices to express a Bayes-optimal classifier. In this more generic setting, the probability of being in the ballistic domain of attraction of the stable fixed points (corresponding to the Bayes optimal classifiers) is which goes to 1 exponentially fast as K grows. This clearly demonstrates the benefits of overparametrizaiton of the landscape in a concrete two-layer network: a random initialization is more likely to to contain the "right" initial signature (corresponding to none of I + µ , I − µ , I + ν , I − ν being empty at initialization) in order to be in the basin of a Bayes optimal classifier as the width grows, and as long as the right signature is present in the nodes at initialization, the SGD will ballistically converge to a global minimizer of the population loss. We refer the reader to Remark 5-6 for more.

5.2.
Diffusive limits at unstable fixed points. As an example of the diffusions that can arise in the rescaled effective dynamics at the unstable fixed points, let us consider the unstable fixed points in which v has the correct signature (two positive, two negative) but for each of those we are at a corresponding quarter-ring. By way of example, we can set K = 4, or equivalently focus on a fixed point where all indices beyond the first four have v i = m µ i = m ν i = 0. Here, the dynamics effectively becomes a pair of 2 two-layer GMM's on quarter-rings (as in Section 8), that are anti-correlated. More precisely, let (a 1,µ , a 2,µ ) be such that a 2 1,µ +a 2 2,µ = C α and (a 3,ν , a 4,ν ) such that a 2 3,ν +a 2 4,ν = C α , for C α = −logit(4α). Take as fixed points about which we expand to be .

Proof of Theorem 2.2
In this section, we prove our main convergence result, namely Theorem 2.2. The drift terms can be seen from a Taylor expansion out to second order, with the role played by δ n -localizability being to justify neglecting certain negligible second order terms, as well as all higher order terms. The identification of the stochastic term can be thought of as a version of the classical martingale problem [48] for summary statistics of stochastic gradient descent in the high-dimensional n → ∞ limit.
Notational remark. For ease of notation, in the following we say that f g if there is some constant C > 0 such that f ≤ Cg and that f a g if there is some constant C(a) > 0 depending only on a such that f ≤ C(a)g. We will often suppress the dependence on n in subscripts, when it is clear from context.

Proof of Theorem 2.2.
Our aim is to establish u n → u weakly as random variables on C([0, ∞)) where u solves (2.4). It is equivalent to show the same on C([0, T ]) for every T > 0.
Let τ n K denote the exit time for the interpolated process u n (t) from E n K . Define its pre-image E * K,n := u −1 n (E n K ) and let L ∞ K,n = L ∞ (E * K,n )). For any function f , we use the shorthand f to denote f (X ). By Taylor's theorem, we have that for any C 3 function f and any ≤ τ n K /δ, where A f and M f are defined by their increments as follows: Observe that A f is pre-visible and M f is a martingale. We bound these for f = u j among u n = (u 1 , ..., u k ).
Recalling Definition 2.1, since u n are δ n -localizable, the error term in (6.1) has Since δ n goes to infinity as n → ∞, we may thus write u j (X ) as where the last term is o(1) in L 1 uniformly for ≤ τ K /δ. Now let us define for s ∈ [0, T ], [s/δ]−1 ) and b j (s) = s 0 b j (s )ds , then recalling that u n (s) is the linear interpolation of (u j ([s/δ])) j , we may write u n (s) = u n (0) + a n (s) + b n (s) + o(1).
where a n (s) = (a j (s)) j and b n (s) = (b j (s)) j .
We now prove that the sequence (u n (s ∧ τ n K )) is tight in C([0, T ]) with limit points which are α-Holder for each K. To this end, let us define v n (s) = a n (s) + b n (s) + u n (0). As the o(1) error above is uniform in t, we have that Thus it suffices to show the claimed tightness and Holder properties of limit points for v n instead of u n . We aim to show that for all 0 ≤ s, t ≤ T , from which we will get that the sequence v n (s ∧ τ K ) is uniformly 1/4-Hölder by Kolmogorov's continuity theorem. Evidently, for all s, t we have v n (s) − v n (t) ≤ a n (s) − a n (t) + b n (s) − b n (t) .
We control these terms in turn. We will do this coordinate wise and, for readability, fix some j ≤ k and let u = u j , a = a j , b = b j etc. For the pre-visible term, we have uniformly over all x ∈ E * K,n . Therefore, the first term in (6.3) satisfies by continuity of h j . For the second term in (6.3), 4 by the definition of δ n -localizability. (Applying this bound for s = 0, t = T , the last term in a is vanishing in the limit for each K whenever δ n = o(1).) Combining the above bounds yields For the martingale term, notice that by independence, where the sum again runs over steps ranging from [s/δ] ∧ τ K to [t/δ] ∧ τ K . Repeatedly using the inequality (x + y + z) 2 x 2 + y 2 + z 2 , it suffices to bound the above quantity for each of the three terms defining the martingale difference M u − M u −1 respectively. For the first term in that martingale difference, observe that where in the second line we used Cauchy-Schwarz and in the last we used δ n -localizability. For the second term in the martingale difference, again by δ n -localizability. Finally, by the same reasoning, for the third term, All of the above terms are O((t − s) 2 ) since 0 ≤ s, t ≤ T . Thus we have the claimed (6.2), and by Kolmogorov's continuity theorem, (v n (s ∧ τ K )) s , are uniformly 1/4-Holder and thus the sequence is tight with 1 /4-Holder limit points. Notice furthermore that if we look at (v n (t ∧ τ K ) − a n (t ∧ τ K )) t , this sequence is also tight and the limits points are continuous martingales. Let us examine their limiting quadratic variations.
Let v K n (t) = v n (t ∧ τ K ) and define a K n (t) and b K n (t) analogously. Furthermore, let v K (t), a K (t) and b K (t) be their respective limits which we have shown to exist and be 1 /4-Holder.
We will compute the limiting quadratic variation for b K (t). For ease of notation, let is a martingale. We therefore need to consider the limit as n → ∞ of the integral above. Write Consider the integrals of δ times each of these four terms separately. By the assumption of (2.3), goes to zero as n → ∞. We now reason that the integrals of the other three terms all go to zero as n → ∞. We will show this for one of them, the arguments for the others being similar. By Cauchy-Schwarz, The first expectation contributes δ −1/2 by the first part of Item 3 of localizability. Also, The first of these terms is at most δ −1 as argued in (6.5). The second is o(δ −3/2 ) by the second part of Item 3 in the definition of localizability. As such, we are able to conclude that goes to zero as n → ∞. Altogether, with the analogous reasoning for the other two terms, Thus, if we consider the continuous martingales given by b K (t), its angle bracket is, by definition, given by By Ito's formula for continuous martingales (see, e.g., [18,Theorem 5 Since, by assumption, h, √ Σ are locally lipschitz-and thus lipschitz on E K -this property uniquely characterizes the solutions to (2.4) (see, e.g., [48,Theorem 6.3.4]). Thus v K converges to the solution of (2.4) stopped at τ K . Thus by a standard localization argument [48, Lemmas 11.1.11-12], every limit point v(t) of v n (t) solves the SDE (2.4) (using here that E K is an exhaustion by compact sets of R k ).

Proofs for matrix and tensor PCA
Our first aim is to establish Proposition 3.1, showing that the summary statistics u n = (m, r 2 ⊥ ) satisfy the conditions of Theorem 2.2 with the desired f , g and Σ. In what follows, for ease of notation we will denote r 2 = r 2 ⊥ and R 2 = m 2 + r 2 . We first establish that the sequence u n is δ n -localizable for any δ n = O(1/n). The localizing sequence E K will simply be centered balls of radius K in R 2 , say. We first check the regularity of the observable pair u n ; express the Jacobian for that pair as To check the regularity of observables, notice that ∇ 2 m = 0, while ∇ 2 r 2 = 2(I − vv T ), whose operator norm is simply 2, and ∇ u i = 0 for all ≥ 3. Next, we verify the regularity of the loss. In this appendix we will do things in the more general setting where we add a ridge penalty to the loss, so that for α > 0 fixed, the loss is given by , and thus H(x) = −2 W, x ⊗k . In the coordinates (m, r 2 ⊥ ), we have Φ(x) = −2λm k + (r 2 ⊥ + m 2 ) k + α 2 (r 2 ⊥ + m 2 ) + c . Observe that ∇Φ = ∂ 1 φ∇m + ∂ 2 φ∇r 2 , where Notice that ∇m, ∇m = 1, ∇m, ∇r 2 = 0, and ∇r 2 , ∇r 2 = 4r 2 . Consider ∇Φ ≤ |∂ 1 φ| ∇m + |∂ 2 φ| ∇r 2 ; the bounding quantity is evidently a continuous function of m, r 2 and therefore as long as x is such that (m, r 2 ) ∈ E K , it is bounded by some C(K). Next, if we consider where the bound on the operator norm of an i.i.d. Gaussian k-tensor can be found, e.g., in [5]. By the same reasoning, for every w, If w = ∇m = v then w = 1 and if w = ∇r 2 = 2(x − mv) then w ≤ C(K), so in both cases this is at most C(k, K)n 2 . Finally, ∇ 2 u is only non-zero if u = r in which case it is I − vv T . Then, by the second item in the definition of localizability, and evidently the right-hand side is o(δ −2 ) if δ n = O(1/n). This concludes the proof of δ n localizability for every δ n = O(1/n).
We now turn to calculating f , g, Σ. Starting with f , by the above, We next turn to calculating the corrector. For this, we first calculate the matrix V = E[∇H ⊗ ∇H].
Alternatively, if u 1 = 0 at a fixed point, then we can simplify further and get λu k−2 1 = R 2k−2 + α/k , and kR 2k−2 = (kR 2k−2 + α)u 2 , so that at the fixed point, For simplicity of calculations, set α = 0 as is the case in Proposition 3.1. Then, we simply get u 2 = c δ . In the case of k = 2, we also find that there is a solution if and only if λ > c δ , in which case R 2 = λ, from which together with R 2 = u 2 1 + u 2 , we also get In the general case of k > 2, we find that This has real solutions (all of which have R ≥ u 2 = c δ as required) whenever λ > λ c (k) defined as (Interpreting 0 0 = 1, this returns λ c (2) = c δ .) With this λ, whenever λ > λ c (k), the equation for R 2 has exactly two real solutions, both of which are at least c δ which we can denote by When λ > λ c (k), ρ † < ρ and when λ = λ c (k), the two are equal. Given this, we can then solve for u 1 at the corresponding fixed point, and find that they occur at Effective dynamics for the population loss. In practice it is often most useful to track the loss, or ideally, the generalization error. In this subsection, we add the generalization error Φ to our set of summary statistics and obtain limiting equations for its evolution. For simplicity of calculations let us stick to α = 0.
Taking limits as n → ∞, as long as λ is fixed in n, we see that f is given by , and fũ 2 = 4kũ k 2 .
We turn to obtaining the correctors in these rescaled coordinates. Evidently δLũ 1 = 0 still by linearity ofũ 1 . Following the calculation for the corrector, it is now given by gũ 2 = 4kũ k−1 2 . Next we consider the volatility of the stochastic process one gets in the limit. Recalling JV J T from (7.3), and noticing that the rescaling J →J multiplies its (1, 1)-entry by n and its off-diagonal entries by √ n, we find that in the new coordinates, Multiplying by δ = 1/n and taking the limit as n → ∞, the only entry of this matrix that survives is from Σ 11 where we get Σ 11 = 4kũ k−1 2 . Putting the above together yields the claimed Proposition 3.3. Regarding the discussion in the k ≥ 3 case when λ n = Λn (k−2)/2 , observe that the first term in Φ, ∇ũ 1 above would not vanish and would instead converge to −4kΛũ k−1 1 .

Proofs for the binary Gaussian mixture model
Recall the cross-entropy loss for the binary GMM with SGD from (4.1), and recall the set of summary statistics u n from (4.2).
Lemma 8.1. The distribution of L((v, W )) depends only on u n from (4.2). In particular, we have that Φ(x) = φ(u n ) for some φ. Furthermore, u n satisfy the bounds in item (1) of Definition 2.1 if E K is the ball of radius K in R 2N +2 .
Proof. Let X µ ∼ N (µ, I/λ) and X −µ ∼ N (−µ, I/λ). Then, notice that Next, notice that as a vector, , and Z 1,⊥ , Z 2,⊥ are jointly Gaussian with means zero and covariance Similarly, the distribution of W X −µ also only depends on Therefore, at a fixed point, the law of L((v, W )) is simply a function of u n (v, W ). This of course implies the same for the population loss Φ. To see that the summary statistics satisfy the bounds of item (1) in Definition 2.
For the higher derivatives, evidently we only have second derivatives in the last 3 variables each of which is given by a block diagonal matrix where only one block is non-zero and is given by an identity matrix. The third derivatives of all elements of u n are zero.
We can now express the loss, the population loss, and their respective derivatives and they (their laws at a fixed point) will evidently only depend on the summary statistics. One arrives at the following expressions for ∇L by direct calculation from (4.1).
In what follows, for an arbitrary vector w ∈ R N , we use the notation is only a function of u n by the same reasoning as used in Lemma 8.1.) Then, we can also easily express 8.6) and for H = L − Φ, Finally, the matrix V can be expressed as follows: Let us conclude this subsection with the following simple preliminary bounds that will be useful towards establishing the conditions of δ n -localizability from Definition 2.1, and the promised limiting equations. The proofs of these are straightforward using Gaussianity and are provided in the supplementary materials.
Fact 8.1. Fix µ ∈ S N −1 (1), and let g(x) = x ∨ 0 and X µ ∼ N (µ, I/λ). There is a function C : For every v i , R ⊥ ij and m i = 0 for i, j = 1, 2, we have 8.1. Verifying the conditions of Theorem 2.2 for fixed λ. Throughout this section we will take µ = e 1 . By rotational invariance of the problem, this is without loss of generality, and only simplifies certain expressions. The δ n -localizability can be seen by application of the moment bounds listed above and we defer it to the supplementary material.
Lemma 8.4. For δ n = O(1/N ) and any fixed λ, the 2-layer GMM with observables u n is δ n -localizable for E K being balls of radius K about the origin in R 7 .
Proof of Proposition 4.1. The convergence of the population drift to f from Proposition 4.1 follows by taking the inner products of ∇L from (8.6) with the rows of J from (8.2), and noticing that A µ i from (4.3) is exactly A i · µ and A ⊥ ij from (4.3) is exactly A i · W ⊥ j . Next consider the convergence of the correctors to the claimed g. The variables u ∈ {v 1 , v 2 , m 1 , m 2 } are linear so L n u = 0 and for these, g u = 0. For u = R ⊥ ij for i, j ∈ {1, 2}, the relevant entries in V are those corresponding to W ⊥ i and W ⊥ j . For ease of notation, in what follows let π = σ(v · g(W X)). For ease of calculation taking µ = e 1 , we have L n R ⊥ ij = k =1 V W ik ,W jk , which by (8.9), and the choice of δ n = c δ /N , is given by Consider the two terms separately. First, rewrite 1 Of course the second term is exactly what we want to be g u , so we will show the first term here goes to zero. By Cauchy-Schwarz, if Z ∼ N (0, I − e ⊗2 1 ), the first term above is at most where we used the fact that for a standard Gaussian, g ∼ N (0, 1), we have E[(g 2 − 1) 2 ] = 2. It remains to show the inner product term in (8.12) goes to zero as n → ∞. For this term, rewrite where X 1 , X 2 are i.i.d. copies of X, and π 1 , π 2 are the corresponding σ(v ·g(W X 1 )) and σ(v ·g(W X 2 )).
This term therefore also vanishes as n → ∞, yielding the desired limit for the corrector, which we emphasize is only a function of u n . We lastly need to show that the diffusion matrix Σ n goes to zero as n → ∞ when δ n = O(1/n). This is straightforward to see by considering any element of JV J T and using Cauchy-Schwarz together with the two bounds of Lemma 10.1 to bound it in absolute value by some C(K) independent of n. Then when multiplying by any δ n = o(1), this entire matrix will evidently vanish.
8.2. The small-noise limit of the effective dynamics. One can now take a λ → ∞ limit to arrive at the ODE system of Proposition 4.2. We defer the proof of this limit to the supplementary materials. The proof of the fixed points is provided below. Evidently, every fixed point must have and therefore every fixed point of the ODE system must have u i = 0, which is to say v i = m i . Therefore, it suffices to characterize the fixed points in terms of (v 1 , v 2 ) as claimed. This reduces Observe first that the point (v 1 , v 2 ) = (0, 0) is a fixed point of this system. If (v 1 , v 2 ) = 0, then dividing out by v i , the above reduces to σ(− v 2 ) = 2α if v 1 v 2 > 0 and σ(−v 2 i ) = 2α otherwise. Recalling that C α = − logit(2α) = log(1 − 2α) − log(2α) we obtain the claimed set of fixed points by inverting these equations (they only have a solution if α < 1/4).
In order to study the stability of the various fixed points, notice first that the ODE system of Proposition 4.2 is a gradient system for the λ = ∞ population loss, Since it is a gradient system, with only the specified fixed points, the stability of a fixed point can be deduced by showing it is the minimizer of Φ. In particular, the values of Φ at its critical points are given by Φ 0 = log 2 at v 1 = v 2 = 0, Φ + = 1 2 (log 2 + log(1 + e −Cα ) + αC α when v 1 v 2 > 0, and Φ − = log(1 + e −Cα ) + 2αC α when v 1 v 2 < 0. It is a simple calculus exercise to show that the smallest of these is Φ 0 when α > 1/4 and Φ − when α < 1/4.
To show that each of the other critical points are all unstable, one can find a direction along which the dynamical system is locally repelled from it. For instance, we will show that the ring of fixed points with v i = m i and R ⊥ ij = 0 with v 1 v 2 ≤ 0 is unstable, by showing a repelling direction arbitrarily close to the point v 1 = − √ C α , v 2 = 0. If v 1 = − √ C α and v 2 = > 0, thenv 2 there reduces to ( σ(− 2 ) 2 − α), and as long as α < 1/4, there exists > 0 such that σ(− 2 ) > 2α sov 2 > 0 for all small enough. 8.3. Rescaled effective dynamics around unstable fixed points. In this section, we consider scaling limits of the rescaled effective dynamics in their noiseless limit, where the rescaling is about the unstable set of fixed points given by the quarter circle v 2 1 +v 2 2 = C α per item (2) of Proposition 4.2. Let δ n = cδ /N, and fix (a 1 , a 2 ) ∈ R 2 + with a 2 1 + a 2 2 = C α , and let u n be the variables of (4 Proof of Proposition 4.3. We start by considering the drift process for these rescaled variables. Notice that the rescaling induces the transformationJ multiplying J by √ N in its entries corresponding to v i , m i . The fact that the rescaled variables satisfy the conditions of Theorem 2.2 follows as in Lemma 8.4 with the only distinction arising in the bound on (10.4), where previously we did not use the δ 2 n factor-in the new coordinates, the factor of √ N raised to the fourth power is cancelled out by δ 2 n as long as δ n = O(1/N ). For the population drift of the new variables, if the variablesṽ i ,m i are in a ball of radius K in R 4 (which we take to be our E K ), the signs of m i agree, and therefore We wish to claim that these expressions have consistent limits whenṽ i ,m i are localized to E K for fixed K. notice that in m i = a i + N −1/2m i and v i = a i + N −1/2ṽ i , and using Now Taylor expanding the sigmoid function, and using the definition of C α , we get Plugging these into the earlier expressions for fṽ i , we see that Taking the limit as n → ∞, this yields exactly the population drift claimed for theṽ i variable. The calculation for fm i is analogous, and the equations for R ⊥ ij are evidently unchanged by the transformation of v i , m i toṽ i ,m i . Furthermore, these variables are still linear so no corrector is introduced.
We now turn to computing the limiting diffusion matrix Σ in the new variablesṽ i ,m i . We first use the following expression for the matrix V when λ = ∞, by taking the λ = ∞ in (8.9): with similar expressions for V v i ,W j and V W i ,W j . Rewriting inṽ andm, we see that in E K , Now multiplying this on both sides byJ, for theũ n variables, the two factors of √ N fromJ cancel out with the choice of δ n = 1/N , and in the n → ∞ limit, leaveΣ v i v j =Σ m i m j =Σ v i m j = α 2 a i a j as claimed. 9. Proofs for the XOR Gaussian mixture model Fix two orthogonal vectors µ, ν ∈ R N and recall the cross-entropy loss with penalty p(v, W ) = α 2 ( v 2 + W 2 ). For the XOR GMM with SGD, the cross-entropy loss is given by where if the class label y = 1, then X is a symmetric binary Gaussian mixture with means ±µ, and if y = 0, then X is a symmetric Gaussian mixture with means ±ν. This has the same form as the loss for the 2-layer binary GMM, and we will find many similarities in the below between them. Indeed, the only difference is in the distribution of X conditionally on the class label y as described, and the fact that v is now in In what follows we take n = KN + K. As such, all the formulae of (8.3)-(8.9) also hold for the XOR GMM, but with the law of (y, X) now understood differently.
Remark 4. We could also have added a bias at each layer, however the Bayes classifier in this problem is an "X" centered at the origin so we can safely take the biases to be 0.
9.1. Summary statistics and localizability. Recall the set of summary statistics u n from (5.1). The next lemma shows that u n form a good set of summary statistics.
Lemma 9.1. The distribution of L((v, W )) depends only on u n from (5.1). In particular, we have that Φ(x) = φ(u n ) for some φ. Furthermore, u n satisfy the bounds in item (1) of Definition 2.1 with an exhaustion by balls of R KN +K .
Proof. Let X w = N (w, I/λ) for w ∈ {µ, −µ, ν, −ν}. Notice that the law of L at a fixed point (v, W ) ∈ R K+KN can be written as Next, notice that as a vector Therefore, at a fixed point (v, W ) the law of L(v, W ) is only a function of u n (v, W ).
To see that the summary statistics satisfy the bounds of item (1) in Definition 2.1, note that the non-zero entries of J = (∇u ) are as follows.
where δ ij is 1 if i = j and 0 otherwise. For higher derivatives, we only have second derivatives in the R ⊥ jk variables, each of which is given by a block diagonal matrix where only one block is non-zero and it is twice an identity matrix. Thus the operator norm of these second derivatives is 2. The third derivatives of all elements of u n are zero.
In the following, let By the same reasoning as in Lemma 9.1, if w ∈ {µ, ν, W i , W ⊥ i }, then w · A i is only a function of u n . We then also have the conclusions of Lemma 10.1 for X distributed according to the XOR GMM by simply decomposing it into two mixtures, and we will therefore appeal to this lemma meaning its analogue for the XOR GMM.
Proposition 9.1. Let u n be as in (5.1) and fix any λ > 0 and δ n = c δ /N . Then u n (t) converges to the solution of the ODE systemu t = −f (u t ) + g(u t ), initialized from lim n (u n ) * µ n with The convergence of the population drift to f from Proposition 4.1 follows by taking the inner products of ∇L from (8.6) with the rows of J from (9.3), and noticing that A µ i is exactly A i · µ, A ν i is exactly ν · A i , and A ⊥ ij is exactly A i · W ⊥ j . We next consider the population correctors. The fact that g v i = g m µ i = g m ν i = 0 follows from the fact that the Hessians of v i , m µ i , m ν i are zero. For the corrector g R ⊥ ij for 1 ≤ i ≤ j ≤ K, the relevant entries of V are those corresponding to W ⊥ i and W ⊥ j . For ease of notation, in what follows let π = σ(v · g(W X)).
Similar to the calculation of (8.12), By the same arguments on the concentration of the norm of Gaussian vectors as used in the binary GMM case, then we deduce from this that Finally, let us establish that the limiting diffusion matrix is all-zero whenever δ n = o(1). This follows exactly as it did in the proof of Proposition 4.1.
9.3. Small noise limit of the effective dynamics. The aim of this section is to establish the following small-noise λ → ∞ limit of the effective dynamics ODE of Proposition 9.1. This will again be quite similar to the analogous proofs for the binary GMM in Section 8, and when these similarities are clear we will omit details.
Proposition 9.2. In the λ → ∞ limit, the ODE from Proposition 9.1 converges tȯ Proof. Let us begin with convergence of A µ i . We claim that it converges to ) . In order to see this, expand The point will be that when taking the inner product with µ, the first two terms here contribute to the limit and the latter two vanish, while when taking the inner product with ν, the first two terms vanish in the λ → ∞ limit while the latter two contribute. Consider e.g., the first of the four terms above, and inner product with µ. In this case, consider , which is precisely the quantity that was exactly shown to go to zero as λ → ∞ in (10.5). To see that the third and fourth terms above go to zero when taking their inner product with µ, observe that they become which by orthogonality of µ and ν is at most λ −1/2 by the reasoning of Lemma 10.1, therefore vanishing as λ → ∞. Together with its analogue for X −ν , this implies the claim for the convergence of A µ i , as well as its analogous limit of A ν i . We next consider the limit as λ → ∞ of A ⊥ ij , which we claim goes to 0. Using the expansion of A i from earlier in this proof, we can consider A ⊥ ij = A i · W ⊥ j as four terms having the form of the terms in (10.6), which were there showed to go to zero as λ → ∞. Since W ⊥ j here is orthogonal both to µ and ν, the same proof applies.
Finally, in order to see that the limit as λ → ∞ of g R ⊥ ij = c δ v i v j λ B ij is zero, which follows from the fact that |B ij | ≤ 1. Proposition 9.3. The fixed points of the ODE system of Proposition 9.2 are classified as follows. If α > 1/8, then the only fixpoint is at u n = 0.
Our classification now reduces to understanding the possible values taken by (v 1 , ..., v K ) given their signs (when non-zero). Fix a partition (I 0 , I + µ , I − µ , I + ν , I − ν ) of {1, ..., K} and consider the set of fixed points having m µ i = m ν i = v i = 0 for i ∈ I 0 , m µ i = v i > 0 on I + µ and so on as designated by Proposition 9.3; by the above any fixed point is of this form. It remains to check that the values of v i on each of these sets are as described by the proposition.
In order to see this, fix e.g., i ∈ I + µ . Then, m µ i = v i and m ν i = 0, and so the fixed point equations reduce to since the only coordinates where g(m µ ) will be non-zero are j ∈ I + µ , where m µ j = v j . Inverting the sigmoid function, this implies exactly the claimed j∈I + µ v 2 j = logit(−4α). The cases of I − µ , I + ν , I − ν are analogous, concluding the proof.
The count of the number of connected components of fixed points this forms is sensitive to K, so for concreteness let us perform it when K = 4. We first notice that the fixed point at (0, ..., 0) is disconnected from all others. Fixed points corresponding to some (I 0 , ..., I − ν ) are part of the same connected component of fixed points if one goes from one to the other by moving an element of I η ι (for some ι ∈ {µ, ν} and η ∈ {±} to I 0 without making I η ι empty, or by moving an element of I 0 to a non-empty I η ι . We turn now to studying the stability of these various sets of fixed points. Observe that in the λ → ∞ limit, the dynamical system of Proposition 9.2 is a gradient system for the population loss Φ = 1 4 log(1 + e −v·g(m µ ) ) + · · · + log(1 + e −v·g(−m ν ) ) At a fixed point (which necessarily has v i = m i , R ⊥ ii = 0, and is characterized by the partition of {1, ..., 4} into I + µ , I − µ , I + ν , I − ν , this reduces to At this point, noticing that i∈I + µ v 2 i is equal to C α = −logit(4α) if I + µ is non-empty and 0 if it is empty, and similarly for I − µ , I + ν , I − ν , this turns into a simple optimization problem over the number of non-empty I + µ , I − µ , I + ν , I − ν . Just as in the binary GMM case, it becomes evident that when α > 1/8, this is minimized at v i = 0 for all i (i.e., they are all empty and I 0 = {1, ..., 4}, whereas when α < 1/8 the above is minimized when every one of I + µ , I − µ , I + ν , I − ν are all non-empty. This yields the global minima of Φ in these coordinates, and ensures the fixed points we claimed were stable are indeed stable.
To show the instability of any other connected set of fixed points, the reasoning goes just as in the binary GMM case: consider a small perturbation of the specified critical region in the direction of the stable fixed points and it can be seen by examining the drifts directly, that the dynamical system has a repelling direction.
Remark 5. When K > 4, the counting of connected components of fixed points of course changes. However, what is still clear by an identical calculation is that the sets of fixed points minimizing Φ will still be (0, ..., 0) when α > 1/8 and will be all fixed points that have all four of I + µ , ..., I − ν being non-empty if α < 1/8. Notice that when α < 1/8 and K > 4, even the set of stable fixed points become connected to form a single stable manifold. 9.4. 3 ⁄32-probability of ballistic convergence to an optimal classifier. We now reason that when K = 4 the ballistic effective dynamics of Proposition 9.2 is such that under an uninformative Gaussian initialization, the probability of being in a basin of attraction of one of the 24 stable fixed points is 3/32. Begin by noticing that if the first layer weights are initialized as W i ∼ N (0, I N /N ) independently for i = 1, ..., 4 and the second layer weights v i are independent standard Gaussians, then the projection onto the coordinate system For the m µ i , m ν i variables one should understand these δ 0 Dirac masses as a 1 2 -1 2 mixture of δ 0 − and δ 0 + . Under the flow of Proposition 9.2, if v i (0) is positive, then m ν i stays fixed at zero, and if m µ i (0) = 0 − then m µ i becomes negative infinitesimally quickly, whereas if m µ i (0) = 0 + then it becomes positive infinitesimally quickly. At any rate, the sign of v i never changes to negative from such an initialization, and similarly if v i (0) is negative, the sign of v i will never change to positive. As such, in order to have a chance at being in the basin of attraction of one of the stable fixed points outlined in Proposition 9.3, it must be the case that two of (v i (0)) i have positive sign and two of them have negative sign; evidently this has probability 4 2 /2 4 = 3/8. Given that two of v i (0) are positive, and two of them are negative-say without loss of generality that i = 1, 2 are the coordinates in which it is positive, and i = 3, 4 are the coordinates in which it is negative-then the dynamical system for (v 1 , v 2 , m µ 1 , m µ 2 ) is exactly the ballistic limit of the two-layer GMM studied in Section 8, for which we found that the probability of converging to a good classifier is 1/2. Similarly, the dynamical system for (v 2 , v 4 , m ν 3 , m ν 4 ) independently gives a further probability 1/2 of converging to its good classifier. Together, these yield a probability of 3/32 of converging to one of the 4! many optimal classifiers for the XOR GMM.
Remark 6. Generically, if K ≥ 4, by a similar reasoning to the above, in order to fall in the basin of attraction of the stable fixed points, it must be the case that the initialization has some four indices each of which initially belong to I + µ , I − µ , I + ν , I − ν . This is the probability that v i (0) are positive for at least two indices, and negative for at least two indices, and then among the indices at which v i (0) is positive, there is at least one index where m µ i is positive and one where it is negative, and similarly with v i (0) negative and m ν i . Doing this combinatorial calculation out, we find that the probability of being in a good initialization is exactly the expression in (5.1). This is easily seen to go to 1 exponentially fast as K → ∞ since the initial choice of v i (0)'s will typically have around K/2 positive and K/2 negative coordinates, and with exponentially high probability those will have both positive and negative m µ i and m ν i . 9.5. Diffusive limit on critical submanifolds. We now consider scaling limits of the rescaled effective dynamics in their noiseless limit, where the rescaling is about the unstable set of fixed points given by the product of two quarter circles where I + µ = {1, 2} and I + ν = {3, 4} (if K > 4, examine the fixed point in which all coordinates after the first four are in I 0 ). In what follows, fix (a 1,µ , a 2,µ ) ∈ R 2 + with a 2 1,µ + a 2 2,µ = C α , and a 2 3,ν + a 2 4,ν = C α , and let u n be the variables of (4.2) with v i , m µ i , m ν i replaced byṽ .
By the choices ofm µ i = 0 andm ν i = 0, we mean that we formally mean that we remove those variables fromũ n , and for us now E K will be the ball of radius K in the other coordinates, and the point {0} for (m µ i ) i=3,4 and (m ν i ) i=1,2 . Proof of Proposition 5.1. The fact that the rescaled variablesũ n satisfy the conditions of Theorem 2.2 follows as in Lemma 9.2 with the only distinction arising in the bound on (10.4), where previously we did not use the δ 2 n factor, but is still satisfied using δ n = O(1/n). We next consider the population drift of the new variablesṽ i ,m µ i andm ν i . If we take these variables to be in E K , and recall the population drifts etc. in the λ = ∞ setting from Proposition 9.2, for i = 1, 2, we have fṽ i is the n → ∞ limit of If we then use the expansion v · g(m µ ) = C α + N −1/2 j=1,2 a j,µ (ṽ j +m µ j ) + O(1/n) from which we obtain σ(−v · g(m µ )) = σ(−C α ) + 1 √ N j=1,2 a j,µ (ṽ j +m µ j ) (4α)(1 − 4α) + O( 1 n ) Plugging these in, and taking the n → ∞ limit we find that for i = 1, 2, a k,µ (ṽ k +m µ k ) .
The claimed equations for fm µ i when i = 1, 2 and fm ν i when i = 3, 4 hold by analogous reasoning, and the equations for f R ⊥ ij are evidently unaffected by the change of variables toũ n . Regarding the population correctors, they are also unaffected (all zero) since the variables that were changed inũ n are all linear.
By a similar reasoning, if i, j ∈ {1, 2}, then and if i ∈ {1, 2} and j ∈ {3, 4}, then Taking the limit as n → ∞, we again recover the claimed limiting diffusion matrix, and similar calculations yield the same for Using the fact that Z is mean zero, and pulling out w · µ, we see that this is at most some universal constant times (w · µ) 4 + λ −2 E[(w · Z) 4 ]. For the second term, notice that w · Z is distributed as z ∼ N (0, w 2 ), implying the desired. The bound on A i goes as follows. Evidently it suffices to let X µ = µ + λ −1/2 Z for Z ∼ N (0, I), and prove the bound on the norm of E[X µ 1 W i ·Xµ≥0 (−1 + σ(g(W X µ )))] = E[(µ + λ −1/2 Z)1 W i ·Xµ≥0 (−1 + σ(g(W X µ )))] . Now decompose Z as Z µ µ + Z 1,⊥ W ⊥ 1 + Z 2,⊥ W ⊥ 2 + Z 3 , where Z µ ∼ N (0, 1) is independent of (Z 1,⊥ , Z 2,⊥ ) which is distributed as N (0, A) with A given by (8.1), which is independent of Z 3 distributed as a standard Gaussian vector orthogonal to the subspace spanned by (µ, W ⊥ 1 , W ⊥ 2 ). By independence of Z 3 from the indicator and the argument of the sigmoid, all those terms contribute nothing to the expectation, and therefore, Here, we used the first inequality of the lemma. This yields the desired.

This in turn is bounded by
If m i = 0, the expectation on the right goes to zero by (10.1). Second, When m i < 0, this is evidently zero; when m i > 0, if G λ ∼ N (0, I/λ), this is which goes to zero as O(λ −1 ) when λ → ∞, by the explicit formula for the moment generating function of the Gaussian W i · G λ , whose variance is (m 2 i + R ⊥ ii )λ −1 .
Proof of Lemma 8.4. The condition on u n was satisfied per Lemma 8. , and recall the expressions for ∇H from (8.7)-(8.8). Using the trivial bound |σ(x)| ≤ 1, and the inequality (a + b) 4 ≤ C(a 4 + b 4 ), for i ∈ {1, 2}, the first term is at most C(E[|X · W i | 4 ] + W i 4 A i 4 ) which is bounded by a constant depending continuously on u n per Lemma 10.1. If we let Z be a standard Gaussian, the second term is evidently governed by Using the well-known bound that E[ Z 4 ] ≤ N 2 , and the fact that δ = O(1/N ), we see that this is at most Cδ −2 . We turn to the third item in the definition of localizability. We next verify the claimed bound that When u i is v i , this is simply a fourth moment bound on ∇ v i H, which follows as the third moment bound did, with no need for the δ 2 n . When u i is m i , or R ⊥ ij , the bound follows from for choices of w being either µ in which case w = 1 or W ⊥ i in which case w = R ⊥ ii . For each K, this is at most some constant C(K) using the two bounds of Lemma 10.1.
Finally, consider the quantity E[ ∇ 2 u, ∇H ⊗ ∇H − V 2 ]. This is only non-zero for u ∈ {R ⊥ ij } for which ∇ 2 u is a block-identity matrix, having operator norm at most 2 in all cases. Therefore, this quantity is at most 4E[ ∇H 2 ] which is at most N by the above proved second item in the definition of localizability. This is therefore O(δ −1 n ) = o(δ −2 n ) as needed.
Proof of Proposition 4.2. We begin with considering lim λ→∞ A µ i : its limiting value will depend on the signs of both m 1 and m 2 . We can express A µ i from (4.3) as