Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks

Introduction

Deep learning has emerged as a successful paradigm for solving challenging machine learning and computational problems across a variety of domains . However, theoretical understanding of the training and generalization of modern deep learning methods lags behind current practice. Ideally, a theory of deep learning would be analytically tractable, efficiently computable, capable of predicting network performance and internal features that the network learns, and interpretable through a reduced description involving desirably initialization-independent quantities.

Several recent theoretical advances have fruitfully considered the idealization of wide neural networks, where the number of hidden units in each layer is taken to be large. Under certain parameterization, Bayesian neural networks and gradient descent trained networks converge to gaussian processes (NNGPs) and neural tangent kernel (NTK) machines in their respective infinite-width limits. These limits provide both analytic tractability as well as detailed training and generalization analysis . However, in this limit, with these parameterizations, data representations are fixed and do not adapt to data, termed the lazy regime of NN training, to contrast it from the rich regime where NNs significantly alter their internal features while fitting the data . The fact that the representation of data is fixed renders these kernel-based theories incapable of explaining feature learning, an ingredient which is crucial to the success of deep learning in practice . Thus, alternative theories capable of modeling feature learning dynamics are needed.

Recently developed alternative parameterizations such as the mean field and the $\mu P$ parameterizations allow feature learning in infinite-width NNs trained with gradient descent. Using the Tensor Programs framework, Yang & Hu identified a stochastic process that describes the evolution of preactivation features in infinite-width $\mu P$ NNs . In this work, we study an equivalent parameterization to $\mu P$ with self-consistent dynamical mean field theory (DMFT) and recover the stochastic process description of infinite NNs using this alternative technique. In the same large width scaling, we include a scalar parameter $\gamma_{0}$ that allows smooth interpolation between lazy and rich behavior . We provide a new computational procedure to sample this stochastic process and demonstrate its predictive power for wide NNs.

Our novel contributions in this paper are the following:

We develop a path integral formulation of gradient flow dynamics in infinite-width networks in the feature learning regime. Our parameterization includes a scalar parameter $\gamma_{0}$ to allow interpolation between rich and lazy regimes and comparison to perturbative methods.

Using a stationary action argument, we identify a set of saddle point equations that the kernels satisfy at infinite-width, relating the stochastic processes that define hidden activation evolution to the kernels and vice versa. We show that our saddle point equations recover at $\gamma_{0}=1$, from an alternative method, the same stochastic process obtained previously with Tensor Programs .

We develop a polynomial-time numerical procedure to solve the saddle point equations for deep networks. In numerical experiments, we demonstrate that solutions to these self-consistency equations are predictive of network training at a variety of feature learning strengths, widths and depths. We provide comparisons of our theory to various approximate methods, such as perturbation theory.

A natural extension to the lazy NTK/NNGP limit that allows the study of feature learning is to calculate finite width corrections to the infinite-width limit. Finite width corrections to Bayesian inference in wide networks have been obtained with various perturbative and self-consistent techniques . In the gradient descent based setting, leading order corrections to the NTK dynamics have been analyzed to study finite width effects . These methods give approximate corrections which are accurate provided the strength of feature learning is small. In very rich feature learning regimes, however, the leading order corrections can give incorrect predictions .

Another approach to study feature learning is to alter NN parameterization in gradient-based learning to allow significant feature evolution even at infinite-width, the mean field limit . Works on mean field NNs have yielded formal loss convergence results and shown equivalences of gradient flow dynamics to a partial differential equation (PDE) .

Our results are most closely related to a set of recent works which studied infinite-width NNs trained with gradient descent (GD) using the Tensor Programs (TP) framework . We show that our discrete time field theory at unit feature learning strength $\gamma_{0}=1$ recovers the stochastic process which was derived from TP. The stochastic process derived from TP has provided insights into practical issues in NN training such as hyper-parameter search . Computing the exact infinite-width limit of GD has exponential time requirements , which we show can be circumvented with an alternating sampling procedure. A projected variant of GD training has provided an infinite-width theory that could be scaled to realistic datasets like CIFAR-10 . Inspired by Chizat and Bach’s work on mechanisms of lazy and rich training , our theory interpolates between lazy and rich behavior in the mean field limit for varying $\gamma_{0}$ and allows comparison of DMFT to perturbative analysis near small $\gamma_{0}$. Further, our derivation of a DMFT action allows the possibility of pursuing finite width effects.

Our theory is inspired by self-consistent dynamical mean field theory (DMFT) from statistical physics . This framework has been utilized in the theory of random recurrent networks , tensor PCA , phase retrieval , and high-dimensional linear classifiers , but has yet to be developed for deep feature learning. By developing a self-consistent DMFT of deep NNs, we gain insight into how features evolve in the rich regime of network training, while retaining many pleasant analytic properties of the infinite-width limit.

Problem Setup and Definitions

where we define base cases $G_{\mu\alpha}^{L+1}(t,s)=1,\Phi^{0}_{\mu\alpha}(t,s)=K^{x}_{\mu\alpha}=\frac{1}{D}\bm{x}_{\mu}\cdot\bm{x}_{\alpha}$. We note that the above formula holds for any data point $\mu$ which may or may not be in the set of $P$ training examples. The above expressions demonstrate that knowledge of the temporal trajectory of the NTK on the $t=s$ diagonal gives the temporal trajectory of the network predictions $f_{\mu}(t)$.

Following prior works on infinite-width networks , we study the mean field limit

As we demonstrate in the Appendix D and N, this is the only $N$-scaling which allows feature learning as $N\to\infty$. The $\gamma_{0}=0$ limit recovers the static NTK limit . We discuss other scalings and parameterizations in Appendix N, relating our work to the $\mu P$-parameterization and TP analysis of , showing they have identical feature dynamics in the infinite-width limit. We also analyze the effect of different hidden layer widths and initialization variances in the Appendix D.8. We focus on equal widths and NTK parameterization (as in eq. (1)) in the main text to reduce complexity.

Self-consistent DMFT

Next, we derive our self-consistent DMFT in a limit where $t,P=\mathcal{O}_{N}(1)$. Our goal is to build a description of training dynamics purely based on representations, and independent of weights. Studying feature learning at infinite-width enjoys several analytical properties:

