Imitating Deep Learning Dynamics via  Locally Elastic Stochastic Differential Equations
Authors: Jiayao Zhang
Hua Wang
Weijie J Su
Venue: NeurIPS 2021 Poster (2021)
Keywords: deep learning
stochastic differential equations
ordinary differential equations
local elasticity

Abstract: Understanding the training dynamics of deep learning models is perhaps a necessary step toward demystifying the effectiveness of these models. In particular, how do training data from different classes gradually become separable in their feature spaces when training neural networks using stochastic gradient descent? In this paper, we model the evolution of features during deep learning training using a set of stochastic differential equations (SDEs) that each corresponding to a training sample. As a crucial ingredient in our modeling strategy, each SDE contains a drift term that reflects the impact of backpropagation at an input on the features of all samples. Our main finding uncovers a sharp phase transition phenomenon regarding the intra-class impact: if the SDEs are locally elastic in the sense that the impact is more significant on samples from the same class as the input, the features of training data become linearly separable---meaning vanishing training loss; otherwise, the features are not separable, no matter how long the training time is. In the presence of local elasticity, moreover, an analysis of our SDEs shows the emergence of a simple geometric structure called neural collapse of the features. Taken together, our results shed light on the decisive role of local elasticity underlying the training dynamics of neural networks. We corroborate our theoretical analysis with experiments on a synthesized dataset of geometric shapes as well as on CIFAR-10.

