Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: 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. 

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

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

The 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. 
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_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: 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. 
main_review: **Originality/Significance:**
The 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.

**Clarity/Quality:**
This 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.


**Major Areas for Improvement:**

 - 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!
 - 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.
- Is $\alpha(m)$ and $\beta(m)$ implicitly assumed to be non-negative? For example, is $\alpha,\beta \ge 0$ in theorem 2.1?
- 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.
- 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.
- 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)?
- Could you provide more intuition for the expressions in equation (11)?
- 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.
- 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?


**Specific Comments:**

- [equation 4] - equation (1) is between m+1 and m, while equation (4) is between m and m-1.  Might consider making these consistent?
- [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$)?
- [line 47] - "essentialingredient" => essential ingredient
- [line 58] - "emerge" => emergence
- [line 74] - ". also appears" => missing a word?
- [line 116] - "that common" => that is common or commonly discussed
- [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?
- [line 196] - "H is SDP" => do you mean SPD (semi-positive definite)?
- [line 325] - "Aa" => "A"
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.

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?

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

Typos: 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.
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_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: 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.
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.

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

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

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

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

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

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

Some other comments:
- 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.
- 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.
- In Fig. (4) caption: “crosses, 1 entering…” is not readable.

limitations_and_societal_impact: Yes
review_title: null

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)