The next section derives these facts from a path-integral formulation of gradient flow dynamics.

We let $Z$ represent the moment generating functional (MGF) for these stochastic fields

2 Deriving the DMFT Equations from the Path Integral Saddle Point

where $\left<\right>$ denotes an average over the stochastic process induced by $\mathcal{Z}$, which is defined below

This theory is easily extended to more general architectures such as networks with varying widths by layer (App. D.8), trainable bias parameter (App. H), multiple (but $\mathcal{O}_{N}(1)$) output channels (App. I), convolutional architectures (App. G), networks trained with weight decay (App. J), Langevin sampling (App. K) and momentum (App. L), discrete time training (App. M). In Appendix N, we discuss parameterizations which give equivalent feature and predictor dynamics and show our derived stochastic process is equivalent to the $\mu P$ scheme of Yang & Hu .

Solving the Self-Consistent DMFT

Examples of the predictions obtained by solving these systems of equations are provided in Figure 2. We see that these DMFT equations describe kernel evolution for networks of a variety of depths and that the change in each layer’s kernel increases with the depth of the network.

2 Feature Learning with L2 Regularization

As we show in Appendix J, the DMFT can be extended to networks trained with weight decay $\frac{d\bm{\theta}}{dt}=-\gamma^{2}\nabla_{\bm{\theta}}\mathcal{L}-\lambda\bm{\theta}$. If neural network is homogenous in its parameters so that $f(c\bm{\theta})=c^{\kappa}f(\bm{\theta})$ (examples include networks with linear, ReLU, quadratic activations), then the final network predictor is a kernel regressor with the final NTK $\lim_{t\to\infty}f(\bm{x},t)=\bm{k}(\bm{x})^{\top}[\bm{K}+\lambda\kappa\mathbf{I}]^{-1}\bm{y}$ where $K(\bm{x},\bm{x}^{\prime})$ is the final-NTK, $[\bm{k}(\bm{x})]_{\mu}=K(\bm{x},\bm{x}_{\mu})$ and $[\bm{K}]_{\mu\alpha}=K(\bm{x}_{\mu},\bm{x}_{\alpha})$. We note that the effective regularization $\lambda\kappa$ increases with depth $L$. In NTK parameterization, weight decay in infinite width homogenous networks gives a trivial fixed point $K(\bm{x},\bm{x}^{\prime})\to 0$ and consequently a zero predictor $f\to 0$ . However, as we show in Figure 3, increasing feature learning $\gamma_{0}$ can prevent convergence to the trivial fixed point, allowing a non-zero fixed point for $K,f$ even at infinite width. The kernel and function dynamics can be predicted with DMFT. The fixed point is a nontrivial function of the hyperparameters $\lambda,\kappa,L,\gamma_{0}$.

Approximation Schemes

We now compare our exact DMFT with approximations of prior works, providing an explanation of when these approximations give accurate predictions and when they break down.

In the $\gamma_{0}\to 0$ limit, we recover static kernels, giving linear dynamics identical to the NTK limit . Corrections to this lazy limit can be extracted at small but finite $\gamma_{0}$. This is conceptually similar to recent works which consider perturbation series for the NTK in powers of $1/N$ (though not identical, see Appendix P.7 for finite $N$ effects). We expand all observables $q(\gamma_{0})$ in a power series in $\gamma_{0}$, giving $q(\gamma_{0})=q^{(0)}+\gamma_{0}q^{(1)}+\gamma_{0}^{2}q^{(2)}+...$ and compute corrections up to $\mathcal{O}(\gamma_{0}^{2})$. We show that the $\mathcal{O}(\gamma_{0})$ and $\mathcal{O}(\gamma_{0}^{3})$ corrections to kernels vanish, giving leading order expansions of the form $\bm{\Phi}=\bm{\Phi}^{0}+\gamma_{0}^{2}\bm{\Phi}^{2}+\mathcal{O}(\gamma_{0}^{4})$ and $\bm{G}=\bm{G}^{0}+\gamma_{0}^{2}\bm{G}^{2}+\mathcal{O}(\gamma_{0}^{4})$ (see Appendix P.2). Further, we show that the NTK has relative change at leading order which scales linearly with depth $|\Delta K^{NTK}|/|K^{NTK,0}|\sim\mathcal{O}_{\gamma_{0},L}(L\gamma_{0}^{2})=\mathcal{O}_{N,\gamma,L}(\frac{\gamma^{2}L}{N})$, which is consistent with finite width effective field theory at $\gamma=\mathcal{O}_{N}(1)$ (Appendix P.6). Further, at the leading order correction, all temporal dependencies are controlled by $P(P+1)$ functions $v_{\alpha}(t)=\int_{0}^{t}ds\Delta^{0}_{\alpha}(s)$ and $v_{\alpha\beta}(t)=\int_{0}^{t}ds\Delta^{0}_{\alpha}(s)\int_{0}^{s}ds^{\prime}\Delta^{0}_{\beta}(s^{\prime})$, which is consistent with those derived for finite width NNs using a truncation of the Neural Tangent Hierarchy . To lighten notation, we focus our main text comparison of our non-perturbative DMFT to perturbation theory in the deep linear case. Full perturbation theory is in Appendix P.2.

Using the timescales derived in the previous section, we find that the leading order correction to the kernels in infinite-width deep linear network have the form

We see that the relative change in the NTK $|\bm{K}^{NTK}-\bm{K}^{NTK}(0)|/|\bm{K}^{NTK}(0)|\sim\mathcal{O}(\gamma_{0}^{2}L)=\mathcal{O}(\gamma^{2}L/N)$, so that large depth $L$ networks exhibit more significant kernel evolution, which agrees with other perturbative studies as well as the non-perturbative results in Figure 2. However at large $\gamma_{0}$ and large $L$, this theory begins to break down as we show in Figure 4.

