Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training

Abstract: The mean field theory of multilayer neural networks centers around a particular infinite-width scaling, in which the learning dynamics is shown to be closely tracked by the mean field limit. A random fluctuation around this infinite-width limit is expected from a large-width expansion to the next order. This fluctuation has been studied only in the case of shallow networks, where previous works employ heavily technical notions or additional formulation ideas amenable only to that case. Treatment of the multilayer case has been missing, with the chief difficulty in finding a formulation that must capture the stochastic dependency across not only time but also depth.In this work, we initiate the study of the fluctuation in the case of multilayer networks, at any network depth. Leveraging on the neuronal embedding framework recently introduced by Nguyen and Pham, we systematically derive a system of dynamical equations, called the second-order mean field limit, that captures the limiting fluctuation distribution. We demonstrate through the framework the complex interaction among neurons in this second-order mean field limit, the stochasticity with cross-layer dependency and the nonlinear time evolution inherent in the limiting fluctuation. A limit theorem is proven to relate quantitatively this limit to the fluctuation realized by large-width networks.We apply the result to show a stability property of gradient descent mean field training: in the large-width regime, along the training trajectory, it progressively biases towards a solution with "minimal fluctuation" (in fact, vanishing fluctuation) in the learned output function, even after the network has been initialized at or has converged (sufficiently fast) to a global optimum. This extends a similar phenomenon previously shown only for shallow networks with a squared loss in the empirical risk minimization setting, to multilayer networks with a loss function that is not necessarily convex in a more general setting.


Reviews and author responses:

Based on the infinite-width mean-field (MF) limit of multi-layer neural networks (NNs) under the neuronal embedding framework proposed previously, this paper studies the fluctuation of a finite-width multi-layer NN (under random sampling and gradient flow training) around its mean-field limit. The authors show that the fluctuation depends on two components - one being a Gaussian process resulting from the Central Limit Theorem, and the other following a second-order MF equation. The authors also show that under certain conditions, the fluctuations tend to vanish as we take the long-time limit.  1. Originality: The novelty of this paper is extending the analysis of the fluctuations of NNs from their infinite-width limits from shallow models to multi-layer ones, based on the neuronal embedding mean-field limit.

2. Clarity: The amount of efforts that the authors spend in formalizing all the theory is quite impressive. However, perhaps for most of the audience at NeurIPS, many of the technical details that allow the framework to be rigorous are not of huge interest, but instead make the paper difficult to penetrate. This is partly due to the abstract nature of the subject, though perhaps there is also some room for the authors to improve the readability of the paper.

3. It would help to comment on the connection between the theory proposed in this paper and the one in ref. [3] when restricting to shallow NNs - is there a similarity in their high-level proof ideas as well as a correspondence between the various objects being defined? For example, to argue for the decay of the fluctuations over time, it seems that both papers rely on showing that a fluctuation term lying in the range of a negative semi-definite linear operator. Are there deeper connections?

4. For the audience at NeurIPS, a main take-home message of this paper seems to be that gradient flow has a variance reduction effect for not only shallow but also multi-layer NNs. It would be more relevant and convincing if this theoretical result can be supported by numerical experiments. And overall, I am not very sure if NeurIPS is the best venue for the mathematical rigor of this paper to be fully appreciated.
 More could be said about the limitations and certain assumptions of the theoretical results. For example, how does the neuron sampling method adopted in this paper compare with the common initialization methods in practice?

 Thanks for your careful responses! The experimental results are helpful, as figures (c) - (e) indeed show the variance reduction effect in actual NNs and also compare the variance obtained by gradient descent with that obtained by MC sampling to illustrate the two components of the variance. This constitutes a nice complement to the theoretical results. I will raise my rating from 5 to 6.

 Many thanks for the thoughtful comment.