Imitating Deep Learning Dynamics via Locally Elastic Stochastic Differential Equations
Jiayao Zhang
Hua Wang
Weijie J. Su
University Pennsylvania
{zjiayao, vanghua, (aaraassswwwwaaaaa upenn. edu
Abstract
Understanding the training dynamics of deep learning models perhaps necessary step toward demystifying the effectiveness of these models. In particular, how do data from different classes gradually become separable in their feature spaces when training neural networks using stochastic gradient descent? In this study, we model the evolution of features during deep learning training using a set of stochastic differential equations (SDEs) that each corresponds to training sample. As crucial ingredient in our modeling strategy, each SDE contains a drift term that reflects the impact of backpropagation an input the features of all samples. Our main finding uncovers sharp phase transition phenomenon regarding the intra-class impact: if the SDEs locally elastic [19] in the sense that the impact is more significant on samples from the same class as the input, the features of the training data become linearly separable, meaning vanishing training loss; otherwise, the features are not separable, regardless of how long the training time is. Moreover, in the presence of local elasticity, analysis of our SDEs shows that the emergence of simple geometric structure called the neural collapse of the features. Taken together, our results shed light on the decisive role of local elasticity in the training dynamics of neural networks. We corroborate our theoretical analysis with experiments on synthesized dataset of geometric shapes and CIFAR-10.
Introduction
Deep learning models have achieved significant empirical success over the past decade across a wide spectrum of domains spanning computer vision, natural language processing, and reinforcement learning [3 43 48 Despite these remarkable achievements at the empirical level, there is still much to learn about deep neural networks, as evidenced by the fact that almost all important advances concerning architecture design and optimization for deep learning are based on heuristics, without much input from theoretical perspective 220 21 27
An important step toward opening these black-box models and unveiling their formidable details is to quantitatively understand the impact of backpropagation in deep learning training. While there has been continued effort to demystify how simple optimization methods give rise to impressive generalization performance, for example, [49 26 4 this by no means an easy problem, perhaps because of the daunting nonconvex nature of neural networks. Accordingly, for near-term purposes, a more practical approach is take a phenomenological viewpoint by relating simple empirical patterns to the effectiveness of deep learning models.
In this spirit, we are interested in how data from different classes gradually become separable in their feature space by repetitively calling backpropagation. From a phenomenological viewpoint, this question can be addressed by first analyzing the impact of single update using a stochastic gradient on the performance of the neural networks. More precisely, imagine that the gradient is evaluated in an image of cat, how does the hidden representation of another image image-say, an image of another
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

1000
(a) GEOMNIST in RR
(b) GEOMNIST in IR.
(c) CIFAR in IR.
(d) CIFAR R.
Figure I: Separation of features (logits). GEOMNIST dataset ((a)-())) and CIFAR dataset (d) trained using = 3 classes on a variant the ALEXNET model. Separation in is done by projecting the logits to v E RR, which is set to be the difference between pair of class means X (T) X (T) for large T, where k, I are chosen heuristically. The dashed lines (b) and (d) are simulated paths from Equation (9) using estimated (t) and (t). More details are given Section and Appendix D.2
cat image of planc-evolver because of the backpropagation? Recent studies answer this question by introducing a phenomenon called local elasticity, which, roughly speaking, means that the impact is generally larger on similar sample (an image of another cat) than dissimilar sample (an image of a plane) [9]
Motivated by the phenomenon of local elasticity, we propose model that captures the interaction between different training samples during deep learning training using a set of stochastic differential equations (SDEs) that reflect local elasticity in neural networks. Characterizing the intra-class and inter-class effects is an essential component of our modeling strategy, each of which contains a drift term that imitates the impact of backpropagation
specific training data of all samples.
Our main finding uncovers sharp phase transition phenomenon regarding the intra-class and inter- class impact; if the SDEs are locally elastic in the sense that the impact is more significant on samples from the same class as the input (the intra-class effect strictly greater than the inter-class effect), the features of the training data are guaranteed to be linearly separable, meaning vanishing training loss; otherwise, the features are not separable, no matter how long the training time is. This result provides convincing theoretical evidence for the presence of local elasticity in deep learning (19). Our model is also quite accurate in simulating the feature dynamics of deep learning. As shown in Figure and detailed in Section 4. the dynamics of the simulated logits are reasonably close to the real dynamics of deep learning on both synthetic and real datasets, indicating a well-suited model for theoretical and practical purposes. Moreover, in the presence of local elasticity, our SDEs also predict the emergence of a simple geometric structure called neural collapse features 38
Taken together, our results shed light on the decisive role of local elasticity in the training dynamics of neural networks. We corroborate our theoretical analysis with experiments on synthetic dataset of geometric shapes, well CIFAR-10. The experimental evidence consistently supported our model, which provides new insights into the dynamics of deep learning training.
1.1 Related Work
Dynamics in Deep Neural Nets. Many properties of linear deep neural nets are relatively well understood, such the loss landscapes (25). trajectory-based convergence [2] [2], and implicit acceleration [3] Exact solutions the training dynamics can be obtained in certain initialization
schemes 440 41 30]. In the presence of non-linearity, various assumptions
generally made. [39 [7] studied the dynamics of shallow neural nets with non-linearity and the neural tangent kernel (NTK) literature (22 4 [4] linearizes the network function of an infinitely wide neural net at initialization, which is similar to that of the deep Gaussian process literature (9] 18 32 (a treatise comparing and contrasting them be found in [47). Although approximations, NTKs are generally used when studying optimization trajectories neural nets such as 34 which also appears implicitly in many works when studying the optimization trajectories of neural net training
45 11 24 7)..
SGD as SDEs Neural Nets. The study of dynamics or trajectories of weights in deep neural nets via SDEs relies on the more precise characterization of stochasticity (33 36). Built on top of

---

this formalism, 44 studied the trade-oft off between batch size and learning rate, [23] analyzed factors influencing the quality of local minima, [5] studied the behavior of the SGD near local minima, and [42] studied the effect of learning rates. Although SGD-SDE approximation requires an infinitesimal learning rate, 35 verified that the SDE approximation can meaningful in practical settings and obtained necessary conditions for the validation of such approximation.
Local Elasticity and Phenomenological Models. Local clasticity is proposed in [9] as a phc- nomenological approach to reasoning the behaviors of neural networks. This phenomenon has inspired several works on generalization bounds [00] and an improvement on the NTK (6.
2 Binary Separation via LE-SDE
2.1 Setup, Notations and Assumptions
Throughout the paper, we work with the following setup and assumptions. For ease of reading, vectors and matrices are written in boldface and we denote by [rx] the set {....., n). When there is no ambiguity, we will write both X (t) and X, for a continuous-time (possibly stochastic) process.
Classification Problem. Consider K-class classification problem with K 2 2, with each class having n training examples. We denote by a E R the i-th sample of the k-th class, and g E [K] its label, where E [n], k E [K]. A neural net a function f Rd + IR K that maps the samples to logits ( pre-activation of the softmax).
Feature Vectors. We denote by X (m) E IP a p-dimensional feature of the i-th sample in the k-th class learned by the neural net at iteration For example, it can be the logits the output of the second-to-lase layer. Assume that the initial values X are i..d. samples from some distribution for each i E [m]] E []. We use j E [n] as indices for an individual sample, k. I E [K] for classes, and capital letters J and L to indicate random samples from Unif and Unif respectively.
Training Dynamics. We model the training dynamics in neural nets under SGD with an emphasis on local elasticity. At the m-th iteration, the Jm -h sample is sampled from the Lm-th class, where Jm ~ Unif((n]) and Im ~ Unif([K). Training on z/... affects the features of another data sample z in the form of
(m) X 1) = h Lm (m) Jm (m 1) Vhet 1)
where i E [n], k E [K], h is the step size, and ç): (m) is the noise term that is modeled as Gaussian noise. The scalar F.R.T... measures the strength of local elasticity that z exerts on 2 at iteration m. We assume (0). c (m) are jointly independent.
Local Elasticity. Clearly, by writing (m,( assume that this effect depends only on the class I, k. and time m. We write the effect matrix as E(m) = (Ex,(m))) For ease of exposition, we assume E only consists of two values (m)) and ((m), with a(m) representing the intra-class effect and ((m) the inter-class effect. To this end, we assume the effective training assumption, that is. as training progresses, the features become more discriminative: features from the same class are more similar, whereas those from different classes are more distinct, measured by some similarity measure in the feature space. We also assume that the LE effect is "proportional" to the feature x/- (m) itself. We generalize this point in Section 3 by introducing transformation matrix H on the features.
2.2 Binary LE-SDE
Our construction of equation (1) emphasizes the effect of intra- and inter-class effects on the dynamics of features, and thus differs from the usual weight dynamics that is common in the literature. Before deriving the general form of our locally elastic SDE (LE-SDE), we shall familiarize the reader with our model by demonstrating this in the case of binary classification (K = 2) with -ieedmmenii features (p- 1) the output of the model to be fed into the softmax function, also called the logit.
Let the intra-class effect be Fn F.z2 = a, and the inter-class effect is F12 F221 = 33, both of which are iime-independent. Expanding equation (1), for 1 < i,j S. and m 0, when train the
3

---

model on the Jm -th training example from the Lm-th class, we have
(m) = - - 1) ax Lm (m 1mm11111 1) + vh c. (m - 1), (x](m = (m h- 3 Jm 1) vh. 5l (m - 1).
In the limit of h + O, we can show that X (m) approximates some continuous- time stochastic processes (1) (under the identification of ( = mmh) governed by the set of stochastic differential equations as follows:
(x) ( xxx(())
(t) + _ (10)))xx0+00) (x)=
dt dw) t2 0,
[K]
[n]
(2)
where (t) = (x) +.... (t)) /n, and W are independent standard Wiener processes. The detailed derivation is given in Appendix A.
Now averaging over for each is in equation (2), obtain the following set of two ordinary differential equations (ODFs) governing the per-class means that (1) for k= k== 1,
a x (t) + B x3 (t)) dt + od (t) +... + (t) 2
dx)((t)
Taking the limit of + we observe that ad ('+++++WW"(0)) O. Thus, the above display converges weakly the following ODE:
(t) a
= 2 (1)+ + 1 x))+
dt
With the initial conditions EX,((0 (0) the solution to the above ODE is
x
(1)
C2 3
c1 C2 a
X2
02 a
C1
+ 02
2
2
In the samite-sample setting, we may replace X1, X2 in the SDE (2) by their deterministic solutions and obtain
(t) 9 9-8 citere o+8t - + x (t) X3(1 = -- - اليي - س
e - 02 + +
We are now ready to derive the condition under which these 2n feature vectors become asymptotically separable, that is, min, X! (t) > max, (t) or max, (t) min,X3 (t as 20.
Theorem 2.1 (Separation in Binary Classification). Given the feature vectors (t),,. x² for i [n), as t + 0o and large n,
1. if a >> B, they are asymptotically separable with probability tending to one,
2. if S B, they are asymptotically separable with probability tending to zero.
This result indicates sharp phase transition when is just above 3, that is, in the regime of local elasticity. As long as the intra-class effect is slightly greater than the inter-class effect, separation is guaranteed. This simple model already captures local elasticity and reveals the important role it plays in the perfect separation of training samples. We can generalize this model to more realistic settings: when there multiple classes, when features high-dimensional, and when the LE matrix E is time-dependent. In the next section, we discuss each of these three generalizations in more depth.
3 General LE-SDE Model
Now, we consider the general case where K 2 2 and the feature vectors are p-dimensional with p 2 K. Inquisitive readers may have already noticed that Theorem 2.1 only asserts the emergence of the separation of features, while being inconclusive to their relative orders at separation, that is, which class converges to where? This drawback is intrinsic to the toy model as neither intra-class nor inter-class effect identifies different classes. In this section, we introduce the general LE-SDE model that alleviates this difficulty with the help of an extra block matrix H with the (i, block H,,; models how features in the j-th class affect those in the i-th class, which also partially defines how

---

classes are separated in higher dimensions. In the local elasticity formalism, H 1.j can be viewed as inducing metric on the feature space under which local clasticity manifests.
As hinted before, in the case of multiple-class features in higher dimensions, we want to guarantee a stronger separation: know which class converges to where, thus incorporating supervision from label information. For example, when the features are logits (outputs of the neural nets) and the model trained under the softmax cross-entropy loss, previous work suggests they separate according to specific geometric structures 338. To this end, we need to adjust the raw feature vectors X with a proper transformation that incorporates the label information into the dynamics. This motivates the following modification of the dynamics (1) by adding an extra transformation H H. RPPP the
features. For & E [K], E [n]. and at iteration m, we have the following:
(4) The H..L.. (m) term models the LE effect as proportional to a linear "transformation" of the features. The dynamics in Equation (1) are special cases when H (..(m)) I for all / [K]. By specifying proper H. we can overcome the limitation in our toy example of not knowing which class converges to where. We specify interesting choices of H in Section 3.2
(m) = x (m 1) + h L., (m) L. J Jm. (m 1) + Vhe 1).
further step of abstraction is write X (m) instead of X (m), to indicate one generic sample from the distribution (m)) of all the features of class i at iteration m. As in Section 2.2, can derive the continuous dynamics of equation (4) in the limit of h - in the same way as
equation (2). Similar to writing X = (X )i=1 E RKP for the concatenation of per-class features, X = (x*) jmm E RK is the concatenation of per-class mean features. Our model (4) approximates the following SDE with identification t = mh as h + O. We term this model the LE-SDE:
dx M,X, dl + E dw..
(5)
where W, is the standard Wiener process in K2, Eg is the covariance matrix, and M. E R Kpx Kp is a K x K block matrix, with each block of sizepxp. The (k:, /)th block of M, is (t) (,)(()) when * k;, and (t) (t)) K when = k. The rationale for dividing K is that we assume that the data are balanced; therefore, each of the K possible classes has an equal chance of being sampled, as proved in Equation (2), where K = 2. In Appendix E. we discuss how we can generalize this to model SGD with mini-batches, imbalanced data, and label corruptions.
Taking expectation with respect to the randomness arising from sampling X, from its distribution, the per-class mean X, satisfics the following system, which we term the LE-ODE:
X M,
(6)
Under the assumptions in Section 2.1, we define (t) min (a(t)- + +K (K))))))),,
A(t) = Jn (T, dr, B(t) = To B(7) dr, and I((t) = min [A(t) - ((t), + +K - 1))(tt)
3.1 The Separation Theorem
Similar to the discussions Theorem 2.1, the LE-SDE allows us to derive the separability result for general K and p 2 K. We say the feature vectors (x (x())) KE[K] are separable if for any
two classes 7 L, there exists a hyperplane in R2P that linearly separates the features of the two classes. To characterize the separation in Theorem 2.1 need conditions on a(t), B(t) therein. Intuitively, when 7(t) < O, the classes cannot be separated, even in a pairwise manner. Therefore, we focus a more interesting case when (t) >> O. We now state the following characterization theorem of separability for general LE-SDE dynamics:
Theorem 3.1 (Separation of LE-SDE). Under our working assumptions in Section 2.1, and in the case of local elasticity (i.e., 7(t) > 0), assume H = (Hi)) is positive semi-definite (PSD) with positive diagonal entries. As + oc, we have
1. if (t) = w (1/t), the features are separable with probability tending to 1;
'Here, >(t) = w (1/t) stands for (t) >> 1/t as t + co. For example, 1/t0.5 = w (1/t) and (t = (1/t) + co.
5

---

2. (t) = (1/t), and the number of per-class-featuer n. tending to oc at an arbitrarily slow rate, the features are asymptotically pairwise separable with probability 0.
This theorem sheds light on the crucial impact of the local elasticity effect for separation in general case. The proof to Theorem 3.1 well discussions the empirically best ways of choosing the universal direction v (i.e., direction that does not depend on the class index) are detailed in Appendix C.2.
3.2
Two Specific Models
We next discuss two specific choices of the H matrix that allows us to analyze X precisely.
3.2.1 Isotropic Feature Learning Model
As a straightforward extension to Section 2.2, we can simply choose H = I to be the identity matrix. This choice of H is PSD, and thus, we can apply Theorem 3.1 to obtain the conditions for asymptotic separation. In this case, the solution (t) the LE-ODE can be computed analytically as given in the following proposition.
Proposition 3.2 (I-model). Let k,t Ip then the solution to the LE-ODE (6) is given by
(t) = ce. A(t) k B(t) + (Ix @ co) ek A(t)+ K-1 B(t)
where c = (c)) E Kp and co E IRP are constants with Σ Cx = 0 E and x()) = c. co-
The derivation of Equation (7) is deferred to Appendix C.2.1. From equation (7), we can easily reconstruct Theorem 3.1 in this special case. The difference between a feature vector from class and that from class I time is given by (Ce- - ci j A(1) B(r) + // (w W provided that the first deterministic term dominates the random second term, thus ensuring separation, which are precisely the conditions specified in Theorem 3.1. We term this model as the isotropic feuture model, or I-model for short; as the H matrix has identity matrices as its blocks and consequently the dynamics do not prescribe any preferred directions for each class.
3.2.2 Logits-as- Logitu-as-Features Model
An important type of features in neural nets is the logits, the outputs of the neural net before the
softmax layer. A logit vector logits) is K- di-imensional, and in this model we identify X (t) as the logits at time f of a generic sample from the k-th class. In a well-trained neural net, the logits of learned data instance from the k-th class should have its k-th logit being the largest, and heuristically, the other coordinates should be approximately equal and negative, As we shall
detail in Appendix B, the exact dynamics of neural net training pushes the logits X by its margin,
d, = softmax X") ) which roughly aligns with the direction of e₁: 1p/ K. This suggests us how to choose the metric under which local elasticity acts: we can choose H such that it always
aligns X in the direction of dg. that is,
H = H = d,d.
E RRXP, ||,
1 1 ER, JE [ = - K
(8)
Roughly speaking, the map x + H a projects æ in the direction of d, and ideally aligns æ with d, after iterative applications; hence, H can be viewed an approximation of the nonlinear transformation in the exact dynamics in the sense that the direction of their stationary point coincides. Furthermore, H' thus defined has operator norm 1; thus, it does not affect the magnitudes, but only directions. Note that H does not satisfy the condition in Theorem 3.1 as it is not symmetric; yet the separation theorem can be easily extended in light of the following proposition.
Proposition 3.3 (L-model), Let H be the same as equation (8), then the solution to the LE-ODE (6) is given by
K-1 Σ
X(1) co Cy A(1) +
C2ifi
A XTR B(t)
where fi's are fixed vectors in KR Cn E RR* constant vector with K(K 1) degrees of freedom, and C1, Cz E I. for I E [K- 1] constants.

---

Figure 2:
Samples from GEOMNIST dataset.
The specific form of fr is not the focus here; equation (9) allows us to prove the statement of Theorem 3.1 under this choice of H. The proof of Proposition 3.3 is deferred to Appendix C.2.2. where we also provide the analytical solution of fi² when K =
We term this model as the logits-as-features model, or L-model for short, because it is an elaborate model specifically for logits. In Section 4, we provide concrete demonstrations of our abstract feature vector X as logits under L-model. We numerically simulated our LE-ODE and compared its predicted dynamics with real deep learning training dynamics. The experimental results provide strong empirical support for the validity of L-model.
3.3 Connection with Neural Collapse
Neural collapse is recent phenomenological finding the geometry of the logits learned by deep neural nets at convergence with the cross-entrops loss [38] (see an explanation of neural collapse in [16] Simply speaking, taking our L-model as example, with balanced training samples, this model asserts that the logit vectors from different classes at convergence form an equiangular tight frame (ETF). ETFs are the best configuration to spread K unit vectors in ambient space of p dimensions. Formally, we say a set of vectors { form an ETF in Rp' if they are the columns of matrix
K S V K- Q
1 K k)
IK
(10)
where Q E R and QTQ = Ig. As a direct corollary of Proposition 3.3, when set according equation (8), find that our L-model also predicts the existence of neural collapse from the local elasticity point view.
Proposition 3.4 (Neural Collapse of the LE-ODE). Under L-model and the same setup as in Theorem 3.1 if (t) > and there exists some T > O such that B(t) < 0 for t 2 T, then
K
x (t)/1lx* (t)l
forms an ETF oc.
4
Experiments
We perform various experiments to test our theory, where we choose logits as our protagonist
4.1
Setup
Datasets and Models. We perform experiments on a synthesized dataset called GEOMNIST containing K types of geometric shapes (RECTANGLE, ELLIPSOID, and TRIANGLE) and on CIFAR-10 ((8)) denoted by CIFAR) with K E [2, 3] classes. A few samples from GEOMNIST are shown in Figure We vary the number of training samples per class and label pollution ratio Per and use variants of the ALEXNET ((9) model. More details can be found in the Appendix.
Training Configurations. All models are trained for T 10th iterations (for GEOMNIST) = 3 x 10S iterations (for CIFAR) with a learning rate of 0.005 and batch size of under the softmax cross-entropy loss. Models on GROMNIST converged with training and validation losses to zero, and those on CIFAR to validation accuracies greater than 90%.
Estimation Procedures. Each experiment is repeated for Periial = 100 independent runs to estimate (1). We use both the isotropic feature learning model (Section 3.2. and the logits-as-feature model (Section 3.2.2), denoted by L-model and I-model respectively, estimate ar( and (()). The
"Cod for reproducing experiments is publicly available atggithub. som:zziaayo/ee_sde git
7

---

Ate)
1.30
30
A13 Rig
A3) BIO
Axx BiO
(a) GEOMNIST (l-model). GEOMNIST (l-model).
(c) CIFAR (l-model).
(d) CIFAR (l-model).
210
200
A10
Atc. Rvg)
AtC, A(0)
A30
Adt
Bi0
Bi0
(e) GEOMNIST (L-model).() GEOMNIST (L-model). (g) CIFAR (L-model).
(h) CIFAR (L-model).
Estimated A(t), B(t), att), and B(t). The first row was estimated using I-model and the second L-model; the first two columns are on GEOMNIST and the last two on CIFAR. The first and third rows show A(t) and Btt) and the other two rows @(t) and 3(t).
Figure 3:
L-model is used only when K = 3. To estimate att) and ((t), we first estimate A(t) and B(t) by
A(t) = avg avgg log
x-111
(I-model)
coe x
K-
X, = avg. X
B(() = avg avgg log
Co
X
(1I)
[t(t) A'(+) 2tB B(() 2(B((() A'()),
A'(t)) = log
X v - 1
(L-model)
B'(1) = log
X
(v2 - (0,)
where vector division is interpreted entry-wise. We write avg, for averaging over the class index, avg for averaging over the coordinates, and define (t)(t) = (t)/x(() We explain how and why to choose the vectors U1 and V2 in Appendix D.1. The main idea is to view the eigenvectors of the K p-by-Kp drift matrix as a concatenation of K vectors of dimension p and construct their linear combinations such that one more independent components in the solution vanishes. With A(t) and B() estimated, we use the Savitzky Golay filter to obtain @()) and 3(1) through numerical differentiation. For GEOMNIST and CIFAR datasets, we choose window sizes of this filter as 191 and 551, respectively, Figure 3, and 21 and 21, respectively, in Figure 5.
We assess the tail of @t) and B(t) by a tail index defined as ra == =sup, {s: : lim a(t) - oo) and ry is defined similarly. We estimate F. by fixing interval [2] with Ti T2 sufficiently log a(t) and similarly for Pg.. We use the estimates from the last 1000 iterations for averaging in our experiments.
large such that we may ignore terms with smaller order and have F. 1 avgn, SiSTa log(1+t))
4.2 Results
Local Elasticity in Neural Net Training. Local elasticity manifests from our model as the heav- iness of the tail of 7(1) = a()) 3(), and in Figure 3. we plot the estimations A()), B()) a(t), and B(t) using both I-model and L-model. We note that (i) The estimations from the two models are visually similar, especially in the late stage of training when f is large; (ii) The major difference lies in the initial stage, where the estimates from L-model behave slightly wilder. This is not surprising because of the effect of the unknown constant offset co in the L-model; (iii) Both @(t) and B(1) behave similarly on both datasets.
Phase Transition Separability. Theorem 3.1 states that separation of features under the LE-SDE takes place when r(t) = a(t)= B(t) w (1/t), or roughly speaking, when r. minfra, re} < 1. Although we cannot directly control (1) and ((),, we can bias them by tuning the label corruption ratio Perr- When Perr ~ Pérr = 2/3, we are in effect assigning labels completely at random and par.- This is indeed the story depicted in Figure 4 Figures 4a and 4b show that the validation loss and accuracy for p 2 Per
thus we expect a phase transition of separability should happen around
8

---

Betioaared Taal Index
Loas
0.2
46 Per (% cormuttd label)
(a) Validation loss versus Perr- (b) Validation accuracy versus Perr-
(c) Tail index versus Perr.
Figure Phase transition of separability. (a)-(b) Validation loss and accuracy suggest separation
fails for Perr 2 Perr 2/3. The dashed line in (a) carries the value at initialization and overlaps with
the case where Perr = 0.6; the dashed line (b) is Parr = 2/3, when labels are assigned completely at random. (c) Tail indices of (tt) and (t). Note that ((t)
att) B(t) crosses the horizontal line r = 1, entering the non-separable regime (shaded in red) from the separable regime (shaded in
green), around the same perr that cross the dashed lines in (a) and (b), as predicted by Theorem 3.1.
(a) GEOMNIST (l-model).bb) GEOMNIST (L-model). (c) CIFAR (l-model).
(d) CIFAR (L-model).
Figure 5: Simulated LE-ODE solutions versus genuine dynamics. We use (t) and 3(t) estimated from I-model ((a) and (c)) or L-model, ((b) and (d)) and numerically simulate the solution under the L-model. The results were overlaid with true dynamics from neural nets. We note L-model in general imitated true dynamies reasonably well.
are not increasing over time and in Figure 4c we observe the minimum tail index of (<)) and ((1) crosses from below around perr 0.6, entering the non-separable regime (shaded in red) from the separable regime (shaded in green). given in Theorem 3.1.
Simulating DNN Dynamics via LE-ODE. Having estimated (t)) and B()), it natural to ask, to what capacity our LE-ODE models recover the real dynamics of deep neural nets? We use the forward Euler method simulate the L-model using @(t) and ((t) estimated from either I-model or L-model, We choose K = 3 and show in Figure 5 the simulated solution (solid line) with error bars depicting one standard deviation over 500 independent runs, overlaying on the real dynamics from DNNs in the background (shaded transparent markers). As the moving average may reduce the magnitudes of (t) and (()), we rescale the simulated paths such that first coordinate is approximately equal to the ground truth at convergence. Note that estimations from L-model can faithfully recover the genuine dynamics from neural nets, whereas those from I-model fail, notably in Figure 5a, where the simulated paths preserve the relative magnitude but fail to identify the correct order three logits.
5 Discussion and Future Works
In this study, we introduce I. LE-SDE/ODE models that draw inspiration from the local elasticity phenomenon. Conditions for sharp phase transition of separability of features are derived. We also show that once the clasticity strengths a(t) and 3(t) are well estimated, our model can faithfully simulate the dynamics of neural nets. We outline few interesting problems for future research while leaving the details in the Appendix. (i) General LE Matrix. A similar result as in Theorem 3.1 may be expected for symmetric but no necessarily semi memi-definte LE matrices E(t). (ii) Mini- batch Training, Imbalanced Datasets, and Label Corruptions. Generalizing the drift matrix to M. = EE X P) H K for K-by-K doubly stochastic matrix can used to model various sampling effects. (iii) Beyond L-model for Imitating Genuine Dynamics of DNNs. Although the L-model shown to able to mimic the real dynamics reasonably well, we postulate that a more precise model might have its (i, -)tt block encode the other directions other than d,.

