Connecting NTK and NNGP: A Unified Theoretical Framework for Neural Network Learning Dynamics in the Kernel Regime

We thank the reviewer for raising interesting questions regarding the learning dynamics. We would like to make several clarifications that may potentially address the reviewer's concern. We would like to point out that that our theory considers a small fixed constant learning rate throughout the dynamics (in the notation of the paper, the learning rate is proportional to $1/\lambda$). The two learning stages emerge not because of reduction in the learning rate but due to the difference in the magnitude of the gradient at different points in the loss landscape. Initially when loss is large the gradient term in the Langevin dynamics dominates. When the system reaches the neighborhood of the energy minimum ($\mathcal{O}(T)$ loss), the residual gradient is small, $\mathcal{O}(\sqrt{T})$, of the same size as the Langevin noise, hence the network executes a slow and noise-driven exploration of the solution space, corresponding to a roughly flat region of the loss landscape. Therefore the two phase phenomenon is not a result of adapting learning rate. It reflects the process of an over-parameterized network reaching and exploring a degenerate space of solutions.

Regarding the discretization error, the theoretical results in the paper including the two phases address the continous time stochastic ODE of Langevin dynamics hence does not incorporate discretization errors. It is true that for comparing the theory with simulations of these ODEs we used discretised time Langevin dynamics which suffers from discretization errors. We have fixed the discretization step size so that further reductions of its value did not change significantly the results. Indeed, despite the inevitable discretization step the numerics agree remarkably well with the theory as shown in Fig.1e and SI Fig.7. Relevant work on discretization error of Langevin dynamics suggests that the Wasserstein distances between the original Langevin diffusion and its discretized version decays as $\mathcal{O}(\eta)$ with $\eta$ denoting the discretization step size under appropriate smoothness conditions (Alfonsi et al, 2015); we include the relevant reference in the revised paper.

Response to questions:

1. Incorporating time dependent learning rate, as suggested by the reviewer, is in principle possible and is an interesting topic for future work.

2. Various previous works have observed that despite operating in a highly over-parameterized regime, learning tend to bias neural networks towards solutions that also generalize well. This type of 'inductive bias' comes in various forms including explicit and implicit regularizations (Bietti et al, 2019). In our paper, we mitigate the problem of overfitting by adding an explicit $L_{2}$ regularizer, as given in Eq.3 and noted below it. In the long time limit, the network explores the degenerate solution space with small ($\mathcal{O}(T)$) training error, biased towards weights with smaller $L_{2}$ norm, and thus reduces the overfitting problem, as has been shown in standard works on NTK and NNGP in the infinite width and finite data size limit. However, we agree with the reviewer that it would be very interesting to extend the current work to incorporate the more realistic 'linear' regime (where data size scales linearly with width). This is a topic of ongoing work, and is discussed in the discussion section.

We would like to thank the reviewer for the additional comment. It is true that in general there may be multiple local minima and steep slopes in the loss landscape. In our paper we focus on the infinite width limit, and previous works [1,2,3] have shown that the loss landscape of infinitely wide neural networks has nice properties which ensure smoothness around the minimum and guarantee the convergence of (S)GD algorithms with a fixed learning rate (i.e., without having to cap the gradient, etc). Furthermore, these works have shown that for over-parameterized networks, the solution space forms a manifold instead of distinct local minima, which allows us to consider our diffusive learning phase as explorations within the solution manifold.

[1] On the linearity of large non-linear models: when and why the tangent kernel is constant, Liu et al, NeurIPS 2020.

[2] Loss landscapes and optimization in over-parameterized non-linear systems and neural networks, Applied and Computational Harmonic Analysis, Liu et al, 2022.

[3] The loss landscape of over-parameterized neural networks, Cooper, 2018.

We are grateful for the reviewer's comments highlighting the unclear sections of our original paper. We agree that the presentation of Eq. 5 in the previous version of the paper may lead to confusion. In our revised version, we present the moment-generating function for the predictor statistics in Eq.5 and Eq.6, and clarify that ${\bf u}(t)$ and ${\bf v}(t)$ are Gaussian auxiliary integration variables that are integrated out. We have also clarified potential ambiguities in the notations of kernels (see response to Reviewer 2CVF). We have also clarfied the statement "the integral equation can be transformed into a linear ODE" in Section 3.1.We revised the manuscript to incorporate references to SI A, which introduces the Markov proximal learning framework to formulate our theory of Langevin dynamics, and the details of the theoretical calculations is in SI B1-B4. With these changes we believe the statements and derivations of our results are clarified.