We think the mean field regime, and the neuronal embedding framework which rigorizes the mathematical foundation for multilayer neural nets in this regime, are a proxy to understand dynamics of neural networks. It is a good proxy, in the sense that nonlinear feature learning is retained, unlike NTK, and yet we can still prove and understand many properties of the dynamics.

Let us highlight an important fact from the work of Yang and Hu: they show that under the standard parameterization, NTK behavior still arises. As such, we think all infinite-width scalings which avoid NTK (including both the mean field scaling, and Yang and Hu's) are proxies of practical finite-width networks. This also motivates the need for a finite-width study, such as our submission.

What is currently known about Yang and Hu’s proposed scaling is much more limited than the mean field scaling. In particular, the work by Yang and Hu, with quite a lot of calculations, only proves that after 1 step of SGD and with i.i.d. Gaussian initialization, the weights move. In contrast, the fact that the weights move in the mean field regime — at any SGD step — is actually immediate from its formulation. The neuronal embedding framework has been used to study what happens at convergence time, what guarantees a good optimization outcome, what simplifications the dynamics may have, under a wide range of (i.i.d. or non-i.i.d.) initializations (see e.g. [Nguyen and Pham 2020], [Pham and Nguyen 2020]). In our work, we further this program to study what happens with large but finite widths, via the fluctuation, and how gradient descent benefits training stability in the long-time horizon. In short, what we have known for the mean field regime is a whole lot more, though not yet a complete theory.

Of course, the mathematics looks more involved. We offer an alternative view to this: an abstract mathematical view could prove to be powerful. For example, [Javanmard et al 2019] proves a convergence rate for two-layer neural networks, thanks to a deep theorem from studies of Wasserstein gradient flows. Such a result is difficult to obtain otherwise.

The gradient vanishing / explosion issue -- at initialization -- can actually be avoided by an appropriate choice of initialization hyperparameters. This is a good question indeed. In the standard scaling with i.i.d. initializations, there have been amazing works, such as [Schoenholz et al 2017], but some calculation has to be done. In the mean field scaling, the calculation with i.i.d. initializations is actually easier. As shown in [Nguyen and Pham 2020] and [Pham and Nguyen 2020], certain sophisticated non-i.i.d. initializations could be much more beneficial than i.i.d. initializations, so we expect the theory does not stop at those simple calculations.

Finally let us re-mention that Lu et al 2020 empirically demonstrated an improved performance over standard Resnets by using the mean field-scaled Resnets.

References:

"Deep information propagation", SS Schoenholz, J Gilmer, S Ganguli, J Sohl-Dickstein, 2017.

"A rigorous framework for the mean field limit of multilayer neural networks", P-M Nguyen, HT Pham, 2020.

"A note on the global convergence of multilayer neural networks in the mean field regime", HT Pham, P-M Nguyen, 2020.

“Analysis of a two-layer neural network via displacement convexity”, A Javanmard, M Mondelli, A Montanari, 2019.

"A mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth", Y Lu, C Ma, Y Lu, J Lu, L Ying, 2020.

 In the case you have any questions after reading our responses, we would love to have an opportunity to address them, so feel free to let us know!

 Thanks for a lot for including new plots. I think they are very valuable to readers. 

I appreciate the authors' efforts in forwarding the neural embedding framework. The current formulation seems to be quite far away from practice, in particular, the 1/width factor in all hidden layer is unnatural. Will this cause a severe gradient vanishing / numerical issue for very deep network >=100 layer if this parameterization method is being used in practice?  The parameterization used in Yang&Hu is much closer to practice and seems to be less technical. What is the limitation of that framework in term of mathematical formulation? Alternatively, what are the benefits of the neural embedding framework compared to Yang&Hu and others? It will be helpful to include such comparison in the paper, in particular, for (NeurIPS, ICML, etc) conference papers, since most of the readers are from a less technical background. 

I also agree with reviewer SSUu that NeurIPS may not be the optimal avenue for publishing this paper, since the mathematical efforts and rigorousness may not be fully appreciated. 

Overall, I tend to keep my score. 

This paper studies the mean-field limit of deep fully connected neural networks. The approach is based on extending the analysis of previous work on “neuronal embedding”. Several results are obtained in the considered setting, including 1) the derivation of an ODE system that describes the the asymptotic fluctuations around the mean field limit when there are finite samples; 2) a result showing that the these fluctuations can be characterized by a Gaussian process via a CLT. This extends the analysis of Ref [13] in the paper, which establishes the validity of the mean-field limit, without explicit characterization of the fluctuations around this limit. As an application of these results, the authors studies the evolution of the limiting variance process, where the authors show that under assumptions on the initial condition or the convergence rate, the variance is non-increasing under gradient descent, alluding to a variance reduction effect.  The topic under investigation is an interesting one, since most existing analysis (be it NTK or MF) usually focus on the limiting process itself, without explicit characterization of the fluctuations and their effects. In this sense, the paper addresses a relevant problem. The theoretical results of the paper appear sound, but I think there is some lost opportunity here, in that the consequences of these results with respect to deep learning are insufficiently explored. As far as I can see, the only consequence of these results explored is Theorem 9 and 10 on the long-time limit of the scaled variance.