---

Acknowledgements
This work was supported in part by NSF through CCF-1934876, an Alfred Sloan Research Fellowship. the Wharton Dean's is Research Fund, and ONR Contract Dan Roth and the Cognitive Computation Group at the University of Pennsylvania for stimulating discussions and for providing computational resources.
We would like to thank
References
[1] Z. Allen-Zhu, Y. Li, and Z. Song. A convergence theory for deep learning via over- over-vveeriiiee prrerarramtttiiaaioonn.. In International Conference on Machine Learning, pages 242-252, 2019.
[2] S. Arora, N. Cohen, N. Golowich, and W. Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2019.
[3] S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameteriiation. In International Conference on Machine Leurning, pages 244-253, 2018.
[4] S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grainee analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference Machine Learning, pages 322-332, 2019.
[5] P. Chaudhari and S. Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10. IEEE, 2018.
[6] S. Chen, H. He, and W. Su. Lahel-aware neural tangent kernel: Toward better generalization and local clasticity. In Advances in Neural Information Processing Systems, volume 33, pages 158477-1558,, 2020.
[7] Z. Chen, Y. Cao, Q. Gu, and T. Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. Advances in Neural Information Processing Systems, 33, 2020.
[8] J. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021.
[9] A. Damianou and N. D. Lawrence. Deep Gaussian processes. In Artificial intelligence and statistics, pages 407-215- 2013.
[10] Z. Deng, H. He, and W. Su. Toward better generalization bounds with locally clastic stability. In Pruceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2590-2600, 2021.
[11] Devlin, M. Chang, Lee, and K. Toutanova. BERT: pre-training of deep bidirectional transform- ers for language understanding. Conference of the North American Chapter of the Association for Computational Linguistics, pages 4171-4186. Association for Computational I Linguistics, 2019.
[12] S. Du and W. Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pages 1655-1664, 2019.
[13] S. Du, J. Lee, H. Li, L.. Wang, and X. Zhai. Gradient descent finds glohal minima of deep neural networks. In International Conference Machine Learning, pages 1675-1685, 2019.
[14] S. S. Du, K. Hou, R. R. Salakhuudinov, B. Poczos, R. Wang, and K. Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kemels. In Advances in Neural Information Processing Systems, volume 32, 2019.
[15] S. Du, X. Zhai, B. Póczos, and Singh. Gradient descent provably optimizes over-parameterized neural networks. In 7th International Conference on Learning Representations. OpenReview.nec, 2019.
[16] C. Fang, H. He, Q. Long, and W. Su. Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 2021.
[17] S. Goldd, M. Advani, A. M. Saxe, F. Krzakala, and L. Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In Advances in Neural Information Processing Systems, volume 32, 2019.
[18] T. Hazan and T. Jaakkola. Steps toward deep kernel methods from infinite neural networks. arXiv preprint 1508 05133, 2015.
10

