High-dimensional SGD aligns with emerging outlier eigenspaces
Authors: Gerard Ben Arous
Reza Gheissari
Jiaoyang Huang
Aukosh Jagannath
Venue: ICLR 2024 spotlight (2024)
Keywords: stochastic gradient descent
Hessian
multi-layer neural networks
high-dimensional classification
Gaussian mixture model
XOR problem
Primary area: learning theory

Abstract: We rigorously study the joint evolution of training dynamics via stochastic gradient descent (SGD) and the spectra of empirical Hessian and gradient matrices. We prove that in two canonical classification tasks for multi-class high-dimensional mixtures and either 1 or 2-layer neural networks, the SGD trajectory rapidly aligns with emerging low-rank outlier eigenspaces of the Hessian and gradient matrices. Moreover, in multi-layer settings this alignment occurs per layer, with the final layer's outlier eigenspace evolving over the course of training, and exhibiting rank deficiency when the SGD converges to sub-optimal classifiers. This establishes  some of the rich predictions that have arisen from extensive numerical studies in the last decade about the spectra of Hessian and information matrices over the course of training in overparametrized networks.

Published as a conference paper at ICLR 2024
HIGH-DIMENSIONAL SGD ALIGNS WITH EMERGING OUTLIER EIGENSPACES
Gérard Ben Arous
Reza Gheissari
Courant Institute
Department of Mathematics
New York University
Northwestern University
New York, NY, USA
Evanston, II., USA
ciiarousscims.oo wyw. nyu.
g
Jiaoyang Huang
Aukosh Jagannath
Department of Statisties and Data Science University of Pennsylvania
Department of Statistics and Actuarial Science Department of Applied Mathematics, and Cheriton School of Computer Science University of Waterloo Waterloo, ON, Canada a.
Philadclphia, PA, USA
muangjy@whare upenn.
ABSTRACT
We rigorously study the joint evolution of training dynamics via stochastic gradi- ent descent (SGD) and the spectra empirical Hessian and gradient matrices. We prove that in two canonical classification tasks for multi-class high-dimensional mixtures and either or 2-layer neural networks, the SGD trajectory rapidly aligns with emerging low-rank outlier eigenspaces of the Hessian and gradient matrices. Moreover, in multi-layer settings this alignment occurs per layer, with the final layer's outlier eigenspace evolving over the course training, and exhibiting rank deficiency when the SGD converges to sub-optimal classifiers. This establishes some of the rich predictions that have arisen from extensive numerical studies the last decade about the spectra of Hessian and information matrices over the course of training in overparametrized networks.
INTRODUCTION
Stochastic gradient descent (SGD) and its many variants, are the backbone of modern machine learning algorithms (see e.g., Bottou (199))). The training dynamics of neural networks, however, are still poorly understood in the non-convex and high-dimensional settings that are frequently en- countered. common explanation for the staggering success of neural networks, especially when overparameterized, is that the loss landscapes that occur in practice have many "flat" directions and a hidden low-dimensional structure within which bulk of training occurs.
To understand this belief, much attention has been paid to the Hessian of the empirical risk (and related matrices formed via gradients) along training. This perspective on the training dynamics of neural networks cet al. (2017a The upshot of these studies was broad understanding of the spectrum of the Hessian of the empirical risk, that summarize follows:
proposed in LeCun et al. (2012), and numerically analyzed in depth in Sagun
1. It has a bulk that is dependent on the network architecture, and is concentrated around O, becoming more-so as the model becomes more overparametriiee:
2. It has (relatively) few outlier eigenvalues that dependent on the data, and evolve non- trivially along training while remaining separated from the bulk;
Since those works, these properties of the Hessian and related spectra over the course of training have seen refined and large-scale experimentation. Papyan| (2019) found hierarchical de- composition to the Hessian for deep networks, attributing the bulk, emergent outliers, and a minibulk three different "parts" of the Hessian. Perhaps most relevant to this work, Gur-Ari al. (2019)

---

Published as a conference paper at ICLR 2024
noticed that gradient descent tends to quickly align with a low-dimensional outlier subspace of the Hessian matrix, and stay in that subspace for long subsequent times. They postulated that this com- mon wwwddimensional structure to the SGD and Hessian matrix may be key to many classification tasks in machine learning. For a sampling of other empirical investigations of spectra of Hessians and information matrices along training. see e.g.. [Ghorbani et al. (2019); Papyan (2020): Li et al.)
Martin & Mahoney (2019); Cohen et al. (202I); Xie et al. (2023).
From a theoretical perspective, much attention has been paid to the Hessians of deep networks using random matrix theory approaches. Most of this work has focused on the spectrum at a fixed point in parameter space, most commonly at initialization. Early works in the direction include Watanabe (2007): et al. (2014). Chhromanska et al. (2015) noted similarities of neural net Hessians to spin glass Hessians, whose complexity (exponential numbers of critical points) has been extensively studied see e.g.. Auffinger & Ben Arous (2013): (Affinger et al. (2013). More recently, the expected complexity has been investigated statistical tasks like tensor PCA and generalized linear estima- tion Ben Arous et al. (2019); Maillard et al. (222)). In Pennington & Worah (2018) and Pennington & Bahri (2017), the empirical spectral distribution of Hessians and information matrices in single layer neural networks were studied at initialization. Liao & Mahoney (202T) studied Hessians of some non-linean models, which they referred to as generalized GLMs, also at initialization.
Under the infinite-wide neural tangent kernel limit of]acot et al. (2018). 8) Fan & Wang (2020) derived the empirical spectral distribution at initialization, and et al. (2020) studied its evolution over the course of training. In this limit the input dimension is kept fixed compared to the parameter dimension, while our interest in this paper is when the input dimension, parameter dimension, and number of samples all scale together.
An important step towards understanding the evolution of Hessians along training in the high- dimensional setting, is of course understanding the training dynamics themselves. Since the clas- sical work of Robbins & Monro (1911). there has been much activity studying limit theorems for stochastic gradient descent. In the nigh-dimensional setting. following [Saad & Solla (1955a b). in- vestigations have focused on finding finite number of functions (sometimes called "observables" or "summary statistics"), "), whose dynamics under the SGD are asymptotically autonomous in the high- dimensional limit. For necessarily small sampling of this rich line of work, we refer to Goldt et al.) (2019); Veiga et al. (2022); Paquette et al. (202T): Arnaboldi et al. (2023); Tan & Vershynin (2019); Ben Arous (2022). Of partieular relevance, it shown by Damian al. (2022) Mousavi Hosseini et al. (202)) that for multi-inder models, SGD predominantly lives in the low- dime-iinniooal
subspace spanned by the ground truth parameters.
A class of tasks whose SGD is amenable to this broad approach is classification of Gaussian mixture models (GMMs). With various losses and linearly separable class structures, the minimizer of the empirical risk landscape with single-layer networks was studied in Miggnacco et al. (2020): Loureiro al. (2021). A well-studied case of Gaussian mixture model needing two-layer network is under an XOR-type class structure; the training dynamics of SGD for this task were studied in Refinetii al. (20TT) and [en Arous al. (2022) and was found to have a particularly rich structure with positive probability of convergence to bad classifiers among other degenerate phenomena.
Still, a simultaneous understanding of high-dimensional SGD and the Hessian and related matrices' spectra along the training trajectory has remained largely open.
1.1 OUR CONTRIBUTIONS
In this paper, study the interplay between the training dynamics (via SGD) and the spectral compositions of the empirical Hessian matrix and an empirical gradient second moment matrix, or simply G-matrix (2.2) (similar in spirit information matrix) over the course of training- We rigorously show the following phenomenology in two canonical high-dimensional classification tasks with k "hidden" classes:
1. Shortly into training, the empirical Hessian and empirical G-matrices have C(k) many outlier eigenvalues, and the SGD itself predominantly lives in their eigenspace. Here ((k) is explicit and depends on the performance of the classifier to which the SGD converges.
2

---

To demonstrate results in more complex multi-layer settings, consider supervised classifica- tion of a GMM version of the famous XOR problem of Minsky & Papert (1969). This is one of the two-layer architecture Published as a conference paper at ICLR 2024
2. In multi-layer settings, this alignment happens within each layer, i..., the first layer parame- ters align with the outlier eigenspaces of the corresponding blocks of the empirical Hessian and G-matrices, and likewise for the second layer parameters.
3. This alignment is not predicated on success at the classification task: when the SGD con- verges to sub-optimal classifier, the empirical Hessian and G matrices have lower rank outlier eigenspaces, and the SGD aligns with those rank deficient spaces.
The first model we consider is the basic example of supervised classification of general k-component Gaussian mixture models with & linearly independent classes by single layer neural network. In Theorem 3.1 we establish alignment of the form of Item above, between each of the & one-vs-all classifiers and their corresponding blocks in the empirical Hessian and G-matrices. See also the depictions in Figures 3.2 In order to show this, we show that the matrices have an outlier- minibulk-bull structure throughout the parameter space, and derive limiting dynamical equations for the trajectory of appropriate summary statistics of the SGD trajectory following the approach Ben Arous et al. (2022). Importantly, the same low-dimensional subspace is at the heart of both the outlier eigenspaces and the summary statistics. At this level of generality, the SGD can behave very differently within the outlier eigenspace. As an example of the refined phenomenology that can arise, we further investigate the special case where the means specifically with the single largest outlier cigenvalue, which itself has separated from the other outliers along training. This is proved in Theorem 3.3 and depicted in Figure 3.3 These results are presented in Section
orthogonal: here the SGD aligns
simplest models that requires a two-layer neural network to solve. We with a second layer of width K. As indicated by Item 2 above, in Theorems 4.1 4.2 the alignment of the SGD with the matrices' outlier eigenspaces occurs within each layer, the first layer having an outlier space of rank two, and the second layer having an outlier space of rank 4 when the dynamics converges to an optimal classifier. This second layer's alignment is especially rich, as when the model is overparametrized (K large), its outlier space of rank 4 is not present at initialization, and only separates from its rank-K bulk over the course of training. This can be interpreted as a dynam- ical version of the well-known spectral phase transition in spiked covariance matrices of Baik et al. (2005): see Figure for a visualization. Moreover, the SGD for this problem is known to converge to sub-optimal classifiers with probability bounded away from zero Ben Arous et al. (2022). and find that in these situations, the aligmment still occurs but the outlier eigenspaces into which the SGD moves rank-deficient compared to the number of hidden classes, 4: see Figure 4.3 These results are presented Section
2 PRELIMINARIES
Before turning to our results, let us begin by introducing the following framework and notation that will be relevant throughout this work. We suppose that we given data from distribution Py over pairs Y = Y) where y E R is a one-hot "label" vector that takes the value 1 on a class (sometimes identified with the element of [k] (1,...,k) k) which it is 1), and Y E Bd is a corresponding feature vector. In training we take as loss a function of the form L(x, Y) Iº * + R where E IR' represents the network parameter. (As we are studying supervised classification, in both settings this loss will be the usual cross-entropy loss corresponding to the architecture used.)
We imagine we have two data sets, a training set (Y*) and a test set (Y all drawn i..dd from Py. Let us first define the stochastic gradient descent trained using )) In order to ensure the SGD doesn't 't go off to infinity we add E? penalty term (as is common in practice) with Lagrange multiplier/ B. The (online) stochastic gradient descent with initialization Xy and learning rate, or step-size, 6, will be run using the training set (Y 1 as follows:
)==
- LV Y*) 8x--1
(2.1)
Our aim is to understand the behavior of SGD with respect to principal subspaces, i.c., outlier eigenvectors, of the empirical Hessian matrix and empirical second moment matrix of the gradient.
3

---

