Imitating Deep Learning Dynamics via Locally Elastic Stochastic Differential Equations
Year: 2021

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 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 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 [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, an analysis of our SDEs shows that the emergence of a 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 a synthesized dataset of geometric shapes and CIFAR-10.

heading: 1 Introduction
text: 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 [31, 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 a theoretical perspective [20, 11, 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 a continued effort to demystify how simple optimization methods give rise to impressive generalization performance, for example, [49, 26, 4], this is 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 to 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 a 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 a cat, how does the hidden representation of another image—say, an image of another
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
X̄ k (T )−X̄ l(T ) for a large T , where k, l are chosen heuristically. The dashed lines in (b) and (d) are
simulated paths from Equation (9) using estimated α̂(t) and β̂(t). More details are given in Section 4 and Appendix D.2.
cat or an image of a plane—evolves 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 a similar sample (an image of another cat) than on a dissimilar sample (an image of a plane) [19].
Motivated by the phenomenon of local elasticity, we propose a 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 on specific training data of all samples.
Our main finding uncovers a sharp phase transition phenomenon regarding the intra-class and interclass 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 is 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 1 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 of 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 a synthetic dataset of geometric shapes, as well as on CIFAR-10. The experimental evidence consistently supported our model, which provides new insights into the dynamics of deep learning training.

heading: 1.1 Related Work
text: Dynamics in Deep Neural Nets. Many properties of linear deep neural nets are relatively well understood, such as the loss landscapes [25], trajectory-based convergence [2, 12], and implicit acceleration [3]. Exact solutions of the training dynamics can be obtained in certain initialization schemes [40, 41, 30]. In the presence of non-linearity, various assumptions are generally made. [39, 17] studied the dynamics of shallow neural nets with non-linearity and the neural tangent kernel (NTK) literature [22, 4, 14] 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 can be found in [47]). Although as approximations, NTKs are generally used when studying optimization trajectories of neural nets such as [34], which also appears implicitly in many works when studying the optimization trajectories of neural net training [45, 1, 13, 24, 7].
SGD as SDEs in 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-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 be meaningful in practical settings and obtained necessary conditions for the validation of such approximation.
Local Elasticity and Phenomenological Models. Local elasticity is proposed in [19] as a phenomenological approach to reasoning the behaviors of neural networks. This phenomenon has inspired several works on generalization bounds [10] and an improvement on the NTK [6].

heading: 2 Binary Separation via LE-SDE

heading: 2.1 Setup, Notations and Assumptions
text: 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 [n] the set {1, . . . , n}. When there is no ambiguity, we will write both X(t) and Xt for a continuous-time (possibly stochastic) process.
Classification Problem. Consider a K-class classification problem with K ≥ 2, with each class having n training examples. We denote by zki ∈ Rd the i-th sample of the k-th class, and yki ∈ [K] its label, where i ∈ [n], k ∈ [K]. A neural net is a function f : Rd → RK that maps the samples to logits ( pre-activation of the softmax).
Feature Vectors. We denote byXki (m) ∈ Rp a p-dimensional feature of the i-th sample in the k-th class learned by the neural net at iteration m. For example, it can be the logits or the output of the second-to-last layer. Assume that the initial valuesXki (0) are i.i.d. samples from some distribution for each i ∈ [n], k ∈ [K]. We use i, j ∈ [n] as indices for an individual sample, k, l ∈ [K] for classes, and capital letters J and L to indicate random samples from Unif([n]) and Unif([K]), 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-th sample is sampled from the Lm-th class, where Jm ∼ Unif([n]) and Lm ∼ Unif([K]). Training on zJmIm affects the features of another data sample zki in the form of
Xki (m)−Xki (m− 1) = h · Ek,Lm(m)XLmJm (m− 1) + √ hζki (m− 1) (1)
where i ∈ [n], k ∈ [K], h is the step size, and ζki (m) is the noise term that is modeled as Gaussian noise. The scalar Ek,Lm measures the strength of local elasticity that z Lm Jm
exerts on zki at iteration m. We assumeXki (0), ζ k i (m) are jointly independent.
Local Elasticity. Clearly, by writing Ek,l(m), we assume that this effect depends only on the class l, k, and time m. We write the effect matrix as E(m) = (Ek,l(m)) K k,l=1. For ease of exposition, we assume E only consists of two values α(m) and β(m), with α(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, as measured by some similarity measure in the feature space. We also assume that the LE effect is “proportional” to the feature XLmJm (m) itself. We generalize this point in Section 3 by introducing a transformation matrixH on the features.

heading: 2.2 Binary LE-SDE
text: 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 a one-dimensional 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 E11 = E22 = α, and the inter-class effect is E12 = E21 = β, both of which are time-independent. Expanding equation (1), for 1 ≤ i, j ≤ n and m ≥ 0, when we train the
model on the Jm-th training example from the Lm-th class, we have{ X1i (m) = X 1 i (m− 1) + h · αXLmJm (m− 1) + √ h · ζLmi (m− 1),
X2j (m) = X 2 j (m− 1) + h · βXLmJm (m− 1) + √ h · ζLmj (m− 1).
In the limit of h → 0, we can show that Xki (m) approximates some continuous-time stochastic processes Xki (t) (under the identification of t = mh) governed by the set of stochastic differential equations as follows:
dXki (t) =
( α
2 X̄k(t) +
β 2 X̄3−k(t)
) dt+ σ dW ki (t), t ≥ 0, k ∈ [K], i ∈ [n] (2)
where X̄k(t) := ( Xk1 (t) + · · ·+Xkn(t) ) /n, and W ki are independent standard Wiener processes. The detailed derivation is given in Appendix A.
Now averaging over i for each k in equation (2), we obtain the following set of two ordinary differential equations (ODEs) governing the per-class means that X̄k(t) for k = 1, 2:
dX̄k(t) =
( α
2 X̄k(t) +
β 2 X̄3−k(t)
) dt+ σ d
W ki (t) + · · ·+W ki (t) n .
Taking the limit of n → ∞, we observe that σdW 1(t)+···+Wn(t)
n ⇒ 0. Thus, the above display converges weakly to the following ODE:
dX̄k(t)
dt = α 2 X̄k(t) + β 2 X̄3−k(t). (3)
With the initial conditions EXki (0) = ck, the solution to the above ODE is
X̄1(t) = c1 − c2
2 e α−β 2 t + c1 + c2 2 e α+β 2 t, X̄2(t) = −c1 − c2 2 e α−β 2 t + c1 + c2 2 e α+β 2 t.
In the finite-sample setting, we may replace X̄1, X̄2 in the SDE (2) by their deterministic solutions and obtain {
X1i (t) = c1−c2 2 e α−β 2 t + c1+c22 e α+β 2 t − c1 +X1i (0) + σW 1i (t), X2j (t) = − c1−c22 e α−β 2 t + c1+c22 e α+β 2 t − c2 +X2j (0) + σW 2j (t).
We are now ready to derive the condition under which these 2n feature vectors become asymptotically separable, that is, miniX1i (t) > maxj X 2 j (t) or maxiX 1 i (t) < minj X 2 j (t) as t→∞.
Theorem 2.1 (Separation in Binary Classification). Given the feature vectors X1i (t), X2j (t) for i, j ∈ [n], as t→∞ and large n,
1. if α > β, they are asymptotically separable with probability tending to one,
2. if α ≤ β, they are asymptotically separable with probability tending to zero. This result indicates a sharp phase transition when α is just above β, 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 are multiple classes, when features are 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.

heading: 3 General LE-SDE Model
text: Now, we consider the general case where K ≥ 2 and the feature vectors are p-dimensional with p ≥ 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 matrixH with the (i, j)-th blockHi,j 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,Hi,j can be viewed as inducing a metric on the feature space under which local elasticity manifests.
As hinted before, in the case of multiple-class features in higher dimensions, we want to guarantee a stronger separation: to 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 is trained under the softmax cross-entropy loss, previous work suggests they separate according to specific geometric structures [38]. To this end, we need to adjust the raw feature vectorsXki 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 transformationHk,Lm ∈ Rp×p to the features. For k ∈ [K], i ∈ [n], and at iteration m, we have the following:
Xki (m) = X k i (m− 1) + h · Ek,Lm(m)Hk,Lm(m)XLmJm (m− 1) + √ hζki (m− 1). (4)
TheHk,Lm(m) term models the LE effect as proportional to a linear “transformation” of the features. The dynamics in Equation (1) are special cases whenHk,l(m) ≡ Ip for all k, l ∈ [K]. By specifying a properH , we can overcome the limitation in our toy example of not knowing which class converges to where. We specify interesting choices ofH in Section 3.2.
A further step of abstraction is to write X̃ k (m) instead ofXki (m), to indicate one generic sample from the distribution Dk(m) of all the features of class k at iteration m. As in Section 2.2, we can derive the continuous dynamics of equation (4) in the limit of h → 0 in the same way as equation (2). Similar to writing X̃ = (X̃
k )Kk=1 ∈ RKp for the concatenation of per-class features,
X̄ = (X̄ k )Kk=1 ∈ RKp is the concatenation of per-class mean features. Our model (4) approximates the following SDE with identification t = mh as h→ 0. We term this model the LE-SDE:
dX̃t = M tX̄t dt+ Σ 1 2 t dW t, (5)
whereW t is the standard Wiener process in RKp, Σt is the covariance matrix, andM t ∈ RKp×Kp is aK×K block matrix, with each block of size p×p. The (k, l)th block ofM t isEk,l(t)Hk,l(t)/K when l 6= k, and El,l(t)H l,l(t)/K when l = 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̃t from its distribution, the per-class mean X̄t satisfies the following system, which we term the LE-ODE:
X̄ ′ t = M tX̄t. (6)
Under the assumptions in Section 2.1, we define γ(t) = min {α(t)− β(t), α(t) + (K − 1)β(t)} , A(t) = ´ t 0 α(τ) dτ , B(t) = ´ t 0 β(τ) dτ , and Γ(t) = min {A(t)−B(t), A(t) + (K − 1)B(t)}.

heading: 3.1 The Separation Theorem
text: Similar to the discussions in Theorem 2.1, the LE-SDE allows us to derive the separability result for a general K and p ≥ K. We say the feature vectors { (Xki )i∈[n] } k∈[K] are separable if for any two classes k 6= l, there exists a hyperplane in Rp that linearly separates the features of the two classes. To characterize the separation as in Theorem 2.1, we need conditions on α(t), β(t) as therein. Intuitively, when γ(t) ≤ 0, the classes cannot be separated, even in a pairwise manner. Therefore, we focus on a more interesting case when γ(t) > 0. 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., γ(t) > 0), assume H = (Hij)ij is positive semi-definite (PSD) with positive diagonal entries. As t→∞, we have1:
1. if γ(t) = ω (1/t), the features are separable with probability tending to 1;
1Here, γ(t) = ω (1/t) stands for γ(t) 1/t as t→∞. For example, 1/t0.5 = ω (1/t) and (t ln t)−1 = o (1/t) as t→∞.
2. if γ(t) = o (1/t), and the number of per-class-feature n tending to∞ 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 a general case. The proof to Theorem 3.1 as well as discussions on the empirically best ways of choosing the universal direction ν (i.e., a direction that does not depend on the class index) are detailed in Appendix C.2.

heading: 3.2 Two Specific Models
text: We next discuss two specific choices of theH matrix that allows us to analyze X̃ precisely.

heading: 3.2.1 Isotropic Feature Learning Model
text: As a straightforward extension to Section 2.2, we can simply choose H lk = Ip to be the identity matrix. This choice ofH is PSD, and thus, we can apply Theorem 3.1 to obtain the conditions for asymptotic separation. In this case, the solution X̄(t) to the LE-ODE can be computed analytically as given in the following proposition.
Proposition 3.2 (I-model). LetHk,l = Ip, then the solution to the LE-ODE (6) is given by
X̄(t) = ce 1 KA(t)− 1 KB(t) + (1K ⊗ c0) e 1 KA(t)+ K−1 K B(t), (7)
where c = (ck)Kk=1 ∈ RKp and c0 ∈ Rp are constants with ∑K k=1 ck = 0 ∈ Rp and X̄(0) = c+c0.
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 k and that from class l at time t is given by (ck − cl)e 1 KA(t)− 1 KB(t) + Σ1/2(W kt −W lt), 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 feature model, or I-model for short; as theH matrix has identity matrices as its blocks and consequently the dynamics do not prescribe any preferred directions for each class.

heading: 3.2.2 Logits-as-Features Model
text: An important type of features in neural nets is the logits, the outputs of the neural net before the softmax layer. A logit vector (or logits) is K-dimensional, and in this model we identify X̃ k (t) as the logits at time t of a generic sample from the k-th class. In a well-trained neural net, the logits of a 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̃ k by its margin, dk := ek − softmax(X̃ k ), which roughly aligns with the direction of ek − 1p/K. This suggests us how to choose the metric under which local elasticity acts: we can chooseH l,k such that it always aligns X̃ k in the direction of dk, that is,
Hij = H̄ j := djd
> j
‖dj‖22 ∈ Rp×p, dj := ej −
1
K 1p ∈ Rp, j ∈ [K]. (8)
Roughly speaking, the map x 7→ H̄jx projects x in the direction of dj and ideally aligns x with dj after iterative applications; hence, H̄
j can be viewed as an approximation of the nonlinear transformation in the exact dynamics in the sense that the direction of their stationary point coincides. Furthermore, H̄j thus defined has operator norm 1; thus, it does not affect the magnitudes, but only directions. Note thatH 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). LetH be the same as in equation (8), then the solution to the LE-ODE (6) is given by
X̄(t) = c0 + C1de 1 KA(t)− 1 KB(t) + ( K−1∑ l=1 C2lf l ) e 1 KA(t)+ 1 K(K−1)B(t), (9)
where f l’s are fixed vectors in RK 2 , c0 ∈ RK 2 is a constant vector with K(K − 1) degrees of freedom, and C1, C2l ∈ R, for l ∈ [K − 1] are constants.
The specific form of f l 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 f l’s when K = 3.
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.

heading: 3.3 Connection with Neural Collapse
text: Neural collapse is a recent phenomenological finding on the geometry of the logits learned by deep neural nets at convergence with the cross-entropy loss [38] (see an explanation of neural collapse in [16]). Simply speaking, taking our L-model as an 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 an ambient space of p dimensions. Formally, we say a set of vectors {si}Ki=1 form an ETF in Rp if they are the columns of a matrix
S =
√ K K − 1Q ( IK − 1 K 1K1 > K ) , (10)
where Q ∈ Rp×K , and Q>Q = IK . As a direct corollary of Proposition 3.3, when H is set according to equation (8), we find that our L-model also predicts the existence of neural collapse from the local elasticity point of view.
Proposition 3.4 (Neural Collapse of the LE-ODE). Under L-model and the same setup as in Theorem 3.1, if γ(t) > 0 and there exists some T > 0 such that B(t) < 0 for t ≥ T , then{ X̄ k (t)/‖X̄k(t)‖ }K k=1 forms an ETF as t→∞.

heading: 4 Experiments
text: We perform various experiments to test our theory, where we choose logits as our protagonist2.

heading: 4.1 Setup
text: Datasets and Models. We perform experiments on a synthesized dataset called GEOMNIST containing K = 3 types of geometric shapes (RECTANGLE, ELLIPSOID, and TRIANGLE) and on CIFAR-10 ([28], denoted by CIFAR) with K ∈ [2, 3] classes. A few samples from GEOMNIST are shown in Figure 2. We vary the number of training samples per class and label pollution ratio perr and use variants of the ALEXNET ([29]) model. More details can be found in the Appendix.
Training Configurations. All models are trained for T = 105 iterations (for GEOMNIST) or T = 3 × 105 iterations (for CIFAR) with a learning rate of 0.005 and a batch size of 1 under the softmax cross-entropy loss. Models on GEOMNIST 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 ntrial = 100 independent runs to estimate X̄(t). We use both the isotropic feature learning model (Section 3.2.1) and the logits-as-features model (Section 3.2.2), denoted by L-model and I-model respectively, to estimate α(t) and β(t). The
2Code for reproducing our experiments is publicly available at github.com:zjiayao/le_sde.git.
0 2000 4000 6000 8000 10000
t
0
5
10
15
20
Â (t
), B̂
(t )
Â(t)
B̂(t)
(a) GEOMNIST (I-model).
0 2000 4000 6000 8000 10000
t
−0.005
0.000
0.005
0.010
0.015
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(b) GEOMNIST (I-model).
0 5000 10000 15000 20000 25000 30000
t
−2.5
0.0
2.5
5.0
7.5
10.0
12.5
15.0
Â (t
), B̂
(t )
Â(t)
B̂(t)
(c) CIFAR (I-model).
0 5000 10000 15000 20000 25000 30000
t
−0.0010
−0.0005
0.0000
0.0005
0.0010
0.0015
0.0020
0.0025
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(d) CIFAR (I-model).
0 2000 4000 6000 8000 10000
t
−20
−15
−10
−5
0
5
10
15
20
Â (t
), B̂
(t )
Â(t)
B̂(t)
(e) GEOMNIST (L-model).
0 2000 4000 6000 8000 10000
t
−0.010
−0.005
0.000
0.005
0.010
0.015
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(f) GEOMNIST (L-model).
0 5000 10000 15000 20000 25000 30000
t
−30
−20
−10
0
10
Â (t
), B̂
(t )
Â(t)
B̂(t)
(g) CIFAR (L-model).
0 5000 10000 15000 20000 25000 30000
t
−0.003
−0.002
−0.001
0.000
0.001
0.002
0.003
0.004
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(h) CIFAR (L-model).
Figure 3: Estimated Â(t), B̂(t), α(t), and β(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 B̂(t) and the other two rows α̂(t) and β̂(t).
L-model is used only when K = 3. To estimate α(t) and β(t), we first estimate A(t) and B(t) by
(I-model) Â(t) = avg avgk log ∣∣∣|X(X̄k−|X)K−1 c0c K−1 k ∣∣∣, B̂(t) = − avg avgk log ∣∣∣ c0ck X̄k−|X|X ∣∣∣, |Xt := avgl X̄ l t,
(L-model)
{ Â(t) = A′(t) + 2B′(t),
B̂(t) = 2(B′(t)−A′(t)), A ′(t) := log ∣∣∣〈X̄>v1 − 1〉∣∣∣ , B′(t) := log ∣∣∣〈X̄> (v2 − 43v1)〉∣∣∣ , (11)
where vector division is interpreted entry-wise. We write avgl for averaging over the class index, avg for averaging over the coordinates, and define 〈X〉(t) := X(t)/X(0). We explain how and why to choose the vectors v1 and v2 in Appendix D.1. The main idea is to view the eigenvectors of the Kp-by-Kp drift matrix as a concatenation of K vectors of dimension p and construct their linear combinations such that one or more independent components in the solution vanishes. With A(t) and B(t) estimated, we use the Savitzky - Golay filter to obtain α̂(t) and β̂(t) through numerical differentiation. For GEOMNIST and CIFAR datasets, we choose window sizes of this filter as 191 and 551, respectively, in Figure 3, and 21 and 21, respectively, in Figure 5.
We assess the tail of α̂(t) and β̂(t) by a tail index defined as rα := sups {s : limt→∞ α(t) · ts <∞} and rβ is defined similarly. We estimate r̂α by fixing an interval [T1, T2] with T1 < T2 sufficiently large such that we may ignore terms with smaller order and have r̂α = 1− avgT1≤t≤T2 logα(t) log(1+t) , and similarly for r̂β . We use the estimates from the last 1000 iterations for averaging in our experiments.

heading: 4.2 Results
text: Local Elasticity in Neural Net Training. Local elasticity manifests from our model as the heaviness of the tail of γ(t) = α(t) − β(t), and in Figure 3, we plot the estimations Â(t), B̂(t), α̂(t), and β̂(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 t 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 c0 in the L-model; (iii) Both α̂(t) and β̂(t) behave similarly on both datasets.
Phase Transition of Separability. Theorem 3.1 states that separation of features under the LE-SDE takes place when γ(t) = α(t)− β(t) = ω(1/t), or roughly speaking, when rγ = min{rα, rβ} < 1. Although we cannot directly control α(t) and β(t), we can bias them by tuning the label corruption ratio perr. When perr ≈ p∗err := 2/3, we are in effect assigning labels completely at random and thus we expect a phase transition of separability should happen around p∗err. This is indeed the story depicted in Figure 4: Figures 4a and 4b show that the validation loss and accuracy for p ≥ perr
0 2000 4000 6000 8000 10000
t
−15
−10
−5
0
5
10
X̄ (t
)
X̄00 X̄01 X̄02
(a) GEOMNIST (I-model).
0 2000 4000 6000 8000 10000
t
−7.5
−5.0
−2.5
0.0
2.5
5.0
7.5
10.0
X̄ (t
)
X̄00 X̄01 X̄02
(b) GEOMNIST (L-model).
0 5000 10000 15000 20000 25000 30000
t
−2
−1
0
1
2
X̄ (t
)
X̄00 X̄01 X̄02
(c) CIFAR (I-model).
0 5000 10000 15000 20000 25000 30000
t
−2
−1
0
1
2
X̄ (t
)
X̄00 X̄01 X̄02
(d) CIFAR (L-model).
Figure 5: Simulated LE-ODE solutions versus genuine dynamics. We use α̂(t) and β̂(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 dynamics reasonably well.
are not increasing over time and in Figure 4c we observe the minimum tail index of α(t) and β(t) crosses 1 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 β(t), it is natural to ask, to what capacity can our LE-ODE models recover the real dynamics of deep neural nets? We use the forward Euler method to 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 β(t), we rescale the simulated paths such that its 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 of three logits.

heading: 5 Discussion and Future Works
text: In this study, we introduce 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 elasticity strengths α(t) and β(t) are well estimated, our model can faithfully simulate the dynamics of neural nets. We outline a 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-definite LE matrices E(t). (ii) Minibatch Training, Imbalanced Datasets, and Label Corruptions. Generalizing the drift matrix to M t = (Et ⊗ P ) ◦H/K for a K-by-K doubly stochastic matrix P can be used to model various sampling effects. (iii) Beyond L-model for Imitating Genuine Dynamics of DNNs. Although the L-model is shown to be able to mimic the real dynamics reasonably well, we postulate that a more precise model might have its (i, j)-th block encode the other directions other than dj .

heading: Acknowledgements
text: This work was supported in part by NSF through CCF-1934876, an Alfred Sloan Research Fellowship, the Wharton Dean’s Research Fund, and ONR Contract N00014-19-1-2620. We would like to thank Dan Roth and the Cognitive Computation Group at the University of Pennsylvania for stimulating discussions and for providing computational resources.

Review summaries:
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.   Numerical 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. 

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.   ===================== After rebuttal: Thanks for the detailed reply!  Overall while I am pleased that the authors made a lot of efforts to clarify, I'm not quite convinced. Some notes: - 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. - 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.  I 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.  Given 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.

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. 

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.