Our DMFT suggests that for networks sufficiently wide for their kernels to concentrate, the dynamics of loss and kernels should be invariant under the rescaling $N\to RN,\gamma\to\gamma/\sqrt{R}$, which keeps $\gamma_{0}$ fixed. To evaluate how well this idea holds in a realistic deep learning problem, we trained CNNs of varying channel counts $N$ on two-class CIFAR classification . We tracked the dynamics of the loss and the last layer $\Phi^{L}$ kernel. The results are provided in Figure 5. We see that dynamics are largely independent of rescaling as predicted. Further, as expected, larger $\gamma_{0}$ leads to larger changes in kernel norm and faster alignment to the target function $y$, as was also found in . Consequently, the higher $\gamma_{0}$ networks train more rapidly. The trend is consistent for width $N=250$ and $N=500$. More details about the experiment can be found in Appendix C.2.

Discussion

We provided a unifying DMFT derivation of feature dynamics in infinite networks trained with gradient based optimization. Our theory interpolates between lazy infinite-width behavior of a static NTK in $\gamma_{0}\to 0$ and rich feature learning. At $\gamma_{0}=1$, our DMFT construction agrees with the stochastic process derived previously with the Tensor Programs framework . Our saddle point equations give self-consistency conditions which relate the stochastic fields to the kernels. These equations are exactly solveable in deep linear networks and can be efficiently solved with a numerical method in the nonlinear case. Comparisons with other approximation schemes show that DMFT can be accurate at a much wider range of $\gamma_{0}$. We believe our framework could be a useful perspective for future theoretical analyses of feature learning and generalization in wide networks.

Though our DMFT is quite general in regards to the data and architecture, the technique is not entirely rigorous and relies on heuristic physics techniques. Our theory holds in the $T,P=\mathcal{O}_{N}(1)$ and may break down otherwise; other asymptotic regimes (such as $P/N,T/\log(N)=\mathcal{O}_{N}(1)$, etc) may exhibit phenomena relevant to deep learning practice . The computational requirements of our method, while smaller than the exponential time complexity for exact solution , are still significant for large $PT$. In Table 1, we compare the time taken for various theories to compute the feature kernels throughout $T$ steps of gradient descent. For a width $N$ network, computation of each forward pass on all $P$ data points takes $\mathcal{O}(PN^{2})$ computations. The static NTK requires computation of $\mathcal{O}(P^{2})$ entries in the kernel which do not need to be recomputed. However, the DMFT requires matrix multiplications on $PT\times PT$ matrices giving a $\mathcal{O}(P^{3}T^{3})$ time scaling. Future work could aim to improve the computational overhead of the algorithm, by considering data averaged theories or one pass SGD . Alternative projected versions of gradient descent have also enabled much better computational scaling in evaluation of the theoretical predictions , allowing evaluation on full CIFAR-10.

Acknowledgments and Disclosure of Funding

This work was supported by NSF grant DMS-2134157 and an award from the Harvard Data Science Initiative Competitive Research Fund. BB acknowledges additional support from the NSF-Simons Center for Mathematical and Statistical Analysis of Biology at Harvard (award #1764269) and the Harvard Q-Bio Initiative.

BB thanks Jacob Zavatone-Veth, Alex Atanasov, Abdulkadir Canatar, and Ben Ruben for comments on this manuscript as well as Greg Yang, Boris Hanin, Yasaman Bahri, and Jascha Sohl-Dickstein for useful discussions.

References

Checklist

Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] As described in the abstract and introduction, we provide a dynamical field theory of deep networks based on kernel evolution.

Did you describe the limitations of your work? [Yes] We have an explicit limitations as the last paragraph of the paper in Section 7.

Did you discuss any potential negative societal impacts of your work? [N/A] This work is theoretical and is very unlikely to present negative social impacts.

Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

If you are including theoretical results…

Did you state the full set of assumptions of all theoretical results? [Yes] We describe that our theory holds for NN architectures in the infinite-width $N\to\infty$ limit.

Did you include complete proofs of all theoretical results? [Yes] All claims made in the main text are supported by derivations in the Appendix.

Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] Code to reproduce experimental results is provided in the supplementary material.

Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] We provide details of all experiments in C.

Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] We provided errorbars in the alignment scores of DMFT as a function of width $N$ in Figure 1. All other runs were over a single wide network, where performance is predicted to concentrate over initialization.

Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We mention our GPU usage in C.2.

If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

If your work uses existing assets, did you cite the creators? [Yes] We cited the creators of Jax, Neural Tangents, and CIFAR-10.

Did you mention the license of the assets? [N/A] These are all open source provided they are appropriately credited in academic research.

Did you include any new assets either in the supplemental material or as a URL? [N/A]

Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

If you used crowdsourcing or conducted research with human subjects…

Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]

Appendix

Appendix A Additional Figures

Appendix B Algorithmic Implementation

The parameter $\beta$ controls recency weighting of the samples obtained at each iteration. If $\beta=1$, then the rank of the kernel estimates is limited to the number of samples $\mathcal{S}$ used in a single iteration, but with $\beta<1$ smaller sample sizes $\mathcal{S}$ can be used to still obtain accurate results. We used $\beta=0.6$ in our deep network experiments. Convergence is usually achieved in around $\sim 15$ steps for a depth 4 ($L=3$ hidden layer) network such as the one in Figure 1 and 7.

Appendix C Experimental Details

All NN training was performed with Jax gradient descent optimizer with fixed learning rate.

C.2 CNN Experiments on CIFAR-10

We define a depth $L$ CNN model with ReLU activations and stride $1$, which is implemented as a pytree of parameters in JAX . We apply global average pooling in the final layer before a dense readout layer. The code to initialize and evaluate the model is provided below.

Appendix D Derivation of Self-Consistent Dynamical Field Theory

In this section, we introduce the dynamical field theory setup and saddle point equations. The path integral theory we develop is based on the Martin-Siggia-Rose-De Dominicis-Janssen (MSRDJ) framework , of which a useful review for random recurent networks can be found here . Similar computations can be found in recent works which consider typical behavior in high dimensional classification on random data .

As discussed in the main text, we consider the following wide network architecture parameterzied by trainable weights $\bm{\theta}=\text{Vec}\{\bm{W}^{0},\bm{W}^{1},...\bm{w}^{L}\}$, giving network output $f_{\mu}$ defined as

Since $K_{NTK}$ is $O_{\gamma}(1)$ at initialization, it is clear that to have $O_{\gamma}(1)$ evolution of the network output at initialization we need $\eta=\gamma^{2}$. With this scaling, we have the following

