Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks
Authors: Blake Bordelon
Cengiz Pehlevan
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Learning Theory
Infinite Width
Kernel Methods

Abstract: We analyze feature learning in infinite-width neural networks trained with gradient flow through a self-consistent dynamical field theory. We construct a collection of deterministic dynamical order parameters which are inner-product kernels for hidden unit activations and gradients in each layer at pairs of time points, providing a reduced description of network activity through training. These kernel order parameters collectively define the hidden layer activation distribution, the evolution of the neural tangent kernel, and consequently output predictions. We show that the field theory derivation recovers the recursive stochastic process of infinite-width feature learning networks obtained from Yang & Hu with Tensor Programs. For deep linear networks, these kernels satisfy a set of algebraic matrix equations. For nonlinear networks, we provide an alternating sampling procedure to self-consistently solve for the kernel order parameters. We provide comparisons of the self-consistent solution to various approximation schemes including the static NTK approximation, gradient independence assumption, and leading order perturbation theory, showing that each of these approximations can break down in regimes where general self-consistent solutions still provide an accurate description. Lastly, we provide experiments in more realistic settings which demonstrate that the loss and kernel dynamics of CNNs at fixed feature learning strength is preserved across different widths on a CIFAR classification task.

Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks
Blake Bordelon & Cengiz Pehlevan
John Paulson School of Engincering and Applied Sciences, Center for Brain Science Harvard University
Cambridge MA, 02138
blake_bordelondg harvard. edu, epehlevangg www.paa................ harvard. edu
Abstract
We analyze feature learning in infinite-width neural networks trained with gradient flow through self-consistent dynamical field theory. We construct collection of deterministic dynamical order parameters which are inner-produs kernlls for hidden unit activations and gradients in each layer at pairs of time points, providing a reduced description of network activity through training. These kernel order parameters collectively define the hidden layer activation distribution, the evolution of the neural tangent kernel, and consequently output predictions. We show that the field theory derivation recovers the recursive stochastic process of infinite-width feature learning networks obtained from Yang & Hu with Tensor Programs [I]. For deep linear networks, these kernels satisfy a set of algebraic matrix equations. For nonlinear networks, we provide alternating sampling procedure to self- consistently solve for the kernel order parameters. We provide comparisons of the self-consistent solution to various approximation schemes including the static NTK approximation, gradient independence assumption, and leading order perturbation theory, showing that each of these approximations can break down in regimes where general self-consistent solutions still provide an accurate description. Lastly, provide experiments in more realistic settings which demonstrate that the loss and kernel dynamics of CNNs at fixed feature learning strength is preserved across different widths CIFAR classification task.
Introduction
Deep learning has emerged as a successful paradigm for solving challenging machine learning and computational problems across variety of domains 2 S
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 quantities.
through a reduced deseription involving desirably ditiialzzationnnddpeennn
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) [4-66 and neural tangent kernel (NTK) machines 779 in their respective infinite-width limits. These limits provide both analytic tractability as well as detailed training and generalization analysis IC 177777 However, in this limit, with these parameterizations, data representations
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 [ 19 The fact that the representation of data is fixed renders these kernel-base theories incapable of explaining feature
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

learning, an ingredient which is crucial to the success of deep learning in practice 120 21 Thus,
alternative theories capable of modeling feature learning dynamics are needed.
Recently developed alternative parameterizations such as the mean field [22] and the jaP 1 param- eterizations allow feature learning in anfiitt-wwdth 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-widit μul NNs M. In this work, we study an equivalent parame- terization to μ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 scalar parameter Yo 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 the following:
We develop a path integral formulation of gradient flow dynamics in infinite-widie width networks the feature learning regime. Our parameterization includes a scalar parameter "oo to allow interpolation between rich and lazy regimes and comparison to perturbative methods.
2. 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 Too = 1, from an alternative method, the same stochastic process obtained previously with Tensor Programs M.
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.
1.1 Related Works
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-widde limit. Finite width corrections to Bayesian inference in wide networks have been obtained with various perturbative [23 and self-consistent techniques 3C. In the gradient descent based setting, leading order corrections to the NTK dynamics have been analyzed to study finite width effects [34 36 27.. 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 G7/ 33
Another approach to study feature learning is to alter NN parameterization in gradeent-based learning to allow significant feature evolution even at wiith,, the mean field limit [2 39 Works on mean field NNs have yielded formal loss convergence results 440 41 gradient flow dynamics to a partial differential equation (PDE)
and shown equivalences of
Our results are most closely related to a set of recent works which studied wiini NNs trained with gradient descent (GD) using the Tensor Programs (TP) framework M1. We show that our discrete time field theory at unit feature learning strength fo 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 hyyer-parameter search 45 Computing the exact infinite-widit limit of GD has exponential time requirements M. which we show can be circumvented with an alternating sampling procedure. A projected variant of GD training has provided an wiinth theory that could be scaled to realistic datasets like CIFAR- 10 46). Inspired by Chizat and Bach's ' work on mechanisms of lazy and rich training [8] our theory interpolates between lazy and rich behavior in the mean field limit for varying " and allows comparison of DMFT to perturbative analysis near small Yo- Further, our derivation of 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 47 53). This framework has been utilized in the theory of random recurrent networks [54] 59).
tensor PCA (60 (1), phase retrieval (62), and high-dimensional linear classifiers (33-66 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-wide width limit.

---

2 Problem Setup and Definitions
Our theory applies to infinite-widde width networks, both fully-connected and convolutional. For notational case we will relegate convolutional results to later sections. For input x E IR,, we define the hidden pre-activation vectors h' E IR for layers & E {1,.... | as
1 VD
fw
TVN w ((h))
h++1
W o'
W
(I)
where 0 = Vec(Wo ..... ) are the trainable parameters of the network and é is a twice differ- entiable activation function. Inspired by previous works on the mechanisms of lazy gradient based training. the parameter 7 will control the laziness or richness of the training dynamics [8 19 42
Each of the trainable parameters are initialized as Gaussian random variables with unit variance W
(0, 1). They evolve under gradient flow 0 = - Vo. The choice of learning rate y2 causes d di Cl=-0 to be independent of 7. To characterize the evolution of weights, introduce backpropaga- ffg ah
tion variables -f. = TV N
= o(h,,) 3 z where z VN W gfft1 is the pre- pre-gradient signal.
The relevant dynamical objects to characterize feature learning are feature and gradient kernels for each hidden layer e E (1,..., L), defined as
1 F (s).
T (t, s) N (f,, (sh) Gjin (t,
(2)
From the kernels (
G we can compute the Neural Tangent Kernel K (t,s) = (ef
Vofa(s) = Eiio G+t po (t, s) (t, s), 7 and the dynamics of the network function fP
d dt fu(t)
= Σ
K NTK (t, (t)t) (t) 0fr -l,,,0))
(3)
a=1
where we define base cases po (1, 1, were termed forward and backward kernels and were theoretically computed at initialization and empirically measured through training [67]. Our DMFT will provide exact formulas for these kernels throughout the full dynamics of feature learning. We note that the above formula holds for any data point which may 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 s diagonal gives the temporal trajectory of the network predictions (t).
10 pa (4,8) = K2 J00 aa In prior work, &e, Ge
Following prior works on infinite-width networks
222
40
we study the mean field limit
N 00.
To
= (1)
(4)
As demonstrate in the Appendix D and N this is the only N-scaling which allows feature learning
as N > oc. The To = limit recovers the static NTK limit 77. We discuss other scalings and parameterizations in Appendix N relating our work the ..-paarameterizzatonn and TP analysis of W. 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 [D8] We focus equal widths and NTK parameterization (as in eq. (T) in the main text to reduce complexity.
3 Self-consistent DMFT
Next, we derive our self-consistent DMFT in a limit where I, = ON(1). Our goal is to build description of training dynamics purely based on representations, and independent of weights. Studying feature learning at infinite-width enjoys several analytical properties:
The kernel order parameters & G concentrate allowing flexible adaptation of features to the task structure.
random initializations but are dynamical,
In each layer f, each neuron's preactivation hy and pregradient 2 become i.i.d. draws from a distribution characterized by a set of order parameters {Φ6 G AC, B²).
The kernels are defined as self consisteet averages (denoted by ()) over this distribution of neurons in each layer of no (1, 8) =
(+))) and G.aaa (1, s) = (gf, (s))
The next section derives these facts from path- aath-integral formulation of gradient flow dynamics.
3

---

3.1 Path Integral Construction
Gradient flow after random initialization of weights defines a high dimensional stochastic process over initalizations for variables {h.g). Therefore, we will utilize DMFT formalism to obtain re- duced description of network activity during training. For simplified derivation of the DMFT for the two-layer (I, (I=1] 1) case, see D.2 Generally, we separate the contribution on each forward/backwart pass between the initial condition and gradient updates to weight matrix W defining new stochastic variables ,5 E RR. as
XF (1)
w (O))
(O 9 (1).
(5)
VN We let Z represent the moment generating functional (MGF) for these stochastic fields
VN
dt (1t(t) (1) (t) () ))
ZHi
Σ E,
exp
(W)())...w=(01
which requires, by construction the normalization condition 2 0J} = 1. We enforce our definition of X. using an integral representation of the eetta-function Thus for each sample μ E [P] and each time t E R4, we multiply Z by
==fr JRN JRN
dx x+1 ++1 (t) R (2
(t)
(x xx+1(0)
1 W (0)000)))))
1 1=
exp
(x
(6)
VN
for X and the respective expression for &. After making such substitutions, we perform integration over initial Gaussian weight matrices arrive at an integral expression for Z. which we derive in the
appendix D.4
We show that Z can be described by set of order-parametts {Ø', $6, G G A B²}
/
ZHi o H
II Epeatts
d (t, s)d ju (t, s)dG jua (t, s)d un (t, s)dA pm (t. s)dB
pm (t, (7)
G. x exp C,
S lua
Σ
dt
ds
Th Wa (1, (1, (1,8) Gha (1, 8) s) na (1, A Ayaa 8) B) (1, s)]
+ In 2 &, G, G.G.A.B.j...)..
where S' is the DMFT action and 2 is single-site MGF, which defines the distribution of fields { over the neural population in each layer. The kernels A and are related to the correlations between feedforward and feedback signals in the network. We provide a detailed formula for Z in the Appendix D.4 and show that it factorizes over different layers 2
l Zz.
3.2 Deriving the DMFT Equations from the Path Integral Saddle Point
As N + 00, the moment- -enneettgnnerating function Z exponentially dominated by the saddle point of S. The equations that define this saddle point also define our DMFT. We thus identify the kernels that render locally stationary (6S 0). The most important equations are those which define {Ø G
85
1
8z
T (t,
Ob (t, (0()) (s)) = 0,
8 uc (t)
Z
s pa (t,
8S
1
8Z
(1,88) (t,a(t)))
= (t,(t) (f) = O,
(9)
Z
sôt ia (1,8)
where () denotes average over the stochastic process induced by Z, which is defined below
(u)(0)EE)P))) ~GP(O, $(-1), (r((())))))))) GP(0,
P
(((11) = (1) + 70 I ds
[ (1, + An) nn (1, 8)] ()))(()))))))
a=l
{{(()) = (t) + [ ds
Σ
[B pu (t, + (s) (t. s)] @(h (s)),
(10)
a=l

---

where we define base cases $ (t, = jos and G (t,s) = 1, = = = = 0. We see that the fields 2", which represent the single site preactivations and pre-gradients, are implicit functionals of the mean-zero Gaussian processes {u', which have covariances (uf, (s)) Dfaaa (t, and (r)()))(())) = Gfa (t) The other saddle point equations give A (t, 8) %o"
(O) (x)
6g. +1 (e) dug (s) which arise due to coupling between
By (t, =
the feedforward and feedback signals. We note that, in the lazy limit To - 0, the fields approach
Gaussian processes h² + u², x² re. Lastly, the final saddle point equations .55 0,
As
64
DG
imply that e' = G = O. The full set of equations that define the DMFT are given in D.7
This theory is easily extended to more general architectures such as networks with varying widths by layer (App- D.8. trainable bias parameter (App. Hi multiple (but (x()) output channels (App. D. convolutional architectures (App- GD. networks trained with weight decay (App. D. Langevin sampling (App. K and momentum (App. D. discrete time training (App. Mp. In Appendix N we
discuss parameterizations which give equivalent feature and predictor dynamics and show our derived stochastic process is equívalent to the μP scheme of Yang & Hu
4 Solving the Self-Consistent DMFT
The saddle point equations obtained from the field theory discussed in the previous section must be solved coff-sonsittetly. By this we mean that, given knowledge of the kernels, we can characterize the distribution of {h ). and given the distribution of {h 2f), we can compute the kernels [68] 64]. In the Appendix B we provide Algorithm|] D numerical procedure based on this idea to efficiently solve for the kernels with an alternating Monte-Carlo strategy. The output of the algorithm are the
dynamical kernels &f (laaa. (1, s), GLo (1, 8), (4, Be (1,, 8), from which any network observable can be computed as discuss Appendix [ We provide
example of the solution to the saddle point equations compared to training a fimite NN in Figure [ We plot & G at the end of training and the sample-trace of these kernels through time. Additionally, we compare the kernels of finite width N network to the DMFT predicted kernels using a cosine-similarity alignment metric A( A($DME DMFT
NNN)=
4 DMFT NN
Additional examples are in Appendix Figures and Figure
4.1 Deep Linear Networks: Closed Form Self-Consistent Equations
Deep linear networks (d(h) = h) are of theoretical interest since they are simpler to analyze than nonlinear networks but preserve non-trivial training dynamics and feature learning (69] 7K
255 K2 23
In deep linear network, we can simplify our saddle point equations to algebraic formulas that close in terms of the kernels H. (t, s) = (h (t) (s)). (t,t, s) (g'( (s)) [. This is a significant simplification since it allows solution of the saddle point equations without sampling procedure.
To describe the result, we first introduce a vectorization notation ht (t),)))e) Like- wise we convert kernels H = Mat H noo (1, s)) into matrices. The inner product under this vectorization is defined as a b -o dt Ep= ag (t)),,)t). In a practical computa- tional implementation, the theory would be evaluated on grid I' time points with discrete time gradient descent, so these kernels H E pTXPT would indeed be matrices of the appropriate size. The fields h h'g are linear functionals of independent Gaussian processes u', u',r', giving D'C' are causal integral operators which depend on {A--1 H and (B', Gt+1] respectively which we define in Appendix [ The saddle point equations which define the kernels are
(I (I- -- 'D)hh D' = u + to (I-28
D'C' "g" r + D'u'. The matrices C and D
H = (h'n" = (1- D") + G 1 D
(11) Examples of the predictions obtained by solving these systems of equations provided in Figure 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.
G (I-
[G D [II DDD'
Unlike many prior results (69 72), our DMFT does not require any restrictions on the structure of the input data but hold for any K" y. However, for whitened data K = I we show in
5

---

3.8
Initial Density
Final Density
4
04
344
100
(a) Lazy Rich Loss Dynamics
(b) Initial Preactivation Density
(c) Final Preactivation Density
o Exp M'taa
Exst 0,
Molugr
Theiry 0
10
(d) Final Kernels "ro
(e) &
Dynamics n = 1.0
(l) 4' Convergence to DMFT
1.0
Therry cl.
(g) Final G kemnels 70 1.0
(h) G Dynamics " 1.0
G Convergence DMFT
Figure 1: Neural network feature learning dynamics is captured by self-consistent dynamical mean field theory (DMFT). (a) Training loss curves on a subsample of P = 10 CIFAR-I0 training points in depth 4 (L N 2500) tanh network (#(h) = (anh(h) trained with MSE. Increasing To accelerates training. (b)-(c) The distribution of preactivations at the beginning and end of training matches predictions of the DMFT. (d) The final 4' (at 100) kernel order parameters match the finite width network. (c) The temporal dynamics of the sample-traee traced kernels E,, D "pp (t, s) matches experiment and reveals rich dynamics across layers. (f) The alignment A DMMT:
& $ N defined cosine similarity, of the kernel of (1,, s) predicted by theory (DMFT) and width N networks for different N but fixed Around N 500 DMFT begins show near perfect agreement with the NN. (g)-(i) The same plots but for the gradient kernel G Whereas finite width effects for 4 are larger at later layers l since variance accumulates on the forward pass, fluctuations G are large in early layers.
N/vN. Errorbars show standard deviation computed over 10 repeats.
Appendix
that our DMFT learning curves interpolate between NTK dynamics and the
sigmoidal trajectories prior works [69 70I as To is increased. For example, in the two layer
(L = 1) linear network with K I, the dynamics of the error norm A(t) = | (f)l takes the form (t) 22 7 (v ((t)) where y = |lyll- These dynamics give the linear convergence rate of the NTK if 70 - 0 but approaches logistic dynamics of (70) as 7o + go. Further, ( = (h' (t) E PxP only grows in the yy direction with H2)() y H()y
At the end of training Htt) I 1111+++ 22²² -1lyy" recovering the
rank spike which recently obtained in the small initialization limit 74 We show this one dimensional system in Figure g
4.2 Feature Learning with L2 Regularization
As we show in Appendix g the DMFT can be extended to networks trained with weight decay di - VoL X0. If neural network homogenous in its parameters so that f(cØ) = c* (0) (examples include networks with linear, ReLU, quadratic activations), then the final network predictor

---

H1
H2
H3
H4
HS
0.4
Theory
(a) Deep Linear Loss Dynamics
(b) Predicted vs Experimental Final H Kermels
f=1
1=2
==3
i= 4
1=5
G
A(H10, 0), HLT,
0.85
6.60
Hitt,s)
2.
23
3.2
3.5
4.2
4.5
3.0
(c) I-Dependent Kernel Movement
(d) == 5 DMFT Temporal Kernels
Figure 2: Deep linear network with the full DMFT. (a) The train loss for NNs of varying L. (b) For 5, 1000 NN, the kernels H' at the end of training compared DMFT theory on P 20 datapoints. (c) The average displacement of feature kernels for different depth networks at same "N0 value. For equal values of "o. deeper networks exhibit larger changes to their features, manifested in lower alignment with their initial O kernels H. (d) The solution to the temporal components of the G (t, s) and E,, Hy (t, s) kernels obtained from the self-consistent equations.
is kernel regressor with the final NTK lim f(x, = k(x) [ + AxI] y where K ((x, x') the final-NTK, (x)))]
(ar, x. and μua K (2 T.). We note that the effective regularization Aσ increases with depth L. In NTK parameterization, weight decay in infinite width homogenous networks gives trivial fixed point K ((x, x²) - and consequently zero predictor f 0 15 However, as we show in Figure 3 increasing feature learning %o 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 nontrivial function of the hyperparameters X, k, L, 70-
=o 0.25
Yo 0.50
Yo 0.75
Yo 1.00
Expt
c(3)
1==025
0.50
b DMFT
1.0000
(a) Loss for varying 7oo
(b) Final Kernlss
Figure 3: Width N 1000 ReLU networks trained with L.2 regularization have nontrivial fixed point in DMFT limit (70 > 0). (a) Training loss dynamics for a L ReLU network with A = 1. In 7o + 0 limit the fixed point is trivial f K 0. The final loss is a decreasing function of 7h- (b) The final kernel is more aligned with target with increasing To. Networks with homogenous activations enjoy representer theorem at fonntee-width as show in Appendix !

---

5
Approximation Schemes
We now compare exact DMFT with approximations of prior works, providing explanation of when these approximations give accurate predictions and when they break down.
5.1 Gradient Independence Ansatz
We can study the accuracy of the ansatz = 0, which is equivalent to treating the weight matrices (0) and (0) _ which appear in forward and backward passes respectively indepen- dent Gaussian matrices. This assumption was utilized in prior works on signal propagation in deep networks in the lazy regime 776-801 A consequence of this approximation is the Gaussianity and statistical independence of x" and 5 (conditional on {Ø², G*) in each layer as we show in Appendix O This ansatz works very well near "lo x O (the static kernel regime) since dh dz dr du ~ (17o) or
around initialization i but begins to fail at larger values of 70, (Figure 4
5.2 Perturbation theory in to at indinite-width
In the To + limit, we recover static kernels, giving linear dynamics identical the NTK limit 77. Corrections to this lazy limit can be extracted at small but finite voo 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 finite N effects). We expand all observables ((70) in
B5
28
a power series in To. giving ((7)) = ") + +oog(1) + 220g(2) + and compute corrections up to O(1)). We show that the ((1g) and (())) corrections to kernels vanish, giving leading order expansions of the form $- = Φ + +1, 22 + +(n)) and G = G + 22 + O(1)) (see Appendix
Further, we show that the NTK has relative change leading order which scales linearly with depth AK |/ ~ O (L))) = -N.. (²*) which is consistent with finite width effective field theory 7 = (1) 26 28 (Appendix P.6 Further, at the leading order correction, all temporal dependencies are controlled by P(P + 1) functions Va (t) = Jdd ds (s) and (t) f (s) So ds' (g')), which consistent with those derived for finite width NNs using truncation of the Neural Tangent Hierarchy 34 33 27 To lighten notation, focus
main text comparison of our pontuertur DMFT to perturbation theory in the deep linear case.
Full perturbation theory is in Appendix P2
Using the timescales derived the previous section, we find that the leading order correction to the kernels in infinite-wide deep linear network have the form
L(L+1)
KNT (1, = (I + 1)K%, +10 2
K Σ
K (a + (s)(s) + vn (+()))(())]
2
ag
+ 70 2
L(L 1)
Σ a,3
Kz na (oglvos()) + upa (x)] +
Σ a,
na Kiga
+ +(()). (12)
We see that the relative change in the NTK KNNKK- K NTK (0)l/ K KNTK (0) O L/N). so that large depth L networks exhibit more significant kernel evolution, which agrees with other perturbative studies 335 27 25 as well as the non non-eetuurbativ results in Figure However large %o and large L, this theory begins to break down as we show in Figure
The DMFT formalism can also be used to extract leading corrections to observables at large but finite width N as we explore in P.7. When deviating from infinite width, the kernels are longer deter- ministic over network initializations. The key observation is that the DMFT action S defines a Gibbs measure over the space of kernel order parameters k Vec & G A B²) } with probability den- sity 7 exp (k]) where Z is a normalization constant. Near infinite width, any observable average (O(k)) = 7 f dk exp (K)) Okk) is dominated by order parameters within a / neighborhood of k.. As a consequence, a perturbative series for (O(k)) can be obtained from simple averages over Gaussian fluctuations in the kernel & ~ (K",,, S(*]+ [29. The components for V S"S] include four point correlations of fields computed over the DMFT distribution.
1
8

---

1.0
Gr Ind. Part. NTC OWFT Eus.
0.8
0.6
NN
0.4
DMFT
Pert.
Gr. Ind.
0.0
100
(a) Loss dynamics
(b) Final H' Kernels "yu = 1.5
alical
Hlical
1.00
NTC. DHFT Expr.
0.95
A(H. A(H.Hww) Haw)
0.90
DMFT
0.85
NTK
Int. Pert.
0.80
Pert.
0.25
0.50
0.75
1.00
1.25
1.50
Yo
(c) H Kernel Dynamics "Yo 1.5
(d) Theory H NN with N 1000
Figure 4: Comparison of DMFT to various approximation schemes in L hidden layer, width N 1000 linear network with % = 1.0 and P 100. (a) The loss for the various approximations do track the trajectory induced by gradient descent in the large To regime. (b)-(c) The feature kernels H. uaa (1,, 8) across each the = 5 hidden layers for each of the theories is compared to width 1000 neural network. Again, we plot the sample-traced dynamics Enn, H. (t, s). (d) The alignment of H compared to the finite NN A(H² H averaged across E (1,,.. 5} for varying y. The predictions of all of these theories coincide in the "(o = O limit but begin to deviate in the feature learning regime. Only the non-perturbative DMFT is accurate over a wide range of yo-
Feature Learning Dynamics is Preserved at Fixed %
Our DMFT suggests that for networks sufficiently wide for their kernels concentrate, the dynamics of loss and kernels should invariant under the rescaling N - RN, y +
7///R. which keeps " fixed. To evaluate how well this idea holds in realistic deep learning problem, we trained CNNs of varying channel counts N on two-class CIFAR classification [8]. We tracked the dynamics of the loss and the last layer &- kernel. The results are provided in Figure We see that dynamics are largely independent of rescaling as predicted. Further, as expected, larger "o leads to larger changes in kernel norm and faster alignment to the target function y, the higher " networks train more rapidly. The trend consistent for width N 250 and 500.
also found in [82]. Consequently,
More details about the experiment can be found in Appendix C.2
7 Discussion
We provided a unifying DMFT derivation of feature dynamics in infinite networks trained with gradient based optimization. Our theory interpolates between lazy infinite-widih behavior of static NTK in %o + 0 and rich feature learning. At 2o = 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 deep linear networks and can efficiently solved with numerical method in the nonlinear case. Comparisons with other approximation schemes show that DMFT can be accurate at much wider range of 70- We believe our framework could useful perspective for future theoretical analyses of feature learning and generalization in wide networks.

---

(40
Test Accuracy
Alignment
cm
197
695
(a) Test MSE
(b) Classification Error
(c) A( yy Dynamies
Figure 5: The dynamics of a depth 5 (L = hidden) CNNs trained on first classes of CIFAR (boat vs plane) exhibit consistency for different channel counts N E (250, 500) for fixed To = 7/ N. (a) We plot the test loss (MSE) and (b) test classification error. Networks with higher "o train more rapidly. Time is measured in every 100 update steps. (c) The dynamics of the last layer feature kernel 2², shown as alignment to the target function. As predicted by the DMFT, higher Yo corresponds to more active kernel evolution, evidenced by larger change in the alignment.
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, (x) and may break down otherwise; other asymptotic regimes (such as P/N, T/ log(N) (x etc) may The computational requirements of our method, while smaller than the exponential time complexity for exact solution D1 are still significant for large PT. In Table II 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 data points takes (PN22 computations. The static NTK requires computation of ((2)) entries in the kernel which do not need be recomputed. However, the DMFT requires matrix multiplications on PT X PT matrices giving a ( time scaling. Future work could aim to improve the computational overhead of the algorithm, by considering data averaged theories (44 or one pass SGD M. Alternative projected versions of gradient descent have also enabled much better computational scaling in evaluation of the theoretical predictions 46 allowing evaluation on
exhibit phenomena relevant to deep learning practice [32 K3
full CIFAR-10.
Requirements Memory for Kernels
Width-N N NN
Static NTK
Perturbative
Full DMFT PPT'2
O(N2)
((P))
O(PIT)
Time for Kernels
O(PN*T)
((P))
O(P'T)
OP333
Time for Final Outputs
O(PN2T)
O(Pa)
O(P4)
((P3T3)
Table 1: Computational requirements to compute kernel dynamics and trained network predictions on P points a depth N neural network on a grid time points trained with data points for various theories. DMFT is faster and less memory intensive than width N network only if N > PT. It
more computationally efficient to compute full DMFT kernels than leading order perturbation theory when T < vP. The expensive scaling with both samples and time are the cost of full-batch non-perturbative theory of gradient based feature learning dynamics.
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 # 16742269 and the Harvard Q-Bio Initiative.
BB thanks Jacob Zavatone-Vett Alex Atanasov, Abdulkadir Canatar, and Ben Ruben for comments on this manuscript as well Greg Yang, Boris Hanin, Yasaman Bahri, and Jascha Sohl-Dickstein for useful discussions.
10

---

References
[I] Greg Yang and Edward Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference Machine Learning, pages 1 11727-11737. PMLR, 2021.
[2] Ian Goodfellow, Yoshua Bengio, and Aaron Courville, Deep learning. MIT press, 2016.
[3] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
[4] Radford Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
[5] Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak. Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018.
[6] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
[7] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. BisacBi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 8571-8588 Curran Associates, Inc., 2018.
[8] Jachoon Lee, Lechao Xiao. Samuel Schoenholz, Yasaman Bahri, Roman Novak. Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[9] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019.
[10] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685. PMIR, 2019.
[11] B. Bordelon, A. Canatar, and C. Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. International Conference of Machine Learning, 2020.
[12] Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communications, 12(1):1-12, 2021.
[13] Omry Cohen, Or Malka, and Zohar Ringel. Learning curves for overparametrized deep neural networks: A field theory perspective. Physical Review Research, 3(2):023334, 2021.
[14] Arthur Jacot, Berfin Simsek, Francesco Spadaro, Clément Hongler, and Franck Gabriel. Kernel alignment risk estimator: Risk prediction from training data. Advances in Neural Information Processing Systems, 33:15568- 15578, 2020.
[15] Bruno Lourciro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mezard, and Lenka Zdeborova. Learning curves of generic features maps for realistic datasets with teacher-student model. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[16] James B Simon, Madeline Dickens, and Michael DeWeese. Neural tangent kernel eigenvalues accurately predict generalization. arXiv preprint arXiv:2110.03322, 2021.
[17] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via preettrratttriiiatinnn. International Conference on Machine Learning, pages 242-252. PMLR, 2019.

---

[18] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program ming. Advances Neural Information Processing Systems, 32, 2019.
[19] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal Statistical Mechanics: Theory and Experiment, (020011):1133001 2020.
[20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901 2020.
[21] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[22] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 1 115(3))EE7665- E7671, 2018.
[23] Laurence Aitchison. Why bigger is not always better: on finite and infinite neural networks, In International Conference Machine Learning, pages 156-164 PMLR, 2020.
[24] Sho Yaida. Non-gaussian processes and neural networks finite widths. In Mathematical and Scientific Machine Learning. pages 165-192. PMLR, 2020.
[25] Jacob Zavatone-Vett Veth, Abdulkadir Canatar, Ben Ruben, and Cengiz Pehlevan. Asymptotics of representation learning in finite bayesian neural networks. Advances in Neural Information Processing Systems, 34, 2021.
[26] Gadi Naveh, Oded Ben David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite deep neural networks trained with noisy gradients. Physical Review E, 104(6);66330,, 2021.
[27] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106. 10165, 2021.
[28] Boris Hanin. Correlation functions in random fully connected neural networks at finite width. arXiv preprint 2rXiv:2204.01058, 2022.
[29] Kai Segadlo, Bastian Epping, Alexander van Meegen, David Dahmen, Michael Krämer, and Moritz Helias. Unified field theory for deep and recurrent neural networks, 2021.
[30] Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects in finite enns. Advances in Neural Information Processing Systems, 34, 2021.
[31] Inbar Seroussi and Zohar Ringel. Separation scales and a thermodynamic description of feature learning in some cnns. arXiv preprint arXiv:2112.15388 2021.
[32] Qianyi and Haim Sompolinsky. Statistical mechanies of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3):031059, 2021.
[33] Jacob A Zavatone- Veth and Cengiz Pehlevan. Depth induces scale-averaging in overparame- terized linear bayesian neural networks, 55th Asilomar Conference on Signals, Systems, and Computers, 2021.
[34] Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International conference on machine learning, pages 4542-4551, PMLR, 2020.
[35] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. arXiv preprint arXiv: 1909.1 1304, 2019.
[36] Anders Andreassen and Fthan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint 20080:0000..8675. 2020.
12

---

[37] Jacob A Zavatone- Veth, William L Tong, and Cengiz Pehlevan. Contrasting random and learned features in deep bayesian linear regression. arXiv preprint arX1v:2203.00577, 2022.
[38] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprini arXiv:2003.02218, 2020.
[39] Dyego Araújo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprini vrXiv:/906.00/99 2019.
[40] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[41] Grant M Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.009) 00915, 2018.
[42] Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-fit free bounds and kernel limit. In Conference on Learning Theory, pages
PMLR, 2019.
[43]
Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprini arXiv: wwt.:/90220.000 2019.
[44] Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean-fiell field framework for over-voeriit werrparameeterize deep neural networks. In Conference on learning theory, pages 1887-1936. PMLR, 2021.
[45] Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34, 2021.
[46] Greg Yang, Michael Santacroce, and Edward J Hu. Efficient computation of deep nonlinear infinite-width neural networks that learn features. In International Conference on Learning Representations, 2022.
[47] Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A, 8(1):423, 1973.
[48] C De Dominicis. Dynamics as a substitute for replicas in systems with quenched random impurities. Physical Review B, 18(9):4933, 1978.
[49] Haim Sompolinsky and Annette Zippelius. Dynamic theory of the spin-glass phase. Physical Review Letters, 47(5):359, 1981.
[50] Haim Sompolinsky and Annette Zippelius. Relaxational dynamics of the edwards-anderson model and the mean-field theory of spin- glansgg Physical Review B, 25(1 ((:1))6, 1982.
[51] G Ben Arous and Alice Guionnet. Large deviations for langevin spin glass dynamics. Probability Theory and Related Fields, 102(4):555500, 1995.
[52] G Ben Arous and Alice Guionnet. Symmetric langevin spin glass dynamics. The Annals of Probability, 25(3):1367-144 1997.
[53] Gérard Ben Arous, Amir Dembo, and Alice Guionnet. Cugliandooo--urrhann equations for dynamics of spin-glasses. Probability theory and related fields, 136(4))61-660, 2006.
[54] A Crisanti and H Sompolinsky. Path integral approach to random neural networks. Physical Review E, 98(6):062120, 2018.
[55] Haim Sompolinsky, Andrea Crisanti, and Hans-Jurgen Sommers. Chaos in random neural networks. Physical review letters, 61(3):259, 1988.
13

---

[56] Moritz Helias and David Dahmen. Statistical Field Theory for Neural Networks. Springer International Publishing, 2020.
[57] Lutz Molgedey, Schuchhardt, and Heinz G Schuster. Suppressing chaos in neural networks by noise. Physical review letters, 69(26):3717, 1992.
[58] M Samuelides and Bruno Cessac. Random recurrent neural networks dynamics. The European Physical Journal Special Topics, (42(1):89 122, 2007.
[59] Kanaka Rajan, LF Abbott, and Haim Sompolinsky, Simendesdependennt suppression of chaos recurrent neural networks. Physical review e, 82(1):011903, 2010.
[60] Stefano Sarao Mannelli, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborova. Passed & spurious: Descent algorithms and local minima in spiked matrix-tensor models. In interna- tional conference on machine learning, pages 4333 4342. PMLR, 2019.
[61] Stefano Sarao Mannelli, Giulio Biroli, Chiara Cammarota, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Marvels and pitfalls of the langevin algorithm in noisy high- dimensional inference. Physical Review X, (0(1):01105,, 2020.
[62] Francesca Mignacco, Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to navigate rough landscapes: comparing gradient-descent-bass algorithms in the phase retrieval problem. Machine Learning: Science and Technology, 2(3):035029, 2021.
[63] Elisabeth Agoritsas, Giulio Biroli, Pierfrancesco Urbani, and Francesco Zamponi. Out-of- equilibrium dynamical mean-field equations for the perceptron model. Journal of Physics A: Mathematical and Theoretical, 51 (1):080500,, 2018.
[64] Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-fiele theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540 95500, 2020.
[65] Michael Celentano, Chen Cheng, and Andrea Montanari. The high digh-dimensiooal asymptotics of first order methods with random data. arXiv preprint arXiv:2112.07572, 2021.
[66] Francesca Mignacco and Pierfrancesco Urbani. The effective noise of stochastic gradient descent. arXiv preprint arXiv:2 10852, 2021.
[67] Yizhang Lou, Chris Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 14248- 14282. PMLR, 17-23 Jul 2022.
[68] Alessandro Manacorda, Grégory Schehr, and Francesco Zamponi. Numerical solution of the dynamical mean field theory of infiniteddimensional equilibrium liquids. The Journal of chemical physics, 152(16):164506, 2020.
[69] Kenji Fukumizu. Dynamics of batch learning in multilayer neural networks. In International Conference on Artificial Neural Networks, pages 189-194. Springer, 1998.
[70] Andrew M Saxe, James McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 2013.
[71] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2019.
[72] Madhu Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks, 132:428-446, 2020.
[73] Arthur Jacot, François Ged, Franck Gabriel, Berfin Şimşek, and Clément Hongler. Deep linear networks dynamics: Low-rank biases induced by initialization scale and 12 regularization. arXiv preprint 2rXiv:2106.159:: 2021.
14

---

[74] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022.
[75] Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with 1_2 regularization. Advances in Neural Information Processing Systems, 33:4790-47999, 2020.
[76] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
[77] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep informa- tion propagation. International Conference of Learning Representations, 2017.
[78] Greg Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. Advances in neural information processing systems, 30, 2017.
[79] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
[80] Greg Yang and Etai Litwwin Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. In International Conference on Machine Learning, pages 1 11762- 11772. PMLR, 2021.
[81] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[82] Haozhe Shan and Blake Bordelon. A theory of neural tangent kernel alignment and its influence on training. 2021.
[83] Stéphane d' Ascoli, Maria Refinetti, and Giulio Biroli. Optimal learning rate schedules in high-dimensional non-convex optimization problems, 2022.
[84] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Handerman-Milng and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[85] Juha Honkonen. Ito and stratonovich calculuses in stochastic field theory. arXiv preprint 11Xiv:1102.15811 2011.
[86] Crispin W Gardiner et al. Handbook of stochastic methods, volume 3. springer Berlin, 1985.
[87] Carl M Bender and Steven Orszag. Advanced mathematical methods for scientists and engineers Asymptotic methods and perturbation theory, volume 1. Springer Science & Business Media, 1999.
[88] John Hubbard. Calculation of partition functions. Physical Review Letters, 3(2):77, 1959.
[89] Charles Stein. A bound for the error in the normal approximation to the distribution of a sum of dependent random variables. In Proceedings of the sixth Berkeley symposium mathematical statistics and probability, volume 2: Probability theory, volume 6, pages 583-603. University of California Press, 1972.
[90] Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. arXiv preprint arXiv: 1810.055 2018.
[91] Greg Yang. Wide feedforward or recurrent neural networks of any architecture gaussian processes. Advances in Neural Information Processing Systems, 32, 2019.
[92] Adam X. Yang. Maxime Robeyns, Edward Milsom, Nandi Schoots, and Laurence Aitchison. A theory of representation learning deep neural networks gives a deep generalisation of kernel methods, 2021.
15

---

[93] Yurii E Nesterov. A method for solving the convex programming problem with convergence rate o (1/k" 2). Dokl. akad. nauk Sssr, volume 269, pages 543-547, 1983.
[94] Yurii Nesterov and Boris T Polyak. Cubic regularization of newton method and its global performance. Mathematical Programming, 108(1):177-205, 2006.
[95] Gabriel Goh. Why momentum really works. Distill, 2017.
[96] Michael Muchlebach and Michael I Jordan. Optimization with momentum: Dynamical, control- theoretic, and symplectic perspectives. Journal of Machine Learning Research, 22(73):1-50, 2021.
[97] Mehran Kardar. Statistical physics of fields. Cambridge University Press, 2007.
Checklist
1. For all authors...
(a) 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 dynamical field theory of deep networks based on kernel evolution.
(b) Did you describe the limitations of your work? [Yes] We have an explicit limitations as the last paragraph of the paper in Section 7
(c) 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.
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? [Yes]
2. If you are including theoretical results...
(a) 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 + co limit.
(b) Did you include complete proofs of all theoretical results? [Yes] All claims made in the main text are supported by derivations in the Appendix.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main experimental re- sults (either in the supplemental material or as a URL)? [Yes] Code to reproduce experimental results is provided in the supplementary material.
(b) 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
(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 DMFT as a function of width N in Figure All other runs over single wide network, where performance is predicted to concentrate over initialization.
(d) 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
you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cited the creators of
Jax, Neural Tangents, and CIFAR-10.
(b) Did you mention the license of the assets? [N/A] These are all open source provided they are appropriately credited in academic research.
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you't 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicaa ble? [N/A]
16

---

(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
17


Reviews:
[{"review_id": "ge5iRWS860y", "reviewer": "Reviewer_DJ54", "summary": "This paper analyzes feature learning in infinite-width neural networks trained with gradient descent. In this limit, the distribution over hidden unit activations and gradients in each layer becomes i.i.d. and can be characterized with some order parameters using self-consistent dynamical field theory. These order parameters are inner-product kernels that are determined self-consistently from the aforementioned distributions. These self-consistent equations can be solved exactly for linear networks, as they reduce to matrix equations in this case, or numerically for nonlinear networks. The authors demonstrate very good agreement between their theory and simulations in finite-width neural networks. Finally, limitations of the theory are discussed, showing where the theory breaks down, and a comparison to other analyses of feature learning via perturbations of the NTK is made.\n", "questions": "- Do the authors have any sense how much of the performance gap between finite-width neural networks and infinite-width neural networks is closed by the feature learning in this limit?\n- For linear networks, is the main obstacle to scaling the experiments to larger datasets in solving the linear system of Equation (12) or are there other challenges?", "limitations": "No concerns about negative societal impact.\n\n", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Understanding feature learning is arguably the most important open theoretical problem for neural networks and is certainly of interest and significance to the NeurIPS community.\n- Obviously, the work builds on previous techniques and analyses but the paper makes a strong original contribution with interesting results.\n- The discussion of connections to previous work are helpful for the reader to build intuition.\n\nWeaknesses\n- Most of the calculations are relegated to the supplement. This is probably expected given their complexity.\n- It is not completely clear what is proved rigorously and what results are formal calculations.\n- The experiments are on extremely small datasets due to the computational limitations discussed in Section 7.\n- There is little interpretation of the features that are learned in this limit.\n\nMinor typos\n- Citation on line 198.\n- \"numerically method\" on line 249.\n- \"Which\" and \"that\" are mixed up several times.\n- Some compound adjectives are missing hyphens, e.g. \"infinite width neural networks\" -> \"infinite-width neural networks\"\n- \"Figures 5 and Figure 6\" on line 165", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "PFDy_5LUVp", "reviewer": "Reviewer_mQQo", "summary": "This work proposes to simulate the dynamics of infinitely-wide neural networks under gradient descent through a self-consistent dynamical field theory. The key idea is to adopt techniques from dynamical kernels, which can govern the evolution of a neural network. The authors provide an analytic result for the linear network while providing a sampling method to approximate the dynamics. Compared to various approximation sketches, the proposed method is shown to obtain consistent solutions across different regime. Lastly, the authors provide experiments in more realistic settings which demonstrate that the method is still valid.", "questions": "We know that NTK methods cannot fully explain the full capabilities of neural networks. What makes me curious is, what is the difference between the method in this work and the finite width neural network in practical performance, especially in the feature learning area?", "limitations": "I appreciate that the authors acknowledge the limitation of expensive computation. I encourage the authors to try to further improve the computational efficiency of the algorithm in future research.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, the work is of high quality. First of all, in terms of writing, this article has a clear structure and is relatively easy to read. From a method perspective, it is a novel idea to use dynamicsl field theory to simulate the dynamics of infinitely wide neural networks, especially the feature learning area. From the simulation results, compared with other basic methods, the method proposed by the author can effectively capture the dynamics of the neural network.\n\nOn the other hand, the method proposed in this work has a relatively large limitation in computational efficiency, and there is still a certain gap from the actual neural network application scenario. However, I also agree that this limitation can be addressed in the next step.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NFRhnntgIlH", "reviewer": "Reviewer_m42e", "summary": "In this paper, the authors continues the line of research that uses dynamical mean field theory (DMFT) to study neural networks. \nUsing statistical physics heuristics, they derive self-consistent equations that governs the evolutions of the kernels, the (internal) representations and the gradients of the networks in the infinite-width feature learning regime for deep neural networks. Although these equations are too expensive to solve in the general setting, the authors propose a sampling-based approach to solve the equations numerically. They show good agreement between theory prediction and simulation. Finally, the authors have a very nice discussion explaining that their framework capture several recently developed perturbation-based approaches to understand neural networks. \n\nOverall, I find the paper interesting, insightful and have valuable contribution to the NeurIPS community. ", "questions": "- Is it possible to formulate the assumptions under which the results are theoretically sound?   \n\n- Yang&Hu (Feature Learning in Infinite-Width Neural Networks) also derives similar (recursive) equations that depend on the distribution of of the (hidden) representations and their derivatives. Followup work (EFFICIENT COMPUTATION OF DEEP NONLINEAR ∞-\nWIDTH NEURAL NETWORKS THAT LEARN FEATURES) also derives efficient approximation scheme to solve the recursive equations, which is scalable to full CIFAR10. Can you add more detailed discussion about these two papers? What are the advantages and disadvantages? Is it also possible to scale up your algorithm to full CIFAR10?\n\n", "limitations": "The paper is purely theoretical. No potential negative societal impact as I can tell. \n\nThe authors discussed the limitation regarding computation. The non-rigorousness nature of the approach should also be discussed. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths \n\n- Generating several previous DFMT related works to multiple layer networks setting\n- Technical contribution in deriving self-consistent equations \n- Proposed framework captures several related existing works, e.g. finite-size correction [26, 27] in the paper.   \n- Good agreement between theory and (small scale) simulation. \n\n\n## Weaknesses \n- The calculations are very far from rigorous. It is totally unclear that under what assumptions the results of the paper are correct. \n- The current presentation is not very friendly to readers without DMFT background. E.g. Section 3 (main theoretical contribution) is very hard to follow and hard to extract key ideas and insights behind the equations and the techniques for deriving them. Walking the readers through the deviation in the simplest possible setting (e.g. 2-layer linear networks) will be much appreciated. \n- The self-consistent equations are not very interpretable (at least at the current form). It is not clear, at least to me, what extra insights regarding feature learning (beyond linear models) can we get from those equations.\n-  Computationally, these equations are much much more expensive to solve than just training the original networks and unscalable (cubic dependence on both training steps and training samples)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept