SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2: We show the accuracy of the lazy-limit ODE in equation (where) compared to a two-layer finite width  $N = 100$  ReLU network trained with  $\gamma = 0.05$  on  $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  $\Delta_k(t) = \Delta(t) \cdot \phi_k$ . These are not ordered simply by magnitude of eigenvalues or the target projections  $y_k = \pmb{y} \cdot \pmb{\phi}_k$ , but rather depend on all eigenvalue gaps  $\lambda_k - \lambda_\ell$  for  $k \neq \ell$  and also the  $\kappa_{k\ell nm}$  tensor. (c) The total variance for all training points  $N \sum_{\mu} \mathrm{Var}\Delta_{\mu}(t) = N \sum_{k} \mathrm{Var}\Delta_{k}(t)$  is also well predicted by the DMFT propagator equations.
(b) Top Eigenmodes Variance
(c) Total Train Variance
An example verifying these dynamics is provided in Figure 2. In the case where the target is an eigenfunction  $y = \psi_{k^*}$ , the covariance has the form  $\Sigma_{k\ell}^{\Delta} (t, s) = \kappa_{k\ell k^* k^*} \frac{\exp(-\lambda_{k^*}(t + s))}{(\lambda_k - \lambda_{k^*})(\lambda_\ell - \lambda_{k^*})}$ . If the kernel is rank one with eigenvalue  $\lambda$ , then the dynamics have the simple form  $\Sigma^{\Delta}(t, s) = \kappa y^{2} t s e^{-\lambda(t + s)}$ . We note that similar terms appear in the prediction dynamics obtained by truncating the Neural Tangent Hierarchy [10, 11], however those dynamics concerned small feature learning
(a) Deviation from DMFT vs  $N$
(b) Average NTK vs DMFT
(c) Average test prediction
(d) Train error variance
(e) Test Prediction Variance
(f) Kernel Variance Dynamics
Figure 3: An ensemble of  $E = 1000$  two layer  $N = 256$  tanh networks trained on a single training point. Dashed black lines are DMFT predictions. (a) The square deviation from the infinite width DMFT scales as  $\mathcal{O}(1 / N)$  for all order parameters. (b) The ensemble average NTK  $\langle K(t)\rangle$  (solid colors) and (c) ensemble average test point predictions  $f_{\star}(t)$  for a point with  $\frac{\boldsymbol{x}\cdot\boldsymbol{x}_{\star}}{D} = 0.5$  closely follow the infinite width predictions (dashed black). (d) The variance (estimated over the ensemble) of the train error  $\Delta (t) = y - f(t)$  initially increases and then decreases as the training point is fit. (e) The variance of  $f_{\star}$  increases with time but decreases with  $\gamma$ . (f) The variance of the NTK during feature learning experiences a transient increase before decreasing to a lower value.
corrections rather than from initialization variance (App. H.1). Corrections to the mean  $\langle \Delta \rangle$  are analyzed in App. H.2. We find that the variance and mean correction dynamics involves non-trivial coupling across eigendirections with a mixture of exponentials with timescales  $\{\lambda_k^{-1}\}$ .
6 Rich Regime in Two-Layer Networks
In this section, we analyze how feature learning alters the variance through training. We show a 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  $\Delta$ . To gain insight, we first study a simplified setting where the data distribution is a single training example  $\pmb{x}$  and single test point  $\pmb{x}_{\star}$  in a two layer network. We will track  $\Delta(t) = y - f(\pmb{x}, t)$  and the test prediction  $f_{\star}(t) = f(\pmb{x}_{\star}, t)$ . 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_{\star}(t)$ . In this case, all order parameters can be viewed as scalar functions of a single time index (unlike the deep network case, see App. E).
Proposition 4 Computing the Hessian of the DMFT action and inverting (App. I), we obtain the following covariance for  $\mathbf{q}_1 = \text{Vec}\{\Delta(t), f_\star(t), K(t), K_\star(t)\}_{t \in \mathbb{R}_+}$
$$
\boldsymbol {\Sigma} _ {q _ {1}} = \left[ \begin{array}{c c c c} \mathbf {I} + \boldsymbol {\Theta} _ {K} & 0 & \boldsymbol {\Theta} _ {\Delta} & 0 \\ - \boldsymbol {\Theta} _ {K _ {\star}} & \mathbf {I} & 0 & - \boldsymbol {\Theta} _ {\Delta} \\ - D & 0 & \mathbf {I} & 0 \\ - D _ {\star} & 0 & 0 & \mathbf {I} \end{array} \right] ^ {- 1} \left[ \begin{array}{c c c c} 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 0 & 0 & \boldsymbol {\kappa} & \boldsymbol {\kappa} _ {\star} ^ {\top} \\ 0 & 0 & \boldsymbol {\kappa} _ {\star} & \boldsymbol {\kappa} _ {\star \star} \end{array} \right] \left[ \begin{array}{c c c c} \mathbf {I} + \boldsymbol {\Theta} _ {K} & 0 & \boldsymbol {\Theta} _ {\Delta} & 0 \\ - \boldsymbol {\Theta} _ {K _ {\star}} & \mathbf {I} & 0 & - \boldsymbol {\Theta} _ {\Delta} \\ - D & 0 & \mathbf {I} & 1 \\ - D _ {\star} & 0 & 0 & \mathbf {I} \end{array} \right] ^ {- 1 ^ {\top}},
$$
where  $[\Theta_K](t,s) = \Theta (t - s)K(s),[\Theta_\Delta ](t,s) = \Theta (t - s)\Delta (s)$  are Heaviside step functions and  $D(t,s) = \left\langle \frac{\partial}{\partial\Delta(s)} (\phi (h(t))^2 +g(t)^2)\right\rangle$  and  $D_{\star}(t,s) = \left\langle \frac{\partial}{\partial\Delta(s)} (\phi (h(t))\phi (h_{\star}(t)) + g(t)g_{\star}(t))\right\rangle$
(a) Offline prediction Variance
(b) Offline leading correction breaks down for large  $P$
(c) Large  $\gamma$  reduces variance
(d) Large  $\gamma$  well captured by DMFT First Order Correction
(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  $P$  becomes comparable to  $N$ . (c)-(d) Richer training reduces variance. (e)-(f) Online learning in a depth 2 linear network has identical dynamics and finite width fluctuations, but with predictor variance  $\sim D / N$  for input dimension  $D$  (Appendix K).
quantify sensitivity of the kernel to perturbations in the error signal  $\Delta(s)$ . Lastly  $\kappa$  and  $\kappa_{\star}$  are the uncoupled variances of  $K(t)$  and  $K_{\star \star}(t)$  and  $\kappa_{\star}$  is the uncoupled covariance of  $K(t), K_{\star}(t)$ .
In Fig. 3, we plot the resulting theory (diagonal blocks of  $\Sigma_{q_1}$  from Equation 8) for two layer neural networks. As predicted by theory, all average squared deviations from the infinite width DMFT scale as  $\mathcal{O}(N^{-1})$ . Similarly, the average kernels  $\langle K\rangle$  and test predictions  $\langle f_{\star}\rangle$  change by a larger amount for larger  $\gamma$  (equation (79)). The experimental variances also match the theory quite accurately. The variance of the train error  $\Delta (t)$  peaks earlier and at a lower value for richer training, but all variances go to zero at late time as the model approaches the interpolation condition  $\Delta = 0$ . As  $\gamma \rightarrow 0$  the curve approaches  $N\operatorname {Var}(\Delta (t))\sim \kappa y^2 t^2 e^{-2t}$ , where  $\kappa$  is the initial NTK variance (see Section 5). 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 [7, 18]. In Fig. A.1, we show that the reduction in variance is not due to a reduction in the uncoupled variance  $\kappa (t,s)$ , which increases in  $\gamma$ . 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  $\{\Delta_{\mu}(t)\}_{\mu = 1}^{P}$  and multiple kernel entries  $K_{\mu \nu}(t)$  (App. E). Each of the errors  $\Delta_{\mu}(t)$  receives a  $\mathcal{O}(N^{-1 / 2})$  fluctuation, the training error  $\sum_{\mu}\left\langle \Delta_{\mu}^{2}\right\rangle$  has an additional variance on the order of  $\mathcal{O}\left(\frac{P}{N}\right)$ . In the case of two-layer linear
(a) Prediction Variance
(b) NTK variance
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  $\langle K(t)\rangle^2$ , 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  $\gamma$ .
(c) Variance by Layer
networks trained on whitened data  $(\frac{1}{D}\pmb{x}_{\mu}\cdot \pmb{x}_{\nu} = \delta_{\mu \nu})$ , the equations for the propagator simplify and one can separately solve for the variance of  $\Delta (t)\in \mathbb{R}^{P}$  along signal direction  $\pmb {y}\in \mathbb{R}^{P}$  and along each of the  $P - 1$  orthogonal directions (App. J). At infinite width, the task-orthogonal component  $\Delta_{\perp}$  vanishes and only the signal dimension  $\Delta_y(t)$  evolves in time with differential equation [9, 46]
$$
\frac {d}{d t} \Delta_ {y} (t) = 2 \sqrt {1 + \gamma^ {2} (y - \Delta_ {y} (t)) ^ {2}} \Delta_ {y} (t), \Delta_ {\perp} (t) = 0. \tag {9}
$$
However, at finite width, both the  $\Delta_y(t)$  and the  $P - 1$  orthogonal variables  $\Delta_{\perp}$  inherit initialization variance, which we represent as  $\Sigma_{\Delta_y}(t,s)$  and  $\Sigma_{\perp}(t,s)$ . In Fig. 4 (a)-(b) we show this approximate solution  $\langle |\Delta (t)|^2\rangle \sim \Delta_y(t)^2 +\frac{2}{N}\Delta_y^1 (t)\Delta_y(t) + \frac{1}{N}\Sigma_{\Delta_y}(t,t) + \frac{(P - 1)}{N}\Sigma_{\perp}(t,t) + \mathcal{O}(N^{-2})$  across varying  $\gamma$  and varying  $P$  (see Appendix J for  $\Sigma_{\Delta_y}$  and  $\Sigma_{\perp}$  formulas). We see that variance of train point predictions  $f_{\mu}(t)$  increases with the total number of points despite the signal of the target vector  $\sum_{\mu}y_{\mu}^{2}$  being fixed. In this model, the bias correction  $\frac{2}{N}\Delta_y^1 (t)\Delta_y(t)$  is always  $\mathcal{O}(1 / N)$  but the variance correction is  $\mathcal{O}(P / N)$ . The fluctuations along the  $P - 1$  orthogonal directions begin to dominate the variance at large  $P$ . Fig. 4 (b) shows that as  $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 as  $\sim D / N$  (Appendix K) as we verify in Figure 4 (e)-(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. E). This leads to Hessian blocks with four time and four sample indices such as  $D_{\mu \nu \alpha \beta}^{\Phi^{\ell}}(t,s,t^{\prime},s^{\prime}) = \frac{\partial}{\partial\Phi_{\alpha\beta}^{\ell - 1}(t^{\prime},s^{\prime})}\left\langle \phi (h_{\mu}^{\ell}(t))\phi (h_{\nu}^{\ell}(s))\right\rangle$ , rendering any numerical calculation challenging. However, in deep linear networks trained on whitened data, we can exploit the symmetry in 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 of results for a depth 4 network is provided in Fig. 5. The variance for the feature kernels  $H^{\ell}$  accumulate finite size variance by layer along the forward pass and the gradient kernels  $G^{\ell}$  accumulate variance on the backward pass. The SNR of the kernels  $\frac{\langle H\rangle^{2}}{N\mathrm{Var}(H)}$  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 in the gradient flow regime, including the catapult [57] and edge of stability (EOS) phenomena [19-21]. In these settings, the kernel undergoes an initial scale
(a) Mean Kernel Dynamics
(b) Error 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  $\gamma$ , the kernel reaches its asymptote before hitting the edge of stability. For large  $\gamma$ , the kernel increases and then oscillates around  $2 / \eta$ . (b)-(c) Remarkably variance due to finite size can reduce during training (for  $\gamma$  smaller and larger than the critical value  $\sim 1 / \eta$ ), showing that infinite width DMFT can be predictive of finite NNs trained with large learning rate.
(c) Kernel Variance
growth before exhibiting either a recovery or a 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  $|\pmb{x}|^2 = D$  and  $y = 1$ . We use learning rate  $\eta$  and feature learning strength  $\gamma$ . The infinite width mean field equations for the prediction  $f_{t}$  and the kernel  $K_{t}$  are (App. M)
$$
f _ {t + 1} = f _ {t} + \eta K _ {t} \Delta_ {t} + \eta^ {2} \gamma^ {2} f _ {t} \Delta_ {t} ^ {2}, K _ {t + 1} = K _ {t} + 4 \eta \gamma^ {2} f _ {t} \Delta_ {t} + \eta^ {2} \gamma^ {2} \Delta_ {t} ^ {2} K _ {t}. \tag {10}
$$
For small  $\eta$ , the equations are well approximated by the gradient flow limit and for small  $\gamma$  corresponds to a discrete time linear model. For large  $\eta \gamma > 1$ , the kernel  $K$  progressively sharpens (increases in scale) until it reaches  $2 / \eta$  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. 6 that the leading order propagator elements decrease even after reaching the EOS threshold, indicating reduced disagreement between finite and infinite width dynamics.
9 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 a vanilla depth-6 ReLU CNN trained on CIFAR-10 (experimental details in App. B, G.2) In Fig. 7, we train an ensemble of  $E = 8$  independently initialized CNNs of each width  $N$ . Wider networks not only have better performance for a single model (solid), but also have lower bias (dashed), measured with ensemble averaging of the logits. Because of faster convergence of wide networks, we observe wider networks have higher variance, but if we plot variance at fixed ensembled training accuracy, wider networks have consistently lower variance (Fig. 7(d)).
We next seek an explanation for why wider networks after ensembling trains at a faster rate. Theoretically, 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.1). Our analysis in App. G.1 suggests that the rate of convergence receives a finite size correction with leading correction  $\mathcal{O}(N^{-1})$  G.2. To test this hypothesis, we fit the ensemble training loss curve to exponential function  $\mathcal{L} \approx C \exp(-R_N t)$  where  $C$  is a constant. We plot the fit  $R_N$  as a function of  $N^{-1}$  result in Fig. 7(e). For large  $N$ , we see the leading behavior is linear in  $N^{-1}$ , but begins to deviate at small  $N$  as a quadratic function of  $N^{-1}$ , suggesting that second order effects become relevant around  $N \lesssim 100$ .
In App. Fig. A.4, we train a smaller subset of CIFAR-10 where we find that  $R_N$  is well approximated by a  $\mathcal{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  $\gamma$  on variance in this task. In App. Fig. A.5, we train  $N = 64$  models with varying  $\gamma$ . Increased  $\gamma$  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  $\gamma$  in the range we considered.
We studied the leading order fluctuations of kernels and predictions in mean field neural networks. Feature learning dynamics can reduce undesirable finite size variance, making finite networks order
(a) Train Loss Across Widths
(b) Test Acc. Across Widths
(c) Alignment Dynamics
(d) Test MSE Ensembling Ratio
(e) Training Rates vs  $N$
(f) Train Acc. Rescaled Time
Figure 7: Depth 6 CNN trained on CIFAR-10 for different widths  $N$  with richness  $\gamma = 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 non-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_N$  of training width  $N$  models as a function of  $N^{-1}$ . We rescale the time axis by  $R_N$  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 of 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  $\Sigma$  does not capture the complete finite size distribution defined in Eq. (3), especially as the sample size becomes comparable to the width. It would be interesting to see if proportional limits of 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(\pmb{q})$  by examining cubic and higher derivatives of  $S$  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 we expect our perturbative expressions to give a precise asymptotic description of finite networks in mean field/  $\mu \mathrm{P}$ , the resulting expressions are not realistically computable in deep networks trained on large dataset size  $P$  for long times  $T$  since the number of Hessian entries scales as  $\mathcal{O}(T^4 P^4)$  and a matrix of this size must be stored in memory and inverted in the general case. Future work could explore solveable special cases such as high dimensional limits.
Code Availability
Code to reproduce the experiments in this paper is provided at https://github.com/ Pehlevan-Group/dmft_fluctuations. Details about numerical methods and computational implementation can be found in Appendices F and N.
CP is supported by NSF Award DMS-2134157, NSF CAREER Award IIS-2239780, and a Sloan Research Fellowship. BB is supported by a 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 at 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.