Published as a conference paper at ICLR 2024
This latter matrix is exactly the information matrix when L is the log- legliiklihood;; in our paper L is taken to be a cross-entropy loss, so we simply refer to this as the G-matrix henceforth. We primarily consider the empirical Hessian and empirical G-matrices generated out of the test data, namely:
M
M
4 {²
L' 8=1
Y
and ((x)
1 M 1=1
Σ
VL(x
Y2)02.
(2.2)
(Notice that in our online setting, is as natural to generate these matrices with test data as with train data. See Remark [.1 how our results extend when training data is used.) When the parameter space naturally splits into subsets of its indices (e.g., the first-layer weights and the second-layer weights), for subset / of the parameter coordinates, we use subscripts V7, and Gr, to denote the block corresponding to that subset.
To formalize the notion alignment between the SGD and the principal directions of the Hessian and G-matrices, introduce the following language. For subspace B, we let PB denote the orthogonal projection onto B: for a vector v, we let ||a|| be its l norm; and for a matrix A, let ||^||| = || op be its (2 l2 operator norm.
Definition 2.1. The alignment of a vector v with a subspace B is the ratio p(e, B) = ||P We say a vector lives in subspace B up to error if p(v, B) 2
For matrix A, we let (A) denote the span of the top & eigenvectors of A, i.e., the span of the k cigenvectors of A with the largest absolute values. We also use the following.
Definition 2.2. We say matrix A lives in subspace B up error e if there exists M such that Im(A- M) C B with |op I Allop where |Aloop denotes the f2-10-22 operator norm.
3 CLASSIFYING LINEARLY SEPARABLE MIXTURE MODELS
We begin by illustrating our results on (arguably) the most basic problem of high-dimensional multi- class classification, namely supervised classification of i component Gaussian mixture model with constant variance and linearly independent means using a single-layer network. (This is sometimes used toy model for the training dynamics of the last layer of a deep network via the common ansatz that the output of the second-to-last layer of a deep network behaves like a linearly separable mixture of Gaussians: see e.g., the neural collapse phenomenon posited by Papyan al. (020).)
3.1
DATA MODEL
Let C = [k] be the collection of classes, with corresponding distinct class means (Ha)aejs E Rdd co- variance matrices la/A, A, where À > 0 can be viewed signal-to to-iose parameter, and corresponding probabilities 0 < (Pa)aEja)
such that Σee(k] Pa The number of classes & O(1) is fixed (here and throughout the paper o(1), O(1) and ?((1) notations are with respect to the dimension parameter d, and may hide constants that are dimension independent such as k, (3).
For the sake of simplicity we take the means to be unit norm. Further, in order for the task to indeed be solvable with the single-layer architecture, we assume that the means are linearly independent, say with a fixed (i.e., d-independent) matrix of inner products (Mat)a,b = (Ju [b))a,-- a. Our data distribution Py is mixture of the form E. PoN (Hey la/A), with an accompanying class label y E & Namely, if Zx NNO0 Ia/A), our data is given Y (v. Y) where:
y ~ Σ Pao1.: and Y Σ YaHa + Zx-
(3.1)
aE(k]
ac(A]
We perform classification by training single layer network formed by k "all-vs-one" classifiers using the cross entropy loss (cquivalently, we are doing multi-class logistic regression):
L(x = Σ yor Y + log
exp(x . Y),
(3.2)
cETA]
cE(M]
where x = )oE) are the parameters, each which is a vector in Rd i.e. x E Rilk, (Note that we can alternatively view l x matrix.)
4

---

Published as a conference paper at ICLR 2024
Correlation with princippal eigerepaces
Correlation with pricipal exgenspaces of Hessian
69
05
09
07-
05-
85-
04
04
03
02
02
35
12
23
Figure 3.1: The alignment of the SGD trajectory x with Eg (V2 R(xx)) (left) and Eg (G (x,)) (right) over time, for I,..., (shown in different colors). The æ--axis is rescaled time, t8. The parameters here are k = 10 classes in dimension d = 1000 with = 10, B = 0.01, and = 1 1/d.
oz
3.2 RESULTS AND DISCUSSION
Our first result is that after some linearly many steps, the SGD finds the subspace generated by the outlier eigenvalues of the Hessian and/or G- G-matrix of the test loss and lives there for future times.
Theorem 3.1. Consider the mixture of k-Gaussians with loss function from (3.2). and SGD (2.)) with learning rate 8 = O(1/d) /d), regularizer 3 > O, initialized from ( Ia/d). There exists on, Ao such that if 2 Ao, and M 2 and, the following hold. For every >> there exists To(e) such that for any time horizon To < T < M/d, with probability 1 - -oa(1),
x lives in Eg (Vxx)(xx)) and in E Ex(C)((X)),,
for every c E [ki], up to O( + X error, for all
This result is demonstrated in Figure 3.1 which plots the alignment of the training dynamics xf with the principal eigenspaces of the Hessian and G for each e E [k]]. As we see the alignment increases to near 1 rapidly for all blocks in both matrices. This theorem, and all our future results, are stated using a random Gaussian initialization, scaled such that the norm of the parameters is O(1) in d. The fact that this is Gaussian is not relevant to the results, and similar results hold for other uninformative initializations with norm of (1)).
Theorem 3.1 follows from the following theorem that describes the SGD trajectory, its Hessian and G-matrix (and their top l eigenspaces). all live up to O(1/A) error in Span(/11,... 14x)-
Theorem 3.2. In the setup of Theorem 3.1 1
the following live in Span([11 -- (_))))">>))))))))))))))))))))>)> Pk) up O( + A-1)
error with probability 1 o1(1):
1. The state of the SGD along training, xi every c;
2. The b, c blocks of the empirical test Hessian, R(xx) for all b, [k];
3. The b, c blocks of the empirical test G-matrix (re(Xe) for all b, E [k].
We demonstrate this result in Figure 3.2
Inside the low-rank space spanned by μi..... [4. the training dynamics, Hessian and G-matrix spec- tra can display different phenomena depending on the relative locations of 41,,.. P and weights Pi....,e Pk- To illustrate the more refined alignment phenomena, let us take as a concrete example Po for all c, and -i,-.., P4 orthonormal.
With this concrete choice, we can analyze the limiting dynamical system of the SGD without much difficulty and its relevant dynamical observables have a single stable fixed point, to which the SGD converges in linearly many, i.e., O steps. This allows us to show more precise alignment that occurs within Span(p1,.... Hx) over the course of training.
Theorem 3.3. In the setting of Theorem 3.1 with the means (Hi, (b) being orthonormal. the estimate xi has ((1), positive, inner product with the top eigenvector of both V2 R(xx) and Gec(xe) (and negative, (111) inner product with the k- next largest eigenvectors). Also, the top eigenvector of 2 R(xx)), as well as that of Geo(Xx) live in Span(fic) up O + x-1 error.
5

---

Published as a conference paper at ICLR 2024
11 thec ef Meston speewedore
300
475
035
1.3
235
800 250 3000
300
400
1300
300
Figure 3.2: From left to right: Plot of entries of x and the &o leading eigenvectors (in different
colors) of Vt R(xe) and (xx( respectively at the end of training, namely & 50 d 25, 000 steps. Here the z-axis represents the coordinate index. The parameters are the same as in Fig. 3.1 and the means are μ4 3j-50-
Eigenvaiie trajectones
product
050
0000
0.05
0.25
Figure 3.3: Left: the eigenvalues (in different colors) of V (1xe)) over the course training. The leading cigenvalues separated from the bulk all times, and the top eigenvalue, corresponding to in separates from the remaining eigenvalues soon after initialization. Right: the inner product of x! with the means /1.-.-,/ (4g undergoes a similar separation over the course of training. Parameters are the same as in preceding figures.
Put together, the above three theorems describe the following rich scenario for classification of the k-GMM. At initialization, and throughout the parameter space in each class-block, the Hessian and G-matrices decompose into rank-k outlier part spanned by ji...., Mik. and a correction term of size O(1/A) in operator norm. Furthermore, when initialized randomly, the SGD not aligned with the outlier eigenspaces, but does align with them in short 0 number of steps. Moreover, when the means are orthogonal, each class block of the SGD xi in fact correlates strongly with the specific mean for its class Me, and simultaneously. in the Hessian and G-matrices along training, the eigenvalue corresponding to Pe becomes distinguished from the other k outliers. We illustrate these last two points in Figure 3.3 We also refer the reader Section F for further numerical demonstrations. Each of these phenomena appear in more general contexts, and in even richer manners, as we will see in the following section.
CLASSIFYING XOR-TYPE MIXTURE MODELS VIA TWO-LAYER NETWORKS
With the above discussion in mind, let us now turn more complex classification tasks that linearly separable and require the corresponding network architecture to be multilayer.
4.1
DATA MODEL
For our multilayer results, we consider the problem of classifying 4-component Gaussian mix- ture whose class labels are in a so-called XOR form. More precisely, consider mixture of four Gaussians with means - -v where || = ||-||| and variances I, There are two classes, class label 1 for Gaussians with mean Iu, and 0 for Gaussians with mean tv. To be more precise, our data distribution Py is
and, say for simplicity, are orthogonal,
and Y [W(,, + (V(-H,1/AN)
y=1
y ~ 00 + 2 2
(4.1)
(iv, + NN(-,,
y y=

---

Published as a conference paper at ICLR 2024
This is Gaussian version of the famous XOR problem of Minsky & Papert (1969). It is one of the simplest examples of a classification task requiring a multi multi--ayer network to express a good classifier.
We therefore use two-layer architecture with the intermediate layer having width A" (any less and a Bayes-optimal classifier would not be expressible), ReLu activation function g a x V (W,v)
and then sigmoid activation function (xx) Itc-z- The parameter space is then
I @Kd+K K where the first layer weights are denoted by W = (x) E Rxxd and the second layer weights are denoted by xxx E K We use the binary cross-entropy loss on this problem,
Lxx Y) yy g(Y) log(1 e">)())))
(4.2)
with applied entrywise. The SGD for this classification task was studied in some detail in Refinetti al. 2021) and Ben Arous et al. (2022), with "critical" step size 5 = (1/dd).
4.2 RESULTS AND DISCUSSION
We begin our discussion with the analogue of Theorem 3.1 in this setting. As the next theorem demonstrates, in this more subtle problem, the SGD still finds and lives in the principal directions of the empirical Hessian and G-matrices] | Here, the principal directions vary significantly across the parameter space, and the alignment phenomenon differs depending on which fixed point the SGD converges to. Moreover, the relation between the SGD and the principal directions of the Hessian and G- G-matrices can be seen per layer.
Theorem 4.1. Consider the XOR GMM mixture with loss function (4.2) and the corresponding SGD 2.1 with 3 E (0, 1/8), learning rate ê = O(1fd), initialized from N (0, Ia/d) There exist ag, Ao such that if X 2 Ao, and M "2 ond, the following hold. For every >> there exists To(e) such that for any time horizon To < TT < M/d, with probability - oa(1), for all i E {1, K).
1. (x(x)) lives in Ez Vii R(xx)) and in E2(G (xi)). and
2. (xxe) lives in EV( R(xe)) and (xvo((
up to O(E + - + error, for all é [T65- 1,,,8-1111
Remark 1. The restriction to 3 < 1/8 in Theorems 4.1 4.2 is because when B 1/8 the regular- ization is too strong for the SGD to be meaningful; in particular, the SGD converges ballistically to the origin in parameter space, with no discernible preference for the directions corresponding to μr, as the other directions. The above theorems are still valid there if the notion of error for living in space from Definition 2.1 were additive instead of multiplicative (i.e., at most e rather than el|e!]).
This theorem is demonstrated in Figure 4.1
There we have plotted the alignment of the rows in the intermediate layer with the space spanned by the top two eigenvectors of the corresponding first-layer blocks of the Hessian and G-matrices, and similarly for the final layer.
As before, the above theorem follows from the following theorem that describes both the SGD trajectory, its Hessian, and its G-matrix, living up to O(c X-1/2) error in their first-layer blocks in Span(u,, 2) and (applying g entrywise) in their second layer blocks in
Span(g (W ((pp)))).((-)))) (X(()))))), (
Theorem 4.2. In the setting of Theorem 4.1 up to O + X 1/2) error with probability 1 - oa(1), the following live in Span(p, (u, v).
The first layer weights, Wx(xe) for each i E {1,..., K}.
The first-layer empirical test Hessian Vi.w, R((x)) for each i E {l,..,
The first-layer empirical test G-matrix Gw,w, (xx) for each E (1,..., K}. and the following live in Span(gl (oon(g(We) (W( (())))
As the second derivative of ReLU only defined in the sense of distributions, care must be taken when studying these objects. This singularity is not encountered by the SGD trajectory almost surely.
In this theorem and in Theorem 4.2 the big-O notation also hides constant dependencies the initial magnitudes of v o(xx), and Ty.
7

---

Published as a conference paper at ICLR 2024
Nig yoar wit Migespun
Agoret proppd cogeaper
(a)
(b)
(c)
(d)
Figure 4.1: (a)-b) depict the alignment of the first layer weights (x(x) ) for i = 1,.,, K (in differ- ent colors) with the principal subspaces of the corresponding blocks of the Hessian and G-matrices, i.e, with VaWW.. R(xx)) and F2, w. (xx)). (c)-(d) plot the second-layer alignment, namely of (xxe) with (V. R(xe)) and E (G Parameters are d 1000, À 10, and K 20
Traectoones vergenvalies
Trojetories of w eigenvalues of llesaan
0000
0.00
0005
0400
0015
0.04
0010
0005
0.00
0.00
Figure 4.2: The eigenvalues (in different colors) of the ov blocks of the Hessian and G-matrices over time from random initialization. Initially, there is one outlier eigenvalue due to the positivity of the ReLU activation. Along training. four outlier cigenvalues separate from the bulk, corresponding to the four "hidden" classes in the XOR problem. Parameters are the same in Figure 4.1
The second layer weights vxxe),
The second-layer empirical test Hessian V R(xe),
The second-layer empirical test G-matrix Geo (Xe).
Let us discuss bit more the phenomenology of the alignment in the second layer. First of all, we observe that the subspace in which the alignment occurs is a random depending on the initial- ization and trajectory (in particular, choice of fixed point the SGD converges to) 4-dimensional subspace of R Furthermore, if we imagine K to be much larger than 4 so that the model is overparametrized, at initialization, unlike the 1-layer case studied Section 3 the Hessian and G- matrices do not exhibit any alignment in their second layer blocks. In particular, the second layer blocks of the Hessian and G-matrices look like (non-spiked) Gaussian orthogonal ensemble and Wishart matrices5) in K dimensions at initialization, and it is only over the course of training that they develop 4 outlier eigenvalues as the first layer of the SGD begins to align with the mean vectors (xx)))))oEE(
and the vectors lization of the last layer around these vectors over the course of training is reminiscent of the neural collapse phenomenon outlined in Papyan et al. 2020) (see also Han et al. (022)) Zhu et al. (2021) )). The simultaneous emergence, along training, of outliers in the Hessian and G-matrix spectra can be seen as dynamical version of what is sometimes referred to as the BBP transition after Baikk et al. (005) (see also Péché (2006)). This dynamical transition is demonstrated in Figure 4.2
in turn get large enough generate outliers. This crystal-
Finally, we recall that Ben Arous et al. (2022) found a positive probability (uniformly in d, but shrinking as the architecture is overparametrized by letting K grow) that the SGD converges to sub- optimal classifiers from a random initialization. When this happens, g (x+)(0 remains small for the hidden classes & E { Зッ} that are not classifiable with SGD output. In those situations, Theorem 4.2) shows that the outlier subspace in the vu-blocks that emerges will have rank smaller than 4, whereas when optimal classifier is found it will be of rank 4: see Figure 4.3 Knowing that
'These are two of the most classical random matrix models, Anderson et al. 2010 for more.
8

---

Published as a conference paper at ICLR 2024
Yeatte
100
400
600
400
400
Figure 4.3: Evolution of eigenvalues in the v component of G over time in rank deficient cases. Here SGD is started from initializations that converge to suboptimal classifiers (this has uniformly positive, K K-dependent, probability under random initialization). From left to right, the SGD's classifier varies the number of hidden classes discerns, from to 4. There is still dynamical spectral transition, now with only a corresponding number of emerging outlier eigenvalues.
the classification task entails a mixture of 4 means, this may provide a method for devising a stopping rule for classification tasks of this form by examining the rank of the outlier eigenspaces of the last layer Hessian or G-matrix. While the probability of such sub-optimal classification is bounded away from zero, the probability exponentially goes to zero as the model is overparametrized via K + oo, and that gives more chances to allow the second layer SGD, Hessian, and G-matrices to exhibit full 4-dimensional principal spaces, This serves as a concrete and provable manifestation of the lottery ticket hypothesis of Frankle & Carbin 2019).
5 OUTLINE AND IDEAS OF PROOF
The proofs of our main theorems break into three key steps.
1. In Sections A we show that the population Hessian and G-matrices, have bulks (and possibly minibulks) that are ((1AA) in operator norm, and finite ((k)-rank parts. We can explicitly characterize the low-rank part's eigenvalues and eigenvectors up to O(1/A) corrections, functions of the parameter space, see exactly where their emergence as outliers occurs depending on the model.
2. In Section C
we analyze the SGD trajectories for the k-GMM and XOR classification problems. We do this using the limiting effective dynamies theorem proven in Ben Arous et al.) (2022) for finite families of summary statistics. We derive ODE limits for these summary statistics, notably for the general k-GMM which was not covered that paper: see Theorem C.7 We then pull back the limiting dynamics to finite d and expand its X- finite trajectory about its À = oo solution. These latter steps involve understanding some of the stability properties of the dynamical system limits of the SGD's summary statistics.
3. In Section D we prove concentration for the empirical Hessian and G-matrices about their population versions, in operator norm, throughout the parameter space. In some related set- tings, including binary mixtures of Gaussians, Mei et al. (2018) established concentration of the empirical Hessian about the population Hessian uniformly in the parameter space assuming polylogarithmic sample complexity. Our proofs are based on c-nets and concen- tration inequalities for uniformly sub-exponential random variables, together with some twists due, for instance, to non-differentiabiiit of the ReLU function.
Remark 5.1. Our results are stated for the empirical Hessian and G-matrices generated using test data, along the SGD trajectory generated from training data. Since we are considering online SGD, the empirical Hessian and G-matrices with training data are no more relevant than those generated from test data, but the reader may still wonder whether the same behavior holds. A straightforward modification of our arguments in Section D extends our results for the k-GMM model to Hessian and G-matrices generated from train data, if we assume that M Z dllog d rather than simply M Ž d. See the full version for a sketch of this extension. It is an interesting mathematical question to drop this extra logarithmic factor. The extension in the XOR case is technically more involved due to the regularity of the ReLU function. See Section |F for numerical demonstrations that the phenomena in the empirical matrices generated from train data are identical to those generated from test data.

---

Published as a conference paper at ICLR 2024
AUTHOR CONTRIBUTIONS
All authors contributed equally this work.
ACKNOWLEDOMENTS
The authors thank anonymous referees for their helpful comments. G.B. acknowledges the support of NSF grant DMS- 2134216. R.G. acknowledges the support of NSF grant DMS-2246780 The research of J.H. is supported by NSF grants DMS-2054835 and DMS-2331 20066.... A.J. acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC) and the Canada Research Chairs programme. Cette recherche a été enterprise grâce. en partie, au sou- tien financier du Conseil de Recherches en Sciences Naturelles et en Génie du Canada (CRSNG), [RGPIN-2020-0455) DGECR-2020- 2020-00199, du Programme des chaires de recherche du Canada.
REFERENCES
Greg W Anderson, Alice Guionnet, and Ofer Zeitouni. An introduction to random matrices. Number 118. Cambridge university press, 2010.
Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high nigh-dimmnssioal & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. In Gergely Neu and Lorenzo Rosasco (eds.), Proceedings of Thirty Sixth Conference Learning Theory, volume 195 of Proceedings of Machine Learning Research, pp. 1 1199-1227, PMLR, 12- 15 Jul 2023. URL
.ttml
Antonio Auffngger and Gérard Ben Arous. Complexity of random smooth functions on the high- dimensional sphere. Ann. Probab., 41(6))4214 4247, 2013. ISSN 0091-1798. doi: 10.1214/
13-A0P862.
Antonio Auffinger, Gérard Ben Arous, and Jifí Černý. Random matrices and complexity of spin glasses. Comm. Pure Appl. Math., 66(2):165-20) 2013. ISSN 0010-3640. doi: 10. 1002/paap 21422.
Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. The Annals of Probability, 3365):1443-167, 2005.
Gérard Ben Arous, Song Mei, Andrea Montanari, and Mihai Nica. The landscape of the spiked tensor model. Comm. Pure Appl. Math., 72(11):222--230) 2019. ISSN 0010-3640. doi: 10.
1002/cpa.21861.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High- nigh-dimensional limit theorems for SGD: Effective dynamics and critical scaling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
net net forum?id widmm3dddddokkii
Léon Bottou. On-Line Learning and Stochastic Approximations. Cambridge University Press, USA, 1999. ISBN 0521652634.
Anna Choromanska, MIkael Henaff, Michael Mathieu, Gerard Ben Arous, and Yann LeCun. The Loss Surfaces of Multilayer Networks. In Guy Lebanon and N. Vishwanathan (eds.), Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics, volume 38 of Proceedings of Machine Learning Research, pp. 192-204, San Diego, Califor- wwt.m.rrrr./eee/
nia, USA, 09-12 May 2015. PMLR. URL https: /ttttps:/prroceediiie Thoromanskal5.htze
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations, 2021. URL https: /eo/opperrree revi ew net. forum? ii htw.::idthhhheeeeeeeeeee
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Po-Ling Loh and Maxim Raginsky (eds.), Proceedings of Thirty Fifth Conference Learning Theory, volume 178 of Proceedings of Machine Learning Research, pp-
10

---

Published as a conference paper at ICLR 2024
5413 5452. PMLR. 02-05 Jul 2022. URL Bttps://proceednngg......eesss// damian22a.hor
Yann N. Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in high-dimensional non-convex opti- mization. In Proceedings of the 27th International Conference on Neural Information Processing Systems Volume 2, NIPS' 14, pp. 2933-2941, Cambridge, MA, USA, 2014, MIT Press.
Zhou Fan and Zhichao Wang. Spectra of the conjugate kernel and neural tangent kernel for linear- width neural networks. In Proceedings of the 34th International Conference on Neural Informa- fion Processing Systems, NIPS 20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019. URI. https:// openreview net/
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2232-2241. PMLR, 09-15 Jun 2019. URL ht..ht]
https: /tpps://rree http::/ppocceeeeee ngs. ml www.ss.pp.essee.e.............ee /v9 h99//horban/n
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace, 2019. URLptttps:/opperevlaa http://wwww..ttttrrrrmooeeeeeeee
X.Y. Han, Vardan Papyan, and David L. Donoho. Neural collapse under MSE loss: Proximity to and dynamics on the central path. International Conference on Learning Representations, 2022. URL
UbdvWH mnet/forrre id=w 1idw:wwwddddwwwww R3
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and gen- eralization neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: attpps://procee:innge www.::sssss.iiri....o..o..... cc/pppe 4be
fi wwwee//ppere21111 htb fa33ee62bb8a
- pdd
Arthur Jacot, Franck Gabriel, and Clement Hongler.
The asymptotic spectrum of the hessian of dnn throughout training. In International Conference on Learning Representations, 2020. URL
/ttpt::////peeeeeeeeeee net/ www.ormmidssssscaaae.e
Yann A. LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient Back- Prop. pp. 9-48.
Springer Berlin Heidelberg, Berlin, Heidelherg, 2012. ISBN 978-3- 642-35289-88 doi: 10. 10.11007/978-3-42--3357
URL httpp://aa htg/10.1007.
Xinyan Li, Qilong Gu, Yingxue Zhou, Tiancong Chen, and Arindam Banerjee. Hessian based analysis of SGD for Deep Nets: Dynamics and Generalization. pp. 190-198. doi: 10.
1137/1.97816111772266222 URL lhttps://epubt sia..iiaaa..... www.am..rg/aae hatgg//oo/abbbb 9781611976236..
Zhenyu Liao and Michael W. Mahoney. Hessian eigenspectra of more realistic nonlinear models. In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan (eds.), Advances in Neu- ral Information Processing Systems, 2021. URL httts:://ppnreviii ew. net. forum? di
RYNVO1 xAB
Bruno Loureiro, Gabriele Sicuro, Cedric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborova. Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL hattps://openrers net. net/forum
11

---