which allows us to compute the NTK in terms of these features

where $K^{x}_{\mu\alpha}=\frac{\bm{x}_{\mu}\cdot\bm{x}_{\alpha}}{D}$ is the input Grammian. We see that the NTK can be built out of the following primitive kernels

We observe that the dynamics of the hidden features is controlled by the factor $\frac{\gamma}{\sqrt{N}}$. If $\gamma=O_{N}(1)$ then we recover static NTK in the limit as $N\to\infty$. However, if $\gamma=O_{N}(\sqrt{N})$ then we obtain $O_{N}(1)$ evolution of our features and we reach a new rich regime. We choose the scaling $\gamma=\gamma_{0}\sqrt{N}$ for our field theory so that $\gamma_{0}>0$ will give a feature learning network.

D.2 Warmup: DMFT for One Hidden Layer NN

where the network predictions evolve as $\frac{\partial}{\partial t}f_{\mu}(t)=\sum_{\alpha}[\Phi_{\mu\alpha}(t,t)+G_{\mu\alpha}(t,t)K^{x}_{\mu\alpha}]\Delta_{\alpha}(t)$ for kernels $\Phi_{\mu\alpha}(t,t)=\frac{1}{N}\phi(\bm{h}_{\mu}(t))\cdot\phi(\bm{h}_{\alpha}(t))$ and $G_{\mu\alpha}(t,t)=\frac{1}{N}\bm{g}_{\mu}(t)\cdot\bm{g}_{\alpha}(t)$. At finite $N$, the kernels $\Phi,G$ will depend on the random initial conditions $\bm{\chi},\bm{\xi}$, leading to a predictor $f_{\mu}$ which varies over initializations. If we can establish that the kernels $\Phi,G$ concentrate at infinite-width $N\to\infty$, then $\Delta_{\mu}$ are deterministic. We now study the moment generating function for the fields

To perform the average over $\bm{\theta}_{0}=\{\bm{W}^{0}(0),\bm{w}^{1}(0)\}$, we enforce the definition of $\bm{\chi}_{\mu},\bm{\xi}$ with delta functions

Though this step may seem redundant in this example, it will be very helpful in the deep network case, so we pursue it for illustration. After mulitplying by these factors of unity and performing the Gaussian integrals, we obtain

We now aim enforce the definitions of the kernel order parameters with delta functions

where the fields $\bm{h}_{\mu}(t),\bm{g}_{\mu}(t)$ are regarded as functions of $\{\bm{\chi}_{\mu}\}_{\mu},\bm{\xi}$ (see Equation (D.2)) and the $\hat{\Phi},\hat{G}$ integrals run over the imaginary axis $(-i\infty,i\infty)$. After this step, we can write

where the DMFT action $S[\Phi,\hat{\Phi},G,\hat{G}]$ is $\mathcal{O}_{N}(1)$ and has the form

The single site moment generating function $\mathcal{Z}[j,v]$ arises from the factorization of the integrals over $N$ different fields in the hidden layer and takes the form

where the $i$-th single site average $\left<\right>_{i}$ of an observable $O(\chi,\hat{\chi},\xi,\hat{\xi})$ is defined as

Since $\hat{\Phi}=\hat{G}=0$ the single site MGF reveals that the initial fields are independent Gaussians $\{\chi_{\mu}\}\sim\mathcal{N}(0,\bm{K}^{x})$ and $\xi\sim\mathcal{N}(0,1)$. At zero source $\bm{j},\bm{v}\to 0$, all single site averages $\left<\right>_{i}$ are equivalent and we may merely write $\Phi_{\mu\alpha}(t,s)=\left<\phi(h_{\mu}(t))\phi(h_{\alpha}(s))\right>\ ,\ G_{\mu\alpha}(t,s)=\left<g_{\mu}(t)g_{\alpha}(s)\right>$, where $\left<\right>$ is the average over the single site distributions for $\bm{j},\bm{v}\to 0$.

Putting all of the saddle point equations together, we arrive at the following DMFT

We see that for $L=1$ networks, it suffices to solve for the kernels on the time-time diagonal. Further in this two layer case $\chi,\xi$ are independent and do not vary in time. These facts will not hold in general for $L\geq 2$ networks, which requires a more intricate analysis as we show in the next section.

D.3 Path Integral Formulation for Deep Networks

Moments of these stochastic fields can be computed through differentiation of $Z$ near zero-source

D.4 Order Parameters and Action Definition

We define the following order parameters which we will show concentrate in the $N\to\infty$ limit

After introducing these order parameters into the definition of the partition function, we have a factorization of the integrals over each of the $N$ sites in each hidden layer. This gives the following partition function

We thus see that the action $S$ consists of inner-products between order parameters $\{\Phi,G,A\}$ and their duals $\{\hat{\Phi},\hat{G},B\}$ as well as a single site MGF $\mathcal{Z}[\{\Phi,\hat{\Phi},G,\hat{G},A,B,j,v\}]$, which is defined as

D.5 Saddle Point Equations

Since the integrand in the moment generating function $Z$ takes the form $e^{NS[\{\Phi,\hat{\Phi},G,\hat{G},A,B\}]}$, the $N\to\infty$ limit can be obtained from saddle point integration, also known as the method of steepest descent . This consists in finding order parameters $\{\Phi,\hat{\Phi},G,\hat{G},A,B\}$ which render the action $S$ locally stationary. Concretely, this leads to the following saddle point equations.

We use the notation $\left<\right>$ to denote an average over the self-consistent distribution on fields induced by the single-site moment generating function $\mathcal{Z}$ at the saddle point. Concretely if $\mathcal{Z}=\int d\chi d\xi d\hat{\chi}d\hat{\xi}\exp\left(-\mathcal{H}[\chi,\xi,\hat{\chi},\hat{\xi}]\right)$ then the single-site self-consistent average of observable $O([\chi,\xi,\hat{\chi},\hat{\xi}])$ is defined as

As we will demonstrate in the next subsection, these correlators must vanish. Lastly, we can calculate the remaining correlators in terms of primal variables

D.6 Single Site Stochastic Process: Hubbard Trick

