Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks
Authors: Blake Bordelon
Cengiz Pehlevan
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Deep Learning Theory
Feature Learning
Dynamics
Ensembles
TL;DR: A theoretical analysis of finite width effects in strongly feature learning neural networks from dynamical mean field theory.

Abstract: We analyze the dynamics of finite width effects in wide but finite feature learning neural networks. Starting from a dynamical mean field theory  description of infinite width deep neural network kernel and prediction dynamics, we provide a characterization of the $\mathcal{O}(1/\sqrt{\text{width}})$ fluctuations of the DMFT order parameters over random initializations of the network weights. Our results, while perturbative in width, unlike prior analyses, are non-perturbative in the strength of feature learning. In the lazy limit of network training, all kernels are random but static in time and the prediction variance has a universal form. However, in the rich, feature learning regime, the fluctuations of the kernels and predictions are dynamically coupled with a variance that can be computed self-consistently. In two layer networks, we show how feature learning can dynamically reduce the variance of the final tangent kernel and final network predictions. We also show how initialization variance can slow down online learning in wide but finite networks. In deeper networks, kernel variance can dramatically accumulate through subsequent layers at large feature learning strengths, but feature learning continues to improve the signal-to-noise ratio of the feature kernels. In discrete time, we demonstrate that large learning rate phenomena such as edge of stability effects can be well captured by infinite width dynamics and that initialization variance can decrease dynamically. For CNNs trained on CIFAR-10, we empirically find significant corrections to both the bias and variance of network dynamics due to finite width.