Published as a conference paper at ICLR 2024
Antoine Maillard, Gérard Ben Arous, and Giulio Biroli. Landscape complexity for the empiri- cal risk of generalized linear models. In Jianfeng Lu and Rachel Ward (eds.), Proceedings of The First Mathematical and Scientific Machine Learning Conference, volume 107 of Pro-
ceedings of Machine Learning Research, pp. 287-327. PMLR, 20-24 Jul 2020. URL https: //procedd ngs ml www.ppress.vvv..
11a rd20a. html
Charles H. Martin and Michael W. Mahoney, Traditional and heavy tailed self regularization in neu- | https: bottps://openreviieee wwwtfoooumiiiie 20--
ral network models, 2019. URL
Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. Ann. Statist, 46(6A)):747-27774 12 2018. doi: 10.1214/77AASS1677.
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynam- ical mean-field theory for stochastic gradient descent in gaussian mixture classification. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 9540-99 9540-9550. Curran Associates, Inc., 2020. URL httts::/prrceeddinn .ew.eeeiddinnnnnnni......s.. nwwrips. www.lpps.ccpaaae cc/paper files/paper/2220/ htlle//ccccccccccbo0055555556663 93 PPa3--aper.ppp .pdf
Marvin Minsky and Seymour Papert. An introduction to computational geometry. Cambridge tiass., HIT, 479:480, 1969.
Alireza Hosseeii, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In The Eleventh International Conference on Learning Representations, 2023. URI. htt.ttpss///ppeeeviieees net fort/? (ttrrmmmiddtttttzgeeeeoe
Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5012-5021, PMLR, 09-15 Jun 2019. URL https: httpt:t:tptcoeeeeeeeeeeee http:///cccccec.........
mlr wwwppppsssss.spepem ha.mmmmllm...mm111
Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research, 21((52): 2020. URL httpp::///ml...rr/pppprrs///////
20 933. html
Vardan Papyan, X. Y. Han, and David L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40); 24662,24663, 2020. doi: 10. 20055// URL URLhttps:: http::://wwwwwaaatt.. org/doi/ abs/10.1013 pnas 201 5509117
Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. SGD in the large: Average- case analysis, asymptotics, and stepsize criticality. In Conference on Learning Theory. pp. 3548- 3626. PMLR, 2021.
S. Péché. The largest eigenvalue of small rank perturbations of hermitian random matrices. Probability Theory and Related Fields, 134(1):127-173, Jan 2006. ISSN 1432-2064. doi: 10. 00.10077/0004400005-046----
Jeffrey Pennington and Yasaman Bahri. Geometry of neural network loss surfaces via random matrix theory. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1798-28006 PMLR, 06-11 Aug 2017. URL https: htttss://proceediin
pernington1 htmmlll
Jeffrey Pennington and Pratik Worah. The spectrum of the fisher information matrix of a single- hidden-layer neural network. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.. 2018. URL ht /twp / wttp://wwww.edeeeee............. fi
12

---

Published as a conference paper at ICLR 2024
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional gaussian mixtures: Where kemee methods fail and neural networks succeed. In International Conference Machine Learning, pp. 9936-8947. PMLR, 2021.
Herbert Robbins and Sutton Monro. A stochastic approximation method. Ann. Math. Statistics, 22: 400-407, 1951. ISSN 0003-4851. doi: 10. 10.1214//amm/1117772986
David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. Advances neural information processing systems, 8, 1995a.
David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995b.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singu- larity and beyond, 2017a. URL httpsss//pee..eeee....eeeeoooomm /tt.:/orrmmmii-11882222 9gx
Levent Sagun, Utku Evci, V. Ugur Güney, Yann N. Dauphin, and Léon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. CoRR, abs/1706.04454, 2017b. URL http:
//arxivv.rr//bb///06. 04454
Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval. arXiv preprint arXiv:1910.128:0 2019.
Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. arXiv preprint arXiv:2202.00020 2022.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Sumio Watanabe. Almost all learning machines are singular. In 2007 IEEE Symposium on Founda- fions of Computational Intelligence, pp. 383-388, 2007. doi: 10.1 109/FOC1..007.3315000
Zeke Xie, Yuan Tang, Mingming Sun, and Ping Li. On the overlooked structure of stochastic gradients. In sevethh Conference on Neural Information Processing Systems, 2023. URL net forum? id H4GsteoLOM
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances Neural Information Processing Systems, 2021. URL https:::: wttpps:/tpppeeeeeeeeeeeeep /pwpe//eeeee/.eeeee/mmo net caat//orumnte
13

---

Published as a conference paper at ICLR 2024
A ANALYSIS OF THE POPULATION MATRICES: 1-LAYER NETWORKS
In this section, we study the Hessian and G-matrix the population loss for the A-GMM problem whose data distribution and loss were given in (3.1 (3.2 Specifically, give the Hessian and G-matrices' X-large expansion, and showing they have low rank structures with top eigenspaces generated by (L1/A) perturbations of the mean vectors ( ((a). In Section D we will show that the empirical matrices are well concentrated about the population matrices.
A.1 PRELIMINARY CALCULATIONS AND NOTATION
helps to first fix some preliminary notation, and give expressions for derivatives of the loss function of 3.2). Differentiating that, we get
exp(x<. Y) ) Eo exp(xx (z Y)
Vv. L = (- +
(A.1)
the Hessian matrix
exp(x*. .))
exp(xx. Y) exp(xx Y) YOY,
V Va.
One
(A.2)
En exp Y)
(Eo exp(ex 8])2
and the G-matrix
V L.
(yelb - TT (b)y MbTy (c) NY
TY (c)ay (b)) ( DY.
(A.3)
Given the above, the following probability distribution over the [k] classes naturally arises, for which reserve the notation IT (-) MM((k]) (which is probability measure on [k)):
exp(x" Y)
(c) = xy (c:x) Enera exp(x* Y)
(A.4)
Note the dependency of TTy on the point in parameter space x. Since in this section x can be viewed fixed, will suppress this dependence from the notation. In the rest of this section, will denote 7 ~ (O, Ta/ A). We also denote
Bx = (IB) = Σ> xx)x((), (x
(x = Σ (x @ x (T)((( -x))
(A.5)
where B RT.
We consider the population Hessian 224 = V'E[L), block by block, using Vi to denote the d block in V2 corresponding V V 4. Then, thanks to (A.2). the be block of the population Hessian of the form
+ = E ( (b)ōse - TT ( OY
We also consider the population G-matrix I E block by block, using Doe to denote the x d block corresponding to V.,., L v., I.J. Then, thanks to (A.3). the be block of the population Hessian of the form
The (yoM Ty (b) yy (c) + (x)
For both population Hessian matrix and G-matrix, We will study the off-diagonal blocks a j and the diagonal ones separately.
A.2 ANALYSIS THE POPULATION HESSIAN MATRIX
We now compute exact expressions for the blocks of the population Hessian as X gets large: see Lemmas
We begin by studying the off diagonal blocks: b # c, for which,
Vo. + E [TTY(b)d TY (b))
E mB [TY (b)T(
(A.6)
14

---

Published as a conference paper at ICLR 2024
where Y = 14 + Z, i... it is distributed like Y given class choice L. It helps here to recall the well-known Gaussian integration by parts formula: For f that is differentiable with derivative of at most exponential growth at infinity,
Ef(ZA)ZA EVI(Z))
Our goal is to show the following.
Lemma A.I. The off-diagonal blocks of the population Hessian satisfy
(x) V be (x) = B TY
(C) (a) (1) (B())))B_T(()eee
+ 1100 +<*=+xx)+ x²)) 22 002
x Ia
A (a x (a /
In particular they are shifts by the identity of matrix of rank at most k2
Proof. We decompose each term in (A66) as three terms
(U) (c) pf 2 (E TT (b) Tr (c) & 1i) + E JY (b) (c) (c) =: (i) + (ii) + (ii)
where we' ve used here multi- lineariii of @. Here Sym(c @ b) = (c 2 b & c)/2. Let's look at this lerm-by-term. We leave term (i) as is. For Term (ii), we notice by Gaussian integration by parts,
(eyy( (b) @ E ( pt
E ((0" + x²) - -2 2 _,, )TY )T, m
[E|Ty (b) (c) + + @ 2E {x²) on (0) @ pr).
where we recall
A.5) and used the derivative calculation
exp(x* (mi + 2)]
(a) Vz
E. exp(x* (μ + Z)] (xx
(A.7)
This tells that
2 (ii) x Sym
(B x²) - Tri (b) 8 @m).
notice that this is of rank most 2 and of order ((1/A).
Finally for term (iii), we integrate by parts again, to get, for every i.j.
E Ty. (b) Ty (c)
Eiaz, (TH, (b) + E|Ty, (()x
(A.8)
x (T), + (b)
so term (iii) is given by
E (TT)())
Examining this second term,
Ei0z,z, (b) (c)) E [02 { (a) + -2xx ny, Jr (b) (b))
=B( (1f xz. 2(x
> m., ((x)) + xz)) 2( m Tr. (b) (c)
[x ( (0)))))))))))) )m))) Tr. (b) (c)]
where we first used
A.7 , then used that
Es r exp[x" (F + 2)]
V Vz(z B 1xy;
Vz
(x X
B Jay
- -10 -)0 ( (A.9)
En expla* (w Z)]
15

---

Published as a conference paper at ICLR 2024
As such,
x (1
1 xX
E (x x -2(10) ) (b)Tm(c)
2B (o :0 1e )av, ay (6) (m)
This term can seen to be of rank at most k22 (each (x is weighted sum of (z*)a).
Combining all three terms above, we get that the off-diagonal block of the population Hessian matrix given by the sum over y E [k] of
2 Sym(EE [(1a + x²) - -2xx0))) )om) mon, (b) (o) @ mi) + Br (b)TT) (c)]Ta
E (TT, my, (c)
82
x
E ((a" + x²) 2 (5,B)) AM To (b) Toy (c)) 2 ( ) Tan (b) T (0)
X²
(A.10)
Summing this in y we get that off-diagonal block of the form
-(A + X A2
B C)
(A.11)
where
A Σ BBB Tay, (c) (a) u1²
B = Σ (T) (c) la + Σ pr Sym Be (2 + +x0 x-==02266 2((BB) mn)) Tr. (c) @ @m).
C= Σ +x" -x- 2xx 2(
102 Tmy (c) (0)
2E
(x (10; : ) (x(c))y(aa))
zx.
In summary to leading order it is A which is rank k. To next order it is O(1AA) and that term is full rank (identity) plus rank most 2k term. To next order it is O(1/x²) and this is covariance-type quantity with respect to the Gibbs probability Ty,, with rank at most 1.2
We can group the expression (A.IT) further as
nr. EmB(Tr
(c) (a) (H4 + Hla + x - 2 (x0)
1)0²
+
la
X²
my
This is of the form rank & plus rank 2.2 shifted by the identity (adding one more eigenvalue). Incor- porating the average over into the expectation, this is exactly the claimed expression.
A.2.1 ON-DIAGONAL BLOCKS
In this section, we study the aa diagonal blocks
Vaa + E (T))(()()T(())))0))) = ΣnE
(A.12)
We prove the following large expansion.
Lemma A.2. The diagonal aa-block of the population Hessian V. aa & equals
1 (2 -(x (x X2 (x B. )nv)
Ty (a) (
{[) (mn 2 (10)
2
(x B B )TV )]
[ Tо (a)) la.
In particular, it is shift by the identity of a rank at-most k22 matrix.
Proof. By A. the diagonal aa-block is given by an average over l [k] of
El (a)(1 Ty (2)) (aTT
2) E[TT,(a) + 8)
16

---

Published as a conference paper at ICLR 2024
Note that the second term was exactly what was computed Lemma A.1 setting b It remains to compute the first. To this end, we proceed as before and for each fixed write the above as
(a)l4 + (ym(B (a) @ (4) + B (a)
We will integrate the second and third terms by-parts. By the gradient calculations of A.7
E Ty; (a) = } E EVV [a); = -
- El( - (o )ark)an)) (Tr,))
(A.13)
and by the calculation of (A.9.
R f (a)
[ + X F myy
X BlTT (a)] E/T (a) ((Fr - (on)) )82 (0";00")
(A.14)
Combining the above expressions yields
1 - (xx)) ((0) )82 + 2
1
E (a)) E| (a) + 2
(00 1x (, ).
1t
On the otherhand, Lemma A.I yields
E E(TTy, (a) (8 Combining these two and averaging over yields the desired.
22 Bi (a) +
2
X (x0 - (a
82
Id
2
(x
;
1
myy )1
X
Xt
A.3 ANALYSIS OF THE POPULATION G-MATRIX
We now compute an exact expansion of the population G-matrix gets large.
Lemma A.3. For any b, the be block of the population G-mulrix can be writlen in the form:
B C
[ 2. Vo. A+
X
where A is in the span of (Pe))e and B, C have operator bounded by 1. In particular,
A S @2 P.Y. (b) + Ā 1 [ (xx - (10) -
+ [00 )T,
(b)TT, (xb xx= -((x)))) )m, 02
B (inePh (poBlTTr (c)] + E. (b))) + Σ prE(TTY (c)r)
C Pxx x my, Tr. (b)) [l i )) Tr,
TY, (c)] 2 Σ (x Bx B) ),,,), n.; (b) (c)
Proof. We recall from A.3 there are four terms in the be block of the G-matrix:
E
Vas L] =: (i) (6i) (iii) (iv),
(A.15)
where
(i) E UeBbb Y ④ Y},
(ti) (b)y YJ, (Mi = ElybTy (c)Y @ YJ. (iv) : ()
The first term in (A.15) easy compute
E & Y = SucPo(1b @ μ4 + A)
(A.16)
The second and third terms similar to each other, will only compute the second term (ii),
(il)
Ely.Ty (b)Y ④ Y = F. [TY. (b) + 2) 2)82]
(A.17)
p.El (b) 22 2p. Sym E|ry, (b) pd] P.E TT, (b)
17

---