To get a better sense of this distribution, we can now simplify the quadratic forms appearing in $\mathcal{Z}$ using the Hubbard trick , which merely relates a Gaussian function to its Fourier transform.

Applying this to the quadratic forms in the single-site MGF $\mathcal{Z}$, we get

D.7 Final DMFT Equations

D.8 Varying Network Widths and Initialization Scales

Using our definition of the kernels and the $\bm{h},\bm{z}$ fields

We also find the usual formula for the NTK

Appendix E Two Layer Networks

In a two layer network, there are no $\bm{A}$ or $\bm{B}$ order parameters, so the fields $\chi^{1}$ and $\xi^{1}$ are always independent. Further, $\chi^{1}$ and $\xi^{1}$ are both constant throughout training dynamics. Thus we can obtain differential rather than integral equations for the stochastic fields $h^{1},z^{1}$ which are

where the average is taken over the random initial conditions $\bm{h}^{1}(0)\sim\mathcal{N}(0,\bm{K}^{x})$ and $\bm{z}^{1}(0)\sim\mathcal{N}(0,\bm{1}\bm{1}^{\top})$. An example of the two layer theory for a ReLU network can be found in Appendix Figure 6. In this two layer setting, a drift PDE can be obtained for the joint density of preactivations and feedback fields $p(\bm{h},z;t)$

which is a zero-diffusion feature space version of the PDE derived in the original two layer mean field limit of neural networks .

Appendix F Deep Linear Networks

Using the formulas which define the fields, we have

The saddle point equations can thus be written as

The kernels $\bm{H}(t)=\left<\bm{h}(t)\bm{h}(t)^{\top}\right>$ and $G(t)=\left<g(t)^{2}\right>$ thus evolve as

where the initial conditions are $\bm{H}(0)=\bm{I}$, $G(0)=1$ and $\bm{\Delta}(0)=\bm{y}$. These equations hold for any choice of data $\bm{K}^{x},\bm{y}$.

For input data which is whitened where $\bm{K}^{x}=\mathbf{I}$, then the dynamics can be simplified even further, recovering the sigmoidal curves very similar to those obtained under a special initialization . In this case we note that the error signal always evolves in the $\bm{y}$ direction, $\bm{\Delta}(t)=\Delta(t)\frac{\bm{y}}{|\bm{y}|}$, and that $\bm{H}$ only evolves in a rank one direction $\bm{y}\bm{y}^{\top}$ direction as well. Let $\frac{1}{|\bm{y}|^{2}}\bm{y}^{\top}\bm{H}(t)\bm{y}=H_{y}(t)$. Let $y=|\bm{y}|$ represent the norm of the target vector, then the relevant scalar dynamics are

We note that this system obeys a conservation law which constrains $(H_{y},y-\Delta)$ to a hyperbola

This conservation law implies that $H_{y}(0)^{2}=1=\lim_{t\to\infty}H_{y}(t)^{2}-\gamma_{0}^{2}y^{2}$ or that the final kernel has the form $\lim_{t\to\infty}\bm{H}(t)=\frac{1}{y^{2}}[\sqrt{1+\gamma_{0}^{2}y^{2}}-1]\bm{y}\bm{y}^{\top}+\mathbf{I}$. The result that the final kernel becomes a rank one spike in the direction of the target function was also obtained in finite width networks in the limit of small initialization and also from a normative toy model of feature learning . We can use the conservation law above $1=H_{y}(t)^{2}-\gamma_{0}^{2}(\Delta(t)-y)^{2}$ to simplify the dynamics to a one dimensional system

where $f=y-\Delta$. We see that increasing $\gamma_{0}$ provides strict acceleration in the learning dynamics, illustrating the training benefits of feature evolution. Since this system is separable, we can solve for the time it takes for the network output norm to reach output level $f$

The NTK limit can be obtained by taking $\gamma_{0}\to 0$ which gives

which recovers the usual convergence rate of a linear model. The right hand side of Equation (F.1.1) has a perturbation series in $\gamma_{0}^{2}$ which converges in the disk $\gamma_{0}<\frac{1}{y}$. The other limit of interest is the $\gamma_{0}\to\infty$ limit where

which recovers the logistic growth observed in the initialization scheme of prior works . The timescale $\tau$ required to learn is only $\tau\sim\frac{1}{\gamma_{0}}\ll 1$, which is much smaller than the $O_{\gamma_{0}}(1)$ time to learn predicted from the small $\gamma_{0}$ expansion. We note that the above leading order asymptotic behavior at large $\gamma_{0}$ considers the DMFT initial condition $\Delta(0)=y$ as an unstable fixed point. For realistic learning curves, one would need to stipulate some alternative initial condition such as $\Delta=y-\epsilon$ for some small $\epsilon>0$ in order to have nontrivial leading order dynamics.

F.2 Deep Linear Whitened Data

Lastly we have the simple evolution equation for the scalar error $\Delta(t)$

Appendix G Convolutional Networks with Infinite Channels

The feature space description of the forward and backward pass relations is

These two order parameters per layer collectively define the neural tangent kernel. Following the computation in D, we obtain the following field theory in the $N\to\infty$ limit:

We see that this field theory essentially multiples the number of sample indices by the number of spatial indices $P\to P|\mathcal{S}|$. Thus the time complexity of evaluation of this theory scales very poorly as $\mathcal{O}(P^{3}|\mathcal{S}|^{3}T^{3})$, rendering DMFT solutions very computationally intensive.

Appendix H Trainable Bias Parameter

Appendix I Multiple Output Channels

We can now compute the NTK for samples $\mu,\alpha$

From these fields, the saddle point equations define the kernels as

This allows studying the multi-class structure of learned representations.

Appendix J Weight Decay in Deep Homogenous Networks

If we train with weight decay, $\frac{d}{dt}\bm{\theta}=-\gamma^{2}\nabla_{\bm{\theta}}\mathcal{L}-\lambda\bm{\theta}$, in a $\kappa$-degree homogenous network ($f(c\bm{\theta})=c^{\kappa}f(\bm{\theta})$), then the prediction dynamics satisfy