Dynamies of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks
Blake Bordelon & Cengiz Pehlevan
John Paulson School of Engineering and Applied Sciences,
Center for Brain Science,
Kempner Institute for the Study of Natural & Artificial Intelligence, Harvard University
Cambridge MA, 02138
blake_ Laake_bordelong harvard. edu, cpehlevanOg harvard. edu
Abstract
We analyze the dynamics of finite width effects in wide but finite feature learning neural networks. Starting from dynamical mean field theory description of infinite width deep neural network kernel and prediction dynamics, we provide a characterization of the (1/ ((1/wwidth //width) fluctuations of the DMFT order parameters over random initializations of the network weights. Our results, while perturbative in width, unlike prior analyses, are non-perturbative learning. In the lazy limit of network training, all kernels are random but static in time and the prediction variance has a universal form. However, in the rich, feature learning regime, the fluctuations of the kernels and predictions are dynamically coupled with variance that can be computed self-consistently. networks, we show how feature learning can dynamically reduce the variance the final tangent kernel and final network predictions. We also show how initialization variance slow down online learning in wide but finite networks. In deeper networks, kernel variance can dramatically accumulate through subsequent layers large feature learning strengths, but feature learning continues improve the signal-to-noise ratio of the feature kernels. In discrete time, we demonstrate that large learning phenomena such edge of stability effects be well captured by infinite width dynamics and that initialization variance can decrease dynamically. For CNNs trained CIFAR-10 10, we empirically find significant corrections to both the bias and variance of network dynamics due to finite width.
the strength of feature
two layer
Introduction
Learning dynamics of deep neural networks are challenging analyze and understand theoretically, but recent progress has been made by studying the idealization of infinite-width networks. Two types of width limits have been especially fruitful. First, the kernel or lazy ffinite-width limit, which arises in the standard or neural tangent kernel (NTK) parameterization, gives prediction dynamies which correspond to capture adaptation of internal features in the neural network, which are thought to be crucial to the success of deep learning in practice. Alternatively, the mean field or i-parameterization allows feature learning at infinite width (6]))
linear model [1.). This limit is theoretically tractable but fails
With a set of well-defined infinite-width limits, prior theoretical works have analyzed finite networks in the NTK parameterization perturbatively. revealing that finite width both enhances the amount of feature evolution (which still small this limit) but also introduces variance in the kernels and the predictions over random initializations [50. Because of these competing effects, in some situations wider networks are better, and others wider networks perform worse [66].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this paper, we analyze finite-width network learning dynamics in the mean field parameterization. In this parameterization, wide networks are empirically observed to outperform narrow networks 7 17 18] Our results and framework provide methodology for reasoning about detrimental finite- size effects in such feature-learning neural networks. We show that observable averages involving kemnlls and predictions obey well-defined power series in inverse width even in rich training regimes. We generally observe that the leading finite-size corrections to both the bias and variance components of the square loss are increased for narrower networks, and diminish performance. Further, we show that richer networks are closer to their corresponding infinite- width mean field limit. For simple tasks and architectures the leading (11 ///ddthh) corrections to the error can be descriptive, while for large sample size or more realistic tasks, higher order corrections appear to become relevant. Concretely, our contributions are listed below:
Starting from a dynamical mean field theory (DMFT) description of infinite-width nonlinear deep neural network training dynamics, we provide complete recipe for computing fluctuation dynamics of DMFT order parameters over random network initializations during training. These include the variance of the training and test predictions and the (1//idthh) variance of feature and gradient kernels throughout training.
2. We first solve these equations for the lazy limit, where no feature learning occurs, recovering a simple differential equation which describes how prediction variance evolves during learning.
We solve for variance in the rich feature learning regime in two-layer networks and deep linear networks. We show richer nonlinear dynamics improve the signal-to-noie to-iee ratio (SNR) of kernlss and predictions, leading to closer agreement with infinite-width mean field behavior.
4. We analyze in a two-layer model why larger training set sizes in the overparameterized regime enhance wiith effects and how richer training can reduce this effect.
5. We show that large learning rate effects such as edge-of-stte -ottbbbiity [19 21 dynamies can be well captured by infinite width theory, with finite size variance accurately predicted by our theory.
6. We test our predictions in Convolutional Neural Networks (CNNs) trained on CIFAR-10 [22]. We observe that wider networks and richly trained networks have lower logit variance as predicted. However, the timescale of training dynamics is significantly altered by finite width even after ensembling. We argue that this is due to a detrimental correction to the mean dynamical NTK.
1.1 Related Works
Infinite-width networks initialization converge Gaussian process with a covariance kernel that is computed with a layerwise recursion 26 13. In the large but finite width limit, these kernels do not concentrate at each layer, but rather propagate finite-size corrections forward through the network [27 [4]. During gradient-based training with the NTK parameterization, hierarchy of differential equations have been utilized to compute small feature learning corrections to the kernel through training [IO 13). J. However the higher order tensors required to compute the theory are initialization dependent, and the theory breaks down for sufficiently rich feature learning dynamics. Various works on Bayesian deep networks have also considered fluctuations and perturbations in the Other relevant work in this domain are 33
kernels at finite width during inference
31
An alternative to standard/NTK parameterization is the mean field or P--limit where features evolve even at infinite width 6 42) Recent studies on two-layer mean field networks trained online with Gaussian data have revealed that finite networks have larger sensitivity to SGD noise 43 44). Here, we examine how finite-width neural networks are sensitive initialization noise. Prior work has studied how the weight space distribution and predictions converge to mean field dynamics with a dynamical error 201//widdts [40] 45). however in the deep case this requires a probability distribution over couplings between adjacent layers. Our analysis, by contrast, focuses on a function and kernel space picture which decouples interactions between layers at infinite width. A starting point for our present analysis of finite-width effects was a previous set of studies p 46 which identified the DMFT action corresponding to randomly initialized deep NNs which generates the distribution over kernel and network prediction dynamics. These prior works discuss the possibility of using a finite-size perturbation series but crucially failed to recognize the role of the network prediction fluctuations on the kernel fluctuations which are necessary to close the self-consistent equations in the rich regime. Using the mean field action to calculate perturbation expansion around DMFT is long celebrated technique to obtain finite size corrections in physics (47450) and has been utilized for random untrained recurrent networks [51 52 and more recently to calculate variance of

---

feature kernels & at initialization t = 0 in deep MLPs or RNNs 53 We extend these prior studies to the dynamics of training and to probe how feature learning alters finite size corrections.
2 Problem Setup
We consider wide neural networks where the number of neurons (or channels for a CNN) N in each layer large. For multi-layer perceptron (MLP), the network defined as a map from input T E RDD to hidden preactivations h E IR- in layers e E {1...., L.} and finally output Sy
f4 1 w (b()) 7N
1
b h
W W W'o(h)) h VD
VN
where 7 is scale factor that controls feature learning strength, with large y leading rich feature learning dynamics and the limit of small y + (or generally if y scales as N for a > 0 as N + Do, NTK parameterization corresponds to 2) gives lazy learning where no features are (W", W ..... w²) are optimized with gradient descent or gradient flow - 0 NN- VoL where C EE {[(x,,, 0), y/.) is a loss computed over dataset D ( /1 (z2, V22 ................................ (æ yP)}. This parameterization and learning rate scaling ensures that d. fz ON, (1) and d
learned 4 7
19..
The parameters
dr h (n))(7) initialization. This is equivalent to maximal update parameterization (8,) which can be easily extended to other architectures including neural networks with trainable bias parameters well as convolutional, recurrent, and attention layers g
3 Review of Dynamical Mean Field Theory
The width training dynamics of feature learning neural networks was described by a DMFT in 9 (6). We first review the DMFT's key concepts, before extending it to get insight into finitewwithhs. To arrive at the DMFT, one first notes that the training dynamics of such net- works can be rewritten in terms of collection of dynamical variables (or order parameters) (ec(Ju()) (t, 8), G (1, ...} 19, which include feature and gradient kernels 99 54]
N * (h),()) ((h((s)),
(4, 1 00.00 (C(s),
where (t) NN (l, the back propppropagated gradient signals. Further, for width-N networks the distribution of these dynamical variables across weight initializations (from a Gaussian distribution (0, 1)) given by p(q) x exp (NV Sq)), where the action S(q) contains interactions between neuron activations and the kernels at each layer [9]
(0)
The DMFT introduced in 99] arises in the N as
limit when p(q) is strongly peaked around the saddle point 9x where 0g 4x O. Analysis of the saddle point equations reveal that the
training dynamics of the neural network can be alternatively described by a stochastic process. A key feature of this process is that it describes the training time evolution of the distribution neuron pre-activations in each layer (informally the histogram of the elements of h (t)) where each neuron's pre-activation behaves as an i.i.d. draw from this single-site stochastic process. We denote these random processes by h (l). Kernels in (2) are now computed as averages over these infinite- wiith pu (t, 8) = (ø(h,, (t)) (s)).. G (t, s) = (gf. (t)g (s)). where averages arise from the N + ao limit of the dot products ( DMFT also provides set of self-consistent
single site processes Φ
equations that describe the complete statistics of these random processes, which depend on the kernels, as well as other quantities. To make notation and terminology clearer for a machine learning audience, we provide forr a definition of the physics terminology in machine learning language.
Order params. ? Concentrating variables
Action S(q)
Propagator Σ
Single Site Density
q log-density
Asymptotic Covariance
Neuron Marginals
Table Relationship between the physics and MI. terminology for the central objects in this paper. The q which concentrate at infinite width, but fluctuate at finite width N. This paper is primarily interested in Σ, the asymptotic covariance of the order parameters.
4 Dynamical Fluctuations Around Mean Field Theory
We are interested in going beyond the indt limit to study more realistic finite-width networks. In this regime, the order parameters fluctuate in O(L/v neighborhood of goo [53] 51 53 46).
3

---

Statistics of these fluctuations can be calculated from a general cumulant expansion (see App D 155 56 5T]. We will focus on the leading-orde order corrections to the -hnnt limit in this expansion.
Proposition 1 The finite- width
winht N average of observable (gq) across initializations, which denote by (O(g)) N. admits an expansion of the form whose leading terms are
edqexp (N)) O(g) ( + +V (O()))
(O(a)) S [dqexp (NS(g)
(3)
where 20 denotes an average over the Gaussian distribution q N (9:0- (V ")
and the function V (g) S(g) S(gs)) 3 q V qoo) contains cubic and higher terms in the Taylor expansion of S around qoo The terms shown include all the leading and sub-leading terms in the series in powers of L/N. The terms ellipses are at least O(N-1)
suppressed compared to the terms provided.
The proof of this statement is given in App. D The central object to characterize finite size effects the unperturbed covariance (the propagator): Σ [V25(q This object can be shown to capture leading order fluctuation statistics (q- q=0 (()-)x))( (q- 9=0)
x E + (App- D. which be used reason about, for example, expected square error over random initializations. Correction terms at finite width may give a possible explanation of the superior performance
of wide networks at fixed y 7 17 18] To calculate such corrections, in App. [ we provide complete description of Hessian V (g)) and its inverse (the propagator) for a depth-L network.
This description constitutes one of main results. The resulting expressions lengthy and
left to App. E. Here, we discuss them at high level. Conceptually there are two primary ingredients for obtaining the full propagator:
Hessian sub-blocks &. which describe the uncoupled variances the kernels, such as
A (,, ('," =
- (1, 8) (1'
(4)
Similar terms also appear in other studies on finite width Bayesian inference 13 and in studics on kernel variance at initialization (27 [4 29 (53).
Blocks which capture the sensitivity of field averages to pertubations of order parameters, such as a (0thE, (6(((()))((()) 04 (V,s')
pua (t,,,t a
Do' u (t, (,t.,'',))
(t.(t) : (5)
where (t) Melf 05p
(e) are error signal for each data point
Order Parameters
Abstractly, we consider the uncoupled variances "sources" of finit--widt noise for each order parameter and the blocks as summarizing directed causal graph which captures how this noise propagates in the network (through layers and network predictions). In Figure I we illustrate this graph showing directed lines that represent causal influences of order parameters on fields and vice versa. For instance, if &e were perturbed, D 4 would quantify resulting perturbation to +f+1 through the fields h&+1
Fields
G²
gf
hfr
Gf+1
off
K
In App. we calculate k and D tensors, and show how to use them calculate the propagator. As an example of our results:
Figure 1: The directed causal graph between DMFT order parameters (blue) and fields (green) defines the D tensors of our theory. Each arrow rep- resents a causal dependence. K denotes the NTK.
Proposition 2 Partition Vec(J, (1), TC (1, Vecl fu & (t, s).... Let R
into primal and conjugate variables
92
8a S(gi. 42) and Dq20q.
D = = S(q,, qzl. then the propagator for qi has the form Eqg D-*x [D-1] (App The variables qi are related to network observables, while conjugates q2 arise Lagrange multipliers in the DMFT calculation. From the propugator Egg we can read off the variance of network observables such as N Var(fu) Ela.

---

The necessary order parameters for calculating the fluctuations are obtained by solving the DMFT using numerical methods introduced in 9 We provide a pseudocode for this procedure in App. F We proceed to solve the equations defining Σ in special cases which are illuminating and numerically feasible including lazy training, two layer networks and deep linear NNs.
5 Lazy Training Limit
To gain some initial intuition about why kernel fluctuations alter learning dynamics, we first analyze the static kernel limit -> O where features are frozen. To prevent divergence of the network in this limit, we use a background subtracted function fæx 0) = fax 0) - fx, 9b) which is identically zero at initialization 4 For mean square error, the N - oo and - O limit governed by af(a)
at Eg'nD A (x') (x')K(( (xx x') with (xx) (()) f f(x) (for MSE) and K is the static (finite width and ran- dom) NTK. The finite-N initial covariance of the NTK has been analyzed in prior works [27 13 14], which reveal dependence depth and nonlinearity. Since the NTK is static in the 7 0 limit, it has constant initialization variance through training. Further, all sensitivity blocks of the Hessian involv-
ing the kernels and the prediction errors A (such as the D .A) vanish. We represent the covariance of the NTK as n( 22 222 223) NCov (ar, 222, ) K ((#3, 224). To identify the dynamics the relate K, the finite width NTK, to Ko which is the deterministic infinite
error A covariance, width NTK Koo- We consider the eigendecomposition of the infinite-width NTK (x x') = Σx A (a)))ax with respect to the training distribution D, and decompose in this basis.
Kikmmn (x( 22, 33, ze 14 (x) (z2)))))), (ma))m (a:4) m...
(6)
where averages are computed over the training distribution D.
Proposition 3 For MSE loss, the prediction error covariance E (t, NNCovo A()) satisfies a differential equation (App. H
Ak
de
E (1, 8) Σ 4 (0)
01
where (t) = exp (-xxt) (1b (Z))(()) (x))
are the errors at infinite width for eigenmode k.
Sctt) (A,1t))
335
Var 6,
1.5-
02-
-32
Ensemble (K= 1601
-1.9
0.0
0.0-
30
120
150
200
20
200
150
200
150
200
(a) Average Prediction Errors
(b) Top Eigenmodes Variance
(c) Total Train Variance
Figure 2: We show the accuracy of the lazy-limit ODE in equation (where) comapared to a two-layer finite width N 100 ReLU network trained with 7 0.05 P = 10 random training data points. (a) The average dynamics over an ensemble of E = 500 networks (solid colors) compared to the infinite width predictions (dashed black). (b) The predicted finite size variance for each eigenmode of the error A = A(t) Φy. These are not ordered simply by magnitude of eigenvalues or the target projections yk y Φk. but rather depend on all eigenvalue gaps A x for k # l and also the Kitno tensor. (c) The total variance for all training points N E, (ar N Σx Var (t) is also well predicted by the DMFT propagator equations.
An example verifying these dynamics is provided Figure 2 In the case where the target is an exp(->g-
eigenfunction = ud,+, the covariance has the form Σ (t, s) = KAek" (Ae- -A If the kernel is rank one with eigenvalue X, then the dynamies have the simple form (1, = wyy² is e -A(t+a) We note that similar terms appear the prediction dynamics obtained by truncating the Neural Tangent Hicrarchy [10] [_ however those dynamics concerned small feature learning
5

---

2011
3.5
((ott) (.(0)))
(K(E))
10
0.5
02-
124
10f
10'
1.54
180
N
(a) Deviation from DMFT vs
(b) Average NTK vs DMFT
(c) Average test prediction
0.30
P-05 0.50
0.10
N VariA)
0.08
2.00
6.00
1.00
2,00
1.75
0.04
0.02
0.25
1.00
0.00
0.000
00
100
150
205
250
5
50
100
150
200
250
150
150
700
250
(d) Train error variance
(e) Test Prediction Variance
(1) Kernel Variance Dynamies
Figure 3: An ensemble of E = 1000 two layer N = 256 tanh networks trained on single training point. Dashed black lines are DMFT predictions. (a) The square deviation from the infinite width DMFT scales O(1 (/) ) for all order parameters. (b) The ensemble average NTK (K())) (solid colors) and (c) ensemble average test point predictions (t) for a point with 2-2, = 0.5 closely follow the infinite width predictions (dashed black). (d) The variance (estimated over the ensemble) of the train error (1)
D
f(t) initially increases and then decreases as the training point fit.
(e) The variance of f. increases with time but decreases with y. (f) The variance of the NTK during feature learning experiences a transient increase before decreasing to lower value.
corrections rather than from initialization variance (App. H.1 Corrections to the mean (A) are analyzed in App. HH.2 We find that the variance and mean correction dynamics involves non-trivial coupling across eigendirections with a mixture of exponentials with timescales (A²1).
Rich Regime in Two-Layer Networks
In this section, we analyze how feature learning alters the variance through training. We show denoising effect where the signal to noise ratios of the order parameters improve with feature learning.
6.1 Kernel and Error Coupled Fluctuations on Single Training Example
In the rich regime, the kernel evolves over time but inherits fluctuations from the training errors A. To gain insight, we first study simplified setting where the data distribution is single training example æ and single test point æ, in a two layer network. We will track A(t) = ffxx and the test prediction f.(() = fxx, f). To identify the dynamics of these predictions we need the NTK K ((t) on the train point, as well as the train-test NTK K, (t). this case, all order parameters can be viewed scalar functions of single time index (unlike the deep network case, App. Eb
Proposition 4 Computing the Hessian of the DMFT action and inverting (App. obtain the following covariance for q Vec A ((t), f+(t) (t),, (x)))
T
0
4
0
0
0
I
0
0
En
θ
K.
0
A
0
+
K.
I
+
A
-
0
1
0
0
K
D
0
1
0
Dr
0
I
K.
K...
D,
0
I
where [0x](1,8) = 0(1 - (s)(s,, (1)](1 = O(1 - (s)(s) are Heaviside step functions
and ((1, 8) [2(3] (6(h(1))2 +g(()²) and (1,,,)) = (لاد)) - (Oh(1))(sh, (1) + (())..1))

---

W 500
P- 1-
P-50
P- 100
10
NY V.Va To
107
ETTOR
24
15-2
10a4
Var. Along
P- 102 Ful Var.
P- 10 Full Van
38-
A-308
33
16-
P- 500 Fut Var.
R--3033
100
400
500
230
226
566
202
425
530
206
320
460
(a) Offline prediction Variance
(b) Offline leading correction breaks down for large
W P
P=
Var 60
1ran Error
10F
Wa
1-10 $.20
10-7
(c) Large 7 reduces variance
(d) Large 7 well captured by DMFT First Order Correction
Enor
38.
Var B)
10
Mong
10-
10
33
20- (DMFT
10
21 DPST
38
39
100
200
300
500
469
500
(e) Online function variance
(f) Online leading correction breaks down for large D
Figure 4: Large input dimension or multiple samples amplify finite size effects in a simple two layer model with unstructured data. Black dashed lines are theory. (a) The variance of offline learning with P training examples in a two layer linear network. (b) The leading perturbative approximation to the train error breaks down when samples becomes comparable N. (c)-(d) Richer training reduces variance. (e)-(f) Online learning in depth 2 linear network has identical dynamics and finite width fluctuations, but with predictor variance D/N for input dimension (Appendix) K
quantify sensitivity of the kernel to perturbations in the uncoupled variances of K(t) and K..(t) (t) and k, is the uncoupled covariance of K K(t,, (t,)),,tt) (t).
signal A(s). Lastly K. and k., the
In Fig. 3 we plot the resulting theory (diagonal blocks of Eg, from Equation S for two layer neural networks. As predicted by theory, all average squared deviations from the infinite width DMFT scale as O(N ). Similarly, the average kernels (K) and test predictions (ff.) change by a larger amount for larger y (equation (79). The experimental variances also match the theory quite accurately. The variance of the train error (t)) peaks earlier and at lower value for richer training, but all variances go to zero at late time as the model approaches the interpolation condition A 0. As - the curve approaches N Var(A(())
K. y f2 2-2 where k is the initial NTK variance (see Section 55.. While the train prediction variance goes to zero, the test point prediction does not, with richer networks reaching a lower asymptotic variance. We suspect this dynamical effect could explain lower variance observed in feature learning networks compared to lazy networks [8]. In Fig. A.I we show that the reduction in variance is not due to a reduction in the uncoupled variance x(t, s), which increases in y. Rather the reduction in variance is driven by the coupling of perturbations across time.
6.2 Offline Training with Multiple Samples or Online Training in High Dimension
In this section we go beyond the single sample equations of the prior section and explore training with multiple P examples. In this case, we have training errors (t) and multiple kernel entries K (1) (App. E). Each of the enrors (l) receives (NN11/2) fluctuation, the training error En (A)) has additional variance on the order of O (). In the case of two-layer linear
7

---

2414
Vo AId FOC0)
5 X
Vovr
(a) Prediction Variance
(b) NTK variance
(c) Variance by Layer
Figure 5: Depth 4 linear network with single training point. Black dashed lines are theory. (a) The variance of the training error along the task relevant subspace. We see that unlike the two layer model, more feature learning can lead to larger peaks in the finite size variance. (b) The variance of the NTK in the task relevant subspace. When properly normalized against the square of the mean (K (K)) the final NTK variance decreases with feature learning. (c) The gap in feature kernel variance across different layers of the network is amplified by feature learning strength
networks trained on whitened data ( 2, =Suv), the equations for the propagator simplify and one can separately solve for the variance of (t)) E IR along signal direction y E R and along each of the P - orthogonal directions (App. At infinite width, the task-orthogonal component vanishes and only the signal dimension A evolves time with differential equation 46
d dt
22 (t)) (t) A (t) =.
(9)
However, at finite width, both the (1) and the P P-1 orthogonal variables A inherit initialization variance, which we represent as C', (t, s) and > (t, In Fig. (a)-(b) we show this approximate solution (Ia (1))) (y(1)² 2 (1) (1) E.A, (1, (P-1) En (t, O(N-2) across varying and varying P (see Appendix for Σs, and Σn formulas). We see that variance of train point predictions (t) increases with the total number of points despite the signal of the target vector E,, y/ being fixed. In this model, the bias correction (t) is always ((1/N) but the variance correction is OPP/N). The fluctuations along the P orthogonal directions begin to dominate the variance at large P. Fig. (b) shows that P increases, the leading order approximation breaks down as higher order terms become relevant. Analysis for online training reveals identical fluctuation statistics, but with variance that scales D/N (Appendix K verify in Figure (f)-(f).
7 Deep Networks
In networks deeper than two layers, the DMFT propagator has complicated dependence on non- diagonal (in time) entries of the feature kernels (see App. Eb. This leads to Hessian blocks with four time and four sample indices such as
8,1',
a-- (e (0(h,, (h))
rendering any numerical calculation challenging. However, in deep linear networks trained on whitened data, we can exploit the symmetry sample space and the Gaussianity of preactivation features to exactly compute derivatives without Monte Carlo sampling as we discuss in App. L. An example set results for depth network provided in Fig. 5 The variance for the feature kernels H accumulate finite size variance by layer along the forward pass and the gradient kernels &
(#)
accumulate variance on the backward pass. The SNR of the kernels Vari( improves with feature learning, suggesting that richer networks will be better modeled by their mean field limits. Examples
of the off-diagonal correlations obtained from the propagator are provided in App. Fig. A.3
8 Variance can be Small Near Edge of Stability
In this section, we move beyond the gradient flow formalism and ask what large step sizes do to finite size effects. Recent studies have identified that networks trained at large learning rates can be qualitatively different than networks the gradient flow regime, including the catapult (57] and edge of stability (EOS) phenomena [1921] In these settings, the kernel undergoes an initial scale

---

10"
100
1/25-
20
Varrr KiB
10:
10t
10-2
4.75
2012
20
10-1
3.2
10-1
10-1
(a) Mcan Kernel Dynamics
(b) Error Variance
(c) Kernel Variance
Figure 6: Edge of stability effects do not imply deviations from infinite width behavior. Black dashed lines are theory. (a) The average kernel over an ensemble of several N = 500 NNs (solid color). For small y, the kernel reaches its asymptote before hitting the edge of stability. For large 7, the kernel increases and then oscillates around 2/n.. (b)-(c) Remarkably variance due to finite size reduce during training (for smaller and larger than the critical value 1/m), showing that infinite width DMFT can be predictive of finite NNs trained with large learning rate.
growth before exhibiting either a recovery or clipping effect. In this section, we explore whether these dynamics are highly sensitive to initialization variance or if finite networks are well captured by mean field theory. Following [57], we consider two layer networks trained on a single example |x/2 = D and y 1. We use learning rate 7, and feature learning strength y. The infinite width mean field equations for the prediction ft and the kernel K2 are (App. M
f2+1 / K, + 2)²g f: K,+1 K, + 4n7 S.A. + ?2
(10)
For small 1), the equations are well approximated by the gradient flow limit and for small corresponds to discrete time linear model. For large 17y > 1. the kernel K progressively sharpens (increases in scale) until it reaches 2/17 and then oscillates around this value. It may be expected that near the EOS, the large oscillations in the kernels and predictions could lead to amplified finite size effects, however, we show in Fig. that the leading order propagator elements decrease even after reaching the EOS threshold, indicating reduced disagreement between finite and infinite width dynamics.
Finite Width Alters Bias, Training Rate, and Variance in Realistic Tasks
To analyze the effect of finite width on neural network dynamics during realistic learning tasks. we studied vanilla depth-6 ReLU CNN trained on CIFAR-10 (experimental details in App. B G.2p Fig. 7 we train an ensemble E 8 independently initialized CNNs of each width N. Wider networks not only have better performance for single model (solid), but also have lower bias (dashed), measured with ensemble averaging of the logits. Because of faster convergence of wide networks, observe wider networks have higher variance, but if we plot variance at fixed ensembled training accuracy, wider networks have consistently lower variance (Fig. d)).
We next seek an explanation for why wider networks after ensembling trains at faster rate. Theo- retically, this can be rationalized by a finite-width alteration to the ensemble averaged NTK, which governs the convergence timescale of the ensembled predictions (App. G.Ip. Our analysis in App. G.T suggests that the rate of convergence receives a finite size correction with leading correction G.2 To test this hypothesis, we fit the ensemble training loss curve to exponential function O(N-1)
R (- Rxyl) where C is a constant. We plot the fit Rx as a function of N-1 result Fig. 7 Ie). For large N, we see the leading behavior is linear in N-1 but begins to deviate at small N as quadratic function of N suggesting that second order effects become relevant around N S 100.
In App. Fig, A.4, we train a smaller subset of CIFAR-IO where we find that Ry is well approximated by a O(N-1) correction, consistent with the idea that higher sample size drives the dynamics out of the leading order picture. We also analyze the effect of on variance in this task. In App. Fig. A.5 we train N 64 models with varying 7. Increased y reduces variance of the logits and alters the representation (measured with kernel-task alignment), the training and test accuracy are roughly insensitive to the richness y in the range we considered.
10 Discussion
We studied the leading order fluctuations of kernels and predictions mean field neural networks. Feature learning dynamics reduce undesirable finite size variance, making finite networks order

---

8-
0.08
Train
Test
9.4
8: 45
N:
0.02
Me56
N-m
N::
f- 120
N-m 220
0.1
256
0.2
No 256
Epochs
Epochs
Epochs
(a) Train Loss Across Widths
(b) Test Acc. Across Widths
(c) Alignment Dynamics
1.20
11
22
0.0175
52
Test MSE Ratio
0.0150
1.35
0.0115
Train Loss
0.0100
0.6975
1.05
0.950
.... 3.057
003354
1.00
0.925
2.037 0.430W 6.5231
0.2
0.4
0.6
0.0
3.0
0.8
0.5
13
1.5
Ens. Train Acc
Rescaled Time
(d) Test MSE Ensembling Ratio
(e) Training Rates N
(l) Train Acc. Rescaled Time
Figure 7: Depth 6 CNN trained CIFAR-10 for different widths N with richness 7 0.2, E = 8 ensembles. (a)-(b) For this range of widths, we find that smaller networks perform worse in train and test error, not only in terms of the single models (solid) but also in terms of bias (dashed). The delayed training of ensembled finite width models indicates that the correction to the mean order parameters (App. G is negligible.
(c) Alignment of the average kernel to test labels is also not conserved across width. (d) The ratio of the test MSE for a single model to the ensembled logit MSE. (e) The fitted rate R of training width N models function of N We rescale the time axis by Ry to allow for a fair comparison of prediction variance for networks at comparable performance levels. (f) In rescaled time, ensembled network training losses (dashed) are coincident.
parameters closer to the infinite width limit. In several toy models, we revealed some interesting connections between the influence of feature learning, depth, sample size, and large learning rate and the variance of various DMFT order parameters. Lastly, in realistic tasks, we illustrated that bias corrections can be significant as rates of learning can be modified by width. Though our full set of equations for the leading finite size fluctuations are quite general in terms of network architecture and data structure, they are only derived at the level of rigor physics rather than a formally rigorous proof which would need several additional assumptions to make the perturbation expansion properly defined. Further, the leading terms in our perturbation series involving only Σ does not capture the complete finite size distribution defined in Eq. (3). especially as the sample size becomes comparable to the width. It would interesting to see if proportional limits the rich training regime where samples and width scale linearly can be examined dynamically (58). Future work could explore in greater detail the higher order contributions from averages involving powers of V (q) by examining cubic and higher derivatives of in Eq. 3 It could also be worth examining in future work how finite size impacts other biologically plausible learning rules, where the effective NTK can have asymmetric (over sample index) fluctuations [46). Also of interest would be computing the finite width effects in other types of architectures, including residual networks with various branch scalings [59] 60) Further, even though expect our perturbative expressions to give a precise asymptotic description of finite networks in mean field/μP, the resulting expressions are not realistically computable in deep networks trained large dataset size P for long times T since the number of Hessian entries scales as (T p4) and matrix of this size must be stored in memory and inverted in the general case. Future work could explore solveable special cases such high dimensional limits.
Code Availability
Code to reproduce the experiments in this paper is provided at https: t/gittub/.ttb..... com/ Pehlevan Schlevan-Group/aaeee /ttdddfffffuuuuationns... Details about numerical methods and computational
implementation can be found in Appendices and N
10

---

Acknowledgements
CP is supported by NSF Award DMS-2134157, NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. BB is supported by Google PhD research fellowship and NSF Award DMS- 2134157. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. The computations in this paper were run on the FASRC cluster supported by the FAS Division of Science Research Computing Group Harvard University. BB thanks Alex Atanasov, Jacob Zavatone- Veth for their comments on this manuscript and Boris Hanin, Greg Yang, Mufan Bill Li and Jeremy Cohen for helpful discussions.
References
[1] 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. Cesa- Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 8571 5571-8580. Curran Associates, Inc., 2018.
[2] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang- On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019.
[3] Jaehoon 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.
[4] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in Neural Information Processing Systems, 32, 2019.
[5] Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. In International Conference on Machine Learning, pages 1 11762- 11772. PMLR, 2021.
[6] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference Learning Theory, pages 2388-2464 PMLR, 2019.
[7] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthicu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment. 2020(11): 13301, 2020.
[8] Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-wine width neural networks. In International Conference on Machine Learning, pages 11727-11737. PMLR, 2021.
[9] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[10] Jiaoyang Huang and Horng-Tzer Yau. Dynamics deep neural networks and neural tangent hierarchy. In International conference on machine learning, pages 4542-450 PMLR, 2020.
[11] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[12] Anders Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint arXiv:2008.0867. 08675, 2020.
[13] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. Cambridge University Press Cambridge, MA. USA, 2022.
[14] Boris Hanin. Random fully connected neural networks as perturbatively solvable hierarchies. arXiv preprint 2rX0v:2204.01008, 2022.

---

[15] Sho Yaida. Meta- veta-principled family of hyperparameter scaling strategies. arXiv preprint arXiv:2210.04990 2022.
[16] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study, Advances in Nearal Information Processing Systems,
2020.
[17] 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.
[18] Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. In The Eleventh International Conference Learning Representations, 2023.
[19] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021.
[20] Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations, 2023.
[21] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. arXiv preprint arXiv:2210.04860, 2022.
[22] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar- 10 (canadian institute for advanced research).
[23] Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pages 29-53. Springer, 1996.
[24] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
[25] Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks gaussian processes. In International Conference Learning Representations, 2018
[26] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. International Conference on Learning Representations, 2018.
[27] Boris Hanin and Mihai Nica. Finite depth and width corrections the neural tangent kernel. In International Conference on Learning Representations, 2020.
[28] Boris Hanin and Mihai Nica. Products of many large random matrices and gradients in deep neural networks. Communications in Mathematical Physics, 376(1)2287-32,, 2020.
[29] Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pages 165-192. PMLR, 2020.
[30] Jacob Zavatone- Veth and Cengiz Pehlevan. Exact marginal prior distributions of finite bayesian neural networks. Advances in Neural Information Processing Systems, 34:3364-3375, 2021.
[31] Jacob Zavatone- Veeth, Abdulkadir Canatar, Ben Ruben, and Cengiz Pehlevan. Asymptotics of representation learning in finite bayesian neural networks. Advances in Neural Information Processing Systems, 34, 2021.
[32] 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),66331,, 2021.
12

---

[33] Adam X. Yang, Maxime Robeyns, Edward Milsom, Ben Anson, Nandi Schoots, and Laurence Aitchison. A theory of representation learning gives a deep generalisation of kemee methods. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 39380-39415. PMLR, 23-29 Jul 2023.
[34] Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications, 14(1):908, 2023.
[35] Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X, 11(3):031059, 2021.
[36] Jacob A Zavatone- and Cengiz Pehlevan. Depth induces scale-averaging in overparame- terized linear bayesian neural networks. 55th Asilomar Conference on Signals, Systems, and Computers, 2021.
[37] Jacob Zavatone- Veth, William I. Tong, and Cengiz. Pehlevan. Contrasting random and learned features in deep bayesian linear regression. Physical Review E, 105(6):064118, 18, 2022.
[38] Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects in finite cnns. Advances in Neural Information Processing Systems, 34, 2021.
[39] Jacob Z Veth, Abdulkadir Canatar, Benjamin S Ruben, and Cengiz Pehlevan. Asymp- totics of representation learning in finite bayesian neural networks*. Journal of Statistical Mechanics: Theory and Experiment, 2022(1 (1::11111100000000 nov 2022.
[40] Grant Rotskoff and Eric Vanden- Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.
[41] 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
[42] Dyego Araújo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.000 00/93, 2019.
[43] Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborova. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[44] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high- dimensional amp: meen-fielld dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 1199-1227. PMLR, 12-15 Jul 2023.
[45] Huy Tuan Pham and Phan-Minh Nguyen. Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training. Advances in Neural Information Processing Systems, 34:4843444444 4855, 2021.
[46] Blake Bordelon and Cengiz Pehlevan. The influence of learning rule on representation dynamics in wide neural networks. In The Eleventh International Conference on Learning Representations, 2023.
[47] Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A, 8(1):423, 1973.
[48] Moshe Moshe and Jean Zinn-Justin. Quantum field theory in the large limit: A review. Physics Reports, 385(3-6):69-22) 220,222822222 2003.
13

---

[49] Jean Zinn-Justin, Quantum field theory and critical phenomena, volume 171. Oxford university press, 2021.
[50] Carson C Chow and Michael Buice. Path integral methods for stochastic differential equations. The Journal of Mathematical Neuroscience (JMN), 5:1-35, 2015.
[51] Moritz Helias and David Dahmen. Statistical Field Theory for Neural Networks. Springer International Publishing, 2020.
[52] A Crisanti and H Sompolinsky. Path integral approach to random neural networks. Physical Review E, 98(6):0 2018.
[53] Kai Segadlo, Bastian Epping, Alexander van Meegen, David Dahmen, Michael Krämer, and Moritz Helias. Unified field theoretical approach to deep and recurrent neuronal networks. Journal of Statistical Mechanics: Theory and Experiment, 2022(10):10340) 2022.
[54] Yizhang Lou, Chris E 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.
[55] Carl M Bender and Steven Orszag. Advanced mathematical methods for scientists and engineers I: Asymptotic methods and perturbation theory, volume 1. Springer Science & Business Media, 1999.
[56] Mehran Kardar. Statisticul physics of fields. Cambridge University Press, 2007.
[57] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large lcarning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218. 2020.
[58] Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and multi-scale theory of generalization. In International Conference on Machine Learning, pages 74-84. PMLR, 2020.
[59] Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Transactions on Machine Learning Research, 2023.
[60] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer residual networks: Dynamics and scaling limit. arXiv preprint wrXiv:2309.16620, 2023.
[61] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023.
[62] Michael Baake and Ulrike Schlaegel. The peano-baker series. Proceedings of the Steklov Institute of Mathematics, 275(1):155-159, 2011.
[63] Roger Brockett. Finite dimensional linear systems. SIAM, 2015.
[64] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners:
The silent alignment effect. In International Conference Learning Representations, 2022.
[65] B. Bordelon, A. Canatar, and C. Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. International Conference of Machine Learning, 2020.
[66] Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint arXiv: 1902.002200000 02880, 2019.
14

---

Appendix
A Additional Figures
y- 0.5
35
1.0
0.06
4=4
Dit. 100)
0.04
0.02
2.0
50
150
200
250
50
100
200
250
(a) Fourth Moment <(t,t)
(b) Kernel sensitivity (tt, 100)
Figure A.1: The x and D functions for varying y in Figure 3
(a) The uncoupled kernel variance k(t, t) increases monotonically with y. This reveals that the dynamical filtering of is what is responsible for the late time decrease measures sensitivity of kernel at time & to perturbation in A at time The (tt, entries also increase with y. This suggests that the reduction in variance of the training error and the kernel are not due to reduction in but rather dynamical filtering effect due to scale growth in Ko and rapid reduction A...
variance during feature learning.
(b) The tensor Dtt, s)
13a
10a
10a
10-1
13
10"4
XV Var
10-1
13
10-4
10-
12-2
10-4
16-4
Exst
DNIT Lime
Load ng Var.
Lead 1g Var 8/6
100
200
200
400
300
100 200
300
420
530
120
220
300
400
500
Figure A.2: A comparison of the bias and variance corrections in the toy model of Figure 4 At small D/N (or P/N for offline training) the leading variance and the leading variance and leading bias both track the experiment. Both the bias and the variance contribute positively towards the total generalization error since the variance correction alone (orange) exceeds the DMFT limiting (dashed) and the variance and bias correction together (green) exceed variance alone (orange). However, for large D/N (or P/N) the leading order picture fails to deseribe the finite width experiment, indicating significant variance possibly at higher order scales (like D2 D2/N2 D D*/N3....)
15

---

y-0.55
y 1.00
V-2.8
2.75
Expt E,n((,
2.6
2.35
20-
2.00
1.00
3.00
5.00
DMFT Sw
2.002
2,00
22003
22
2,00
2.0
2.25
2.00
3.59
3.55
2.0
3.75
y-0.5
y-1.0
y-2.0
V-3.0
Expt Iaelt, s)
82-
19
5.5
33
10
35.0
45
22.
12
4,22
23
20
4.5
30
30
30
10
20
44
5.5
DMFT DMFT
4.22
4.0
4.0
Expt Zaitt,
4.0
10
35
DMFT (t.
40
1.100
32
20-4
20
10
22
15
Figure A.3: The covariance of kernel entries across pairs of time points (t,
N Cov(H t), (s, s)) for depth 4 linear network trained on whitened data. The variance becomes increasingly localized in time as feature learning y increases.
16

---

Eno. Train Loss
101
20
2x30-4
N 61
N 123
254
201
512
100
230
300
100
200
9.500 0 0.330 0.05 6,200 0,025 5/19
Steps
Steps
(a) Ens. Averaged Train Loss
Raleigh Quotient
AA A
KA
(c) Rates Function of N
Figure A.4: The ensemble averaged train loss for the same depth 6 CNN trained on random subsample P = 64 CIFAR-10 points. Training is full batch gradient descent with y = 0.05. (a) The ensemble train accuracy for this subset of CIFAR-10 is well modeled C(t x exp -Ryt) with rate Ry that depends on width. (b) The projection of the errors A on the average NTK (K) (which is related to the rate of decay of the training loss, see Appendix G reveals that wider networks are more aligned with their instantaneous error signals. (c) The rates R are indeed linear functions of N-1, with Ry R.o + It N. consistent with the average NTK (K) receiving N correction. Using ensembling to find scaling law like that above can thus allow one extrapolate the training rate infinite width mean field models.
exponential in time
0.55
0050
0.34
0.00
0.200
0.300
0.08
0.23
202 Epochs
250
Epochs
Epochs
(a) Training MSE
(h) Kernel-Task Alignment
(c) Test Ace.
1.10
Troin MSE Racio
MSE Ratio
1.08
150
200
1.00 Epochs
Epochs
(d) Single vs Ensembled Ratio
(c) Single vs Ensembled Ratio
Figure A.5: Width = 64 depth 6 CNNs trained on the full CIFAR-10 with MSE. An ensemble of
size E 10 randomly initialized networks are trained. (a) Training MSE for varying 7. (b) Final layer kernel-task alignment does strongly depend on 7, despite similar train dynamics. (c) Top-1 classification test accuracy only slightly different across 7. A small benefit from ensembling is visible late in training. (c) Initialization variance (measured by the ratio of single model to ensembled MSE) for training and test losses. Richer networks have lower variance throughout training. (b) Networks have distinct kernel dynamics when trained with different evidenced by the alignment (cosine similarity) between the final layer feature kernel L and the target test labels y.
17

---

B CIFAR-10 Experimental Details
We trained the following depth 6 CNN architecture in the mean field parameterization using FLAX [6 on single GPU. The bias parameters were zero each hidden Conv layer. but were used for the readout weights. The networks
trained with MSE loss centered 10 dimensional targets Dn E IR-0 for E [P]. Each convolution was followed by an average pooling operation. To obtain mean field behavior, NTK parameterization with modified final layer is used 97.
from flax import linen as nn import jax numpy as jnp
class CNN (nn. Module): (nn.Meee))e)):
width: int
def setup(self):
nn. wt..iniaaliiess... horizerrse.rmaalstdde.... 1.0) O_N entries
10
convi mn. Conv (features self width, kernel_init kif, use bias = False, kernel_size (3,3))
self. = mn. cm.Con
self width, ternel_init = kif, bias False, kernel_size (3,3))
12
self width, kernel_init kif, use_bias False, kernel_size (3,3))
conv3 mn. Con
self width, sernel_init kif, use_bias False, kernel_size (3,3))
self.comv = mn. Conr
14
conv. nn Conv (featuuuee self. vidth kernel_init kif, use_bias False, kernel_size (3,3))
15
self readout nn Dense (features = 10, use_bias True, kernel_init = kif)
17
18
def
(self, x, train True): N self width
19
3
self. convicoovlloo (x) jmp. sqrt
22
jnp. sap.sqrt(ce (2.0) relu
25
nn avg_pool(x, vindov_shape= strides (2,2)) # 32 x 32
-> 16 16
self celf.conv2e. (x) jmp jqp. explicit N-(-1/2}
25
jup. snp.sqrt (2.0) reluul
26
x nn avg_poll vindou_shape-(2,22 strides (2,2)) # 16 x 16
> 8 x
self. conv3 (x) ju))
jup. aqrt.qrtde (2. 0) relu
29
nn avg-pool (z, windov_shape=(22, strides (2,2))
x 4
30
self. conv4
jap sqrt
(2.0)
relu
32
nn aagavvopooll
vindow_shape=(2,25 strides (2,2)) x
> x
33
(x)
jmp. sqrt
(2. 0)
rnluelu(x)
avg_pool (I, windov_ s2ado__shape= strides (2,2)) #
-> x
36
x rerhap ((x, shape shape [0] -1)) # flatten
37
self readout (x) / N for mean field scaling
return
All models were trained with standard SGD with a batch size of 256. Each element the ensemble of E networks is trained identical batches presented in identical order. For the Figure T experimmess the raw learning rate scaled as y 10N / with 7 = 0.2 (note that mean field theory requires scaling the raw learning rate linearly with N since the raw NTK O(N-1) (9). For Figure A.5| the learning rate is 7 5N /7. We find that choosing 7 V gives approximately conserved traiiigg
18

---

times across (though distinct representation dynamics). The Figure A.4 shows the dynamics of fitting 64 training points with full batch gradient descent and y 0.1.
C Review of DMFT: Deriving the Action
In this section we derive the DMFT action which contains all of the necessary statistical information about randomly initialized finite width N networks. From the action S the DMFT saddle point and the propagator can be computed. This derivation follows closely the original derivation by Bordelon We start by writing the gradient flow dynamics weight matrices
& Pehlevan 9
P Σ ==1
W (1) dt
An(()
(O)(
(11)
VN
where 4 (t)
(f, are the enror signals and (t) signals. The prediction dynamics satisfy
Ny
are the nack-propagation 8h,,
8f.(t)
dt fult) (t
(12)
where (t) the instantaneous neural tangent kernel (NTK). At finite width N all of the above quantities depend on the precise initialization of the network. We transform the weight dynamics into an integral equation and use the recurrence for h. obtain the following
h (1) 1 w' (0))((() +7 S ds Σ
Φe mi (t,8)
VN
(t) 8(h (t)
1 w (0) (1) + 1/ ds (t,
=1
(13)
VN
where we introduced the feature and gradient kernels
1
Ln (1,
N 6 (1) (s),, Ge (1, ND (s).
(14)
Written this way, that the source of the disorder which depends the initial random weights (0) comes through the fields
++ (1) = VN (O() (1)
('
(15)
VN
If can characterize the distribution of the fields (t) and (t). then consequently characterize the distribution of (t). (t. We therefore choose to study the moment generating functional
exp
Σ Ey
dt [1)(()) x(( + of,()) xx
(16)
O.
Moments of these fields can be computed through differentiation with respect to the sources j. near zero-source (j = V = 0)
(x (t.) (tm) (t.)...ffm (tm))
8 8
5
ZK35, o' Flb====0-
(17)
8jth (t1) 814, (ta) doft (t1)
8offm (tm)
To average over the initial weights, introduce Fourier representation of the Dirac-Dekh Delta function dads 2n exp(i2z). We perform this transformation for each of the fields to enforce
1 = J dzo(2)=
19

---

their definition
= dxt((t
(1 W (0))((((0))
21.11 (2.(0) VN
(2TT)N exp
VN
(t)
(t) (0)
1 VN
(t)()) (5,(0)
('
9p c (00
(2T) exp
VN
(18)
We insert these Dirac delta functions so that directly average over the weights
In E 10) exp (o) / dt E (t)o + 9 -ftt
(t)
- Σ Σ / dids [xx* (1= xx=1(0)0 (1,8) ()) (t)
Σ f dids( ...ggg."" (8)) (+))
(19)
where we introduced the kernels कh G. We next introduce the order parameter
1 (1,8) N (h (1)) (s)
(20)
To enforce the definitions of the new order parameters { G, A) we again introduce Dirac-delta functions
(N (l, ((, ((h(()))
= -f d (t, exp _ (1 [N (L (1,s) o(((() ((,(())))
(21)
2Tí
Analogous constraints for G and A are enforced with conjugate variables G, B. After introducing these variables, we find that the moment generating functional has the form
d& ue (t, s)d (t, 8) de s)GG (t, di (n a (t, s)dB (t, 8)
2 =
II Ejerts
2ri
2xi
2Ti
exp (NSH(T &6 ",,,",""" AA,
(22)
where S is the (1)) DMFT action which defines the statistical distribution over the dynamics. The action takes the form
f
S
Σ Eys
dt ds (1 (1) (1, Afr (1, (B
N In Zeltil
(23)
E=1
where Z, is the single site stochastic process for layer É which defines the marginal distribution of X.E. with the following form
(1) (()) //I -
II
d (1)xx (1) dét (1)
( (/
exp
2T
2T
1
Σ
dtds (ts) ( (2
exp
2
Σ I
dtds Bí. st) (40 (000000)0000011
exp
exp
dt
(24)
20

---

where in the above, the {h. g) fields should be regarded as functionals of {x.E). At zero source je. + 0 this function S can be regarded as the log density for the complete collection of order parameters = 1 4, G. G, A, B} which collectively control the dynamics. Concretely, we have that p(q) x exp (N)( In the next section we explore an approximation scheme for averages over this distribution at large N.
D Cumulant Expansion of Observables
We are interested in a principled power series expansion (in 1/N) of any observable average (O(g))
that depends on DMFT order parameters q. At any width N the observable average takes the form
(O(g)) = _ dxx exp (())) O(q)
(25)
dq exp (NS(g))
As discussed in the main text, the + oo limit gives (0(g)) ~ O(gs)) where as
By gm =
by a steepest descent argument [35]. We assume that S's Hessian is negative semidefinite so that Σ = (2S(g)) 9- - 0 and Taylor expand S(q) around the saddle point q giving S()) = S(gx)) 5 qox) T>> 7 V(q q=o). We note that the remainder function V contains only cubic and higher powers of qx = 6/vN. The variable ō will be order O(1). This will allow us verify that additional terms suppressed in powers of 1/N. Expanding both the numerator and denominator's 's integrands in powers of V, we find
(O())) ddqexp (- N (q 9x0) 1
(xx) + NV -=0)) 0gg)
dd exp (- 2
N (q 900) - 1- - 96) NV 9:0))
S dō exp (- 1 1- ()
Nº V2 00. + +N --1/28) f expexp 1- (1+ [V+NV NV + -// V2
NV (V+
(0) (VO) 0 2
N² (V20)
x2 (V80) +...
1+ N (V) + 2 (V²) 8c
N
+ 3
(V3) +...
1 IVO (0) + N²
(V20) 110)
+ Na उ.
(V'0)</10)) +...
= (0)
1 N (V)
No
(V2)
No
(V3)
(26)
where (. represents an average over the Gaussian fluctuation N 9x0, -
(25(g
We
see that the series in the denominator contains terms of the form N² (V²) while the numerator (v*0)=o/ (0)) In either of these power series, the k-th term AT
depends on terms of the form N²
can contribute at most
N (V* N² (V*) (0)
SO(N (++1/)) (((2+122)
h odd h even
(27)
since V contributes only cubic and higher terms. Thus each term the numerator and denominator's series contains increasing powers of 1/N. Concretely, each of the two series have terms of order N,N N
Thus any quantity of the form TO)= admits ratio of power series in powers of 1/N. One could truncate each of the series in the numerator and denominator to desired order in N. Alternatively, the denominator could be expanded giving single series (the
(0)
cumulant expansion [56]). The first few terms in the cumulant expansion have the form
(0) N =00) + +NOV) - -0) (V)ce 2
+ 2 (120) - (VV)) (V) + 2 (0), (V3) (0)²) +
(28)
In this work, we mainly are interested in the leading order correction to (0) which can always be obtained with the truncation after the terms linear in V for any observable O.
21

---

D.1 Square Deviation from DMFT
We will now analyze the fluctuation statistics of our order parameters around the saddle point ((q - 9:0) - 900)") )r which has the form
((9-900)1 q 900) + + (V 90-9 )00))) +...
(la- -
xxx) -
N+N(V), +...
- I 1+0((--))
1 N E+
(29)
as stated in the main text and verified empirically in Figure numerator involving V can be larger than ( comes from vanishing of odd moments for - 9:c in the unperturbed distribution. Thus the leading expression for ((q 9xx 9x0)" only
3 (a). The reason that the terms in the
depends on E and not on V.
D.2 Mean Deviation from DMFT
Although the square displacement from DMFT only depended on E and not on V, we note that the average order parameter displacement (q - q:0) does receive O(1/N) correction that depends on the perturbed potential V
(9 g00) + N N(q 900 V) + Nº +/4 -a--900 Vo) +..
(q- --0O
1+ 1 N (V) + (V2) +... 0V
Σ
+N
0q
E
+ +NN-3)).
(30)
1+
0q
where in the last line we used Stein's lemma (Gaussian integration by parts) for the Gaussian av distribution over q. Note that 84
O ( since the derivative of the cubic term in gives quadratic function of q - 9oo whose average must be O(N-1). this work, we focus primarily on the structure of the propagator, but outline general recipe for getting the leading mean correction in Appendix G and H.2
D.3
Covariance of Order Parameters
Lastly, we combine the previous two observations to reason about the scaling of the order parameter covariance over initializations. We note that the leading covariance of the order parameters over random initializations is also given by the propagator: Cov(g) J O(N-?), since
Cov(g)
(g (q) (q (4) (4)x)
(q goo) (a -oo) ) ( (g (9) (q)
Σ O
(31)
- and that go0 - (q)xx O(. Therefore, in the leading order picture, it is safe associate with the covariance of order parameters over random initializations of the network weights.
due to the arguments above which showed that
E Propagator Structure for the full DMFT Action
In this section, we examine the propagator structure for the full DMFT action. This action is modified from other prior works (9 46 to include the evolution of the network prediction errors A(t). Those prior works noted that A and the NTK are deterministic functions of deterministic order parameters (&L G+) in the N + oo limit so those authors did not explicitly include A or K in the action. At finite width N. including A, K in the action crucial as the fluctuation in prediction errors A has significant consequences for dynamical fluctuations of kernels through the preactivation and
22

---

pre-gradient fields. In this section, we will mainly focus on gradient flow, but we describe large step size Appendix M
S= Σ
/
dtds [ (t)
(() (t,, + Gun (t, st (t, (t, ( A Aeu (t (t,
Σ / (t) Ax(0) 1:(0)-9a- yu + Σ ds - 5) (s)
Σ / dl Kmu(t) - Σ (1)0
M4M
+
In Zz A $6 Ge G" k-
1 A A B
(32)
where the single site moment generating functionals 2y have the form
Ze -((g(())))))))) exp (-E)
/ dtds
k7((0 +7) ['¹ ds Σ (s.s
GP(0,)
=1(00 -2((0) +7 +7 1 ds Σ G (1,8) + B (1, ((1(()))) (1000 GP G
(33)
with 00,1)) ((())==((0(()))(((00) The saddle point equations give the infinite width evolution of our order
parameters.
0s
T (1, (00()0)))(((())))
06 f.o (t,s)
0S
(t, (s)()))((ss) =0
aG (t,s)
as
oo(h) 0 Or((s)
0At (s, 222B( BB (1, +
0s
0gf,(t)
-72 +
= 0
YB (s,tt)
Oul (s)
as
(t) - - Get (t, (t, (t,tt
0k
as
(t) 3p 1'" ds Σ (s)
(34)
a8,(t)
These equations exactly recover the mean field description obtained 99. Note that () for field averages an average defined by Z, and distinct from the types averages () we have been considering over the order parameters q. The complementary set of equations for the primal variables, such as as De (t.s) 0, give that = = = & 0 at the saddle point. We now set out to compute the
Hessan V To simplify the set of expressions, we will only explicitly write out the nonvanishing
23

---

blocks. We will start with second derivatives involving only pairs of dual variables { G, A, B)
05 O. (u, s), ae (') - (0001100000550000000000000000000
(c, (4)
(1,
05
(11009900)))) (
a6t Mn (1, 8)0 ds (U,8'
n (t,
0S
- of (1, (u Gas))
ahe jio (t, S) de (t,,s')
= (t, s,t's'
02S
(0(11, (t))
(1h (H((()))
aB ne (t, s)aA 1-00A--110 (s'
Ouiz(t))
Box
((L (0())(())
sa(l))
Oulg(s")
oCh, (0( (g(((1)
2 fu (t, () (t, (t, 5)
0ul(())
8²B² st,
= Ym
( (0(0
025
abe me (1,8)BBB
3a 11,1)
Ori(s')
o(h((tt) 2o(h)(()) Org(s') ((ha(tt)
(011(10)66(((001 00(1)(())
2 (t, (l
Org(s')
Th 2'd's (1)
05 age (t, (t, ( 8
( (m) 96 oolooso.())
(0)
( dgt(t)
(
Ba (s, (s)
()))())))))) (g
22 (tu 8) B) (1,
0ul/(5)/ 7t p
G²B (t,s)
0S
(g
? (2a00) agt(t)
Doito (
(1(
00 (1, (),B Ba (8,11)
ory(st)
(1))))(())) (o(h)((1))
7 G ce (t, s)
0r)(s')
= yk C'd's (t,
82
2 bhadngdd 11)8(s-5)
a (1, B) Ba (8',1)
025
02 Dul (s)
(000006)0))) -20B// (1,8) A)aaaaaaa ("')8))
1s, t)Bt
Ba (8',8)
(1, (1,8,1'88) = k "va
B
(35)
24

---

Next, we consider the second derivatives involving only primal variables K, which all vanish
05 ahe mu (1, 8)0 e
a8 (1' ,8)
08
dG mu (t, s)OG 0g (t', 5
OS
0
aa. (t, S)OG ag (t', s')
02S
0
O.. (t, s)0K (55)
02S
Ju (t, s)0K (0Ka(s)))
02S
0
Ob (t, S)0A, (s')
02S
Ju (t, s) (8))
025
0
OK (t)aK a. (8)
02S
0K pDu (t)as, (s)
0PS
04 (t) (8) 0
(36)
Now we consider all derivatives which involve one of the dual variables ( Ge, AC, ,B and the primal variable A
02S
16( (6(C) =- (t,
Jhe uu (1, 8)0A((1)
(1)
025
(1) 1o (g((())(())))
= -D pva (t, (t,
OGt uu (1,8))))))) 8)0Aa(1)
SS
JA, (V)(()) S 12(1) DDB i. (t, 8,t')
a
AALT (a, 1)0Aa(1)
0
(t))rrr ((h,(0)) = 7D4's
uura (t, s,
DB (s, t) (t')
25

---

Now, we consider the second derivatives involving one derivative on dual variable ( A.B. and one of the primal variables
8
pp 8(t t')((- 5s)
2 (t, s)a a0 (t',5')
2e 08 (1, (0()(())((((())))
d2,0 (11 1' (')(( - 5') - - Davas (1, (1,
05
0 d(1 (')o(s 5') -d4+1e 0t (di(1)5(())
86 ju (t, 8)0G
az 81)
= (')0(s - 8)) - -. Dvag (1,8,1,88)
05 a
0
(6(H((0)))((()))) -D Get
uvaB st,s,t',s))
(1, 8)00
08 (U,
8G (t', s
0
(g)g)))(()))
DD -0- wua8
(t,),t'8')
O0 (1, (1 aB (U,
05 1-1 aB (t', 8)
0PS
(1)
= DBB 1,0e umod
(t,
aA (s, t) 1-6--111
ag st',))
=7 1 (t' s')
(s)
0S
((b(M,(11)
DD' .0--1 peap (1,),t's8)
0 (s, t)aa 1-1
up og (t', s²)
=7 1apt-11
aB
(t', s')
@uk(s)
05
0g%,(t)
DDBB- ,t+1 (1,8,1',5')
DA
(s, 03 (V'
4+1
7 00 (1',5)
(r)
=
7DD
uvaB
S
Y acttt (1'
(o(h,
Dd puas
(t,s,t','')
(37)
B (8, 1) a3 8'
@uf(s)
We note that terms such 34 as (V'8) (thh can be further decomposed since the
average over the (1) ~ GP(O, xe-1) and 's explicit dynamics both depend on Φ'-1
aB (1,)) (6(h), (0())))))))((((0)))))))))))(())) 2
0
dug (1))o((())))) o(h ((L()))((()))
0²
00 (1, 2( ((,)))))))))
(38)
where the first term comes from differentiating the Gaussian probability density for u² (e.g. Price's theorem) and the second term is an explicit derivative of the preactivation fields with u' treated as constant. Next we consider the nonvanishing terms which involve {Ó, K, A, K) which give
O2S
5pab(t - 5) + O - 5) ()))(())))
a (t) (s)
02
- (pa
0Â (t)8
025
8node88011 - (')
0k n0 (t)aKap(1))
0P
= dun a/
Get (1, - - ss)
aK nN (t)O 0g (1,88)
05
6 8- (t, s')((t +¹)5(t - 5s)
(39)
0k (t) (t'
26

---

This enumerates all possible non-vanishing terms in the Hessian. We can now construct block
matrix of these Hessians by partitioning our order parameters q [q q2l where
T Vec ( ( 8); (l,(l,)), s), K (1), p (1, s), 6 (1, s), (1);
(40)
42 = Vee ( (A (1, 8), Bj (1, (1,8)}-
(41)
This choice will become apparent shortly. V
V2
s
V2 V
1 42
S
V 91
(42)
S
To calculate the full propagator E = - [
we will assume invertibility of the upper block
E = V2 S] 1 and use this in the Schur complement
Σ [
E1
Σ12 E22
Σ21
Σ11 Ed ~ [ 91 S| (V 92 S + (V (g S)E (V 8))
8) [ 22 " r
E Σn E [2 91 $1 (V2 0:
S (V2 22: (V ())
-1
Eg (V2 S (V v= 5) ~ V 11 5))
(43)
To perform this inverse, we again partition di into two sets of order parameters QI [g] 9i where of = Vec pe (t, s), G (1, s), kw (1)
We now need to solve for Σ = [2 q: S]
and qी = Vec ue (t, 8), G nu (t, 8), Kn (1), (t)
0
V U
2 K
S
U
qq s
(44)
K
q
We seek a physically sensible inverse where the variance of qf is vanishing 53. This leads to the following sub-propagator E
Σ - 12, - U-I
- ---
-U-1
(45)
Thus given K. U, we can solve for E and ultimately for the full propagator Σ. The relevant entries in & and U are given by those second derivatives caleulated above. We note that each of the field derivatives needed for U can be computed implicitly from the field dynamics. For example, for the A (1) derivatives we have
① (t) (L(() = 70(t (') HV (t, (t))
+7 1/' ds Σ [A (t, 8) + + (t,s) (g((s) dD (t')
0 (t') (t) = 70(t - (t,t))
+r 1/" ds Σ (t, + G (t, 8) [s)]
(p(((())) (t.
(46)
These can then be used in the averages such as ( (2s ohh (t) ()))))(()))))))))))))))))))))))) Similarly, we can compute terms such σ4 "s
al through the following closed equations
(t)
70(1- 0(1 O(L- (')
OΦ ng
+7 l' ds Σ (1.) (1,6]]
0gf(()
De (
(1)
S'' ds Σ
(0()((s))
(s) w0 [1)
(47)
0$ f-1
09 (t',') s')
04 p (t,
These terms can then be used to compute quantities like a
27

---

F Solving for the Propagator
In this section we sketch out the required steps to obtain the propagator Σ.
Step 1: Solve the infinite width DMFT equations for Too which include the prediction error dynamics A,( (f). the feature kernels f., s), gradient kernels G (t, s). This step corresponds to algorithm in Bordelon & Pehlevan '2 and defines the dynamics one would expect at infinite width (9. See below for more detail.
Step 2: Compute the entries of the Hessian of S evaluated at the Qoo computed the first step. Some of these entries look like fourth cumulants of features like = ((()))))))))))) and some of them measure sensitivity of one order parameter to a perturbation in another order parameter D 8 (d) The averages () used to calculate k and D should be performed over the infinite width stochastic processes for preactivations he which are defined in equation (19).
Step 3: After populating the entries of the block matrix for the Hesssian V we then calculate the propagator Σ with a matrix inversion. Since we discretized time, this is a finite dimensional matrix.
The step above demands solution to the infinite width DMFT equations (solving for the saddle point (sx). We will now give detailed set of instructions about how the infinite width limit for 9xc is solved (step I above). This corresponds to the algorithm of Bordelon & Pehlevan 2022 to solve the saddle point equations 0 (gg) 9x = 0 19]
0g
Step I: Start with guess for the kernels (, 8), G (t, (t. and for the predictions through time
(t). We usually use the lazy limit (e.g. (i (1, 8) = & (0, ...) as an initial guess.
Step 2: Sample Gaussian sources (t) and (((()) (1) based on the current covariances &f and G
Step 3: For each sample, solve integral equations for h()) and =())
kt((t) (t) + L. ds (t,
t/ (t. ()))((((s)))))
=((t=)((t +7 ' ds EiB (t, + Gt s) (s)(((s)
(48)
These will be samples from the single site distribution for z
Step 4: Average over the Monte Carlo samples to produce a new estimate of the kernels: (t, s) (1)))(1)) (s))). A similar procedure is performed for G and the response functions A B'.
Step 5: Compute the NTK estimate K(t) E, G (t, t) t) and then integrate prediction dynamics from the dynamics of the NTK (x()) E, K (t).
Repeat steps 2-5 until the order parameters converge.
Below we provide pseudocode algorithm solve for the propagator elements.
Algorithm 1: Propagator Solver
Data: K y, Initial Guesses { G. HEL (A BOL- Sample count S, Update Speed Result: Propagator Matrix Σ
Solve DMFT equations with Algorithm for order parameters (t), & (a s)....
2 Draw samples {uf (t) ~ GP & " fre p.n (t)} ~ GP(O, G++1);
Integrate dynamics for each sample to get {hy f (t), 2 (t)}
Estimate functions with Monte Carlo integration, for instance
5 ve "noogtt, st,,''
3 Ene[s] (1))(k (()))))
- Df ff (t, s) (t';
an a and kernels
For each sample, compute field sensitivities to error signals, such as JA,(s)
ah", (t)
0ol (t" implicitly using equations 46 47
Use these sensitivities to compute the necessary D tensors such
a DELA Ene(s) (.. (O) (s)]:
Invert U matrix and compute Eg in equation 45
Compute the Schur-complement in equation
43 to handle the response functions: :
28

---

The above propagator solver builds on the solution to the DMFT equations which is provided below.
Algorithm 2: Alternating Monte Carlo Solution to Saddle Point Equations
Data: K", y, Initial Guesses { G'V
===1" {A B Sample count S, Update Speed
Result: Final Kernels {Φ² G 12-1: {A BE I=1 Network predictions through training fa
T = @ 11T, G 11
while Kernels Not Converged do
From ( G compute K KNTK (t, t) and solve $fu(t) = Eo (a (t,t);
{= 1; while 8 do
Draw 8 samples (t),n(t) GP(O, " (n)((t)) ( GP(O, x1)
Integrate dynamics for each sample to get ( 2/10 (t))
Compute new tº G¹ estimates:
Th (t,s) Ene(s] ( (t)) (s)). (t, 1/2 Ene(s) (t)n())) ())))n(s))))
Solve for Jacobians on each sample (oth.) 0g Br
10
Bu
Compute new A B estimates:
11
Á' $ EnE(s)
0e(h Gr
B
EnE[[S]
Bg
12
ou
8+ l + 1;
end
15
t= 1;
16
while 8 < ++ do
Update feature kernels: o6 <- (1 3) 3 G + (1- (1-3)G'+ BG:; if L then Update A t (1 - A + BA B + (1 8) B + 3B end
17
18
19
21
t 1-Et1 ++1
22
end
end
return H (A', B (f,(0)
24
G Leading Correction to the Mean Order Parameters
In this section we use the propagator structure derived in the last section to reason about the leading finite size correction to (q) at width N. Letting the indices é, j, k, n enumerate all entries of the order parameters in q (technically this is sum samples and an integral over time for gradient flow), we find the leading Pade Approximant for the mean has the form (App
N ((g: - ) x + N²
2 ((9. 9: VV2)
x0
1 + N (V) + N2 (V2)
00
05
(5,65ōk51) O(N-2).
(49)
3IN
dq Oq 0q:
jkl
SSS
Eig En + O(N-2)
(50)
2N
dq: jkl
where 6; VN -55)) and the derivatives are computed at the saddle point. In the last line, we a² utilized Wick's theorem and the permutation symmetry of the third derivative Ba, to evaluate the four point averages in terms of the propagator Eij. which was provided in the preceding section practice computing even the full set of second derivatives
the DMFT action to get Σ quite challenging. Despite the challenge of computing the mean order parameter correction, these corrections are relevant in practice and crucially distinguish the training timescales of deep networks at different widths as we show in Figures and A.4
29

---

G.1 Correction to Mean Predictions and Full MSE Correction
Supposing that we solved for the propagator Σ, using the formalism in the preceeding section, we can compute the O(N-)) conrection to the average network prediction enror due to finite size. We let AA (t)) represent the average of errors over an ensemble of width N networks.
d dt (A)( = -- (K (t))v(t))
= Σ (K)) (A.(1))
(Ka ((())
1 Σ (K) (Av(t)
+ O(N-3)
(51)
N
where E (t,t) is the leading covariance (propagator element) between the kernel (t)(t) and prediction enror (t). We see that the average kernel ( (t)) (which depends the finite width N) plays an important role in characterizing the timescales of the average prediction dynamics. Once this equation is solved for (A,(t)), the square loss width N and time & has the form
Σ (A)())) ~
~ (1 (1-
2 N
E (t)
2 ((An(t) N
0c 47
(t)
1 Σ N
(t))
E (t,t) O
(52)
We will now comment on the structure of the cross term in this above solution. First, if (K) > K and y is negligible then the average errors at finite width will decay more rapidly than the infinite width model. However, we suspect that in general, (K) K contains many negative eigenvalues since signal propagation at finite width tends to reduce the scale of feature kernels [14. We suspect that this is the cause of the slower dynamics of ensembled predictors for narrower networks in Figure 7 and Figure A.4 Additionally, the term involving ΣイΛ will generically increase the cross term since the dynamics of A cause its fluctuations to become anti-correlated with the fluctuations in K. In general, it is challenging to make strong definitive statements about the relative scale of these competing effects on the cross term. However, we can say more about this solution in the lazy limit, where we find that the cross term will generically be positive, leading to larger MSE (Appendix
G.2 Perturbation Theory in Rates rather than Predictions
In experiments deep CNNs trained on CIFAR-10 in 7 and A.4. we find that the loss curves for the ensemble averaged predictors are effectively time rescaled by a function of network width. In this section, argue that a proper way to account for this is to compute perturbation expansion in the exponent which defines the rate of decay of the training errors. To illustrate the point, we first consider the case of single training example before describing larger datasets. In this case, consider the change of variables A(t) -r(yg. y. We now treat an order parameter of the theory with dynamics
d T.
K(1)
(53)
Note that this equation is now linear relation between two order parameters (r(t),((tt) whereas the relation was previously quadratic. In the lazy limit, if K + K e then r - r ct, giving an effective rescaling of training time by 1 R.
For multiple training examples, we introduce the notion of a transition matrix T(t) E PxP which has dynamics
a T())=-K())T()) T(O) = I. di
(54)
The solution to the training prediction errors can be obtained at any time / by multiplying the initial condition A (0) y with the transition matrix A (1) = T()),, where y are the training targets. In this case, the relevant rate matrix, which would be an alternative order parameter
R(t)= = log (t))
(55)
30

---

where log is the matrix logarithm function. Note that in general (t) admits Peano-Bake series solution [62 64). In the special case where (t) commutes with (t) 1 .o ds<(s), obtain the following simplified formula for the rate matrix R
R(t) [ ds (()
(56)
The benefit of this representation is the elimination of coupled order parameter dynamics which are quadratic in fluctuations (in A and K) into linear dynamical relation between order parameters R and K. An expansion
R will thus give better predictions at long times than a direct expansion in A. In the lazy 7 + O limit, the constancy of K (() K gives the further simplification R = K Working with this representation, we have the following finite width expression for the training loss
(1a()))) = y (oxp (oxp(-2R(())) y
RR(())) y
1
~ exp
-2 (t) +
N
O²
+ 2 st t)
Σ
DR
y
exp (-2R) y n- + +(N-22 (57)
RRy
jwa
where (R) ~ R.x + + 2) is the leading correction to the mean R. In this representation, it is clear that finite width can alter the timescale of the dynamics through a correction the of R. as well as contribute an additive correction from fluctuations. This justifies the study perturbation analysis of rates Ry as function of 1/N in Figures 7 and A.4
H Variance in the Lazy Limit
We can simplify the propagator equations in the lazy - 0 limit. To demonstrate how to use formalism, we go through the complete process of inverting the Hessian, however, for this case, this procedure is bit cumbersome. A simplified derivation for the lazy limit can be found below in section H.T which relies only on linearizing the dynamics around the infinite width solution. In the y + 0 limit, all of the D tensors vanish and the tensors constant in time. Thus, it suffices to analyze the kemels restricted to = 0 and study the evolution of the prediction variance A (l)..
dt (att) Aa(()) Bu + / ds>ore O(t K)
S
Se De
+ Ge
Gh Gre Gfrr] + EKhu
Ruu Kmn Σ
EEot Getlos of
+ + In Ze
Z: E (ubhirfft exp
(-I
th (in
EGno Ju
Sf. = -r((((())))
(58)
where {u1} ~ Ni0 x 1, ~N(0, ((0,Ge++1 give terms of the form
Taking two derivatives with respect to ( G+)
Te Te ue ax
K (DlAioiigh)) G G.
= (o(u())(u)))) a
le
Ge
(59)
a3
31

---

Given these we also have the relevant non-vanishing sensitivity tensors
Da 0²
DG uuag = aG (ohol)
a
a
a3
θ
Davas
(9/90)
(60)
6 PupG DRaa
uuaB
Dude (1) dsO(l - 50)nn
(61)
As before we let 91 Vec An Tf uas Gue Kew] and 92 Vecf $ He Go propagator has the form
Kno }. The
I I+
DAK
0
U 7 1291
==
0
0
I D
D
I- DGG
0
V
9292
-=
0
K
0,0
G
G
GG
0
0
0
- Dkk
KG
I
0
0
0.
(62)
The propagator of interest is En = =U V find an inverse
S
U-IT.
We can exploit the block structure of U
9292
U AA
U
U
AG
AK 0
U
U-1
(63)
G4
U@
0
UK
UKG
I
where each sub-block can be computed with the Schur-complement formula. Altogether, we multiply through get the propagator
U- U
U +U1 GG U U T-I D
U
Σ
0
U-1 to
UGG -1 Ko
U + UGd
0
0
KKK
U U-I 4G
+ KG
0
U
0
0
0
U
U+1
UUGhT T
x
(64)
U
AG
0
U
AK
0
Two of these blocks corresponding to K, A are especially important for characterizing the fluctuations of network predictions. The covariance structure for K has the form
ΣK = U 4:4 UN! + U
[RO] + URb 4G
URL + U GG CURR
(65)
Next we use the fact that U -Unk Unt K& and that UNG U block structure of U. Consequently we arrive at the identity
AK KG. which follows from the
UR!
Es- - 34h U + UNd UNI +U
G
A
+U C U
= - Σ U
(66)
Lastly, we note that, by the Schur-complement formula that U - - 0x+ DAK Thus,
writing (I + ①K) Es (I + Ok)T DAK Σx ( as an integral equation, we find
DA (1, + ('' di' d'''n Kya DA (1' + 1-¹ ds' Σ
Kus
3
(6¹ dt' 1 ds' Kna s') = Σ E/ Aa((') S ds's ds' ds'
(67)
ag
32

---

Diffcrentiation with respect to and gives a simple differential equation
0²
Ötds
2 (t,8)
+
Σ
K
pa Os
E (t, s) +
Σ Kus
Ot (t,
K na ve (t,
(
(68)
Let { be the eigenvectors of the kernel matrix K. Projecting these dynamics on the eigenspace [te(t, 8) = // [t, s) recovers the equation in the main text
(0r
0 + Av 0t
8 Os
+ xe)
s)
Σ (A((t)Anss)itee
(69)
k'8'
Replacing EK recovers the equation 7 in the main text.
H.1 Perturbed Linear System
In this section, we provide a simpler derivation of the lazy limit training error variance dynamics. In this case, merely perturb the dynamics around its infinite width value A Ax (t) c (t) and = Ko + ex and keep terms only linear these perturbations. The perturbation k is fixed in time and the dynamics of e (1) are
d (t) -ox
(t) EK (t)
(70)
di
Projecting this equation the eigenspace of Koc gives
Σ EZ)
(71)
dt
This immediately recovers the final result of the last section
)
a +
= ( + +xx))
(0+ 2o ())
N
Ak
ot
An
(G)
0i
00
Σ Re
(72)
Qualitatively, the process of computing this linear correction (in <k) to the dynamics of A is identical the argument utilized in prior work perturbative feature learning corrections [11]. In that context, the perturbation is caused by small amounts of feature learning, rather than initialization fluctuations.
H.2 Mean Prediction Error Correction in the Lazy Limit
Using similar heuristic as in the preceeding section, consider the correction to the mean predictor (t)) in the lazy limit. Taylor expanding (A(t)) in powers of 1/N, we find
d
1
dt (a))) = di
00 (t)
+
A
(t) +..
d
((K(K -K 0 00
K +
(A
00 +A*))
K0
(A A
((K Ko)) A x
{(K
K
A A²)
~ K
K
1
(E
(E)EA) +0(N-2) (73)
N
N
From the previous section we have that
d à --= KA
J
(1)==
ds dsexp (-K s)) k exp (- kmss)
(74)
33

---

Projecting these dynamics onto the eigenspace of the kernel gives
Aet c
Ve ( (d + Nx ak(()= dt
(t)
Σ
K ERE
(75)
where should be seen as the limit where 0ς - de of the above. Thus we find that the leading mean correction to the error solves the following differential equation
cd - e
E
Kkeyce Cheeeee-eee
Art EE EEKee
ye'
d - -x
e- Art
Σve ye Art [-K + E
+
KK "kllll
ye
(76)
de
A
We see that at late sufficiently large t, that the terms involving L will dominate. We can gain more intuition by considering the special case of single training data point where the mean error correction has the form
d di 1) A - ye -ve-Ae[ [-K + +
(1) - [[-
+ 1 b 2 +2
e-
(A(t)?) ~ (t) + 2y -2
-K' -k'+ - + (t +O(N=3
2
(1)² + N f 221 [-K' K2 + 1d O(N-3
(77)
While the term involving Σ is positive for all t, K' could be positive or negative for a given K¹
architecture. If K' is positive, then MSE is initially improved at early times but after t > Σ the MSE is worse than the infinite width. On the other hand, if K is negative (as we suspect is typically the case), then the MSE will strictly decrease with network width for any time l.
I Two Layer Equations and Time/Time Diagonal
In this section, we analyze two layer networks in greater detail. Unlike the deep network case, two layer networks can analyzed on the time-time diagonal: ie the dynamics only depend on ( t) and Gtt, t) rather than on all possible off- off-diagonal pairs of time points. Further, there are no response functions A Be which complicate the recipe for calculating the propagator (Appendix
I.1 Single Training Point
For two layer network trained on single training point with norm constraint |00| D, we have the following DMFT action
((), (()))(()))))
(78)
= ff dt [(t)((()+ + )(t) (A()) リー ((t)-v-1
ds -))0)))))))(()))
+ 2[2,, [), 2 Enng exp (-) [ +
The saddle point equations are
as
= K(t) ((o(h(t))² + =0
(K())
as
A(t) y f dss
0å(t)
as
R(s) + A(s)) / ((t))(( 5)
OK(s)
0S
A(S) + K((s) / dt Á(D)(( 3) 0
(79)
0A(3)
34

---

From these equations, we can compute the entries in the Hessian of the DMFT action S. Letting
(t))
[ (t)
((t) = ((t)
and ((t)
K(O)
0
(a()) = 0
025
+
O(t Pit-ap s)(ss) 5(t 8)
aq(t)
(
TD(A) (o(((t))))
0
0
O
(80)
(
k(t,,)
where xit, s) (()(h(t)) ((t)) (())))))))) the NTK's fourth cumulant.
We now vectorize order parameters time (ec(g(t) LER and q (tc) and express the full Hessian
S =
aq
8
= - = a000
8²s )-1
aqa4
02 8 )-1
0q04
0q0q
028
3-1
8' Baja
8 BQ84
8 g
0
ōgeq
(81)
The covariance matrix of interest (for (t)) is thus
I+ O O
0
I O
O I
Σg
(82)
K
D
where [0x] O(L (s)) and (I,](1,)) O 8) as) The above equations allow
one to use the infinite width DMFT dynamics for k(tt. A (t) to compute the finite size fluctuation dynamics of the kernel K and the error signal A.
I.1.1 Computing Field Sensitivities
In this section, we compute (t,, s) by solving for the sensitivity of order parameters. We start with the DMFT field equations
h((1 u + y / ds =(1) + L. dsL()))(h(())
(83)
Now, differentiating both sides with respect A(s') gives
@h(t) 70(1- - 5')g('') + [ ds 0g(s)
(A(s')
0A(s')
dz(t)
do(h(s))
(A(s)) (0 - -())('(((())))) +7 (s)(s) (A(s')
(84)
We can compute D Monte carlo by iteratively solving the above equations for each sampled trajectory
(h(1), x 65 (6). Averaging the necessary fields over the Monte Carlo samples will give us the final expressions for stt,s).
a (AA(s) (O(((t))+
(tt,s)
(85)
Similarly, the uncoupled kernel variance x(t, s) can be evaluated via Monte Carlo sampling for nonlincar networks.
I.2 Test Point Fluctuation Dynamics
We now are position to calculate the test/train kernel and test prediction fluctuations. To do this systematically, we augment S with the test point prediction f. and field h, and introduce the kerel
35

---

(.(t) (tt) g(t)g, The test prediction f. and field h, have dynamics
(t) = uh 7 L (s)()))(( (m-u) K
5 = (k,())a(t), K.(t) (o(h(())0(44))) +g(t)g,(0))
(86)
The augmented action for this DMFT has the form
S / dt ft()) [+(()) dsO(t (30((=))((6))()))))) + (.(()K)())
+ / dt Â(1) (A(1) + |ds ())((())())(())) + / R(OK((
+ - R (100000000000000000000000 +(00)())))))))) We let q(t) = (A(t), (t) ((), K.)
+ Bexp
-
(87)
0
0
T+
0
0
0
0
θ K.
1
-
Sq..
V Sl..9l
0
0
K
D
0
0
LO
Ki
D.
0
I
D(1,8)
(0(1(1))2 +1(1))))
(88)
(A(s)
D,(1,8)= =
(89)
(s)
Our total covariance matrix / propagator is thus
+
OK
0
O
0
1
0
0
0
0
I
O K θ
0
O
-IT
4 K.
I
0
0
0
I
0
-OA
Σ
D
0
0
0
K.
T
D
1
D.
0
I
[
0
K++
- D.
0
0
I
(90)
This is the equation provided in the main text Equation 8
I.3
Two Layer Linear Network Closed Form
For a linear network on single data point, we can compute D(t, s) and k(t, analytically, We start from the field equations
dh(()) >)(1)=(1) dz(!)
7)(1)(())
(91)
dt
dt
We make change of variables (1) (h(1) + =(()) and v1) v2 (h(1) 2(1). We note that 4 (0) ( + r) and v0) J/2 r) are independent Gaussians. These functions u4 (1), u- (t) satisfy dynamics
de (t()) (t), dv_(t) -t)((t)u_(t)
dt
dt
U+ = exp ( dsA(s) u4(0) = (t)
= -t)+()) -
dA(s)
= v- exp (~ l dsa(s) v- - @v-(1)
(t)0(t -8) (92)
0A(s)
36

---

Now, we use the fact that (0) = /2(u normal random variables to compute K(1) - + + = (v+(1)² + (1)
1 (u + and e- (0) = v2 (u
(2(1- ru) are independent standard
a
s(,,)) a5(s)
(n(1)² + +z(()) = (()²)))) (v-( (()-s) s)
27 I doL(o)) - exp (-27) dos(s)) ]
= =27 exp
0 - -)
(93)
This operator is causal (D),,,) 0 for > 1) as expected and vanishes as t - 0. If we take - 0. we have (tt, 8) -> which agrees with our reasoning that fields h, 2 only depend on A in the feature learning regime. Since all fields are Gaussian in the linear network case, we can use Wick's theorem to obtain the exact uncoupled kernel variance in the two layer case.
x(t, 8) ((h(1)² + =00))) + + K(I)((s) + 2 (2(00)))))++(()))))))))))))+))++
+2(()))))))))))
= (u+ (s) +
(s) - - (t)())u(s))
(94)
The o+ (t) functions are those given above. Using the fact that (v+ (u_(0)?) allows us to easily compute the single site average above.
J Multiple Samples with Whitened Data
In this section, we analyze the role that sample number plays in dynamics in simplified model of two layer linear network trained on whitened data. Concretely, we assume that &e the 6pp. The field equations for preactivations (t) and pregradients 2t) obey
d h(((())))))))){{)))))) = 1h.,())=~A<0000)) -(,())),,, d =(1) ==1() = 7 (00Om
(95)
dt
We will assume the targets have unit norm 1/ 1 and we define the projection of A onto the target as (t) y (t). The other P - 1 orthogonal components are denoted A Ai(t) so that A Ay (t)y + A (t) with A (t) y = At infinite width, A_ and our field equations become
d (,),,())))))))) (1 0, h1 N( 1)
dt ht)(t)= (t))())
(96)
dt
However, at finite width N. the off-target predictions A fluctuate over random initialization. To model all of the fluctuations simultaneously, we consider the following action
(f (AISaa(atinn)))))
5 dt (A)())( - yu) + In Expexp
s -
(97)
which enforces the constraint that (t) = yu - - (x(t))a (t)) at infinite width. The Hessian over
order parameters q = (t)) (1)} has the form
V2 S
(II+
( D)
Dmu((,
a
=(0),,(1)
(98)
D+D
(s)
We thus get the following covariance for predictions Es = (7I+ -'R compute the necessary components of the tensor
K [(>I D)
D) -1] We now
(t) 0(t (00(t=s)(s))))))))) +7 1 dt' (t)),(t) dz((t)
(s)
0s.(s)
dz(t)
70(1 - s))b(() +7 [ du' "'AA(" (t'
(s)
- -70(t-s)hu(s) +7 S dt' dt',, (t') (hy (A)(s)
=
(99)
37

---

In the last line, we used the fact that these equations are to be evaluated at the mean field infinite width
stochastic process where (t) = O. To compute the sensitivity tensor D, we find the following equations for our correlators of interest:
(t)
(5) 2t))
S s) (2(s))(t)) p. H.v #y
Oz(t)
(t)
70(t - s)omr H, 4y
(100)
(8)
(t)
70(t () (z(s))(t)) + +7 S. dt' (tt')
(2)((5) (2S)()2(()) z(1)
Oz(t')
2(t)
(s)(s))
(z(
(t") ( (((1))
70(t - 8) (1/(6)) +7 di'
(8)
We therefore see that the components of D decouple over indices. In the direction, we have the following equations
Oz(() Py((, 8) = (s)(s) 2(1) + ( (t)())
(h)
(101)
where the correlators must be solved senf-consist We will provide this solution in one moment, but first, will look at the orthogonal directions. For the P - orthogonal directions, we obtain the explicit formula for D in each of these directions
())()) 2(1) + 0z(1) (s) h (1)
Oh (1)
(1,
70 - (=())=()) + (0(t 8)
(102)
Now, we return Dy. To solve these equations we utilize the change of variables employed in single sample case (1) vz (t) + z(t)), = (t) z()) (see Appendix L.3 This orthogonal transformation decouples the dynamics
d =0,00)= (1) = ())))))0
(103)
As a consequence, the field derivatives close
0Ay(s) 70(t s)0+ (s) + f dt' d'')yt') 0u'(t')
0u+ (t)
0Dy(s)
@v_(t)
70 s)w [ di' (t)) (t)
(104)
(s,(s)
(Dy(s)
The correlator of interest is
1
(h,(())(()) = ([u
2/ (1)2 - 0- (t)2)
(105)
So we get that
(t, 2 (s)
(1+ - -- (1)
0c, (t)
do (t)
(4)
v_
(106)
0A,
Similarly, can derive the on on-target and off-target uncoupled variances Nyy (t, s) and ki (t, s), which satisfy
(1, (v+(1)) (s) + + (1)= (s)) + (1+(1)44 (s) v- =1))
1
Ki (t, = (v+ (t)u (s) + (t)
(107)
2
38

---

Using these functions, we arrive at the following variance for each of the P dimensions
Σa, = - Dy) Kyy (7I + Dy)
Σa. +>I+DD))+ m_ (>I D1)
(108)
Using the fact that all A variables are independent and identically distributed under the leading order picture, the expected training loss has the form
2
N (t,t) (P (P-1) (t. OON (109)
(0)4()+)))+
where (Ay A (1)()) + O(N We note that the bias correction if (NN while the variance is (P N). We compare the above leading order theory with and without the bias correction in Appendix Figure
K Online Learning
Our technology for computing finite size effects can easily be translated setting where the neural network is trained in an online fashion, disregarding the effect of SGD noise. At each step, we compute the gradient over the full data distribution p(a)). Focusing MSE loss, we study the following equation
d A(x,t) = -Ex K (2x 1)
(110)
dt
where K(x, :t; is the dynamic NTK and A(x, t) = y(x) - f(x, t) is the prediction error. In uncountable set of possible inputs To remedy this, we utilize a countable orthonormal basis of functions for the data distribution (xx(()) For example, if p(æ) were the isotropic Gaussian density for A/ I). then tt2 could be Hermite polynomials. We expand and K in this basis vh. and arrive at the following differential equation
general the distribution involves integration
d dt A+())= -- Σ (et(t)
(111)
By orthonormality, the average turned into a sum over all possible orthonormal functions {th) }. We note that since K is evolving in time, there is not generally fixed basis of functions that diagonalize K, resulting in the couplings across eigenmodes in Equation (IIT). Since, in online learning, there is distinction between the training and test distribution,
error of interest is simply C(t = Σ (t)?. To obtain the finite size corrections to this quantity, we compute the joint propagator for all variables {Kee(t), Ax (t)). If we wanted to pursue perturbation theory in rates (Appendix G.2 we could again define a transition matrix T and rate matrix R(() as
R(t) log T(t), dt (t) = Σ KRa (t)Teee)) TRe(O) = dee k
(112)
We can then obtain A = (()()))))) where y/k = En (a (x))(x)) Since R has a finite size mean correction and finite size fluctuations, too does the error A+ and the loss L (Appendix G.2).
K.1 Two Layer Networks
In the two layer case, instead of tracking kernels, we could instead deal with the distribution over read-in vectors w E R and readout scalars a E R as in the original works on mean field networks 666. When training on the population risk equations for x (O,
d dt w al
A ( x) = E. 0s(x)
o(w x)) + EA(x) (A()))(w) w x)w
dx
(113)
di
The action has the form
-)x))) + In Basw exp (J (tdnn((t m) (114)
s 5
39

---

The Hessian over (t), ((t)) is
0
I+ D
S I+DA
(115)
R
where (t 8; = (ن(.. ( ( ((t))(w((1) (2))))))))))))))))))))))))))(()) We can use the following implicit rule
@a(1)
= 70(t s)p(z)))(())) - x)) + t. S at' x')((w a' x')x'
Our(t)
(5(
Oww(t)
7011 115-sspppe s))(())(())((()))))))) x)x
(8)
+ 7B S/" d'' x')
Ba(t') (A a) s(w x') + a(t')o(w x') aw((t) x' 0s x)
(116)
The above equations could be solved and then used to compute Da (t, az 8, x') which must then be inverted to get the observed prediction variance.
K.2
Linear Activations
Using the ideas in the preceding sections, we can make more progress in the case of two layer linear network in the online learning setting. The key idea is to track the kernel and prediction error projections onto the space of linear functions. this case we get the following DMFT over the order
parameter ((t)= N W a E D.
d at) 7B, - B(t) w((t
dt
d
dt wt)) (a(t)(3, B(t))
B(t) = (a(t)w(tt)
(117)
At infinite width, we see that the dynamics can be reduced to tracking the projection the weights and B on the B, direction. The D 1 off tar-tt dimensions vanish B1 (t) = 0. At infinite width, we arrive at the alignment dynamics studied in prior work [64]]] 9
d ((t = (()(33 B(t))
d
M(I) = (B, - 3())) + + - 3(1))
di
+ +222(99 - B(1) 3()I
(118)
We note that ((t) = (t)),,, and that M has only one special eigenvector 3, with eigenvalue m, m,()). (t). It thus suffices to track evolution this single direction
d dt (() m. (0)a)) 3(t)), dt
d m. ())))=)))(((( B(t)
(119)
We note that this equation is identical the differential equation for single training example in Appendix J Here 8, - (()) plays the role of (t) and m. (f) plays the role of the kernel K, (t). d
A key observation is the conservation law 4-2 d (t)) di (t)?. from which it follows that
(t)² - 4 3())(() 9
di 3(1 2 B(t))
(120)
This is identical to the differential equations for single sample (producing prediction [() and kernel K (()) if the following substitutions are made
f(t) + 3(1), K(t) ++ m. (t)
(121)
40

---

We now proceed to compute finite size corrections starting from the action
$ $ / (t,((t) ((t) + Exxp (-) dtB(tt) w(t)a(t)
(122)
The necessary ingredients are
x(t, s) (a(()a()))w(()x((s) -
= (a(t)a(s)) (w(t) ((ss (a(s)w(t)) (a(t)w(())) E R DXD
(123)
Similarly we have compute the sensitivity tensor
a D(1,8) (B((s)T (())(()) (((rs)))(()001)) E E
(124)
We start from the dynamics
d w(() (3, B(t)), dt a B(t])),,,((t,,,
d (B, 3(t)) w(t)
(125)
dt
Next, we have to calculate causal derivatives for fields
aB(s) (1) -70(t - s)a(s)I+ L'¹ di' - 3((')) Da(t')
(3(())
( a(t) = - - s)w(s) + L. dt, B(t)) 0B(s)
Ow(t')
(126)
Following an identical argument as in we see that D has block diagonal structure with Da, (t, s) on the T direction and Dj s) in any of the D remaining directions
(3(s) w()) DR, (1,, (
(0) DI
(s) (()))1(()))) (1)
(127)
Similarly, ((t,)) 8) has similar decomposition
K69, (1 (a())a(s) (wa, (fws, (8)) + (a(s)mz, (1)) (a(t)ws,
ki s) (a(t)a(s)) (wi(()w_(())) (s)) + (a(s)wi(t) (a(t)w_(s))
(128)
The processes have the following equations at infinite width
d
_ ws, (1,(1) (1) di = ra(1)(3. - 3(1), di at) = yws, (v) - B(1)
di w (1)
(129)
As consequence we note that (w (I)a(()))) = 0 so that ii (1,, 8) = (a(I)a()) Letting v4 (1) v (wa, (t) +at)) and v- (t) J (wg. (t) + a (t)), find the same decoupled stochastic processes in Appendix
=
dt (t) = r(3. B(t))u+ (t).
dt v- --(3-------- - (((())=-((BB(()oo__( ((t))u-(t)
(130)
We can use these equations to perform the necessary averages for Kgg and Dzs, Lastly, we use
( w (t) = - s)a(s)
(131)
evaluate D1(t,s) s). The observed covariances just
Σg -7I Da.) )a.)_ Kg. (>I D2.)- Σ1 = (7I -11)- -((1))==( -1)- (132)
We note that these expressions are identical to those in Appendix Junder the substitution 3, - B(t) (t) and D > P. Thus the expected test risk is
(I)(1)) (.l3) (B(1) B.)2 +
1 Eg. (1,1) N
1
(D (1-1)
(1, O(N-3)
(133)
N
This recovers the variance obtained in the multiple-sample whitened data case
41

---

K.3
Connections to Offline Learning in Linear Model
The finite size variance of generalization error in an online learning setting with linear target function 3 x has an identical form as the model described above. In this setting, we sample infinitely many fresh data points x N(. I) at each step leading to the flow ät
Remark
d (t((t)
(a,(()),, and d dt [t) = (t) (x))))) The order parameter of interest in this setting is (()) = Ei w. (t). The precise correspondence between this setting and the offline setting summarized in T We note that this argument could be extended to higher degree monomial activations as well, at the cost of tracking higher degree tensors (eg for quadratic activations
M -- - 0, E DxD is sufficient).
Setting
Order Params.
Target
Off-target Dims.
Loss
Variance
Infinite Quantity
Offline
A A-y-f
P-1
Train
O
D
Online
3.-B
3.
D-1
Test
O
P
Table 2: Summary of the equivalence between the leading I/N correction in the offline setting and the online setting for two layer linear networks. In the offline training setting, the order parameters are the errors A y f E Rf while in the online case they are B. B E RD.
As in the offline case, in Fig. 4 (c) and (d) we see that the variance contribution to test loss | 3__ 1 increases with input dimension D. We note that this perturbative effect to the loss dynamics is reminiscent of the deviations from mean field behavior studied in SGD [43 44], though this present work concerns fluctuations driven by initialization variance rather than stochastic sampling of data. In Fig. 4 (e) we show that richer networks have lower variance at fixed N. Similarly, leading order theory for richer networks more accurately captures their dynamics as D/N increases (Fig. (f)).
L Deep Linear Networks
For deep linear networks, the fields hf, (t) are Gaussian and have the following self-consistent equations
1 ds A (t, (ss) ( wf((1) GP(O,
(2(() = (4,(() -
(f(() = (t) + S." ds Σ (t, + (t.s)) hE()), r((()) GP(00.
(134)
where H H1r (1, ((1(()))()) and (1, (6f.0)))((600066 and A (1,
on (e: 3..(2)
and
Be (1,
(0
9 Therefore, express the action as a differentiable function of the
dru(s)
order parameters by integrating over the Gaussian field distribution. For concreteness, we vectorize our fields over time and samples h (e (mE[)))E_)))) g Vcc((ff,tt) we consider the contribution of single hidden layer.
Ze = / di' dg' dh dg' exp
1 h
h + ih' (h² - -C'g¹) - 1 h it'
2
g وو + ig' (g) D'h') - (g
exp
where Che((,,) = 70(1 - 8) [4
[h (1, + uu
(1,
8) (r,))) and (1, = 70(1
s) (t, + Ght (t, (s) Performing the joint Gaussian integrals over (h', g. of, h² h'.g") we find
0
I
DET
1 In det
0
G
CC
I
In Z, 2
(135)
I
C
Σn
0 Σ,
I
0
42

---

We can then automatically differentiate the DMFT action to get the propagator. For example, for three layer linear network, the full DMFT action has the form
5= Tr
H
H'
+
H
G
G
TTAB
2
H
0
I
-DTT
0
G
-C
I
det
2
I
-C
11TT
0
D
I
0
G
H
0
I
-D
det
0
G
-C
I
(136)
2
-
C
H
0
D2
I
0
11T
where C = ~ and = @ H + and D = 70, o G + and D = 7
This above example can be extended to deeper networks. The total size of the block matrices which compute determinants over is 4PT x 4PT for dataset of size trained for T steps.
M Discrete Time Dynamics and Edge of Stability Effects
Large step size effects can induce qualitatively different dynamics in neural network training. For instance, if the step size exceeds that required for linear stability with the initial kernel, the kernel can decrease in order to stabilize the dynamics [57]. Alternatively, during training the kernel may exhibit "progressive sharpening" phase where its top eigenvalue grows before reaching stability bound set by the learning rate [99]. It is therefore well motivated to study how dynamics in this regime alter finite size effects in neural networks. We will first solve special model which was considered in prior work [ a two layer linear network trained on single training point. We will then provide the full DMFT equations for the discrete time case and provide an outline for how one could obtain finite size effects in that picture.
M.1 Two Layer Linear Equations
In two layer linear network, the DMFT equations are
h(t + 1) h(t) + m)(())=(t) + = =(t) m))(t)h(t
S(() (=(()(())
(137)
The NTK has the form K(I) = (I((1)² +2(1))) We can easily show that the kernel and error have coupled dynamics
J(L 1) -[(() (M(1) +001) A(I + )())(1) (1(()=(1)
= =s(() (1))(((())()) + +,²)"((1))())
(138)
K(t+ + 1) = ((t)+ 4m)))) (h(t)((t))++ ( ((t) (h(t)² + +
= ((I)) + +m77 + +))))
(139)
These equations define the infinite width evolution of A(t) and K ((t). Already at this level of analysis, we can reason about the evolution of K (t). In the small r) limit, we could disregard terms order (17)) and arrive at the following gradient flow approximation for K(t) ~ 2 2/11 + x2 ft)² (9). This evolution will not reach the edge of stability provided that < /I+ッ,5= For large and 1, this
leads to the constraint 177 < 1. However, if r exceeds this bound, the gradient flow approximation is no longer reasonable and the system reaches an edge of stability effect as shown in Figure 6
To calculate the finite size effects, need to compute and D(1, 8) eل( (N(1)² + =()))). To
a
evaluate these quantities utilize the same change of variables employed in Appendix [3 In discrete time, these decoupled equations are
(t + 1) (t) + At)u,, v (t et()) - 7))(()) (t).
(140)
43

---

Given (t), these can be expressed as linear systems of equations. Now, we can easily compute the uncoupled kernel variance
<(4, 8) - +
(2()))())))(())+)+(((((()()))))) +(u+(1)0+ (8) (1)) + (1)(0))++ v- (1)v (s)(ss))
(141)
Similarly, we can calculate (1,, 8) by using the fact (1(1)2 + +1(1)) = (0+ (1) + =--())²)
((t,s) = 2 (t( Be+ (1)
v- 0u (1)
+
(
(A(s)
Ov,(t) (t) = 70(t - s)u (8) + Σ (t (t)
dA(s)
0A(s)
Mit
@v-(t) - (s) Σ
A(') (t)
(142)
0s(s)
01(s)
P<t
These can be directly solved as linear system of equations.
N Computing Details
Experiments for Figures and were conducted on a Google Colab GPU with JAX. Experiments for Figures 5 A.3 were performed on a NVIDIA -000-80BB GPU. The total compute required for all Figures in the paper took around 4 hours. Jupyter Notebooks to reproduce plots can be found at htupttpss//ggieee
44


Reviews:
[{"review_id": "T8Lxj2cns9", "reviewer": "Reviewer_VrD5", "summary": "The paper addresses the problem of analytical description of the rich (feature learning) dynamics of neural networks. To achieve this, the authors use previously introduced dynamical mean field theory (DMFT), which identifies several key characteristics of the problem - order parameters - defines their probability distribution throughout the training using a path integral representation. The paper considers $\\mu P$ parametrization of the network, which is known to display feature learning behavior at infinite width (in contrast to popular NTK and Standard parametrizations) with feature learning strength controlled by parameter $\\gamma$.\n\nIn the paper, the authors focus on leading width $O(N^{-1})$ corrections to the infinite width limit. Technically, this is achieved by taking into account quadratic (Gaussian) fluctuation around saddle point $\\mathbf{q}_\\infty$ of DMFT action $S(\\mathbf{q})$. After deriving general equations governing $O(N^{-1})$ corrections, the authors consider a number of simplified scenarios where these equations can be solved analytically. Also, the authors validate qualitative conclusions of their theory in the non-synthetic experiments with CNN trained on CIFAR10.  ", "questions": "* What dataset size $P$ were used for main CIFAR10 experiment of fig. 10? Due to mentioned $O(P^4T^4)$ memory requirements of storing the propagators, it seems to be impossible to work with the full ($P=50000$) CIFAR10 size. Also, what is the computational complexity of solving DMFT equations? Both for saddle point $\\mathbf{q}_\\infty$ and $O(N^{-1})$ corrections if they are different.\n* While in most of the presented plots DMFT accurately describes the experiments, for the big sample sizes (fig. 3) there is a significant disagreement between theory and experiment. Do you think this a fundamental limitation DMFT on the level of leading order correction (ignoring $O(N^{-2})$ and beyond)? Or maybe it is because in the experiment you considered whitened data, whereas in the more realistic scenarios, the data typically has a low effective dimension (e.g. measured by the decay of data covariance eigenvalues)? ", "limitations": "The authors discuss the limitations of their current approach. The limitations mainly come from 1) the inability to numerically solve DMFT equations for large-scale problems (e.g. due to $O(P^4 T^4)$ propagator size) and  2) the need to consider higher order expansion around the saddle point $\\mathbf{q}_\\infty$. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The approach chosen by the authors - identifying saddle point $\\mathbf{q}_\\infty$ of the system's action $S(\\mathbf{q})$ and investigating it together with Gaussian fluctuation around $\\mathbf{q}_\\infty$ - is a workhorse of various branches of theoretical physics where behavior of a complex system needs to be analyzed. In physics, this approach not only provides a very sizable portion of available analytical results but, for example, also provides SOTA results for numerical modeling of realistic strongly-correlated materials [Kotliar 2004](https://pubs.aip.org/physicstoday/article-abstract/57/3/53/755526/Strongly-Correlated-Materials-Insights-From?redirectedFrom=fulltext), [Vollhardt 2019](https://arxiv.org/abs/1910.12650). Thus, realizing this general strategy is one of the fundamental directions within deep learning theory. Also, as noted by the authors, the DMFT approach is non-perturbative w.r.t. feature learning strength  - a feature that is mostly absent in other approaches to NN dynamics away from kernel regime.", "weaknesses": "* In many cases, solving DMFT equations analytically seems to be intractable. This significantly renders the main purpose of the proposed theory - obtaining analytical insights into the network dynamics.\n* The DMFT equations are bulky, which could make working with them quite exhaustive.\n* I believe it is hard to understand main DMFT ingredients - order parameters $\\mathbf{q}$ and their action $S(\\mathbf{q})$ - from the current paper alone. Most probably, a careful reading of the original DMFT paper [9] is required to understand this paper. However, this is not the author's fault but rather a consequence of the chosen approach. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful reading and supportive comments. Below we address the weaknesses pointed out and attempt to answer the reviewer's questions.\n\n**Response to Weaknesses**\n\nThe reviewer is correct to point out that the DMFT equations are difficult to solve numerically and that the equations are bulky in the general case. We acknowledge these issues in the limitations section and will stress them further in the updated draft. Despite their complexity, we think it is still useful to derive the equations in full generality to show that, in principle, the finite size corrections to DMFT can be computed with similar methods to compute the infinite limit. \n\nThat being said, we think that a good deal of intuition can be gained by studying the special cases that are tractable numerically to analyze how various aspects of the problem (feature learning, sample size, depth etc) alter the dynamics and fluctuations. \n\nSince some of the reviewers point out that we use physics jargon without comparison to ML terminology to describe some of the objects that appear in our study, we will add a table (see attached PDF in global response) which compares the terms (order parameter, propagator, action etc) which show up in the calculation to terms more familiar to a ML audience (concentrating variable, asymptotic covariance, log density, etc). \n\nWe will also add a short self-contained primer and derivation of the DMFT action in the Appendix so that readers do not need to refer to [9] to understand the present paper.  \n\n**Response to Questions**\n1. We apologize for not being clear about this section. The CIFAR-10 plots are purely empirical. In these plots, we do not attempt solving the DMFT equations or the propagator for this setting as both timesteps and samples are prohibitively large. The purpose of this section was to illustrate that the qualitative findings from the simpler analytically tractable models carry over in more realistic cases, specifically the accumulation of finite size effects over training and that the finite size effects depend on a low degree polynomial in $1/N$. We will be sure to clarify this in the draft. \n2. This is a great question. We are mainly showing in this toy example (which has a very simplified whitened data model) that the leading order corrections can be accurate for some problem settings ($P<N$ in this problem) but can underestimate finite size effects in other settings ($P>N$). We suspect that in this example, higher order corrections (like $P^2/N^2, ...$ etc) may be necessary to accurately capture it. We agree with the reviewer's intuition about the correlations in natural data could be reducing the scale of finite size noise in more realistic settings, making finite size networks closer to the infinite width limit. For instance, if the data matrix were low rank with rank $P_{eff} < P$ then the finite width effect in this two-layer, linear network example is actually $P_{eff} / N$. We will add a proof of this and discuss it in the main text. We leave open for future works how the realistic structure of natural data alters the scale of finite width corrections."}, {"review_id": "tbH5pqd0yr", "reviewer": "Reviewer_2Zma", "summary": "Building on past work which set up a DMFT (Dynamical Mean Field Theory) for fully connected networks in the infinite width limit (where the width of each layer tends to infinity), this paper reasons about the *fluctuations* around the infinite width limit. This is important because for finite sized neural networks, these fluctuations are large enough that they are an important part of the network training dynamics. The description of the DMFT is theoretically complicated and cannot be solved exactly, but they enable simulations which confirm that the DMFT captures the behaviour of finite sized real networks quite accurately. As an application, this DMFT theory is used to understand bias, training rates, and variance in realistic tasks and the special case of 2-layer networks (where the theory is quite a bit simpler) is investigated in more detail. This theory has the potential to open the door to many potential future uses that explain how neural networks learn. ", "questions": "* Suggestion: Eqn (3) seems like a very important main result: a bit more discussion on the next order term and why its actually size 1/N (even though it starts with an N) would be helpful. \n* I am wondering if you had any numerical simulations where you checked the rate of covergence between the theory here and simulations (which should presumably be like 1/N^2?). Something like Figure 2a but comparing to the fluctuation predictions rather than to $q_\\infty$.\n* One high-level question I had: It seems that if the learning rate is fixed and not scaled, the the fluctuations and the effect of the learning rate are both on the same scale $1/\\sqrt{N}$. 1. Is that correct? and 2. Does your theory work to analyze what's going on in that case?\n\nHere is a list of other minor errors/suggestions I found while reading:]\n* Eqn (2): use definition equal with three lines to be consistent with definitions later on\n* Line 179 vs 185: Is there a difference between $K$ and $K_\\infty$? If so what is it?\n* Line 188: What does the subscript 0 in \"$Cov_0$\" mean here?\n*  Section 6.1: I was able to more-or-less piece this together, but I think it would be a lot more understandable if you gave explit definitions for $K(t)$ and $K_\\ast(t)$\n* Line 232-235: I think it would be a lot clearer to write out the definitions here of the two new $\\Delta$'s in terms of $\\Delta_\\mu$. \n* References [15],[21],[33], [46], [53] has only author/title but not where published/where to find.\n* Reference [22] missing a title?\n* I would also check the arxiv only referenes e.g. [11],[13],[20],[27],[34],[57],[63] to make sure there isnt a conference or journal version that is now published.", "limitations": "A potential criticism is the physics level of rigor used in Appendix C which is used to establish the main results. The manipulations carried out in the proof of Appendix C certainly seem plausible, and I believe the community as a whole is ok with this level of rigour, but the authors could be a bit more clear about what they mean by \"proof\" in the main paper. It is not a mathematically rigour proof (which would involve all sorts of techincal assumptions), but rather a physics-type statement that holds assuming the usual expansions can be carried out without obstructions. To reiterate: I think the actual work is fine, but they could be a bit more honest about how it is \"proven\" and the level of rigour in the main paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper lays out a framework for theoretical understanding of deep neural networks that incorporates the effect of finite width. This is a  problem that has received a lot of attention since the original NTK/infinite width limits came out, and as far as I know, the approach here is novel and powerful. The paper is well written for the most part although familiarity with the infinite width DMFTs is assumed. \n\nOverall, the fact that the theory and the simulations agree very well is quite impressive, and I think the ideas in the paper are quite ambitious because they can be used for almost any kind of question one might have about the theoretical evolution of the DNN. This paper has the potential to be the basis for future work which uses the theory developed here to investigate questions about how DNNs learn. ", "weaknesses": "The main weakness of the paper is that its a bit spread thin at times: both the theory and a few different applications are covered, but it seems like the authors were trying to make it all fit and I would have liked more detail in a few spots. This is largely due to the page limit of the submission. I personally would have found it to be a stronger paper if a single really clear example was presented in a lot of detail. (Although again, I completely understand that this is largely pressure from the conference format to try and do a lot of stuff)  \n\nThe other main \"weakness\" of the paper (which is strictly speaking a limitation of the audience of the paper) is that to understand it, you need to be familiar with the previous DMFT on which this paper is built. The authors include a very short section called \"Review of Dynamical Mean Field Theory\" citing [9],[46] as a review, but this section is extremely sparse for actually understanding what is going on. I essentially had to read [9] in its entirety first to understand what was going on in this paper. (Also the reference [46] could not be found since only authors and title are given...where would one find this reference?) In my view, this weakness could be mitigated by just being more honest with the reader up front about this...for example [9],[46] should be cited at point 1 in the list of contributions to make it more clear the dependence and what is/isn't actually explained in this paper. \n\nAnother (related) \"weakness\" is that the paper relies quite heavily on physics technology and jargon to reach its conclusions. The fact that the results are so heavily entrenched in physics jargon like \"order parameters\" or \"propagator\" makes this paper less likely to have a broad impact on the deep learning community. The authors would add a lot of value to the work by attempting to make a \"translation guide\" to help people who don't have the same physics background understand what is going on in more detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for appreciating the strengths of our approach and its applicability to wide DNN dynamics. Below we provide some responses to the weaknesses, questions and limitations. \n\n**Responses to Weaknesses**\n\nWe agree that the paper is spread a bit thin at times. Based on the detailed comments and questions of this reviewer, we have added more detailed definitions of the terms which appear in the main text equations. Given an additional page for a final draft, we will expand our comments on the setup of the calculation, the statement of the main results and their implications. \n\nBased on the comments of all reviewers, we will add a more self-contained derivation of DMFT in the Appendix which will introduce the main concepts and derive the action $S$ which plays such a central role in this paper. We also will add a table, as the first reviewer suggested which translates the physics jargon for the objects into more ML theory friendly language. \n\nWe agree with the reviewer that we do not provide formally rigorous proofs of our results, but derive them at a physical level of rigor. We will add a sentence in the disussion acknowledging this limitation and will leave open for future works to provide rigorous proofs of these dynamical expansions. \n\n**Responses to Questions**\n\n1. This is a good question/suggestion. Indeed the raw covariance of the order parameters is size $\\mathcal{O}(1/N)$. We define the propagator $\\Sigma \\sim N \\text{Cov}(q) \\sim \\mathcal{O}(1)$ entries to be $N$ times this covariance so that it behaves as an $\\mathcal{O}(1)$ quantity. Thus equation 7 can be solved for $\\Sigma^{\\Delta} \\sim \\mathcal{O}(1)$ once for all possible $N$. When we want to compare to simulations of a finite width $N$ network, we could multiply the empirically observed covariance by $N$ and compare to $\\Sigma$. As the reviewer points out, the covariance $N \\text{Cov}(q) \\sim \\Sigma + \\mathcal{O}(N^{-1})$ is correct up to a subleading term. This is established in Appendix 3.1 and 3.3, specifically equations 15 and 17 (disregard the typo in the sentence above equation 17). We will comment on this near equation 3. \n2. When we submitted the paper, we did not have simulations showing that the covariance predicted $\\frac{1}{N}\\Sigma$ is correct up to $\\mathcal{O}(N^{-2})$. To accurately measure the deviations between the order parameter covariance and the theoretically predicted propagator $\\Sigma$, we add a simulation in the attached PDF which shows that this rate is accurate. Estimating this error rate requires a very large number of neural networks (many more than to estimate $\\text{Cov}(q)$) so we focus on the variance of ReLU feature and gradient kernels at initialization, where we can exactly compute $q_\\infty$ and $\\Sigma$. \n3.  On the learning rate scaling with $N$ question, we have a few comments. First, if the reviewer is asking about the learning rate in the gradient flow $\\frac{d}{dt}\\theta = N \\gamma^2 \\nabla \\mathcal{L}$, if the factor of $N$ is removed it will take a time $t \\sim \\mathcal{O}(N)$ to make progress on the training loss, but in gradient flow this is just a rescaling of the time axis in all our plots (so our theory still applies). What happens if other parameters (such as feature learning rate $\\gamma$ or discrete time step size $\\eta$) depend on $N$? First, what if $\\gamma$ is depenent on width $N$? In the classic NTK parameterization where $\\gamma = \\frac{1}{\\sqrt N}$, we can still compute a predicted dynamics for the kernels and fluctuations. First, each finite width $N$ network has a corresponding infinite width (feature learning) network with order parameters $q_{\\infty}(\\gamma = \\frac{1}{\\sqrt N})$. Second, the propagator can be computed as $\\Sigma(\\gamma= \\frac{1}{\\sqrt N})$, in both cases evaluating the dynamics at feature learning velocity $\\gamma$ which depends on the width. The predicted covariance would be $\\text{Cov}(q) \\sim \\frac{1}{N} \\Sigma(\\gamma =\\frac{1}{\\sqrt N})$. This makes understanding the effect of width possible but more complicated in NTK parameterization since decreasing $N$ has two effects (increase feature learning and fluctuation variance). In mean field parameterization/ $\\mu P$, the feature learning parameter $\\gamma$ is fixed with width so that the same $q_{\\infty}(\\gamma)$ and $\\Sigma(\\gamma)$ can be used to estimate finite width effects for different widths $N$ simply by multiplying by $1/N$. What if the raw learning rate $\\eta$ in discrete time is scaled differently with $N$? This will also change the dynamics, but it will lead to a badly behaved $q_\\infty$. For instance, if $\\eta$ is rescaled by $1/\\sqrt{N}$, the neural network will not fit the data in finite time at infinite width and if $\\eta$ is multiplied by a higher power of width $N$, the dynamics will become unstable in discrete time. We add a comment on this in section 4. \n\n**Minor Comments**\n\n1. We will be sure to use $\\equiv$ when we define new terms (such as in Eq 2).\n2. Yes, there is a difference between $K$ and $K_\\infty$. The $K$ should be thought of as the random finite width NTK while $K_\\infty$ is the infinite width kernel. We will add a sentence clarifying this in this section.  \n3. In this section by $\\text{Cov}_0$, we meant the leading order covariance (neglecting $1/N^2$ and smaller terms). We will remove this notation and simply explain that we are computing the asymptotic covariance. \n4. We will add an explicit definition of $K$ and $K_\\star$ which represent the train-train NTK and the train-test NTK respectively.\n5. We will explicitly define $\\Delta_y$ and $\\Delta_{\\perp}$. In words, $\\Delta_y$ is the projection of the vector $\\mathbf \\Delta$ on the label direction $\\mathbf y$ and $\\Delta_{\\perp}$ is the projection on orthogonal directions. \n6. We thank the reviewer for catching the issues with the references we will go through them and add the most recent (journal versions) citations for the papers. \n"}, {"review_id": "5YNCwK7sPJ", "reviewer": "Reviewer_fkFE", "summary": "The authors study a finite width correct to the Dynamical Mean Field Theory (DMFT) of finite depth neural networks in the feature learning regime. While I will be the first to admit that I am not an expert on the DMFT calculations, the authors did produce very convincing simulations capturing interesting properties of finite size networks. In particular, the edge of stability behaviour at large learning rates seems to be modelled by the finite width correction. \n\nI hope the authors can help me clarify some questions regarding the implementation of the solver, after which I would be happy to raise my score to accept. ", "questions": "I would like the authors to provide some intuitions about how the DMFT equations are being solved numerically. I understand this may be the contents of a previous paper (Bordelon and Pehlevan, 2022), but I am having trouble understanding this part. \n\n1. Based on my understanding, all the elements of $q$ are processes index by time $t$. So what does it mean to compute elements of $q$ or $\\Sigma$? Are the authors time discretizing these processes? \n2. How should I understand the role of the propagator $\\Sigma$, and how the authors solve for this object first independently of the elements of $q$? \n3. How should I understand the saddle point solver in Algorithm 2? Here I am just asking for what is happening within the algorithm, as I find reading the pseudocode quite difficult to understand. \n\nI would be happy to continue this discussion during the rebuttal period, and once again I would raise my score once I find these questions addressed. ", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors develop a strong theory of finite width neural networks capable of making accurate predictions. ", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank their reviewer for their support and good questions. We hope to make our methods more understandable and self-contained in the paper. Below we provide some more explanations about how we solve our self-consistent equations. \n\n### Response to Questions \n\n1. This is a great question that we will address more in the work (Appendix E). First, the $\\mathbf q$ vector represents the collection of the order parameters (kernels $\\Phi^\\ell_{\\mu\\nu}(t,s)$, predictions $f_\\mu(t)$, response functions $A^\\ell_{\\mu\\nu}(t,s)$, etc) at *all possible times through training*. Theoretically, for gradient flow this is an uncountably infinite set. For discrete time gradient descent, this is a countable set. In a practical numerical solution, we need to discretize time to compute it. For subtle theoretical reasons, the path integral is also initially derived using a discretization of all the dynamics in time (with the Ito convention). After integration over the initial weights, we can then take a continuous time limit to arrive at the action $S$ provided in Appendix D. Concretely we will add the following text to Appendix E \n    *For nonlinear deep networks, we build on the Monte-carlo approach developed by Bordelon & Pehlevan 2022 which computes the saddle point equations for all order parameters $q_\\infty$. For a practical numerical algorithm, we discretize the time steps so that we store finite matrices (such as $\\Phi^\\ell_{\\mu\\nu}(t,s)$) with layer, sample, and time indices.  Our approach allows us to estimate the entries of the action's Hessian $\\nabla^2 S$ and ultimately invert it to obtain the propagator $\\Sigma = \\left[ - \\nabla^2 S(q_\\infty) \\right]^{-1}$. To do this, we need to use sampling to estimate the fourth feature moments $\\kappa$ and the sensitivity blocks $D$. We evaluate these averages over the stochastic process defined by $q_\\infty$.*\nFor a detailed set of steps to compute $q_\\infty$ and $\\Sigma$, see the third bullet point below.\n3. The propagator $\\Sigma$ can be thought of as the covariance of the order parameters near infinite width. Its name in physics derives from the fact that fluctuations in the order parameters such as the kernel $\\Phi^\\ell$ at time $t$ can propagate noise to other order parameters (such as $\\Phi^{\\ell+1}$) at time $t' > t$. When we say we compute $q_\\infty$ it means that we are solving the saddle point equations, which can be solved using the methods of Bordelon & Pehlevan 2022. **To compute $\\Sigma$, we need to have already first solved for $q_\\infty$**. We evaluate the second derivative of $S$ at this value of the order parameters $q_\\infty$. This involves using similar methods (either monte carlo for nonlinear networks or closed form expressions for linear networks.) Again, in practice, we discretize time in a numerical algorithm to evaluate these expressions. (See below for more detail.)\n4. The Algorithm pseudocode can be summarized by the following high level instructions\n* Step 1: First solve the infinite width DMFT equations for $q_\\infty$ which include the prediction error dynamics $\\Delta_\\mu(t)$, the feature kernels $\\Phi^\\ell_{\\mu\\nu}(t,s)$, gradient kernels $G^\\ell_{\\mu\\nu}(t,s)$. This step corresponds to algorithm in Bordelon & Pehlevan 22 and defines the dynamics one would expect at infinite width.\n* Step 2: Compute the entries of the Hessian of $S$ evaluated at the $q_\\infty$ computed in the first step. Some of these entries look like fourth cumulants of features like $\\kappa = \\left< \\phi(h)^4 \\right> - \\left< \\phi(h)^2 \\right>^2$ and some of them measure sensitivity of one order parameter to a perturbation in another order parameter $D^{\\Phi^\\ell} = \\frac{\\partial}{\\partial \\Phi^{\\ell-1}} \\left< \\phi(h^\\ell)^2 \\right>$. The averages $\\left< \\right>$ used to calculate $\\kappa$ and $D^{\\Phi^\\ell}$ should be performed over the infinite width stochastic processes for preactivations $h^\\ell$ which are defined in equation (19).  \n* Step 3: After populating the entries of the block matrix for the Hesssian $\\nabla^2 S$, we then calculate the propagator $\\Sigma$ with a matrix inversion. Since we discretized time, this is a finite dimensional matrix. \n\nThe above text will also be added to the Appendix of our work.\n\n4. Now we will give a detailed set of intuitions about how the infinite width limit for $q_\\infty$ is solved (step 1 above). This corresponds to the algorithm of Bordelon & Pehlevan 2022 to solve the saddle point equations $\\frac{\\partial}{\\partial q} S(q)|_{q_\\infty} = 0$.\n* Step 1: Start with a guess for the kernels $\\Phi^\\ell_{\\mu\\nu}(t,s), G^\\ell_{\\mu\\nu}(t,s)$ and for the predictions through time $f_\\mu(t)$. We usually use the lazy limit as an initial guess.\n* Step 2: Sample gaussian sources $u^\\ell_\\mu(t) \\sim \\mathcal{GP}(0,\\Phi^{\\ell-1})$ and $r^\\ell_\\mu(t) \\sim \\mathcal{GP}(0,G^{\\ell+1})$ based on the current estimated covariances $\\Phi^\\ell$ and $G^\\ell$ respectively. \n* Step 3: For each sample, solve integral equations for $h^\\ell_\\mu(t)$ and $z^\\ell_\\mu(t)$. \nThese will be samples from the single site distribution for $h^\\ell, z^\\ell$. In a discretization, the integrals will be replaced with sums.\n* Step 4: Average over the Monte-carlo samples to produce a new estimate of the kernels, for instance $\\Phi^\\ell_{\\mu\\nu}(t,s) = \\left< \\phi(h^\\ell_\\mu(t)) \\phi(h^\\ell_\\nu(s)) \\right>$. A similar procedure is performed for $G^\\ell$ and the response functions $A,B$.\n* Step 5: Compute the NTK estimate $K(t) = \\sum_\\ell G^{\\ell+1}(t,t) \\Phi^\\ell(t,t)$ and then integrate prediction dynamics from the dynamics of the NTK $\\frac{d}{dt} f_\\mu(t) = \\sum_\\nu K_{\\mu\\nu}(t) \\Delta_\\nu(t)$. \n* Repeat steps 2-5 until the order parameters converge. \n\nPlease let us know if these explanations helped clarify the numerical methods and if there are any remaining questions or concerns.\n\nAs we mentioned in the global response, we are also going to add a more self-contained derivation of the DMFT action (similar to Bordelon & Pehlevan 2022). "}]

Author rebuttal:
We thank all of the reviewers for their detailed reading and comments. We appreciate the general support for this paper and the comments on the paper's strengths and weaknesses. Many concerns were shared among reviewers which has caused us to make the following updates to the paper
1. We spend more space defining all of the mathematical terms which appear in our paper and relate the physics terminology with more standard machine learning terminology. A table summarizing the map between our terminology and more standard terms is provided in the PDF. This will hopefully reduce the obscurity of our writing. 
2. We provided a new experiment (for ReLU kernels at initialization) that the asymptotic covariance $\frac{1}{N} \Sigma$ predicted by our theory is accurate up to order $1/N^2$ (see attached PDF). This provides additional support for our approach, as we can characterize the error of our theoretical covariance predictions (though empirical estimates of the covariance error require simulating a very large number of networks).  
3. We are now spending more space to define each of the terms (like $K_\infty, K, K_\star, f, f_\star, \Delta_y, \Delta_{\perp}$, etc) which appear in our equations. This will hopefully make the paper easier to parse and allow the reader to more easily interpret our results.
4. We will provide a self-contained derivation of the DMFT action $S$ in the Appendix so that we do not force the reviewer to read [9] Bordelon & Pehlevan 2022. We will also show how one can use $S$ to find the saddle point equations for $q_\infty$. 
5. We expand in the Appendix our section which explains how to numerically solve the self-consistent equations for the saddle point $q_\infty$ and the propagator $\Sigma$ (see response to reviewer fkFE). 
6. We will acknowledge in the limitations section that our paper operates at the level of rigor of a physics calculation rather than a fully rigorous proof which would need several additional assumptions to make the expansion properly defined. 
7. We will fix all of the issues with the citations to make sure they are up to date and contain the appropriate journal and update paper citations to Arxiv preprints which are now published. 
8. We will clarify that the CIFAR-10 experiments are purely experimental to see whether the dynamics in a more realistic setting has qualitatively similar dynamics to our solveable examples. 

Overall, we aim to make the paper more understandable and readable. With the additional page, we will be able to expand the writing and exposition in the main text. 

Meta-review: This paper develops a new DMFT (Dynamical Mean Field Theory) theory that admits a finite width approximation around the infinite width limit while the existing work mainly focused on the infinite width limit. This enables us to describe the behavior of the network precisely such as bias, variance, and training rates. Overall, this paper provides a solid theoretical contribution on the DMFT analysis and the theory yields a impressively precise estimate of the actual behavior. All the reviewers are positive on this paper. I also think that this is a good paper. Hence I am happy to recommend this paper to accept.
Decision: Accept (spotlight)