Published as a conference paper at ICLR 2024
Using the calculations of (A.13 and (A.14) we conclude that
(it) PlTY (4 02 + La/ A)
Sym(( 1Tx @ pe]]
(A.18)
22
(a k.C 002 ; mvy [b)]
For the last term (iv) in A.15 it has been computed in (A.10) that (iv) is the expectation over of Tr, yy (c) 14 82
Sym((B [1e²+ -
x
)mn) 20(x(()))))
2(xx
Tr, (b)
@ mi) + 1 Bly (A.19)
E ((0* + x²) - 2((x0)) mx, )ol TH. (b)TT) (c) - 21E (xr ;x ) Thy ())(())))
+
X
By plugging 16 (A.18)
its analogue for (iii), and (A.19) into (A.15) E[G2ya. of the population G-matrix is given by
we conclude that the block
B'
C'
G = A'
X
where
A' OsePh1h 82 PoTE.TT. (b)lpe - -pollTTV ( Σ (0
and
B'
pIBlTT, = 8bcPs (DoLE(TYs + PoBi (b))) + EmbbTyy (0)xx((0)) LT
2pe Sym(B((x - (o moy )TY, (b)) @ Mc)] - 2pb Sym(B[(s be
)r (c)] @ mb)]
Σ 2pr Sym(LE ((*' + x²) 2(x 2(2")xn)) )mn) Ty. (b) 2 m).
and
C² =ppe[((<>> (xr 102_))22110 :x Tv )TYY (b)) - Pc - (1r)) -2 10: )TY, (c)
+ Eei E ((x" + x²) 2(x ) 02 nY, (b)TT, (c) 2E (0";>") )an Tyi (b)TT)
pi
Grouping tensor-squares obtain the desired decompostion in terms of A, B, and C.
B ANALYSIS OF POPULATION MATRICES: THE 2-LAYER CASE
In this section analyze the population Hessian and G-matrices for the 2- XOR model, and especially its À large behavior by viewing perturbation of its oo value. Specifically, we compute its large À expansion, and uncover underlying low-rank structure. We will show that the empirical Hessian and G-matrices concentrate about their population versions Section D
B. PRELIMINARY CALCULATIONS
Recall the data model for the 2-layer XOR GMM from and the loss function 4.2 with being the sigmoid function and being ReLU. The 2-layer architecture has intermediate layer width K, that the first layer weights W form K d matrix, and the second layer weights v form K --vector. Observe that
V.L = -v--v)gWWY and
Vw, -yy - -Y)
where
4 = (
= olv =o(v.g(WY))
(B.1)
1+ (
18

---

Published as a conference paper at ICLR 2024
The vw-block of the G-matrix is the K x K matrix given by
V 02 (y (WY)02
(B.2)
The W,W, block of the G-matrix corresponding to W is given by
Vw Vw, L = (y 00- u,zyg' (W, Y)g (W; Y)r#2.
(B.3)
The per-layer Hessian for the loss given sample is then given by V = (l)g)(W))) D( - Ø) V.., L (- (W Y) ū) + oyuyg' (W; ()g(W, Y) - 0)) (B.4)
Note that we may also write the diagonal block in W of the form,
Vi L. = (y diagg(rig"
Y)) & y + 7(1 diagl' (WY))) 82 e y
We assume that W; # for all i, a set which we call W Then recall that for ReLu activation g. we have g" = 6g in the sense of distributions. Thus as long as W E Wo, we can drop g" in B.4 for any finite < co.
Our aim is to approximate the expected values of (B.4 by their "A = xx" versions, and show that they are within 0(2-1/2) of one another. Let F(t) P(G > t) where G ~ NN 1) be the tail probability of standard Gaussian. Let us also introduce the notations for 1 E {= +v).
m W and R₁ W W for 1 S i,j S K.
These notations will reappear as summary statistics for the analysis of the SGD the following sections.
Lemma B.1. For X ~ NO0, la/A) and W K x d matrix, we have that
[[[]((X)) 62 g (W 3 Sx 202 W,
(B.5)
where
(2,0,W = max
( (m
3
3/2 3/2 R//2,.112 22/ 1/3, (mm)) + + -3/2 R 1//3/1113
isijsk
mm (1 ))
1/3
x F V
+ F |m!l
(V
A R3j
W W
(B.6)
Rig
Proof. By the equivalence of norms in finite dimensional vector spaces, it suffices control the norm entry-wise at the price of a constant that depends at most on K.
Case 1: W 0, W 0 2 0. In this case have that
W; W;
Elg(W,XX)g(N)) X)] g(W, - (W - 8) = - E[((, ([(W,-0(WW, I -w xooww,xxed X
The absolute value of the first term is bounded, by Holder's inequality and union bound, by
B[[W, B|j,, (P(,, X 0) + P(W, X 03
Case 2: W, 1 < O (the case W, < 0 is symmetrical). Here we have Elg(W, X)s(W, X)] E[(W, [[(,, X)Imxxw3.x50), which, by Holder's inequality, is bounded by
El[W; - X/////"/ //P X> 03
combining the two cases yields the desired.
Lemma B.2. There exists a universal c such that for ~ NN00 La/N) and W kxdd
W W
BillgWXX) 1 1/2 <<>>i | + +
2 |
(B.7)
where || 11 the Frobenius norm.
19

---

Published as a conference paper at ICLR 2024
Proof. Let Z, W Z. Then
B(W; - - = [m + =1) = (m²)² 6(mg)2
x2
so that
|1ppl - (lllo(W))]l) < || + 11
( X
x
with the bound the following from the fact that |1-1|p 11-1la for p 2 q.
Lemma B.3. For (1, Ta/A) and v E R R
,W E kRd we have
B[(((v g(WX)) ov 9 (W0))) V((o ((X)) d' ((W0)))) Sx < (
(B.8)
where vz(x)= = x(1 + x)).
Proof. Observe that
(o(v:((W)))>- o g(W)))
<d'(u g(Wo))) (g(WX) x) ([W)))++0()) g( ((W)))3
(B.9)
- -(W))) | (g(WX))
and similarly for (o'(v ()W d'(v g(W where we used that or, o', a" are all uniformly bounded. By Jensen's inequality, it will suffice to bound the expectation of the quadratic terms. We begin by noting that
(9(WW)) o(Wo))l
(W-22)² ISiCK
Σ
so that
Bi(v (o(WX) s(W < |||| [
- ||0|| X W 12.
Consequently, taking the expectation of the right-hand side of B.9 squared, gives B[[a(v g(W av ((W x B(( (o(WX)
(W
(o(WX) ((W)))
ZK ||| X | = usl 1| |1
and the analogue of it for El(o'(v x(WX)) (v ((W)))) follows similarly.
B.2 POPULATION HESSIAN ESTIMATES
In this section, we study the population Hessian matrices. Our large approximations will be uniform compact sets in parameter space. Thus, we will use B to denote ball in parameter space, so that any constant dependencies on the choice of B are just dependencies on its radius. Lemma B.4. For (v, W) E B the first layer block of the population Hessian matrix satisfy
E V 1
Σ
d'(v ((Wo))
SK.B 22( W. A) + X }
(B.10)
(E(+tutt))
where 1h2 was defined in (B.6). For (v, W) E B with W E Wo. the diagonal second layer blocks of the population Hessian satisfy
Ammg
B W.W. D
alvg(m))) - 1alll11-ougggge o(ug(m))))
002
V
Rei
ME
1/2
A
1
1
F
mm²l
(B.11)
(E(+
VX
R
where we remind the reader that m standard Gaussian respectively.
Wt . R₁ = W, W,, and F and F are the edf and tail of
20

---

Published as a conference paper at ICLR 2024
Proof. We begin with the estimate the block. Recall from (B.4 that
B[ Elo'(v (v g(WY) ((WY))((
(B.12)
By conditioning on the value of v, writing F, for the conditional expectation, it suffices to bound |Ba[o'(u (u (WW)))(()))))) - (' g(Wo))
for each 0 E (+t,,t) To this end, we write
Ey [o' ((V-)((WY))(WWY))) - d'vv g(W
=Eglo'(w ((((((u)(WW))) ((W))) (1) + (ii)
(o(WY) - 82) + Eol(o' g(W)) d'(o (to-((W)))g(WY)e
By uniform boundedness of o', we then have by B.5 that [1()]]ap SK (0, W. On the other hand, by (B.7 and (B.8).
oo I (Bl(o'(v g(WY)) - d' (v. ((W)))
12/2 BB. X
We now turn to the W,W, block of the population Hessian. Recall from (B.4) and the fact that W E Wg. that
F w.w, LJ = E - Desg (W, Y)y
(since -z)))=='((x)) We now aim to show the following two bounds, the second of which will give the desired bound of (B.11: for each i, we have that
(v A Rii
1/2
[|Elotg (W: -
Σ DE (E(N
m!
082115
max #E #E(#udU)
VT
1 F
REi
Im
+
1
(B.13)
and
(V
[ D)
2)r
Σ
((W))( a(u o(WO)))F
my
00211
Ris
)
1/2
max (E(Ad))
F
Im
+
(B.14)
Rai
Conditioning on 9 and pulling out of BB.33. suffices to bound
Eolg' Y)Y (W, Y 0)002 + Sym Eg BowAw, 2 =: (i) + (ti) + (iki).
0
Eo
Eo
202
Term (i) is exactly the term to which we are comparing. In particular
(V
V2 Rui
Po Y 0) = P(W. 0 W, Z) F
We thus wish to bound the operator norms of (ii) and (iii). For term (ii). observe that for u E Sd-1 since Eg (u, Z) O, we have that
Lo[Aw,Y>>e (u, 2)] - -E Lw Z)].
Thus we may bound this, by Cauchy-Schwarz by
1 F
(V mm3
1/2
[Em[Lw,, (u, Z> < (u, 2)3 Y) sggmm))))) < 00)1
Vx
Ruj
follows, using that |||||| = that
(V
(
lop sup (u, (ii) <
F
mmm
Vx
u
21

---

Published as a conference paper at ICLR 2024
Finally, for term (dit) we do naive bound to get that for every E S4-1,
1 x
(u, (ii))))
[w, (u, 2) (E 21411/2
In order to show B.14 we can rewrite it as
[ Dy (W, V) ( (o((o
V
mi
182ll
Rii
(B.15)
= =I (0 M)' (W; 8) c 2 ( (1)) ov 9(W (W, ryy
+o(vg(Wo))( (1000)110(00()))))) ( (W, r)r - (v X
Rii m
+ (1i)
Since (vg(WW))(( - (W))) S 1, (ii) on the right-hand side in (B.15) can be bounded as in B.13) (without the i, terms). Term (i) on the right-hand side in (B.15). can be bounded as
(1) &
sup Bol(u, 8y Illllas1
where we used that Iy, Y)I S 1. Since E, (u, Yy x 1 + xt for ||u||2 S 1, we have by B.8 that [|(:]]| SB 1. from which the desired follows.
B.3 BOUNDS ON POPULATION G-MATRIX
We now turn obtaining analogous X-large approximations the blocks of the population G- matrix. In what follows, let ye = if 0 E {+/}} and yo = if E (+v} be the label given that the mean is 8.
Lemma B.S. For all (v, W) B, the first layer block of the population G-matrix satisfies
| V 12
Σ
(vo olv (e((ug((WO)))) ((W))))>((WD))) 21 SK.B 22 W, (2(), x. (B.16)
DE(+mANt
and the second layer blocks satisfies,
1/2
(V (m?))
|L[V
-1AAl1 S VA (=(=
max F
+
(B.17)
where
A
1
Σ -
(va-o(v
olv s(Wo)))2
(
R mm
y82
(B.18)
4
dE(Lp,Lo)
Proof. We begin with the v-block. Recalling B.2). and conditioning on the mean being 0, it suffices to show for i' (#,,
E (o (
alv o(W))) 9(W ((W))) SK.B (22 W,
1 X
and similarly with negative signs in the sigmoids, to account for the ye term when relevant. The proof's are similar, so we just do the first, with a fixed choice of v. We begin by writing
Eg [o(( ((WY))?g(WYY))
- olv ((W)))'((W0))
= (v g(WY))² o(u ((W))""l(WY))0)) 02) + Esfo(v g(Wo)) s(W)) (i) (ii).
We begin with bounding the operator norm of term (i). To this end, by Cauchy-Schwarz, then B.7 and B.S).
[10)) lp 5 <Eollo(WY)
1 111 1/2 1/2 (o())))()))))))))))) ale (W0))))3112 Ss
22

---

Published as a conference paper at ICLR 2024
For term (ii), since |o| S 1,
lop I [| ((W)) llo SK (22 X),
by B.5).. Combining these two and averaging over yields (6.16..
Recalling (B.3). and conditioning on the mean being 0, it suffices to control the operator norm of E [7 ya)² gWW; ; Y) Ag
where Ag is the corresponding summand in (B.18.. Fix & with yo = 0, the choices being analo- gous. In this case wish bound the difference,
Eolo(v- g( -)) - -o (oo-g(W)) P(y
A
mmm) -() (1) +ii),
Rii
where
= Esl(a(v g(WY))² ou )g' r)r
(i
A Rsi
(ii) = a(v ((W)))
Eold'(,, r) F
m
1622
By BB13)) and the boundedness of
1 L.F VX
1/2
(VV
|m?1
+
aE(+p.
Ris
A
For (i), we bound its operator norm by Cauchy-Schwae as l1())lop s [[(((v g(W, Y))² olv g(W, -0)) 33)3/ |1wlaa51
sup Rol(u, r,
where we used that Ig'(W; - Y) S 1. Since Eg (u, r) $ 1 1 x2 for ||u||2 S 1, we have by B.8 that [|(:)] ZB 3. from which the desired follows.
C ANALYSIS OF THE SGD TRAJECTORY
Our goal in this section is to show the following results for the SGD for two classes of classifi- cation tasks. The first is our main result on stochastic gradient descent for the mixture of k-GMM's via a single-layer network.
Proposition C.1. For every 3,3 > 0, there exists To such that for all T3 > To and all X large, the SGD for the k-GMM with step size 6 ((1d)), initialized from (0, Ia/d), does the following for all E [k]] with probability 1 - o1()):
1. There exists L(3) (independent of,, >) such that ||x|]| S for all é E (Toō-1,Tyō"*
2. There exists nt()) > 0 (independent of c, A) such that x2 is within O(c + 1 distance of point in Span(p1,.... H&) having ||x||| 17
The following our analogous main result for the XOR GMM with two layer networks.
Proposition C.2. For every > O, there exists To such that for all T7 > To and all X large, the SGD for the XOR GMM with width K fixed, with step-size = O (1/d), initialized from ( Ia/d) in its first layer weights, and ((, i.id in its second layer weights, with probability oa(1), satisfies:
1. There is an L(K, 3, v(xo)) (independent of c, A) such that [[aa S for all < K, and v(xe) S for all & E [Tu .T,5-13
2. B < 1/8, there exists 1 (K, 8, vxxo) > O (independent of e, A) such that xę has
[xx()))) 7) and max W. 0 > n for all E E (To _T,ō and i < K. aE(#u.tu)
and furthermore, it has Wx(xe) lives in Span(u,,) v) and v(xe) lives in (pan(((((xx)))))ee3)) up error X112
23

---

Published as a conference paper at ICLR 2024
Note, the dependencies of L and the absolute initial second layer value |u(x))] continuous.
Our approach goes by first recalling the main result of Ben Arous al. (2022) regarding limit theorem as d + oo for the trajectories summary statistics stochastic gradient descent. We then apply this result to the task of classifying & Gaussian mixtures, obtaining ballistic limits for the SGD) in Theorem C.7 that may be of independent interest. We then analyze the ballistic limits to establish recall the ballistic limits for the XOR Gaussian mixture from Ben Arous et al. (2022) and similarly establish Proposition C.2 Since we are imagining the dimensions of the parameter space and dimension space, and the step size to be scaling with relation to one another, we use dummy index n, in this section to encode their mutual relationship, namely d = da P and 8 5, and then + oc.
Proposition C.I Then,
Also, to compress notation throughout this section, we will combine the loss function with the regularizer 2 ||xx| so that the stochastic gradient descent updates are indeed making gradient updates with respect to the new loss:
X x 13--1-VV(xx-1) Y where L(x) - L(x) + 11
C.I RECALLING THE EFFECTIVE DYNAMICS FOR SUMMARY STATISTICS
Suppose that we are given a sequence of functions u, E C2 (RRP::R IR*) for some fixed k; where (x) = (uf(z), ..., uf (x)). and our goal is understand the evolution of u, (xx).
In what follows, let (x, Y) (x, Y) & (x), where (x) EJL,, Y)J. (Note that since the regularizer term is non-random, H is the same with L in place of L.) Throughout the following, suppress the dependence of H on Y and instead view H as random function of x, denoted (x). We let V(x) = Ey [V (x) & (x))) denote the covariance matrix for V H at x. Definition C.3. A triple (u,, L, P..) is -,.lllccie with localizing sequence (Eg) if there is an exhaustion by compacts (Ek) K of Rt, and constants CK (independent of n) such that
1. max; sup suP (Ex) (Ex)
/2 1lop S Cx 1/2 and max; sup supxeun (Ex) | llop Cki
2. 2uPxenn'(EE,
|| S Ck, and sup (F V CxO,t
3. max; SUP E V < O and
max; sup SUP u V & VH o(
Define the following first and second-order differential operators,
An Σ 0,00, and Ln
ViJ
(C.1)
Alternatively written, Ar = (V., V) and Cn = (V. V Let Jn, denote the Jacobian matrix Vun-
Definition C.4. family of summary statistics (u,) are asymptotically closable for step-size ön if (Mn, T., Pw) are o,c-localie with localizing sequence (Fx) K. and furthermore there exist locally fipschitzfumcttios functions h IR + Rk and IR kk such that
sup xEu, ((Ex)
(1( An, + 8n Cn
un un (x) h(,, + +
(C.2)
sup xEn,'(FF
18 J.J.V Ex,, 0.
(C.3)
In this case we call h the effective drift, and the effective volatility.
For a function f and measure I we let fajt denote the push-forward of p. The main result of Ben Arous et al. (022) the following limit theorem for SGD trajectories - 00.
24

---

Published as a conference paper at ICLR 2024
Theorem 2.2)). Let (x(")) be stochastic gradient descent initialized from X6 ~ An for pn E M1 (IPP)) ) with learning rate on for the loss Ln(-s.) and data
Theorem C.5 (Ben Arous et al. 2022
distribution Pn. For family of summary statistics u, (ut) let (u,(t)); be the linear inter-
polation of (u, (un(x(1s, [ 1))e
Suppose that Un are asymptotically closable with learning rate ons effective drift h, and effective volatility Σ, and that the pushforward of the initial data has (u,),Rn - v weakly for some v e Mi (IR*). Then (un(t)) + (Ut)r weakly as + where un solves
du, = huu dt + ddu).BB
(C.4)
initialized from v, where B, a standard Brownian motion in R
As result, we can read off in the ballistic regime the following finite-n approximation result. Lemma C.6. In the setting of Theorem C.S if the limiting dynamics have Σ then for every T,
we have with probability 1 oI()),
||u(x|6 44)) uelccoo.11 oa(1).
C.2 1-LAYER NETWORKS: BALLISTIC LIMITS
Our first aim in this section is to show that the limit theorem of Theorem C.5 indeed applies to the mixture of & Gaussians.
Recall the definitions of the data distribution, and the cross-entropy loss for x E Redd (i.e., Ta E R for E ( from B.2).. Recall that Meb Mo Ms, and that in order for the problem be linearly classifiable (and therefore solvable with this loss) no mean should be the weighted barycenter of the other means, i.e., Pafte + & Es Potts. For this task, the input dimension
n, the parameter dimension Px = kn and the step-size will be taken to be 5 = O(1/n) = O(1/d), and we will tend use d rather than the dummy index n. It will be helpful write Y, Ha + Zx, and recall for E [k)), the Gibbs probability TY, (c) = Tr, (c; x) as defined in A.4
C.2.1 THE SUMMARY STATISTICS
We will show that the following family of functions form a set of localizable summary statistics u, (x) ((mab():) (.)),,b (Rd))a))
m (Mab)a,be[k) where where R = x. 2%
mab No. : 14b
R = (Rih)
where vn denotes the part of xr orthogonal to Span([i, (RR), i.e., ant Ptr., Ra. where we use P1 to be the projection operator into the orthogonal complement of Span(/11, .... px))
not hard see that the law of the full sequence (my, (c))a.e (and therefore its moments etc...) only depend on through m and R and therefore they have no finite d dependence, The following describes the ODEs obtained by taking the simultaneous d > and 8 O(1/d), limit of the SGD trajectory in summary statistics.
Theorem C.7. If the step sizes 8 cs/d. then the summary statistics u, (m, R are 5- localizable, and satisfy the following ballistic limit as d + oc:
maa = Paa Bmabb - Σ (m QZau))
(C.5)
cE/A]
Rt (t) = -3Rd
Σ cE(K]
(aaQaRiu + Po
Ri
Σ cE(R]
Sxpe (TTY. 1a= (TTY (b) 1b=c)
(C.6)
where Pe, Qeo are the following Gaussian integrals:
(m R E[[TY. (a)) and Q ) -( (a))
(C.7)
and where if L Span( (4:) we use the shorthand Q for r = || since for such v, Q%*"
only depends v through || The case where d = o(1 ((//d) d) is read-off by formally setting Ca = 0.
25

---

Published as a conference paper at ICLR 2024
Remark 2. While priori, may appear that the quantity Qa" in C.7 depends the dimension
d, if v is one of (Hi,.... (x), or if it is simply orthogonal to Span P1,... Mg), then Q% does not depend on d. Such cases are the only ones appearing in (C.5) (C.6 The same be said for the expectation in the last term (C.7
Proof. Theorem will follow from an application of Theorem C.5 to the summary statisties m, R. for the K-GM,, so we start by verifying that the problem fits the assumptions of 8- localizability and asymptotic closability.
Verifying 6-localizability Our aim is to verify the conditions of á localizability from Definition C.3 Let us begin with some calculations, recalling that the Jacobian matrix J Vu. Let V. denote the derivative in the Rd coordinates corresponding to Ia. For a, b, c E [K].
# h =c
S Lb 0
if else
Fb
b b a
V Mab
and
(C.8)
2rd
a=b=c
(0
else
Continuing, all higher derivatives of aas are zero. The Hessian of Ry is given by
(p-
if a
V aed = 0 and Rib
(C.9)
2P
a
(and other blocks are zero), and higher derivatives of R. are also zero.
Let us also express the loss function equal in distribution to random variable whose law depends only on m and R-- For ease of notation, let
Ya 1aa + Zx 3.2 and adding in the regularizer, we can write for each fixed x,
Rub
(Taa .ib)
and
Recalling
L(x, Y) (d)
yo Mmaa + Ga) +
Σ DE
Ub log
mab+ +
E aEIM
Roa
(C.10)
aE[k]
aE[k]
where (Ga)a a Gaussian vector with covariance matrix (Rab), and y uniformly drawn 1-hot vector in R Observe that the law of L only depends on x through the values of the summary statistics.
Lemma C.8. Suppose 8 = O(1/d) > fixed (or growing with d) and B > 0 is fixed. Then the family of summary statistics (m, R-) are 8-localizable with balls.
Proof. Hem (1) of b-localizability: The first part is easily seen to be satisfied by (C.9 since the Hessians of statisties in m are all zero, and the Hessian of R is bounded in operator norm by 2 - 2k using the triangle incquality and || = 1.
Item (2) of ò-localizability. We begin with the bound the population loss's norm. By taking expectation (C.10)
(x))
Σ
Po
Σ P.
log
Σ
(an))a))
Σ
Bua
(C.11)
(E(K)
nE(A]
bE(M]
aE(M]
Taking the derivative of this, we get for each E [k].
vel(x) -Pelee + Σ (A 8x
(C.12)
aE(M]
where Y, = + Zx and π is as in A.4
Considering the norm of / we have
|| < max V S +
112 1/2 + 3 max Baa)
26

---

Published as a conference paper at ICLR 2024
which is bounded by constant C(K) for m, R in ball of radius K.
Moving on to the bound on ||| V *), first observe using V.H = V V. that
H = (Petle (e + Zx)) + -B(Tc))yBBBTTo))
(C.13)
Taking the norm and the 8'th moment, we use the fact that || f X (I|u ||u|()). ). and the fact that Ty, (.) is probability mass function and therefore at most 1, to upper bound
max [|] cE|M]
sup pupl|mm. |a
Z < 1 ++ ((d/A)").
As long as 6n ((/d), since À is uniformly bounded away from zero, this will be bounded by constant times 0,z4 4 as required.
Item (3) of 6-localizability. We next turn to bounding the fourth moments of the directional deriva- tives, starting with the statistics mab:
(V.H, Vamab) = ((pe Mo) pb) - Zx mb) + po (c) - (c)
+ Σ [c) Zx)I - (c) ZAi).
[E[R]
Taking the fourth moment, again bounding things by the fourth moments of the terms individually up to universal constant, then taking expected value, the second by the fourth moment of a Gaussian random variable with variance 1/X, i.e., by C/X², the third and fourth by since is probability distribution, and the summands individually have fourth moments bounded by C/ for the same reason. Altogether, we get
that the first term is bounded by 1,
HH, V m & 1 O(1/^²),
satisfying the requisite bound with room to spare. Turning to the directional derivative in the direc- tion of R note that derivatives of R.,, are orthogonal to (J1,....,14 .4k) so for a o b,
H, ) H, Ro (=) Za)l TY (c) Z>)) ye( Zx)
Considering the fourth moment of the above, using that π is bounded by 1, and (x) ZA) is dis- tributed Gaussian random variable with variance R2 /X, we obtain
HH
S (Raa/A)².
which is bounded by C(K) while R is bounded by K. The diagonal a = b is the same up to factor of 2. The last thing to check is the second part of item (3) in 8.-localizability. For this purpose, recall that V(x) = E(V H722, and notice from equation C.13 that for the k-GMM we have
V.H @ V.H = ((TY(() (c) y.) Y - (C( yc)Y @ ((TY(()) (d) ya)Y B[( (d) - ya)
Taking expected values, we get that the ed-block of V is given by
Vad E V & V Valf- = Cov (TY (c) v.)Y, (TT ua)Y)
(C.14)
where Cov is the covariance matrices associated the vectors inside it.
For the second part of item (3) in localizability, we only need to consider the statistics R since the second derivatives of mab all zero. For the ballistic limit it is sufficient to the bound
E Re,, V.H VaHl Vea)a] R H[I
The first term is bounded by per (C.9. The second can be seen to be bounded via the 8th moment above by O((d/ 2)2 which is o(d when 8 = O(1dd).
Remark 3. The reader might notice that there is a lot of room in the bounds above as compared to the allowed thresholds from the 9-localizability conditions. The weaker conditions in localizability are to allow taking diffusive limit theorems about saddle points by rescaling the summary statistics by d-dependent factors, which help with understand timescales to escape fixed point regions of the limits of Theorem C.7 This was explored in much detail for matrix and tensor PCA in Ben Arous
et al. (2022).
27

---

Published as a conference paper at ICLR 2024
Calculating the drift and corrector Now that have verified that the 5-localizability conditions apply to the summary statistics u = (m, R we compute the limiting ODE one gets in the d - oo limit for these statistics. We will establish individual convergence of Au to some fu( and convergence of oL to some 9u (u) for each u E u, whence h from Definition C. equals -f +
Recall the differential operator A from C.Ip. D. the expression for V from C.12 and consider
Amab
Σ (V. V Mama)) (V. (Pataa Mb) + E Ey
+ 8(za, mb)
cE[R]
Recalling that Mas = (Ma, Hb), we get
Amab Mab + Brab Σ pr (m,.ETy + (o)x, Pb)). [E[K]
Notice that the two expected values are Gaussian expectations that only depend on x through (ma)) and (R In particular, we can take the limit as - oo to get the limiting drift function
Ima (m, R-)) = aaalaa + Bm +
Σ pr myy Pd + Q%
(E(M)
where Pl and 2s are defined per C.7
The contribution coming from 5Lmab vanishes in the oo limit since the second derivative of mab is zero, i.e., Im..
For the drift function for R since (ard mu) for all a, b, get
AR
R +
Σ Lc
(paQuat + +
In particular, we get
fR.. = 3R +
Σ
(paQua + R.a
(E(K)
using Q. = ot since Q::" only depended on through its norm when 1 Span(H1,... Mb)
There is also contribution from SCR recall that L 1 (V, v2) and notice that
CR ((Vas. P + (Vo., P pT
Plugging in for Vas from CC.14 and expanding this out, calculation yields
CR ab
Σ cE(R]
[[[ - 1 ( (b) 10=c)]
(B(Z(TTY - 1,,c)) (b) 1o=.)))) + 0(1),
where to see that the extra terms are O(1), we notice that the inner products of the j's with each other and y's with Z's are order By Cauchy-Suaie Schhar, the inner product of the expectations most O( O(Vd/A) A) and will vanish when multiplied by = ((1/d) and the - oo limit taken; on the other hand, the second moment term is of order d/ A. In particular, if 8 = cs/d, we get
SCR C5 Σ P.B dt 1/1
1 1.=e) 10=e) + ol)).
<E(k]
We claim that the d + oo limit of this gives
C5
Σ P. (ay la=o) 1 b==) cE[R]
In
(C.15)
Notice that the expectation is bounded by in absolute value and thus the above goes to as - co. In order to show (C.15). let consider the term coming from ny, (a) (b), the other terms being
28

---

Published as a conference paper at ICLR 2024
analogous even easier since the indicator is deterministic. Using the fact that probabilities and therefore bounded by 1, note that
1/2 B||ld ny. (b) Bl (b)] S ||dd|//2 2 1
which goes to from the standard fact that for standard Gaussian vector Ga in R. one has 1) 11
O(1dd).
The last thing to check is that in the ballistic regime where our summary statistics are not rescaled, the limiting dynamics are indeed ODE, i.e., there is no limiting stochastic part. Towards that purpose, notice that in any ball of values for (m, R),
JJV O(1),
using an O(1) bound on the operator norm of V, and noticing that || is bounded by a constant plus the sums of R. Therefore when multiplied by o(1) this vanishes in the limit and therefore the diffusion matrix E is identically zero.
The following confines this X-finite dynamical system to compact set for all times.
Lemma C.9. For every B > O, there exists L(3) such that for all X large, the dynamics of Theo-
rem C.7 stays inside the balball of radius L for all time.
Proof. Notice that a naive bound on Qd from CC7 is at most v V|ue| ( since Ty, (a) is bounded, and pd is always at most Plugging these bounds into Theorem C.7 together with the definition of Mab and the fact that || for all a, yields the inequalities for all a.b.
(t)a((t) + Bmasl E 1 V1/A,
RR + (R 2 V R / -
By these, for À larger than a fixed constant, we have Red, < 1 R 2 which Gronwall's inequality ensures will be bounded by + R.., (0) for all times 0. Similarly, we get that for À sufficiently large, ama (t) is bounded by 3 + (a).
C.2.2 THE ZERS-NOISE LIMIT
We now send À - oo, or simply Taylor expand in the large limit, to understand the behavior of the limiting ODE's derived when À is large. Let
(a) = (a;
(C.
Eve(H)
emas
This is the "> = xo" value of (a). The aim of this subsection is to establish the following.
Proposition C.10. The X + co limit of the ODE system from Theorem C.7 is the following dynam- ical system:
maa Pla Bmabb
Σ
(C.17)
cE(M]
R (t) -3
(C.18)
Moreover, at large finite X, the difference of the drifts in C.55
C.6) to the above drifts is O(x-1).
The main thing to prove is the following behavior of integrals of TY. (a) as À + co.
Lemma C.11. Recalling Pa and Q from (C.7 we have
pe = (a) + -(1/A)
Q O(1e)l/A)-
29

---

Published as a conference paper at ICLR 2024
Proof. By Taylor expanding, write
TY. (a) + (Ia Z)e te
@ to (EDERA (26 2)
O Tb 2)
ESE(M 00
bElk) exosee )2
be
Taking an expectation of the right-hand side, noting that 2b right-hand side after va(aa is O(1/a) for Rig that is O(1).
Z NO0 Lmn
everything on the
For Qa", using the Gaussian integration by parts formula and A.7
2% [[ (alaa-e)mye D)TTY, (a)] E[(x u) Tr, (a)l
Since and Tb S Rb |o||, this is easily seen to be O(1|v|1 Righ/ A) claimed.
Proof of Proposition C.10 For any K, uniformly over all m, R a ball of radius about the origin, claim that the limit of the drifts for each of those variables converge
the claimed + oo limiting drifts. This is obtained by applying the above lemma to the P and Q terms in the drifts in Theorem C.7 and finally the observation that Var(ir, (a) TT) v (b)) S 1 so that taking À the last two terms in the drift for R. in Theorem C.7 also vanish.
The following gives quantitative approximation of the ODE by its X = 0 limit.
Corollary C.12. The trajectories of the ODEs of Theorem C.7 and Proposition C.10 are within distance (t//X) of one another.
Proof. Let u and u be the two solutions to the À finite and = oa ballistic dynamics respectively. Then, while [u|.. [ < L, we have by Proposition 1h
û - ا|| St X-1
Per Lemma C.9 both dynamics remain confined for a large enough L(3) for all times, and therefore integrating the above gives the claim.
C.3 LIVING IN SUBSPACE SPANNED BY THE MEANS
We now wish to show that the SGD trajectory lives in the span of the means. This can be done by showing that on the one hand, R- will stay as small as we want after an O(1) burn-in time, and on the other hand, towards the error being multiplicative in Definition 2. for every a, x9 needs to be non-negligible.
We first establish that this happens for the À = dynamics, then pull it back to the dynamics finite but large via Corollary C.I2
Lemma C.13, The solution to the dynamical system of Proposition C.10 is such that for all t 2 To(e), it is within distance of a point having R2 = 0 and max; |mabl 0g 0 every a E [A:].
Proof. By the expression from Proposition C.10 for the drift of Ruo for E [k]. the dynamical system has R.l (t) e- R,! (0). In particular, for any O, the has Raa
all 2 Te(.).
We need to show for every a, in the solution of the dynamical system, some (mas)) is bounded away from zero after some small time. Let Mg = U, Na {Mac = 0} be the set of (m) values we would like to ensure the dynamics stays away from. First observe that the À oo dynamical system of Proposition i. is a gradient system for the energy function
H(m, R )= = Σ Maa +
Σ (E(R]
Pa log
Σ DE[R
Mud +
(||m)
|R 113)
aE(k]
30

---

Published as a conference paper at ICLR 2024
so it has recurrent orbits. It thus suffices to show that for every point in Mo. the quantity max; |aabl has a drift strictly bounded away from zero. If we show that, then the dynamics is guaranteed to leave eg neighhbbborooo of the set MM in finite time (uniform by continuity consid- erations and we are already guaranteed that the dynamics stays in a compact set by Lemma C.9
Consider point such that mab for all E [R:]. There, by Proposition C.10
Σ 0E(A]
mab (PaHa
We need to show that the maximum over b, of the absolute values of these, is non-zero. Indeed, if it
were for all b E [k], then Palla E. Pe because the difference of these vectors would be in Span(j4, ..., M&x) while being orthogonal to all of p1,.... In turn, however, this is impossible by
our assumption that Pı.... j4g are linearly independent. Therefore, in the ball of radius L(3) about the origin, for every a, at least one of the drifts ma is bounded away from zero uniformly.
O
Proof of Proposition C.I We have shown that the limit dynamics for the summary statistics of SGD initialized from N (0, Ia/d) the dynamical system of Theorem C.7 initialized from the deterministic initialization mab ~ 6u and Ryo óg if a & b and 81 if a = b.
By Lemma C.9 there exists L(3) such that that dynamical system confined to a ball of radius at most L for all time. Since ||=1|2 is encoded by smooth function of the summary statistics Re (meb)s appearing in the dynamical system of Theorem C.7 this transferred to the SGD x² via Lemma C.6 for different constant L(8) for d sufficiently large.
For the second part, by Lemma C.T3 the solution to the dynamical system of Theorem thmnk- GMM--alliisti-liiitt is distance O(A-1) from the solution of the dynamical system (with the
same initialization) of Proposition 10 for which Lemma C.13 applies. By Lemma (.6ma these get pulled back the summary statistics applied the SGD itself (at d sufficiently large depending on A). We therefore deduce that for all é To(E)o- steps, the SGD has (xe < 3( A and has max; (x5 |666l > Cg > These imply the claim using that ||x*|| Z max; [(x", Mb01 and that R
(xx) ) is by definition the projection of x orthogonal to (J1..... ..)).
C.3.1 FIXED POINT ANALYSIS WITH ORTHOGONAL MEANS
The behavior of the ODE system of Proposition C.10 can be sensitive to the relative location of the means (41))..,/))). In order to be able make more precise statements about the alignment of the SGD with specific eigenvectors rather than just living in Span((1,.... p&)). we specialize to the case where the means form orthonormal family of vectors. Here, fixed points of Proposition C.10
explicitly characterize the
Namely, in this subsection, assume that Mab la=b- Then, by (C.17) any fixed point must satisfy
mas
3mab = Pala=b - Pb E. ema
(C.19)
At fixed point, the function E. meb thus must equal O. Also, if a, c # b then
mas
Meb
Mab
Meb
Eyeua
Since the function + oe is strictly increasing, the only solutions to this mas meb (so long as p6 > O). Combining the above two observations, at a fixed point,
1
Mab
mas
for all a + b.
Plugging this to C.19 we find that at a fixed point
pb
mb6
m&b
3
(k 1)e
x= maa emas
This can easily be seen have unique solution, and that solution must have mis (0, 81)
31

---

Published as a conference paper at ICLR 2024
Therefore, long (po)be(s) all positive, the dynamical system of Proposition C.10 has unique fixed point, call u, (mad)a, b as above, and R2 for all a, b. As observed in the proof of Lemma C.13 the dynamical system never leaves ball of some radius L(5) and is also gradient system for an energy function H. Combining these facts with continuity of the drift functions, means that for every > O, there is (c) such that the solution to the SGD gets within distance To of that unique fixed point and stays there for all t > To.
Altogether, this leaves us with the following stronger form of Proposition C.I
Proposition C.14. When the means (p1,,.,,1ss, 14k.) are orthonormal, beyond Proposition C.I. we fur- ther have that is distance O x-= of a point x, such that for each x,, has positive (bounded away from zero uniformly in X) inner product with pe and negative inner product with (Pa) Mafe
(Po)cE(s) are assumed all be equal, then furthermore at X..
1 k- =c(c))
=b(a)
for all a, b, e such that + b.
C.4 2-LAYER NETWORKS
We now turn to the analysis of the SGD in the case of multilayer networks for the XOR GMM problem (4.1) In this problem, the input dimension is d = n, the parameter dimension is
Kd + K, and the step-size is again taken be á O (1)d). Consider the following family of 4K + (K)
summary statistics of x: for sisj S K and 0 E {u,v).
US
m
We
R Wt+ W
(C.20)
where Wi- = W, - (ve( m 0. Use u = (v, m m", R for these families.
In Ben Arous et al. (0022 it was shown that this family of summary statistics is 5-localizable and asymptotically closable with respect to the loss for the XOR GMM of (4.2). and the following
convergence to ODE's was established. For a point x = (v, () E RK kkkKKd define the quantity
A; E[y f[Ylw,r20(- y alv g(WY)))]
(C.21)
(where we recall that g is the ReLU function and a is the sigmoid function) and let
A 2 = 0 Ai,
A W A.
Furthermore, let
Bij
aov ((WY)))
(C.22)
It can be observed that these functions expressible as functions of alone.
Proposition 5.1 of Ben Arous et al. (002 proved the following effective ballistic dynamics.
Proposition C.15. Let un be as in
C.20 and fix any À > and cs/d. Then u, u,)()) converges weakly to the solution of the ODE system ú, =f((u) + g(u)). initialized from (i,,, (u,)=Pn with
fe, m. A/ (u) m Au(u) + A! S...
fm; u.A, + Sm
fR. = u (1) + My (u) 2,
Jm = - A,A + 3.m
and correctors 9o, = 9m 9m = O. and 9R = C3 21 8 o(1/d) is read-off by formally setting Cg = 0.
Bij for 1 si S < K. The case where
x
C.4.1 LARGE À BEHAVIOR
We now wish to investigate the large À behavior of the dynamical system in Proposition C.15 Our approach to doing this is to give a large À approximation to the drifts in the above, and then use that to show that the trajectory is close to its À oo version.
The first thing we do is give large À approximations to the quantities A, and B In what follows, use F(x) to denote the cumulative distribution function of a standard Gaussian random variable.
32

---

Published as a conference paper at ICLR 2024
Lemma C.16. Suppose Y ~ 0 + Zx for fixed vector & {ur -, -v} and fix unit vector b E RR and a vector U E R Then
[[[[b Y)1w,r
g(WY)) ([t-)>/oooo[g]]]]]--) (b- (..0))' (m²
A Rio
o(v g(m"))
Rui )e (m93 +((-)))
S max(( + my)(m, +
Proof. Let us start with Taylor expansion of a: For simplicity, let us dy and dg denote (WY)) and oo ((W0)) respectively and similarly for a'. We then have σy og = oo (e (o(W)) g(WY))) + to)()) (o(W0) -00-))((>)))
(C.23)
for some point o between v g(WY) and 0) Therefore, the left-hand side of the lemma is
E
8)o
E Iw, 2)00| + E Il M)ob(u (o(W)) s(WY))))
E 1w. (o)(() (g(W - s(WY)))²)]
(C.24)
The first term of C.24) is exactly the term comparing to in the left-hand side of the lemma statement. Since E[ 2)] = 0, the absolute value of the second term of (C.24) is bounded via
Cauchy-Schwane as follows:
E 1 wubb Z) E (Iw,Y>o - I mm 2000
F(-lm; VX/R,)) 1/2
VA
For the next two terms we can start by rewriting
g(W, 0) aWW, Y) = (mg +W, Z)1w, Z) commf30 (W, -Y) w,yYy,, c0 <WWW Z)1
my >0 (C.25)
Using this, the third expectation in C.24 is
(b. H'
V) Σ uy (mf+W, 2) v - (W, I w, Ys (W, Z) mm330)
The first and second terms in the parentheticals are such that the Cauchy-Schwarz incquality can be applied to bound their total contributions by
max (mm))² )² + Rig F A/ -
Kollool J
X
The last will contribute (up to a net sign)
B (W; Z) - 1 m d' Z)(W, 2
The absolute value of the first term is bounded similarly an earlier one by Cauchy-Schware Schwarz. The absolute value of the second term is bounded by dropping the indicator functions and applying Cauchy-Schward to see that (1/AA) so long as [[b]]. [W, = 0(1).
Finally, for the fourth term of C.24 the square allows to put absolute values on every term, and immediately apply the Cauchy-Schwarz Schwarz inequality, to bound its absolute value by
2/2 max E[(g(W, 0) - g(W, -))
The fourth moment in the above expression is bounded, up to constant, by the fourth moments of each of the individual terms in (C25). The first two of those will be at most some constant times F( F(-Imş) / /A/R,)) For the last of them, we use E[(W, 2)//1/2 < O(1A)).
Combining all of the above bounds, and naively bounding the edf of the Gaussian via
F(-lm5) / A/ se (mm)) (mm]))22(166))))) x/(161;))
arrive the claimed bound.
33

---

Published as a conference paper at ICLR 2024
As consequence of Lemma C.16 can deduce that the quantities appearing above satisfy the following large À behavior.
Corollary C.17. For every x such that m m 2 log for all i,
VX
m A (0(m))) (oo0(((mm))))) 4 gl 10-mm; )o - g(-m"))) O(A-1)
1
m A. 1 jo(m))o(o (omi))((((mm))))) + 1 ( jo(-mp)o(( g( 9(-m")) + O(x-1)
A O(
Without the assumption on m. m., the same holds with O(x-1) replaced by O(X-1/2) as long as the indicators from g(m²) = m, m >0 are replaced by the "soft indicators" F(m / A/Ra).
Proof. We begin with the estimate on m A.". This quantity can be split into four terms correspond- ing to whether the mean chosen for Y is ,t -- 14, v. Namely, if we let Y, _ + Zx.
200 A = [Bl(u Y,) ww) g (WWY,) + [E(( w,, g(WY-
+ + Y.) Y.) (WYY)) + E[((u Y.,) g(WY
Now notice that each of the four quantities are of the form of Lemma C.16 with b 14, and = 14, -1, 14, V4 -v respectively (the change of the sigmoid possibly having a negative sign on its argument is realized by switching the sign of u since that is the only place it appears).
It is easily seen that so long mm!) 2 (log A) V,, then
F(m ())R..) = 1 fm > 0} + 0(1A)-
By similar bound on the error term on the right of Lemma C.16
B(uu- W
g(WY,)] = 1 m >00 (-v g(m))) + O(1/A)
and the cases where it is Y, contribute 0 + O(A-1) since bounds for my A and A. this gives the first part of the corollary.
O. Together with the analogous
If we drop the assumption on mf. we find from the general inequality are -e-=2L S C for some
VI. O(1/v)) as
uniform C, that the errors in the right-hand side of Lemma C.16 become at most claimed. Leaving the soft indicator in place, this gives the second part of the corollary.
We deduce from the above approximation and trivial bound of O(1/A) B A limiting dynamics. However, one has to be slightly careful about the À + do limit near the hyperplanes where m. or my = 0; Therefore, in what follows we envision a different limit associated to each orthant of the parameter space (associated to the signs of m,"m"") This is reasonable because in each orthant, the = oo dynamical system initialized from that orthant stays in that orthant. The following proposition captures that limiting dynamics (all orthants being written into the below simultaneously, with ambiguities at m being therefore omitted).
Proposition C.18. In the À limit, the ODE frofrom Proposition C.I5 converges to
é = (9(m)) g(m²)) + g(-mm))))) ( g( )mm (olmm) (o g(m")) g(- (o o(-m²)) BBu
m. = 2/
4 (1m )(m"))mm)))))) 1 (0000000(0000 g(-m")) Bm
m = 4 (1m g(m")) Amp-ooo-r o(-m")) 3m v
2/
and R1J = R, for lsisisK.
Moreover, for X large, the difference in the drifts above and those of Proposition C.15 are of order 0(1/2) is (1 (mm are all at least V(log X)/X, and O(1 (1/VVX everywhere.
34

---

Published as a conference paper at ICLR 2024
C.4.2 CONFINING THE SGD TO COMPACT SET
Similar to the k-GMM case, our first aim is confine the SGD to a compact set so long as B > Lemma C.19. For every 3 > 0, there exists L(3, v(xo)) such that for all À large, the dynamics of Proposition C.15 stays inside the (2-ball of radius L for all time.
Proof. We start by considering the evolution of the e--oorm of all the parameters, (v, m/ m²). If we use the shorthand
g" = g(m"), 9 g""=g(-m"), r" = o o== oovv. g
and similar quantities with instead of n. short calculation shows that we get in the oo dynamical system of Proposition C.18
| 1 0".g"O" + ....g"o tu.
g" to -
2
Notice that 6 goes to 0 as v ge co, and moreover, (- 0 as co. Therefore, there is a uniform bound of 1 (in fact, (1/e) where W is the product log function) on v g Similar uniform bounds apply to the other three terms, that in total.
| < 2 2/
which implies in particular that drift negative when ||v||² is at least L(3). Similar arguments go through mutatis mutandis for (m and (m These then imply that the drift when À is finite are similarly negative when ||v||2 = L(3) as long as À is sufficiently large per the approximation from Proposition C.18 Altogether, this implies that for the dynamical system of Proposition C.I5 mmI], stay bounded by some L(3) for all time.
all of [|v]l, m²] and
Using that, if we consider the expression for R2 and use the boundedness of A; B, the drift in Proposition C.15
gives
|R TR-11
< B 1p 11 44
The right-hand side is at most C2 bounded by some constant L'(3) for all time under the dynamics of Proposition C.I5 Altogether these prove the lemma.
2 R|22 for some Cl, whence by Gronwall, R | also stays
Corollary C.20. The trajectories of the ODEs of Proposition C.13 and Proposition C.18 are within distance ( /VX) of one another.
Proof. Let and ü be the two solutions to the À finite and À ballistic limits respectively. Then
|| ûl ZK (1 25) ا|| + ((-1//2),
where we used that L bounds the norm of and ü for all times by Lemma C.19 and that the Lips- chitz constant of the sigmoid is 1. By Gronwall's inequality, this implies that u ( as claimed for some C depending only on 33.
C.4.3 LIVING IN THE PRINCIPAL DIRECTIONS NEAR FIXED POINTS
The last thing to conclude is that the fixed points the oo dynamical system are indeed living in the principal directions desired by Theorem 4.2 By the above arguments, for all 2 To (e) the dynamics is within distance e of one of the fixed points of Proposition C.18 Let us recall the exact locations of those fixed points from Ben Arous al. (2022).
0 < B < 1/8, then let (Iu, It - - II,,,) be any disjoint (possibly empty) subsets whose union - 1,),, ), is a set of fixed points that have {1, (1....,A K)- Corresponding to that tuple (To, It, u -,,,,,,
R 0 for all i, j. and have
1. mf m n2 for lo.
2. mf e; 0 such that Eier, of logit( (44)) and my 0 for all i e 1t+
35

---

Published as a conference paper at ICLR 2024
3. -m, 0 such that Eier, 2/ Oogit(-43) and m 0 for all E I
4. m = V; such that Eier of logit(-4.5) and my = 0 for all i 17.
5. - = 02 < 0 such that Σier, of = logit( 4:3) and m = 0 for all E 15.
The following observation is easy to see from the fixed point characterization described above. Observation C.21. Suppose that x4 is a fixed point amongst the above. Then
(xx.) E Span(μ, v).
* v(x.) Span(g" (ppan(g"(xx).5)
(with no error).
The first claim that (x.) E Span(p, v) follows from the fact that R = 0 for all Furthermore, if 1, 1, are empty, then it lives in Span(u), and similarly if J 4 are empty, then lives in Span(v).
Proof.
For the next claim, observe that (x.)) (sx.) is the vector that is m coordinates belonging to It and 0 in all other coordinates. g g=xxx, is the vector that is -m" in coordinates belonging to I=p- and zero in others. g (1x)) is m² for coordinates in T zero else, and (xx)) is -m² on It-
Since the I-sets are a partition of (1...., K} it evident that we can express
v(x.) = (x.) + (xx. 9 (x.) - (x.), implying that indeed v(x.) lives in Span (x (x.)).- ((x)), ( (5.),
Section 9.4) that there is a transition at 3 = 1/8 the regularization, and long as 1/8, with probability with respect to the random initialization, the SGD converges to a fixed point having 1o- = 0.
was further argued in (Ben Arous et al. 2022
Lemma C.22. Suppose 3 <. 1/8 8 and consider the initialization
vy ((0)~~(0,1))
and
(0), m (0), R1)(0) 80 and R ~ 81,
(C.26)
where since the dynamical system of Proposition C.10 defined per orthant, the do are understood as 1/2-1 /2 mixtures of dg- and 8y-, With probability over this initialization, the dynamical system of Proposition convergee fixed point as characterized above, above having lo 0.
Putting the above together, we can conclude our proof of Proposition C.2
Per Proposition C.15 the limit of the summary statistics of the SGD along training the solution of that dynamical system initialized from m² mn, , Roy ~ 6y for i + (with equal probability of m m. being 6g+ and d0- we need to distinguish which orthant is initialized in), R ~ dj. and 2 ~ (O, 1) i..dd
Proof of Proposition C.
For the first part, notice that the norms || [x R + (m +m,)) and [ Σk uy are expressible in terms of the summary statistics. Therefore, their boundedness follows from Lemma C.19 pulled back to the summary statistics of the SGD per Lemma C.6
For the second item, by Corollary C.20 and Lemma C.22 together with Observation C.21 for every >>> O. there is a Tp such that for every Ty. for all t E [To, Tyl, the dynamical sys- tem of Proposition [.155 is within distance + --2 of point u, having (v:(!) > and max (Im?", [m]]) > 1) in and having that its first layer lives in Span(/4, v) and its second layer lives in Span((g (u,) (e())tt)) This is pulled back to the summary statistics applied the SGD
trajectory per Lemma C.6 (with the observation that (g" (x))a are functions only the summary statistics).
36

---

Published as a conference paper at ICLR 2024
D CONCENTRATION OF HESSIAN AND G-MATRICES
We recall the general forms of the empirical test Hessian matrix V V² R(x) and G-matrix G(x) from 2.2 respectively. Our aim in this section is to establish concentration of those empirical matrices about their population versions throughout the parameter space.
D.1
HESSIAN AND G-MATRIX: 1-LAYER
We first prove the concentration of the empirical Hessian and G-matrix for the k-GMM problem about their population versions, which we analyzed in depth in Section [ This concentration will be uniform over the entire parameter space. Namely, our aim in this section is to show the following.
Theorem D.1. Consider the k--MMM data model of (3.1 and sample complexity M/d= a. There are constants = c(k), C Ckk (independent of A) such that for all t >> O, the empirical Hessian
matrix concentrates as
sup P(I xcRéd
((((x))=[[((x))) ((I (R(x) - E( > t) S exp ( [eaat C -c
(D.1)
and so does the empirical G-matrix
sup P([[((x) [[[(()]]lop > S exp (-)
(D.2)
XERAI
Proof. In the following we fix x E Rdd and simply write R(x), G(x) as R,G. Let Ā = (Y .rl...yM denote the test data matrix and let
(b)) D diag(Ty- (c)dbe (c)
(D.3)
DC = (V' (b) T T + (C)())
(D.4)
where (c) was defined in AA.4p. Then D D are M x M diagonal matrices for each pair be E [k]². We denote the (kM) x (kM) matrices D (D and DG (D%) and the
dk x kM matrix
x 4. @ A. With these notations, per (.2)). we can rewrite the Hessian and G-matrices as
22A
7 x D DD (AKAyT
(D.5)
M
â 1
x D (AxxyTT
(D.6)
M
To prove that the operator norm of 2 unit ball in dd to show that the following concentrates
E R( concentrates, we'll use net argument over the
sup vE(R*)).Ilv[1==
(v,
(D.7)
where v (Ve)ce(k) E (Rd)+
By plugging D.5) into D.7 we want concentration estimate for F(v) = (v, v2 - - which we can rewrite as follows.
F(v)
1 (v, T X
"D (A(XXTTVV) -
M
Σ
Vay ÄDH T m))
(va: EE| AT 7t A=|wbb
M
(D.8)
b.b
de(a, (v,, 80) - E [de b)( (') 84].
4.6
37

---

Published as a conference paper at ICLR 2024
where de(a, b) [Dable Recalling Veershynin 2018 Section 4.4.1) that for e-net Nr of {v ||v| = 1}, and any real symmetric matri H the tollowing holds
sup kv, Hv)l 26 VEN.
|op
sup 1(v, Hv}l 2 sup ((v, Hvy [|v|l==1
veN.
So by union bound and the e-covering number of {v E Rdd |||||| = 1}. we have
PP sup
[P())>() S [&l sup P((P(v)) 1/2) S |
(D.9)
vEN.
so long as < 1/4, say.
To control the last quantity Pv] > 1/2), we notice that F(v) is a sum of O(1) many (in fact h22 many) terms (corresponding to each pair a, b) so it suffices by union bound to control the concentration of each summand. That is, suffices to understand the concentration of
1 Σ
(dela, (", 82) B[dr(a, () ( 8-1)
(D.10)
M
We recall that the test data
ye= Σ (E(M]
Vaka 2x == n² +zx
(D.11)
where z are i.idd. (O Ta/A), and jf Eaelk] Pa5 In this case, note that de(a, b) =
Ty. (b) (() are uniformly bounded by 2, and that for s s M
(Vos ye) - (Var (') + (vas 23)
(D.12)
are i.idd sub-Gaussian (with norm O(1)) for fixed Vas since ||uall2 1 and |elll2 for all c, and À 2 say. Thus, for all a, b, the products (de(a, () (16, Y")) are i..d. uniformly sub- exponential random variables. As such, (D.10) is sum of wi.d. centered uniformly sub-exponemtial random variables, and Bernstein's inequality yields that there exists a small constant c c(k)
P(] > 1/2) S exp(-cMM A E)).
(D.13)
follows from plugging
D.13 into D.9 and taking e = 1/8, that
[F(v)] > t/2) < exp (
sup ||v 1a
which yields the concentration bound (D.1) for empirical Hessian matrix, by noticing that M/d a.
For the proof of the concentration bound (D.2) for empirical G-matrix, thanks to (D6).. we can
define F(v) (as in D.8)) using DG instead of D H Noticing that the entries D ab are bounded by 4, the concentration of this new F(v) follows from Bernstein's 's inequality by the same argument. Then the concentration estimates of this new F(vv) together with a epsilon-net argument gives (.2)
D.2 HESSIAN AND G-MATRIX: 2-LAYER GMM MODEL
For the reader's convenience, Section B Our main aim in this subsection to prove the following analogue of Theorem D. for the XOR GMM with a 2-layer network. There is small problem with differentiating the ReLU activation twice exactly at zero, so let us define the set Wi = -nskk (W; * 0}.
recall the empirical Hessian and G-matrix from the beginning of
Theorem D.2. We consider the 2-layer XOR GMM model from (4. (4.2) with sample complexity
M/d= =dd=oo a. For any L. there are constants = c(K, C(K, L) such that for all 1 >> 0, the empirical Hessian matrix concentrates as
sup (((VV((((e,W)) W) Ei W) > t) S -x Cld},
(D.14)
SL
W WEWS
and the empirical G-matrix concentrates as
sup P([]@(e, W) E[Â(v, W)llbop >t) <c))(-(on((At] Cld}. W
(D.15)
38

---

Published as a conference paper at ICLR 2024
Proof. We begin by recalling expressions. Letting T' o( ((((g((W))) as in B.1 by (B.4)
we have
M
V
R
- 9") (WWP3
M
1=1
M
V
R
1
"o (W, F') --") + 1,,0001 -' (") ( P900000
W.W,
M
1=1
Vww, R
(0' T') Yee, + , 3 (W, g' 8))) &
M
(D.16)
and by B.2 B.3
we have
A7
G
T7' - 8') -
,(WY")) & ((WPg)
M
M
Gw.w,
10' gu) (W (', p (D.17)
=
M
Gow, M
u (W, Y Do v
Recalling the data distribution from long W Wi. almost surely, all the coordinates of WY are nonzero: as long as X < x0.
P(Ii, W 0, for all W W
(D.18)
For any W We thus can ignore the second derivative term g"(W; Yt) in W.w, R. so a.s.
M
V W.w,RR
R
Σ M E=1
(oid' (W, (oss' (1)
(D.19)
With these calculations in hand, the proof is similar to that of Theorem D.1 the main differences. Let B the ball of radius in parameter space. Fix (v, W) E B and simply
we will only emphasize
write R(e W), Gev, W) as R, G. Our goal is to prove concentration of the operator norm
(c. (a))),,
sup (R, ( [|
where a = (a), 02 ax) and u= (i, u2, (u)) E (IR)), As the proof Theorem D.1 specifically (D.9, by an epsilon-net argument, we only need prove concentration estimate for the following quantity
F(a. u) = ((a, u), (R- (a, u)),
(D.20)
individually per (a, u) (a, u)ll The inner product on the right-hand side of (D.20) splits into (K + +1) terms corresponding the terms in (D.6)) and (D.19):
Ti Σ
(a. V2 Ra)
01 o(W2))²
M
1
(ui, Vir W.w, Ruy))
(ngg (W. (") (W, 89)5"(1
-1 P"),,, V9), (D.21)
M
t=1
M
(a, 22 Rew,
Ruy)
1
Σ
(agg' (W, 89 - 56) + D'g 89 (9)) ,u))
=
M
t=1
39

---

Published as a conference paper at ICLR 2024
The concentration bound for F(a, u) follows from concentration bounds for each of the above terms about their respective expected values. Each such term minus its expected value will be a sum of M i.idd centered random variables. It remains to show that the summands are uniformly (in À and BL) nab-exponential. Then the concentration of F(a, u) follows from Bernstein's inequality as in (D.13).
7 are bounded by 1, and U; vj and llg'lloo are all bounded by C(L). By the same argument as in the proof of Theorem D.1 we have that (ug, Ye) and (ug, Y are uniformly sub-Gause Gaussian..
To that end, we notice that (0') (11--87 )))
Next we show that (OWW)) aa) is also sub- Gaussian. In law, conditionally on the mean choice among tu,, HV, we have wýee & W(+ Zx) or wÿe á W(# + Zx). Since (a, u) E B and || = |||| = 1, we have by Cauchy-Schwarz that [|W
< L. Then we can write
(e(WY'), a) = (g(W Zx), a} + (o(WW - g( a). By the uniform Lipschitz continuity of g. we have
Ko(WW*) g( Zx), a)l = [|]] + || Al) 2.
We also notice that if 11 viZx, then
K
1 Vx |00ll
L
Zz, (o Z2), a)
ogg (W:))): <
|||||||| W 2-
VX
implying that it is uniformly L-Lipschitz function of standard Gaussians, from which it follows that ( Zx) is uniformly sub- Sab-Gaussian (see (Vershynin) 2018 Section 5.2. For the expectation we have
B(o(W2/VV)) a\]l s [F)( a}]l + B(g(W//V - g(0), a)ل
- lø(0) V K |al2 + |1/ B ||| / //V/lalllall Allaalaall
- (o( V K |a/2 + - 1lo' | 1 /T// C(K,
We conclude that (WWY)), a) is sub-Gaussian with norm bounded by 0(1) (depending only on K, L.) As consequence, each summand in (D.21) is sub-exponential, with norm uniformly bounded by ((1) (depending only on K, L.) Bernstein's inequality gives that there exists small constant e L)
P( u) > t/2) S exp(-cM(( A t*)
(D.22)
follows from D.22 and the epsilon-net argument over the unit ball of (a, u) E R x Rkd as in D.9. that
P
sup
|F(a, u) > t) < exp (- -
[a.w)]==1
which yields the concentration bound (D.14) for empirical Hessian matrix, by noticing that M //d- a.
For the proof of the concentration bound (.15 for the empirical G-matrix, thanks to (D177, define F(v) (as in (D20 using & instead of 2 R. Then we need concentration bounds for the following quantities
M T
1
(a. Gvea)
M
M
(us, Gw.wyug) =
1
TJ - Jf) (W, D'g 8d) 8' F9.
(D.23)
M
M
(a, Gow, uz)
1
3 - 3 uW, (a, {9)) 84.
M
40

---

Published as a conference paper at ICLR 2024
By the same argument that after (D211) each summand in (D233 is uniformly (with constant only depending on K, L) sub-exponential. The concentration of this new F(a. u) follows from Bernstein's 's inequality by the same argument. Then the concentration estimates of this new F(a, u) together with a epsilon-net argument gives (D.15.
E PROOFS OF MAIN THEOREMS
In this section we put together the ingredients we have established in the preceding sections to deduce our main theorems, Theorems 3.1
E. I-LAYER NETWORK FOR MIXTURE OF GAUSSIANS
We begin with the theorems for classification of the k-GMM with single layer network.
We prove the alignment (up to multiplicative error O(c+>-1)) of the SGD trajectory, the Hessian, and the G-matrix, with Span( .... (4x)
Proof of Theorem 3.2
time.
The results for the SGD were exactly the content of item (2) of Proposition
Namely, that
item tells that the part of x; orthogonal to Span( -..; (bk) has norm most Ole x-1 while
||x]|| 2 7 (e + )- 2 2//2 for an vj independent of e, A. Absorbing 1) into the big-O, this is exactly the definition of living a subspace per Definition 2.1
We recall from Lemma A.I for b # c, the be-block of the population Hessian is given by
R(x)) = Σ 2 + ER
(E.1)
LE(M]
where IT's = (c) for Y, μr + Zx. and where ER is matrix with operator norm bounded by O(1/A). By Lemma A.2 the aa diagonal block of the population Hessian given by
E V R(x))
Σ pill'arf ER [E(1]
(E.2)
where IT = (a) - Ty, (a))) and ER is matrix with norm bounded by ((1/A). lo
The expressions 1E.T) and [..2.. together with the concentration of the empirical Hessian matrix Theorem D.I liplies that blocks of the empirical Hessian concentrate around low rank matrices; i.e., for each x, we have
(2(((x)= = Σ
T 602 + +EB + (V v(E[((()))
1Sick
(E.3)
Σ 151<k
pill
32
ER
where for each fixed x, except with probability (ccos"-C) we have | S c Since the test and training data are independent of one another, this also means for any realization of the SGD trajectory (Xe)egT,8-1, and so long T cld)) by union bound get
PP
max 8-T
V N(xe)) - Σ p (x) (x 2 C(e C(++ = oa(1).
(E.4)
LE(R
where we' we've put in the X to emphasize the dependence of II',, the location in parameter space.
Moreover, we recall from item (1) Proposition C.I that there exists a constant such that for all o1),,, the SGD trajectory has ||xx||| < L for all To8-1 1 < l < T,8-1 follows by definition и that there exists a constant = -((.) > 0, such that the coefficients in E. and (E.2) are lower bounded: IT (xe). IT (x²) 2 c for all 1.5---6 T Thus the first sum E1stee p 82 the righthand side of E.3 positive definite, and its norm is lower bounded by (uniformly in e, A). Together with E conclude that the b, blocks of the test
A large, with probability
41

---

rewrite the Published as a conference paper at ICLR 2024
Hessian V (xx)) live in Span(u1, Ma,... ... .Mx) up to emror O + 1) Namely, this is because it satisfies Definition 2.2 with the choice of M ER after absorbing e= into the big-O.
By the same argument, thanks to Lemma A.3
the be block of the population G-matrix is given by
dbcPo 00 - PeEltyy pa ty, (c)lu8 + Σ
+ &
(E.5)
together with the concentration of empirical G-matrix Theorem D.1, and a union bound over TT implies that blocks of the empirical G-matrix concentrate around
where ER is a matrix with norm bounded by (11/A). The expression E.5
Vhl G(x,) = 6 @2 PrElTy, (b)l - po + Σ (
EG (x;),
(E.6)
where except with probability T -C)d we have [xx) S Se+
+ X for & & TT
Namely, the analogue of (E.4) will apply V2.G about the low-rank part of the above. In order to deduce the claimed alignment Theorem 3.2 it remains to show that each block of the matrix
in (E.6) (minus &G has operator norm bounded away from zero uniformly in e, A. Towards this, recall that we can work the event that ||xe|| S L. In the on-diagonal blocks,
part of E.6 that is not EG as
PL TY, (b)
Σ P.B Ty (b) 19th
(E.7)
This matrix is positive definite and since ||ee|| < L, there exists c(L.) such that the coefficients of pe are all bounded away from zero by c, so that the norm of the above is bounded from below by zero for all (x))),-'-(ETTT5
Thus by C.I the coefficients P, Ty, (b)) p/ (b) in (E.7 are lower bounded. Thus is positive definite, there exists a constant e >> O, such that its norm is lower bounded by c.
For * can lower bound the operator norm of the matrix
-PoBlTy. Ty. (c))lue (T) - Ty; (b))lu 82 + Σ (1)
tyb.c
(E.8)
while |xe|| < T, using that if k: then the last contributes some positive portion outside of Span(He; (b) which can be used to lower bound the operator norm by some o( L) and if k = 2, then the first two terms are the only two, are negative definite, and have coefficients similarly bounded away from zero by some -c(L).
Proof of Theorem 3.7 The theorem follows from Theorem together with the observation that the matrices in E.I and (E.5) that not the error portion ToT and E respectively, are of rank k since they are sums rank-1 matrices and as explained in the above proof, each of their eigenval- ues bounded away from uniformly in c, X, (in manner depending only on L.).
Proof of Theorem 3.3
The claims regarding the SGD are proved in Proposition C.
It remains prove alignment of the top eigenvectors of the ec-blocks of the Hessian and G-matrices with Me. Recall from EE3) the decomposition of the empirical Hessian matrix
2 R(x) =: Σ Pi
662 + ER
(E.9)
ISICK
where IT'
ea = EfTTy ny; (a)) and with probability e- we have | S &
+ X-
By our assumption that μ11, /22, 14 orthonormal, thus the first part in the decomposition EE can be viewed as an orthogonal decomposition. V2 R(x) perturbation of Eistsh PilI by agg- In turn, by the same reasoning as in Lemma C.11
IT!
II + +0(--)) where
π ([m((a((1 ((()))
42

---

Published as a conference paper at ICLR 2024
for T((a)= E Eo e as in CC.66. By Lemma C.19 there exists c(B) > 0 such that for all (Xe)TTB- (xe))o8-1stsT/ 1, all the coefficients pill'a (E lower bounded: pıll 2 Thus EisisA 02 has positive eigenvalues, {prlly hisiss and each of them is lower bounded by c. The associated eigenvectors are given by {1iss We furthermore claim that the one corresponding to / is separated from the others uniformly in A. This follows from the fact that we derived in Proposition C.14 that X₂ is within O(c + X-1) distance of a point x, such that TT(a) - (1 (c)) as long a b. This implies that (t(a) for + is closer to 0 than (a) is close 1, whence T. (a) Fa (a)) > Tr(a)(1
5r (a)) by an amount that is uniform in A. This ensures that along (xx) the largest eigenvector in E, Maa u is the one with eigenvector are those corresponding to (fi)ta Altogether, by eigenvector stability, the aa R(xx) lives in Span(pa) up to error O(e O X ) and the next k - 1 all live in Span((4))tza) up an error OE + L/A) for all
Na and the next
top eigenvector of V
"Tb8-1,T663
The statement for the empirical G-matrix is established similarly by using (E.6) as input. The part that different from the above to see that its top eigenvector in the aa-block is the one corre- sponding approximating by Ha and its next k - 1 are those approximating (14 (H4))taa For that, recall the expression (E.7) and use that since (1 R. (a)) > (a(a), the coefficient of n in the aa-block strictly larger than the coefficients hn for + a.
E.2 2-LAYER NETWORK FOR XOR GAUSSIAN MIXTURE
We now turn to proving our main theorems for the XOR Gaussian mixture model with 2-layer network of width K.
We prove the alignment individually for each of the SGD, the Hessian ma- trix, and the G-matrix. For the SGD trajectory, the claimed alignment was exactly the content of item (2) in Proposition C.2
Proof of Theorem 4.2
Letting F denote the cdf of a standard Gaussian, thanks to Lemma B.4
B
Σ
o'lo -(Wo)) ER
(E.10)
4
de(+
n de dE(#,ttt))
V
m
0o2
cR
W.W,
(E.11)
Ris
where the two error matrices satisfy ||ER| |Ie O(IvA).. The two expressions E. 10 and (E.II). together with the concentration of empirical Hessian matrix Theorem D.2limply that for every L. every fixed x = (v, W) with & Wj. the blocks of the empirical Hessian matrix satisfy
V2 A(o
Σ
d'(v ((W)))
En Con
(E.12)
dE (=b
V Rlw,
Σ
m
A Rii
002 +
ER W.W.
(E.13)
where except with probability e (cns2-C7) we have
1
I (v. Willop:
ER (v, Wo Set
(E.14)
VĀ
Using independence of the test and training data, recalling from item (1) of Proposition C.2 that the SGD stays confined ball of radius parameter space for all < Ty8- and from item (2) of Proposition C.2 that |W]]| 0 for all i for all TØ8-1 < & < Tyō-1, taking a union bound over
43

---

Published as a conference paper at ICLR 2024
{ $ T, we get for ag.
P
22 R(xe)
1
Σ DE(=
d'v g(Wo)) 2 O(e 0(E+1-1/2)) =[
max
Toó StsTyo 6--1
(B.15)
www,
R(xx)
Σ #E(#)=)))
F
mi
A Rii
0
2 C(s (((=+11133) o(1)
ToE E
(E.16)
where in the quantity that the blocks of the empirical Hessian are being compared v, W evaluated along the SGD trajectory X₃.
It remains to show that the low-rank matrices in E..5 E. have some operator norm uniformly bounded away from zero to deduce the aligmment of the empirical test Hessian with the claimed vectors. By item (2) of Proposition C.2 there exists some constant >> 0 uniform in X, such that with probability 0a(1), for all T the SGD x₃ is such that
max d' (v- g(W 0) ())lo()))) (Etrtt)))
o' E ".g(m")) > c, (E.17)
max de (Etrddd))
Thus the deterministic matrix in (5.15) and its norm is lower bounded away from 0 for all (Xt)ssT,8-1. Together with E. we conclude
Eve (tv, ('v g( (W is positive definite
that the second-layer test Hessian V ((xe)) live in (pan(g (X(())))) +))e(+n.tv)).
For (6, by item (2) of Proposition C.2 there exists 0 (independent of e, such that with probability 1 - o1(1), for every i, it must be the case that
vf F [mm
>C A )
max (E(
V Ra
for all (xe)egT,8-1.
(E.18)
Then the deterministic matrix in 16 is positive definite and has norm uniformly lower bounded away from 0 for all x; for é < Tyo Together with (E.16), we conclude that the first-layer test Hessian V w.w. R(x;) lives in Span(p.
By the same argument, thanks to Lemma B.5 together with the concentration of the empirical G-matrix from Theorem D.2 the blocks of the empirical G-matrix concentrate around low rank matrices
(v, W) 1
Σ
((O)) vo) 9)W EG
de (=(
(B.
Gw.w( W) u A W,W,,
where recalling B.18
A =
(yo -oo ((to:)(WW))))))) F(mm V F(m!
0002
(E.20)
Ra
dE (=)
and for every= x = (u, W) of norm at most L, except with probability e (once we have Union bounding over the T 1 points along the trajectory of the
| 11 w. S E
SGD, as in the lead-up to EE.5 E. we get
1
P max [XGe)
(alv- g(W)) - 9 2 C(e = o(1)
EST,8
4
de (
(E.21)
P max [|G..w, (x:) - Allp C(=+ o()) egT,
(E.22)
where in the quantity that the blocks of the empirical G-matrix are being compared to, v, W are evaluated along the SGD trajectory Xe.
44

---

Published as a conference paper at ICLR 2024
We recall from item (1) of Proposition C.2 the SGD stays inside the b-ll of radius L for all time. Thus for x = Xg, the coefficients (()))() and (E.20) are lower bounded away from O. By the same argument for the test Hessian matrx, conclude that the second-layer test G-matrix Gve (xe) lives in Span(g (x:))) its first layer lives in Span(p, v), up to error +
(((ug((W))) (l--o((((WW))) in (E.19)
and
Proof of Theorem The first layer alignment follows from Theorem 4.2 together with the ob- servation that the part (3.13) besides ER w,W, is a rank-2 matrix with eigenvectors μ., and having both eigenvectors bounded away from zero uniformly in (c. A); this latter fact comes from the ob- always, and o; being bounded away from zero per part (2) Proposition C.2 A similar behavior ensures the same
servation that the sum of the two cdf's F(a) + (faa)
G-matrix's first E.20 and uniform lower bound the sigmoid function over all possible parameters a ball of radius L, as guaranteed by part (1) Proposition C.2
layer top eigenvectors per (E.19
For the alignment of the second layer. it follows from Theorem together with the follow- ing. First observe that the part of (2.12) that not ER rank-4 matrix with eigenvectors (g(Wo))o((txxtt)> To reason that all of the corresponding eigenvectors are uniformly (in e, X) bounded away from zero, use the uniform lower bound on o' while the parameters are a ball of radius T. about the origin as promised by part (1) of Proposition C.2 together with the uniform lower bound, for each on one of (W; (WV--)ee/==tte +)eE(=p,to) which holds for the SGD after To(e) per the proof of part (2) of Proposition C.2
F ADDITIONAL FIGURES
F. ADDITIONAL FIGURES FOR K-GMM
In this section, we include more figures depicting the spectral transitions for the h-GMM model.
0.2
0.3
0.22
0.01
0.01
-0
-0.03
3050
130
1.0
075
0.7
050
0.0
025
800
020
1000
Figure F.1: x² and x at initial and finial trianing times. Initially they are random and then they correlate (and anti-corellate) per Theorem 3.2 Parameters the same in Figure 3.1
45

---

Published as a conference paper at ICLR 2024
1 blcc hew speciuuee
Stecir lecann caedmnn
Heosly
12-
Figure F22: Further illustration of dynamical BBP from Fig. 3.3 Here that there is initially two components of the spectrum, then as time evolves, the top eigenvalue and the next 9 separate from each other.
0.07
0.00
0.01
0.00
0.00
-0 02
0.00
200
400
600
8000
1000
200
400
620
2003
3000
1.5
150
125
125
1.00
075
025
000
0.0
-025
-325
200
400
500
800
1000
200
1000
Figure F3: x² and 2 at initial and finial training times. Initially they random and then they correlate (and anti-correlatel) as per Theorem 3.2 Parameters are the same as in Figure 3.1
F.2 TRAINING DATA
We include here few plots the preceding discussion run with training data, as opposed test data. In the following, the plots are generated using the Hessian of the empirical risk computed over the entire training data set and similarly the G-matrix generated using the entire training data set, as opposed to the test data as in earlier plots.
F.2.1 K-COMPONENT MIXTURE MODEL
coreeaaiin
eigenspace blocks
correlation with
of blocks
03
05
05
0.
0.4
02
Figure F.4; Parameters are the same as in Fig. 3.1
46

---

Published as a conference paper at ICLR 2024
Evo usion of Elgenvalves
x mean
0005
23
0015
0010
05
0.00
Figure F.5: Training data analogue of Fig. 3.3 the cigenvalues of 2 (xx)) over the course of training (but with training data. The leading k eigenvalues are separated from the bulk at all times, and the top eigenvalue, corresponding to μa separates from the remaining eigenvalues soon after initialization. Right: the inner product of x with the means Mi, Mi undergoes similar separation over the course of training. Parameters are the same in preceding figures.
F.2.2 XOR MIXTURE MODEL
70
81
74
241
45
(a)
(b)
(c)
(d)
Figure F6: Training data analogue of Fig. 4.1 and (b) depict the alignment the first layer weights (xxx)) with the principal subspaces of the corresponding blocks of the Hessian and G-matrices fromed using training data, i.e., with (V w.w. R(xx)) and Ez(G w. (xx)). (c) and (d) plot the second-layer alignment, namely of (xx) with Eg (V3 R(xx)) and E (Grv (xx)). Parameters are 500, À 10, and K 20
algerwalues vs
eigenvaluuss time
0006
0014
0012
0.00
0.64
Figure F.:: Training data analogue of Fig. 4.2 The eigenvalues of the ou blocks of the Hessian and G-matrices time from random initialization. Along training, the top four eigenvalues separate from the bulk, and four outlier eigenvalues emerge, corresponding to the four "hidden" classes in the XOR problem. Parameters the same in Figure 4.1
47


Reviews:
[{"review_id": "I8h9lVXwGB", "reviewer": "Reviewer_BqTJ", "summary": "This paper studies SGD applied to two problems, namely multi-class logistic regression and XOR classification with a two-layer network. It is shown that in both scenarios, the training trajectories of SGD aligns with both the hessian and the Fisher information during training. Moreover, for the XOR problem, the alignment occurs in both layers of the network.", "questions": "- I'm trying to understand Theorem 3.1. Intuitively, as $\\varepsilon$ approaches 0, $T_0$ should grow larger. Is it possible that $T_0$ becomes so large that the theorem becomes vacuous, i.e., $T_0>M/d$?\n- Is there a counterpart of Theorem 3.3 for the XOR setting?\n- Is there any specific reason why the authors study the Hessian and Fisher information that is obtained from the **test data** instead of from the **expectation w.r.t. the data distribution**?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The results of alignment are technically strong.\n- The reasoning is solid and convincing.", "weaknesses": "1. The major weakness of this work is that the motivation of studying the evolution of the Hessian/Fisher information is not super clear. According to the introduction and the related works, from my perspective, the best thing we know about the relationship between the Hessian and the effectiveness of SGD is that \"this common low-dimensional structure to the SGD and Hessian matrix may be key to many classification tasks in machine learning\" (Gur-Ari et al., 2019). From my point of view, there is no solid evidence indicating that the main results of this work imply an advantage or a disadvantage of SGD. I will consider raising my score if this question is properly settled.\n2. Figure 3.3  and Figure 4.3 are a bit confusing without legends.\n3. In the paragraph following Theorem 4.2, the terms \"GOE\" and \"Wishart matrix\" may appear unfamiliar to many.\n4. In the notation of covariance matrix $I_d/\\lambda$, the $\\lambda$, which is typically a weight decay parameter, is used as the inverse of variance, causing some confusions to me.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qsGqZzytrW", "reviewer": "Reviewer_K4sY", "summary": "This paper studies high-dimensional SGD by examining the alignment between the SGD iterates with the emerging outlier eigenspaces of the Hesian and gradient matrices. It is shown that after a short period of training, the SGD iterates start to align with the low-rank outlier eigenspace with the empirical Hessian and empirical G-matrices, and that the alignment may happen layer-wise for multi-layer architectures. The main results are proved for two settings: 1) learning a Gaussian mixture model with linearly independent classes by a single layer neural network, and 2) learning a Gaussian mixture model version of the XOR problem. Numerical evidence is also provided to illustrate the theoretical results.", "questions": "- The learning rate $\\delta = O(1/d)$ seems to be on a very small scale, especially for large $d$. Could the authors comment on this requirement for the learning rate?\n\n- The main theorems state that the alignment happens for $\\ell \\in [T_0\\delta^{-1}, T_f \\delta^{-1}]$. What's the corresponding generalization performance of the model on the test data?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper is well written and easy to follow. The results are novel and solid, and the presentation of the technical results are very clear with  the numerical illustrations of the theoretical predictions.", "weaknesses": "- The results are proved for two somewhat restricted settings, and it seems not clear how to extend the argument to more general settings. It would be helpful if the authors can comment on the limitations of the current proof.\n\n- It seems that the main technical tool is from Ben Arous et al. (2022). It would be helpful to add some discussion on the technical novelty.\n\n- The results are valid for online SGD. It is worth a comment whether this is necessary, as well as what happens for multi-pass SGD.\n\n- A small typo: In the sentence above Theorem 4.1, \"principal direcitons\" -> \"principal directions\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "okkxR3X9dP", "reviewer": "Reviewer_WpLy", "summary": "The paper investigates the interplay between the training dynamics of one-pass SGD and the spectral decompositions of the Hessian matrix and G-matrix. The authors study this interaction on two high-dimensional classification tasks with cross-entropy loss. \n\nThe authors start with the classification of k-component Gaussian mixture models by a single-layer network. The outlier-minibulk-bulk structure is shown by deriving limiting dynamical equations for the trajectory of summary statistics of the SGD trajectory. The authors also study a specific case where means are orthogonal and show the one-pass SGD aligns with the largest outlier eigenvalue.\n\nThe authors further study the classification of the XOR problem on the Gaussian mixture model via a two-layer network. It is shown that the alignment between SGD and outlier eigenspaces is present in each layer. The results match the spectral phase transition in spiked covariance matrices.", "questions": "I would like to ask the following questions to the authors:\n\n- Is it possible to extend the theory based on cross-entropy loss to general function class?\n- The authors mention that SGD finds the subspace generated by the outlier eigenvalues for any uninformative initializations with norm $O(1)$ in Section 3.2, does the similar property hold for XOR-type mixture models via two-layer networks?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Provide theoretical understandings of the interaction between training dynamics of one-pass SGD and the outlier-bulk structure of Hessian matrix and G-matrix.\n- Detailed explanations help understand the theory and make the manuscript easy to follow.\n- Numerical experiments verify the theorems.", "weaknesses": "- Comparison and contribution over the prior works are not very clear.\n- Some technical part is hard to parse, it might be helpful to introduce the intuition for the main proof steps.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jbw7gYdFKl", "reviewer": "Reviewer_SJuj", "summary": "This manuscript investigates the alignement of one-pass SGD trajectory with the leading eigenvectors of the Hessian in the high-dimensional limit. In particular, it shows that for two different tasks:\n 1. multi-class logistic regression on a k-Gaussian mixture model;\n 2. classifying a XOR Gaussian mixture with a two-layer neural network   \n\nfor sufficiently high signal-to-noise ratio, the SGD trajectory of the model parameters approximately lies in the span of the top eigenvectors of the test Hessian after a number of steps which is linear in the input data dimension.", "questions": "- **[Q1]**: In Thm 3.1, how does the critical variance $\\lambda_{0}$ for SGD alignement compares with the BBP threshold for the recovery of the means $\\mu_{a}$ from $Y = y\\mu + Z_{\\lambda}$?\n\n- **[Q2]**: In both examples considered here, the timescale for the trajectory to correlate with the outlier subspace is linear in the dimension. How general should we expect this to be? For instance, would be authors expect this to hold in problems where escaping a fixed point would take longer than linear time in the dimension?\n\n- **[Q3]**: On the same spirit of the question above, would the authors expect the theorem to remain true at longer timescales? In general, should we expect these subspaces to be stable or the SGD trajectories to eventually escape these subspaces?\n\n- **[Q4]**: From the statement of the results, it seems they hold for both the G-matrix and the Hessian. Is there any reason for looking at the Hessian instead of the computationally simpler G-matrix? Are the subspaces spanned by the top eigenvectors of both matrices equivalent, e.g. in the sense of Def. 2.2?\n\n- **[Q5]**: Are the authors aware of converse examples to their result. For instance, a problem where the underlying target function does depend only on a few relevant directions, but that the outliers in the Hessian/G-matrix are not necessarily aligned with it?  \n\n**Minor comments**\n\n- The reference in the top of Page 2:\n> \"*Liao & Mahoney (2021) studied Hessians of some non-linear models beyond generalized linear models, also at initialization.*\"\n\nis misleading. The main result in Liao & Mahoney (2021) concerns the Hessian of a loss function of the type (see eq. 3):\n$$L(w)=\\frac{1}{n}\\sum\\limits_{i=1}^{n}\\ell(y_{i}, w^{\\top}x_{i})$$ under the assumption $y_{i}\\sim f(y|w_{\\star}^{\\top}x_{i})$ (see eq. 1). Despite their unusual choice of terminology for this model (\"G-GLM\"), this is mostly commonly known as a generalized linear model in the literature... when just not simply refereed to as \"linear model\".\n\n- The following references:\n> *ODE limits of the SGD training with single layer networks have been derived and numerically solved in Mignacco et al. (2020); Loureiro et al. (2021).*\n\nare not accurate. Indeed, Mignacco et al. (2020) and Loureiro et al. (2021) study the classification of Gaussian mixtures with a single layer networks. However, they derive exact expressions for the training and test errors of the minimiser of the empirical risk. Indeed, this should correspond to the asymptotic performance of one-pass SGD when $t\\to\\infty$ of the example considered in Section 3.\n\n- The figures in the manuscript are not very readable. First, even if this is described in the caption, it would be good to label the (x,y)-axis for clarity. Second, they are small, and since the format is not vectorial they get pixelated when zoomed over. Third, what does the colour scheme means? Maybe add a colour bar?\n\n- I understand the authors might not want to dwell 30 years of literature on SGD scaling limits. But I would encourage them to mention at least some of the recents works on this line which are contemporary to [Ben Arous et al. 2022] to represent the diversity of this literature, see e.g. the list below.\n\n**Typos**:\n\n- Page 7: \"*direcitons*\" -> \"*directions*\"\n\n**References**\n- **[Goldt et al. 2019]** Sebastian Goldt, Madhu Advani, Andrew M. Saxe, Florent Krzakala, Lenka Zdeborová. \"*Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup*\". Part of Advances in Neural Information Processing Systems 32 (NeurIPS 2019)\n\n- **[Chen et al. 2019]** Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. \"*Gradient descent with random initialization: fast global convergence for nonconvex phase retrieval*\". Mathematical Programming, 176(1):5–37, Jul 2019\n\n- **[Veiga et al. 2022]** Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, Lenka Zdeborová. *Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks*. Part of Advances in Neural Information Processing Systems 35 (NeurIPS 2022).\n\n- **[Paquette et al. 2022]** Courtney Paquette, Elliot Paquette, Ben Adlam, Jeffrey Pennington. \"*Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties*\". arXiv:2205.07069 [math.ST]\n\n- **[Arnaboldi et al. 2023]** Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro. *From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks*. Proceedings of Machine Learning Research vol 195:1–29, 2023.\n\n- **[Shuo & Vershynin 2023]** Yan Shuo Tan and Roman Vershynin. \"*Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval*\". Journal of Machine Learning Research, 24(58):1–47, 2023", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Characterising the geometry of SGD, and in particular finding the relevant subspaces where the low-dimensional subspaces where the high-dimensional trajectory lies is an important problem. In many theoretical setups, the relevant subspaces can be guessed from symmetries of the problem, and can be used to derive low-dimensional scaling limits for the projection of the trajectory in these subspaces - an approach that has been employed in many works old and recent works in the literature, c.f. [Saad & Solla 1995a,b; Goldt et al. 2019; Chen et al. 2019; Refinetti et al. 2021; Veiga et al 2022; Ben Arous et al. 2022; Paquette et al. 2022; Arnaboldi et al 2023; Shuo & Vershynin 2023]. However, an important limitation of this approach is that it is hard to generalise to more general setups where the relevant statistics are less clear to guess. This work proposes a roadmap to overcome this limitation, consisting of taking the spam of the top outlier directions in the Hessian/G-matrix, and shows this correlates with the standard summary statistics for these examples. If this proves more general, it could open up an important new theoretical tool in the study of SGD.", "weaknesses": "The main limitation is that the two classification tasks have a particular structure. Indeed, in both the linear case and the two-layer case considered here, the Hessian/G-matrix of the first layer weights is proportional to $Y\\otimes Y$. For a Gaussian mixture $Y = y\\mu + Z_{\\lambda}$, indeed the relevant subspace spanned by the means will naturally pop as an outlier of this matrix when the SNR $\\lambda$ is large enough: this is the classical BBP transition. Therefore, it is not surprising this also happens at the level of the Hessian/G-matrix. Maybe proving this requires a tour de force, but it is fair to doubt on the generality of the conclusion. For instance, it would be desirable to show a similar result for tasks where the structure is on the labels, and not on the inputs, e.g. teacher-student models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z7njO6i64u", "reviewer": "Reviewer_frpP", "summary": "This paper looks at the task of classification for high-dimensional gaussian mixtures using 1 or 2-layer neural networks. The authors show that the SGD trajectories rapidly aligns with the low-rank outlier eigenspaces of the Hessian and gradient matrices.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "There has been many previous papers that study the empirical Hessian during the training process of SGD. It has been observed that the hessian spectrum can often be separated into a bulk component depending on the network architecture, and outlier eigenvalues which depend on the data. For a simple Gaussian mixture model, this paper gives a rigorous proof of this phenomenon for SGD.\n\nOverall I think this is a strong theoretical result that characterizes the implicit bias of SGD and explains some of the previous empirical observations about the empirical Hessian.", "weaknesses": "Overall I found this to be a very solid work. I did not find any major weaknesses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "divdx2bXbK", "reviewer": "Reviewer_tvSQ", "summary": "Background: Empirical findings in [1] show that the iterates of SGD, used for training deep learning models for classification, converges to a neighborhood of a tiny subspace of the parameter space - in particular to the eigenspace of the top eigenvalues of the Hessian and the covariance matrix of gradients - which justifies why deep learning algorithms does not suffer the curse of dimensionality.\n\nSummary: This manuscript rigorously proves that this observation should hold in multi-class logistic regression and 2-layer neural networks trained on Gaussian mixture models. Moreover, they also show that this alignment between the SGD iterates and the top eigenspace of the Hessian holds layer-wise and does not depend on the model's success.\n\n[1] Gur-Ari, G., Roberts, D.A., & Dyer, E. (2018). Gradient Descent Happens in a Tiny Subspace. ArXiv, abs/1812.04754.", "questions": "* I checked the proofs corresponding to Theorems 3.1, 3.2, and 3.3 carefully.  There are some minor mistakes (which should be fixable):\n   - In Eq. (C.10), the second term in the RHS should not be correct.\n   - In Eq. (C.8), you missed the case b \\neq a = c.\n   - In the fifth equation on Page 26 in Appendix, (the one starts with \\langle \\nabla_c H, \\nable_a R_aa^{\\perp} \\rangle), I think you missed a term in the RHS.\n   - In the equation above (C.14) (the one starts with \\nabla_c H \\otimes  \\nabla_d H), I think many terms are missing in the RHS. Please check that part carefully.\n   - Because of the previous point, Eq. (C.14) should be incorrect, which eventually makes the correctness of the correction terms in Theorem 5.7 questionable. Can you check this part again as well?\n\n* I skimmed through the rest of the proofs. I could not see a major mistake.\n\n* As the last question for the authors: Your results crucially depend on the high-dimensional limit for SGD proven in [1]. In that work, it was shown that there is a critical learning rate that causes additional correction terms in the limit. Can you comment on how the existence of the correction terms affects your results?\n\n[1] Arous, G.B., Gheissari, R., & Jagannath, A. (2022). High-dimensional limit theorems for SGD: Effective dynamics and critical scaling. ArXiv, abs/2206.04030.", "limitations": "", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* It provides theoretical evidence for an important observation. \n\n* Although the previous studies considered the single and/or multi-index models [1,2] suggest that such alignments are expected, this work goes beyond them and proves more, such as the layer-wise alignment, and that alignment occurs even if the model fails to learn the ground truth.\n\n* Overall, it is a very good work! (However, there are some minor mistakes in the proof. Please see the Questions section)\n\n[1] Damian, A., Lee, J.D., & Soltanolkotabi, M. (2022). Neural Networks can Learn Representations with Gradient Descent. ArXiv, abs/2206.15144.\n[2] Mousavi-Hosseini, A., Park, S., Girotti, M., Mitliagkas, I., & Erdogdu, M.A. (2022). Neural Networks Efficiently Learn Low-Dimensional Representations with SGD. ArXiv, abs/2209.14863.", "weaknesses": "NA", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)