Response to questions:

1. We define optimal early stopping time by the time the network reaches the minimal generalization error in the entire learning trajectory. In Fig. 3, we plot $\Delta t$, defined as the difference between optimal early stopping time and the time when the network first reaches the NNGP equilibrium. We have included a more careful definition in Section 4.2 “The role of initialization and regularization and early stopping phenomena”, and the caption of Figure 3.

2. Replica: The reviewer is right that in most applications of replica method in statistical mechanics of learning and memory, it is required in the linear regime where the data size and network width are proportional to each other. Relatedly, the use of replica in these works is used to average over the data statistics (assumed to obey simple statistics, e.g., iid gaussian). Our use of the replica method is very different. We do not make assumptions about the statistics of the data at all, hence replica method is not required for dealing with the data. In contrast, we apply the replica method to handle the normalization factor $Z(\theta_{t-1})$ in the transition matrix $\mathcal{T}(\theta_{t}|\theta_{t-1})$ (see details in SI B1), analogous to the application of the replica method in calculation of the Franz-Parisi potential (Franz et al, 1992). This is needed even in the infinite width, finite data limit, in our Markov chain formalism. As the reviewer commented, for this paper we focus on the regime where the network width N tends to infinity, while the amount of data remains finite, consistent with the typical NTK/NNGP regime.

3. In our revised paper, we add some further clarification for the replica method in SI Eq. 28.

We thank the reviewer for questions regarding the clarity of our paper. We have modified the paper to clarify the notations, equations and references to the appendix.

Response to questions:

1. We acknowledge your concern. In our revision, we will carefully address this issue and make necessary adjustments to reduce the use of first person plural and emphasize the significance of the findings.

2 and 3. We appreciate the feedback and we agree that Eq. 5 may lead to confusion in its previous version. In our revised version, we present the moment-generating function (MGF) for the statistics of the predictor $f({\bf x},t)$ in Eq.5 and Eq.6, and ${\bf u}(t)$ and ${\bf v}(t)$ are auxiliary integration variables which are integrated out. Eq. 5 is important because the predictor statistics, and in particular the mean predictor can be derived from it by evaluating its derivative w.r.t. the field $\ell(t)$ at $\ell(t)=0$. We have modified the paper and incorporate necessary explanations of Eq.5 to highlight its importance and its connection to the kernel functions and predictor statistics.

4. The typo in Eq. 9 of the original paper has been fixed, it is now Eq.10 in the revised paper. We apologize for the confusion of multiple notations for the kernel function applied on training/testing data. In the revised version, we clarify the kernel notations by using bold capital letters ${\bf K}(t,t')$ to denote $P\times P$ kernel matrices, ${\bf k}(t,t')$ to denote $P$ dimensional kernel vector, and $K(t,t',\boldsymbol{x},\boldsymbol{x}')$ to denote functions of two arbitrary inputs ${\bf x}$ and ${\bf x}'$. We also denote by ${\bf \tilde{K}}(t,t')$ $(P+1)\times(P+1)$ or $(P+1)\times P$ matrices (where the $P+1$ index refers to the test point). These notations are detailed at the end of Section 2.2 and Section 2.4.

5. The derivative kernel is defined as Eq.13 in our revised paper.

6. Eq. 17 is the limit of Eq.12 which relates the NDK to the NNGP kernel. It results from the definition of the NDK and not from the equation for the predictor. This non-trivial identity is proven in SI Sec. C4.

7. NDK with $t=t'$ has the form of a time-dependent averaged NTK(Eq.7). In the infinite width limit, the NTK theory states that the learning dynamics is govened by the NTK at initiliaztation (Jacot et al, 2018), hence does not have time dependence. Indeed, as we show in Eq.16, in the relevant time scale appropriate for the NTK theory, NDK with $t=t'$ is independent of time. However, for longer time scales NDK with $t=t'$ is not equivalent to NTK (Eq.12), specifically, our NDK is explicitly time dependent since we averaged over the weights $\Theta(t)$, while the NTK implicitly depends on time through $\Theta(t)$. Also, as shown in Eq. 15, the predictor dynamics depend on the NDK with general $t$ and $t'$.

