Imitating Deep Learning Dynamics via  Locally Elastic Stochastic Differential Equations

1 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 [ 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 . 1.1 Related Work . 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 ] . 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 [ 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 . 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 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 .

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