Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks

Summary

The authors study a finite width correct to the Dynamical Mean Field Theory (DMFT) of finite depth neural networks in the feature learning regime. While I will be the first to admit that I am not an expert on the DMFT calculations, the authors did produce very convincing simulations capturing interesting properties of finite size networks. In particular, the edge of stability behaviour at large learning rates seems to be modelled by the finite width correction. I hope the authors can help me clarify some questions regarding the implementation of the solver, after which I would be happy to raise my score to accept.

Strengths

The authors develop a strong theory of finite width neural networks capable of making accurate predictions.

Weaknesses

N/A

Questions

I would like the authors to provide some intuitions about how the DMFT equations are being solved numerically. I understand this may be the contents of a previous paper (Bordelon and Pehlevan, 2022), but I am having trouble understanding this part. 1. Based on my understanding, all the elements of $q$ are processes index by time $t$. So what does it mean to compute elements of $q$ or $\Sigma$? Are the authors time discretizing these processes? 2. How should I understand the role of the propagator $\Sigma$, and how the authors solve for this object first independently of the elements of $q$? 3. How should I understand the saddle point solver in Algorithm 2? Here I am just asking for what is happening within the algorithm, as I find reading the pseudocode quite difficult to understand. I would be happy to continue this discussion during the rebuttal period, and once again I would raise my score once I find these questions addressed.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

Building on past work which set up a DMFT (Dynamical Mean Field Theory) for fully connected networks in the infinite width limit (where the width of each layer tends to infinity), this paper reasons about the *fluctuations* around the infinite width limit. This is important because for finite sized neural networks, these fluctuations are large enough that they are an important part of the network training dynamics. The description of the DMFT is theoretically complicated and cannot be solved exactly, but they enable simulations which confirm that the DMFT captures the behaviour of finite sized real networks quite accurately. As an application, this DMFT theory is used to understand bias, training rates, and variance in realistic tasks and the special case of 2-layer networks (where the theory is quite a bit simpler) is investigated in more detail. This theory has the potential to open the door to many potential future uses that explain how neural networks learn.

Strengths

This paper lays out a framework for theoretical understanding of deep neural networks that incorporates the effect of finite width. This is a problem that has received a lot of attention since the original NTK/infinite width limits came out, and as far as I know, the approach here is novel and powerful. The paper is well written for the most part although familiarity with the infinite width DMFTs is assumed. Overall, the fact that the theory and the simulations agree very well is quite impressive, and I think the ideas in the paper are quite ambitious because they can be used for almost any kind of question one might have about the theoretical evolution of the DNN. This paper has the potential to be the basis for future work which uses the theory developed here to investigate questions about how DNNs learn.

Weaknesses

The main weakness of the paper is that its a bit spread thin at times: both the theory and a few different applications are covered, but it seems like the authors were trying to make it all fit and I would have liked more detail in a few spots. This is largely due to the page limit of the submission. I personally would have found it to be a stronger paper if a single really clear example was presented in a lot of detail. (Although again, I completely understand that this is largely pressure from the conference format to try and do a lot of stuff) The other main "weakness" of the paper (which is strictly speaking a limitation of the audience of the paper) is that to understand it, you need to be familiar with the previous DMFT on which this paper is built. The authors include a very short section called "Review of Dynamical Mean Field Theory" citing [9],[46] as a review, but this section is extremely sparse for actually understanding what is going on. I essentially had to read [9] in its entirety first to understand what was going on in this paper. (Also the reference [46] could not be found since only authors and title are given...where would one find this reference?) In my view, this weakness could be mitigated by just being more honest with the reader up front about this...for example [9],[46] should be cited at point 1 in the list of contributions to make it more clear the dependence and what is/isn't actually explained in this paper. Another (related) "weakness" is that the paper relies quite heavily on physics technology and jargon to reach its conclusions. The fact that the results are so heavily entrenched in physics jargon like "order parameters" or "propagator" makes this paper less likely to have a broad impact on the deep learning community. The authors would add a lot of value to the work by attempting to make a "translation guide" to help people who don't have the same physics background understand what is going on in more detail.

Questions

* Suggestion: Eqn (3) seems like a very important main result: a bit more discussion on the next order term and why its actually size 1/N (even though it starts with an N) would be helpful. * I am wondering if you had any numerical simulations where you checked the rate of covergence between the theory here and simulations (which should presumably be like 1/N^2?). Something like Figure 2a but comparing to the fluctuation predictions rather than to $q_\infty$. * One high-level question I had: It seems that if the learning rate is fixed and not scaled, the the fluctuations and the effect of the learning rate are both on the same scale $1/\sqrt{N}$. 1. Is that correct? and 2. Does your theory work to analyze what's going on in that case? Here is a list of other minor errors/suggestions I found while reading:] * Eqn (2): use definition equal with three lines to be consistent with definitions later on * Line 179 vs 185: Is there a difference between $K$ and $K_\infty$? If so what is it? * Line 188: What does the subscript 0 in "$Cov_0$" mean here? * Section 6.1: I was able to more-or-less piece this together, but I think it would be a lot more understandable if you gave explit definitions for $K(t)$ and $K_\ast(t)$ * Line 232-235: I think it would be a lot clearer to write out the definitions here of the two new $\Delta$'s in terms of $\Delta_\mu$. * References [15],[21],[33], [46], [53] has only author/title but not where published/where to find. * Reference [22] missing a title? * I would also check the arxiv only referenes e.g. [11],[13],[20],[27],[34],[57],[63] to make sure there isnt a conference or journal version that is now published.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