---

[19] H. He and W. Su. The local elasticity of neural networks. In International Conference on Learning Representations, 2020.
[20] K. He, X. Zhang. S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 770- 0700, 2016.
[21] S. loffe C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 448-456, 2015.
[22] A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, volume 31, 2018.
[23] S. Jastrzebski, Z. Kenton, D. Arpit, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey, Three factors influeneing minima in SGD. arXiv preprint arXiv:/711.04623, 2017.
[24] Z. Ji and M. Telgarsky. Polylogarithmic width suffices for gradient descent achieve arbitrarily small test error with shallow relu networks. International Conference on Learning Representations, 2020.
[25] K. Kawaguchi. Deep learning without poor local minima. In Advances in Neural Information Processing Systems, volume 29, 2016.
[26] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2017.
[27] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Reprexentations, 2015.
[28] A. Krizhevsky. Learning multiple layers of features from tiny images, 2009.
[29] A. Krizhevsky, Sutskever, and E. Hinton. Imagenet classification with deep convolutional neural networks. Advances in Neural Information Processing Systems, 25:1097-1105, 2012.
[30] A. K. Lampinen and S. Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. In International Conference on Learning Representations, 2019.
[31] Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. Nature, 521(755)):4364444 2015.
[32] J. Lee, Y. Bahri, R. Novak. S. S. Schoenholz, J. Pennington, and J. Sohl- Sohl-Dickstein. Deep neural networks gaussian processes. arXiv preprint arXiv:1711.00165 2017.
[33] Q Li, C. Tai, and F. Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference Machine Learning, pages 2101-2110, 2017.
[34] Y. Li and Y. Liang. Learning overparameterrized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems, volume 31, 2018.
[35] 7. Li, Malladi, and S. Arora. the validity of modeling sgd with stochastic differential equations (sdes). arXiv preprint 102: 12470, 2021.
[36] S. Mandt, M. D. Hoffman, and D. M. Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, (8(1):487344907, 2017.
[37] A. W. Marshall, I. Olkin, and B. C. Arnold. Inqqualities: Theory of Majorization and its Applications, volume 143. Springer, second edition, 2011.
[38] V. Papyan, X. Y. Han, and D. L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40))2465-22463, 2020.
[39] D. Saad and S. Solla. Dynamics on-line gradient descent learning for multilayer neural networks. Advances in Neural Information Processing Systems, pages 302-308, 1996.
[40] A. M. Saxe. J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep lincar neural networks. In International Conference on Learning Representations, 2014.
[41] A. M. Saxe, J. L. McClelland, and S. Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings the National Academy of Sciences, 116/23):11537 -1111111157766666611111111111111 2019.
11