This holds by the following identity $\frac{\partial}{\partial c}f(c\bm{\theta})=\frac{\partial}{\partial c}c^{\kappa}f(\bm{\theta})$, which when evaluated at $c=1$ gives $\frac{\partial}{\partial\bm{\theta}}f(\bm{\theta})\cdot\bm{\theta}=\kappa f(\bm{\theta})$. This identity was utilized in a prior work which studied L2 regularization in the lazy regime . For a $L$-hidden layer ReLU network $\phi(h)=\max(0,h)$, the degree is $\kappa=L+1$, while rectified power law nonlinearities $\phi(h)=\max(0,h)^{q}$ give degrees $\kappa=\frac{q^{L+1}-1}{q-1}$. We note that the fixed point of the function dynamics above gives a representer theorem with the final NTK

where $[\bm{k}(x)]_{\mu}=\lim_{t\to\infty}K(\bm{x},\bm{x}_{\mu},t)$ and $K_{\mu\alpha}=\lim_{t\to\infty}K(\bm{x}_{\mu},\bm{x}_{\alpha},t)$. The prior work of Lewkowycz et al considered NTK parameterization $\gamma_{0}=0$. In this limit, the kernel (and consequently output function) decay to zero at large time, but if $\gamma_{0}>0$, then the network converges to a nontrivial fixed point as $t\to\infty$. In the DMFT limit we can determine the final kernel by solving the following field dynamics

We see that the contribution from initial conditions is exponentially suppressed at large time $t$ while the second term contributes most when the system has equilibrated. We provide an example of the weight decay DMFT showing its validity in a two layer ReLU network in Figure 3.

Appendix K Bayesian/Langevin Trained Mean Field Networks

Rather than studying exact gradient flow, many works have considered Langevin dynamics (gradient flow with white noise process on the weights) of neural network training . This setting is of special theoretical interest since the distribution of parameters converges at long times to a Gibbs equilibrium distribution which has a Bayesian interpretation . The relevant Langevin equation for our mean field gradient flow is

where $\lambda$ is a ridge penalty which controls the scale of parameters, and $d\bm{\epsilon}(t)$ is a Brownian motion term which has covariance structure $\left<d\bm{\epsilon}(t)d\bm{\epsilon}(t^{\prime})^{\top}\right>=\delta(t-t^{\prime})\mathbf{I}$. The parameter $\beta$, known as the inverse temperature controls the scale of the random Gaussian noise injected into this stochastic process. The dynamical treatment of the $\beta\to\infty$ limit will coincide with our usual DMFT while the $\beta\ll\infty$ will exhibit a nontrivial balance between the usual DMFT feature updates and the random Langevin noise. At late times, such a system will equilibrate to its Gibbs distribution.

In this section we analyze the dynamical mean field theory for these Langevin dynamics. First we note that the effect of regularization can be handled with a simple integrating factor

Our aim is now to compute the moment generating function for the $\{\bm{\chi},\bm{\xi},\bm{\chi}^{\epsilon},\bm{\xi}^{\epsilon}\}$ fields which causally determine $\{\bm{h},\bm{z}\}$. This MGF has the form

We insert Dirac-delta functions in the usual way to enforce the definitions of $\bm{\chi},\bm{\xi},\bm{\chi}^{\epsilon},\bm{\xi}^{\epsilon}$ and then average over $\bm{\theta}_{0},\bm{\epsilon}(t)$. These averages can be performed separately with the $\bm{\theta}_{0}$ average giving the identical terms as derived in previous sections. We focus on the average over Brownian disorder

The order parameters can be determined by the saddle point equations. These equations for $\Phi,\hat{\Phi},G,\hat{G},A,B$ are the same as before. The new equations are

As before, we make the substitutions $\bm{B}\to\gamma_{0}^{-1}{\bm{B}}^{\top}$ and $\bm{A}\to\gamma_{0}^{-1}\bm{A}$ and arrive at the final DMFT equations

K.2 Weak Feature Learning, Long Time Limit

On the temporal diagonal $\tau=0$, these equations give the usual recursions used to compute the NNGP kernels at initialization , though with initialization variance $C_{\lambda,\beta}(0)=\lambda^{-1}$, set by the weight decay term in the Langevin dynamics. This indicates that the long time Langevin dynamics at $\gamma_{0}\to 0$ simply rescales the Gaussian weight variance based on $\lambda$. It would be interesting to explore fluctuation dissipation relationships at finite $\gamma_{0}$ within this framework which we leave to future work.

K.3 Equilibrium Analysis

We see that $p(\bm{f}|\mathcal{D})\propto\int d\Phi d\hat{\Phi}\exp\left(NS[\Phi,\hat{\Phi}]\right)$ where

which implies that $f_{\mu}$ at the fixed point satisfies the following equations

The last layer’s dual kernel has the form $\hat{\Phi}^{L}_{\mu\alpha}=-\frac{\gamma_{0}^{2}\beta^{2}}{2\lambda}\Delta_{\mu}\Delta_{\alpha}$, which we see vanishes as feature learning strength is taken to zero $\gamma_{0}\to 0$, while for non-negligible $\gamma_{0}$, we see that the last layer features are non-Gaussian. We thus see that the moment generating function for the last layer field has the form

In the $\gamma_{0}\to 0$ limit, the non-Gaussian component of this density vanishes. Now that we have this form, we can compute $\Phi^{L}$ conditional on $\Phi^{L-1}$. Next, we calculate $\hat{\Phi}^{L-1}_{\mu\alpha}=\left<\hat{h}^{L}_{\mu}\hat{h}^{L}_{\alpha}\right>$, giving

Again, we note that in the $\gamma_{0}\to 0$ limit, since $\left<\bm{h}^{L}\bm{h}^{L}\right>\sim\lambda^{-1}\bm{\Phi}^{L-1}$, so that $\hat{\bm{\Phi}}^{L-1}=0$, implying that the $h^{L-1}$ fields are also Gaussian in this $\gamma_{0}\to 0$ limit. For arbitrary $\gamma_{0}$, this recursive argument can be completed going backwards using

For deep linear networks, the distributions are all Gaussian, allowing one to close algebraically, the saddle point equations for $\Phi,\hat{\Phi}$ .

Appendix L Momentum Dynamics