8. We already show the NNGP equilibrium in Fig. 2, as shown by the red dashed lines.

9. We appreciate the reviewer for pointing out the unclear aspects of the definition of the optimal early stopping point, we have correspondingly modified the relevant paragraph in section 4.2 for clarity. The optimal early stopping point for a set of test data points, is defined as the time where generalization error on this set achieves its minimal value (as a function of time).

10. We will increase the font of labels, legends and titles of figures in our finalized version.

11. For linear activation the expressions for the kernels are straightforward as shown in SI Eq. 78-80. For ReLU activation, we are refering to Eq. 12 and 13 in Cho and Saul 2009, with $n=1$, derived in SI A of the paper, and also shown in SI Eq. 81-85 in our paper. For error function, we refer to Eq. 11 in Williams 1996, and also shown in SI Eq. 86-87 in our paper.

12. We have included a paragraph at the beginning of SI Sec. B outlining the content of this section.

13. It is unclear what the reviewer is asking, but we have changed the sentence to “predictor converges fast to the target output” for clarification.

14. We agree that this sentence is confusing. We have changed it to “our theory for the Langevin dynamics suggests a possible mechanism of representational drift”.

15. We thank the reviewer for noticing this, we have fixed the margin violation problem.

I have reviewed the author response. However, I do not believe the anticipated revisions will be enough to satisfy the substance of my concerns.

We appreciate the reviewer for pointing out the several interesting and relevant works, we have modified the introduction part of the paper to recognize and connect to existing contributions, and included relevant references as the reviewer suggests. With respect to previous single-time NTK kernel (now cited in the Introduction) we note that the time dependence of NTK is not explicit but arises through a weight dependent kernel (i.e., before any averaging). In contrast, our NDK has an explicit time dependence that arises when the weights are averaged. As we show, the full dynamics requires averaging over two sets of weights (i.e., at two time points) hence two times.

Response to questions:

1. Multiple outputs: Thanks for asking. As we comment now in discussion section of the paper, for $M$ outputs where $M\sim\mathcal{O}(1)$, one can simply replace the $P$ dimensional vector ${\bf y}$ by a $P\times M$ dimensional matrix in Eq.14 and 15.

2&3. We thank the reviewer for these interesting questions regarding extensions of our current results. $L_{1}$ regularization is challenging as the integration over the weight priors is harder. Effects of batch norm and dropout can in principle be incorporated by adding state dependent noise (Yang, 2019; Schoenholz et al.,2017). These interesting extensions are topics for future work.

4. The transition from phase 1 to phase 2 happens after the network reaches the NTK equilibrium. As Eq.16 (revised) and previous work (Jacot et al, 2018) show, the convergence time of the NTK dynamics depends on the spectrum of the NTK, which is determined by the data structure, number of training examples, network nonlinearities, and network depth. This allows us to evaluate where the transition happens. However, the early stopping time, defined as the time where the network reaches minimum generalization error across the entire learning trajectory is more complicated. As we show in Fig.2c and Fig.3, it may occur after the transition to phase 2, and we do not have a simple way of determining the early stopping point without numerically solving Eq.14 and 15.

I thank the authors for further clarifications. I just noticed that in the revised manuscript under the contributions sections, there are two no 4's, please fix that. Also, I'd like to stress a previous comment that I made, regarding the connection between the diffusive properties of Langevin dynamics and representational drift that was discovered in previous work (check the two papers I cited under "Strengths". If you agree with this, please fix the point 4 contribution and rewrite it as reaffirming the existing connection as opposed to being the first ones to make this connection.

I maintain my score and I believe the work deserves to be published.

We thank the reviewer for stressing this point. We have made the changes, and added a point 5 contribution to discuss relation with previous works. We added reference to the relevant literature, and phrased it as "Our work complements these previous findings by providing theoretical analysis of the two phases under Langevin dynamics, reaffirming the connections between the diffusive learning stage and representational drift in neuroscience as established in previous works".