Below I list some explicit questions that the authors may wish to address:

1. Pertaining to theorems 9 and 10: the limit here studied is $N\rightarrow\infty$, and then $t\rightarrow\infty$, yielding 0 asymptotic variance. However, in practice, we hold $N$ constant, and training gives $t\rightarrow\infty$, thus it appears the correct limit to consider when discussing variance reduction is $t\rightarrow\infty$ at finite $N$, or at least $N\rightarrow\infty$ after taking $t\rightarrow \infty$. Does your results apply to the exchanged limits? 
2. One of the main results here is the derivation of the limiting ODEs in Section 3.1,3.2 and the corresponding characterizations of the finite-sample process as a Gaussian perturbation of this. Hence, it may be useful to discuss any insights from these results. For example, are there anything different about the modes of convergence involved with the deep neural network that is different from its shallow counterparts? What do the results in 3.2 tell us about the nature of the fluctuations (R_i)? As written, the paper reads like a collection of technical results with little summary of its implications.
3. Other issues
   1. Eq (2) may require more explanation without the need to read ref [13], how is it fundamental? It appears to be a L^2 law of large numbers type of result (with a convergence rate). What does the universal constant $c_*$ capture?
   2. What is $O_{D,T}$? in definition 1?
   3. Theorem 9: at initialization, why would one expect the derivative to have zero expectation? Does it only apply when one initialize it at the global optimum. Also, doesn’t Theorem 10 includes Theorem 9 as a special case? Is there a need to present 9 then?

---

I thank the authors for addressing the comments - most of the minor questions have been addressed. While the issue with the exchanged limits cannot yet to addressed, there is some experimental evidence that it holds - if these new results can be considered for review then I am happy to recommend acceptance. The main limitation is, as rightly identified by the other reviewers, is limited motivation for the practical readers. Not applicable

 We thank the reviewer for finding the study interesting and relevant!