Standard gradient descent often converges slowly and requires careful tuning of learning rate. Momentum, in contrast can, be stable under a wider range of learning rates and can benefit from acceleration on certain problems . In this section we show that our field theory is still valid when training with momentum; simply altering the field definitions appropriately gives the infinite-width feature learning behavior.

which implies the following field evolution

We see that in the $\tau\to 0$ limit, the $t^{\prime\prime}$ integral is dominated by the contribution at $t^{\prime\prime}\sim t^{\prime}$ recovering usual gradient descent dynamics. For $\tau\gg 0$, we see that the integral accumulates additional contributions from the past values of fields and kernels.

Appendix M Discrete Time

We treat $f_{\mu}(t)$ as a potentially random variable and insert

Noting that $\bm{w}^{L}(0)$ is involved in the definition of both $f_{\mu}(t)$ and $\bm{\xi}_{\mu}^{L}(t)$, we see that the average over $\bm{w}^{L}(0)$ now takes the form

We extend our definition as before $iA^{L}_{\mu\alpha}(t,s)=\frac{1}{N\gamma_{0}}\phi(\bm{h}^{L}_{\mu}(t))\cdot\bm{\xi}^{L}_{\alpha}(s)$. Proceeding with the calculation as usual, we find that

The saddle point equations can now be analyzed. In addition to the usual order parameters, we note that $f,\hat{f}$ also generate saddle point equations

We also obtain saddle point equations for the new $A^{L},B^{L}$ order parameters.

which implies $B^{L}_{\mu\alpha}(t,s)=0$ and $A^{L}=\gamma_{0}^{-1}\left<\frac{\phi(h^{L}_{\mu}(t))}{\partial r^{L}_{\alpha}(s)}\right>$. This gives the following DMFT

We leave it to future work to verify that a continuous time limit of the above DMFT recovers function evolution governed by the NTK.

Appendix N Equivalent Parameterizations

We will now derive constraints on $(a,b,c,d)$ which give desired large width behavior. We will identify a one-dimensional family of parameterizations which satisfy three desiderata of network training 1. finite preactivations, 2. learning in finite time, 3. feature learning.

As before we let the NTK be the matrix which defines network prediction dynamics $\partial_{t}f_{\mu}=\sum_{\alpha}K^{NTK}_{\mu\alpha}\Delta_{\alpha}$. We demand that this matrix be $\mathcal{O}_{N}(1)$ so that the network predictions have $\mathcal{O}_{N}(1)$ evolution

Having fields which are $\mathcal{O}_{N}(1)$ can be ensured at initialization provided that

Since we already demanded that $2a_{L}+b_{L}=1$, this inner product will be $\mathcal{O}_{N}(1)$. Now we will see whether it remains $\mathcal{O}_{N}(1)$ under its recursion

which is indeed $\mathcal{O}_{N}(1)$ as desired provided that $2a+b=1$.

N.4 Putting Constraints Together

We now let $\gamma=\gamma_{0}N^{d}$. We see that the set of parameterizations which yield $\mathcal{O}(1)$ feature evolution are those for which

Now that we have established that the parameterization we consider here (modified NTK parameterization) is equivalent to $\mu P$, (modified standard parameterization), we will now demonstrate that the stochastic process which we obtained through a stationary action principle applied to our DMFT action $S$ is equivalent to the stochastic process derived from the Tensor Programs framework of Yang . Using the notation from Appendix H of Yang and Hu , they give the following evolution equations for the preactivations in a hidden layer in one pass SGD

After these substitutions are made, we see that the equations above match the one-pass SGD version of the DMFT Equations in Appendix M. A similar identification can be made for the backward pass. This shows that both Tensor Programs and DMFT, though alternative derivations, give identical descriptions of the stochastic processes induced by random initializations + GD in infinite neural networks.

Appendix O Gradient Independence

Appendix P Perturbation Theory

In this section we analyze the leading corrections in a small $\gamma_{0}$ expansion of our DMFT theory. All fields are expanded in power series in $\gamma_{0}$.

Our goal is to calculate all corrections to the kernels up to $\mathcal{O}(\gamma_{0}^{3})$ to show that the leading correction is $\mathcal{O}(\gamma_{0}^{2})$ and the subleading correction is $\mathcal{O}(\gamma_{0}^{4})$. It will again be convenient to utilize the vector notation defined in D.

We note that unlike other works on perturbation theory in wide networks, we do not attempt to characterize fluctuation effects in the kernels due to finite width, but rather operate in a regime where the kernels are concentrating and their variance is negligible. For a more thorough discussion of perturbative field theory in finite width networks, see .

Note that $[\bm{C}^{0}\bm{g}]_{\mu t}=\int_{0}^{t}dt^{\prime}\sum_{\beta}H^{0}_{\mu\beta}(t,t^{\prime})\Delta_{\beta}(t^{\prime})g(t^{\prime})=\sum_{\beta}K^{x}_{\mu\beta}\int_{0}^{t}dt^{\prime}\Delta_{\beta}(t^{\prime})g(t^{\prime})$ and note that $[\bm{D}\bm{h}]_{t}=\int_{0}^{t}dt^{\prime}G^{0}(t,t^{\prime})\sum_{\alpha}\Delta_{\alpha}(t^{\prime})h_{\alpha}(t^{\prime})=\sum_{\alpha}\int_{0}^{t}dt^{\prime}\Delta_{\alpha}(t^{\prime})h_{\alpha}(t^{\prime})$.

We can simplify the notation by introducing functions $v_{\alpha}(t)=\int_{0}^{t}\Delta_{\alpha}(t^{\prime})$ and $v_{\alpha\beta}(t)=\int_{0}^{t}dt^{\prime}\Delta_{\alpha}(t^{\prime})\int_{0}^{t^{\prime}}dt^{\prime\prime}\Delta_{\beta}(t^{\prime\prime})$.

P.2 Nonlinear Perturbation Theory

We start with the formula which implicitly defines $\bm{h},\bm{z}$

We proceed under the assumption of a power series in $\gamma_{0}$

We start in the first layer where $\bm{u}^{1}\sim\mathcal{GP}(0,\bm{K}^{x}\otimes\bm{1}\bm{1}^{\top})$ (note that this is $\mathcal{O}_{\gamma_{0}}(1)$) and compute the expansion of $\Phi^{1}$ in $\gamma_{0}$

