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

Part I Introduction and main results

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 , 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 . 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 . 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 . 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., .

In recent years, there has been considerable interest in understanding the high-dimensional setting, where one is constrained in the amount of data or the run-time of the algorithm due to the high-dimensional nature of the data and the complexity of the model being trained. In these regimes, 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 .

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 and special classes of non-convex objectives . There has also been important work on scaling limits as the dimension tends to infinity for the specific problems of linear regression , Online PCA , and phase retrieval from random starts, and teacher-student networks and two-layer networks for XOR Gaussian mixtures 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 .

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.3 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 and very recently investigated in detail in .

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 mean-reverting and mean-repellent Ornstein–Uhlenbeck (OU) 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 .) 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 $\nicefrac{{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$; this is a mathematically rigorous example of the lottery ticket hypothesis of .

Before delving into the XOR problem, we first analyze the classification 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 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\geq 2$. We leave this to future investigation. In this paper, we only consider the simplest variant of SGD, namely online SGD; we leave other variants involving batching and re-use to future works.

Main result

To develop a scaling limit, we need some regularity 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

$\max_{i}\sup_{x\in\mathbf{u}_{n}^{-1}(E_{K})}\lvert\lvert\nabla^{2}u_{i}^{n}\rvert\rvert_{\operatorname{op}}\leq C_{K}\cdot\delta_{n}^{-1/2}$, and $\max_{i}\sup_{x\in\mathbf{u}_{n}^{-1}(E_{K})}\lvert\lvert\nabla^{3}u_{i}^{n}\rvert\rvert_{\operatorname{op}}\leq C_{K}$;

We now turn to our second assumption, that the limiting evolution equations for the family of summary statistics chosen close. Define the following first and second-order differential operators,

Alternatively written, $\mathcal{A}_{n}=\langle\nabla\Phi,\nabla\rangle$ and $\mathcal{L}_{n}=\frac{1}{2}\langle V,\nabla^{2}\rangle$.

In this case we call $\mathbf{h}$ the effective drift, and $\mathbf{\Sigma}$ the effective volatility.

We are now ready to present our main result. For a function $f$ and measure $\mu$ we let $f_{*}\mu$ denote the push-forward of $\mu$.

The proof of Theorem 2.3 is provided in Section 6 and can be seen as a version of the classical martingale problem (see ) for high-dimensional stochastic gradient descent. We call the solution to (2.4) the effective dynamics of the summary statistics $\mathbf{u}_{n}$. The fact that $\mathbf{h},\mathbf{\Sigma}$ are locally Lipschitz ensures that this solution is unique.

We end this subsection with discussion of the various scalings appearing in Definition 2.1.

Turning to item (2) of Definition 2.1, we comment that the regularity assumptions made on $\Phi,L$ here are less restrictive than uniform Lipchitz assumptions common to the literature. In particular, we do not assume the population loss is Lipschitz everywhere, as we may have that $\bigcup_{K}\mathbf{u}_{n}^{-1}(E_{K})$ does not cover $\mathcal{X}_{n}$, nor does it imply uniform smoothness of $H$ (and in turn $L$) as we may (and will) be taking $\delta_{n}\to 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 $\nabla u$ and $\nabla^{2}u$ at $x$. As a testbed, suppose that $\nabla H(x)$ is a random vector with i.i.d. entries all of order $1$. If $u$ is a rescaled linear statistic, e.g., $\delta_{n}^{-1/2}\langle x,e_{1}\rangle$ 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., $\delta_{n}^{-1/2}\|x\|^{2}$, again assuming for maximal simplicity that $\nabla 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.

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.3 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\to\infty$ limit. In principle the second assumption in item (3) of Definition 2.1 could be dropped; in that case, the same quantity is still ensured to be $O(\delta^{-3})$ by the other localizability assumptions. Then Theorem 2.3 would still apply, but the limiting diffusion matrix would be the $n\to\infty$ limit (assuming it exists) of

as opposed to simply the limit of $\delta JVJ^{T}$.

in which case, evidently (2.2) holds with $\mathbf{h}=-\mathbf{f}+\mathbf{g}$. When (2.5) and (2.6) both hold, we call $\mathbf{f},\mathbf{g}$ and $\Sigma$ the population drift, the population corrector, and the diffusion matrix of $\mathbf{u}$ respectively. From the fixed dimensional perspective, when (2.5) holds, one predicts $\mathbf{u}$ to solve

with initial data $\mathbf{u}_{0}\sim\mathbf{u}_{*}\mu$. as this is its evolution under gradient descent on the population loss $\Phi$. Evidently this perspective only applies in the high-dimensional limit of Theorem 2.3 if both the population corrector $\mathbf{g}$ and the diffusion matrix $\Sigma$ are zero. We find that for any triple $(\mathbf{u}_{n},L_{n},P_{n})$, there is a scaling of the learning rate $\delta_{n}$ with $n$ below which $\mathbf{g}=\Sigma=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 $\Sigma$ may be non-zero, and the high-dimensionality induces non-trivial corrections to $\mathbf{f}$. (In the case of teacher–student networks, the terms $\mathbf{f}$ and $\mathbf{g}$ can be compared to the “learning" and “variance" terms in Eq. (14a) of .)

To see this, notice that if the triple $(\mathbf{u}_{n},L_{n},P_{n})$ is $\delta_{n}$-localizable for some $\delta_{n}\to 0$, then it is also $\delta_{n}^{\prime}$-localizable for every sequence $\delta^{\prime}_{n}=O(\delta_{n})$. If furthermore (2.3) and (2.5)–(2.6) hold for $\delta_{n}$ with some $\mathbf{f},\mathbf{g}$ and $\Sigma$, then these limits also exists for $\delta_{n}^{\prime}=o(\delta_{n})$ with the same $\mathbf{f}$ but with $\mathbf{g}=\Sigma=0$. As such, there can be exactly one scaling of $\delta_{n}$ with $n$ at which $\mathbf{g}$ or $\Sigma$ may be non-zero, and for all smaller scales of $\delta_{n}$, the fixed-dimensional perspective of (2.7) applies.Note that if $\delta_{n}=o(\delta_{n}^{\prime})$, then limiting $\mathbf{g},\Sigma$ may not exist for $\delta^{\prime}_{n}$, so there is no super-critical regime.

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 $\mathbf{h}\neq\mathbf{f}$. We call this the ballistic limit. In this case, the effective dynamics of the summary statistics is given by the ODE system

In these settings, the phase portrait of the summary statistics is asymptotically that of this flow.

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 number of iterations that are order $1/\delta_{n}$. If one is then interested in the evolution of $\mathbf{u}_{n}$ in microscopic $o(1)$ neighborhoods of the fixed points of the ballistic effective dynamics of (2.8), Theorem 2.3 also allows one to develop separate diffusive limits there.

This then leads to the rescaled effective dynamics of the summary statistics $\mathbf{u}_{n}$ near $\mathbf{u}_{\star}$:

In , it was empirically observed that the best training for neural networks does not occur when step-sizes are small enough for the classical gradient flow approximation to be valid. Instead, it occurs at the edge of stability where the step size is just small enough for the training to remain stable. Here, the loss fluctuates for some time before eventually converging to lower values than it would with smaller step size. This critical step size scaling is defined via the sharpness, namely the largest eigenvalue of the training loss Hessian. For a selection of recent theoretical investigations of this phenomenon see, e.g., .

While sharpness and edge of stability do not have direct analogues in the context of online SGD, a qualitatively similar phenomenon can be seen by taking the population loss as a summary statistic. The critical scaling of the learning rate with dimension discussed in Sections 2.1–2.2 constrains the step size in terms of the top eigenvalue of the Hessian of the loss. With this scaling, the population loss fluctuates near critical regions of its ballistic flow, allowing it to escape the critical region, whereas with a sub-critical learning rate the population loss stays stuck. We leave more detailed investigation of this connection to edge-of-stability phenomena for SGD to future investigation.

Part II Examples

In the following sections, we demonstrate Theorem 2.3 on a range of popular examples of high-dimensional statistical tasks. We begin first in Section 3 by presenting an application to a widely studied problem of high-dimensional estimation: namely, de-noising a rank one tensor that has been corrupted additively by Gaussian noise. We then turn to classification. Our aim in these examples is to demonstrate the applicability of our result to the analysis of multi-layer neural networks. To this end we analyze the training dynamics of a two-layer neural network for two canonical classification tasks, namely classification of a symmetric, binary gaussian mixture model (Section 4) and classification of a Gaussian analogue of the XOR problem of Minsky–Papert (Section 5).

Consider the problem of de-noising a rank one tensor that has been corrupted additively by Gaussian noise via SGD. A popular statistical model of this task is the spiked tensor model . Suppose that we are given i.i.d. samples of data of the form

In the case $k=2$, this is a version of the well-known spiked matrix model of PCA for which there is, by now, a substantial literature regarding the statistical thresholds. For a necessarily small selection see, e.g., . For related work on online learning in this context see, e.g., . Of particular interest in this direction is the well-known phase transition at $\lambda=1$ for estimation in this problem, which was determined first for Wishart ensembles in and subsequently for this setting in . As we will see in Section 3.3 below, we find a dynamical analogue of this transition at $\lambda=1$.

The case $k\geq 3$ was introduced by Montanari and Richard as a natural generalization of the spiked matrix models for estimation (and testing) problems where the data has multiple indices or requires higher moments. Here there has been a large literature on the statistical thresholds for estimation and testing, see, e.g., . In this setting, there has also been a tremendous literature on the computational aspects of this problem as it is viewed as a important example of a model with a statistical-computational gap, namely, a setting where there is a gap between the regimes of statistical and computational tractability. See, e.g., .

We begin with these examples as their effective dynamics are particularly simple to analyze. In particular, they are 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. In particular, as mentioned above, we find dynamical phase transitions corresponding to the aforementioned thresholds in these models. For our analysis we will focus exclusively on the most interesting, critical step-size scaling which corresponds to the proportional asymptotics regime from the random matrix theory literature.

2. Analysis

We take as loss the (negative) log-likelihoodNote that one might also add additional penalty terms. The case of a ridge penalty is treated in Section 7. namely,

are such that $\Phi(x)=-2\lambda m^{k}+(r_{\perp}^{2}+m^{2})^{k}+c$, and the law of $L$ only depends on them.

In our normalization with $\lambda>0$ fixed, the regime $\delta_{n}=o(1/n)$ is sub-critical and the regime $\delta_{n}=\Theta(1/n)$ is critical. Note that with different scalings of $\lambda_{n}$, the critical learning rate changes. We focus our presentation on the most interesting regime, namely the critical scaling regime of $\delta_{n}=c_{\delta}/n$ for some constant $c_{\delta}$. Recalling the relation between number of samples and step-size, we see that this regime corresponds to the proportional asymptotics regime most studied in the random matrix theory literature where the above-mentioned transition for the top eigenvalue occurs. Note, however, that the limits in the subcritical regime are in all cases recovered by taking the $c_{\delta}\downarrow 0$ limits of the ODE’s/SDE’s of the critical regime.

For notational simplicity, let $R^{2}:=m^{2}+r_{\perp}^{2}$. We consider the pair $\mathbf{u}_{n}=(u_{1},u_{2})=(m,r_{\perp}^{2})$, for which Theorem 2.3 yields the following effective dynamics.

Fix $k\geq 2$, $\lambda>0$, $c_{\delta}>0$ and let $\delta_{n}=\nicefrac{{c_{\delta}}}{{n}}$. Then $\mathbf{u}_{n}(t)$ converges as $n\to\infty$ to the solution of the following ODE initialized from $\lim_{n\to\infty}(\mathbf{u}_{n})_{*}\mu_{n}$:

We are able to identify and classify the set of fixed points of this effective dynamics. We focus on the critical step-size regime with $c_{\delta}=1$ where one sees from (3.1) that $r_{\perp}^{2}\to 1$, where the problem in the matrix case is most directly related to an eigenvalue problem (see Section 7 for the generic $c_{\delta}$ dependencies). Throughout the following, we use the following notion of stability/unstability of a set of fixed points of an ODE.

We call a set of fixed points $U$ for an ODE stable if for every $\epsilon>0$, for every $u\in B_{\epsilon}(U)$, the solution of the ODE with initialization $u$ converges to some point in $U$ as $t\to\infty$. Otherwise, we call $U$ unstable.

Eq. (3.1) has isolated fixed points classified as follows. Let $\lambda_{c}(k)$ be as in (7.6) and $m_{\dagger}(k,\lambda)\leq m_{\star}(k,\lambda)$ be as in (7.7) (if $k=2$, $\lambda_{c}=1$ and $m_{\dagger}=m_{\star}=\sqrt{\lambda-1}$):

An unstable fixed point at $(0,0)$ and a fixed point at $(0,1)$; if $k=2$, $(0,1)$ is stable if $\lambda<\lambda_{c}(2)$ and unstable if $\lambda>\lambda_{c}(2)$; if $k>2$ $(0,1)$ is always stable.

If $\lambda>\lambda_{c}(k)$: when $k=2$, two stable fixed points at $(\pm m_{\star}(2),1)$. When $k\geq 3$, two unstable fixed points at $(\pm m_{\dagger}(k),1)$ and two stable fixed points at $(\pm m_{\star}(k),1)$.

The presence of two pairs of fixed points when $k\geq 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 $\lambda$ is finite. This is an interesting consequence of the corrector in (3.1) and if one tracks the $c_{\delta}$ dependence in the above, the fixed point $m_{\dagger}$ goes to zero as $c_{\delta}\to 0$ and this barrier to recovery from warm starts vanishes as one approaches sub-critical step-sizes.

3. A dynamical analogue of the BBP transition

4. On the sample complexity of tensor PCA

which transitions between mean-reverting and mean-repellent at $\Lambda_{c}(k)=1$, as in $k=2$.

We only considered a few specific choices of summary statistics in the above, and the strength of Theorem 2.3 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 $\Phi(x)$. The population loss can be taken added to the family of summary statistics in our $\delta_{n}$-localizable triple; in the case of $k$-tensor PCA, this yields,

5. A finer diffusive limit theorem at a random start

Interestingly, with this double rescaling, the $n\to\infty$ 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 Figures 4–5.

Two-layer networks for classifying a binary Gaussian mixture

As a warm-up to the XOR problem that we will consider in Section 5, we consider the problem of supervised classification of a binary Gaussian mixture model (binary GMM) which is defined as follows. Suppose that we are given i.i.d. samples of the form $Y=(y,X)$, where $y$ is a $\{0,1\}$-valued $Ber(1/2)$ random variable and, conditionally on $y$, we have

It is classical that the Bayes optimal estimator in this setting is given by $\hat{y}=\operatorname{sgn}(\mu\cdot x)$. Furthermore, this estimator can be achieved by (a rounding of) the output of a single layer neural network trained using the binary-cross-entropy loss (4.1). This is also called logistic regression. The single-layer setting can be easily analyzed via our framework. Our focus here, however, is to demonstrate our analysis on multi-layer neural networks.

To that end, we consider now the same setting, except that we will estimate the class labels using a simple two-layer neural network. (Note that the Bayes’ optimal estimator is still expressible by this architecture.) At first glance, this may seem an elementary setting with little to say. However as we will see, even in this simple setting surprising behaviour can occur in the high-dimensional setting which runs counter to common intuition. Furthermore, as we will see in Section 5, the phenomena occurring here also appear in richer problems such as the XOR problem.

2. Analysis

where $g$ is applied component wise and $p(v,W):=(\alpha/2)(\lvert\lvert v\rvert\rvert^{2}+\lvert\lvert W\rvert\rvert^{2})$.

It can be shown (see Lemma 8.1) that the law of the loss at a given point, $(v,W)\in\mathcal{X}_{n}$, depends only on the $7$ summary statistics,

where $m_{i}=W_{i}\cdot\mu$ and $R_{ij}^{\perp}=W_{i}^{\perp}\cdot W_{j}^{\perp}$ with $W_{i}^{\perp}=W_{i}-m_{i}\mu$ denoting the part of $W_{i}$ orthogonal to $\mu$. For a point, $(v,W)\in\mathcal{X}_{n}$, let

By similar reasoning to Lemma 8.1, it can be seen that these are functions only of $\mathbf{u}_{n}$, and we denote them as such, e.g., $\mathbf{A}_{i}^{\mu}=\mathbf{A}_{i}^{\mu}(\mathbf{u}_{n})$. See Section 8. The critical scaling for $\delta$ is then of order $\Theta(1/n)$ and we obtain the following.

Let $\mathbf{u}_{n}$ be as in (4.2) and fix any $\lambda>0$ and $\delta_{n}=\nicefrac{{c_{\delta}}}{{N}}$. Then $\mathbf{u}_{n}(t)$ converges to the solution of the ODE system, $\dot{\mathbf{u}}_{t}=-\mathbf{f}(\mathbf{u}_{t})+\mathbf{g}(\mathbf{u}_{t})$, initialized from $\lim_{n\to\infty}(\mathbf{u}_{n})_{*}\mu_{n}$, with:

and correctors ${g}_{v_{i}}=g_{m_{i}}=0$, ${g}_{R_{ij}^{\perp}}=c_{\delta}\frac{v_{i}v_{j}}{\lambda}\mathbf{B}_{ij}$ for $i,j=1,2$.

3. Low variance asymptotics

Due to the Gaussian integrals defining $\mathbf{f},\mathbf{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 $\lambda\to\infty$ corresponding to a small noise regime for the Gaussian mixture. We emphasize that this limit is taken after $n\to\infty$ and therefore is still approximately on the critical scale of $\lambda=\Theta(1)$ at which there is a transition in the existence of any fixed point which is a good classifier. In particular, if $\lambda=\lambda_{n}$ is any diverging sequence, then the limiting effective dynamics would exactly match that attained by now sending $\lambda\to\infty$. In Figure 6, we demonstrate numerically that the following predicted fixed points from the $\lambda\to\infty$ limit match those arising at finite large $N$ and $\lambda>0$.For large $\lambda$, this is indeed a quantitative approximation as $\mathbf{f},\mathbf{g}$ exhibit locally Lipschitz dependence on $\lambda^{-1}$, so the corresponding dynamics converges as $\lambda\to\infty$ by classical well-posedness results (see, e.g., )

The $\lambda\to\infty$ limit of the ODE system of Proposition 4.1 is given by

and $\dot{R}_{ij}^{\perp}=-2\alpha R_{ij}^{\perp}$. The fixed points of this system are classified as follows. All fixed points have $R_{ij}^{\perp}=0$ and $m_{i}=v_{i}$ for $i,j=\{1,2\}$. In $(v_{1},v_{2})$, the coordinates are classified by

A fixed point at $(v_{1},v_{2})=(0,0)$ that is stable if $\alpha>\nicefrac{{1}}{{4}}$;

If $\alpha<\nicefrac{{1}}{{4}}$, two unstable sets of fixed points at the quarter-circles given by $(v_{1},v_{2})$ having $v_{1}v_{2}>0$ such that $v_{1}^{2}+v_{2}^{2}=C_{\alpha}$ for $C_{\alpha}:=\log(1-2\alpha)-\log(2\alpha)$.

If $\alpha<\nicefrac{{1}}{{4}}$, two stable fixed points at $(v_{1},v_{2})$ equals $(\sqrt{C_{\alpha}},-\sqrt{C_{\alpha}})$ and $(-\sqrt{C_{\alpha}},\sqrt{C_{\alpha}})$.

If $\mu_{n}$ is e.g., given by $(v_{1},v_{2})\sim\mathcal{N}(0,I_{2})$ and $W_{1},W_{2}\sim\mathcal{N}(0,I_{N}/(\lambda N))$ then $\nu:=\lim(\mathbf{u}_{n})_{*}\mu_{n}$ is $\mathcal{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 $\nicefrac{{1}}{{2}}$ and the basin of attraction of the stable fixed points of (3) with probability $\nicefrac{{1}}{{2}}$.

4. Convergence to spurious solutions

Let us pause to interpret this result. The stable fixed points when $\alpha<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$ (with respect to 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 $\lambda$.

5. Degeneracy of diffusive limits

Two-layer networks for the XOR Gaussian mixture

This data model is a Gaussian mixture model analogue of the (in)famous XOR problem of Minsky–Papert . In particular, it is easy to see that the optimal decision boundary is not expressible by a single-layer neural network as the data is not linearly separable. That said, it is also straightforward to see that this decision boundary is realizable by simple two-layer networks.In the notation of the following subsection, this can be realized by taking $K=4$, $W_{1}=-W_{2}=\mu$, $W_{3}=W_{4}=\nu$, and $v_{i}=c$ for $i=1,\ldots,4$ for some $c>0$.

We focus on this example as a demonstration of the applicability of our techniques to the analysis of the training dynamics for two-layer neural networks on natural data models. While this model is arguably the simplest model requiring a multi-layer network to solve, it nevertheless exhibits very complex phenomenology. We mention that some of these complexities were also observed in a very similar setup in where ballistic limits from warm starts were derived.

2. Analysis

Consider the corresponding classification problem using a two-layer neural network, taking as our estimator of the class label $\hat{y}(X)$ to be the natural rounding of $\sigma(v\cdot g(WX))$, where $\sigma$ and $g$ are the sigmoid and ReLU as in Section 4. We take $W$ to be a $K\times N$ matrix and $v$ to be a $K$-vector.

where again $\sigma,g$ are applied component wise and again $p(v,W):=(\alpha/2)(\lvert\lvert v\rvert\rvert^{2}+\lvert\lvert W\rvert\rvert^{2})$.

In Lemma 9.1 below, we show that the law of the loss at a point $(v,W)$ depends only on the following $4K+\binom{K}{2}$ variables: for $1\leq i\leq j\leq K$,

By similar reasoning, it can be shown that these functions are expressible as functions of $\mathbf{u}_{n}$ alone (see Section 9 below). We then find the following effective ballistic dynamics.

Let $\mathbf{u}_{n}$ be as in (5.1) and fix any $\lambda>0$ and $\delta_{n}=c_{\delta}/N$. Then $\mathbf{u}_{n}(t)$ converges to the solution of the ODE system $\dot{\mathbf{u}}_{t}=-\mathbf{f}(\mathbf{u}_{t})+\mathbf{g}(\mathbf{u}_{t})$, initialized from $\lim_{n}(\mathbf{u}_{n})_{*}\mu_{n}$ with

and correctors $g_{v_{i}}=g_{m_{i}^{\mu}}=g_{m_{i}^{\nu}}=0$, and $g_{R_{ij}^{\perp}}=c_{\delta}\frac{v_{i}v_{j}}{\lambda}\mathbf{B}_{ij}$ for $1\leq i\leq j\leq K$.

3. Low variance asymptotics

As with the binary GMM, one can develop the large $\lambda$ limit of these asymptotics after $n\to\infty$. The effective dynamics in this regime are noticeably more tractable. We defer the precise expressions of these dynamics to Proposition 9.1 below. Let us instead classify the corresponding fixed points.

The fixed points of the ODE system of Proposition 9.1 are classified as follows. If $\alpha>1/8$, then the only fixed point is at $\mathbf{u}_{n}=\boldsymbol{0}$.

If $0<\alpha<1/8$, then let $(I_{0},I_{\mu}^{+},I_{\mu}^{-},I_{\nu}^{+},I_{\nu}^{-})$ be any disjoint (possibly empty) subsets whose union is $\{1,...,K\}$. Corresponding to that tuple $(I_{0},I_{\mu}^{+},I_{\mu}^{-},I_{\nu}^{+},I_{\nu}^{-})$, is a set of fixed points that have $R_{ij}^{\perp}=0$ for all $i,j$, and have

$m_{i}^{\mu}=m_{i}^{\nu}=v_{i}=0$ for $i\in I_{0}$,

$m_{i}^{\mu}=v_{i}>0$ such that $\sum_{i\in I_{\mu}^{+}}v_{i}^{2}=\mbox{logit}(-4\alpha)$ and $m_{i}^{\nu}=0$ for all $i\in I_{\mu}^{+}$,

$-m_{i}^{\mu}=v_{i}>0$ such that $\sum_{i\in I_{\mu}^{-}}v_{i}^{2}=\mbox{logit}(-4\alpha)$ and $m_{i}^{\nu}=0$ for all $i\in I_{\mu}^{-}$,

$m_{i}^{\nu}=v_{i}<0$ such that $\sum_{i\in I_{\nu}^{+}}v_{i}^{2}=\mbox{logit}(-4\alpha)$ and $m_{i}^{\mu}=0$ for all $i\in I_{\nu}^{+}$,

$-m_{i}^{\nu}=v_{i}<0$ such that $\sum_{i\in I_{\nu}^{-}}v_{i}^{2}=\mbox{logit}(-4\alpha)$ and $m_{i}^{\mu}=0$ for all $i\in I_{\nu}^{-}$.

In the $K=4$ case, these form $39$ connected sets of fixed points, and of which $4!=24$ are fixed points that are stable, corresponding to the possible permutations in which each of $I_{\mu}^{+},I_{\mu}^{-},I_{\nu}^{+},I_{\nu}^{-}$ are singletons.

Similar to the binary GMM, in Figures 9–10, we demonstrate numerically that the following predicted fixed points from the $\lambda\to\infty$ limit match those arising at finite large $n$ and $\lambda>0$.

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 $\mu_{n}$ where $v_{i}\sim\mathcal{N}(0,1)$ and $W_{i}\sim\mathcal{N}(0,I_{N}/N)$ independently, this converges to $\nicefrac{{3}}{{32}}$. We refer the reader to Section 9.4 for the proof.

4. 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 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_{\mu}^{+},I_{\mu}^{-},I_{\nu}^{+},I_{\nu}^{-}$ 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. This is a rigorous example of the well-known lottery ticket hypothesis of . Roughly speaking, the lottery ticket hypothesis proposes that the reason for the success of overparametrized networks is that they give more attempts for a sufficiently expressive subnetwork to be initialized well, and succeed at the task on its own.

5. 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}^{\mu}=m_{i}^{\nu}=0$. 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,\mu},a_{2,\mu})$ be such that $a_{1,\mu}^{2}+a_{2,\mu}^{2}=C_{\alpha}$ and $(a_{3,\nu},a_{4,\nu})$ such that $a_{3,\nu}^{2}+a_{4,\nu}^{2}=C_{\alpha}$, for $C_{\alpha}=-\mbox{logit}(4\alpha)$. Take as fixed points about which we expand to be $v_{i}=m_{i}^{\mu}=a_{i,\mu}>0$ and $v_{i}=m_{i}^{\nu}=a_{i,\nu}<0$ for $i=3,4$. Namely, we let

Numerical simulations in Figure 11 confirm these degenerate diffusive limits at finite $\lambda$.

Part III Proofs

In this section, we prove our main convergence result, namely Theorem 2.3. The drift terms can be seen from a Taylor expansion out to second order, with the role played by $\delta_{n}$-localizability being to justify neglecting certain negligible second order terms, as well as all higher order terms. The identification of the stochastic term is via the classical martingale problem for summary statistics of stochastic gradient descent in the high-dimensional $n\to\infty$ limit.

Our aim is to establish $\mathbf{u}_{n}\to\mathbf{u}$ weakly as random variables on $C([0,\infty))$ where $\mathbf{u}$ solves (2.4). It is equivalent to show the same on $C([0,T])$ equipped with the sup-norm for every $T>0$.

Recalling Definition 2.1, since $\mathbf{u}_{n}$ are $\delta_{n}$-localizable, the error term in (6) has

and similarly $b_{j}(s)=\int_{0}^{s}b^{\prime}_{j}(s^{\prime})ds^{\prime}$, then recalling that $\mathbf{u}_{n}(s)$ is the linear interpolation of $(u_{j}([s/\delta]))_{j}$, we may write

where $\mathbf{a}_{n}(s)=(a_{j}(s))_{j}$ and $\mathbf{b}_{n}(s)=(b_{j}(s))_{j}$.

We now prove that the sequence $(\mathbf{u}_{n}(s\wedge\tau_{K}^{n}))$ is tight in $C([0,T])$ with limit points which are $(1/4)$-Holder for each $K$. To this end, let us define

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 $\mathbf{v}_{n}$ instead of $\mathbf{u}_{n}$. We aim to show that for all $0\leq s,t\leq T$,

from which we will get that the sequence $\mathbf{v}_{n}(s\wedge\tau_{K})$ is uniformly $1/4$-Hölder by Kolmogorov’s continuity theorem. Evidently, for all $s,t$ we have that

We control these terms in turn. We will do this coordinate wise and, for readability, fix some $j\leq k$ and let $u=u_{j}$, $a=a_{j}$, $b=b_{j}$ etc.

Let $\mathbf{h}=(h_{j})_{j\leq k}$ be as in (2.2). Then the first term in (6) satisfies

by continuity of $h_{j}$. For the second term in (6),

which is $\lesssim_{K}\delta^{2}(t-s)^{4}$ by items (1)–(2) $\delta_{n}$-localizability. (Applying this bound for $s=0,t=T$, the last term in $a$ is vanishing in the limit for each $K$ whenever $\delta_{n}=o(1)$.) Combining the above bounds yields

For the martingale term, notice that by Burkholder’s inequality,

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 item (3) of $\delta_{n}$-localizability.

For the second term in the martingale difference,

by items (1)–(2) of $\delta_{n}$-localizability. Finally, by the same reasoning, for the third term,

All of the above terms are $O((t-s)^{2})$ since $0\leq s,t\leq T$. Thus we have the claimed (6.2), and by Kolmogorov’s continuity theorem, $(\mathbf{v}_{n}(s\wedge\tau_{K}))_{s}$, are uniformly $\nicefrac{{1}}{{4}}$-Holder and thus the sequence is tight with $\nicefrac{{1}}{{4}}$-Holder limit points. Notice furthermore that if we look at $(\mathbf{v}_{n}(t\wedge\tau_{K})-\mathbf{a}_{n}(t\wedge\tau_{K}))_{t}$, this sequence is also tight and the limits points are continuous martingales. Let us examine their limiting quadratic variations.

Let $\mathbf{v}_{n}^{K}(t)=\mathbf{v}_{n}(t\wedge\tau_{K})$ and define $\mathbf{a}_{n}^{K}(t)$ and $\mathbf{b}_{n}^{K}(t)$ analogously. Furthermore, let $\mathbf{v}^{K}(t)$, $\mathbf{a}^{K}(t)$ and $\mathbf{b}^{K}(t)$ be their respective limits which we have shown to exist and be $\nicefrac{{1}}{{4}}$-Holder.

is a martingale. We therefore need to consider the limit as $n\to\infty$ of the integral above. Write

Consider the integrals of $\delta$ times each of these four terms separately. For the first term,

goes to zero as $n\to\infty$ by the assumption in (2.3).

We now reason that the integrals of $\delta$ times the other three terms in (6.7) all go to zero as $n\to\infty$. The second and third are identical: by Cauchy–Schwarz,

The first expectation contributes $\delta^{-1/2}$ by the first part of item (3) of localizability. Also,

The first of these terms is at most $\delta^{-1}$ as argued in (6). The second is $o(\delta^{-3/2})$ by the second part of item (3) in the definition of localizability. As such, we are able to conclude that

The integral of $\delta$ times the fourth term in (6.7) is handled similarly using Cauchy–Schwarz and the bound of $o(\delta^{-3/2})$ on (6.8).

Thus, if we consider the continuous martingales given by $\mathbf{b}^{K}(t)$, its angle bracket is, by definition, given by

Proofs for matrix and tensor PCA

In this section, we prove the results of Section 3. We will state them in the more general setting where we add a ridge penalty to the loss, so that for $\alpha\geq 0$ fixed, the loss is given by

where $c(Y)$ only depends on $Y$. Note that $H(x)=-2\langle W,x^{\otimes k}\rangle$.

Our first aim is to establish Proposition 3.1, showing that the summary statistics $\mathbf{u}_{n}=(m,r_{\perp}^{2})$ satisfy the conditions of Theorem 2.3 with the desired $\mathbf{f},\mathbf{g}$ and $\Sigma$. We begin by checking localizability for $\mathbf{u}_{n}$. In what follows, for ease of notation we will denote $r^{2}=r_{\perp}^{2}$ and $R^{2}=m^{2}+r^{2}$. In these coordinates,

We check the items in Definition 2.1 one by one, beginning with item (1). Express the derivatives for $\mathbf{u}_{n}$ as

For item (2), differentiating (7.2), $\nabla\Phi=\partial_{1}\phi\nabla m+\partial_{2}\phi\nabla r^{2}$, where

Notice that $\left\langle\nabla m,\nabla m\right\rangle=1,\left\langle\nabla m,\nabla r^{2}\right\rangle=0,$ and $\left\langle\nabla r^{2},\nabla r^{2}\right\rangle=4r^{2}$. Consider

the bounding quantity is evidently a continuous function of $m,r^{2}$ and therefore as long as $x$ is such that $(m,r^{2})\in 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 [9, Lemma 4.7]. Moving on to item (3), by the same reasoning, for every $w$,

If $w=\nabla m=v$ then $\|w\|=1$ and if $w=\nabla r^{2}=2(x-mv)$ then $\|w\|\leq C(K)$, so in both cases this is at most $C(k,K)n^{2}$. Finally, $\nabla^{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(\delta^{-2})$ if $\delta_{n}=O(1/n)$. ∎

Having checked localizability for $\mathbf{u}_{n}$, we apply Theorem 2.3. To compute $\mathbf{f}$, by the above,

In particular, for $\delta=c_{\delta}/n$, we have $\delta\mathcal{L}^{\delta}m=0$ and

from which we obtain in the limit that $n\to\infty$ that $g_{m}=0$ and $g_{r^{2}}=4c_{\delta}kR^{2k-2}$.

Together, these yield the ODE system of (3.1),

which in the $\alpha=0$ case matches Proposition 3.1. Finally, to see that $\Sigma=0$, consider

which when multiplied by $\delta=O(1/n)$ evidently vanishes. ∎

We now turn to analyzing the ODE of Proposition 3.1.

At the fixed points of the ODE in Proposition 3.1,

If $u_{1}=0$, then $R^{2}=u_{2}$ and there are two possible fixed points: either $u_{2}=0$ or $u_{2}$ solves

Notice that if $k=2$, this has a nontrivial solution of the form $c_{\delta}-\frac{\alpha}{2}=u_{2}$, provided $\alpha<\alpha_{c}(2):=2c_{\delta}$, and if $k>2$, this has a nontrivial solution provided $\alpha\leq\max_{x\geq 0}kx^{k-2}(2c_{\delta}-2x)$ at $c_{\delta}(k-2)x^{k-3}-(k-1)x^{k-2}=0$ i.e., $\frac{c_{\delta}(k-2)}{k-1}=x$. This gives

Evidently when we take $\alpha=0$, then its non-trivial solution is at $u_{2}=1$ for all $k\geq 2$.

Alternatively, if $u_{1}\neq 0$ at a fixed point, then we can simplify further and get

For simplicity of calculations, set $\alpha=0$ as is the case in Proposition 3.1. Then, we simply get $u_{2}=c_{\delta}$. In the case of $k=2$, we also find that there is a solution if and only if $\lambda>c_{\delta}$, in which case $R^{2}=\lambda$, from which together with $R^{2}=u_{1}^{2}+u_{2}$, we also get $u_{1}=\pm\sqrt{\lambda-c_{\delta}}$.

In the general case of $k>2$, we find that $R^{2}={c_{\delta}}+\lambda^{-\frac{2}{k-2}}R^{\frac{4(k-1)}{k-2}}$. This has real solutions (all of which have $R\geq u_{2}=c_{\delta}$ as required) whenever $\lambda>\lambda_{c}(k)$ defined as

(Interpreting $0^{0}=1$, this returns $\lambda_{c}(2)=c_{\delta}$.) With this $\lambda$, whenever $\lambda>\lambda_{c}(k)$, the equation for $R^{2}$ has exactly two real solutions, both of which are at least $c_{\delta}$ which we can denote by

2. Effective dynamics for the population loss

In practice, one is interested in tracking the loss, or ideally, the generalization error. In this subsection, we add the generalization error $\Phi$ to our set of summary statistics and obtain limiting equations for its evolution from (3.4).

Recalling (7.2), the fact that $\Phi$ is a localizable summary statistic follows from the facts that $\|\nabla m\|,\|\nabla r^{2}\|\leq C(K)$, and the fact that $\Phi$ is a smooth $n$-independent function of $m,r^{2}$.

For simplicity of calculations let us stick to $\alpha=0$.

Next, consider the corrector for $\Phi$. For this, notice that

Recalling $V$ from (7.4), and taking $\delta=c_{\delta}/n$, all the terms in $\sum_{ij}V_{ij}\partial_{i}\partial_{j}\Phi$ vanish in the limit except the contribution from the $\nabla^{2}r^{2}$, which yields $g_{\Phi}=\lim_{n\to\infty}\delta\mathcal{L}^{\delta}\Phi=4c_{\delta}k^{2}R^{4(k-1)}$ Finally, we wish to compute the volatility for the stochastic part of the evolution of $\Phi$. For this, consider $\nabla\Phi V\nabla\Phi^{T}$ and notice that all the entries of that matrix are continuous functions of $\mathbf{u}_{n}$ and thus go to zero when multiplied by $\delta=O(1/n)$.

3. Diffusive limits at the equator

Taking limits as $n\to\infty$, as long as $\lambda$ is fixed in $n$, we see that $\mathbf{f}$ is given by

4. Diffusive limit for the radius

where we used in the first inequality that the law of $H$ is rotation invariant and $H$ is a $k$-homogenous function. For the second part of item (3),

The $\partial_{i}H$ are Gaussian with mean zero, and by (7.4), variance $C_{k}^{\prime}R^{2(k-2)}x_{i}^{2}+C_{k}R^{2(k-1)}$ and covariance $C_{k}x_{i}x_{j}R^{2(k-2)}$. Recall the following fact about Gaussians: if $X,Y$ are Gaussians with variances $\sigma^{2}$ and covariance $t$, then $\mbox{Cov}(X^{2},Y^{2})\leq Ct^{2}\sigma^{4}$ for some universal constant $C$. Also, $\mbox{Var}(X^{2})\leq C\sigma^{4}$. Applying this to $\partial_{i}H$, we get

Combined with the above, this gives a bound of $n^{2}=O(\delta^{-2})$ on the second part of item (3).

We now calculate the resulting drifts. For $\mathbf{f}$, write

Combining terms and sending $n\to\infty$, we obtain

Multiplying by $\delta=1/n$ and taking the limit as $n\to\infty$, the two entries of this matrix that survive are $\Sigma_{11}$ and $\Sigma_{22}$, where $\Sigma_{11}=4k$ and $\Sigma_{22}=4k(k-1)$. All in all, we obtain (3.5).

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 $\mathbf{u}_{n}$ from (4.2).

Let $X_{\mu}\sim\mathcal{N}(\mu,I/\lambda)$ and $X_{-\mu}\sim\mathcal{N}(-\mu,I/\lambda)$. Then, notice that

Next, notice that as a vector, $(W_{1}X_{\mu},W_{2}X_{\mu})$ is distributed as $(m_{1}+Z_{1,\mu}m_{1}+Z_{1,\perp},m_{2}+Z_{2,\mu}m_{2}+Z_{2,\perp})$, where $Z_{1,\mu},Z_{2,\mu}$ are i.i.d. $\mathcal{N}(0,\lambda^{-1})$, and $Z_{1,\perp},Z_{2,\perp}$ are jointly Gaussian with means zero and covariance

Similarly, the distribution of $WX_{-\mu}$ also only depends on $(m_{i},R_{ij}^{\perp})_{i,j}$. Finally,

Therefore, at any point $(v,W)$, the law of $L((v,W))$, and thus $\Phi$, is simply a function of $\mathbf{u}_{n}(v,W)$. To see that the summary statistics satisfy the bounds of item (1) in Definition 2.1, write $\nabla=(\partial_{v_{1}},\partial_{v_{2}},\nabla_{W_{1}},\nabla_{W_{2}})$. Then

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 $\mathbf{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 $\nabla L$ by direct calculation from (4.1).

(Notice that if $w\in\{\mu,W_{i},W_{i}^{\perp}\}$, then $\mathbf{A}_{i}\cdot w$ is only a function of $\mathbf{u}_{n}$ by the same reasoning as used in Lemma 8.1.) Then, we can also easily express

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 $\delta_{n}$-localizability from Definition 2.1, and the promised limiting equations. The proofs of these are straightforward using Gaussianity and are provided in Section 10 for completeness.

For each $i$, for every $R_{ii}^{\perp}<\infty$ and every $m_{i}>0$, we have

For every $v_{i},R_{ij}^{\perp}$ and $m_{i}\neq 0$ for $i,j=1,2$, we have

Throughout this section we will take $\mu=e_{1}$. By rotational invariance of the problem, this is without loss of generality, and only simplifies certain expressions. The $\delta_{n}$-localizability can be seen by application of the moment bounds listed above.

The condition on $\mathbf{u}_{n}$ was satisfied per Lemma 8.1. Recalling $\nabla\Phi$ from (8.12), one can verify that the norm of each of the four terms in $\nabla\Phi$ is individually bounded, using the Cauchy–Schwarz inequality together with the bound of Lemma 8.2 on $\|\mathbf{A}_{i}\|$.

When $u_{i}$ is $v_{i}$, this is simply a fourth moment bound on $\nabla_{v_{i}}H$, which follows from the $8$’th moment by Jensen’s inequality. When $u_{i}$ is $m_{i}$, or $R_{ij}^{\perp}$, the bound follows from

for choices of $w$ being either $\mu$ in which case $\|w\|=1$ or $W_{i}^{\perp}$ in which case $\|w\|=R_{ii}^{\perp}$. For each $K$, this is at most some constant $C(K)$ using the two bounds of Lemma 8.2.

The convergence of the population drift to $\mathbf{f}$ from Proposition 4.1 follows by taking the inner products of $\nabla L$ from (8.12) with the rows of $J$ from (8.8), and noticing that $\mathbf{A}_{i}^{\mu}$ from (4.3) is exactly $\mathbf{A}_{i}\cdot\mu$ and $\mathbf{A}_{ij}^{\perp}$ from (4.3) is exactly $\mathbf{A}_{i}\cdot W_{j}^{\perp}$.

Next consider the convergence of the correctors to the claimed $\mathbf{g}$. The variables $u\in\{v_{1},v_{2},m_{1},m_{2}\}$ are linear so $\mathcal{L}_{n}u=0$ and for these, $\mathbf{g}_{u}=0$. For $u=R_{ij}^{\perp}$ for $i,j\in\{1,2\}$, the relevant entries in $V$ are those corresponding to $W_{i}^{\perp}$ and $W_{j}^{\perp}$. For ease of notation, in what follows let $\pi=\sigma(v\cdot g(WX))$.

For ease of calculation taking $\mu=e_{1}$, we have $\mathcal{L}_{n}R_{ij}^{\perp}=\sum_{k\neq 1}V_{W_{ik},W_{jk}}$, which by (8), and the choice of $\delta_{n}=c_{\delta}/N$, is given by

which we emphasize is only a function of $\mathbf{u}_{n}$. We lastly need to show that the diffusion matrix $\Sigma_{n}$ goes to zero as $n\to\infty$ when $\delta_{n}=O(1/n)$. This is straightforward to see by considering any element of $JVJ^{T}$ and using Cauchy–Schwarz together with the two bounds of Lemma 8.2 to bound it in absolute value by some $C(K)$ independent of $n$. Then when multiplying by any $\delta_{n}=o(1)$, this entire matrix will evidently vanish. ∎

2. The small-noise limit of the effective dynamics

One can now take a $\lambda\to\infty$ limit to arrive at the ODE system of Proposition 4.2.

We begin with considering $\lim_{\lambda\to\infty}\mathbf{A}_{i}^{\mu}$: its limiting value will depend on the signs of both $m_{1}$ and $m_{2}$. We can express $\mathbf{A}_{i}^{\mu}$ from (4.3) as

We claim that the two terms on the right-hand side converge to $-\frac{1}{2}\mathbf{1}_{m_{i}>0}\sigma(-v\cdot g(m))$ and $-\frac{1}{2}\mathbf{1}_{m_{i}<0}\sigma(v\cdot g(-m))$ respectively. This follows by e.g., writing the difference as

at which point, we see that if $m_{1},m_{2}\geq 0$, this becomes $\frac{1}{2}\sigma(-v\cdot m)$, as it is if $m_{1},m_{2}\leq 0$. If $m_{1}\geq 0$ and $m_{2}\leq 0$, then you get $\lim_{\lambda}\mathbf{A}_{1}^{\mu}=-\frac{1}{2}\sigma(-v_{1}m_{1})$ and $\lim_{\lambda}\mathbf{A}_{2}^{\mu}=-\frac{1}{2}\sigma(-v_{2}m_{2})$ and likewise if $m_{1}\leq 0$ and $m_{2}\geq 0$.

Next consider the limit as $\lambda\to\infty$ of $\mathbf{A}_{ij}^{\perp}$ from (4.3), which we claim converges to . Write

Finally, since $|\mathbf{B}_{ij}|\leq 1$, the quantity $g_{R_{ij}^{\perp}}=c_{\delta}\frac{v_{i}v_{j}}{\lambda}\mathbf{B}_{ij}$ evidently goes to zero as $\lambda\to\infty$. ∎

The above argument used $m_{i}\neq 0$ for the limit of $\mathbf{A}_{i}^{\mu}$. If one considers the cases when $m_{i}=0$, the limiting drifts still apply. For this, it suffices to show that if $m_{i}=0$, then $\mathbf{A}_{i}^{\mu}$ converges to zero. Without loss of generality, suppose $m_{1}=0$ and consider

This is zero independently of $\lambda$ by independence of $Z_{1,\mu}$ from the other Gaussians in the expectation.

Evidently, every fixed point must have $R_{ij}^{\perp}=0$. Furthermore, if we let $u_{i}=v_{i}-m_{i}$, then

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 to $v_{i}\sigma(-\|v\|^{2})=2\alpha v_{i}v_{1}v_{2}>0$ if $v_{1}v_{2}>0$ and $v_{i}\sigma(-v_{i}^{2})=2\alpha v_{i}$ otherwise. Observe first that the point $(v_{1},v_{2})=(0,0)$ is a fixed point of this system. If $(v_{1},v_{2})\neq 0$, then dividing out by $v_{i}$, the above reduces to $\sigma(-\|v\|^{2})=2\alpha$ if $v_{1}v_{2}>0$ and $\sigma(-v_{i}^{2})=2\alpha$ otherwise. Recalling that $C_{\alpha}=-\operatorname{logit}(2\alpha)=\log(1-2\alpha)-\log(2\alpha)$ we obtain the claimed set of fixed points by inverting these equations (they only have a solution if $\alpha<1/4$).