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