---

[42] B. Shi, W. J. Su, and M. I. Jordan. On learning rates and Schrödinger operators. arXiv preprint arXiv:2004.08 06977, 2020.
[43] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schritwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, al. Mastering the game of go with deep neural networks and tree search. Bature, 529(7587):484 489, 2016.
[44] S. L. Smith, P.J. Kindermans, and Q. V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018
[45] D. Soudry, E. Hoffer, M. S. Nacson, S. Gunasekar, and N. Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1)28822-23 28782828288, 2018.
[46] S. Särkkä and Solin. Applied Stochastic Differential Equations. Institute of Mathematical Statistics Textbooks, Cambridge University Press, 2019.
[47] G. Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint wVXiv:/902.047000 2019.
[48] T. Young, D. Hazarika, S. Poria, and F. Cambria. Recent trends in deep learning based natural language processing, IEEE Computational Intelligence Magazin, 13(3):55-75, 2018
[49] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization, In International Conference on Learning Reprexentations, 2017.
12


Reviews:
[{"review_id": "kXbNBc_N8k3", "reviewer": "Reviewer_imQZ", "summary": "The paper proposes a proxy tractable dynamics to study how the latent features of neural networks evolve inter- and intra-class during training. The core of this dynamics is a linear ODE, which is shown to exhibit inter-class separability once a “locally elastic” condition is met. The paper then examines this dynamics for two choices of the H matrix (which is a parameter of the dynamics), and concludes with experiments to demonstrate the relevance of the proposed theory with the actual simulation.\n\n\n=====================\nAfter rebuttal:\nThanks for the detailed reply!\n\nOverall while I am pleased that the authors made a lot of efforts to clarify, I'm not quite convinced. Some notes:\n- I agree that time-dependence of E(t) is a good point; my concern about linearization (which I believe is also the concern of some other reviewers) is actually about why at each time t, F(X(t),t), in the authors' response, is modeled by a linear relationship with X(t). Why ignoring higher-order terms? Are they much smaller than the linear term? I do not see why the higher-order terms can be absorbed into the noise term.\n- I do not agree with the argument that $c_1\\neq c_2$ assumption can be alleviated by taking the results \"in probability\". Clearly the derivation in the paper requires taking a law-of-large-numbers average.\n\nI agree (and well understood) that the paper does not aim to mimic exactly the dynamics; the aim is to provide a reasonable proxy dynamics that can explain interesting properties. However I feel the paper needs a careful justification to convince readers that this proxy is not oversimplifying.\n\nGiven that the paper may attract interests, I'm willing to raise the score, but not fully convinced of the quality to recommend acceptance. If a reviewer is willing to champion the paper, accepting the paper would not be a bad outcome to me.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "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": "", "main_review": "The subject of the paper is timely and the result, if convincingly shown to be a good proxy to study neural net dynamics, would be very interesting. However I do not find the proposal to have good potential (or have been sufficiently demonstrated to have good potential), and in fact, I suspect the proxy dynamics could be over-simplifying. More comments follow.\n\nThe paper models the evolution of the features by a linear class-based approximation, plus some noise. Firstly the paper does not give an intuition or an argument why the relation should be linear. I think this is an important modeling assumption: linear approximation is typically good only if the effect size is small. While it is intuitive near the end of training when changes are small, it is unclear why it should hold throughout all training time, especially at time closer to initialization when a lot of the training takes place. This shortcoming could be critical, since the paper models the whole evolution, and so any mis-modeling around initialization could lead to bad results eventually by the end of training time.\n\nSecondly the paper does not justify why the influence onto one feature could be simplified per class (by restricting the H matrix to be dependent only on the classes). Similar to above, this may not hold closer to initialization when different classes could occupy well-overlapped regions in the feature space. The only piece of argument I can find from the paper is the expectation of locally elastic behavior, but it is important to note that the notion of (class-based) local elasticity in this paper is different from papers by W. Su an co-authors.\n\nThirdly — a point which is mainly to point out the modeling downside, but perhaps not as critical as previous points — one can see that the stochastic part in SDE modeling does not matter at all. Rather it is the linear ODE, resulted from the averaging over samples within the same class, that dictates the picture. The fact that it is a linear ODE makes it straightforward to obtain separability. However it may seem to be oversimplifying neural nets, which are known to be highly nonlinear.\n\nOne specific comment I have on the theorems (in particular, Theorem 2.1, which is easier to discuss) is that separability cannot hold if $c_1 = c_2$. Recall that $c_1$ is the average of the features in class 1 at initialization, and the same for $c_2$. With the typical initialization in neural nets, I doubt this holds. For example, if the feature $X = Wx$ (linear model) where W is the weight matrix, with $W$ being standard zero-mean Gaussian, wouldn’t the average be zero and hence $c_1 = c_2$? Also note that the assumption $c_1\\neq c_2$ does not make sense at all: an untrained unbiased network should not be able to do any classification yet.\n\nAnother concern is on the guesses on H (the I-model and L-model). It is clear that I-model is insufficient. While it seems L-model is good as per Fig. 5, this guess is based on a well-trained net, and so it is unclear how well it models the behavior near initialization. This is not clear from Fig. 5, since already around t=3000 (where it is not easy to spot any deviation from Fig. 5), the validation accuracy is 90%, i.e. training is close to being finished; in other words, the good match after t>3000 is an easy test for the theory. In fact, one can observe quite big deviation in Fig D.8. This raises a doubt on the proposed guesses of H.\n\nFig 4(c) is also unconvincing. Firstly we know that I-model is not a good choice, so this should have been done for L-model instead. But when looking at one for L-model (Fig. D.2(b)), one sees a mixed picture. This figure seems to suggest the non-separable regime should start around p=0.5, but this is at odds with firstly the correct threshold p=0.66 and secondly Fig. 4(b) where the validation accuracy for p=0.5 is significantly better than random guessing.\n\nSome other comments:\n- Corollary 3.4 does not seem to be correct: for example, if B(t) = -0.1 and A(t) = t, then the term ${\\bf d}e^{A(t)/K - B(t)/K}$ is not dominating the other term.\n- The proofs are not fully rigorous. For instance, to deduce the equations after line 133 on page 4, one has to do careful probabilistic argument to prove it holds at all t, as n tends to infinity.\n- In Fig. (4) caption: “crosses, 1 entering…” is not readable.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "azoxhcmhl0z", "reviewer": "Reviewer_J1us", "summary": "This paper proposed a stochastic differential equation model based on the local elasticity assumption of features of neural networks.  The local elasticity phenomenon says samples have greater influence to other samples in the same class than those in different classes. Starting from this assumption, a system of linear SDEs is proposed for features corresponding to each data point, with the coefficients encoding the difference of inter and intra classes impact. Two types of features are studied---the isotropic feature learning model, and the logics-as-features model. For the two models, dynamics of the class means are derived as a system of ODEs and the solutions are found, depending on the coefficients about local elasticity. From the solutions, the separability of classes is theoretically connected with local elasticity. Specifically, the classes are separable if the local elasticity effect is nonzero, while not separable otherwise. The transition between the two cases is sharp. \n\nNumerical experiments are conducted to justify the effectiveness of the proposed model to characterize the actual dynamics of features. Experiments are done on deep convolutional networks on the GeoMNIST and CIFAR10 dataset. Firstly, the coefficients in the SDEs (and the ODEs for the class means) are estimated using real feature dynamics given by the training process (I understand this step as a kind of \"fitting\"). Then, the ODEs with estimated coefficients are simulated, and the results are compared with the real feature dynamics. It shows that in many cases (but not all) the trajectories produced by the SDEs recover the real trajectories given by the training. ", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "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": "", "main_review": "This paper provides a new way to understand the training process of deep neural networks. Specifically, the dynamics of features, rather than the parameters, is considered. Due to the stochastic nature of the training, the dynamics of features is formulated as SDEs. The work is not purely theoretical, since its modeling and analysis is built on the local elasticity assumption, which is observed in practice but not proven theoretically. However, considering the local elasticity is a special property for nonlinear machine learning models, such as deep neural network, and does not appear for linear models, the LE-SDEs still provide some insights to the interaction between different data during the training process of neural networks. Numerical experiments also show that the class means given by LE-SDEs with properly calculated coefficients match the real class means well. \n\nMy major concern about the work lies on the linear formulation of the LE-SDE. By the formulation, the influence of one data to other data is linear. (this influence actually happens on features, of course). However, no explanation or justification is provided for this linear formulation. Does it come from the stochastic gradient descent update, or come from the choices of features? Will the linear formulation still work if other optimization algorithms are used, or other quantities in the neural network are picked as features? The authors should discuss the rationality and potential limitation of the linear formulation.\n\nThe paper is generally well written and easy to read. Though, there are still some typos and seemingly wrong expressions (perhaps I'm wrong). To mention one, in the equations between line 121 and 122, \\alpha and \\beta are fixed in the expressions, regardless of the classes L_m. I believe this is not the right formulation. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have address the limitations of the work. Since this is a fundamental study not directly related to any applications, I do not think societal impact is necessary. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}, {"review_id": "6PwELcNBrLj", "reviewer": "Reviewer_qrk4", "summary": "The paper provides a simple phenomenological model of neural networks for K-class classification. Each sample is assigned a feature modeling NN features, and these features are evolved via SDEs in which the drift term of a sample is a linear transformation of the average features of each class. Separation, which means the case in which samples from different classes can be separated by hyperplanes, holds w.h.p. when the intra-class drift coefficients larger than the inter-class coefficients in a certain sense, and it does not hold otherwise. Two specific choices of features are considered: isotropic features and logit features. The latter are shown to be better than the former at modeling the feature evolution for NN on CIFAR and a dataset of their own creation.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "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": "", "main_review": "The paper is well written and clear. Its greatest contribution is the finding that such a simple model is able to describe the feature evolution of neural networks, as shown in experiments. This work shows, with its limitations, that the phenomenological approach can in some settings produce results that closely resemble the real NN. It opens up several questions that I hope can be answered in future works: When is this model valid? Can we somehow infer the drift terms from the neural network weights, and then design architectures that favor separation?\n\nWhen presenting the experiments, the plots for $\\hat{\\alpha}$ and $\\hat{\\beta}$ are very noisy and it is unclear what the decay of this quantities is.  If possible, it would be interesting to denoise the plots and show them in a log-log scale, to see whether indeed the separation results claimed in Theorem 3.1 translate to real NN.\n\nTypos: The paper needs to be proofread for typos. I point out some. In line 133 $c_x, c_y$ should be $c_1, c_2$. In line 184, “semi-positive definite” should be “positive semi-definite”.  In line 188, I believe that the sentence should be “the features are asymptotically pairwise separable with probability 0,” as the current sentence means that the features are asymptotically pairwise separable with probability 1.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper has no potential negative societal impact. The limitations of the paper are clear from the beginning: the paper follows a phenomenological approach, which precludes a deep understanding of the factors at play but might be beneficial in terms of the simplicity and predictive power of the model. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}, {"review_id": "0g2HB2RmLkR", "reviewer": "Reviewer_zWV1", "summary": "This paper builds off a recently proposed phenomenon of deep learning models, local elasticity, which, says the that impact in feature space of a gradient update from an input is in general larger for feature of data in the same class as the input rather than a different class. In this work the authors propose a set of SDE's modeling SGD that captures the inter and intra class effect of back propagation. They find that there is a sharp phase transition in the dynamics governed by these SDEs when they do or don't display local elasticity. In particular, when the SDEs are locally elastic then the features of the training data are guaranteed to become linearity separable, while this is not the case otherwise. They demonstrate this theoretical analysis empirically with CIFAR-10 and connect there results to another recently proposed phenomenon, neural collapse. ", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "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": "", "main_review": "**Originality/Significance:**\nThe authors seek to understand how neural networks are able to achieve linear separability explaining the recent empirical phenomenon of neural collapse.  They conjecture that this is caused by the network possessing local elasticity.  They construct a SDE in feature space that under certain assumptions will provably lead to feature separability when local elasticity is met.  My greatest concern with this work was that it was not clear why the SDE they use should reflect the actual feature dynamics of a real neural network.  This rational was not made clear theoretically and as far as I could understand the empirical explorations required training a DNN to estimate parameters of their model, which was then simulated, and then compared back to the empirical data used in the estimation.  Even accepting the validity of this model to track the feature dynamics it was not clear how this understanding should inform future work or answer questions like why is depth needed or non-linearities to achieve feature separability.  Despite these comments I thought the direct modeling of feature dynamics rather than just parameter dynamics was interesting and could be a promising direction for future work.\n\n**Clarity/Quality:**\nThis work has grammatical and spelling mistakes throughout (some listed below) and should be more carefully edited. The overall structure was fine and I appreciated the authors use of a simple setting for introducing the theory before diving into the more general expressions. I found the empirical setup (in particular the estimation procedure) very difficult to follow. The authors did not spend enough time focusing on assumptions applied in their theory or limitations of their empirical approach to \"validating\" the theory. Finally, future work and related work should be expanded, with particular emphasis on how this model could inform algorithmic improvement or answer open questions.\n\n\n**Major Areas for Improvement:**\n\n - Equation (1) is the basis of your analysis, but you made very little effort to build it up or discuss the assumptions applied.  For example, why should we expect that the dynamics of the features are driven by a linear drift term? How should I think of equation (1) being derived or implied from the SGD parameter updates? All the related work you discuss in \"SGD as SDEs\" is about SDEs in parameter space, your SDE derived from eq. 1 is operating in feature space. This is a significant difference and should be discussed!\n - Discussion on \"Local Elasticity and Phenomenological Models\" could be much more detailed.  This is clearly the closest related work section and it would be good to go into detail about their claims and the implications of follow up works.\n- Is $\\alpha(m)$ and $\\beta(m)$ implicitly assumed to be non-negative? For example, is $\\alpha,\\beta \\ge 0$ in theorem 2.1?\n- What is the relationship between equation (2) and an Ornstein-Uhlenbeck process?  Imagine the setting of $n=1$ and let $Z(t) = [X^1(t), X^2(t)]$, then doesn't equation (2) describe the OU process $dZ(t) = AZ(t) + \\sigma dW(t)$ where $A = \\frac{1}{2}\\begin{bmatrix}\\alpha &\\beta\\newline \\beta & \\alpha \\end{bmatrix}$?  If this is the case then how does theorem 2.1 relate to the eigenvalues of $A$?  I think this perspective/connection to a classical SDE would be helpful for understanding.\n- In equation (5) the potentially anisotropic diffusion $\\Sigma_t$ kind of came out of nowhere.  So in your SDE model you do assume the noise is anisotropic and it is potentially temporally dependent?  This is not the case in equation (2)?  Discussing what is being generalized here would be helpful.\n- The analysis and experiments assume a batch size of 1?  How would the results change for a realistic scenario with mini batches (batch size = m > 1)?\n- Could you provide more intuition for the expressions in equation (11)?\n- To estimate $\\alpha(t)$ and $\\beta(t)$, \"each experiment is repeated for ntrial = 100 independent runs\".  Does this mean you train a neural network 100 independent times keeping track of the logits through time? The estimated $\\hat{\\alpha}(t)$ and $\\hat{\\beta}(t)$ are then used to simulate your SDE and these dynamics are displayed over one of the trials of training in figure 5? Can you confirm that your simulated DNN dynamics require training a DNN?  I apologize if I didn't follow the experimental setup.\n- You write \"This link between a tractable mathematical model (LE-SDE/ODE) and an expressive yet complicated black box would potentially yield greater benefits towards demystifying the magic of deep models\" but the future directions seem to be strategies for expanding the model rather than providing explanation for how to use the model.  For example, does the model suggest hyperparameter settings or architecture design that would facilitate separability?\n\n\n**Specific Comments:**\n\n- [equation 4] - equation (1) is between m+1 and m, while equation (4) is between m and m-1.  Might consider making these consistent?\n- [equation 4] - In the current notation it wasn't clear to me why equation (1) can't capture equation (4) (i.e. why can't I pull the $H$ into the definition of $E$)?\n- [line 47] - \"essentialingredient\" => essential ingredient\n- [line 58] - \"emerge\" => emergence\n- [line 74] - \". also appears\" => missing a word?\n- [line 116] - \"that common\" => that is common or commonly discussed\n- [line 190-191] \"discussions on the empirically best ways of choosing the universal direction $\\nu$ are detailed in the Appendix\"... what is the universal direction?  Where is $\\nu$ discussed earlier?\n- [line 196] - \"H is SDP\" => do you mean SPD (semi-positive definite)?\n- [line 325] - \"Aa\" => \"A\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the training dynamics of neural network by considering a stochastic differential equation model based on the local elasticity assumption of features of neural networks. The proposed model is justified by numerical experiments and sheds new light on the training dynamics. Overall the referees all find this paper interesting and novel, and hence the meta-reviewer would recommend acceptance of the paper as a poster. 
Decision: Accept (Poster)