Regarding the switching of the limits $N\to\infty$ and $t\to\infty$, we are unfortunately unable to do this yet for the fluctuation. A rare instance where this is possible is the optimizing loss, which is known to decrease monotonically with $t$ regardless of $N$ and hence the limit switching is easy. As can be seen in our numerical illustration (the details of which are in the general comment above (https://openreview.net/forum?id=jg9LM8QItms&noteId=OibEz6PQV-B)), the long-time horizon behavior of the output variance, predicted by the theorems, can already be observed well beyond convergence for reasonable width $N$.

Regarding the derivation in Section 3, as explained in the beginning of the section, the main message is that the fluctuation can be essentially broken into two pieces: the Gaussian component $G$ and the second-order MF limit $R$. The MF scaling is typically associated with the following interpretation: neurons at any time act as if they are sampled from some distribution. The fluctuation that arises from this interpretation is $G$. This does not account for the MF interaction among neurons during GD training, whose fluctuation is represented by $R$. As seen from our numerical illustration above, the output variance due to $G$ increases with time, whereas the overall output variance decreases with time. This signifies the role of MF interactions in reducing the variance.

Eq (2) indeed refers to a law of large numbers type of result, so that the readers understand the connection between the MF limit and the neural network. The constant $c_*$ is just an absolute constant that depends on the constants $K$ in Assumption 1. In Definition 1, $O_{D,T}$ is the usual big O notation with implicit constants that may depend on $T$ and $D$. Regarding Theorem 9, the derivative of the loss is zero and the infinite-width network does not move (although the finite-width network still moves on the order of $O(1/\sqrt{N})$), which is congruent with the assumption that we initialize at a global optimum. Theorem 10 does not include Theorem 9 as a special case, since Theorem 9 predicts monotonic decrease of the output variance to zero with stronger assumption on the initialization, while Theorem 10 only predicts eventual convergence to zero variance in more general setting. We will be clearer in our description. 


 We thank the reviewer for finding the results interesting! We will provide some numerical evidence, the details of which are discussed in the general comment above (https://openreview.net/forum?id=jg9LM8QItms&noteId=OibEz6PQV-B).

Regarding the scaling, the mean field scaling (even in the shallow case) follows roughly the $1/N$ scaling, instead of the “standard” $1/\sqrt{N}$ scaling. This scaling enforces nonlinear feature learning as one takes the width to infinity. We note that the NTK behavior can show up for the standard $1/\sqrt{N}$ scaling, and yet NTK does not exhibit feature learning. The paper Mei et al 2019 suggests that other scalings (which may closely resemble practical nets) behave like the mean field scaled nets in the long-time horizon. So while the mean field scaling is an idealized assumption, it allows to gain some understanding without too much oversimplification.

With regards to the position of the paper within the literature, we tend to think of the paper as mainly a theoretical perspective to understand neural nets. Arguably none of the infinite-width regimes (including mean field) capture precisely the behavior of practical nets, and we wish to work with certain idealized assumptions (but not too oversimplifying, like NTK). The mathematics become heavy, which is possibly unavoidable when it comes to the difficult non-NTK multilayer nets. We discuss this more in the general comment above. A longer-term goal is to find theoretically grounded designs. For example, Lu et al 2020 showed some interesting empirical improvement with mean field resnet. But results of this kind are rare, since even a satisfying theoretical picture is still far away.

Thanks for spotting the typo below line 101. For the equation below line 102, there is no $1/N_{i-1}$ factor, so yes, the quantities are rescaled appropriately. In Eq (2), $\mathbb{E}_J$ is the expectation w.r.t. the $J_i$’s random variables, as explained in lines 97-98.

Reference:

S. Mei, T. Misiakiewicz, A. Montanari, Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit, COLT 2019.

Y. Lu, C. Ma, Y. Lu, J. Lu, L. Ying, A mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth, ICML 2020.

 We thank the reviewer for the detailed reading and helpful review! We will add some numerical illustration, the details of which are discussed in the general comment above (https://openreview.net/forum?id=jg9LM8QItms&noteId=OibEz6PQV-B). We hope it would aid understanding of the main results for more general audience of NeurIPS. 

As to compare with reference [3], we note that while the derivation of the fluctuation is different, in both our work and [3], the end result is described in terms of a system of linear ODEs, from which the dynamics of the output fluctuation defined via the linear operator in line 303 can be derived. The main remaining step in the argument is the analysis of the operators, where the difference lies. As explained in Section 4, this is quite immediate for shallow nets, but not obvious for multilayer nets due to the complex covariance structure. Determining the covariance structure in the multilayer case is the key input to Theorem 9 concerning variance reduction when the initialization is at a global optimum. Another important technical difference is between Theorem 10 and its analog in [3] concerning the limiting variance when initializing away from a global optimum. Specifically, due to the interaction among weights across layers in multilayer nets, our argument requires much more delicate estimates for uniform control on weight movements. 

Regarding the theoretical assumptions, some of our assumptions are quite mild (for instance, the initialization assumption should cover most initial weight distributions of interest). One thing to note is that the neural net should follow the mean field scaling, and hence so does its initialization. While this does not strictly follow the usual initialization practice, it allows to retain nonlinear feature learning, as opposed to NTK-like behavior (which actually follows the usual initialization practice). We refer to the general comment above, addressed to all reviewers, for some discussion.


 We thank the reviewers for their careful reading and very constructive comments, especially when the paper is unavoidably heavily technical. Below we address two main common comments.

### Experimental verification

While we generally avoid unhelpful experiments, we agree that some numerical simulations could be useful for verification as well to quickly aid understanding. We provide some results and codes in the following link: https://anonymous.4open.science/r/nn-clt-FF2B/

In particular, this is a simple experiment with MNIST, classifying digits 0 and 4 versus 5 and 9 with 3-layer fully-connected networks of varying width $N$, equipped with Huber loss and tanh activation. We use only 100 images for training. The small scale of this experiment is necessary for us to run full-batch gradient descent in reasonable time, since the paper focuses on the mean-field fluctuation and avoids the effect of SGD (discussed in our Limitations section).

Following the figure in this link: https://anonymous.4open.science/r/nn-clt-FF2B/figure.png

- Fig (a): This is the evolution of the training loss, for $N=50, 100, 200, 400$. Note that the global optimum is found at around iteration $10^4$, and the trajectory indicates nonlinear dynamics.
- Fig (b): This is the histogram of the output fluctuation $\sqrt{N}(\hat{\bf y}(t,x_{test}) - \hat{y}(t,x_{test}))$ at a particular test sample $x_{test}$ and iteration $10^3$. The agreement among various $N$ verifies the existence of a limiting Gaussian-like behavior, predicted by Theorems 5 and 6.
- Fig (c): This is the plot of the output variance ${\bf E}\mathbb{E}_Z[|\sqrt{N}(\hat{\bf y}(t,X) - \hat{y}(t,X))|^2]$ versus training iteration. The variance decreases quickly past iteration $10^4$. This is predicted by Theorem 10. 
- Fig (d): This is the plot of the Gaussian component $\tilde{G}^y$ in the output variance. Note that after iteration $10^4$, this component no longer moves since the global optimum is reached. By contrasting this plot with Fig (c), we see clearer the role of gradient descent training in reducing the variance.
- Fig (e): This is similar to Fig (c), except that we initialize at a global optimum. The variance monotonically decreases on the entire period. This is predicted by Theorem 9.

All simulations are run with a single NVIDIA Tesla P100 GPU. We will provide a full description of the experiment in the updated manuscript, and please feel free to ask for any further detail (e.g. how estimation is done to obtain the plots, etc). The experiment can be repeated by running the provided codes.

### Significance

The significance of the mean field scaling is that as we know, large-width neural nets are prone to be very degenerate and unrealistic (e.g. in NTK regime), and the MF scaling is among the few scaling regimes that we know neural nets are not degenerate and retain some realistic features of practical nets. Analyzing a nontrivial neural net is of course a nontrivial task (see for instance the very recent work by Yang and Hu that provides some preliminary analysis of another nontrivial scaling regime). It becomes increasingly harder for multilayer networks.

One such difficult task is to understand what happens near the end of training. This is arguably an important task. Yet it is expected that the landscape near convergence is typically complex and chaotic, and the infinite-width first-order approximation may no longer be excellent. As such, we think it is important to study a second-order approximation via the (up-scaled) fluctuation. In fact, in our opinion, the finding that gradient descent reduces the output variance (Theorems 9 and 10) is quite encouraging: it suggests that the GD-trained network is not too unstable despite the complex loss landscape.

Reference:
G. Yang and E. Hu, Feature learning in infinite-width neural networks, ICML 2021.

Building upon on the neural embedding framework for multi-layer mean field limit of neural networks [13], the authors derive a system of ODEs that characterized the second order fluctuation of large width network to its MF limit. Interestingly, the authors also show that in the large depth large time limit, variance of the learned output vanishes. The paper contains no experimental results and it is purely theoretical.   Theoretical understanding of neural networks is a central topic in machine learning. The line of research that applies mean field approximation has attracted considerable amount of interest recently. The MF approximation in shallow networks is quite successful. However, the multiple layer setting seems to be quite different and challenging. Building upon on the formulation in [13], the current work derived the second order fluctuation to the mean field limit. Roughly speaking, the authors obtained approximation of several quantities of interest (e.g. outputs, parameters, derivatives of parameters), which are of the following general form 
\begin{align}
N^{1/2 }(f_n -f_{\infty})\approx  R_{\infty}
\end{align} 
Where $N^{1/2 }(f_n -f_{\infty})$ is the scaled correction term and $R_\infty$ is the residual, which is usually characterized by a system of ODE. 

Overall, I think the theoretical results of the paper is quite interesting, in particularly, for readers who are interested in MF approximation of neural networks. However, the current work relies heavily on [13] which I am not familiar with and I could not justify the correctness of its proof. If other reviewers couldn't do so either, I would recommend the authors to submit the paper to a journal of which the reviewers have the interest and bandwidth to go through all the technical details, regardless whether the paper will be accepted by NeurIPS or not. Finally, the MF approach to (multilayer) neural networks has become more theoretical/technical and less practicable, what can be achieved by this direction? Will it lead us to a mathematical foundation for deep learning that explains the fundamental principle behind its great success? A more principle-guided approach for architecture design/ optimization? Or is it just a kind of well-defined mathematical formulation of pure theoretical interest?  


Other comments/ suggestions.

My biggest question is: how relevant is the current formulation to practice? In addition, what audiences are you targeting? People with familiarity of multiple layer mean field or just shallow networks or none? It seems to me that the authors assume the readers have some familiarity with [13]. 

Since this is a heavy theory paper, I would appreciate the authors to provide empirical support to major claims made in their paper.  


Notation in (1) is a bit dense for readers who are not familiar with [13].  What is the connection to practical networks, what parameterization is used? Standard (e.g. He initialization), NTK style parameterization, or the recent feature parameterization claimed in Yang&Hu. From the expectation operator defined on the second line of (1), it seems like there is a $1/ N_i$ normalization factor (rather than the standard $1/\sqrt {N_i}$) which is different from all parameterization approaches mentioned above. Authors should connect (1) to practically used networks in order for readers to understand if this is closely related to practical networks or it is just a mathematical formulation(not necessary relevant to practice).  

Below 101: lowercase $n_i$ should be uppercase $N_i$? 
Second Equation below 102: update for $\omega_i$, should there be a $1/N_{i-1}$ factor in the updates? Or the updates are rescaled by a layer wise learning rate which is suggested in line 104? 

Eq 2; what is $J$ under the $\mathbb E$ ? 

Yang&Hu Feature Learning in Infinite-Width Neural Networks
 The authors have addressed the limitations of the work. As far as I know, this is a pure theory paper and has no negative societal impact. 

Meta-review: The paper studies the fluctuation of a finite-width multi-layer neural network around its mean-field limit, based on the infinite-width mean-field limit under the neuronal embedding framework. While referees have concerns on the practicality of the parametrization, they find the result interesting and novel, especially with the newly added numerical experiments. The meta-reviewer recommends acceptance of the paper as a poster. 

Decision: accept