A potential criticism is the physics level of rigor used in Appendix C which is used to establish the main results. The manipulations carried out in the proof of Appendix C certainly seem plausible, and I believe the community as a whole is ok with this level of rigour, but the authors could be a bit more clear about what they mean by "proof" in the main paper. It is not a mathematically rigour proof (which would involve all sorts of techincal assumptions), but rather a physics-type statement that holds assuming the usual expansions can be carried out without obstructions. To reiterate: I think the actual work is fine, but they could be a bit more honest about how it is "proven" and the level of rigour in the main paper.

Summary

The paper addresses the problem of analytical description of the rich (feature learning) dynamics of neural networks. To achieve this, the authors use previously introduced dynamical mean field theory (DMFT), which identifies several key characteristics of the problem - order parameters - defines their probability distribution throughout the training using a path integral representation. The paper considers $\mu P$ parametrization of the network, which is known to display feature learning behavior at infinite width (in contrast to popular NTK and Standard parametrizations) with feature learning strength controlled by parameter $\gamma$. In the paper, the authors focus on leading width $O(N^{-1})$ corrections to the infinite width limit. Technically, this is achieved by taking into account quadratic (Gaussian) fluctuation around saddle point $\mathbf{q}_\infty$ of DMFT action $S(\mathbf{q})$. After deriving general equations governing $O(N^{-1})$ corrections, the authors consider a number of simplified scenarios where these equations can be solved analytically. Also, the authors validate qualitative conclusions of their theory in the non-synthetic experiments with CNN trained on CIFAR10.

Strengths

The approach chosen by the authors - identifying saddle point $\mathbf{q}_\infty$ of the system's action $S(\mathbf{q})$ and investigating it together with Gaussian fluctuation around $\mathbf{q}_\infty$ - is a workhorse of various branches of theoretical physics where behavior of a complex system needs to be analyzed. In physics, this approach not only provides a very sizable portion of available analytical results but, for example, also provides SOTA results for numerical modeling of realistic strongly-correlated materials [Kotliar 2004](https://pubs.aip.org/physicstoday/article-abstract/57/3/53/755526/Strongly-Correlated-Materials-Insights-From?redirectedFrom=fulltext), [Vollhardt 2019](https://arxiv.org/abs/1910.12650). Thus, realizing this general strategy is one of the fundamental directions within deep learning theory. Also, as noted by the authors, the DMFT approach is non-perturbative w.r.t. feature learning strength - a feature that is mostly absent in other approaches to NN dynamics away from kernel regime.

Weaknesses

* In many cases, solving DMFT equations analytically seems to be intractable. This significantly renders the main purpose of the proposed theory - obtaining analytical insights into the network dynamics. * The DMFT equations are bulky, which could make working with them quite exhaustive. * I believe it is hard to understand main DMFT ingredients - order parameters $\mathbf{q}$ and their action $S(\mathbf{q})$ - from the current paper alone. Most probably, a careful reading of the original DMFT paper [9] is required to understand this paper. However, this is not the author's fault but rather a consequence of the chosen approach.

Questions

* What dataset size $P$ were used for main CIFAR10 experiment of fig. 10? Due to mentioned $O(P^4T^4)$ memory requirements of storing the propagators, it seems to be impossible to work with the full ($P=50000$) CIFAR10 size. Also, what is the computational complexity of solving DMFT equations? Both for saddle point $\mathbf{q}_\infty$ and $O(N^{-1})$ corrections if they are different. * While in most of the presented plots DMFT accurately describes the experiments, for the big sample sizes (fig. 3) there is a significant disagreement between theory and experiment. Do you think this a fundamental limitation DMFT on the level of leading order correction (ignoring $O(N^{-2})$ and beyond)? Or maybe it is because in the experiment you considered whitened data, whereas in the more realistic scenarios, the data typically has a low effective dimension (e.g. measured by the decay of data covariance eigenvalues)?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

The authors discuss the limitations of their current approach. The limitations mainly come from 1) the inability to numerically solve DMFT equations for large-scale problems (e.g. due to $O(P^4 T^4)$ propagator size) and 2) the need to consider higher order expansion around the saddle point $\mathbf{q}_\infty$.

Response

I apologize for the delayed response. I have been traveling and wrestling with a heavy review load this year. I believe my concerns have been addressed, and I will raise my score to accept.