We also obtain, by a similar argument, that the cubic $\mathcal{O}(\gamma_{0}^{3})$ term vanishes. To see this, note that $\bm{h}^{1,3}$ only contains odd powers of $\bm{r}^{1}$. Next, $\bm{h}^{1,1}\bm{h}^{1,2}$ contains only odd powers of $\bm{r}$, and $(\bm{h}^{1,1})^{3}$ is cubic in $\bm{r}$. Since all odd moments of a mean-zero Gaussian vanish, all averages of these terms over $\bm{r}$ annihilate, causing the $\gamma_{0}^{3}$ terms to vanish. Thus $\bm{\Phi}^{1}=\bm{\Phi}^{1,0}+\gamma_{0}^{2}\bm{\Phi}^{1,2}+\mathcal{O}(\gamma_{0}^{4})$.

The analogous argument for $\bm{G}^{L}$ now can be provided. First note that $\bm{r}^{L}$ is independent of $\bm{u}^{L}$ and of $\gamma_{0}$. Thus we can find that $\bm{G}^{L}$ has no linear-in-$\gamma_{0}$ term in its expansion since

each term contains only odd powers of $\bm{r}^{L}$ and odd moments of Gaussian variables vanish. After much more work, one can verify that $\bm{G}^{L,3}$ also must vanish since all terms contain odd powers of $\bm{r}$.

First, note that $\bm{g}^{L,0}$ is linear in $\bm{r}$. Next, note that $\bm{g}^{L,1}$ only depends on even powers of $\bm{r}$ since $\bm{g}^{L,1}=\dot{\phi}(\bm{u})\bm{z}^{L,1}+\ddot{\phi}(\bm{u})\bm{h}^{L,1}\bm{r}$. Next, we have

which only depends on odd powers of $\bm{r}$. Lastly, we have $\bm{g}^{L,3}$

which we see only contains even powers of $\bm{r}$. Thus $\bm{g}^{L,3}\bm{g}^{L,0}$ will be odd in $\bm{r}$. Looking at the expansion for $\bm{G}^{L,3}$, we see that all terms are odd in $\bm{r}$ and so the averages vanish under the Gaussian integrals.

P.6 Form of the Leading Corrections

The first term requires additional expansion to extract the corrections in $\gamma_{0}^{2}$

Now the correction to the NTK has the form

P.7 Perturbation Theory in Width N𝑁N (Finite Size Corrections)

The infinite-width DMFT is characterized by the set of saddle point equations which are $\nabla_{\bm{k}}S[\bm{k}]|_{\bm{k}=\bm{k}^{*}}=0$. Let the saddle point be $\bm{k}^{*}$. To identify corrections to the observable average $\left<O(\bm{k})\right>$ due to finite size, we now Taylor expand $S$ around $\bm{k}^{*}$

The linear component vanishes at the saddle point since $\nabla_{\bm{k}}S[\bm{k}]|_{\bm{k}=\bm{k}^{*}}=0$. Our observable average is thus

where we made the change of variables $\bm{\delta}=\sqrt{N}(\bm{k}-\bm{k}^{*})$. The function $U(\bm{\delta})$ contains all higher order terms (cubic and higher) in the Taylor expansion of $NS[\bm{k}]$. Since the leading power in $U$ is cubic in $(\bm{k}-\bm{k}^{*})=N^{-1/2}\bm{\delta}$, the leading behavior of this remainder is $U=\mathcal{O}(N^{-1/2})$ so it can be regarded as a perturbation to the Gibbs distribution. Taylor expanding the exponential $\exp\left(-\frac{1}{2}\bm{\delta}\nabla^{2}S[\bm{k}^{*}]\bm{\delta}-U(\bm{\delta})\right)=\exp\left(-\frac{1}{2}\bm{\delta}\nabla^{2}S[\bm{k}^{*}]\bm{\delta}\right)[1-U+\frac{1}{2}U^{2}+...]$ in both numerator and denominator, we eliminate the presence of the higher order terms in the Gibbs measure. Lastly, we let $\left<\cdot\right>_{0}$ represent an average over the unperturbed Gaussian potential $\bm{\delta}\sim\mathcal{N}(0,[\nabla^{2}S[\bm{k}^{*}]]^{-1})$. For notational simplicity, we let $\epsilon=N^{-1/2}$ and obtain

where $\left<\right>_{0}^{c}$ represents a connected cumulant . The first two connected correlations have the form

If one is interested only in the leading order correction to the observable $\left<O(\bm{k})\right>$, this can be obtained with the following correction

Since $U=\mathcal{O}(N^{-1/2})$ both corrections are of order $1/N$. This analysis shows that the leading order correction of the kernel distributions is $\mathcal{O}(N^{-1})$ and can be approximated by performing averages over a Gaussian distribution for $\bm{k}$ determined by the saddle point solution $\bm{k}^{*}$ and covariance given by $\frac{1}{N}\left[\nabla^{2}_{\bm{k}}S[\bm{k}]|_{\bm{k}=\bm{k}^{*}}\right]^{-1}$. We derive expressions for the components of this Hessian in P.7.1. These fluctuations have standard deviation $\mathcal{O}(N^{-1/2})$. This technique is a common approach to identifying finite size effects and was recently employed in Bayesian inference setting for networks in the lazy regime .

Before computing the Hessian terms, we can compare finite size effects under NTK scaling $\gamma=\mathcal{O}_{N}(1)$, and the mean field scaling $\gamma=\mathcal{O}(\sqrt{N})$. Concretely, we are interested in the feature learning component of the kernel change which is $\mathcal{O}(\frac{\gamma^{2}}{N})$. Let $\left<\Delta\bm{k}\right>$ represent the change in the kernel through training, which we showed in P.2 is of size $\mathcal{O}(\gamma_{0}^{2})=\mathcal{O}\left(\frac{\gamma^{2}}{N}\right)$. We will now define the signal to noise ratio of feature learning as

For NTK regime, this is vanishing as $N\to\infty$, while for the DMFT regime, this is goes as $\mathcal{O}(\sqrt{N})$ since kernel evolution is always $\mathcal{O}(1)$ but variance is $\mathcal{O}(N^{-1})$.